Clear DarkCluster Buffer on Refresh (#1132)

* Clear DarkCluster Buffers On Refresh

* Improvements

* Simplify

* Additional unit test, remove logs, add back comment

* Fix test

* Bump version and update CHANGELOG

---------

Co-authored-by: Jason Leezer <jleezer@linkedin.com>

Author: jcleezer

CHANGELOG.md

@@ -14,6 +14,9 @@ and what APIs have changed, if applicable.
 
 ## [Unreleased]
 
+## [29.81.2] - 2026-01-09
+- Clear DarkCluster Buffer on Refresh
+
 ## [29.81.1] - 2025-11-04
 - Fix xds client active wait time metric also add more logs.
 
@@ -5933,7 +5936,8 @@ patch operations can re-use these classes for generating patch messages.
 
 ## [0.14.1]
 
-[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.81.1...master
+[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.81.2...master
+[29.81.2]: https://github.com/linkedin/rest.li/compare/v29.81.1...v29.81.2
 [29.81.1]: https://github.com/linkedin/rest.li/compare/v29.81.0...v29.81.1
 [29.81.0]: https://github.com/linkedin/rest.li/compare/v29.80.3...v29.81.0
 [29.80.3]: https://github.com/linkedin/rest.li/compare/v29.80.2...v29.80.3

darkcluster/src/main/java/com/linkedin/darkcluster/api/DarkClusterStrategy.java

@@ -33,4 +33,9 @@ public interface DarkClusterStrategy
    * @return  true if at least one request was sent.
    */
   boolean handleRequest(final RestRequest originalRequest, final RestRequest darkRequest, final RequestContext requestContext);
+
+  default void shutdown()
+  {
+    // no-op by default
+  }
 }

darkcluster/src/main/java/com/linkedin/darkcluster/impl/ConstantQpsDarkClusterStrategy.java

@@ -81,6 +81,12 @@ public boolean handleRequest(RestRequest originalRequest, RestRequest darkReques
     return addRequest(originalRequest, darkRequest, requestContext);
   }
 
+  @Override
+  public void shutdown()
+  {
+    _rateLimiter.cancelAll(new RuntimeException("Shutting down ConstantQpsDarkClusterStrategy"));
+  }
+
   /**
    * We won't create this strategy if this config isn't valid for this strategy. For instance, we don't want to create
    * the ConstantQpsDarkClusterStrategy if any of the configurables are zero, because we'd be doing pointless work on every getOrCreate.
@@ -155,6 +161,7 @@ private float getSendRate()
    */
   private boolean addRequest(RestRequest originalRequest, RestRequest darkRequest, RequestContext requestContext)
   {
+    final BaseDarkClusterDispatcher baseDispatcher = _baseDarkClusterDispatcher;
     _rateLimiter.submit(new Callback<None>()
     {
       @Override
@@ -166,7 +173,7 @@ public void onError(Throwable e)
       @Override
       public void onSuccess(None result)
       {
-        _baseDarkClusterDispatcher.sendRequest(originalRequest, darkRequest, requestContext, NUM_REQUESTS_TO_SEND_PER_RATE_LIMITER_CYCLE);
+        baseDispatcher.sendRequest(originalRequest, darkRequest, requestContext, NUM_REQUESTS_TO_SEND_PER_RATE_LIMITER_CYCLE);
       }
     });
     return true;

darkcluster/src/main/java/com/linkedin/darkcluster/impl/DarkClusterStrategyFactoryImpl.java

@@ -278,14 +278,30 @@ public void onSuccess(DarkClusterConfigMap updatedDarkConfigMap)
 
             for (int partitionId : partitions)
             {
+
               Map<String, DarkClusterStrategy> darkStrategyMap = new ConcurrentHashMap<>();
               for (Map.Entry<String, DarkClusterConfig> entry : updatedDarkConfigMap.entrySet())
               {
                 String darkClusterToAdd = entry.getKey();
                 darkStrategyMap.put(darkClusterToAdd, createStrategy(darkClusterToAdd, entry.getValue(), partitionId));
-                LOG.info("Created new strategy for dark cluster: " + darkClusterToAdd + ", partition: " + partitionId + ", source cluster: " + _sourceClusterName);
+                LOG.info("Created new strategy for dark cluster: " + darkClusterToAdd + ", partition: " + partitionId
+                    + ", source cluster: " + _sourceClusterName);
               }
-              _partitionToDarkStrategyMap.put(partitionId, darkStrategyMap);
+              _partitionToDarkStrategyMap.compute(partitionId, (key, oldDarkStrategyMap) -> {
+                // Shutdown any old strategies that are being replaced.
+                if (oldDarkStrategyMap != null)
+                {
+                  for (Map.Entry<String, DarkClusterStrategy> oldEntry : oldDarkStrategyMap.entrySet())
+                  {
+                    final DarkClusterStrategy oldStrategy = oldEntry.getValue();
+                    if (oldStrategy != null)
+                    {
+                      oldStrategy.shutdown();
+                    }
+                  }
+                }
+                return darkStrategyMap;
+              });
             }
           }
         });
