problem: not subscribing to an address

build.gradle

@@ -60,7 +60,7 @@ dependencies {
     compile "org.springframework.security:spring-security-core:$springVersion"
     compile "org.springframework.security:spring-security-web:$springVersion"
     compile "org.springframework.security:spring-security-config:$springVersion"
-    compile 'io.projectreactor:reactor-core:3.2.9.RELEASE'
+    compile "io.projectreactor:reactor-core:$reactorVersion"
     compile 'io.projectreactor.addons:reactor-extra:3.2.3.RELEASE'
     compile 'io.projectreactor.kotlin:reactor-kotlin-extensions:1.0.0.M1'
     compile 'com.salesforce.servicelibs:reactor-grpc:0.10.0'
@@ -97,6 +97,7 @@ dependencies {
     testCompile 'cglib:cglib-nodep:3.2.12'
     testCompile "org.spockframework:spock-core:$spockVersion"
     testCompile "io.grpc:grpc-testing:${grpcVersion}"
+    testCompile "io.projectreactor:reactor-test:$reactorVersion"
 }
 
 compileKotlin {

gradle.properties

@@ -12,6 +12,7 @@ protobufVersion=3.7.1
 # Core
 springBootVersion=2.1.4.RELEASE
 springVersion=5.1.4.RELEASE
+reactorVersion=3.2.9.RELEASE
 
 # Our Libs
 etherjarVersion=0.7.0-SNAPSHOT

src/main/kotlin/io/emeraldpay/dshackle/Config.kt

@@ -4,13 +4,17 @@ import com.fasterxml.jackson.core.Version
 import com.fasterxml.jackson.databind.DeserializationFeature
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.databind.module.SimpleModule
+import org.springframework.beans.factory.annotation.Qualifier
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
 import org.springframework.scheduling.annotation.EnableAsync
 import org.springframework.scheduling.annotation.EnableScheduling
 import org.springframework.scheduling.annotation.Scheduled
+import reactor.core.scheduler.Scheduler
+import reactor.core.scheduler.Schedulers
 import java.text.SimpleDateFormat
 import java.util.*
+import java.util.concurrent.Executors
 
 @Configuration
 @EnableScheduling
@@ -31,4 +35,9 @@ open class Config {
         return objectMapper
     }
 
+    @Bean @Qualifier("upstreamScheduler")
+    open fun upstreamScheduler(): Scheduler {
+        return Schedulers.fromExecutorService(Executors.newFixedThreadPool(16))
+    }
+
 }

src/main/kotlin/io/emeraldpay/dshackle/rpc/TrackAddress.kt

@@ -17,22 +17,25 @@ import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
 import reactor.core.publisher.TopicProcessor
 import reactor.core.publisher.toFlux
-import reactor.math.sum
+import reactor.core.scheduler.Scheduler
 import java.lang.Exception
 import java.time.Duration
 import java.time.Instant
 import java.util.*
 import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.atomic.AtomicLong
 import javax.annotation.PostConstruct
 
 @Service
 class TrackAddress(
         @Autowired private val upstreams: Upstreams,
-        @Autowired private val availableChains: AvailableChains
+        @Autowired private val availableChains: AvailableChains,
+        @Autowired private val upstreamScheduler: Scheduler
 ) {
 
     private val log = LoggerFactory.getLogger(TrackAddress::class.java)
     private val clients = HashMap<Chain, ConcurrentLinkedQueue<TrackedAddress>>()
+    private val seq = AtomicLong(0)
 
     @PostConstruct
     fun init() {
@@ -40,7 +43,7 @@ class TrackAddress(
             if (!clients.containsKey(chain)) {
                 clients[chain] = ConcurrentLinkedQueue()
                 upstreams.getUpstream(chain)?.getHead()?.let { head ->
-                    head.getFlux().subscribe { verifyAll(chain) }
+                    head.getFlux().subscribe { updateBalancesAll(chain) }
                 }
             }
         }
@@ -60,7 +63,19 @@ class TrackAddress(
         }
     }
 
-    fun initializeSimple(request: BlockchainOuterClass.BalanceRequest): Flux<SimpleAddress> {
+    private fun startTracking(client: TrackedAddress) {
+        clients[client.chain]?.add(client) ?: log.warn("Chain ${client.chain} is not available for tracking")
+    }
+
+    private fun stopTracking(client: TrackedAddress) {
+        clients[client.chain]?.remove(client) ?: log.warn("Chain ${client.chain} is not available for tracking")
+    }
+
+    fun isTracked(chain: Chain, address: Address): Boolean {
+        return clients[chain]?.any { it.address == address } ?: false
+    }
+
+    private fun initializeSimple(request: BlockchainOuterClass.BalanceRequest): Flux<SimpleAddress> {
         val chain = Chain.byId(request.asset.chainValue)
         if (!availableChains.supports(chain)) {
             return Flux.error(Exception("Unsupported chain ${request.asset.chainValue}"))
@@ -81,35 +96,37 @@ class TrackAddress(
         }
     }
 
-    fun initializeSubscription(request: BlockchainOuterClass.BalanceRequest, observer: TopicProcessor<BlockchainOuterClass.AddressBalance>): Flux<TrackedAddress> {
+    private fun initializeSubscription(request: BlockchainOuterClass.BalanceRequest, observer: TopicProcessor<BlockchainOuterClass.AddressBalance>): Flux<TrackedAddress> {
         return initializeSimple(request)
                 .map {
-                    it.asTracked(observer)
+                    it.asTracked(observer, seq.incrementAndGet())
                 }
     }
 
-    fun send(request: BlockchainOuterClass.BalanceRequest, addresses: List<TrackedAddress>): Mono<Long> {
-        val chain = Chain.byId(request.asset.chainValue)
-        return verify(chain, addresses)
-                .map { updated -> notify(updated); 1 }
-                .sum()
-    }
-
     fun subscribe(requestMono: Mono<BlockchainOuterClass.BalanceRequest>): Flux<BlockchainOuterClass.AddressBalance> {
         return requestMono.flatMapMany { request ->
-            val chain = Chain.byId(request.asset.chainValue)
+            val bus = TopicProcessor.create<BlockchainOuterClass.AddressBalance>()
-            val sender = TopicProcessor.create<BlockchainOuterClass.AddressBalance>()
+            initializeSubscription(request, bus)
-            initializeSubscription(request, sender)
+                    .flatMap { tracked ->
-                    .doOnNext { tracked -> clients[chain]?.add(tracked) }
+                        val current = getBalance(tracked).map {
-                    .thenMany(sender)
+                            tracked.withBalance(it)
+                        }.doOnNext {
+                            startTracking(it)
+                        }.map {
+                            buildResponse(it)
+                        }
+                        Flux.merge(current, bus).doFinally { stopTracking(tracked) }
+                    }.doOnError { t ->
+                        log.warn("Failed to process subscription", t)
+                    }
         }
     }
 
     fun getBalance(requestMono: Mono<BlockchainOuterClass.BalanceRequest>): Flux<BlockchainOuterClass.AddressBalance> {
         return requestMono.flatMapMany { request ->
             initializeSimple(request)
-                    .flatMap { getBalance(it) }
+                    .flatMap { a -> getBalance(a).map { a.withBalance(it) } }
-                    .map { process(it) }
+                    .map { buildResponse(it) }
         }
     }
 
@@ -121,33 +138,32 @@ class TrackAddress(
         )
     }
 
-    private fun verifyAll(chain: Chain) {
+    private fun updateBalancesAll(chain: Chain) {
         clients[chain]?.let { all ->
             all.toFlux()
                     .buffer(20)
                     .map { group ->
-                        verify(chain, group).subscribe { updated -> notify(updated) }
+                        updateBalances(chain, group).subscribe { updated -> notify(updated) }
                     }
                     .subscribe()
         }
     }
 
-    fun getBalance(addr: SimpleAddress): Mono<SimpleAddress> {
+    fun getBalance(addr: SimpleAddress): Mono<Wei> {
         val up = upstreams.getUpstream(addr.chain) ?: return Mono.error(Exception("Unsupported chain: ${addr.chain}"))
         return up.getApi()
                 .executeAndConvert(Commands.eth().getBalance(addr.address, BlockTag.LATEST))
                 .timeout(Duration.ofSeconds(15))
-                .map { value ->
-                    addr.withBalance(value)
-                }
     }
 
-    private fun verify(chain: Chain, group: List<TrackedAddress>): Flux<TrackedAddress> {
+    private fun updateBalances(chain: Chain, group: List<TrackedAddress>): Flux<TrackedAddress> {
         val up = upstreams.getUpstream(chain) ?: return Flux.empty<TrackedAddress>()
         return group.toFlux()
+                .parallel(8).runOn(upstreamScheduler)
                 .flatMap { a ->
-                    getBalance(a).map { Update(a, it.balance!!) }
+                    getBalance(a).map { Update(a, it) }
                 }
+                .sequential()
                 .filter {
                     it.addr.balance == null || it.addr.balance != it.value
                 }
@@ -159,27 +175,26 @@ class TrackAddress(
                 }
     }
 
-    private fun process(address: SimpleAddress): BlockchainOuterClass.AddressBalance {
+    private fun buildResponse(address: SimpleAddress): BlockchainOuterClass.AddressBalance {
         return BlockchainOuterClass.AddressBalance.newBuilder()
                 .setBalance(address.balance!!.amount!!.toString(10))
                 .setAsset(Common.Asset.newBuilder()
                         .setChainValue(address.chain.id)
-                        .setCode("ETHER")
+                        .setCode("ETHER"))
-                )
                 .setAddress(Common.SingleAddress.newBuilder().setAddress(address.address.toHex()))
                 .build()
     }
 
     private fun notify(address: TrackedAddress) {
-        address.stream.onNext(process(address))
         address.lastPing = Instant.now()
+        address.stream.onNext(buildResponse(address))
     }
 
     class Update(val addr: TrackedAddress, val value: Wei)
 
     open class SimpleAddress(val chain: Chain, val address: Address, var balance: Wei? = null) {
-        fun asTracked(stream: TopicProcessor<BlockchainOuterClass.AddressBalance>): TrackedAddress {
+        fun asTracked(stream: TopicProcessor<BlockchainOuterClass.AddressBalance>, id: Long): TrackedAddress {
-            return TrackedAddress(chain, stream, address, balance = this.balance)
+            return TrackedAddress(chain, stream, address, balance = this.balance, id = id)
         }
 
         open fun withBalance(balance: Wei) = SimpleAddress(chain, address, balance)
@@ -189,8 +204,13 @@ class TrackAddress(
                          val stream: TopicProcessor<BlockchainOuterClass.AddressBalance>,
                          address: Address,
                          var lastPing: Instant = Instant.now(),
-                         balance: Wei? = null
+                         balance: Wei? = null,
+                         val id: Long
     ): SimpleAddress(chain, address, balance) {
-        override fun withBalance(balance: Wei) = TrackedAddress(chain, stream, address, lastPing, balance);
+        override fun withBalance(balance: Wei) = TrackedAddress(chain, stream, address, lastPing, balance, id)
+
+        override fun equals(other: Any?): Boolean {
+            return other != null && other is TrackedAddress && other.id == id
+        }
     }
 }

src/test/groovy/io/emeraldpay/dshackle/rpc/TrackAddressSpec.groovy

@@ -0,0 +1,129 @@
+package io.emeraldpay.dshackle.rpc
+
+import io.emeraldpay.api.proto.BlockchainOuterClass
+import io.emeraldpay.api.proto.Common
+import io.emeraldpay.dshackle.test.EthereumApiMock
+import io.emeraldpay.dshackle.test.TestingCommons
+import io.emeraldpay.dshackle.upstream.AggregatedUpstreams
+import io.emeraldpay.dshackle.upstream.AvailableChains
+import io.emeraldpay.dshackle.upstream.EthereumHead
+import io.emeraldpay.dshackle.upstream.Upstreams
+import io.emeraldpay.grpc.Chain
+import io.infinitape.etherjar.domain.Address
+import io.infinitape.etherjar.domain.BlockHash
+import io.infinitape.etherjar.rpc.RpcClient
+import io.infinitape.etherjar.rpc.json.BlockJson
+import reactor.core.publisher.Mono
+import reactor.core.publisher.TopicProcessor
+import reactor.core.scheduler.Schedulers
+import reactor.test.StepVerifier
+import spock.lang.Specification
+
+import java.time.Duration
+
+class TrackAddressSpec extends Specification {
+
+    AvailableChains availableChains = new AvailableChains()
+    Upstreams upstreams
+    TrackAddress trackAddress
+
+    def chain = Common.ChainRef.CHAIN_ETHEREUM
+    def address1 = "0xe2c8fa8120d813cd0b5e6add120295bf20cfa09f"
+    def address1Proto = Common.SingleAddress.newBuilder()
+            .setAddress(address1)
+    def etherAsset = Common.Asset.newBuilder()
+            .setChain(chain)
+            .setCode("ETHER")
+
+
+    def setup() {
+        upstreams = Mock(Upstreams)
+        trackAddress = new TrackAddress(upstreams, availableChains, Schedulers.immediate())
+    }
+
+    def start() {
+        trackAddress.init()
+        availableChains.add(Chain.ETHEREUM)
+        availableChains.add(Chain.TESTNET_KOVAN)
+    }
+
+    def "get balance"() {
+        setup:
+        def req = BlockchainOuterClass.BalanceRequest.newBuilder()
+                .setAsset(etherAsset)
+                .setAddress(Common.AnyAddress.newBuilder().setAddressSingle(address1Proto).build())
+                .build()
+        def exp = BlockchainOuterClass.AddressBalance.newBuilder()
+            .setAddress(address1Proto)
+            .setAsset(etherAsset)
+            .setBalance("1234567890")
+            .build()
+
+        def upstreamMock = Mock(AggregatedUpstreams)
+        def apiMock = new EthereumApiMock(Mock(RpcClient), TestingCommons.objectMapper(), Chain.ETHEREUM)
+        apiMock.answer("eth_getBalance", ["0xe2c8fa8120d813cd0b5e6add120295bf20cfa09f", "latest"], "0x499602D2")
+        _ * upstreams.getUpstream(Chain.ETHEREUM) >> upstreamMock
+        _ * upstreamMock.getApi() >> apiMock
+        start()
+        when:
+        def flux = trackAddress.getBalance(Mono.just(req))
+        then:
+        StepVerifier.create(flux)
+            .expectNext(exp)
+            .expectComplete()
+            .verify(Duration.ofSeconds(3))
+        !trackAddress.isTracked(Chain.ETHEREUM, Address.from(address1))
+    }
+
+    def "recheck address after each block"() {
+        setup:
+        def req = BlockchainOuterClass.BalanceRequest.newBuilder()
+                .setAsset(etherAsset)
+                .setAddress(Common.AnyAddress.newBuilder().setAddressSingle(address1Proto).build())
+                .build()
+        def exp1 = BlockchainOuterClass.AddressBalance.newBuilder()
+                .setAddress(address1Proto)
+                .setAsset(etherAsset)
+                .setBalance("1234567890")
+                .build()
+        def exp2 = BlockchainOuterClass.AddressBalance.newBuilder()
+                .setAddress(address1Proto)
+                .setAsset(etherAsset)
+                .setBalance("65432")
+                .build()
+
+        def block2 = new BlockJson().with {
+            it.number = 1
+            it.totalDifficulty = 100
+            it.hash = BlockHash.from("0xa0e65cbc1b52a8ca60562112c6060552d882f16f34a9dba2ccdc05c0a6a27c22")
+            return it
+        }
+
+        def blocksBus = TopicProcessor.create()
+        def upstreamMock = Mock(AggregatedUpstreams)
+        def headMock = Mock(EthereumHead)
+        def apiMock = new EthereumApiMock(Mock(RpcClient), TestingCommons.objectMapper(), Chain.ETHEREUM)
+        apiMock.answerOnce("eth_getBalance", ["0xe2c8fa8120d813cd0b5e6add120295bf20cfa09f", "latest"], "0x499602D2")
+        apiMock.answerOnce("eth_getBalance", ["0xe2c8fa8120d813cd0b5e6add120295bf20cfa09f", "latest"], "0xff98")
+        _ * upstreams.getUpstream(Chain.ETHEREUM) >> upstreamMock
+        _ * upstreamMock.getApi() >> apiMock
+        _ * upstreamMock.getHead() >> headMock
+        _ * headMock.getFlux() >> blocksBus
+        start()
+        when:
+        def flux = trackAddress.subscribe(Mono.just(req))
+        then:
+        StepVerifier.create(flux)
+                .expectNext(exp1)
+                .then {
+                    assert trackAddress.isTracked(Chain.ETHEREUM, Address.from(address1))
+                }
+                .then {
+                    blocksBus.onNext(block2)
+                }
+                .expectNext(exp2)
+                .thenCancel()
+                .verify(Duration.ofSeconds(3))
+        !trackAddress.isTracked(Chain.ETHEREUM, Address.from(address1))
+    }
+}

src/test/groovy/io/emeraldpay/dshackle/test/EthereumApiMock.groovy

@@ -25,8 +25,12 @@ class EthereumApiMock extends EthereumApi {
         this.objectMapper = objectMapper
     }
 
-    EthereumApiMock answer(@NotNull String method, List<Object> params, Object result) {
+    EthereumApiMock answerOnce(@NotNull String method, List<Object> params, Object result) {
-        predefined << new PredefinedResponse(method: method, params: params, result: result)
+        return answer(method, params, result, 1)
+    }
+
+    EthereumApiMock answer(@NotNull String method, List<Object> params, Object result, Integer limit = null) {
+        predefined << new PredefinedResponse(method: method, params: params, result: result, limit: limit)
         return this
     }
 
@@ -40,6 +44,8 @@ class EthereumApiMock extends EthereumApi {
             log.error("Method ${method} with ${params} is not mocked")
             json.error = new RpcResponseError(-32601, "Method ${method} with ${params} is not mocked")
         }
+        predefined.onCalled()
+        predefined.print()
         return Mono.just(objectMapper.writeValueAsBytes(json))
     }
 
@@ -61,8 +67,14 @@ class EthereumApiMock extends EthereumApi {
         String method
         List params
         Object result
+        Integer limit
 
         boolean isSame(int id, String method, List<?> params) {
+            if (limit != null) {
+                if (limit <= 0) {
+                    return false
+                }
+            }
             if (method != this.method) {
                 return false
             }
@@ -71,5 +83,15 @@ class EthereumApiMock extends EthereumApi {
             }
             return this.params == params
         }
+
+        void onCalled() {
+            if (limit != null) {
+                limit--
+            }
+        }
+
+        void print() {
+            println "Execute API: $method ${params ? params : '_'} >> $result"
+        }
     }
 }