CASSJAVA-92: Local DC provided for nodetool clientstats

patch by Lukasz Antoniak; reviewed by Bret McGuire and Abe Ratnofsky for CASSJAVA-92

Author: lukasz-antoniak

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

@@ -24,6 +24,7 @@
 import com.datastax.oss.driver.api.core.tracker.RequestTracker;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
+import java.util.Collections;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Queue;
@@ -76,6 +77,12 @@ default Optional<RequestTracker> getRequestTracker() {
    */
   void init(@NonNull Map<UUID, Node> nodes, @NonNull DistanceReporter distanceReporter);
 
+  /** Returns map containing details that impact C* node connectivity. */
+  @NonNull
+  default Map<String, ?> getStartupConfiguration() {
+    return Collections.emptyMap();
+  }
+
   /**
    * Returns the coordinators to use for a new query.
    *

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

@@ -216,8 +216,8 @@ public class DefaultDriverContext implements InternalDriverContext {
       new LazyReference<>("metricIdGenerator", this::buildMetricIdGenerator, cycleDetector);
   private final LazyReference<RequestThrottler> requestThrottlerRef =
       new LazyReference<>("requestThrottler", this::buildRequestThrottler, cycleDetector);
-  private final LazyReference<Map<String, String>> startupOptionsRef =
-      new LazyReference<>("startupOptions", this::buildStartupOptions, cycleDetector);
+  private final LazyReference<StartupOptionsBuilder> startupOptionsRef =
+      new LazyReference<>("startupOptionsFactory", this::buildStartupOptionsFactory, cycleDetector);
   private final LazyReference<NodeStateListener> nodeStateListenerRef;
   private final LazyReference<SchemaChangeListener> schemaChangeListenerRef;
   private final LazyReference<RequestTracker> requestTrackerRef;
@@ -335,16 +335,15 @@ public DefaultDriverContext(
   }
 
   /**
-   * Builds a map of options to send in a Startup message.
+   * Returns builder of options to send in a Startup message.
    *
    * @see #getStartupOptions()
    */
-  protected Map<String, String> buildStartupOptions() {
+  protected StartupOptionsBuilder buildStartupOptionsFactory() {
     return new StartupOptionsBuilder(this)
         .withClientId(startupClientId)
         .withApplicationName(startupApplicationName)
-        .withApplicationVersion(startupApplicationVersion)
-        .build();
+        .withApplicationVersion(startupApplicationVersion);
   }
 
   protected Map<String, LoadBalancingPolicy> buildLoadBalancingPolicies() {
@@ -1013,7 +1012,8 @@ public ProtocolVersion getProtocolVersion() {
   @NonNull
   @Override
   public Map<String, String> getStartupOptions() {
-    return startupOptionsRef.get();
+    // startup options are calculated dynamically and may vary per connection
+    return startupOptionsRef.get().build();
   }
 
   protected RequestLogFormatter buildRequestLogFormatter() {

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

@@ -19,24 +19,34 @@
 
 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.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 +129,7 @@ public Map<String, String> build() {
     if (applicationVersion != null) {
       builder.put(APPLICATION_VERSION_KEY, applicationVersion);
     }
+    driverBaggage().ifPresent(s -> builder.put(DRIVER_BAGGAGE, s));
 
     return builder.build();
   }
@@ -142,4 +153,21 @@ 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()) {
+      Map<String, ?> config = entry.getValue().getStartupConfiguration();
+      if (!config.isEmpty()) {
+        builder.put(entry.getKey(), config);
+      }
+    }
+    try {
+      return Optional.of(mapper.writeValueAsString(builder.build()));
+    } catch (Exception e) {
+      LOG.warn("Failed to construct startup driver baggage", e);
+      return Optional.empty();
+    }
+  }
 }

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

@@ -45,6 +45,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;
@@ -155,10 +156,38 @@ public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String
    * Before initialization, this method always returns null.
    */
   @Nullable
-  protected String getLocalDatacenter() {
+  public String getLocalDatacenter() {
     return localDc;
   }
 
+  @NonNull
+  @Override
+  public Map<String, ?> getStartupConfiguration() {
+    ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
+    if (localDc != null) {
+      builder.put("localDc", localDc);
+    } else {
+      // Local data center may not be discovered prior to connection pool initialization.
+      // In such scenario, return configured local data center name.
+      // Note that when using DC inferring load balancing policy, startup configuration
+      // may not show local DC name, because it will be discovered only once control connection
+      // is established and datacenter of contact points known.
+      Optional<String> configuredDc =
+          new OptionalLocalDcHelper(context, profile, logPrefix).configuredLocalDc();
+      configuredDc.ifPresent(d -> builder.put("localDc", d));
+    }
+    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/loadbalancing/helper/OptionalLocalDcHelper.java

@@ -65,20 +65,14 @@ public OptionalLocalDcHelper(
   @Override
   @NonNull
   public Optional<String> discoverLocalDc(@NonNull Map<UUID, Node> nodes) {
-    String localDc = context.getLocalDatacenter(profile.getName());
-    if (localDc != null) {
-      LOG.debug("[{}] Local DC set programmatically: {}", logPrefix, localDc);
-      checkLocalDatacenterCompatibility(localDc, context.getMetadataManager().getContactPoints());
-      return Optional.of(localDc);
-    } else if (profile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER)) {
-      localDc = profile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER);
-      LOG.debug("[{}] Local DC set from configuration: {}", logPrefix, localDc);
-      checkLocalDatacenterCompatibility(localDc, context.getMetadataManager().getContactPoints());
-      return Optional.of(localDc);
+    Optional<String> localDc = configuredLocalDc();
+    if (localDc.isPresent()) {
+      checkLocalDatacenterCompatibility(
+          localDc.get(), context.getMetadataManager().getContactPoints());
     } else {
       LOG.debug("[{}] Local DC not set, DC awareness will be disabled", logPrefix);
-      return Optional.empty();
     }
+    return localDc;
   }
 
   /**
@@ -138,4 +132,19 @@ protected String formatDcs(Iterable<? extends Node> nodes) {
     }
     return String.join(", ", new TreeSet<>(l));
   }
+
+  /** @return Local data center set programmatically or from configuration file. */
+  @NonNull
+  public Optional<String> configuredLocalDc() {
+    String localDc = context.getLocalDatacenter(profile.getName());
+    if (localDc != null) {
+      LOG.debug("[{}] Local DC set programmatically: {}", logPrefix, localDc);
+      return Optional.of(localDc);
+    } else if (profile.isDefined(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER)) {
+      localDc = profile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER);
+      LOG.debug("[{}] Local DC set from configuration: {}", logPrefix, localDc);
+      return Optional.of(localDc);
+    }
+    return Optional.empty();
+  }
 }

core/src/test/java/com/datastax/oss/driver/internal/core/addresstranslation/FixedHostNameAddressTranslatorTest.java

@@ -26,7 +26,6 @@
 import com.datastax.oss.driver.internal.core.context.DefaultDriverContext;
 import com.datastax.oss.driver.internal.core.context.MockedDriverContextFactory;
 import java.net.InetSocketAddress;
-import java.util.Optional;
 import org.junit.Test;
 
 public class FixedHostNameAddressTranslatorTest {
@@ -36,7 +35,7 @@ public void should_translate_address() {
     DriverExecutionProfile defaultProfile = mock(DriverExecutionProfile.class);
     when(defaultProfile.getString(ADDRESS_TRANSLATOR_ADVERTISED_HOSTNAME)).thenReturn("myaddress");
     DefaultDriverContext defaultDriverContext =
-        MockedDriverContextFactory.defaultDriverContext(Optional.of(defaultProfile));
+        MockedDriverContextFactory.defaultDriverContext(defaultProfile);
 
     FixedHostNameAddressTranslator translator =
         new FixedHostNameAddressTranslator(defaultDriverContext);

core/src/test/java/com/datastax/oss/driver/internal/core/addresstranslation/SubnetAddressTranslatorTest.java

@@ -30,7 +30,6 @@
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
 import java.net.InetSocketAddress;
 import java.util.Map;
-import java.util.Optional;
 import org.junit.Test;
 
 @SuppressWarnings("resource")
@@ -148,6 +147,6 @@ public void should_fail_on_default_address_without_port() {
   private static DefaultDriverContext context(Map<String, String> subnetAddresses) {
     DriverExecutionProfile profile = mock(DriverExecutionProfile.class);
     when(profile.getStringMap(ADDRESS_TRANSLATOR_SUBNET_ADDRESSES)).thenReturn(subnetAddresses);
-    return MockedDriverContextFactory.defaultDriverContext(Optional.of(profile));
+    return MockedDriverContextFactory.defaultDriverContext(profile);
   }
 }

core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContextTest.java

@@ -42,7 +42,7 @@ private DefaultDriverContext buildMockedContext(Optional<String> compressionOpti
     DriverExecutionProfile defaultProfile = mock(DriverExecutionProfile.class);
     when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none"))
         .thenReturn(compressionOption.orElse("none"));
-    return MockedDriverContextFactory.defaultDriverContext(Optional.of(defaultProfile));
+    return MockedDriverContextFactory.defaultDriverContext(defaultProfile);
   }
 
   private void doCreateCompressorTest(Optional<String> configVal, Class<?> expectedClz) {