Resubscribe to newHeads immediately after reconnection (#189)

* Resubscribe to newHeads immediately after reconnection

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

@@ -0,0 +1,37 @@
+package io.emeraldpay.dshackle.upstream
+
+import io.emeraldpay.dshackle.upstream.ethereum.WsConnection
+import io.emeraldpay.dshackle.upstream.ethereum.WsSubscriptions
+import reactor.core.publisher.Flux
+
+class WebsocketConnectionStatesHandler(
+    private val wsSubscriptions: WsSubscriptions,
+    private val onConnected: () -> Unit
+) {
+    private var connectionId: String? = null
+    private var tryResubscribe = false
+
+    init {
+        wsSubscriptions.connectionInfoFlux()
+            .subscribe {
+                if (it.connectionId == connectionId && it.connectionState == WsConnection.ConnectionState.DISCONNECTED) {
+                    connectionId = null
+                    if (tryResubscribe) {
+                        tryResubscribe = false
+                    }
+                } else if (connectionId == null && it.connectionState == WsConnection.ConnectionState.CONNECTED) {
+                    if (!tryResubscribe) {
+                        tryResubscribe = true
+                        onConnected()
+                    }
+                }
+            }
+    }
+
+    fun subscribe(method: String): Flux<ByteArray> =
+        wsSubscriptions.subscribe(method)
+            .also {
+                connectionId = it.connectionId
+                tryResubscribe = false
+            }.data
+}

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/EthereumWsHead.kt

@@ -23,6 +23,7 @@ import io.emeraldpay.dshackle.data.BlockContainer
 import io.emeraldpay.dshackle.reader.JsonRpcReader
 import io.emeraldpay.dshackle.upstream.BlockValidator
 import io.emeraldpay.dshackle.upstream.Lifecycle
+import io.emeraldpay.dshackle.upstream.WebsocketConnectionStatesHandler
 import io.emeraldpay.dshackle.upstream.forkchoice.ForkChoice
 import io.emeraldpay.dshackle.upstream.rpcclient.JsonRpcRequest
 import io.emeraldpay.dshackle.upstream.rpcclient.JsonRpcResponse
@@ -40,11 +41,12 @@ class EthereumWsHead(
     forkChoice: ForkChoice,
     blockValidator: BlockValidator,
     private val api: JsonRpcReader,
-    private val wsSubscriptions: WsSubscriptions,
+    wsSubscriptions: WsSubscriptions,
     private val skipEnhance: Boolean
 ) : DefaultEthereumHead(upstreamId, forkChoice, blockValidator), Lifecycle {
 
     private var subscription: Disposable? = null
+    private val wsConnectionStatesHandler = WebsocketConnectionStatesHandler(wsSubscriptions, this::onNoHeadUpdates)
 
     override fun isRunning(): Boolean {
         return subscription != null
@@ -67,7 +69,7 @@ class EthereumWsHead(
     }
 
     fun listenNewHeads(): Flux<BlockContainer> {
-        return wsSubscriptions.subscribe("newHeads")
+        return wsConnectionStatesHandler.subscribe("newHeads")
             .map {
                 Global.objectMapper.readValue(it, BlockJson::class.java) as BlockJson<TransactionRefJson>
             }

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/WsConnection.kt

@@ -25,7 +25,18 @@ interface WsConnection : AutoCloseable {
 
     val isConnected: Boolean
 
+    fun connectionId(): String
     fun getSubscribeResponses(): Flux<JsonRpcWsMessage>
     fun callRpc(originalRequest: JsonRpcRequest): Mono<JsonRpcResponse>
     fun connect()
+    fun connectionInfoFlux(): Flux<ConnectionInfo>
+
+    data class ConnectionInfo(
+        val connectionId: String,
+        val connectionState: ConnectionState
+    )
+
+    enum class ConnectionState {
+        CONNECTED, DISCONNECTED
+    }
 }

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/WsConnectionImpl.kt

@@ -18,6 +18,8 @@ package io.emeraldpay.dshackle.upstream.ethereum
 import io.emeraldpay.dshackle.Defaults
 import io.emeraldpay.dshackle.Global
 import io.emeraldpay.dshackle.config.AuthConfig
+import io.emeraldpay.dshackle.upstream.ethereum.WsConnection.ConnectionState.CONNECTED
+import io.emeraldpay.dshackle.upstream.ethereum.WsConnection.ConnectionState.DISCONNECTED
 import io.emeraldpay.dshackle.upstream.rpcclient.JsonRpcError
 import io.emeraldpay.dshackle.upstream.rpcclient.JsonRpcException
 import io.emeraldpay.dshackle.upstream.rpcclient.JsonRpcRequest
@@ -52,6 +54,7 @@ import java.net.URI
 import java.time.Duration
 import java.time.Instant
 import java.util.Base64
+import java.util.UUID
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.Executors
 import java.util.concurrent.ScheduledFuture
@@ -114,17 +117,26 @@ open class WsConnectionImpl(
         .multicast()
         .directBestEffort<Instant>()
 
+    private val connectionInfo = Sinks
+        .many()
+        .multicast()
+        .directBestEffort<WsConnection.ConnectionInfo>()
+
     private val currentRequests = ConcurrentHashMap<Int, Sinks.One<JsonRpcResponse>>()
 
+    private val connId = UUID.randomUUID().toString()
     private val sendIdSeq = AtomicInteger(IDS_START)
     private val sendExecutor = Executors.newSingleThreadExecutor()
     private var keepConnection = true
+    private var firstConnect = true
     private var connection: Disposable? = null
     private val reconnecting = AtomicBoolean(false)
 
     override val isConnected: Boolean
         get() = connection != null && !reconnecting.get()
 
+    override fun connectionId(): String = connId
+
     fun setReconnectIntervalSeconds(value: Long) {
         reconnectBackoff = FixedBackOff(value * 1000, FixedBackOff.UNLIMITED_ATTEMPTS)
         currentBackOff = reconnectBackoff.start()
@@ -135,6 +147,10 @@ open class WsConnectionImpl(
         connectInternal()
     }
 
+    override fun connectionInfoFlux(): Flux<WsConnection.ConnectionInfo> =
+        connectionInfo.asFlux()
+            .distinctUntilChanged { it.connectionState }
+
     private fun tryReconnectLater() {
         if (!keepConnection) {
             return
@@ -143,6 +159,10 @@ open class WsConnectionImpl(
         if (alreadyReconnecting) {
             return
         }
+        connectionInfo.tryEmitNext(WsConnection.ConnectionInfo(connId, DISCONNECTED))
+        if (firstConnect) {
+            firstConnect = false
+        }
 
         // rpcSend is already CANCELLED, since the subscription owned by the previous connection is gone
         // so we need to create a new Sink. Emit Complete is probably useless, and just in case
@@ -184,6 +204,13 @@ open class WsConnectionImpl(
                     tryReconnectLater()
                 }
             }
+            .doOnConnected {
+                if (!firstConnect) {
+                    connectionInfo.tryEmitNext(WsConnection.ConnectionInfo(connId, CONNECTED))
+                } else {
+                    firstConnect = false
+                }
+            }
             .doOnError(
                 { _, t ->
                     log.warn("Failed to connect to $uri. Error: ${t.message}")
@@ -379,6 +406,7 @@ open class WsConnectionImpl(
 
     override fun close() {
         log.info("Closing connection to WebSocket $uri")
+        connectionInfo.tryEmitComplete()
         keepConnection = false
         connection?.dispose()
         connection = null

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/WsConnectionMultiPool.kt

@@ -20,6 +20,9 @@ import io.emeraldpay.dshackle.upstream.DefaultUpstream
 import io.emeraldpay.dshackle.upstream.UpstreamAvailability
 import org.springframework.util.backoff.BackOffExecution
 import org.springframework.util.backoff.ExponentialBackOff
+import reactor.core.Disposable
+import reactor.core.publisher.Flux
+import reactor.core.publisher.Sinks
 import java.time.Duration
 import java.util.concurrent.ScheduledExecutorService
 import java.util.concurrent.TimeUnit
@@ -50,6 +53,8 @@ class WsConnectionMultiPool(
     private var adjustLock = ReentrantReadWriteLock()
     private val index = AtomicInteger(0)
     private var connIndex = 0
+    private val connectionInfo = Sinks.many().multicast().directBestEffort<WsConnection.ConnectionInfo>()
+    private val connectionSubscriptionMap = mutableMapOf<String, Disposable>()
 
     var scheduler: ScheduledExecutorService = Global.control
 
@@ -73,8 +78,13 @@ class WsConnectionMultiPool(
         return next
     }
 
+    override fun connectionInfoFlux(): Flux<WsConnection.ConnectionInfo> =
+        connectionInfo.asFlux()
+
     override fun close() {
         adjustLock.write {
+            connectionSubscriptionMap.values.forEach { it.dispose() }
+            connectionSubscriptionMap.clear()
             current.forEach { it.close() }
             current.clear()
         }
@@ -106,6 +116,9 @@ class WsConnectionMultiPool(
                             }
                         }.also {
                             it.connect()
+                            connectionSubscriptionMap[it.connectionId()] = it.connectionInfoFlux().subscribe { info ->
+                                connectionInfo.emitNext(info) { _, res -> res == Sinks.EmitResult.FAIL_NON_SERIALIZED }
+                            }
                         }
                     )
                     SCHEDULE_GROW
@@ -116,6 +129,7 @@ class WsConnectionMultiPool(
                 current.removeIf {
                     if (!it.isConnected) {
                         // DO NOT FORGET to close the connection, otherwise it would keep reconnecting but unused
+                        connectionSubscriptionMap.remove(it.connectionId())?.dispose()
                         it.close()
                         true
                     } else {

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/WsConnectionPool.kt

@@ -15,7 +15,10 @@
  */
 package io.emeraldpay.dshackle.upstream.ethereum
 
+import reactor.core.publisher.Flux
+
 interface WsConnectionPool : AutoCloseable {
     fun connect()
     fun getConnection(): WsConnection
+    fun connectionInfoFlux(): Flux<WsConnection.ConnectionInfo>
 }

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/WsConnectionSinglePool.kt

@@ -17,6 +17,7 @@ package io.emeraldpay.dshackle.upstream.ethereum
 
 import io.emeraldpay.dshackle.upstream.DefaultUpstream
 import io.emeraldpay.dshackle.upstream.UpstreamAvailability
+import reactor.core.publisher.Flux
 
 class WsConnectionSinglePool(
     ethereumWsConnectionFactory: EthereumWsConnectionFactory,
@@ -36,6 +37,9 @@ class WsConnectionSinglePool(
         return connection
     }
 
+    override fun connectionInfoFlux(): Flux<WsConnection.ConnectionInfo> =
+        connection.connectionInfoFlux()
+
     override fun close() {
         connection.close()
     }

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/WsSubscriptions.kt

@@ -37,5 +37,12 @@ interface WsSubscriptions {
     /**
      * Subscribe on remote
      */
-    fun subscribe(method: String): Flux<ByteArray>
+    fun subscribe(method: String): SubscribeData
+
+    fun connectionInfoFlux(): Flux<WsConnection.ConnectionInfo>
+
+    data class SubscribeData(
+        val data: Flux<ByteArray>,
+        val connectionId: String
+    )
 }

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/WsSubscriptionsImpl.kt

@@ -33,7 +33,7 @@ class WsSubscriptionsImpl(
 
     private val ids = AtomicLong(1)
 
-    override fun subscribe(method: String): Flux<ByteArray> {
+    override fun subscribe(method: String): WsSubscriptions.SubscribeData {
         val subscriptionId = AtomicReference("")
         val conn = wsPool.getConnection()
         val messages = conn.getSubscribeResponses()
@@ -41,7 +41,7 @@ class WsSubscriptionsImpl(
             .filter { it.result != null } // should never happen
             .map { it.result!! }
 
-        return conn.callRpc(JsonRpcRequest("eth_subscribe", listOf(method), ids.incrementAndGet()))
+        val messageFlux = conn.callRpc(JsonRpcRequest("eth_subscribe", listOf(method), ids.incrementAndGet()))
             .flatMapMany {
                 if (it.hasError()) {
                     log.warn("Failed to establish ETH Subscription: ${it.error?.message}")
@@ -51,5 +51,10 @@ class WsSubscriptionsImpl(
                     messages
                 }
             }
+
+        return WsSubscriptions.SubscribeData(messageFlux, conn.connectionId())
     }
+
+    override fun connectionInfoFlux(): Flux<WsConnection.ConnectionInfo> =
+        wsPool.connectionInfoFlux()
 }

src/main/kotlin/io/emeraldpay/dshackle/upstream/ethereum/subscribe/WebsocketPendingTxes.kt

@@ -34,6 +34,7 @@ class WebsocketPendingTxes(
 
     override fun createConnection(): Flux<TransactionId> {
         return wsSubscriptions.subscribe(EthereumEgressSubscription.METHOD_PENDING_TXES)
+            .data
             .timeout(Duration.ofSeconds(60), Mono.empty())
             .map {
                 // comes as a JS string, i.e., within quotes

src/test/groovy/io/emeraldpay/dshackle/upstream/ethereum/EthereumWsHeadSpec.groovy

@@ -26,6 +26,7 @@ import io.emeraldpay.etherjar.rpc.json.BlockJson
 import io.emeraldpay.etherjar.rpc.json.TransactionRefJson
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
+import reactor.core.publisher.Sinks
 import reactor.test.StepVerifier
 import spock.lang.Specification
 
@@ -60,7 +61,9 @@ class EthereumWsHeadSpec extends Specification {
         def apiMock = TestingCommons.api()
         apiMock.answerOnce("eth_getBlockByHash", ["0x3ec2ebf5d0ec474d0ac6bc50d2770d8409ad76e119968e7919f85d5ec8915200", false], block)
 
-        def ws = Mock(WsSubscriptions)
+        def ws = Mock(WsSubscriptions) {
+            1 * it.connectionInfoFlux() >> Flux.empty()
+        }
 
         def head = new EthereumWsHead("fake", new AlwaysForkChoice(), BlockValidator.ALWAYS_VALID, apiMock, ws, false)
 
@@ -73,9 +76,9 @@ class EthereumWsHeadSpec extends Specification {
         act.transactions[0].toHexWithPrefix() == "0x29229361dc5aa1ec66c323dc7a299e2b61a8c8dd2a3522d41255ec10eca25dd8"
         act.transactions[1].toHexWithPrefix() == "0xebe8f22a55a9e26892a8545b93cbb2bfa4fd81c3184e50e5cf6276025bb42b93"
 
-        1 * ws.subscribe("newHeads") >> Flux.fromIterable([
+        1 * ws.subscribe("newHeads") >> new WsSubscriptions.SubscribeData(
-                headBlock
+                Flux.fromIterable([headBlock]), "id"
-        ])
+        )
     }
 
     def "Restart ethereum ws head"() {
@@ -105,7 +108,11 @@ class EthereumWsHeadSpec extends Specification {
         apiMock.answerOnce("eth_blockNumber", [], Mono.empty())
 
         def ws = Mock(WsSubscriptions) {
-            2 * subscribe("newHeads") >>> [Flux.fromIterable([firstHeadBlock]), Flux.fromIterable([secondHeadBlock])]
+            1 * it.connectionInfoFlux() >> Flux.empty()
+            2 * subscribe("newHeads") >>> [
+                    new WsSubscriptions.SubscribeData(Flux.fromIterable([firstHeadBlock]), "id"),
+                    new WsSubscriptions.SubscribeData(Flux.fromIterable([secondHeadBlock]), "id")
+            ]
         }
 
         def head = new EthereumWsHead("fake", new AlwaysForkChoice(), BlockValidator.ALWAYS_VALID, apiMock, ws, true)
@@ -122,4 +129,55 @@ class EthereumWsHeadSpec extends Specification {
                 .thenCancel()
                 .verify(Duration.ofSeconds(1))
     }
+
+    def "Restart ethereum ws head immediately after reconnection"() {
+        setup:
+        def block = new BlockJson<TransactionRefJson>()
+        block.timestamp = Instant.now().truncatedTo(ChronoUnit.SECONDS)
+        block.number = 103
+        block.parentHash = parent
+        block.hash = BlockHash.from("0x3ec2ebf5d0ec474d0ac6bc50d2770d8409ad76e119968e7919f85d5ec8915200")
+        def secondBlock = new BlockJson<TransactionRefJson>()
+        secondBlock.parentHash = parent
+        secondBlock.timestamp = Instant.now().truncatedTo(ChronoUnit.SECONDS)
+        secondBlock.number = 105
+        secondBlock.hash = BlockHash.from("0x29229361dc5aa1ec66c323dc7a299e2b61a8c8dd2a3522d41255ec10eca25dd8")
+
+        def firstHeadBlock = block.with {
+            Global.objectMapper.writeValueAsBytes(it)
+        }
+        def secondHeadBlock = secondBlock.with {
+            Global.objectMapper.writeValueAsBytes(it)
+        }
+
+        def apiMock = TestingCommons.api()
+        def connectionInfoSink = Sinks.many().multicast().directBestEffort()
+        apiMock.answerOnce("eth_getBlockByHash", ["0x3ec2ebf5d0ec474d0ac6bc50d2770d8409ad76e119968e7919f85d5ec8915200", false], null)
+        apiMock.answerOnce("eth_getBlockByHash", ["0x29229361dc5aa1ec66c323dc7a299e2b61a8c8dd2a3522d41255ec10eca25dd8", false], null)
+        apiMock.answerOnce("eth_blockNumber", [], Mono.empty())
+        apiMock.answerOnce("eth_blockNumber", [], Mono.empty())
+
+        def ws = Mock(WsSubscriptions) {
+            1 * it.connectionInfoFlux() >> connectionInfoSink.asFlux()
+            2 * subscribe("newHeads") >>> [
+                    new WsSubscriptions.SubscribeData(Flux.fromIterable([firstHeadBlock]), "id"),
+                    new WsSubscriptions.SubscribeData(Flux.fromIterable([secondHeadBlock]), "id")
+            ]
+        }
+
+        def head = new EthereumWsHead("fake", new AlwaysForkChoice(), BlockValidator.ALWAYS_VALID, apiMock, ws, true)
+
+        when:
+        def act = head.getFlux()
+
+        then:
+        StepVerifier.create(act)
+                .then { head.start() }
+                .expectNext(BlockContainer.from(block))
+                .then { connectionInfoSink.tryEmitNext(new WsConnection.ConnectionInfo("id", WsConnection.ConnectionState.DISCONNECTED)) }
+                .then { connectionInfoSink.tryEmitNext(new WsConnection.ConnectionInfo("id", WsConnection.ConnectionState.CONNECTED)) }
+                .expectNext(BlockContainer.from(secondBlock))
+                .thenCancel()
+                .verify(Duration.ofSeconds(1))
+    }
 }

src/test/groovy/io/emeraldpay/dshackle/upstream/ethereum/WsConnectionMultiPoolSpec.groovy

@@ -16,6 +16,7 @@
 package io.emeraldpay.dshackle.upstream.ethereum
 
 import io.emeraldpay.dshackle.upstream.DefaultUpstream
+import reactor.core.publisher.Flux
 import spock.lang.Specification
 
 import java.util.concurrent.ScheduledExecutorService
@@ -24,7 +25,9 @@ class WsConnectionMultiPoolSpec extends Specification {
 
     def "create connection when less than required"() {
         setup:
-        def conn = Mock(WsConnection)
+        def conn = Mock(WsConnection) {
+            1 * it.connectionInfoFlux() >> Flux.empty()
+        }
         def up = Mock(DefaultUpstream)
         def factory = Mock(EthereumWsConnectionFactory)
         def pool = new WsConnectionMultiPool(factory, up, 3)
@@ -40,9 +43,15 @@ class WsConnectionMultiPoolSpec extends Specification {
 
     def "create connection until target"() {
         setup:
-        def conn1 = Mock(WsConnection)
+        def conn1 = Mock(WsConnection) {
-        def conn2 = Mock(WsConnection)
+            1 * it.connectionInfoFlux() >> Flux.empty()
-        def conn3 = Mock(WsConnection)
+        }
+        def conn2 = Mock(WsConnection) {
+            1 * it.connectionInfoFlux() >> Flux.empty()
+        }
+        def conn3 = Mock(WsConnection) {
+            1 * it.connectionInfoFlux() >> Flux.empty()
+        }
         def up = Mock(DefaultUpstream)
         def factory = Mock(EthereumWsConnectionFactory)
         def pool = new WsConnectionMultiPool(factory, up, 3)
@@ -84,10 +93,18 @@ class WsConnectionMultiPoolSpec extends Specification {
 
     def "recreate connection after failure"() {
         setup:
-        def conn1 = Mock(WsConnection)
+        def conn1 = Mock(WsConnection) {
-        def conn2 = Mock(WsConnection)
+            1 * it.connectionInfoFlux() >> Flux.empty()
-        def conn3 = Mock(WsConnection)
+        }
-        def conn4 = Mock(WsConnection)
+        def conn2 = Mock(WsConnection) {
+            1 * it.connectionInfoFlux() >> Flux.empty()
+        }
+        def conn3 = Mock(WsConnection) {
+            1 * it.connectionInfoFlux() >> Flux.empty()
+        }
+        def conn4 = Mock(WsConnection) {
+            1 * it.connectionInfoFlux() >> Flux.empty()
+        }
         def up = Mock(DefaultUpstream)
         def factory = Mock(EthereumWsConnectionFactory)
         def pool = new WsConnectionMultiPool(factory, up, 3)

src/test/groovy/io/emeraldpay/dshackle/upstream/ethereum/WsSubscriptionsImplSpec.groovy

@@ -36,7 +36,9 @@ class WsSubscriptionsImplSpec extends Specification {
                 ]
         )
 
-        def conn = Mock(WsConnection)
+        def conn = Mock(WsConnection) {
+            1 * it.connectionId() >> "id"
+        }
         def pool = Mock(WsConnectionPool) {
             getConnection() >> conn
         }
@@ -44,6 +46,7 @@ class WsSubscriptionsImplSpec extends Specification {
 
         when:
         def act = ws.subscribe("foo_bar")
+            .data
             .map { new String(it) }
             .take(3)
             .collectList().block(Duration.ofSeconds(1))
@@ -71,7 +74,9 @@ class WsSubscriptionsImplSpec extends Specification {
                 ]
         )
 
-        def conn = Mock(WsConnection)
+        def conn = Mock(WsConnection) {
+            1 * it.connectionId() >> "id"
+        }
         def pool = Mock(WsConnectionPool) {
             getConnection() >> conn
         }
@@ -79,6 +84,7 @@ class WsSubscriptionsImplSpec extends Specification {
 
         when:
         def act = ws.subscribe("foo_bar")
+                .data
                 .map { new String(it) }
                 .take(3)
                 .collectList().block(Duration.ofSeconds(1))

src/test/groovy/io/emeraldpay/dshackle/upstream/ethereum/subscribe/WebsocketPendingTxesSpec.groovy

@@ -40,7 +40,9 @@ class WebsocketPendingTxesSpec extends Specification {
                 .collectList().block(Duration.ofSeconds(1))
 
         then:
-        1 * ws.subscribe("newPendingTransactions") >> Flux.fromIterable(responses)
+        1 * ws.subscribe("newPendingTransactions") >> new WsSubscriptions.SubscribeData(
+                Flux.fromIterable(responses), "id"
+        )
         txes.collect {it.toHex() } == [
                 "0xa61bab14fc9720ea8725622688c2f964666d7c2afdae38af7dad53f12f242d5c",
                 "0x911548eb0f3bf353a54e03a3506c7c3e747470d6c201f03babbc07ff6e14cd6e",