Dark Cluster Partition Support

d2-int-test/src/test/java/com/linkedin/d2/loadbalancer/TestDynamicClient.java

@@ -75,6 +75,7 @@
 public class TestDynamicClient extends D2BaseTest {
   private static final String D2_CONFIG_FILE = "d2_config_example.json";
   private static final String ZK_HOST = "127.0.0.1";
+  private static final String ECHO_SERVER_HOST = "127.0.0.1";
   private static final int ECHO_SERVER_PORT_START = 2851;
   private static final int NUMBER_OF_HOSTS = 5;
   private static final int NUMBER_OF_THREADS = 10;

darkcluster-test-api/src/main/java/com/linkedin/darkcluster/MockClusterInfoProvider.java

@@ -40,7 +40,7 @@ public class MockClusterInfoProvider implements ClusterInfoProvider
   public int getClusterCount(String clusterName, String scheme, int partitionId)
     throws ServiceUnavailableException
   {
-    return 0;
+    return clusterHttpsCount.getOrDefault(clusterName, 1);
   }
 
   @Override
@@ -50,6 +50,7 @@ public int getHttpsClusterCount(String clusterName)
     return clusterHttpsCount.getOrDefault(clusterName, 1);
   }
 
+
   @Override
   public DarkClusterConfigMap getDarkClusterConfigMap(String clusterName)
     throws ServiceUnavailableException

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

@@ -27,9 +27,10 @@ public interface DarkClusterStrategyFactory
   /**
    * get retrieves the {@link DarkClusterStrategy} corresponding to the darkClusterName.
    * @param darkClusterName darkClusterName to look up
+   * @param partitionId partition id to scope the strategy to
    * @return {@link DarkClusterStrategy}
    */
