Fix OA-1097: evict ServerRoutingStatsManager entries when servers are decommissioned

pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java

@@ -325,9 +325,10 @@ private void processInstanceConfigChange() {
     long updateRoutingEntriesEndTimeMs = System.currentTimeMillis();
 
     // Remove new disabled servers from _enabledServerInstanceMap after updating all routing entries to ensure it
-    // always contains the selected servers
+    // always contains the selected servers. Also evict their routing stats to prevent unbounded map growth.
     for (String newDisabledInstance : newDisabledServers) {
       _enabledServerInstanceMap.remove(newDisabledInstance);
+      _serverRoutingStatsManager.removeServerStats(newDisabledInstance);
     }
 
     LOGGER.info("Processed instance config change in {}ms "

pinot-common/src/main/java/org/apache/pinot/common/utils/ExponentialMovingAverage.java

@@ -21,6 +21,7 @@
 
 import com.google.common.base.Preconditions;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.ThreadSafe;
@@ -41,6 +42,9 @@ public class ExponentialMovingAverage {
   private volatile double _average;
   private long _lastUpdatedTimeMs;
 
+  @Nullable
+  private ScheduledFuture<?> _decayTaskFuture;
+
   /**
    * Constructor
    *
@@ -70,7 +74,7 @@ public ExponentialMovingAverage(double alpha, long autoDecayWindowMs, long warmU
     if (_autoDecayWindowMs > 0) {
       // Schedule a task to automatically decay the average if updates are not performed in the last _autoDecayWindowMs.
       Preconditions.checkState(periodicTaskExecutor != null);
-      periodicTaskExecutor.scheduleAtFixedRate(new Runnable() {
+      _decayTaskFuture = periodicTaskExecutor.scheduleAtFixedRate(new Runnable() {
         @Override
         public void run() {
           if (_lastUpdatedTimeMs > 0 && (System.currentTimeMillis() - _lastUpdatedTimeMs) > _autoDecayWindowMs) {
@@ -81,6 +85,16 @@ public void run() {
     }
   }
 
+  /**
+   * Cancels the periodic auto-decay task if one is running. Safe to call from any thread.
+   * Has no effect if auto-decay was not configured or has already been cancelled.
+   */
+  public void cancel() {
+    if (_decayTaskFuture != null) {
+      _decayTaskFuture.cancel(false);
+    }
+  }
+
   /**
    * Returns the exponentially weighted moving average.
    */

pinot-core/src/main/java/org/apache/pinot/core/transport/server/routing/stats/ServerRoutingStatsEntry.java

@@ -100,4 +100,14 @@ public void updateNumInFlightRequestsForResponseArrival() {
   public void updateLatency(double latencyMs) {
     _latencyMsEMA.compute(latencyMs);
   }
+
+  /**
+   * Cancels the periodic auto-decay tasks scheduled by both EMA instances.
+   * Should be called when this entry is permanently removed to release the scheduled tasks
+   * from the shared executor and allow the entry to be garbage-collected.
+   */
+  public void cancel() {
+    _inFlighRequestsEMA.cancel();
+    _latencyMsEMA.cancel();
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/transport/server/routing/stats/ServerRoutingStatsManager.java

@@ -165,6 +165,27 @@ private void updateStatsAfterQuerySubmission(String serverInstanceId) {
     }
   }
 
+  /**
+   * Removes and cancels the routing stats entry for a server that has been removed from the cluster.
+   * This prevents unbounded growth of {@code _serverQueryStatsMap} on long-running brokers as servers
+   * are decommissioned and replaced over time. The entry's periodic auto-decay tasks are cancelled
+   * so that the entry and its EMA objects can be garbage-collected.
+   *
+   * <p>Safe to call concurrently with stats recording. Any in-flight recording tasks that obtained the
+   * entry before this removal will complete normally; subsequent recordings for this server will create
+   * a fresh entry via {@code computeIfAbsent}.
+   */
+  public void removeServerStats(String serverInstanceId) {
+    if (!_isEnabled) {
+      return;
+    }
+    ServerRoutingStatsEntry removed = _serverQueryStatsMap.remove(serverInstanceId);
+    if (removed != null) {
+      removed.cancel();
+      LOGGER.info("Removed and cancelled routing stats for decommissioned server: {}", serverInstanceId);
+    }
+  }
+
   /**
    * Called when a query response is received from the server. Updates stats related to query completion.
    */

pinot-core/src/test/java/org/apache/pinot/core/transport/server/routing/stats/ServerRoutingStatsManagerTest.java

@@ -317,6 +317,52 @@ public void testQuerySubmitAndCompletionStats() {
     assertEquals(score, 54.0);
   }
 
+  @Test
+  public void testRemoveServerStats() {
+    Map<String, Object> properties = new HashMap<>();
+    properties.put(CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_ENABLE_STATS_COLLECTION, true);
+    properties.put(CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_AUTODECAY_WINDOW_MS, -1);
+    properties.put(CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_WARMUP_DURATION_MS, 0);
+    properties.put(CommonConstants.Broker.AdaptiveServerSelector.CONFIG_OF_AVG_INITIALIZATION_VAL, 0.0);
+    ServerRoutingStatsManager manager = new ServerRoutingStatsManager(new PinotConfiguration(properties),
+        _brokerMetrics);
+    manager.init();
+
+    int requestId = 0;
+
+    // Record stats for two servers
+    manager.recordStatsForQuerySubmission(requestId++, "server1");
+    manager.recordStatsForQuerySubmission(requestId++, "server2");
+    waitForStatsUpdate(manager, requestId);
+
+    assertNotNull(manager.fetchNumInFlightRequestsForServer("server1"));
+    assertNotNull(manager.fetchNumInFlightRequestsForServer("server2"));
+    assertEquals(manager.fetchNumInFlightRequestsForAllServers().size(), 2);
+
+    // Remove server1 (simulating decommission)
+    manager.removeServerStats("server1");
+
+    assertNull(manager.fetchNumInFlightRequestsForServer("server1"),
+        "Stats for removed server should be null");
+    assertNotNull(manager.fetchNumInFlightRequestsForServer("server2"),
+        "Stats for active server should still exist");
+    assertEquals(manager.fetchNumInFlightRequestsForAllServers().size(), 1,
+        "Only one server should remain in the stats map");
+
+    // Removing an already-absent server must not throw
+    manager.removeServerStats("server1");
+    manager.removeServerStats("nonExistentServer");
+
+    // Remove the last server
+    manager.removeServerStats("server2");
+    assertTrue(manager.fetchNumInFlightRequestsForAllServers().isEmpty(),
+        "Stats map should be empty after all servers are removed");
+
+    // removeServerStats on a disabled manager must be a no-op and not throw
+    manager.shutDown();
+    manager.removeServerStats("server1");
+  }
+
   private void waitForStatsUpdate(ServerRoutingStatsManager serverRoutingStatsManager, long taskCount) {
     TestUtils.waitForCondition(aVoid -> {
       return (serverRoutingStatsManager.getCompletedTaskCount() == taskCount);