CASSJAVA-92: Local DC provided for nodetool clientstats

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

@@ -0,0 +1,34 @@
+/*
+ * 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.NonNull;
+import edu.umd.cs.findbugs.annotations.Nullable;
+import java.util.Map;
+
+/** 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();
+
+  /** Returns map containing details that impact C* node connectivity. */
+  @NonNull
+  Map<String, ?> getStartupConfiguration();
+}

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

@@ -19,24 +19,35 @@
 
 import com.datastax.dse.driver.api.core.config.DseDriverOption;
 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.driver.shaded.guava.common.collect.ImmutableMap;
 import com.datastax.oss.protocol.internal.request.Startup;
 import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import edu.umd.cs.findbugs.annotations.Nullable;
 import java.util.Map;
+import java.util.Optional;
 import java.util.UUID;
 import net.jcip.annotations.Immutable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 @Immutable
 public class StartupOptionsBuilder {
 
   public static final String DRIVER_NAME_KEY = "DRIVER_NAME";
   public static final String DRIVER_VERSION_KEY = "DRIVER_VERSION";
+  public static final String DRIVER_BAGGAGE = "DRIVER_BAGGAGE";
   public static final String APPLICATION_NAME_KEY = "APPLICATION_NAME";
   public static final String APPLICATION_VERSION_KEY = "APPLICATION_VERSION";
   public static final String CLIENT_ID_KEY = "CLIENT_ID";
 
+  private static final Logger LOG = LoggerFactory.getLogger(StartupOptionsBuilder.class);
+  private static final ObjectMapper mapper = new ObjectMapper();
+
   protected final InternalDriverContext context;
   private UUID clientId;
   private String applicationName;
@@ -119,6 +130,8 @@ public Map<String, String> build() {
     if (applicationVersion != null) {
       builder.put(APPLICATION_VERSION_KEY, applicationVersion);
     }
+    // do not cache local DC as it can change within LBP implementation
+    driverBaggage().ifPresent(s -> builder.put(DRIVER_BAGGAGE, s));
 
     return builder.build();
   }
@@ -142,4 +155,27 @@ protected String getDriverName() {
   protected String getDriverVersion() {
     return Session.OSS_DRIVER_COORDINATES.getVersion().toString();
   }
+
+  private Optional<String> driverBaggage() {
+    ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
+    for (Map.Entry<String, LoadBalancingPolicy> entry :
+        context.getLoadBalancingPolicies().entrySet()) {
+      getDriverBaggage(entry.getValue()).ifPresent(baggage -> builder.put(entry.getKey(), baggage));
+    }
+    try {
+      return Optional.of(mapper.writeValueAsString(builder.build()));
+    } catch (Exception e) {
+      LOG.warn("Failed to construct startup driver baggage", e);
+      return Optional.empty();
+    }
+  }
+
+  private Optional<Map<String, ?>> getDriverBaggage(LoadBalancingPolicy loadBalancingPolicy) {
+    if (loadBalancingPolicy instanceof LocalDcAwareLoadBalancingPolicy) {
+      LocalDcAwareLoadBalancingPolicy dcAwareLbp =
+          (LocalDcAwareLoadBalancingPolicy) loadBalancingPolicy;
+      return Optional.of(dcAwareLbp.getStartupConfiguration());
+    }
+    return Optional.empty();
+  }
 }

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;
@@ -45,6 +46,7 @@
 import com.datastax.oss.driver.internal.core.util.collection.QueryPlan;
 import com.datastax.oss.driver.internal.core.util.collection.SimpleQueryPlan;
 import com.datastax.oss.driver.shaded.guava.common.base.Predicates;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
 import com.datastax.oss.driver.shaded.guava.common.collect.Lists;
 import com.datastax.oss.driver.shaded.guava.common.collect.Sets;
 import edu.umd.cs.findbugs.annotations.NonNull;
@@ -99,7 +101,7 @@
  * DefaultLoadBalancingPolicy}</b>.
  */
 @ThreadSafe
-public class BasicLoadBalancingPolicy implements LoadBalancingPolicy {
+public class BasicLoadBalancingPolicy implements LocalDcAwareLoadBalancingPolicy {
 
   private static final Logger LOG = LoggerFactory.getLogger(BasicLoadBalancingPolicy.class);
 
@@ -155,10 +157,28 @@ public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String
    * Before initialization, this method always returns null.
    */
   @Nullable
-  protected String getLocalDatacenter() {
+  @Override
+  public String getLocalDatacenter() {
     return localDc;
   }
 
+  @NonNull
+  @Override
+  public Map<String, ?> getStartupConfiguration() {
+    ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
+    builder.put("localDc", localDc);
+    if (!preferredRemoteDcs.isEmpty()) {
+      builder.put("preferredRemoteDcs", preferredRemoteDcs);
+    }
+    if (allowDcFailoverForLocalCl) {
+      builder.put("allowDcFailoverForLocalCl", allowDcFailoverForLocalCl);
+    }
+    if (maxNodesPerRemoteDc > 0) {
+      builder.put("maxNodesPerRemoteDc", maxNodesPerRemoteDc);
+    }
+    return ImmutableMap.of(BasicLoadBalancingPolicy.class.getSimpleName(), builder.build());
+  }
+
   /** @return The nodes currently considered as live. */
   protected NodeSet getLiveNodes() {
     return liveNodes;

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

@@ -34,6 +34,7 @@
 import com.datastax.oss.driver.internal.core.util.collection.QueryPlan;
 import com.datastax.oss.driver.internal.core.util.collection.SimpleQueryPlan;
 import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
 import com.datastax.oss.driver.shaded.guava.common.collect.MapMaker;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
@@ -350,4 +351,13 @@ private boolean hasSufficientResponses(long now) {
       return this.oldest - threshold >= 0;
     }
   }
+
+  @NonNull
+  @Override
+  public Map<String, ?> getStartupConfiguration() {
+    Map<String, ?> parent = super.getStartupConfiguration();
+    return ImmutableMap.of(
+        DefaultLoadBalancingPolicy.class.getSimpleName(),
+        parent.get(BasicLoadBalancingPolicy.class.getSimpleName()));
+  }
 }

core/src/main/java/com/datastax/oss/driver/internal/core/util/Strings.java

@@ -81,6 +81,18 @@ public static String doubleQuote(String value) {
     return quote(value, '"');
   }
 
+  /**
+   * Double quote the given string; double quotes are escaped. If the given string is null, this
+   * method returns ({@code null}).
+   *
+   * @param value The value to double quote.
+   * @return The double quoted string.
+   */
+  public static String doubleQuoteNullable(String value) {
+    if (value == null) return null;
+    return quote(value, '"');
+  }
+
   /**
    * Unquote the given string if it is double quoted; double quotes are unescaped. If the given
    * string is not double quoted, it is returned without any modification.