-  DarkClusterStrategy get(String darkClusterName);
+  DarkClusterStrategy get(String darkClusterName, int partitionId);
 
   /**
    * Do any actions necessary to start the DarkClusterStrategyFactory.

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

@@ -55,12 +55,15 @@ public class ConstantQpsDarkClusterStrategy implements DarkClusterStrategy
   private final ClusterInfoProvider _clusterInfoProvider;
   private final ConstantQpsRateLimiter _rateLimiter;
 
+
+
   private static final long ONE_SECOND_PERIOD = TimeUnit.SECONDS.toMillis(1);
   private static final int NUM_REQUESTS_TO_SEND_PER_RATE_LIMITER_CYCLE = 1;
 
   public ConstantQpsDarkClusterStrategy(@Nonnull String originalClusterName, @Nonnull String darkClusterName,
       @Nonnull Float darkClusterPerHostQps, @Nonnull BaseDarkClusterDispatcher baseDarkClusterDispatcher,
-      @Nonnull Notifier notifier, @Nonnull ClusterInfoProvider clusterInfoProvider, @Nonnull ConstantQpsRateLimiter rateLimiter)
+      @Nonnull Notifier notifier, @Nonnull ClusterInfoProvider clusterInfoProvider,
+      @Nonnull ConstantQpsRateLimiter rateLimiter)
   {
     _originalClusterName = originalClusterName;
     _darkClusterName = darkClusterName;
@@ -74,6 +77,7 @@ public ConstantQpsDarkClusterStrategy(@Nonnull String originalClusterName, @Nonn
   @Override
   public boolean handleRequest(RestRequest originalRequest, RestRequest darkRequest, RequestContext requestContext)
   {
+    // ClusterInfoProvider is already partition-aware, so we can use it directly
     float sendRate = getSendRate();
     // set burst in such a way that requests are dispatched evenly across the ONE_SECOND_PERIOD
     int burst = (int) Math.max(1, Math.ceil(sendRate / ONE_SECOND_PERIOD));
@@ -134,7 +138,6 @@ private float getSendRate()
       {
         return (numDarkClusterInstances * _darkClusterPerHostQps) / numSourceClusterInstances;
       }
-
       return 0F;
     }
     catch (ServiceUnavailableException e)
@@ -171,4 +174,6 @@ public void onSuccess(None result)
     });
     return true;
   }
+
+
 }

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

@@ -144,13 +144,14 @@ public boolean handleDarkRequest(RestRequest originalRequest, RequestContext ori
         RestRequest reqCopy = originalRequest.builder().build();
         RequestContext newRequestContext = new RequestContext(originalRequestContext);
         DarkClusterConfigMap configMap = _facilities.getClusterInfoProvider().getDarkClusterConfigMap(_sourceClusterName);
+        int partitionId = getPartitionId(originalRequest);
         for (String darkClusterName : configMap.keySet())
         {
           if (_darkGateKeeper.shouldDispatchToDark(originalRequest, originalRequestContext, darkClusterName))
           {
             RestRequest newD2Request = rewriteRequest(reqCopy, darkClusterName);
             // now find the strategy appropriate for each dark cluster
-            DarkClusterStrategy strategy = _darkClusterStrategyFactory.get(darkClusterName);
+            DarkClusterStrategy strategy = _darkClusterStrategyFactory.get(darkClusterName, partitionId);
             darkRequestSent |= strategy.handleRequest(reqCopy, newD2Request, newRequestContext);
           }
         }
@@ -164,6 +165,20 @@ public boolean handleDarkRequest(RestRequest originalRequest, RequestContext ori
     return darkRequestSent;
   }
 
+  private int getPartitionId(RestRequest request)
+  {
+    try
+    {
+      String serviceName = com.linkedin.d2.balancer.util.LoadBalancerUtil.getServiceNameFromUri(request.getURI());
+      com.linkedin.d2.balancer.util.partitions.PartitionAccessor accessor = _facilities.getPartitionInfoProvider().getPartitionAccessor(serviceName);
+      return accessor.getPartitionId(request.getURI());
+    }
+    catch (Throwable t)
+    {
+      return com.linkedin.d2.balancer.util.partitions.DefaultPartitionAccessor.DEFAULT_PARTITION_ID;
+    }
+  }
+
   /**
    * isSafe returns true if the underlying HttpMethod has the expectation of only doing retrieval with no side effects. For further details,
    * see {@link HttpMethod}

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

@@ -23,6 +23,8 @@
 import java.util.Random;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import com.linkedin.d2.balancer.ServiceUnavailableException;
+import com.linkedin.d2.balancer.clusterfailout.FailoutConfig;
 
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
@@ -62,7 +64,9 @@ public class DarkClusterStrategyFactoryImpl implements DarkClusterStrategyFactor
   private final DarkClusterDispatcher _darkClusterDispatcher;
   private final Notifier _notifier;
 
-  private final Map<String, DarkClusterStrategy> _darkStrategyMap;
+  // Map of partition ID to dark cluster name to list of dark cluster strategies for that partition
+  private final Map<Integer, Map<String, DarkClusterStrategy>> _partitionToDarkStrategyMap;
+  private volatile boolean _sourceClusterPresent = false;
   private final Random _random;
   private final LoadBalancerClusterListener _clusterListener;
   private final DarkClusterVerifierManager _verifierManager;
@@ -79,7 +83,7 @@ public DarkClusterStrategyFactoryImpl(@Nonnull Facilities facilities,
     _facilities = facilities;
     _sourceClusterName = sourceClusterName;
     _notifier = notifier;
-    _darkStrategyMap = new ConcurrentHashMap<>();
+    _partitionToDarkStrategyMap = new ConcurrentHashMap<>();
     _random = random;
     _darkClusterDispatcher = darkClusterDispatcher;
     _verifierManager = verifierManager;
@@ -138,16 +142,44 @@ public void shutdown()
    * @return darkClusterStrategy to use.
    */
   @Override
-  public DarkClusterStrategy get(@Nonnull String darkClusterName)
+  public DarkClusterStrategy get(@Nonnull String darkClusterName, int partitionId)
   {
-    return _darkStrategyMap.getOrDefault(darkClusterName, NO_OP_DARK_CLUSTER_STRATEGY);
+    Map<String, DarkClusterStrategy> darkMap = _partitionToDarkStrategyMap.computeIfAbsent(partitionId, k -> new ConcurrentHashMap<>());
+    DarkClusterStrategy existing = darkMap.get(darkClusterName);
+
+    if (existing != null)
+    {
+      return existing;
+    }
+
+    try
+    {
+      // Lazily create the strategy if it doesn't exist.
+      DarkClusterConfigMap darkClusterConfigMap = _facilities.getClusterInfoProvider().getDarkClusterConfigMap(_sourceClusterName);
+      if (darkClusterConfigMap != null && darkClusterConfigMap.containsKey(darkClusterName))
+      {
+        DarkClusterConfig config = darkClusterConfigMap.get(darkClusterName);
+        DarkClusterStrategy strategy = createStrategy(darkClusterName, config, partitionId);
+        darkMap.put(darkClusterName, strategy);
+        return strategy;
+      }
+    }
+    catch (Throwable t)
+    {
+      LOG.warn("Unable to get DarkClusterConfigMap for source cluster: " + _sourceClusterName, t);
+    }
+    return NO_OP_DARK_CLUSTER_STRATEGY;
   }
 
   /**
    * In the future, additional strategies can be added, and the logic here can choose the appropriate one based on the config values.
    */
