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 {
-            it.hash
-        }.filter { block ->
-            val curr = head.get()
-            curr == null || curr.difficulty < block.difficulty
-        }
+        return source
+                .distinctUntilChanged {
+                    it.hash
+                }.filter { block ->
+                    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()
     }
 }