From 7af2dc03fa84dc4ef845dc6cc28c206c1cab3f83 Mon Sep 17 00:00:00 2001
From: Martin Grotzke <martin.grotzke@inoio.de>
Date: Tue, 30 Jan 2024 14:30:09 +0100
Subject: [PATCH] Fix #178: Support multiple FastestServerLB instances

---
 .../scala/io/ino/solrs/LoadBalancer.scala     | 25 ++++++++++++++-----
 .../io/ino/solrs/FastestServerLBSpec.scala    |  7 ++++++
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/main/scala/io/ino/solrs/LoadBalancer.scala b/src/main/scala/io/ino/solrs/LoadBalancer.scala
index 475ea89..6b499ec 100644
--- a/src/main/scala/io/ino/solrs/LoadBalancer.scala
+++ b/src/main/scala/io/ino/solrs/LoadBalancer.scala
@@ -5,7 +5,6 @@ import java.util.concurrent.Executors
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.function.IntUnaryOperator
-
 import io.ino.solrs.LoadBalancer.NoSolrServersAvailableException
 import io.ino.solrs.ServerStateChangeObservable.Removed
 import io.ino.solrs.ServerStateChangeObservable.StateChange
@@ -23,8 +22,11 @@ import org.apache.solr.client.solrj.SolrResponse
 import org.apache.solr.client.solrj.request.IsUpdateRequest
 import org.apache.solr.client.solrj.request.QueryRequest
 import org.apache.solr.client.solrj.response.QueryResponse
+import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 
+import javax.management.InstanceAlreadyExistsException
+import javax.management.InstanceNotFoundException
 import scala.collection.concurrent.TrieMap
 import scala.collection.mutable
 import scala.concurrent.duration._
@@ -184,8 +186,6 @@ class FastestServerLB[F[_]](override val solrServers: SolrServers,
 
   import FastestServerLB._
 
-  private val logger = LoggerFactory.getLogger(getClass)
-
   private var client: AsyncSolrClient[F] = _
 
   private val scheduler = Executors.newSingleThreadScheduledExecutor()
@@ -532,12 +532,24 @@ trait FastestServerLBJmxSupport[F[_]] extends FastestServerLBMBean { self: Faste
   import FastestServerLB._
   import FastestServerLBJmxSupport._
 
+  protected val logger: Logger = LoggerFactory.getLogger(getClass)
+
+  private lazy val objectName = newObjectName()
+
   def initJmx(): Unit = {
-    ManagementFactory.getPlatformMBeanServer.registerMBean(this, ObjName)
+    try {
+      ManagementFactory.getPlatformMBeanServer.registerMBean(this, objectName)
+    } catch {
+        case e: InstanceAlreadyExistsException => logger.warn("Error while registering MBean", e)
+    }
   }
 
   def shutdownJmx(): Unit = {
-    ManagementFactory.getPlatformMBeanServer.unregisterMBean(ObjName)
+    try {
+      ManagementFactory.getPlatformMBeanServer.unregisterMBean(objectName)
+    } catch {
+      case e: InstanceNotFoundException => logger.warn("Error while unregistering MBean", e)
+    }
   }
 
   override def averagesPerSecond(collection: String): TabularData = {
@@ -615,6 +627,7 @@ trait FastestServerLBJmxSupport[F[_]] extends FastestServerLBMBean { self: Faste
 
 object FastestServerLBJmxSupport {
 
-  val ObjName = new ObjectName("io.ino.solrs:type=FastestServerLB")
+  private val instanceCounter = new AtomicInteger(0)
+  def newObjectName(): ObjectName = new ObjectName(s"io.ino.solrs:type=FastestServerLB,instance=${instanceCounter.getAndIncrement()}")
 
 }
\ No newline at end of file
diff --git a/src/test/scala/io/ino/solrs/FastestServerLBSpec.scala b/src/test/scala/io/ino/solrs/FastestServerLBSpec.scala
index e1fb19b..382f9a6 100644
--- a/src/test/scala/io/ino/solrs/FastestServerLBSpec.scala
+++ b/src/test/scala/io/ino/solrs/FastestServerLBSpec.scala
@@ -63,6 +63,13 @@ class FastestServerLBSpec extends StandardFunSpec {
 
   describe("FastestServerLB") {
 
+    it("should support multiple instances") {
+      val servers = IndexedSeq(SolrServer("host1"))
+      this.cut = newDynamicLB(new StaticSolrServers(servers), q, clock)
+      val cut2 = newDynamicLB(new StaticSolrServers(servers), q, clock)
+      cut2.shutdown()
+    }
+
     it("should return a Failure if no solr server matches") {
       val nonMatchingServers = new SolrServers {
         override def all: Seq[SolrServer] = Nil