@@ -300,6 +316,20 @@ public void onClusterRemoved(String clusterName)
     {
       if (_sourceClusterName.equals(clusterName))
       {
+        for (Map.Entry<Integer, Map<String, DarkClusterStrategy>> entry : _partitionToDarkStrategyMap.entrySet())
+        {
+          Map<String, DarkClusterStrategy> strategies = entry.getValue();
+          if (strategies != null)
+          {
+            for (DarkClusterStrategy strategy : strategies.values())
+            {
+              if (strategy != null)
+              {
+                strategy.shutdown();
+              }
+            }
+          }
+        }
         _partitionToDarkStrategyMap.clear();
         _sourceClusterPresent = false;
       }

darkcluster/src/test/java/com/linkedin/darkcluster/TestConstantQpsDarkClusterStrategy.java

@@ -35,10 +35,14 @@
 import java.util.List;
 import java.util.function.Supplier;
 import java.util.stream.IntStream;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
 import org.testng.Assert;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
+import static org.mockito.Mockito.verify;
+
 public class TestConstantQpsDarkClusterStrategy
 {
   private static final String SOURCE_CLUSTER_NAME = "FooCluster";
@@ -176,4 +180,35 @@ static EvictingCircularBuffer getBuffer(Clock clock)
   {
     return new EvictingCircularBuffer(TEST_CAPACITY, TEST_TTL, TEST_TTL_UNIT, clock);
   }
+
+  @Test
+  public void testShutdown()
+  {
+    // Setup mocks
+    ConstantQpsRateLimiter mockRateLimiter = Mockito.mock(ConstantQpsRateLimiter.class);
+    BaseDarkClusterDispatcherImpl mockDispatcher = Mockito.mock(BaseDarkClusterDispatcherImpl.class);
+    MockClusterInfoProvider mockClusterInfoProvider = new MockClusterInfoProvider();
+    mockClusterInfoProvider.putHttpsClusterCount(SOURCE_CLUSTER_NAME, 10);
+    mockClusterInfoProvider.putHttpsClusterCount(DARK_CLUSTER_NAME_ONE, 5);
+
+    // Create strategy instance
+    ConstantQpsDarkClusterStrategy strategy = new ConstantQpsDarkClusterStrategy(
+        SOURCE_CLUSTER_NAME,
+        DARK_CLUSTER_NAME_ONE,
+        50.0f,
+        mockDispatcher,
+        new DoNothingNotifier(),
+        mockClusterInfoProvider,
+        mockRateLimiter);
+
+    // Call shutdown
+    strategy.shutdown();
+
+    // Verify that cancelAll was called with a RuntimeException containing the expected message.
+    ArgumentCaptor<Throwable> throwableCaptor = ArgumentCaptor.forClass(Throwable.class);
+    verify(mockRateLimiter).cancelAll(throwableCaptor.capture());
+    Throwable throwable = throwableCaptor.getValue();
+    Assert.assertTrue(throwable instanceof RuntimeException);
+    Assert.assertEquals(throwable.getMessage(), "Shutting down ConstantQpsDarkClusterStrategy");
+  }
 }

darkcluster/src/test/java/com/linkedin/darkcluster/TestDarkClusterStrategyFactory.java

@@ -26,6 +26,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import com.linkedin.d2.DarkClusterConfig;
 import com.linkedin.d2.DarkClusterStrategyNameArray;
@@ -36,6 +37,7 @@
 import com.linkedin.darkcluster.api.DarkClusterStrategy;
 import com.linkedin.darkcluster.api.DarkClusterStrategyFactory;
 import com.linkedin.darkcluster.api.NoOpDarkClusterStrategy;
+import com.linkedin.darkcluster.impl.ConstantQpsDarkClusterStrategy;
 import com.linkedin.darkcluster.impl.RelativeTrafficMultiplierDarkClusterStrategy;
 import com.linkedin.darkcluster.impl.DarkClusterStrategyFactoryImpl;
 import com.linkedin.darkcluster.impl.DefaultDarkClusterDispatcher;
@@ -314,6 +316,79 @@ public void testStrategyZeroMultiplier()
     Assert.assertTrue(strategy instanceof NoOpDarkClusterStrategy);
   }
 
