From 72b8dba5a5aa0a8834728fd58cf0246a33f429ac Mon Sep 17 00:00:00 2001 From: Mikita Hradovich Date: Tue, 28 Jul 2026 18:19:23 +0200 Subject: [PATCH] feat: report driver configuration to the cluster at connection time (stage 1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stage 1 of driver configuration reporting: the mechanism to report a client's effective configuration to the cluster at connection time, so operators can inspect driver settings while investigating incidents. The server stores the reported STARTUP options in system.clients.client_options on ScyllaDB and in system_views.clients.client_options on Apache Cassandra (added in 4.1). Gated behind a new option, advanced.driver-config-reporting.enabled (constant DRIVER_CONFIG_REPORTING_ENABLED), which ships disabled — when off there is zero change on the wire. When enabled: - SESSION_ID: a dedicated, driver-generated per-session UUID, sent on every connection (control and pool) so the server can group a session's connections. Independent of the user-settable CLIENT_ID. - DRIVER_CONFIG: a compact JSON blob, sent only on the control connection. Stage 1 emits only {"version":1} (version metadata only); the full configuration report follows in stage 2. Both entries are produced by a new DriverConfigReporter, invoked per connection from ProtocolInitHandler and evaluated fresh (so toggling the flag at runtime takes effect on new connections). Reporting is fail-safe: any failure while building the report is swallowed and never breaks connection initialization. The config option is registered consistently across DefaultDriverOption, TypedDriverOption, OptionsMap.driverDefaults() and reference.conf. Tests: - DriverConfigReportingSimulacronIT pins the exact wire contract (SESSION_ID on every connection with one shared value; DRIVER_CONFIG only on the control connection; nothing on the wire when disabled) and asserts the payload parses to version == 1. - DriverConfigReportingCcmIT is the real-server counterpart, reading the stored options back on both backends — ScyllaDB >= 2026.1 (system.clients) and Apache Cassandra >= 4.1 (system_views.clients) — and asserting identical behavior. Verified live against ScyllaDB 2026.1.9 and Cassandra 4.1.10. Co-Authored-By: Claude Opus 4.8 --- .../api/core/config/DefaultDriverOption.java | 17 +- .../driver/api/core/config/OptionsMap.java | 1 + .../api/core/config/TypedDriverOption.java | 5 + .../core/channel/DriverChannelOptions.java | 23 +- .../core/channel/ProtocolInitHandler.java | 5 + .../context/DefaultDriverConfigReporter.java | 131 +++++++++++ .../core/context/DefaultDriverContext.java | 12 ++ .../core/context/DriverConfigReporter.java | 53 +++++ .../core/context/InternalDriverContext.java | 7 + .../core/control/ControlConnection.java | 1 + core/src/main/resources/reference.conf | 20 ++ .../core/channel/ChannelFactoryTestBase.java | 2 + .../core/channel/ProtocolInitHandlerTest.java | 68 ++++++ .../DefaultDriverConfigReporterTest.java | 141 ++++++++++++ .../config/DriverConfigReportingCcmIT.java | 187 ++++++++++++++++ .../DriverConfigReportingSimulacronIT.java | 204 ++++++++++++++++++ 16 files changed, 874 insertions(+), 3 deletions(-) create mode 100644 core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverConfigReporter.java create mode 100644 core/src/main/java/com/datastax/oss/driver/internal/core/context/DriverConfigReporter.java create mode 100644 core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverConfigReporterTest.java create mode 100644 integration-tests/src/test/java/com/datastax/oss/driver/core/config/DriverConfigReportingCcmIT.java create mode 100644 integration-tests/src/test/java/com/datastax/oss/driver/core/config/DriverConfigReportingSimulacronIT.java diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java index 56a6fe9c2be..3a6e4ed69bb 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java @@ -1173,7 +1173,22 @@ public enum DefaultDriverOption implements DriverOption { * *

Value-Type: boolean */ - ADDRESS_TRANSLATOR_RESOLVE_ADDRESSES("advanced.address-translator.resolve-addresses"); + ADDRESS_TRANSLATOR_RESOLVE_ADDRESSES("advanced.address-translator.resolve-addresses"), + /** + * Whether the driver reports its effective configuration to ScyllaDB at connection time. + * + *

When {@code true}, the driver adds two entries to the CQL {@code STARTUP} options, which + * ScyllaDB stores in {@code system.clients.client_options} so operators can inspect driver + * settings while investigating incidents: a {@code SESSION_ID} on every connection (so the server + * can group a session's connections) and a compact JSON payload under the {@code DRIVER_CONFIG} + * key on the control connection only. At this stage the {@code DRIVER_CONFIG} payload carries + * only schema-version metadata ({"version":1}); reporting of the effective + * configuration fields is planned for a later stage. When {@code false}, neither entry is sent + * and there is no change on the wire. + * + *

Value type: boolean + */ + DRIVER_CONFIG_REPORTING_ENABLED("advanced.driver-config-reporting.enabled"); private final String path; diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java index 60190fc1cce..1d906caa985 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/OptionsMap.java @@ -400,6 +400,7 @@ protected static void fillWithDriverDefaults(OptionsMap map) { // values) with no sensible scalar default, analogous to how CONFIG_RELOAD_INTERVAL is omitted. map.put(TypedDriverOption.CLIENT_ROUTES_NATIVE_TRANSPORT_PORT, 9042); map.put(TypedDriverOption.CLIENT_ROUTES_SHARD_AWARENESS_ENABLED, false); + map.put(TypedDriverOption.DRIVER_CONFIG_REPORTING_ENABLED, false); } @Immutable diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java index e7c606cade8..e412b99b404 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/config/TypedDriverOption.java @@ -976,6 +976,11 @@ public String toString() { new TypedDriverOption<>( DefaultDriverOption.CLIENT_ROUTES_SHARD_AWARENESS_ENABLED, GenericType.BOOLEAN); + /** Whether the driver reports its configuration to ScyllaDB at connection time. */ + public static final TypedDriverOption DRIVER_CONFIG_REPORTING_ENABLED = + new TypedDriverOption<>( + DefaultDriverOption.DRIVER_CONFIG_REPORTING_ENABLED, GenericType.BOOLEAN); + private static Iterable> introspectBuiltInValues() { try { ImmutableList.Builder> result = ImmutableList.builder(); diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannelOptions.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannelOptions.java index 208cf52ac22..378fd2dc0b8 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannelOptions.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/DriverChannelOptions.java @@ -47,15 +47,24 @@ public static Builder builder() { public final String ownerLogPrefix; + /** + * Whether the channel should report the driver configuration ({@code DRIVER_CONFIG}) in its + * {@code STARTUP} options. Set only for the control connection, so the (potentially large) config + * blob is sent once per session rather than on every pooled connection. + */ + public final boolean reportConfig; + private DriverChannelOptions( CqlIdentifier keyspace, List eventTypes, EventCallback eventCallback, - String ownerLogPrefix) { + String ownerLogPrefix, + boolean reportConfig) { this.keyspace = keyspace; this.eventTypes = eventTypes; this.eventCallback = eventCallback; this.ownerLogPrefix = ownerLogPrefix; + this.reportConfig = reportConfig; } public static class Builder { @@ -63,6 +72,7 @@ public static class Builder { private List eventTypes = Collections.emptyList(); private EventCallback eventCallback = null; private String ownerLogPrefix = null; + private boolean reportConfig = false; public Builder withKeyspace(CqlIdentifier keyspace) { this.keyspace = keyspace; @@ -82,8 +92,17 @@ public Builder withOwnerLogPrefix(String ownerLogPrefix) { return this; } + /** + * Marks this channel as the one that reports {@code DRIVER_CONFIG} (the control connection). + */ + public Builder reportConfig(boolean reportConfig) { + this.reportConfig = reportConfig; + return this; + } + public DriverChannelOptions build() { - return new DriverChannelOptions(keyspace, eventTypes, eventCallback, ownerLogPrefix); + return new DriverChannelOptions( + keyspace, eventTypes, eventCallback, ownerLogPrefix, reportConfig); } } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandler.java index 85bffdb016e..cf782096f18 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandler.java @@ -194,6 +194,11 @@ Message getRequest() { if (featureStore != null) { featureStore.populateStartupOptions(startupOptions); } + // Adds SESSION_ID on every connection and DRIVER_CONFIG on the control connection + // (options.reportConfig); no-op when driver config reporting is disabled. + context + .getDriverConfigReporter() + .populateStartupOptions(startupOptions, options.reportConfig); return request = new Startup(startupOptions); case GET_CLUSTER_NAME: return request = CLUSTER_NAME_QUERY; diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverConfigReporter.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverConfigReporter.java new file mode 100644 index 00000000000..2cd1b5a8560 --- /dev/null +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DefaultDriverConfigReporter.java @@ -0,0 +1,131 @@ +/* + * 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.internal.core.context; + +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.uuid.Uuids; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.util.Map; +import java.util.UUID; +import net.jcip.annotations.ThreadSafe; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Default {@link DriverConfigReporter}: serializes the driver configuration to the cross-driver + * {@code DRIVER_CONFIG} JSON shape and adds it to the control connection's {@code STARTUP} options. + * + *

The blob is (re)built on demand every time the control connection initializes, so it always + * reflects the current (possibly reloaded) configuration without any caching. + */ +@ThreadSafe +public class DefaultDriverConfigReporter implements DriverConfigReporter { + + private static final Logger LOG = LoggerFactory.getLogger(DefaultDriverConfigReporter.class); + + /** STARTUP option key under which the config JSON is sent. */ + public static final String DRIVER_CONFIG_KEY = "DRIVER_CONFIG"; + + /** STARTUP option key under which the per-session identifier is sent. */ + public static final String SESSION_ID_KEY = "SESSION_ID"; + + /** + * Major schema version. Adding keys is backward-compatible and does not bump this; only + * changing/removing the meaning of an existing key does. + */ + static final int SCHEMA_VERSION = 1; + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + protected final InternalDriverContext context; + + // Dedicated, driver-generated identifier for this session. Not derived from the (user-settable, + // Insights-oriented) CLIENT_ID, so that it is guaranteed unique per session as the grouping key + // requires. The reporter is a per-session singleton (built once via LazyReference), so this value + // is stable and shared across all of the session's connections. + private final UUID sessionId = Uuids.random(); + + public DefaultDriverConfigReporter(InternalDriverContext context) { + this.context = context; + } + + @Override + public void populateStartupOptions( + Map startupOptions, boolean reportDriverConfig) { + // Configuration reporting is a best-effort diagnostic aid: it runs on the connection + // initialization path, so any failure here (a bad config read, a misbehaving policy while + // introspecting, a serialization error) must be swallowed rather than allowed to break the + // connection — which would prevent the session from establishing or reconnecting. + try { + if (!isEnabled()) { + return; + } + // SESSION_ID on every connection so the server can group a session's connections. + startupOptions.put(SESSION_ID_KEY, sessionId.toString()); + // DRIVER_CONFIG blob only on the control connection. + if (reportDriverConfig) { + String json = buildJson(); + if (json != null) { + startupOptions.put(DRIVER_CONFIG_KEY, json); + } + } + } catch (RuntimeException e) { + LOG.warn( + "Error while building the driver configuration report; skipping driver config reporting", + e); + } + } + + private boolean isEnabled() { + return context + .getConfig() + .getDefaultProfile() + .getBoolean(DefaultDriverOption.DRIVER_CONFIG_REPORTING_ENABLED, false); + } + + /** + * Builds the compact, single-line JSON configuration report. + * + *

Stage 1 emits only the schema {@code version}; the individual configuration groups are + * populated in {@link #populateConfig(ObjectNode, DriverExecutionProfile)} in a later stage. + */ + protected String buildJson() { + ObjectNode root = OBJECT_MAPPER.createObjectNode(); + root.put("version", SCHEMA_VERSION); + populateConfig(root, context.getConfig().getDefaultProfile()); + try { + return OBJECT_MAPPER.writeValueAsString(root); + } catch (JsonProcessingException e) { + // An in-memory node tree should never fail to serialize; never let it break connection setup. + LOG.warn("Failed to serialize driver configuration report; skipping DRIVER_CONFIG", e); + return null; + } + } + + /** + * Populates the configuration groups onto the report root. Placeholder in Stage 1; Stage 2 fills + * in {@code connection}, {@code socket}, the policy groups, {@code query_defaults}, {@code tls}, + * etc. + */ + protected void populateConfig(ObjectNode root, DriverExecutionProfile config) { + // Stage 2: populate configuration groups from `config` and the context's policies. + } +} 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 52d9cdb264c..3d1d5b82b87 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 @@ -232,6 +232,8 @@ public class DefaultDriverContext implements InternalDriverContext { new LazyReference<>("requestThrottler", this::buildRequestThrottler, cycleDetector); private final LazyReference> startupOptionsRef = new LazyReference<>("startupOptions", this::buildStartupOptions, cycleDetector); + private final LazyReference driverConfigReporterRef = + new LazyReference<>("driverConfigReporter", this::buildDriverConfigReporter, cycleDetector); private final LazyReference nodeStateListenerRef; private final LazyReference schemaChangeListenerRef; private final LazyReference requestTrackerRef; @@ -369,6 +371,10 @@ protected Map buildStartupOptions() { .build(); } + protected DriverConfigReporter buildDriverConfigReporter() { + return new DefaultDriverConfigReporter(this); + } + protected Map buildLoadBalancingPolicies() { return Reflection.buildFromConfigProfiles( this, @@ -1226,6 +1232,12 @@ public Map getStartupOptions() { return startupOptionsRef.get(); } + @NonNull + @Override + public DriverConfigReporter getDriverConfigReporter() { + return driverConfigReporterRef.get(); + } + protected RequestLogFormatter buildRequestLogFormatter() { return new RequestLogFormatter(this); } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/DriverConfigReporter.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DriverConfigReporter.java new file mode 100644 index 00000000000..d614793e9d0 --- /dev/null +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/context/DriverConfigReporter.java @@ -0,0 +1,53 @@ +/* + * 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.internal.core.context; + +import java.util.Map; + +/** + * Adds the client-configuration-reporting entries to a connection's CQL {@code STARTUP} options, so + * ScyllaDB can store them in {@code system.clients.client_options} and operators can inspect a + * client's effective driver settings while investigating incidents. + * + *

Two entries are produced, both governed by {@code advanced.driver-config-reporting.enabled}: + * + *

+ */ +public interface DriverConfigReporter { + + /** + * Adds the reporting entries to the given startup options: {@code SESSION_ID} on every + * connection, plus {@code DRIVER_CONFIG} when {@code reportDriverConfig} is true (the control + * connection). Does nothing when configuration reporting is disabled. + * + *

Called from the protocol-initialization handler for every connection. + * + *

Implementations must not throw: this runs on the connection initialization path, so a + * failure to build the report must be swallowed (and logged) rather than propagated, otherwise it + * would prevent the session from establishing or reconnecting. + * + * @param reportDriverConfig whether this connection should also carry the full {@code + * DRIVER_CONFIG} blob; true only for the control connection. + */ + void populateStartupOptions(Map startupOptions, boolean reportDriverConfig); +} diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/context/InternalDriverContext.java b/core/src/main/java/com/datastax/oss/driver/internal/core/context/InternalDriverContext.java index 81349b0c665..87933b7385e 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/context/InternalDriverContext.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/context/InternalDriverContext.java @@ -176,6 +176,13 @@ public interface InternalDriverContext extends DriverContext { @NonNull Map getStartupOptions(); + /** + * The component that builds the driver-configuration report ({@code DRIVER_CONFIG}) sent on the + * control connection's Startup message. + */ + @NonNull + DriverConfigReporter getDriverConfigReporter(); + /** * A list of additional components to notify of session lifecycle events. * diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/control/ControlConnection.java b/core/src/main/java/com/datastax/oss/driver/internal/core/control/ControlConnection.java index 52b32255c4b..9da3a2a8aa2 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/control/ControlConnection.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/control/ControlConnection.java @@ -327,6 +327,7 @@ private void init( DriverChannelOptions.builder() .withEvents(eventTypes, ControlConnection.this) .withOwnerLogPrefix(logPrefix + "|control") + .reportConfig(true) .build(); Queue nodes = diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf index 721afa1cd76..8a3a444319e 100644 --- a/core/src/main/resources/reference.conf +++ b/core/src/main/resources/reference.conf @@ -1204,6 +1204,26 @@ datastax-java-driver { } + advanced.driver-config-reporting { + + # Whether the driver reports its effective configuration to ScyllaDB at connection time. + # + # When true, the driver adds two entries to the CQL STARTUP options, which ScyllaDB stores in + # system.clients.client_options so operators can inspect driver settings while investigating + # incidents: a SESSION_ID on every connection (so the server can group a session's connections) + # and a compact JSON payload under the DRIVER_CONFIG key on the control connection only. At this + # stage the DRIVER_CONFIG payload carries only schema-version metadata ({"version":1}); reporting + # of the effective configuration fields is planned for a later stage. When false, neither entry + # is sent and there is no change on the wire. + # + # Required: no + # Modifiable at runtime: yes, the new value will be used for connections initialized after the change. + # Overridable in a profile: no + # Default: false + enabled = false + + } + # Whether to resolve the addresses passed to `basic.contact-points`. # # If this is true, addresses are created with `InetSocketAddress(String, int)`: the host name will diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryTestBase.java b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryTestBase.java index b25a1e9ad71..2780d5bdec9 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryTestBase.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ChannelFactoryTestBase.java @@ -141,6 +141,8 @@ public void setup() throws InterruptedException { when(context.getEventBus()).thenReturn(eventBus); when(context.getWriteCoalescer()).thenReturn(new PassThroughWriteCoalescer(null)); when(context.getCompressor()).thenReturn(compressor); + // The init handler consults the config reporter for every connection; default to a no-op. + when(context.getDriverConfigReporter()).thenReturn((startupOptions, reportDriverConfig) -> {}); // Start local server ServerBootstrap serverBootstrap = diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandlerTest.java index 8057ccdc48e..a7051ac466d 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandlerTest.java @@ -41,6 +41,7 @@ import com.datastax.oss.driver.internal.core.DefaultProtocolVersionRegistry; import com.datastax.oss.driver.internal.core.ProtocolVersionRegistry; import com.datastax.oss.driver.internal.core.TestResponses; +import com.datastax.oss.driver.internal.core.context.DefaultDriverConfigReporter; import com.datastax.oss.driver.internal.core.context.InternalDriverContext; import com.datastax.oss.driver.internal.core.metadata.TestNodeFactory; import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList; @@ -102,6 +103,9 @@ public void setup() { when(defaultProfile.getDuration(DefaultDriverOption.HEARTBEAT_INTERVAL)) .thenReturn(Duration.ofSeconds(30)); when(internalDriverContext.getProtocolVersionRegistry()).thenReturn(protocolVersionRegistry); + // The init handler consults the config reporter for every connection; default to a no-op. + when(internalDriverContext.getDriverConfigReporter()) + .thenReturn((startupOptions, reportDriverConfig) -> {}); channel .pipeline() @@ -153,6 +157,70 @@ public void should_initialize() { assertThat(connectFuture).isSuccess(); } + // Mirrors the real reporter: SESSION_ID on every connection, DRIVER_CONFIG only when asked. + private void stubConfigReporter() { + when(internalDriverContext.getDriverConfigReporter()) + .thenReturn( + (startupOptions, reportDriverConfig) -> { + startupOptions.put(DefaultDriverConfigReporter.SESSION_ID_KEY, "test-session-id"); + if (reportDriverConfig) { + startupOptions.put( + DefaultDriverConfigReporter.DRIVER_CONFIG_KEY, "{\"version\":1}"); + } + }); + } + + @Test + public void should_report_session_id_and_driver_config_on_control_connection() { + stubConfigReporter(); + channel + .pipeline() + .addLast( + ChannelFactory.INIT_HANDLER_NAME, + new ProtocolInitHandler( + internalDriverContext, + DefaultProtocolVersion.V4, + null, + END_POINT, + DriverChannelOptions.builder().reportConfig(true).build(), + heartbeatHandler, + false)); + + channel.connect(new InetSocketAddress("localhost", 9042)); + + Frame requestFrame = readOutboundFrame(); + assertThat(requestFrame.message).isInstanceOf(Startup.class); + Startup startup = (Startup) requestFrame.message; + assertThat(startup.options).containsKey(DefaultDriverConfigReporter.SESSION_ID_KEY); + assertThat(startup.options).containsKey(DefaultDriverConfigReporter.DRIVER_CONFIG_KEY); + } + + @Test + public void should_report_session_id_but_not_driver_config_on_pool_connection() { + stubConfigReporter(); + channel + .pipeline() + .addLast( + ChannelFactory.INIT_HANDLER_NAME, + new ProtocolInitHandler( + internalDriverContext, + DefaultProtocolVersion.V4, + null, + END_POINT, + // DriverChannelOptions.DEFAULT has reportConfig = false (a pool connection). + DriverChannelOptions.DEFAULT, + heartbeatHandler, + false)); + + channel.connect(new InetSocketAddress("localhost", 9042)); + + Frame requestFrame = readOutboundFrame(); + assertThat(requestFrame.message).isInstanceOf(Startup.class); + Startup startup = (Startup) requestFrame.message; + assertThat(startup.options).containsKey(DefaultDriverConfigReporter.SESSION_ID_KEY); + assertThat(startup.options).doesNotContainKey(DefaultDriverConfigReporter.DRIVER_CONFIG_KEY); + } + @Test public void should_query_supported_options() { channel diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverConfigReporterTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverConfigReporterTest.java new file mode 100644 index 00000000000..b303e6db5a5 --- /dev/null +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/context/DefaultDriverConfigReporterTest.java @@ -0,0 +1,141 @@ +/* + * 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.internal.core.context; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.datastax.oss.driver.api.core.config.DefaultDriverOption; +import com.datastax.oss.driver.api.core.config.DriverConfig; +import com.datastax.oss.driver.api.core.config.DriverExecutionProfile; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.util.HashMap; +import java.util.Map; +import java.util.UUID; +import org.junit.Before; +import org.junit.Test; + +public class DefaultDriverConfigReporterTest { + + private InternalDriverContext context; + private DriverExecutionProfile profile; + private DefaultDriverConfigReporter reporter; + + @Before + public void setup() { + context = mock(InternalDriverContext.class); + DriverConfig config = mock(DriverConfig.class); + profile = mock(DriverExecutionProfile.class); + when(context.getConfig()).thenReturn(config); + when(config.getDefaultProfile()).thenReturn(profile); + reporter = new DefaultDriverConfigReporter(context); + } + + private void enableReporting(boolean enabled) { + when(profile.getBoolean(DefaultDriverOption.DRIVER_CONFIG_REPORTING_ENABLED, false)) + .thenReturn(enabled); + } + + @Test + public void should_add_nothing_when_disabled() { + enableReporting(false); + Map options = new HashMap<>(); + reporter.populateStartupOptions(options, /* reportDriverConfig= */ true); + assertThat(options).doesNotContainKey(DefaultDriverConfigReporter.SESSION_ID_KEY); + assertThat(options).doesNotContainKey(DefaultDriverConfigReporter.DRIVER_CONFIG_KEY); + } + + @Test + public void should_add_session_id_and_driver_config_on_control_connection() { + enableReporting(true); + Map options = new HashMap<>(); + reporter.populateStartupOptions(options, /* reportDriverConfig= */ true); + // SESSION_ID is a valid, driver-generated UUID. + String sessionId = options.get(DefaultDriverConfigReporter.SESSION_ID_KEY); + assertThat(sessionId).isNotNull(); + assertThat(UUID.fromString(sessionId)).isNotNull(); // does not throw => valid UUID + // Stage 1 emits only the schema version; the value must be valid compact JSON. + assertThat(options.get(DefaultDriverConfigReporter.DRIVER_CONFIG_KEY)) + .isEqualTo("{\"version\":" + DefaultDriverConfigReporter.SCHEMA_VERSION + "}"); + } + + @Test + public void should_add_session_id_only_on_pool_connection() { + enableReporting(true); + Map options = new HashMap<>(); + reporter.populateStartupOptions(options, /* reportDriverConfig= */ false); + assertThat(options).containsKey(DefaultDriverConfigReporter.SESSION_ID_KEY); + assertThat(options).doesNotContainKey(DefaultDriverConfigReporter.DRIVER_CONFIG_KEY); + } + + @Test + public void should_use_a_stable_session_id_across_connections() { + enableReporting(true); + Map control = new HashMap<>(); + Map pool = new HashMap<>(); + reporter.populateStartupOptions(control, true); + reporter.populateStartupOptions(pool, false); + assertThat(pool.get(DefaultDriverConfigReporter.SESSION_ID_KEY)) + .isEqualTo(control.get(DefaultDriverConfigReporter.SESSION_ID_KEY)); + } + + @Test + public void should_use_a_distinct_session_id_per_reporter() { + enableReporting(true); + Map first = new HashMap<>(); + reporter.populateStartupOptions(first, false); + // A second session (new reporter instance) must get a different SESSION_ID. + Map second = new HashMap<>(); + new DefaultDriverConfigReporter(context).populateStartupOptions(second, false); + assertThat(second.get(DefaultDriverConfigReporter.SESSION_ID_KEY)) + .isNotEqualTo(first.get(DefaultDriverConfigReporter.SESSION_ID_KEY)); + } + + /** Reporting must never break the connection: a failed config read is swallowed entirely. */ + @Test + public void should_not_throw_when_reading_the_flag_fails() { + when(profile.getBoolean(DefaultDriverOption.DRIVER_CONFIG_REPORTING_ENABLED, false)) + .thenThrow(new IllegalStateException("config blew up")); + Map options = new HashMap<>(); + reporter.populateStartupOptions(options, true); // must not throw + assertThat(options).doesNotContainKey(DefaultDriverConfigReporter.SESSION_ID_KEY); + assertThat(options).doesNotContainKey(DefaultDriverConfigReporter.DRIVER_CONFIG_KEY); + } + + /** + * Reporting must never break the connection: a failure while building the config groups (as a + * Stage 2 policy introspection might) is swallowed. SESSION_ID is still emitted (it is added + * before, and independently of, the DRIVER_CONFIG blob); only DRIVER_CONFIG is omitted. + */ + @Test + public void should_keep_session_id_but_skip_driver_config_when_building_config_groups_fails() { + enableReporting(true); + DefaultDriverConfigReporter throwingReporter = + new DefaultDriverConfigReporter(context) { + @Override + protected void populateConfig(ObjectNode root, DriverExecutionProfile config) { + throw new IllegalStateException("policy introspection blew up"); + } + }; + Map options = new HashMap<>(); + throwingReporter.populateStartupOptions(options, true); // must not throw + assertThat(options).containsKey(DefaultDriverConfigReporter.SESSION_ID_KEY); + assertThat(options).doesNotContainKey(DefaultDriverConfigReporter.DRIVER_CONFIG_KEY); + } +} diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/config/DriverConfigReportingCcmIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/config/DriverConfigReportingCcmIT.java new file mode 100644 index 00000000000..4710f9d4efe --- /dev/null +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/config/DriverConfigReportingCcmIT.java @@ -0,0 +1,187 @@ +/* + * 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.core.config; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.config.DefaultDriverOption; +import com.datastax.oss.driver.api.core.cql.Row; +import com.datastax.oss.driver.api.testinfra.ccm.CcmBridge; +import com.datastax.oss.driver.api.testinfra.ccm.CcmRule; +import com.datastax.oss.driver.api.testinfra.requirement.BackendRequirement; +import com.datastax.oss.driver.api.testinfra.requirement.BackendType; +import com.datastax.oss.driver.api.testinfra.session.SessionRule; +import com.datastax.oss.driver.api.testinfra.session.SessionUtils; +import com.datastax.oss.driver.categories.ParallelizableTests; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.time.Duration; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.rules.RuleChain; +import org.junit.rules.TestRule; + +/** + * Verifies driver-configuration reporting end-to-end against a live server (via CCM), by reading + * back what the server actually stored for each connection's {@code STARTUP} options. + * + *

This is the real-server counterpart of {@link DriverConfigReportingSimulacronIT}: where + * Simulacron only proves what the driver sends, this confirms that a real server + * accepts the extra {@code STARTUP} keys and stores them, so that (a) {@code + * SESSION_ID} is present on every one of the session's connections with a single shared value, and + * (b) {@code DRIVER_CONFIG} is stored for exactly one connection (the control connection). + * + *

Runs on both backends, asserting identical behavior — only the table that exposes the stored + * options differs: ScyllaDB uses {@code system.clients.client_options}, while Apache Cassandra + * exposes it in {@code system_views.clients.client_options} (added in Cassandra 4.1). + */ +@Category(ParallelizableTests.class) +@BackendRequirement( + type = BackendType.SCYLLA, + minInclusive = "2026.1.0", + description = "system.clients.client_options is a ScyllaDB feature") +@BackendRequirement( + type = BackendType.CASSANDRA, + minInclusive = "4.1", + description = "system_views.clients.client_options was added in Cassandra 4.1") +public class DriverConfigReportingCcmIT { + + private static final String DRIVER_NAME = "ScyllaDB Java Driver"; + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + private static final CcmRule CCM_RULE = CcmRule.getInstance(); + + private static final SessionRule SESSION_RULE = + SessionRule.builder(CCM_RULE) + .withConfigLoader( + SessionUtils.configLoaderBuilder() + .withBoolean(DefaultDriverOption.DRIVER_CONFIG_REPORTING_ENABLED, true) + .withDuration(DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofSeconds(30)) + .build()) + .build(); + + @ClassRule + public static final TestRule CHAIN = RuleChain.outerRule(CCM_RULE).around(SESSION_RULE); + + /** + * The table exposing connected clients and their {@code STARTUP} options: {@code system.clients} + * on ScyllaDB, {@code system_views.clients} on Apache Cassandra (the only per-backend difference; + * the columns and {@code client_options} keys are identical). + */ + private static String clientsTable() { + return CcmBridge.isDistributionOf(BackendType.SCYLLA) + ? "system.clients" + : "system_views.clients"; + } + + @Test + public void should_store_session_id_on_all_connections_and_driver_config_on_control() { + CqlSession session = SESSION_RULE.session(); + + // The session opens a control connection plus at least one pooled connection; wait until the + // server reflects them (the clients table is updated asynchronously as connections are set up). + await() + .atMost(60, TimeUnit.SECONDS) + .pollInterval(500, TimeUnit.MILLISECONDS) + .until(() -> driverConnections(session).size() >= 2); + + List rows = driverConnections(session); + + // Evidence for the record: dump exactly what the server stored per connection. + System.out.println(clientsTable() + " rows for this driver session:"); + for (Row row : rows) { + System.out.printf( + " %s:%s stage=%s client_options=%s%n", + row.getObject("address"), + row.getObject("port"), + row.getString("connection_stage"), + row.getMap("client_options", String.class, String.class)); + } + + // (a) Every connection carries SESSION_ID, and all of them share a single value (one session). + Set sessionIds = + rows.stream().map(row -> clientOptions(row).get("SESSION_ID")).collect(Collectors.toSet()); + assertThat(sessionIds).doesNotContainNull().hasSize(1); + + // (b) DRIVER_CONFIG is stored for exactly one connection (the control connection), and its + // value round-trips through the server intact as the stage-1 payload: valid JSON carrying + // exactly the schema version. + List driverConfigs = + rows.stream() + .map(row -> clientOptions(row).get("DRIVER_CONFIG")) + .filter(Objects::nonNull) + .collect(Collectors.toList()); + assertThat(driverConfigs).hasSize(1); + assertStageOnePayload(driverConfigs.get(0)); + } + + /** + * Asserts that a {@code DRIVER_CONFIG} value is the stage-1 payload: well-formed JSON whose + * {@code version} is the integer {@code 1}. Guards against an incorrect schema version or a + * malformed blob slipping through a mere key-presence check. + */ + private static void assertStageOnePayload(String driverConfig) { + JsonNode root; + try { + root = OBJECT_MAPPER.readTree(driverConfig); + } catch (JsonProcessingException e) { + throw new AssertionError("DRIVER_CONFIG is not valid JSON: " + driverConfig, e); + } + assertThat(root.path("version").isInt()) + .as("version is an integer in %s", driverConfig) + .isTrue(); + assertThat(root.path("version").intValue()).isEqualTo(1); + } + + /** + * The rows in the clients table that belong to this driver session's connections: this driver, in + * a {@code READY} state, and carrying the reporting {@code SESSION_ID}. Transient + * protocol-version negotiation attempts (no driver identity, closed immediately) are excluded, + * and their absence here is itself the confirmation that they leave no lingering session rows. + */ + private List driverConnections(CqlSession session) { + return session + .execute( + "SELECT address, port, connection_stage, driver_name, client_options FROM " + + clientsTable()) + .all() + .stream() + .filter(row -> DRIVER_NAME.equals(row.getString("driver_name"))) + // connection_stage casing differs across backends; compare case-insensitively. + .filter(row -> "READY".equalsIgnoreCase(row.getString("connection_stage"))) + .filter(row -> clientOptions(row).containsKey("SESSION_ID")) + .collect(Collectors.toList()); + } + + private Map clientOptions(Row row) { + Map options = row.getMap("client_options", String.class, String.class); + return options == null ? Collections.emptyMap() : options; + } +} diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/config/DriverConfigReportingSimulacronIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/config/DriverConfigReportingSimulacronIT.java new file mode 100644 index 00000000000..f1306fc12c9 --- /dev/null +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/config/DriverConfigReportingSimulacronIT.java @@ -0,0 +1,204 @@ +/* + * 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.core.config; + +import static com.datastax.oss.driver.internal.core.context.DefaultDriverConfigReporter.DRIVER_CONFIG_KEY; +import static com.datastax.oss.driver.internal.core.context.DefaultDriverConfigReporter.SESSION_ID_KEY; +import static com.datastax.oss.driver.internal.core.context.StartupOptionsBuilder.CLIENT_ID_KEY; +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +import com.datastax.oss.driver.api.core.CqlSession; +import com.datastax.oss.driver.api.core.config.DefaultDriverOption; +import com.datastax.oss.driver.api.core.config.DriverConfigLoader; +import com.datastax.oss.driver.api.testinfra.session.SessionUtils; +import com.datastax.oss.driver.api.testinfra.simulacron.SimulacronRule; +import com.datastax.oss.driver.categories.ParallelizableTests; +import com.datastax.oss.protocol.internal.request.Register; +import com.datastax.oss.protocol.internal.request.Startup; +import com.datastax.oss.simulacron.common.cluster.ClusterSpec; +import com.datastax.oss.simulacron.common.cluster.QueryLog; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.net.SocketAddress; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import org.junit.Before; +import org.junit.ClassRule; +import org.junit.Test; +import org.junit.experimental.categories.Category; + +/** + * End-to-end check of driver-configuration reporting against a mock server, asserting on the actual + * CQL {@code STARTUP} frames the driver sends. + * + *

Simulacron records every inbound frame with its originating client connection, so we can + * verify that when {@code advanced.driver-config-reporting.enabled} is: + * + *

+ * + *

The control connection is identified independently of the reported options: it is the only + * connection that issues a {@code REGISTER} frame (to subscribe to cluster events). + * + *

Assertions are scoped to the session's real connections, identified by the {@code CLIENT_ID} + * startup option that the driver always sends (independently of config reporting). This excludes + * the short-lived connections opened by protocol-version negotiation: when the protocol version is + * not pinned (the default), the driver first tries higher versions (DSE_V2, DSE_V1, V5) that + * ScyllaDB rejects before the handshake completes, and Simulacron records each such rejected + * attempt as a bare {@code STARTUP} ({@code CQL_VERSION} only) that carries no driver identity. + */ +@Category(ParallelizableTests.class) +public class DriverConfigReportingSimulacronIT { + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + // A single node yields one dedicated control connection plus a pool connection (local.size + // defaults to 1), i.e. at least two distinct session connections of which only the control one + // registers for events. + @ClassRule + public static final SimulacronRule SIMULACRON_RULE = + new SimulacronRule(ClusterSpec.builder().withNodes(1)); + + @Before + public void clearLogs() { + SIMULACRON_RULE.cluster().clearLogs(); + } + + @Test + public void should_report_session_id_on_all_connections_and_driver_config_only_on_control() { + DriverConfigLoader loader = + SessionUtils.configLoaderBuilder() + .withBoolean(DefaultDriverOption.DRIVER_CONFIG_REPORTING_ENABLED, true) + .build(); + try (CqlSession session = SessionUtils.newSession(SIMULACRON_RULE, loader)) { + awaitControlAndPoolConnected(); + + SocketAddress controlConnection = controlConnection().orElseThrow(AssertionError::new); + List startups = sessionStartups(); + + // Sanity: we are actually observing more than one connection (control + at least one pool), + // otherwise the "only on the control connection" assertion below would be vacuous. + assertThat(distinctConnections(startups)).isGreaterThanOrEqualTo(2); + + // SESSION_ID is present on every session connection and has the same value everywhere. + assertThat(startups).allSatisfy(log -> assertThat(options(log)).containsKey(SESSION_ID_KEY)); + assertThat(startups.stream().map(log -> options(log).get(SESSION_ID_KEY)).distinct()) + .hasSize(1); + + // DRIVER_CONFIG is present on exactly one connection, and that connection is the control one. + List withDriverConfig = + startups.stream() + .filter(log -> options(log).containsKey(DRIVER_CONFIG_KEY)) + .collect(Collectors.toList()); + assertThat(withDriverConfig).hasSize(1); + assertThat(withDriverConfig.get(0).getConnection()).isEqualTo(controlConnection); + + // The payload is the stage-1 report: valid JSON carrying exactly the schema version. + assertStageOnePayload(options(withDriverConfig.get(0)).get(DRIVER_CONFIG_KEY)); + } + } + + /** + * Asserts that a {@code DRIVER_CONFIG} value is the stage-1 payload: well-formed JSON whose + * {@code version} is the integer {@code 1}. Guards against an incorrect schema version or a + * malformed blob slipping through a mere key-presence check. + */ + private static void assertStageOnePayload(String driverConfig) { + JsonNode root; + try { + root = OBJECT_MAPPER.readTree(driverConfig); + } catch (JsonProcessingException e) { + throw new AssertionError("DRIVER_CONFIG is not valid JSON: " + driverConfig, e); + } + assertThat(root.path("version").isInt()) + .as("version is an integer in %s", driverConfig) + .isTrue(); + assertThat(root.path("version").intValue()).isEqualTo(1); + } + + @Test + public void should_report_nothing_when_disabled() { + DriverConfigLoader loader = + SessionUtils.configLoaderBuilder() + .withBoolean(DefaultDriverOption.DRIVER_CONFIG_REPORTING_ENABLED, false) + .build(); + try (CqlSession session = SessionUtils.newSession(SIMULACRON_RULE, loader)) { + awaitControlAndPoolConnected(); + + List startups = sessionStartups(); + assertThat(distinctConnections(startups)).isGreaterThanOrEqualTo(2); + + // Neither option is sent on any session connection: zero change on the wire when disabled. + assertThat(startups) + .allSatisfy( + log -> + assertThat(options(log)) + .doesNotContainKey(SESSION_ID_KEY) + .doesNotContainKey(DRIVER_CONFIG_KEY)); + } + } + + /** + * The {@code STARTUP} frames of the session's real connections (control + pool), identified by + * the always-present {@code CLIENT_ID} option; excludes protocol-version negotiation attempts. + */ + private List sessionStartups() { + return SIMULACRON_RULE.cluster().getLogs().getQueryLogs().stream() + .filter(log -> log.getFrame().message instanceof Startup) + .filter(log -> options(log).containsKey(CLIENT_ID_KEY)) + .collect(Collectors.toList()); + } + + /** The {@code STARTUP} options carried by a recorded frame. */ + private Map options(QueryLog log) { + return ((Startup) log.getFrame().message).options; + } + + /** The connection that issued a {@code REGISTER} frame, i.e. the control connection. */ + private Optional controlConnection() { + return SIMULACRON_RULE.cluster().getLogs().getQueryLogs().stream() + .filter(log -> log.getFrame().message instanceof Register) + .map(QueryLog::getConnection) + .findFirst(); + } + + private long distinctConnections(List logs) { + return logs.stream().map(QueryLog::getConnection).distinct().count(); + } + + /** + * Waits until both the control connection (identified by its {@code REGISTER}) and at least one + * pool connection have sent their {@code STARTUP}, since pool connections may finish initializing + * slightly after the session builder returns. + */ + private void awaitControlAndPoolConnected() { + await() + .pollInterval(100, TimeUnit.MILLISECONDS) + .atMost(30, TimeUnit.SECONDS) + .until( + () -> controlConnection().isPresent() && distinctConnections(sessionStartups()) >= 2); + } +}