problem: race condition on a new block on head

src/main/kotlin/io/emeraldpay/dshackle/data/BlockContainer.kt

@@ -63,6 +63,10 @@ class BlockContainer(
         }
     }
 
+    override fun toString(): String {
+        return "Block $height = $hash"
+    }
+
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
         if (javaClass != other?.javaClass) return false

src/main/kotlin/io/emeraldpay/dshackle/data/BlockId.kt

@@ -25,6 +25,11 @@ class BlockId(
 ) : HashId(value) {
 
     companion object {
+        @JvmStatic
+        fun from(hash: ByteArray): BlockId {
+            return BlockId(hash)
+        }
+
         @JvmStatic
         fun from(hash: BlockHash): BlockId {
             return BlockId(hash.bytes)

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

@@ -31,15 +31,24 @@ abstract class AbstractHead : Head {
 
     private val head = AtomicReference<BlockContainer>(null)
     private val stream: TopicProcessor<BlockContainer> = TopicProcessor.create()
+    private val beforeBlockHandlers = ArrayList<Runnable>()
 
     fun follow(source: Flux<BlockContainer>): Disposable {
-        return source.distinctUntilChanged {
+        return source
-            it.hash
+                .distinctUntilChanged {
-        }.filter { block ->
+                    it.hash
-            val curr = head.get()
+                }.filter { block ->
-            curr == null || curr.difficulty < block.difficulty
+                    val curr = head.get()
-        }
+                    curr == null || curr.difficulty < block.difficulty
+                }
+                .doFinally {
+                    // close internal stream if upstream is finished, otherwise it gets stuck
+                    // but technically is should never happen during normal work, only when the Head
+                    // is stopping
+                    stream.onComplete()
+                }
                 .subscribe { block ->
+                    notifyBeforeBlock()
                     val prev = head.getAndUpdate { curr ->
                         if (curr == null || curr.difficulty < block.difficulty) {
                             block
@@ -54,6 +63,20 @@ abstract class AbstractHead : Head {
                 }
     }
 
+    fun notifyBeforeBlock() {
+        beforeBlockHandlers.forEach { handler ->
+            try {
+                handler.run()
+            } catch (t: Throwable) {
+                log.warn("Before Block handler error", t)
+            }
+        }
+    }
+
+    override fun onBeforeBlock(handler: Runnable) {
+        beforeBlockHandlers.add(handler)
+    }
+
     override fun getFlux(): Flux<BlockContainer> {
         return Flux.merge(
                 Mono.justOrEmpty(head.get()),

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

@@ -25,4 +25,7 @@ class EmptyHead : Head {
     override fun getFlux(): Flux<BlockContainer> {
         return Flux.empty()
     }
+
+    override fun onBeforeBlock(handler: Runnable) {
+    }
 }

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

@@ -20,6 +20,20 @@ import io.emeraldpay.dshackle.data.BlockContainer
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
 
+/**
+ * Subscription to listen to updates to the head of a blockchain.
+ */
 interface Head {
+
+    /**
+     * @return stream of all new blocks, starts from the current block (i.e., first item should be available immediately).
+     */
     fun getFlux(): Flux<BlockContainer>
+
+    /**
+     * Add handler that is going to be called each time _before_ a new block is submitted to stream of new blocks.
+     * Supposed to be used for cleanup/preparation before actual block data will come, to avoid race condition.
+     * @see getFlux
+     */
+    fun onBeforeBlock(handler: Runnable)
 }

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

@@ -54,7 +54,6 @@ open class EthereumReader(
         private val log = LoggerFactory.getLogger(EthereumReader::class.java)
     }
 
-    private var headListener: Disposable? = null
     private val balanceCache = CurrentBlockCache<Address, Wei>()
 
     val extractBlock = Function<BlockContainer, BlockJson<TransactionRefJson>> { block ->
@@ -252,18 +251,17 @@ open class EthereumReader(
     }
 
     override fun isRunning(): Boolean {
-        return this.headListener != null
+        //TODO should be always running?
+        return up.isRunning
     }
 
     override fun start() {
-        this.headListener = up.getHead().getFlux().subscribe {
+        val evictCaches: Runnable = Runnable {
             balanceCache.evict()
         }
+        up.getHead().onBeforeBlock(evictCaches)
     }
 
     override fun stop() {
-        val headListener = this.headListener
-        this.headListener = null
-        headListener?.dispose()
     }
 }

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

@@ -109,7 +109,6 @@ class TrackEthereumAddressSpec extends Specification {
                 .expectNext(exp1).as("First block")
                 .then {
                     upstreamMock.nextBlock(BlockContainer.from(block2, TestingCommons.objectMapper()))
-                    Thread.sleep(50)
                 }
                 .expectNext(exp2).as("Second block")
                 .thenCancel()

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

@@ -18,6 +18,7 @@ package io.emeraldpay.dshackle.test
 
 import io.emeraldpay.dshackle.data.BlockContainer
 import io.emeraldpay.dshackle.upstream.Head
+import org.jetbrains.annotations.NotNull
 import org.reactivestreams.Publisher
 import reactor.core.publisher.Flux
 import reactor.core.publisher.Mono
@@ -28,8 +29,12 @@ class EthereumHeadMock implements Head {
     private TopicProcessor<BlockContainer> bus = TopicProcessor.create()
     private Publisher<BlockContainer> predefined = null
     private BlockContainer latest
+    private List<Runnable> handlers = []
 
     void nextBlock(BlockContainer block) {
+        handlers.forEach {
+            it.run()
+        }
         assert block != null
         println("New block: ${block.height} / ${block.hash}")
         latest = block
@@ -52,4 +57,9 @@ class EthereumHeadMock implements Head {
             return Flux.concat(Mono.justOrEmpty(latest), bus).distinctUntilChanged()
         }
     }
+
+    @Override
+    void onBeforeBlock(@NotNull Runnable handler) {
+        handlers.add(handler)
+    }
 }

src/test/groovy/io/emeraldpay/dshackle/upstream/AbstractHeadSpec.groovy

@@ -0,0 +1,117 @@
+/**
+ * Copyright (c) 2020 EmeraldPay, Inc
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *        http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.emeraldpay.dshackle.upstream
+
+import io.emeraldpay.dshackle.data.BlockContainer
+import io.emeraldpay.dshackle.data.BlockId
+import reactor.core.publisher.Flux
+import reactor.core.publisher.TopicProcessor
+import reactor.test.StepVerifier
+import spock.lang.Specification
+
+import java.time.Duration
+import java.time.Instant
+
+class AbstractHeadSpec extends Specification {
+
+    def blocks = [1L, 2, 3, 4].collect { i ->
+        byte[] hash = new byte[32]
+        hash[0] = i as byte
+        new BlockContainer(i, BlockId.from(hash), BigInteger.valueOf(i), Instant.now(), false, null, null, [])
+    }
+
+    def "Calls beforeBlock on each block"() {
+        setup:
+        TopicProcessor<BlockContainer> source = TopicProcessor.create()
+        def head = new TestHead()
+        def called = false
+        when:
+        head.follow(Flux.from(source))
+        head.onBeforeBlock {
+            called = true
+        }
+        def act = head.flux
+        source.onNext(blocks[0])
+        then:
+        StepVerifier.create(act)
+                .expectNext(blocks[0])
+                .then {
+                    assert called
+                    called = false
+                    source.onNext(blocks[1])
+                }
+                .expectNext(blocks[1])
+                .then {
+                    assert called
+                    source.onComplete()
+                }
+                .expectComplete()
+                .verify(Duration.ofSeconds(1))
+    }
+
+    def "Follows source"() {
+        setup:
+        TopicProcessor<BlockContainer> source = TopicProcessor.create()
+        def head = new TestHead()
+        when:
+        head.follow(Flux.from(source))
+        def act = head.flux
+        source.onNext(blocks[0])
+        then:
+        StepVerifier.create(act)
+                .expectNext(blocks[0])
+                .then { source.onNext(blocks[1]) }
+                .expectNext(blocks[1])
+                .then { source.onNext(blocks[2]) }
+                .expectNext(blocks[2])
+                .then { source.onNext(blocks[3]) }
+                .expectNext(blocks[3])
+                .then { source.onComplete() }
+                .expectComplete()
+                .verify(Duration.ofSeconds(1))
+    }
+
+    def "Ignores block will less difficulty"() {
+        setup:
+        TopicProcessor<BlockContainer> source = TopicProcessor.create()
+        def head = new TestHead()
+        def wrongblock = new BlockContainer(
+                blocks[1].height, BlockId.from(blocks[1].hash.value.clone().tap { it[1] = 0xff as byte }),
+                blocks[1].difficulty - 1,
+                Instant.now(),
+                false, null, null, []
+        )
+        when:
+        head.follow(Flux.from(source))
+        def act = head.flux
+        source.onNext(blocks[0])
+        then:
+        StepVerifier.create(act)
+                .expectNext(blocks[0])
+                .then { source.onNext(blocks[1]) }
+                .expectNext(blocks[1])
+                .then { source.onNext(wrongblock) }
+                .then { source.onNext(blocks[3]) }
+                .expectNext(blocks[3])
+                .then { source.onComplete() }
+                .expectComplete()
+                .verify(Duration.ofSeconds(1))
+    }
+
+    class TestHead extends AbstractHead {
+
+    }
+}