Merge pull request #93 from p2p-org/subscribe_node_status

subscribe node status

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

@@ -47,6 +47,7 @@ class BlockchainRpc(
     private val describe: Describe,
     private val subscribeStatus: SubscribeStatus,
     private val estimateFee: EstimateFee,
+    private val subscribeNodeStatus: SubscribeNodeStatus,
     @Qualifier("rpcScheduler")
     private val scheduler: Scheduler
 ) : ReactorBlockchainGrpc.BlockchainImplBase() {
@@ -220,6 +221,10 @@ class BlockchainRpc(
             .doOnError { failMetric.increment() }
     }
 
+    override fun subscribeNodeStatus(request: Mono<BlockchainOuterClass.SubscribeNodeStatusRequest>): Flux<BlockchainOuterClass.NodeStatusResponse> {
+        return subscribeNodeStatus.subscribe(request).subscribeOn(scheduler).doOnError { failMetric.increment() }
+    }
+
     class RequestMetrics(val chain: Chain) {
         val nativeCallMetric = Counter.builder("request.grpc.request")
             .tag("type", "nativeCall")

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

@@ -0,0 +1,162 @@
+package io.emeraldpay.dshackle.rpc
+
+import io.emeraldpay.api.proto.BlockchainOuterClass
+import io.emeraldpay.api.proto.BlockchainOuterClass.NodeDescription
+import io.emeraldpay.api.proto.BlockchainOuterClass.NodeStatus
+import io.emeraldpay.api.proto.BlockchainOuterClass.NodeStatusResponse
+import io.emeraldpay.api.proto.BlockchainOuterClass.SubscribeNodeStatusRequest
+import io.emeraldpay.api.proto.Common
+import io.emeraldpay.dshackle.Chain
+import io.emeraldpay.dshackle.upstream.CurrentMultistreamHolder
+import io.emeraldpay.dshackle.upstream.Upstream
+import io.emeraldpay.dshackle.upstream.UpstreamAvailability
+import org.springframework.stereotype.Service
+import reactor.core.publisher.Flux
+import reactor.core.publisher.Mono
+import reactor.core.publisher.SignalType
+import reactor.core.publisher.Sinks
+import java.time.Duration
+import java.util.concurrent.ConcurrentHashMap
+import java.util.function.Consumer
+
+@Service
+class SubscribeNodeStatus(
+    private val multistreams: CurrentMultistreamHolder
+) {
+
+    companion object {
+        private val RETRY_TIMEOUT = Duration.ofSeconds(10)
+    }
+
+    fun subscribe(req: Mono<SubscribeNodeStatusRequest>): Flux<NodeStatusResponse> =
+        req.flatMapMany { request ->
+            val knownUpstreams = ConcurrentHashMap<String, Boolean>()
+            val duration = Duration.ofMillis(request.timespan)
+            // send known upstreams details immediately
+            val descriptions = Flux.fromIterable(
+                multistreams.all()
+                    .flatMap { multiStream ->
+                        multiStream.getAll().map { up ->
+                            NodeStatusResponse.newBuilder()
+                                .setNodeId(up.getId())
+                                .setDescription(buildDescription(multiStream.chain, up))
+                                .setStatus(buildStatus(up.getStatus(), up.getHead().getCurrentHeight()))
+                                .build()
+                        }
+                    }
+            )
+
+            // subscribe on head/status updates for known upstreams
+            val upstreamUpdates = Flux.merge(
+                multistreams.all()
+                    .flatMap { ms ->
+                        ms.getAll().map { up ->
+                            knownUpstreams[up.getId()] = true
+                            subscribeUpstreamUpdates(ms.chain, up, duration) { r -> knownUpstreams.remove(r) }
+                        }
+                    }
+            )
+
+            // subscribe on head/status updates for just added upstreams
+            val multiStreamUpdates = Flux.merge(
+                multistreams.all()
+                    .map { ms ->
+                        ms.subscribeAddedUpstreams()
+                            .distinctUntilChanged {
+                                it.getId()
+                            }
+                            .filter {
+                                !knownUpstreams.getOrDefault(it.getId(), false)
+                            }
+                            .flatMap {
+                                knownUpstreams[it.getId()] = true
+                                Flux.concat(
+                                    Mono.just(
+                                        NodeStatusResponse.newBuilder()
+                                            .setNodeId(it.getId())
+                                            .setDescription(buildDescription(ms.chain, it))
+                                            .setStatus(buildStatus(it.getStatus(), it.getHead().getCurrentHeight()))
+                                            .build()
+                                    ),
+                                    subscribeUpstreamUpdates(ms.chain, it, duration) { r -> knownUpstreams.remove(r) }
+                                )
+                            }
+                    }
+            )
+
+            Flux.concat(descriptions, Flux.merge(upstreamUpdates, multiStreamUpdates))
+        }
+
+    private fun subscribeUpstreamUpdates(
+        chain: Chain,
+        upstream: Upstream,
+        timespan: Duration,
+        onUnavailable: Consumer<String>
+    ): Flux<NodeStatusResponse> {
+        val retry = Sinks.many().multicast().directBestEffort<Boolean>()
+        val cancel = Sinks.many().multicast().directBestEffort<Boolean>()
+
+        val heads = Mono.just(upstream)
+            .repeatWhen {
+                retry.asFlux()
+            }
+            .sample(RETRY_TIMEOUT)
+            .takeUntilOther(cancel.asFlux()).flatMap { up ->
+                up.getHead().getFlux()
+                    .takeUntilOther(cancel.asFlux())
+                    .map { block ->
+                        NodeStatusResponse.newBuilder()
+                            .setNodeId(up.getId())
+                            .setStatus(buildStatus(up.getStatus(), block.height))
+                            .build()
+                    }.doFinally {
+                        // retry when subscribed head stopped
+                        if (it == SignalType.ON_COMPLETE) {
+                            retry.tryEmitNext(true)
+                        }
+                    }
+            }.sample(timespan)
+
+        val statuses = upstream.observeStatus()
+            .distinctUntilChanged()
+            .takeUntil { it == UpstreamAvailability.UNAVAILABLE }
+            .map {
+                if (it == UpstreamAvailability.UNAVAILABLE) {
+                    onUnavailable.accept(upstream.getId())
+                    // cancel head subscription & reconnections when upstream becomes unavailable
+                    cancel.tryEmitNext(true)
+                }
+                NodeStatusResponse.newBuilder()
+                    .setNodeId(upstream.getId())
+                    .setStatus(buildStatus(it, upstream.getHead().getCurrentHeight()))
+                    .setDescription(buildDescription(chain, upstream))
+                    .build()
+            }.sample(timespan)
+
+        return Flux.merge(heads, statuses)
+    }
+
+    private fun buildDescription(chain: Chain, up: Upstream): NodeDescription.Builder =
+        NodeDescription.newBuilder()
+            .setChain(Common.ChainRef.forNumber(chain.id))
+            .addAllNodeLabels(
+                up.getLabels().map { nodeLabels ->
+                    BlockchainOuterClass.NodeLabels.newBuilder()
+                        .addAllLabels(
+                            nodeLabels.map {
+                                BlockchainOuterClass.Label.newBuilder()
+                                    .setName(it.key)
+                                    .setValue(it.value)
+                                    .build()
+                            }
+                        )
+                        .build()
+                }
+            )
+            .addAllSupportedMethods(up.getMethods().getSupportedMethods())
+
+    private fun buildStatus(status: UpstreamAvailability, height: Long?): NodeStatus.Builder =
+        NodeStatus.newBuilder()
+            .setAvailability(BlockchainOuterClass.AvailabilityEnum.forNumber(status.grpcId))
+            .setCurrentHeight(height ?: 0)
+}

src/main/kotlin/io/emeraldpay/dshackle/upstream/CurrentMultistreamHolder.kt

@@ -45,6 +45,9 @@ open class CurrentMultistreamHolder(
         return chainMapping.getValue(chain).isAvailable()
     }
 
+    override fun all(): List<Multistream> =
+        chainMapping.values.toList()
+
     @PreDestroy
     fun shutdown() {
         log.info("Closing upstream connections...")

src/main/kotlin/io/emeraldpay/dshackle/upstream/Multistream.kt

@@ -39,6 +39,7 @@ import org.springframework.core.annotation.Order
 import reactor.core.Disposable
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
+import reactor.core.publisher.Sinks
 import java.time.Duration
 import java.time.Instant
 import java.util.Locale
@@ -76,6 +77,9 @@ abstract class Multistream(
     private var capabilities: Set<Capability> = emptySet()
     private val removed: MutableMap<String, Upstream> = HashMap()
     private val meters: MutableMap<String, Meter.Id> = HashMap()
+    private val addedUpstreams = Sinks.many()
+        .multicast()
+        .directBestEffort<Upstream>()
 
     init {
         UpstreamAvailability.values().forEach { status ->
@@ -350,6 +354,7 @@ abstract class Multistream(
                         if (!started) {
                             start()
                         }
+                        addedUpstreams.tryEmitNext(event.upstream)
                         log.info("Upstream ${event.upstream.getId()} with chain $chain has been added")
                     }
                 }
@@ -364,6 +369,9 @@ abstract class Multistream(
         return upstreams.any { matcher.matches(it) }
     }
 
+    fun subscribeAddedUpstreams(): Flux<Upstream> =
+        addedUpstreams.asFlux()
+
     // --------------------------------------------------------------------------------------------------------
 
     class UpstreamStatus(val upstream: Upstream, val status: UpstreamAvailability, val ts: Instant = Instant.now())

src/main/kotlin/io/emeraldpay/dshackle/upstream/MultistreamHolder.kt

@@ -25,4 +25,6 @@ interface MultistreamHolder {
     fun getUpstream(chain: Chain): Multistream
     fun getAvailable(): List<Chain>
     fun isAvailable(chain: Chain): Boolean
+
+    fun all(): List<Multistream>
 }

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

@@ -86,6 +86,11 @@ class MultistreamHolderMock implements MultistreamHolder {
         return upstreams.containsKey(chain)
     }
 
+    @Override
+    List<Multistream> all() {
+        return upstreams.values().toList()
+    }
+
     static class EthereumMultistreamMock extends EthereumPosMultiStream {
 
         EthereumCachingReader customReader = null