-  private DarkClusterStrategy createStrategy(String darkClusterName, DarkClusterConfig darkClusterConfig)
+  private DarkClusterStrategy createStrategy(String darkClusterName, DarkClusterConfig darkClusterConfig, int partitionId)
   {
+    // Create partition-aware ClusterInfoProvider that filters cluster information for this specific partition
+    com.linkedin.d2.balancer.util.ClusterInfoProvider partitionAwareProvider =
+        new PartitionAwareClusterInfoProvider(_facilities.getClusterInfoProvider(), partitionId);
+
     if (darkClusterConfig.hasDarkClusterStrategyPrioritizedList())
     {
       DarkClusterStrategyNameArray strategyList = darkClusterConfig.getDarkClusterStrategyPrioritizedList();
@@ -162,7 +194,8 @@ private DarkClusterStrategy createStrategy(String darkClusterName, DarkClusterCo
                   new BaseDarkClusterDispatcherImpl(darkClusterName, _darkClusterDispatcher, _notifier, _verifierManager);
               return new RelativeTrafficMultiplierDarkClusterStrategy(_sourceClusterName, darkClusterName,
                                                                       darkClusterConfig.getMultiplier(), baseDarkClusterDispatcher,
-                                                                      _notifier, _facilities.getClusterInfoProvider(), _random);
+                                                                      _notifier, partitionAwareProvider,
+                                                                      _random);
             }
             break;
           case IDENTICAL_TRAFFIC:
@@ -172,7 +205,8 @@ private DarkClusterStrategy createStrategy(String darkClusterName, DarkClusterCo
                   new BaseDarkClusterDispatcherImpl(darkClusterName, _darkClusterDispatcher, _notifier, _verifierManager);
               return new IdenticalTrafficMultiplierDarkClusterStrategy(_sourceClusterName, darkClusterName,
                                                                        darkClusterConfig.getMultiplier(), baseDarkClusterDispatcher,
-                                                                       _notifier, _facilities.getClusterInfoProvider(), _random);
+                                                                       _notifier, partitionAwareProvider,
+                                                                       _random);
             }
             break;
           case CONSTANT_QPS:
@@ -191,7 +225,7 @@ private DarkClusterStrategy createStrategy(String darkClusterName, DarkClusterCo
               rateLimiter.setBufferTtl(darkClusterConfig.getDispatcherBufferedRequestExpiryInSeconds(), ChronoUnit.SECONDS);
               return new ConstantQpsDarkClusterStrategy(_sourceClusterName, darkClusterName,
                   darkClusterConfig.getDispatcherOutboundTargetRate(), baseDarkClusterDispatcher,
-                  _notifier, _facilities.getClusterInfoProvider(), rateLimiter);
+                  _notifier, partitionAwareProvider, rateLimiter);
             }
             break;
           default:
@@ -215,6 +249,7 @@ public void onClusterAdded(String updatedClusterName)
       // pertinent dark cluster strategy properties are contained there.
       if (_sourceClusterName.equals(updatedClusterName))
       {
+        _sourceClusterPresent = true;
         _facilities.getClusterInfoProvider().getDarkClusterConfigMap(_sourceClusterName, new Callback<DarkClusterConfigMap>()
         {
           @Override
@@ -227,23 +262,23 @@ public void onError(Throwable e)
           @Override
           public void onSuccess(DarkClusterConfigMap updatedDarkConfigMap)
           {
-            Set<String> oldDarkStrategySet = _darkStrategyMap.keySet();
-            Set<String> updatedDarkClusterConfigKeySet = updatedDarkConfigMap.keySet();
-            // Any old strategy entry that isn't in the "updated" set should be removed from the strategyMap.
-            oldDarkStrategySet.removeAll(updatedDarkClusterConfigKeySet);
-            for (String darkClusterToRemove : oldDarkStrategySet)
+            // Determine partitions to (re)build. If none exist yet, ensure default partition is initialized.
+            java.util.Set<Integer> partitions = new java.util.HashSet<>(_partitionToDarkStrategyMap.keySet());
+            if (partitions.isEmpty())
             {
-              _darkStrategyMap.remove(darkClusterToRemove);
-              LOG.info("Removed dark cluster strategy for dark cluster: " + darkClusterToRemove + ", source cluster: " + _sourceClusterName);
+              partitions.add(com.linkedin.d2.balancer.util.partitions.DefaultPartitionAccessor.DEFAULT_PARTITION_ID);
             }
 
-            // Now update/add the dark clusters.
-            for (Map.Entry<String, DarkClusterConfig> entry : updatedDarkConfigMap.entrySet())
+            for (int partitionId : partitions)
             {
-              String darkClusterToAdd = entry.getKey();
-              // For simplicity, we refresh all strategies since we expect cluster updates to be rare and refresh to be cheap.
-              _darkStrategyMap.put(darkClusterToAdd, createStrategy(darkClusterToAdd, entry.getValue()));
-              LOG.info("Created new strategy for dark cluster: " + darkClusterToAdd + ", source cluster: " + _sourceClusterName);
+              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);
+              }
+              _partitionToDarkStrategyMap.put(partitionId, darkStrategyMap);
             }
           }
         });
