Disable dark traffic dispatching during dark warmup (#1032)

CHANGELOG.md

@@ -14,6 +14,9 @@ and what APIs have changed, if applicable.
 
 ## [Unreleased]
 
+## [29.61.0] - 2024-10-24
+- Disable dark traffic dispatching during dark warmup
+
 ## [29.60.0] - 2024-10-17
 - Restore the old constructor to avoid incompatible issue
 
@@ -5749,7 +5752,8 @@ patch operations can re-use these classes for generating patch messages.
 
 ## [0.14.1]
 
-[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.60.0...master
+[Unreleased]: https://github.com/linkedin/rest.li/compare/v29.61.0...master
+[29.61.0]: https://github.com/linkedin/rest.li/compare/v29.60.0...v29.61.0
 [29.60.0]: https://github.com/linkedin/rest.li/compare/v29.59.0...v29.60.0
 [29.59.0]: https://github.com/linkedin/rest.li/compare/v29.58.11...v29.59.0
 [29.58.11]: https://github.com/linkedin/rest.li/compare/v29.58.10...v29.58.11

d2/src/main/java/com/linkedin/d2/balancer/servers/ZooKeeperAnnouncer.java

@@ -105,7 +105,7 @@ public class ZooKeeperAnnouncer implements D2ServiceDiscoveryEventHelper
 
   // Field to indicate if warm up was started. If it is true, it will try to end the warm up
   // by marking down on ZK if the connection goes down
-  private boolean _isWarmingUp;
+  private volatile boolean _isWarmingUp;
 
   // Field to indicate whether the mark up operation is being retried after a connection loss
   private boolean _isRetryWarmup;
@@ -826,4 +826,11 @@ private ImmutablePair<String, String> getZnodePathAndData(String cluster) {
     }
     return new ImmutablePair<>(nodePath, nodeData);
   }
+
+  /**
+   * Indicates whether the announcement is currently made to the dark warmup cluster.
+   */
+  public boolean isWarmingUp() {
+    return _isWarmingUp;
+  }
 }

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

@@ -16,6 +16,7 @@
 
 package com.linkedin.darkcluster.impl;
 
+import com.linkedin.d2.balancer.servers.ZooKeeperAnnouncer;
 import com.linkedin.darkcluster.api.DarkGateKeeper;
 import com.linkedin.darkcluster.api.DarkRequestHeaderGenerator;
 import com.linkedin.r2.message.rest.RestRequestBuilder;
@@ -61,6 +62,7 @@ public class DarkClusterManagerImpl implements DarkClusterManager
   private final DarkGateKeeper _darkGateKeeper;
   private final List<DarkRequestHeaderGenerator> _darkRequestHeaderGenerators;
   private Map<String, AtomicReference<URIRewriter>> _uriRewriterMap;
