solution: refactor head streaming to simple flux subscription

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

@@ -3,7 +3,6 @@ package io.emeraldpay.dshackle.rpc
 import com.google.protobuf.ByteString
 import io.emeraldpay.api.proto.BlockchainOuterClass
 import io.emeraldpay.api.proto.Common
-import io.emeraldpay.dshackle.upstream.UpstreamServices
 import io.emeraldpay.dshackle.upstream.Upstreams
 import io.emeraldpay.grpc.Chain
 import io.infinitape.etherjar.domain.TransactionId
@@ -13,11 +12,6 @@ import org.springframework.beans.factory.annotation.Autowired
 import org.springframework.stereotype.Service
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
-import reactor.core.publisher.TopicProcessor
-import reactor.core.publisher.toFlux
-import java.util.concurrent.ConcurrentLinkedQueue
-import javax.annotation.PostConstruct
-import kotlin.collections.HashMap
 
 @Service
 class StreamHead(
@@ -25,84 +19,27 @@ class StreamHead(
 ) {
 
     private val log = LoggerFactory.getLogger(StreamHead::class.java)
-    private val clients = HashMap<Chain, ConcurrentLinkedQueue<TopicProcessor<BlockchainOuterClass.ChainHead>>>()
-
-    @PostConstruct
-    fun init() {
-        upstreams.observeChains().subscribe { chain ->
-            if (clients.containsKey(chain)) {
-                return@subscribe
-            }
-            clients[chain] = ConcurrentLinkedQueue()
-            subscribe(chain)
-        }
-    }
-
-    private fun subscribe(chain: Chain) {
-        upstreams.getUpstream(chain)!!.let { up ->
-            up.getHead()
-                .getFlux()
-                .doOnComplete {
-                    log.info("Closing streams for ${chain.chainCode}")
-                    clients.replace(chain, ConcurrentLinkedQueue())!!.forEach { client ->
-                        try {
-                            client.dispose()
-                        } catch (e: Throwable) {
-                        }
-                    }
-                }
-                .subscribe { block -> onBlock(chain, block) }
-        }
-    }
-
-    private fun onBlock(chain: Chain, block: BlockJson<TransactionId>) {
-        upstreams.getUpstream(chain)?.let { up ->
-            UpstreamServices.onceOk(up).subscribe {avail ->
-                if (avail) {
-                    clients[chain]!!.toFlux()
-                            .subscribe { stream ->
-                                notify(chain, block, stream)
-                            }
-                }
-            }
-        }
-    }
 
     fun add(requestMono: Mono<Common.Chain>): Flux<BlockchainOuterClass.ChainHead> {
         return requestMono.map { request ->
             Chain.byId(request.type.number)
-        }.filter {
-            it != Chain.UNSPECIFIED && clients.containsKey(it)
         }.flatMapMany { chain ->
-            val sender = TopicProcessor.create<BlockchainOuterClass.ChainHead>()
+            val up = upstreams.getUpstream(chain)
-            clients[chain]!!.add(sender)
+                    ?: return@flatMapMany Flux.error<BlockchainOuterClass.ChainHead>(Exception("Unavailable chain: $chain"))
-            notify(chain, sender)
+            up.getHead()
-            sender
+                    .getFlux()
+                    .map { asProto(chain, it) }
         }
     }
 
-    fun notify(chain: Chain, client: TopicProcessor<BlockchainOuterClass.ChainHead>) {
+    fun asProto(chain: Chain, block: BlockJson<TransactionId>): BlockchainOuterClass.ChainHead {
-        val upstream = upstreams.getUpstream(chain) ?: return
+        return BlockchainOuterClass.ChainHead.newBuilder()
-        val head = upstream.getHead().getFlux().next()
-        head.subscribe { block ->
-            UpstreamServices.onceOk(upstream).subscribe { avail ->
-                if (avail) {
-                    notify(chain, block, client)
-                }
-            }
-        }
-    }
-
-    fun notify(chain: Chain, block: BlockJson<TransactionId>, client: TopicProcessor<BlockchainOuterClass.ChainHead>) {
-        val data = BlockchainOuterClass.ChainHead.newBuilder()
                 .setChainValue(chain.id)
                 .setHeight(block.number)
                 .setTimestamp(block.timestamp.time)
                 .setWeight(ByteString.copyFrom(block.totalDifficulty.toByteArray()))
                 .setBlockId(block.hash.toHex().substring(2))
                 .build()
-        client.onNext(data)
     }
 
-
 }

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

@@ -0,0 +1,76 @@
+package io.emeraldpay.dshackle.rpc
+
+import com.google.protobuf.ByteString
+import io.emeraldpay.api.proto.BlockchainOuterClass
+import io.emeraldpay.api.proto.Common
+import io.emeraldpay.dshackle.test.EthereumUpstreamMock
+import io.emeraldpay.dshackle.test.UpstreamsMock
+import io.emeraldpay.dshackle.upstream.EthereumApi
+import io.emeraldpay.dshackle.upstream.Upstream
+import io.emeraldpay.grpc.Chain
+import io.infinitape.etherjar.domain.Address
+import io.infinitape.etherjar.domain.BlockHash
+import io.infinitape.etherjar.domain.TransactionId
+import io.infinitape.etherjar.rpc.json.BlockJson
+import reactor.core.publisher.Mono
+import reactor.test.StepVerifier
+import spock.lang.Specification
+
+import java.time.Duration
+
+class StreamHeadSpec extends Specification {
+
+    def "Errors on unavailable chain"() {
+        setup:
+        def upstreams = new UpstreamsMock(Chain.ETHEREUM, Stub(Upstream))
+        def streamHead = new StreamHead(upstreams)
+        when:
+        def flux = streamHead.add(
+                Mono.just(Common.Chain.newBuilder().setType(Common.ChainRef.CHAIN_ETHEREUM_CLASSIC).build())
+        )
+        then:
+        StepVerifier.create(flux)
+                .expectError()
+                .verify(Duration.ofSeconds(1))
+    }
+
+    def "Subscribes through upstream head"() {
+        setup:
+
+        def blocks = (100..105).collect { i ->
+            return new BlockJson<TransactionId>().with {
+                it.number = i
+                it.hash = BlockHash.from("0xa0e65cbc1b52a8ca60562112c6060552d882f16f34a9dba2ccdc05c0a6a27${i}")
+                it.totalDifficulty = i * 1000
+                it.timestamp = new Date(1566000000000 + i * 10000)
+                return it
+            }
+        }
+
+        def heads = blocks.collect {
+            return BlockchainOuterClass.ChainHead.newBuilder()
+                .setChain(Common.ChainRef.CHAIN_ETHEREUM)
+                .setTimestamp(it.timestamp.time)
+                .setBlockId(it.hash.toHex().substring(2))
+                .setWeight(ByteString.copyFrom(it.totalDifficulty.toByteArray()))
+                .setHeight(it.number)
+                .build()
+        }
+
+        def upstream = new EthereumUpstreamMock(Chain.ETHEREUM, Mock(EthereumApi))
+        def upstreams = new UpstreamsMock(Chain.ETHEREUM, upstream)
+        def streamHead = new StreamHead(upstreams)
+        when:
+        def flux = streamHead.add(
+                Mono.just(Common.Chain.newBuilder().setType(Common.ChainRef.CHAIN_ETHEREUM).build())
+        )
+        then:
+        StepVerifier.create(flux.take(2))
+                .then { upstream.nextBlock(blocks[0]) }
+                .expectNext(heads[0])
+                .then { upstream.nextBlock(blocks[1]) }
+                .expectNext(heads[1])
+                .expectComplete()
+                .verify(Duration.ofSeconds(1))
+    }
+}