@@ -258,8 +293,58 @@ public void onClusterRemoved(String clusterName)
     {
       if (_sourceClusterName.equals(clusterName))
       {
-        _darkStrategyMap.clear();
+        _partitionToDarkStrategyMap.clear();
+        _sourceClusterPresent = false;
       }
     }
   }
+
+  /**
+   * Partition-aware wrapper around a {@link ClusterInfoProvider} that filters cluster information
+   * for a specific partition before forwarding to the strategies.
+   */
+  private static final class PartitionAwareClusterInfoProvider implements com.linkedin.d2.balancer.util.ClusterInfoProvider {
+    private final com.linkedin.d2.balancer.util.ClusterInfoProvider _delegate;
+    private final int _partitionId;
+
+    PartitionAwareClusterInfoProvider(com.linkedin.d2.balancer.util.ClusterInfoProvider delegate, int partitionId) {
+      _delegate = delegate;
+      _partitionId = partitionId;
+    }
+
+    @Override
+    public int getHttpsClusterCount(String clusterName) throws ServiceUnavailableException {
+      return getClusterCount(clusterName, com.linkedin.d2.balancer.properties.PropertyKeys.HTTPS_SCHEME, _partitionId);
+    }
+
+    @Override
+    public int getClusterCount(String clusterName, String scheme, int partitionId) throws ServiceUnavailableException {
+      return _delegate.getClusterCount(clusterName, scheme, partitionId);
+    }
+
+    @Override
+    public DarkClusterConfigMap getDarkClusterConfigMap(String clusterName) throws ServiceUnavailableException {
+      return _delegate.getDarkClusterConfigMap(clusterName);
+    }
+
+    @Override
+    public void getDarkClusterConfigMap(String clusterName, Callback<DarkClusterConfigMap> callback) {
+      _delegate.getDarkClusterConfigMap(clusterName, callback);
+    }
+
+    @Override
+    public void registerClusterListener(LoadBalancerClusterListener clusterListener) {
+      _delegate.registerClusterListener(clusterListener);
+    }
+
+    @Override
+    public void unregisterClusterListener(LoadBalancerClusterListener clusterListener) {
+      _delegate.unregisterClusterListener(clusterListener);
+    }
+
+    @Override
+    public FailoutConfig getFailoutConfig(String clusterName) {
+      return _delegate.getFailoutConfig(clusterName);
+    }
+  }
 }

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

@@ -107,8 +107,7 @@ private int getNumDuplicateRequests(RequestContext requestContext)
   {
     try
     {
-      // Only support https for now. http support can be added later if truly needed, but would be non-ideal
-      // because potentially both dark and source would have to be configured.
+      // ClusterInfoProvider is already partition-aware, so we use the default partition ID
       int numDarkClusterInstances = _clusterInfoProvider.getHttpsClusterCount(_darkClusterName);
       int numSourceClusterInstances = _clusterInfoProvider.getHttpsClusterCount(_originalClusterName);
       float randomNumber;

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

@@ -110,8 +110,6 @@ private int getNumDuplicateRequests()
   {
     try
     {
-      // Only support https for now. http support can be added later if truly needed, but would be non-ideal
-      // because potentially both dark and source would have to be configured.
       int numDarkClusterInstances = _clusterInfoProvider.getHttpsClusterCount(_darkClusterName);
       int numSourceClusterInstances = _clusterInfoProvider.getHttpsClusterCount(_originalClusterName);
       if (numSourceClusterInstances != 0)