Lookup data center from LBP

Author: lukasz-antoniak

core/src/main/java/com/datastax/oss/driver/api/core/loadbalancing/LocalDcAwareLoadBalancingPolicy.java

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.datastax.oss.driver.api.core.loadbalancing;
+
+import edu.umd.cs.findbugs.annotations.Nullable;
+
+/** Load balancing policy taking into account local datacenter of the application. */
+public interface LocalDcAwareLoadBalancingPolicy extends LoadBalancingPolicy {
+
+  /** Returns the local datacenter name, if known; empty otherwise. */
+  @Nullable
+  String getLocalDatacenter();
+}

core/src/main/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilder.java

@@ -20,6 +20,8 @@
 import com.datastax.dse.driver.api.core.config.DseDriverOption;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
+import com.datastax.oss.driver.api.core.loadbalancing.LoadBalancingPolicy;
+import com.datastax.oss.driver.api.core.loadbalancing.LocalDcAwareLoadBalancingPolicy;
 import com.datastax.oss.driver.api.core.session.Session;
 import com.datastax.oss.driver.api.core.uuid.Uuids;
 import com.datastax.oss.protocol.internal.request.Startup;
@@ -43,7 +45,6 @@ public class StartupOptionsBuilder {
   private UUID clientId;
   private String applicationName;
   private String applicationVersion;
-  private String applicationLocalDc;
 
   public StartupOptionsBuilder(InternalDriverContext context) {
     this.context = context;
@@ -122,9 +123,8 @@ public Map<String, String> build() {
     if (applicationVersion != null) {
       builder.put(APPLICATION_VERSION_KEY, applicationVersion);
     }
-    if (applicationLocalDc == null) {
-      applicationLocalDc = localDc(config);
-    }
+    // do not cache local DC as it can change within LBP implementation
+    String applicationLocalDc = localDc(config);
     if (applicationLocalDc != null) {
       builder.put(DRIVER_LOCAL_DC, applicationLocalDc);
     }
@@ -154,10 +154,16 @@ protected String getDriverVersion() {
 
   private String localDc(DriverExecutionProfile profile) {
     String dc = context.getLocalDatacenter(profile.getName()); // DC set programmatically
+    if (dc == null) {
+      // DC from load balancing policy
+      LoadBalancingPolicy lbp = context.getLoadBalancingPolicy(DriverExecutionProfile.DEFAULT_NAME);
+      if (lbp instanceof LocalDcAwareLoadBalancingPolicy) {
+        dc = ((LocalDcAwareLoadBalancingPolicy) lbp).getLocalDatacenter();
+      }
+    }
     if (dc == null && profile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER)) {
-      dc =
-          profile.getString(
-              DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER); // DC from configuration
+      // DC from configuration
+      dc = profile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER);
     }
     return dc;
   }

core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java

@@ -24,6 +24,7 @@
 import com.datastax.oss.driver.api.core.context.DriverContext;
 import com.datastax.oss.driver.api.core.cql.Statement;
 import com.datastax.oss.driver.api.core.loadbalancing.LoadBalancingPolicy;
+import com.datastax.oss.driver.api.core.loadbalancing.LocalDcAwareLoadBalancingPolicy;
 import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance;
 import com.datastax.oss.driver.api.core.loadbalancing.NodeDistanceEvaluator;
 import com.datastax.oss.driver.api.core.metadata.Node;
@@ -99,7 +100,7 @@
  * DefaultLoadBalancingPolicy}</b>.
  */
 @ThreadSafe
-public class BasicLoadBalancingPolicy implements LoadBalancingPolicy {
+public class BasicLoadBalancingPolicy implements LocalDcAwareLoadBalancingPolicy {
 
   private static final Logger LOG = LoggerFactory.getLogger(BasicLoadBalancingPolicy.class);
 
@@ -155,7 +156,8 @@ public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String
    * Before initialization, this method always returns null.
    */
   @Nullable
-  protected String getLocalDatacenter() {
+  @Override
+  public String getLocalDatacenter() {
     return localDc;
   }