solution: subscribe to bitcoin address balance

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

@@ -16,6 +16,7 @@ import reactor.util.function.Tuples
 import java.math.BigDecimal
 import java.math.BigInteger
 import java.util.*
+import kotlin.collections.HashMap
 
 @Service
 class TrackBitcoinAddress(
@@ -30,35 +31,42 @@ class TrackBitcoinAddress(
         return BlockchainType.fromBlockchain(chain) == BlockchainType.BITCOIN && upstreams.isAvailable(chain)
     }
 
-    override fun getBalance(req: BlockchainOuterClass.BalanceRequest): Flux<BlockchainOuterClass.AddressBalance> {
+    fun allAddresses(request: BlockchainOuterClass.BalanceRequest): List<String>? {
-        if (!req.hasAddress()) {
+        if (!request.hasAddress()) {
-            return Flux.error(SilentException("Address not provided"))
+            return null
         }
-        val chain = Chain.byId(req.asset.chainValue)
+        return when {
-        val upstream = upstreams.getUpstream(chain)?.castApi(BitcoinApi::class.java)
+            request.address.hasAddressSingle() -> {
-                ?: return Flux.error(SilentException.UnsupportedBlockchain(req.asset.chainValue))
+                listOf(request.address.addressSingle.address)
-        val addressesAll = when {
-            req.address.hasAddressSingle() -> {
-                listOf(req.address.addressSingle.address)
             }
-            req.address.hasAddressMulti() -> {
+            request.address.hasAddressMulti() -> {
-                req.address.addressMulti.addressesList
+                request.address.addressMulti.addressesList
                         .map { addr -> addr.address }
+                        .sorted()
             }
-            else -> {
+            else -> null
-                return Flux.error(SilentException("Unsupported address"))
+        }
-            }
+    }
+
+    fun requestBalances(chain: Chain, api: BitcoinApi, addresses: List<String>): Flux<AddressBalance> {
+        return api.executeAndResult(0, "listunspent", emptyList(), List::class.java)
+                .flatMapMany { unspents ->
+                    val result = getTotal(chain, addresses, unspents)
+                    Flux.fromIterable(result)
+                }
+    }
+
+    override fun getBalance(request: BlockchainOuterClass.BalanceRequest): Flux<BlockchainOuterClass.AddressBalance> {
+        val chain = Chain.byId(request.asset.chainValue)
+        val upstream = upstreams.getUpstream(chain)?.castApi(BitcoinApi::class.java)
+                ?: return Flux.error(SilentException.UnsupportedBlockchain(request.asset.chainValue))
+        val addresses = allAddresses(request) ?: return Flux.error(SilentException("Unsupported address"))
+        if (addresses.isEmpty()) {
+            return Flux.empty()
         }
         val result = upstream.getApi(Selector.empty).flatMapMany { api ->
-            val addresses = addressesAll.sorted()
+            requestBalances(chain, api, addresses)
-            val results = api.executeAndResult(0, "listunspent", emptyList(), List::class.java)
+                    .map(this@TrackBitcoinAddress::buildResponse)
-                    .flatMapMany { unspents ->
-                        val result = getTotal(chain, addresses, unspents)
-                        Flux.fromIterable(result)
-                    }
-            results.map { addr ->
-                buildResponse(addr)
-            }
         }
         return result
     }
@@ -98,23 +106,49 @@ class TrackBitcoinAddress(
 
 
     override fun subscribe(request: BlockchainOuterClass.BalanceRequest): Flux<BlockchainOuterClass.AddressBalance> {
-        return Flux.error(SilentException("Not Implemented"))
+        val chain = Chain.byId(request.asset.chainValue)
+        val upstream = upstreams.getUpstream(chain)?.castApi(BitcoinApi::class.java)
+                ?: return Flux.error(SilentException.UnsupportedBlockchain(request.asset.chainValue))
+        val addresses = allAddresses(request) ?: return Flux.error(SilentException("Unsupported address"))
+        if (addresses.isEmpty()) {
+            return Flux.empty()
+        }
+        val initial = upstream.getApi(Selector.empty).flatMapMany { api ->
+            requestBalances(chain, api, addresses)
+        }
+        val following = upstream.getHead().getFlux()
+                .flatMap { block ->
+                    upstream.getApi(Selector.empty).flatMapMany { api ->
+                        requestBalances(chain, api, addresses)
+                    }
+                }
+        val last = HashMap<Address, BigInteger>()
+        val result = Flux.merge(initial, following)
+                .filter { curr ->
+                    val prev = last[curr.address]
+                    val updated = prev == null || curr.balance != prev
+                    last[curr.address] = curr.balance
+                    updated
+                }
+
+        return result.map(this@TrackBitcoinAddress::buildResponse)
     }
 
     private fun buildResponse(address: AddressBalance): BlockchainOuterClass.AddressBalance {
         return BlockchainOuterClass.AddressBalance.newBuilder()
                 .setBalance(address.balance.toString(10))
                 .setAsset(Common.Asset.newBuilder()
-                        .setChainValue(address.chain.id)
+                        .setChainValue(address.address.chain.id)
                         .setCode("BTC"))
-                .setAddress(Common.SingleAddress.newBuilder().setAddress(address.address))
+                .setAddress(Common.SingleAddress.newBuilder().setAddress(address.address.address))
                 .build()
     }
 
-    open class AddressBalance(val chain: Chain, val address: String, var balance: BigInteger = BigInteger.ZERO) {
+    open class AddressBalance(val address: Address, var balance: BigInteger = BigInteger.ZERO) {
-        open fun withBalance(balance: BigInteger) = AddressBalance(chain, address, balance)
+        constructor(chain: Chain, address: String, balance: BigInteger) : this(Address(chain, address), balance)
-
-        open fun plus(other: AddressBalance) = AddressBalance(chain, address, balance + other.balance)
 
+        fun plus(other: AddressBalance) = AddressBalance(address, balance + other.balance)
     }
+
+    data class Address(val chain: Chain, val address: String)
 }

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

@@ -1,12 +1,29 @@
 package io.emeraldpay.dshackle.rpc
 
+import io.emeraldpay.api.proto.BlockchainOuterClass
+import io.emeraldpay.api.proto.Common
+import io.emeraldpay.dshackle.data.BlockContainer
+import io.emeraldpay.dshackle.data.BlockId
 import io.emeraldpay.dshackle.test.TestingCommons
+import io.emeraldpay.dshackle.upstream.AggregatedUpstream
+import io.emeraldpay.dshackle.upstream.Head
+import io.emeraldpay.dshackle.upstream.Upstream
 import io.emeraldpay.dshackle.upstream.Upstreams
+import io.emeraldpay.dshackle.upstream.bitcoin.BitcoinApi
 import io.emeraldpay.grpc.Chain
+import reactor.core.publisher.Flux
+import reactor.core.publisher.Mono
+import reactor.core.publisher.TopicProcessor
+import reactor.test.StepVerifier
 import spock.lang.Specification
 
+import java.time.Duration
+import java.time.Instant
+
 class TrackBitcoinAddressSpec extends Specification {
 
+    String hash1 = "0xa0e65cbc1b52a8ca60562112c6060552d882f16f34a9dba2ccdc05c0a6a27c22"
+
     def "Correct sum from multiple"() {
         setup:
         def json = this.class.getClassLoader().getResourceAsStream("bitcoin/unspent-one-addr.json")
@@ -17,8 +34,8 @@ class TrackBitcoinAddressSpec extends Specification {
 
         then:
         total.size() == 1
-        total[0].chain == Chain.BITCOIN
+        total[0].address.chain == Chain.BITCOIN
-        total[0].address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
+        total[0].address.address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
         total[0].balance.toString() == "32928461"
     }
 
@@ -32,8 +49,8 @@ class TrackBitcoinAddressSpec extends Specification {
 
         then:
         total.size() == 1
-        total[0].chain == Chain.BITCOIN
+        total[0].address.chain == Chain.BITCOIN
-        total[0].address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
+        total[0].address.address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
         total[0].balance.toString() == "32928461"
     }
 
@@ -43,18 +60,18 @@ class TrackBitcoinAddressSpec extends Specification {
         def unspents = TestingCommons.objectMapper().readValue(json, List)
         TrackBitcoinAddress track = new TrackBitcoinAddress(Stub(Upstreams))
         when:
-        def total = track.getTotal(Chain.BITCOIN, ["1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK", "35hK24tcLEWcgNA4JxpvbkNkoAcDGqQPsP"], unspents).sort { it.address }
+        def total = track.getTotal(Chain.BITCOIN, ["1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK", "35hK24tcLEWcgNA4JxpvbkNkoAcDGqQPsP"], unspents).sort { it.address.address }
 
         then:
         total.size() == 2
         with(total[0]) {
-            chain == Chain.BITCOIN
+            address.chain == Chain.BITCOIN
-            address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
+            address.address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
             balance.toString() == "32928461"
         }
         with(total[1]) {
-            chain == Chain.BITCOIN
+            address.chain == Chain.BITCOIN
-            address == "35hK24tcLEWcgNA4JxpvbkNkoAcDGqQPsP"
+            address.address == "35hK24tcLEWcgNA4JxpvbkNkoAcDGqQPsP"
             balance.toString() == "25550215615737"
         }
     }
@@ -67,8 +84,8 @@ class TrackBitcoinAddressSpec extends Specification {
 
         then:
         total.size() == 1
-        total[0].chain == Chain.BITCOIN
+        total[0].address.chain == Chain.BITCOIN
-        total[0].address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
+        total[0].address.address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
         total[0].balance.toString() == "0"
     }
 
@@ -78,18 +95,157 @@ class TrackBitcoinAddressSpec extends Specification {
         def unspents = TestingCommons.objectMapper().readValue(json, List)
         TrackBitcoinAddress track = new TrackBitcoinAddress(Stub(Upstreams))
         when:
-        def total = track.getTotal(Chain.BITCOIN, ["16rCmCmbuWDhPjWTrpQGaU3EPdZF7MTdUk", "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"], unspents).sort { it.address }
+        def total = track.getTotal(Chain.BITCOIN, ["16rCmCmbuWDhPjWTrpQGaU3EPdZF7MTdUk", "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"], unspents).sort { it.address.address }
 
         then:
         total.size() == 2
         with(total[0]) {
-            address == "16rCmCmbuWDhPjWTrpQGaU3EPdZF7MTdUk"
+            address.address == "16rCmCmbuWDhPjWTrpQGaU3EPdZF7MTdUk"
             balance.toString() == "0"
         }
         with(total[1]) {
-            address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
+            address.address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
             balance.toString() == "32928461"
         }
     }
 
+    def "One address for single provided"() {
+        setup:
+        TrackBitcoinAddress track = new TrackBitcoinAddress(Stub(Upstreams))
+        def req = BlockchainOuterClass.BalanceRequest.newBuilder()
+                .setAddress(
+                        Common.AnyAddress.newBuilder()
+                                .setAddressSingle(
+                                        Common.SingleAddress.newBuilder()
+                                                .setAddress("16rCmCmbuWDhPjWTrpQGaU3EPdZF7MTdUk")
+                                )
+                )
+                .build()
+        when:
+        def act = track.allAddresses(req)
+        then:
+        act == ["16rCmCmbuWDhPjWTrpQGaU3EPdZF7MTdUk"]
+    }
+
+    def "Sorted addresses for multiple provided"() {
+        setup:
+        TrackBitcoinAddress track = new TrackBitcoinAddress(Stub(Upstreams))
+        def req = BlockchainOuterClass.BalanceRequest.newBuilder()
+                .setAddress(
+                        Common.AnyAddress.newBuilder()
+                                .setAddressMulti(
+                                        Common.MultiAddress.newBuilder()
+                                                .addAddresses(Common.SingleAddress.newBuilder().setAddress("16rCmCmbuWDhPjWTrpQGaU3EPdZF7MTdUk"))
+                                                .addAddresses(Common.SingleAddress.newBuilder().setAddress("3BMqADKWoWHPASsUdHvnUL6E1jpZkMnLZz"))
+                                                .addAddresses(Common.SingleAddress.newBuilder().setAddress("1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"))
+                                                .addAddresses(Common.SingleAddress.newBuilder().setAddress("bc1qdthqvt6cllzej7uhdddrltdfsmnt7d0gl5ue5n"))
+                                )
+                )
+                .build()
+        when:
+        def act = track.allAddresses(req)
+        then:
+        act == ["16rCmCmbuWDhPjWTrpQGaU3EPdZF7MTdUk", "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK", "3BMqADKWoWHPASsUdHvnUL6E1jpZkMnLZz", "bc1qdthqvt6cllzej7uhdddrltdfsmnt7d0gl5ue5n"]
+    }
+
+    def "Null for no address provided"() {
+        setup:
+        TrackBitcoinAddress track = new TrackBitcoinAddress(Stub(Upstreams))
+        def req = BlockchainOuterClass.BalanceRequest.newBuilder()
+                .build()
+        when:
+        def act = track.allAddresses(req)
+        then:
+        act == null
+    }
+
+    def "Build proto for common balance"() {
+        setup:
+        TrackBitcoinAddress track = new TrackBitcoinAddress(Stub(Upstreams))
+        def balance = new TrackBitcoinAddress.AddressBalance(Chain.BITCOIN, "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK", BigInteger.valueOf(123456))
+        when:
+        def act = track.buildResponse(balance)
+        then:
+        act.address.address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
+        act.balance == "123456"
+        act.asset.chain.number == Chain.BITCOIN.id
+        act.asset.code == "BTC"
+    }
+
+    def "Build proto for zero balance"() {
+        setup:
+        TrackBitcoinAddress track = new TrackBitcoinAddress(Stub(Upstreams))
+        def balance = new TrackBitcoinAddress.AddressBalance(Chain.BITCOIN, "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK", BigInteger.ZERO)
+        when:
+        def act = track.buildResponse(balance)
+        then:
+        act.address.address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
+        act.balance == "0"
+        act.asset.chain.number == Chain.BITCOIN.id
+        act.asset.code == "BTC"
+    }
+
+    def "Build proto for all bitcoins"() {
+        setup:
+        TrackBitcoinAddress track = new TrackBitcoinAddress(Stub(Upstreams))
+        def balance = new TrackBitcoinAddress.AddressBalance(Chain.BITCOIN, "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK", BigInteger.valueOf(21_000_000).multiply(BigInteger.TEN.pow(8)))
+        when:
+        def act = track.buildResponse(balance)
+        then:
+        act.address.address == "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK"
+        act.balance == "2100000000000000"
+        act.asset.chain.number == Chain.BITCOIN.id
+        act.asset.code == "BTC"
+    }
+
+    def "Get update for a balance"() {
+        setup:
+
+        BitcoinApi api = Mock(BitcoinApi) {
+            2 * executeAndResult(0, "listunspent", [], List) >>> [
+                    Mono.just([]), Mono.just([[address: "1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK", amount: 0.0123]])
+            ]
+        }
+        def blocks = TopicProcessor.create()
+        Head head = Mock(Head) {
+            1 * getFlux() >> Flux.from(blocks)
+        }
+        Upstream upstream
+        upstream = Mock(AggregatedUpstream) {
+            _ * getApi(_) >> Mono.just(api)
+            _ * getHead() >> head
+            _ * castApi(_) >> { return upstream }
+        }
+        Upstreams upstreams = Mock(Upstreams) {
+            _ * getUpstream(Chain.BITCOIN) >> upstream
+        }
+        TrackBitcoinAddress track = new TrackBitcoinAddress(upstreams)
+
+        when:
+        def resp = track.subscribe(BlockchainOuterClass.BalanceRequest.newBuilder()
+                .setAsset(Common.Asset.newBuilder().setChain(Common.ChainRef.CHAIN_BITCOIN))
+                .setAddress(
+                        Common.AnyAddress.newBuilder().setAddressSingle(
+                                Common.SingleAddress.newBuilder().setAddress("1K7xkspJg7DDKNwzXgoRSDCUxiFsRegsSK")
+                        )
+                )
+                .build()
+        ).map {
+            it.balance
+        }
+
+        then:
+        StepVerifier.create(resp)
+                .expectNext("0")
+                .then {
+                    blocks.onNext(new BlockContainer(1L, BlockId.from(hash1), BigInteger.ONE, Instant.now(), false, null, []))
+                }
+                .expectNext("1230000")
+                .then {
+                    blocks.onComplete()
+                }
+                .expectComplete()
+                .verify(Duration.ofSeconds(1))
+
+    }
 }