StakeSquid/dshackle
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Resubscribe to newHeads immediately after reconnection (#189)

Browse Source 
* Resubscribe to newHeads immediately after reconnection
This commit is contained in:
KirillPamPam
2023-03-30 13:32:02 +04:00
committed by GitHub
parent dcf360a7d7
commit 97fa616e19
14 changed files with 216 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
src/main/kotlin/io/emeraldpay/dshackle/upstream/WebsocketConnectionStatesHandler.kt Normal file
View File

@@ -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
}

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

@@ -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>
}

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

@@ -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
}
}

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

@@ -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

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

@@ -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 {

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

@@ -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>
}

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

@@ -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()
}

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

@@ -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
)
}

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

@@ -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()
}

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

@@ -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

68
src/test/groovy/io/emeraldpay/dshackle/upstream/ethereum/EthereumWsHeadSpec.groovy
View File

@@ -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([
headBlock
])
1 * ws.subscribe("newHeads") >> new WsSubscriptions.SubscribeData(
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))
}
}

33
src/test/groovy/io/emeraldpay/dshackle/upstream/ethereum/WsConnectionMultiPoolSpec.groovy
View File

@@ -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 conn2 = Mock(WsConnection)
def conn3 = Mock(WsConnection)
def conn1 = Mock(WsConnection) {
1 * it.connectionInfoFlux() >> Flux.empty()
}
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 conn2 = Mock(WsConnection)
def conn3 = Mock(WsConnection)
def conn4 = Mock(WsConnection)
def conn1 = Mock(WsConnection) {
1 * it.connectionInfoFlux() >> Flux.empty()
}
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)

10
src/test/groovy/io/emeraldpay/dshackle/upstream/ethereum/WsSubscriptionsImplSpec.groovy
View File

@@ -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))

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

@@ -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",