diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/loadbalancing/LoadBalancingPolicy.java b/core/src/main/java/com/datastax/oss/driver/api/core/loadbalancing/LoadBalancingPolicy.java index d890ae6c100..de0d9db4ebd 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/loadbalancing/LoadBalancingPolicy.java +++ b/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 getRequestTracker() { */ void init(@NonNull Map nodes, @NonNull DistanceReporter distanceReporter); + /** Returns map containing details that impact C* node connectivity. */ + @NonNull + default Map getStartupConfiguration() { + return Collections.emptyMap(); + } + /** * Returns the coordinators to use for a new query. * diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java index a24b632f640..0d7db27dfbe 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContext.java +++ b/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 requestThrottlerRef = new LazyReference<>("requestThrottler", this::buildRequestThrottler, cycleDetector); - private final LazyReference> startupOptionsRef = - new LazyReference<>("startupOptions", this::buildStartupOptions, cycleDetector); + private final LazyReference startupOptionsRef = + new LazyReference<>("startupOptionsFactory", this::buildStartupOptionsFactory, cycleDetector); private final LazyReference nodeStateListenerRef; private final LazyReference schemaChangeListenerRef; private final LazyReference 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 buildStartupOptions() { + protected StartupOptionsBuilder buildStartupOptionsFactory() { return new StartupOptionsBuilder(this) .withClientId(startupClientId) .withApplicationName(startupApplicationName) - .withApplicationVersion(startupApplicationVersion) - .build(); + .withApplicationVersion(startupApplicationVersion); } protected Map buildLoadBalancingPolicies() { @@ -1013,7 +1012,8 @@ public ProtocolVersion getProtocolVersion() { @NonNull @Override public Map getStartupOptions() { - return startupOptionsRef.get(); + // startup options are calculated dynamically and may vary per connection + return startupOptionsRef.get().build(); } protected RequestLogFormatter buildRequestLogFormatter() { diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilder.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilder.java index 684d6b01b9c..89a9266b3ac 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilder.java +++ b/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 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 driverBaggage() { + ImmutableMap.Builder builder = new ImmutableMap.Builder<>(); + for (Map.Entry entry : + context.getLoadBalancingPolicies().entrySet()) { + Map 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(); + } + } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java index 587ef4183bd..a02a5eb3148 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/BasicLoadBalancingPolicy.java +++ b/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 getStartupConfiguration() { + ImmutableMap.Builder 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 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; diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java index 0f03cbb3643..9c31b606f18 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/DefaultLoadBalancingPolicy.java +++ b/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 getStartupConfiguration() { + Map parent = super.getStartupConfiguration(); + return ImmutableMap.of( + DefaultLoadBalancingPolicy.class.getSimpleName(), + parent.get(BasicLoadBalancingPolicy.class.getSimpleName())); + } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/helper/OptionalLocalDcHelper.java b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/helper/OptionalLocalDcHelper.java index d470f96c42c..c6143f3fa16 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/helper/OptionalLocalDcHelper.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/loadbalancing/helper/OptionalLocalDcHelper.java @@ -65,20 +65,14 @@ public OptionalLocalDcHelper( @Override @NonNull public Optional discoverLocalDc(@NonNull Map 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 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 nodes) { } return String.join(", ", new TreeSet<>(l)); } + + /** @return Local data center set programmatically or from configuration file. */ + @NonNull + public Optional 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(); + } } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/addresstranslation/FixedHostNameAddressTranslatorTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/addresstranslation/FixedHostNameAddressTranslatorTest.java index 92800998056..3bb9c4bc291 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/addresstranslation/FixedHostNameAddressTranslatorTest.java +++ b/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); diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/addresstranslation/SubnetAddressTranslatorTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/addresstranslation/SubnetAddressTranslatorTest.java index 2aa6ae75bc2..420170654dc 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/addresstranslation/SubnetAddressTranslatorTest.java +++ b/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 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); } } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContextTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContextTest.java index baf101508d4..6d4585cb4d7 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContextTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverContextTest.java @@ -42,7 +42,7 @@ private DefaultDriverContext buildMockedContext(Optional 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 configVal, Class expectedClz) { diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/context/MockedDriverContextFactory.java b/core/src/test/java/com/datastax/oss/driver/internal/core/context/MockedDriverContextFactory.java index 06817326844..a8b25193f54 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/context/MockedDriverContextFactory.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/context/MockedDriverContextFactory.java @@ -24,44 +24,45 @@ import com.datastax.oss.driver.api.core.config.DriverConfig; import com.datastax.oss.driver.api.core.config.DriverConfigLoader; 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.NodeDistanceEvaluator; +import com.datastax.oss.driver.api.core.metadata.Node; import com.datastax.oss.driver.api.core.metadata.NodeStateListener; import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListener; import com.datastax.oss.driver.api.core.session.ProgrammaticArguments; import com.datastax.oss.driver.api.core.tracker.RequestTracker; +import com.datastax.oss.driver.internal.core.ConsistencyLevelRegistry; +import com.datastax.oss.driver.internal.core.loadbalancing.DefaultLoadBalancingPolicy; +import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap; import com.datastax.oss.driver.shaded.guava.common.collect.Maps; +import edu.umd.cs.findbugs.annotations.NonNull; +import edu.umd.cs.findbugs.annotations.Nullable; import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.Map; import java.util.Optional; +import java.util.UUID; public class MockedDriverContextFactory { public static DefaultDriverContext defaultDriverContext() { - return defaultDriverContext(Optional.empty()); + return defaultDriverContext(MockedDriverContextFactory.defaultProfile("datacenter1")); } public static DefaultDriverContext defaultDriverContext( - Optional profileOption) { - - /* If the caller provided a profile use that, otherwise make a new one */ - final DriverExecutionProfile profile = - profileOption.orElseGet( - () -> { - DriverExecutionProfile blankProfile = mock(DriverExecutionProfile.class); - when(blankProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none")) - .thenReturn("none"); - when(blankProfile.getDuration(DefaultDriverOption.METRICS_NODE_EXPIRE_AFTER)) - .thenReturn(Duration.ofMinutes(5)); - when(blankProfile.isDefined(DefaultDriverOption.METRICS_FACTORY_CLASS)) - .thenReturn(true); - when(blankProfile.getString(DefaultDriverOption.METRICS_FACTORY_CLASS)) - .thenReturn("DefaultMetricsFactory"); - return blankProfile; - }); + DriverExecutionProfile defaultProfile, DriverExecutionProfile... profiles) { /* Setup machinery to connect the input DriverExecutionProfile to the config loader */ final DriverConfig driverConfig = mock(DriverConfig.class); final DriverConfigLoader configLoader = mock(DriverConfigLoader.class); when(configLoader.getInitialConfig()).thenReturn(driverConfig); - when(driverConfig.getDefaultProfile()).thenReturn(profile); + when(driverConfig.getDefaultProfile()).thenReturn(defaultProfile); + when(driverConfig.getProfile(defaultProfile.getName())).thenReturn(defaultProfile); + + for (DriverExecutionProfile profile : profiles) { + when(driverConfig.getProfile(profile.getName())).thenReturn(profile); + } ProgrammaticArguments args = ProgrammaticArguments.builder() @@ -71,6 +72,89 @@ public static DefaultDriverContext defaultDriverContext( .withLocalDatacenters(Maps.newHashMap()) .withNodeDistanceEvaluators(Maps.newHashMap()) .build(); - return new DefaultDriverContext(configLoader, args); + + return new DefaultDriverContext(configLoader, args) { + @NonNull + @Override + public Map getLoadBalancingPolicies() { + ImmutableMap.Builder map = ImmutableMap.builder(); + map.put( + defaultProfile.getName(), + mockLoadBalancingPolicy( + this, + defaultProfile.getName(), + defaultProfile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER))); + for (DriverExecutionProfile profile : profiles) { + map.put( + profile.getName(), + mockLoadBalancingPolicy( + this, + profile.getName(), + profile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER))); + } + return map.build(); + } + + @NonNull + @Override + public ConsistencyLevelRegistry getConsistencyLevelRegistry() { + return mock(ConsistencyLevelRegistry.class); + } + }; + } + + public static DriverExecutionProfile defaultProfile(String localDc) { + return createProfile(DriverExecutionProfile.DEFAULT_NAME, localDc); + } + + public static DriverExecutionProfile createProfile(String name, String localDc) { + DriverExecutionProfile defaultProfile = mock(DriverExecutionProfile.class); + when(defaultProfile.getName()).thenReturn(name); + when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none")) + .thenReturn("none"); + when(defaultProfile.getDuration(DefaultDriverOption.METRICS_NODE_EXPIRE_AFTER)) + .thenReturn(Duration.ofMinutes(5)); + when(defaultProfile.isDefined(DefaultDriverOption.METRICS_FACTORY_CLASS)).thenReturn(true); + when(defaultProfile.getString(DefaultDriverOption.METRICS_FACTORY_CLASS)) + .thenReturn("DefaultMetricsFactory"); + when(defaultProfile.getString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER)) + .thenReturn(localDc); + return defaultProfile; + } + + public static void allowRemoteDcConnectivity( + DriverExecutionProfile profile, + int maxNodesPerRemoteDc, + boolean allowRemoteSatisfyLocalDc, + List preferredRemoteDcs) { + when(profile.getInt(DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_MAX_NODES_PER_REMOTE_DC)) + .thenReturn(maxNodesPerRemoteDc); + when(profile.getBoolean( + DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_ALLOW_FOR_LOCAL_CONSISTENCY_LEVELS)) + .thenReturn(allowRemoteSatisfyLocalDc); + when(profile.getStringList(DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS)) + .thenReturn(preferredRemoteDcs); + } + + private static LoadBalancingPolicy mockLoadBalancingPolicy( + DefaultDriverContext driverContext, String profile, String localDc) { + LoadBalancingPolicy loadBalancingPolicy = + new DefaultLoadBalancingPolicy(driverContext, profile) { + @NonNull + @Override + protected Optional discoverLocalDc(@NonNull Map nodes) { + return Optional.ofNullable(localDc); + } + + @NonNull + @Override + protected NodeDistanceEvaluator createNodeDistanceEvaluator( + @Nullable String localDc, @NonNull Map nodes) { + return mock(NodeDistanceEvaluator.class); + } + }; + loadBalancingPolicy.init( + Collections.emptyMap(), mock(LoadBalancingPolicy.DistanceReporter.class)); + return loadBalancingPolicy; } } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilderTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilderTest.java index 33811b2793a..d12e50b7e8e 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilderTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/context/StartupOptionsBuilderTest.java @@ -26,10 +26,10 @@ 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.session.Session; +import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList; import com.datastax.oss.protocol.internal.request.Startup; import com.tngtech.java.junit.dataprovider.DataProvider; import com.tngtech.java.junit.dataprovider.DataProviderRunner; -import java.util.Optional; import org.junit.Test; import org.junit.runner.RunWith; @@ -41,7 +41,8 @@ private DefaultDriverContext buildMockedContext(String compression) { DriverExecutionProfile defaultProfile = mock(DriverExecutionProfile.class); when(defaultProfile.getString(DefaultDriverOption.PROTOCOL_COMPRESSION, "none")) .thenReturn(compression); - return MockedDriverContextFactory.defaultDriverContext(Optional.of(defaultProfile)); + when(defaultProfile.getName()).thenReturn(DriverExecutionProfile.DEFAULT_NAME); + return MockedDriverContextFactory.defaultDriverContext(defaultProfile); } private void assertDefaultStartupOptions(Startup startup) { @@ -94,4 +95,44 @@ public void should_fail_to_build_startup_options_with_invalid_compression() { new Startup(ctx.getStartupOptions()); }); } + + @Test + public void should_include_all_local_dcs_in_startup_message() { + + DefaultDriverContext ctx = + MockedDriverContextFactory.defaultDriverContext( + MockedDriverContextFactory.defaultProfile("us-west-2"), + MockedDriverContextFactory.createProfile("oltp", "us-east-2"), + MockedDriverContextFactory.createProfile("olap", "eu-central-1")); + Startup startup = new Startup(ctx.getStartupOptions()); + assertThat(startup.options) + .containsEntry( + StartupOptionsBuilder.DRIVER_BAGGAGE, + "{\"default\":{\"DefaultLoadBalancingPolicy\":{\"localDc\":\"us-west-2\"}}," + + "\"oltp\":{\"DefaultLoadBalancingPolicy\":{\"localDc\":\"us-east-2\"}}," + + "\"olap\":{\"DefaultLoadBalancingPolicy\":{\"localDc\":\"eu-central-1\"}}}"); + } + + @Test + public void should_include_all_lbp_details_in_startup_message() { + + DriverExecutionProfile defaultProfile = MockedDriverContextFactory.defaultProfile("dc1"); + DriverExecutionProfile oltpProfile = MockedDriverContextFactory.createProfile("oltp", "dc1"); + MockedDriverContextFactory.allowRemoteDcConnectivity( + oltpProfile, 2, true, ImmutableList.of("dc2", "dc3")); + DefaultDriverContext ctx = + MockedDriverContextFactory.defaultDriverContext(defaultProfile, oltpProfile); + + Startup startup = new Startup(ctx.getStartupOptions()); + + assertThat(startup.options) + .containsEntry( + StartupOptionsBuilder.DRIVER_BAGGAGE, + "{\"default\":{\"DefaultLoadBalancingPolicy\":{\"localDc\":\"dc1\"}}," + + "\"oltp\":{\"DefaultLoadBalancingPolicy\":{" + + "\"localDc\":\"dc1\"," + + "\"preferredRemoteDcs\":[\"dc2\",\"dc3\"]," + + "\"allowDcFailoverForLocalCl\":true," + + "\"maxNodesPerRemoteDc\":2}}}"); + } }