Adding AllowedDC List to DatacenterAwarePolicy

Ericsson · Aug 16, 2024 · 645d2ca · 645d2ca
1 parent ba16141
commit 645d2ca
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 8 deletions.
diff --git a/...m/ericsson/bss/cassandra/ecchronos/connection/impl/builders/DistributedNativeBuilder.java b/...m/ericsson/bss/cassandra/ecchronos/connection/impl/builders/DistributedNativeBuilder.java
@@ -339,6 +339,7 @@ private static ProgrammaticDriverConfigLoaderBuilder loaderBuilder(
                         SCHEMA_REFRESHED_KEYSPACES);
         if (builder.myType.equals(ConnectionType.datacenterAware))
         {
+            DataCenterAwarePolicy.setAllowedDcs(builder.myDatacenterAware);
             loaderBuilder.withString(DefaultDriverOption.LOAD_BALANCING_POLICY_CLASS,
                 builder.myDatacenterAwarePolicy.getCanonicalName());
             loaderBuilder.withInt(DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_MAX_NODES_PER_REMOTE_DC,

diff --git a/.../src/main/java/com/ericsson/bss/cassandra/ecchronos/connection/DataCenterAwarePolicy.java b/.../src/main/java/com/ericsson/bss/cassandra/ecchronos/connection/DataCenterAwarePolicy.java
@@ -52,12 +52,21 @@ public class DataCenterAwarePolicy extends DefaultLoadBalancingPolicy
 
     private final ConcurrentMap<String, CopyOnWriteArrayList<Node>> myPerDcLiveNodes = new ConcurrentHashMap<>();
     private final AtomicInteger myIndex = new AtomicInteger();
+    private static List<String> myAllowedDcs;
 
     public DataCenterAwarePolicy(final DriverContext context, final String profileName)
     {
         super(context, profileName);
     }
 
+    public static void setAllowedDcs(final List<String> allowedDcs)
+    {
+        if (allowedDcs != null)
+        {
+            myAllowedDcs = allowedDcs;
+        }
+    }
+
     @Override
     public final void init(final Map<UUID, Node> nodes, final DistanceReporter distanceReporter)
     {
@@ -168,6 +177,10 @@ private Queue<Node> getQueryPlan(final String datacenter, final Set<Node> replic
     public NodeDistance distance(final Node node, final String dataCenter)
     {
         String dc = getDc(node);
+        if (!getLocalDatacenter().equals(dc) && myAllowedDcs != null && !myAllowedDcs.contains(dc))
+        {
+            return NodeDistance.IGNORED;
+        }
         if (dc.equals(dataCenter))
         {
             return NodeDistance.LOCAL;
@@ -184,7 +197,11 @@ public NodeDistance distance(final Node node, final String dataCenter)
 
     private Queue<Node> getFallbackQueryPlan(final String dataCenter)
     {
-        CopyOnWriteArrayList<Node> localLiveNodes = myPerDcLiveNodes.get(dataCenter);
+        CopyOnWriteArrayList<Node> localLiveNodes = null;
+        if (getLocalDatacenter().equals(dataCenter) || myAllowedDcs == null || myAllowedDcs.contains(dataCenter))
+        {
+            localLiveNodes = myPerDcLiveNodes.get(dataCenter);
+        }
         final List<Node> nodes = localLiveNodes == null ? Collections.emptyList() : cloneList(localLiveNodes);
         final int startIndex = myIndex.getAndIncrement();
         int index = startIndex;

diff --git a/.../test/java/com/ericsson/bss/cassandra/ecchronos/connection/TestDataCenterAwarePolicy.java b/.../test/java/com/ericsson/bss/cassandra/ecchronos/connection/TestDataCenterAwarePolicy.java
@@ -31,20 +31,14 @@
 import com.datastax.oss.driver.internal.core.ConsistencyLevelRegistry;
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.metadata.MetadataManager;
+import java.util.*;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 
 import java.nio.ByteBuffer;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Queue;
-import java.util.Set;
-import java.util.UUID;
 import java.util.concurrent.CopyOnWriteArrayList;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -57,6 +51,7 @@ public class TestDataCenterAwarePolicy
 {
     private final String myLocalDc = "DC1";
     private final String myRemoteDc = "DC2";
+    private final String[] myAllowedDcs = {"DC1", "DC2"};
 
     @Mock
     private Session mySessionMock;
@@ -138,6 +133,7 @@ public void setup()
     @Test
     public void testDistanceHost()
     {
+        DataCenterAwarePolicy.setAllowedDcs(null);
         DataCenterAwarePolicy policy = new DataCenterAwarePolicy(myDriverContextMock, "");
         policy.init(myNodes, myDistanceReporterMock);
 
@@ -154,6 +150,23 @@ public void testDistanceHost()
         assertThat(distance5).isEqualTo(NodeDistance.IGNORED);
     }
 
+    @Test
+    public void testDistanceHostWithAllowedDcs()
+    {
+        DataCenterAwarePolicy.setAllowedDcs(List.of(myAllowedDcs));
+
+        DataCenterAwarePolicy policy = new DataCenterAwarePolicy(myDriverContextMock, "");
+        policy.init(myNodes, myDistanceReporterMock);
+
+        NodeDistance distance1 = policy.distance(myNodeDC3Mock, myLocalDc);
+        NodeDistance distance2 = policy.distance(myNodeDC1Mock, myLocalDc);
+        NodeDistance distance3 = policy.distance(myNodeDC2Mock, myLocalDc);
+
+        assertThat(distance1).isEqualTo(NodeDistance.IGNORED);
+        assertThat(distance2).isEqualTo(NodeDistance.LOCAL);
+        assertThat(distance3).isEqualTo(NodeDistance.REMOTE);
+    }
+
     @Test
     public void testNewQueryPlanWithNotPartitionAwareStatement()
     {