+  private final List<ZooKeeperAnnouncer> _announcers;
 
   public DarkClusterManagerImpl(@Nonnull String sourceClusterName, @Nonnull Facilities facilities,
                                 @Nonnull DarkClusterStrategyFactory strategyFactory, String whiteListRegEx,
@@ -88,6 +90,20 @@ public DarkClusterManagerImpl(@Nonnull String sourceClusterName,
       @Nonnull Notifier notifier,
       DarkGateKeeper darkGateKeeper,
       List<DarkRequestHeaderGenerator> darkRequestHeaderGenerators)
+  {
+    this(sourceClusterName, facilities, strategyFactory, whiteListRegEx, blackListRegEx, notifier, darkGateKeeper,
+        darkRequestHeaderGenerators, Collections.emptyList());
+  }
+
+  public DarkClusterManagerImpl(@Nonnull String sourceClusterName,
+      @Nonnull Facilities facilities,
+      @Nonnull DarkClusterStrategyFactory strategyFactory,
+      String whiteListRegEx,
+      String blackListRegEx,
+      @Nonnull Notifier notifier,
+      DarkGateKeeper darkGateKeeper,
+      List<DarkRequestHeaderGenerator> darkRequestHeaderGenerators,
+      @Nonnull List<ZooKeeperAnnouncer> announcers)
   {
     _whiteListRegEx = whiteListRegEx == null ? null : Pattern.compile(whiteListRegEx);
     _blackListRegEx = blackListRegEx == null ? null : Pattern.compile(blackListRegEx);
@@ -99,11 +115,17 @@ public DarkClusterManagerImpl(@Nonnull String sourceClusterName,
     // if null, initialize this to a noop which returns true always
     _darkGateKeeper = darkGateKeeper == null ? DarkGateKeeper.NO_OP_DARK_GATE_KEEPER : darkGateKeeper;
     _darkRequestHeaderGenerators = darkRequestHeaderGenerators == null ? Collections.emptyList() : darkRequestHeaderGenerators;
+    _announcers = announcers;
   }
 
   @Override
   public boolean handleDarkRequest(RestRequest originalRequest, RequestContext originalRequestContext)
   {
+    for (ZooKeeperAnnouncer announcer : _announcers) {
+      if (announcer.isWarmingUp()) {
+          return false;
+      }
+    }
     String uri = originalRequest.getURI().toString();
     boolean darkRequestSent = false;
     try

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

@@ -1,5 +1,6 @@
 package com.linkedin.darkcluster;
 
+import com.linkedin.d2.balancer.servers.ZooKeeperAnnouncer;
 import com.linkedin.darkcluster.api.DarkGateKeeper;
 import com.linkedin.darkcluster.api.DarkRequestHeaderGenerator;
 import java.net.URI;
@@ -19,6 +20,7 @@
 import static com.linkedin.darkcluster.TestDarkClusterStrategyFactory.DARK_CLUSTER_NAME;
 import static com.linkedin.darkcluster.TestDarkClusterStrategyFactory.DARK_CLUSTER_NAME2;
 import static com.linkedin.darkcluster.TestDarkClusterStrategyFactory.SOURCE_CLUSTER_NAME;
+import static org.mockito.Mockito.*;
 
 import java.util.Collections;
 import java.util.Optional;
@@ -151,10 +153,10 @@ public void testWithDarkHeaders() {
     DarkClusterConfig darkClusterConfig = createRelativeTrafficMultiplierConfig(1.0f);
     clusterInfoProvider.addDarkClusterConfig(SOURCE_CLUSTER_NAME, DARK_CLUSTER_NAME, darkClusterConfig);
 
-    DarkClusterStrategyFactory mockStrategyFactory = Mockito.mock(DarkClusterStrategyFactory.class);
-    DarkClusterStrategy mockDarkStrategy = Mockito.mock(DarkClusterStrategy.class);
+    DarkClusterStrategyFactory mockStrategyFactory = mock(DarkClusterStrategyFactory.class);
+    DarkClusterStrategy mockDarkStrategy = mock(DarkClusterStrategy.class);
 
-    DarkRequestHeaderGenerator darkRequestHeaderGenerator = Mockito.mock(DarkRequestHeaderGenerator.class);
+    DarkRequestHeaderGenerator darkRequestHeaderGenerator = mock(DarkRequestHeaderGenerator.class);
     Mockito.when(mockStrategyFactory.get(DARK_CLUSTER_NAME)).thenReturn(mockDarkStrategy);
     Mockito.when(darkRequestHeaderGenerator.get(DARK_CLUSTER_NAME))
         .thenReturn(Optional.of(new DarkRequestHeaderGenerator.HeaderNameValuePair("header", "value")));
@@ -179,6 +181,45 @@ public void testWithDarkHeaders() {
     Assert.assertTrue(status);
   }
 
+  @Test
+  public void testDarkWarmup()
+  {
+    MockClusterInfoProvider clusterInfoProvider = new MockClusterInfoProvider();
+    Facilities facilities = new MockFacilities(clusterInfoProvider);
+    MockStrategyFactory strategyFactory = new MockStrategyFactory();
+    DarkClusterManager darkClusterManager = new DarkClusterManagerImpl(SOURCE_CLUSTER_NAME,
+        facilities,
+        strategyFactory,
+        null,
+        null,
+        new DoNothingNotifier(),
+        null);
+    ZooKeeperAnnouncer mockZkAnnouncer = mock(ZooKeeperAnnouncer.class);
+    when(mockZkAnnouncer.isWarmingUp()).thenReturn(true);
+    DarkClusterManager darkClusterManager2 = new DarkClusterManagerImpl(SOURCE_CLUSTER_NAME,
+        facilities,
+        strategyFactory,
+        null,
+        null,
+        new DoNothingNotifier(),
+        null,
+        null,
+        Collections.singletonList(mockZkAnnouncer));
+
+
+    strategyFactory.start();
+
+    // This configuration will choose the RelativeTrafficMultiplierDarkClusterStrategy
+    DarkClusterConfig darkClusterConfig = createRelativeTrafficMultiplierConfig(1.0f);
+    clusterInfoProvider.addDarkClusterConfig(SOURCE_CLUSTER_NAME, DARK_CLUSTER_NAME, darkClusterConfig);
+
+    RestRequest restRequest = new RestRequestBuilder(URI.create("/test")).setMethod(METHOD_SAFE).build();
+    Assert.assertTrue(darkClusterManager.handleDarkRequest(restRequest, new RequestContext()));
+    Assert.assertFalse(darkClusterManager2.handleDarkRequest(restRequest, new RequestContext()));
+    when(mockZkAnnouncer.isWarmingUp()).thenReturn(false);
+    Assert.assertTrue(darkClusterManager2.handleDarkRequest(restRequest, new RequestContext()));
+  }
+
   private static class MockStrategyFactory implements DarkClusterStrategyFactory
   {
     // Always return true from the strategy so that we can count reliably

gradle.properties

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