Fix #178: Support multiple FastestServerLB instances

inoio · Jan 30, 2024 · 7af2dc0 · 7af2dc0
1 parent d020606
commit 7af2dc0
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 6 deletions.
diff --git 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()}")
 
 }
diff --git 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