+  @Test
+  public void testRefreshShutsDownExistingStrategiesPerPartition()
+  {
+    // Build real per-partition strategies (like testUpdateStrategyDarkClusterChange), then refresh and assert the old
+    // strategies were shut down.
+    //
+    // Note: RelativeTrafficMultiplierDarkClusterStrategy has no observable shutdown behavior. ConstantQpsDarkClusterStrategy
+    // does (it calls ConstantQpsRateLimiter.cancelAll), so we use CONSTANT_QPS here with a tracking rate limiter.
+    MockClusterInfoProvider clusterInfoProvider = new MockClusterInfoProvider();
+    Facilities facilities = new MockFacilities(clusterInfoProvider);
+    DarkClusterDispatcher darkClusterDispatcher = new DefaultDarkClusterDispatcher(new MockClient(false));
+    ClockedExecutor executor = new ClockedExecutor();
+
+    AtomicInteger cancelAllInvocations = new AtomicInteger(0);
+    Supplier<ConstantQpsRateLimiter> rateLimiterSupplier = () ->
+        new TrackingConstantQpsRateLimiter(executor, executor, executor,
+            TestConstantQpsDarkClusterStrategy.getBuffer(executor), cancelAllInvocations);
+
+    DarkClusterStrategyFactory strategyFactory = new DarkClusterStrategyFactoryImpl(facilities,
+        SOURCE_CLUSTER_NAME,
+        darkClusterDispatcher,
+        new DoNothingNotifier(),
+        new Random(SEED),
+        new CountingVerifierManager(),
+        rateLimiterSupplier);
+    strategyFactory.start();
+
+    DarkClusterStrategyNameArray strategyList = new DarkClusterStrategyNameArray();
+    strategyList.addAll(Collections.singletonList(CONSTANT_QPS));
+    DarkClusterConfig constantQpsConfig = new DarkClusterConfig()
+        .setDarkClusterStrategyPrioritizedList(strategyList)
+        .setDispatcherOutboundTargetRate(1.0f)
+        .setDispatcherMaxRequestsToBuffer(10)
+        .setDispatcherBufferedRequestExpiryInSeconds(10);
+    clusterInfoProvider.addDarkClusterConfig(SOURCE_CLUSTER_NAME, DARK_CLUSTER_NAME, constantQpsConfig);
+
+    // Trigger a refresh to build partition 0, then lazily create partition 1 so refresh will rebuild both.
+    clusterInfoProvider.notifyListenersClusterAdded(SOURCE_CLUSTER_NAME);
+    DarkClusterStrategy oldP0 = strategyFactory.get(DARK_CLUSTER_NAME, 0);
+    DarkClusterStrategy oldP1 = strategyFactory.get(DARK_CLUSTER_NAME, 1);
+    Assert.assertTrue(oldP0 instanceof ConstantQpsDarkClusterStrategy);
+    Assert.assertTrue(oldP1 instanceof ConstantQpsDarkClusterStrategy);
+
+    // Refresh again: should rebuild both partitions and shut down the old strategies.
+    clusterInfoProvider.notifyListenersClusterAdded(SOURCE_CLUSTER_NAME);
+    Assert.assertEquals(cancelAllInvocations.get(), 2,
+        "Expected existing strategies to be shut down once per partition on refresh");
+
+    Assert.assertTrue(strategyFactory.get(DARK_CLUSTER_NAME, 0) instanceof ConstantQpsDarkClusterStrategy);
+    Assert.assertTrue(strategyFactory.get(DARK_CLUSTER_NAME, 1) instanceof ConstantQpsDarkClusterStrategy);
+  }
+
+  private static final class TrackingConstantQpsRateLimiter extends ConstantQpsRateLimiter
+  {
+    private final AtomicInteger _cancelAllInvocations;
+
+    TrackingConstantQpsRateLimiter(ScheduledExecutorService scheduler,
+        java.util.concurrent.Executor executor,
+        com.linkedin.util.clock.Clock clock,
+        com.linkedin.r2.transport.http.client.EvictingCircularBuffer callbackBuffer,
+        AtomicInteger cancelAllInvocations)
+    {
+      super(scheduler, executor, clock, callbackBuffer);
+      _cancelAllInvocations = cancelAllInvocations;
+    }
+
+    @Override
+    public void cancelAll(Throwable throwable)
+    {
+      _cancelAllInvocations.incrementAndGet();
+    }
+  }
+
   private static class DeletingClusterListener implements LoadBalancerClusterListener
   {
     // handle to MockClusterInfoProvider so it can call triggerCluster actions.

gradle.properties

@@ -1,4 +1,4 @@
-version=29.81.1
+version=29.81.2
 group=com.linkedin.pegasus
 org.gradle.configureondemand=true
 org.gradle.parallel=true