From b450cc9a9407742e0cae62f2b5c197d4978205b5 Mon Sep 17 00:00:00 2001
From: a10zn8 <a1.cyb3r@gmail.com>
Date: Thu, 17 Nov 2022 13:52:36 +0400
Subject: [PATCH] use only chain map

---
 .../upstream/CurrentMultistreamHolder.kt      | 21 +++++--------------
 .../dshackle/upstream/Multistream.kt          |  4 ++++
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/src/main/kotlin/io/emeraldpay/dshackle/upstream/CurrentMultistreamHolder.kt b/src/main/kotlin/io/emeraldpay/dshackle/upstream/CurrentMultistreamHolder.kt
index b4865951..e4f6029d 100644
--- a/src/main/kotlin/io/emeraldpay/dshackle/upstream/CurrentMultistreamHolder.kt
+++ b/src/main/kotlin/io/emeraldpay/dshackle/upstream/CurrentMultistreamHolder.kt
@@ -19,31 +19,23 @@ package io.emeraldpay.dshackle.upstream
 import io.emeraldpay.grpc.Chain
 import org.slf4j.LoggerFactory
 import org.springframework.stereotype.Component
-import reactor.core.publisher.Sinks
-import java.util.*
-import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.locks.ReentrantLock
 import javax.annotation.PreDestroy
-import kotlin.concurrent.withLock
 
 @Component
 open class CurrentMultistreamHolder(
-    private val multistreams: List<Multistream>
+    multistreams: List<Multistream>
 ) : MultistreamHolder {
 
     private val log = LoggerFactory.getLogger(CurrentMultistreamHolder::class.java)
 
-    private val chainMapping = ConcurrentHashMap<Chain, Multistream>().apply {
-        multistreams.forEach { this[it.chain] = it }
-    }
-    private val updateLock = ReentrantLock()
+    private val chainMapping = multistreams.associateBy { it.chain }
 
     override fun getUpstream(chain: Chain): Multistream? {
         return chainMapping[chain]
     }
 
     override fun getAvailable(): List<Chain> {
-        return multistreams.asSequence()
+        return chainMapping.values.asSequence()
             .filter { it.isAvailable() }
             .map { it.chain }
             .toList()
@@ -56,11 +48,8 @@ open class CurrentMultistreamHolder(
     @PreDestroy
     fun shutdown() {
         log.info("Closing upstream connections...")
-        updateLock.withLock {
-            chainMapping.values.forEach {
-                it.stop()
-            }
-            chainMapping.clear()
+        chainMapping.values.forEach {
+            it.stop()
         }
     }
 }
diff --git a/src/main/kotlin/io/emeraldpay/dshackle/upstream/Multistream.kt b/src/main/kotlin/io/emeraldpay/dshackle/upstream/Multistream.kt
index 54fb2dfd..30fa4f24 100644
--- a/src/main/kotlin/io/emeraldpay/dshackle/upstream/Multistream.kt
+++ b/src/main/kotlin/io/emeraldpay/dshackle/upstream/Multistream.kt
@@ -324,6 +324,10 @@ abstract class Multistream(
         log.info("State of ${chain.chainCode}: height=${height ?: '?'}, status=[$statuses], lag=[$lag], weak=[$weak]")
     }
 
+    fun test(event: UpstreamChangeEvent): Boolean {
+        return event.chain == this.chain
+    }
+
     @EventListener
     @Order(Ordered.HIGHEST_PRECEDENCE)
     fun onUpstreamChange(event: UpstreamChangeEvent) {