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/DarkClusterManagerImpl.java

@@ -16,7 +16,9 @@
 
 package com.linkedin.darkcluster.impl;
 
+import com.linkedin.d2.balancer.ServiceUnavailableException;
 import com.linkedin.d2.balancer.servers.ZooKeeperAnnouncer;
+import com.linkedin.d2.balancer.util.partitions.PartitionAccessException;
 import com.linkedin.darkcluster.api.DarkGateKeeper;
 import com.linkedin.darkcluster.api.DarkRequestHeaderGenerator;
 import com.linkedin.r2.message.rest.RestRequestBuilder;
@@ -149,21 +151,37 @@ public boolean handleDarkRequest(RestRequest originalRequest, RequestContext ori
           if (_darkGateKeeper.shouldDispatchToDark(originalRequest, originalRequestContext, darkClusterName))
           {
             RestRequest newD2Request = rewriteRequest(reqCopy, darkClusterName);
+            int partitionId = getPartitionId(newD2Request);
             // 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);
           }
         }
 
       }
     }
-    catch (Throwable e)
+    catch (RuntimeException | ServiceUnavailableException e)
     {
       _notifier.notify(() -> new RuntimeException("DarkCanaryDispatcherFilter failed to send request: " + uri, e));
     }
     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 (RuntimeException | PartitionAccessException | ServiceUnavailableException e)
+    {
+      _log.error("Cannot find partition id for request: {}, defaulting to 0", request.getURI(), e);
+      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,51 @@ 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 strategy = darkMap.computeIfAbsent(darkClusterName, k -> {
+      if (_sourceClusterPresent)
+      {
+        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);
+            return createStrategy(darkClusterName, config, partitionId);
+          }
+        }
+        catch (RuntimeException | ServiceUnavailableException t)
+        {
+          LOG.warn("Unable to get DarkClusterConfigMap for source cluster: " + _sourceClusterName, t);
+        }
+      }
+      return null;
+    });
+
+    if (strategy == null)
+    {
+      LOG.debug("No strategy found for dark cluster: " + darkClusterName + ", partition: " + partitionId
+          + ", source cluster: " + _sourceClusterName + ". Returning NO_OP strategy.");
+      return NO_OP_DARK_CLUSTER_STRATEGY;
+    }
+    return 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 +201,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 +212,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 +232,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 +256,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 +269,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 +300,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);
+    }
+  }
 }