diff --git a/bom/pom.xml b/bom/pom.xml index dd76153a9b1..66f75e89490 100644 --- a/bom/pom.xml +++ b/bom/pom.xml @@ -78,7 +78,7 @@ com.datastax.oss native-protocol - 1.5.2 + 1.5.3-SNAPSHOT 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 2900e897cce..93ce5c018f6 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 @@ -1041,7 +1041,15 @@ 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 to register for GRACEFUL_DISCONNECT events from the server (CEP-59). When enabled and + * the server advertises support, the driver will gracefully drain connections when a node shuts + * down. + * + *

Value-type: boolean + */ + GRACEFUL_DISCONNECT_ENABLED("advanced.connection.graceful-disconnect-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 98faf3e590c..7163dda0b43 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 @@ -272,6 +272,7 @@ protected static void fillWithDriverDefaults(OptionsMap map) { map.put(TypedDriverOption.CONNECTION_MAX_REQUESTS, 1024); map.put(TypedDriverOption.CONNECTION_MAX_ORPHAN_REQUESTS, 256); map.put(TypedDriverOption.CONNECTION_WARN_INIT_ERROR, true); + map.put(TypedDriverOption.GRACEFUL_DISCONNECT_ENABLED, true); map.put(TypedDriverOption.RECONNECT_ON_INIT, false); map.put(TypedDriverOption.RECONNECTION_POLICY_CLASS, "ExponentialReconnectionPolicy"); map.put(TypedDriverOption.RECONNECTION_BASE_DELAY, Duration.ofSeconds(1)); 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 182753300e7..089984c40e3 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 @@ -914,6 +914,9 @@ public String toString() { new TypedDriverOption<>( DefaultDriverOption.ADDRESS_TRANSLATOR_RESOLVE_ADDRESSES, GenericType.BOOLEAN); + public static final TypedDriverOption GRACEFUL_DISCONNECT_ENABLED = + new TypedDriverOption<>(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, GenericType.BOOLEAN); + /** * Ordered preference list of remote dcs optionally supplied for automatic failover and included * in query plan. This feature is enabled only when max-nodes-per-remote-dc is greater than 0. diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultNodeMetric.java b/core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultNodeMetric.java index 0e9934c7034..918621161f1 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultNodeMetric.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultNodeMetric.java @@ -51,6 +51,7 @@ public enum DefaultNodeMetric implements NodeMetric { SPECULATIVE_EXECUTIONS("speculative-executions"), CONNECTION_INIT_ERRORS("errors.connection.init"), AUTHENTICATION_ERRORS("errors.connection.auth"), + GRACEFUL_DISCONNECTS("pool.graceful-disconnects"), ; private static final Map BY_PATH = sortByPath(); diff --git a/core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultSessionMetric.java b/core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultSessionMetric.java index 63027a23fe7..3dc9dc2b5fd 100644 --- a/core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultSessionMetric.java +++ b/core/src/main/java/com/datastax/oss/driver/api/core/metrics/DefaultSessionMetric.java @@ -32,6 +32,7 @@ public enum DefaultSessionMetric implements SessionMetric { THROTTLING_QUEUE_SIZE("throttling.queue-size"), THROTTLING_ERRORS("throttling.errors"), CQL_PREPARED_CACHE_SIZE("cql-prepared-cache-size"), + GRACEFUL_DISCONNECTS("graceful-disconnects"), ; private static final Map BY_PATH = sortByPath(); diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelFactory.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelFactory.java index 66a5c4edc0e..c420c53ff13 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelFactory.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelFactory.java @@ -39,6 +39,7 @@ import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting; import com.datastax.oss.driver.shaded.guava.common.base.Preconditions; import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap; +import com.datastax.oss.protocol.internal.ProtocolConstants; import io.netty.bootstrap.Bootstrap; import io.netty.channel.Channel; import io.netty.channel.ChannelFuture; @@ -339,6 +340,11 @@ protected void initChannel(Channel channel) { options.eventCallback, options.ownerLogPrefix); HeartbeatHandler heartbeatHandler = new HeartbeatHandler(defaultConfig); + // Channels that register for GRACEFUL_DISCONNECT always query OPTIONS, so that support + // can be checked against this channel's own SUPPORTED response. + boolean querySupportedOptions = + productType == null + || options.eventTypes.contains(ProtocolConstants.EventType.GRACEFUL_DISCONNECT); ProtocolInitHandler initHandler = new ProtocolInitHandler( context, @@ -347,7 +353,7 @@ protected void initChannel(Channel channel) { endPoint, options, heartbeatHandler, - productType == null); + querySupportedOptions); ChannelPipeline pipeline = channel.pipeline(); context diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java index 90b02f358cd..137a576477e 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java @@ -32,7 +32,9 @@ import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSet; import com.datastax.oss.protocol.internal.Frame; import com.datastax.oss.protocol.internal.Message; +import com.datastax.oss.protocol.internal.ProtocolConstants; import com.datastax.oss.protocol.internal.request.Query; +import com.datastax.oss.protocol.internal.response.Event; import com.datastax.oss.protocol.internal.response.result.SetKeyspace; import io.netty.channel.ChannelDuplexHandler; import io.netty.channel.ChannelFuture; @@ -218,6 +220,12 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception if (streamId < 0) { Message event = responseFrame.message; + if (event instanceof Event + && ProtocolConstants.EventType.GRACEFUL_DISCONNECT.equals(((Event) event).type)) { + // Start draining this channel first, so that the drain is not compromised if the + // callback below misbehaves. + startGracefulShutdown(ctx); + } if (eventCallback == null) { LOG.debug("[{}] Received event {} but no callback was registered", logPrefix, event); } else { 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 8a426f7b368..c53ea4e609b 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 @@ -36,6 +36,7 @@ import com.datastax.oss.driver.internal.core.protocol.SegmentToFrameDecoder; import com.datastax.oss.driver.internal.core.util.ProtocolUtils; import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions; +import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting; import com.datastax.oss.protocol.internal.Message; import com.datastax.oss.protocol.internal.ProtocolConstants; import com.datastax.oss.protocol.internal.ProtocolConstants.ErrorCode; @@ -55,7 +56,9 @@ import io.netty.channel.ChannelHandlerContext; import io.netty.channel.ChannelPipeline; import java.nio.ByteBuffer; +import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Objects; import net.jcip.annotations.NotThreadSafe; import org.slf4j.Logger; @@ -140,6 +143,27 @@ protected boolean setConnectSuccess() { return result; } + /** + * Whether a SUPPORTED response advertises the CEP-59 graceful disconnect capability. The server + * may send the key with an explicit {@code "false"} value when the feature is disabled. + */ + @VisibleForTesting + static boolean supportsGracefulDisconnect(Map> supportedOptions) { + if (supportedOptions == null) { + return false; + } + List values = supportedOptions.get(ProtocolConstants.EventType.GRACEFUL_DISCONNECT); + if (values == null) { + return false; + } + for (String value : values) { + if ("false".equalsIgnoreCase(value)) { + return false; + } + } + return true; + } + private enum Step { OPTIONS, STARTUP, @@ -157,10 +181,14 @@ private class InitRequest extends ChannelHandlerRequest { private Message request; private Authenticator authenticator; private ByteBuffer authResponseToken; + // The event types to register for; GRACEFUL_DISCONNECT is removed if this channel's SUPPORTED + // response does not advertise it (capability is negotiated per connection). + private List eventTypes; InitRequest(ChannelHandlerContext ctx) { super(ctx, timeoutMillis); this.step = querySupportedOptions ? Step.OPTIONS : Step.STARTUP; + this.eventTypes = options.eventTypes; } @Override @@ -183,7 +211,7 @@ Message getRequest() { case AUTH_RESPONSE: return request = new AuthResponse(authResponseToken); case REGISTER: - return request = new Register(options.eventTypes); + return request = new Register(eventTypes); default: throw new AssertionError("unhandled step: " + step); } @@ -204,7 +232,13 @@ void onResponse(Message response) { ProtocolUtils.opcodeString(response.opcode)); try { if (step == Step.OPTIONS && response instanceof Supported) { - channel.attr(DriverChannel.OPTIONS_KEY).set(((Supported) response).options); + Map> supportedOptions = ((Supported) response).options; + channel.attr(DriverChannel.OPTIONS_KEY).set(supportedOptions); + if (eventTypes.contains(ProtocolConstants.EventType.GRACEFUL_DISCONNECT) + && !supportsGracefulDisconnect(supportedOptions)) { + eventTypes = new ArrayList<>(eventTypes); + eventTypes.remove(ProtocolConstants.EventType.GRACEFUL_DISCONNECT); + } step = Step.STARTUP; send(); } else if (step == Step.STARTUP && response instanceof Ready) { @@ -303,7 +337,7 @@ void onResponse(Message response) { if (options.keyspace != null) { step = Step.SET_KEYSPACE; send(); - } else if (!options.eventTypes.isEmpty()) { + } else if (!eventTypes.isEmpty()) { step = Step.REGISTER; send(); } else { @@ -311,7 +345,7 @@ void onResponse(Message response) { } } } else if (step == Step.SET_KEYSPACE && response instanceof SetKeyspace) { - if (!options.eventTypes.isEmpty()) { + if (!eventTypes.isEmpty()) { step = Step.REGISTER; send(); } else { 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 5c29a9b704b..637d8fc9eec 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 @@ -26,13 +26,17 @@ import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance; import com.datastax.oss.driver.api.core.metadata.Node; import com.datastax.oss.driver.api.core.metadata.NodeState; +import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric; +import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric; import com.datastax.oss.driver.internal.core.channel.ChannelEvent; import com.datastax.oss.driver.internal.core.channel.DriverChannel; import com.datastax.oss.driver.internal.core.channel.DriverChannelOptions; import com.datastax.oss.driver.internal.core.channel.EventCallback; import com.datastax.oss.driver.internal.core.context.InternalDriverContext; +import com.datastax.oss.driver.internal.core.metadata.DefaultNode; import com.datastax.oss.driver.internal.core.metadata.DefaultTopologyMonitor; import com.datastax.oss.driver.internal.core.metadata.DistanceEvent; +import com.datastax.oss.driver.internal.core.metadata.GracefulDisconnectEvent; import com.datastax.oss.driver.internal.core.metadata.MetadataManager; import com.datastax.oss.driver.internal.core.metadata.NodeStateEvent; import com.datastax.oss.driver.internal.core.metadata.TopologyEvent; @@ -190,6 +194,9 @@ public void onEvent(Message eventMessage) { case ProtocolConstants.EventType.SCHEMA_CHANGE: processSchemaChange(event); break; + case ProtocolConstants.EventType.GRACEFUL_DISCONNECT: + processGracefulDisconnect(); + break; default: LOG.warn("[{}] Unsupported event type: {}", logPrefix, event.type); } @@ -242,6 +249,34 @@ private void processSchemaChange(Event event) { }); } + private void processGracefulDisconnect() { + LOG.info( + "[{}] Received GRACEFUL_DISCONNECT event on control connection, " + + "the server is shutting down gracefully", + logPrefix); + context + .getMetricsFactory() + .getSessionUpdater() + .incrementCounter(DefaultSessionMetric.GRACEFUL_DISCONNECTS, null); + // Fire an internal event to notify other components (particularly the ChannelPool) + DriverChannel currentChannel = channel; + if (currentChannel != null) { + context + .getMetadataManager() + .getMetadata() + .findNode(currentChannel.getEndPoint()) + .ifPresent( + node -> { + if (node instanceof DefaultNode) { + ((DefaultNode) node) + .getMetricUpdater() + .incrementCounter(DefaultNodeMetric.GRACEFUL_DISCONNECTS, null); + } + context.getEventBus().fire(new GracefulDisconnectEvent(node)); + }); + } + } + private class SingleThreaded { private final InternalDriverContext context; private final DriverConfig config; @@ -292,7 +327,13 @@ private void init( } initWasCalled = true; try { - ImmutableList eventTypes = buildEventTypes(listenToClusterEvents); + boolean gracefulDisconnectEnabled = + context + .getConfig() + .getDefaultProfile() + .getBoolean(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, true); + ImmutableList eventTypes = + buildEventTypes(listenToClusterEvents, gracefulDisconnectEnabled); LOG.debug("[{}] Initializing with event types {}", logPrefix, eventTypes); channelOptions = DriverChannelOptions.builder() @@ -606,7 +647,8 @@ private boolean isAuthFailure(Throwable error) { return true; } - private static ImmutableList buildEventTypes(boolean listenClusterEvents) { + private static ImmutableList buildEventTypes( + boolean listenClusterEvents, boolean gracefulDisconnectEnabled) { ImmutableList.Builder builder = ImmutableList.builder(); builder.add(ProtocolConstants.EventType.SCHEMA_CHANGE); if (listenClusterEvents) { @@ -614,6 +656,9 @@ private static ImmutableList buildEventTypes(boolean listenClusterEvents .add(ProtocolConstants.EventType.STATUS_CHANGE) .add(ProtocolConstants.EventType.TOPOLOGY_CHANGE); } + if (gracefulDisconnectEnabled) { + builder.add(ProtocolConstants.EventType.GRACEFUL_DISCONNECT); + } return builder.build(); } } diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/GracefulDisconnectEvent.java b/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/GracefulDisconnectEvent.java new file mode 100644 index 00000000000..a595541cc8f --- /dev/null +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/metadata/GracefulDisconnectEvent.java @@ -0,0 +1,59 @@ +/* + * 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.metadata; + +import com.datastax.oss.driver.api.core.metadata.Node; +import java.util.Objects; +import net.jcip.annotations.Immutable; + +/** + * Indicates that a node announced a graceful shutdown (CEP-59): a {@code GRACEFUL_DISCONNECT} + * protocol event was received on one of its connections. + */ +@Immutable +public class GracefulDisconnectEvent { + + /** The node that is shutting down. */ + public final Node node; + + public GracefulDisconnectEvent(Node node) { + this.node = node; + } + + @Override + public boolean equals(Object other) { + if (other == this) { + return true; + } else if (other instanceof GracefulDisconnectEvent) { + GracefulDisconnectEvent that = (GracefulDisconnectEvent) other; + return Objects.equals(this.node, that.node); + } else { + return false; + } + } + + @Override + public int hashCode() { + return Objects.hash(this.node); + } + + @Override + public String toString() { + return "GracefulDisconnectEvent(" + node + ")"; + } +} diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardNodeMetricUpdater.java b/core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardNodeMetricUpdater.java index 2e5e6c8db3d..4669eb074cb 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardNodeMetricUpdater.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardNodeMetricUpdater.java @@ -71,6 +71,7 @@ public DropwizardNodeMetricUpdater( initializeCounter(DefaultNodeMetric.SPECULATIVE_EXECUTIONS, profile); initializeCounter(DefaultNodeMetric.CONNECTION_INIT_ERRORS, profile); initializeCounter(DefaultNodeMetric.AUTHENTICATION_ERRORS, profile); + initializeCounter(DefaultNodeMetric.GRACEFUL_DISCONNECTS, profile); initializeHdrTimer( DefaultNodeMetric.CQL_MESSAGES, diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardSessionMetricUpdater.java b/core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardSessionMetricUpdater.java index 94e10ad6936..0bc2b3bf242 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardSessionMetricUpdater.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/metrics/DropwizardSessionMetricUpdater.java @@ -45,6 +45,7 @@ public DropwizardSessionMetricUpdater( initializeCounter(DefaultSessionMetric.CQL_CLIENT_TIMEOUTS, profile); initializeCounter(DefaultSessionMetric.THROTTLING_ERRORS, profile); + initializeCounter(DefaultSessionMetric.GRACEFUL_DISCONNECTS, profile); initializeCounter(DseSessionMetric.GRAPH_CLIENT_TIMEOUTS, profile); initializeHdrTimer( diff --git a/core/src/main/java/com/datastax/oss/driver/internal/core/pool/ChannelPool.java b/core/src/main/java/com/datastax/oss/driver/internal/core/pool/ChannelPool.java index 6b7d06045bd..42d16b779a5 100644 --- a/core/src/main/java/com/datastax/oss/driver/internal/core/pool/ChannelPool.java +++ b/core/src/main/java/com/datastax/oss/driver/internal/core/pool/ChannelPool.java @@ -28,23 +28,31 @@ import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance; import com.datastax.oss.driver.api.core.metadata.Node; import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric; +import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric; import com.datastax.oss.driver.internal.core.channel.ChannelEvent; import com.datastax.oss.driver.internal.core.channel.ChannelFactory; import com.datastax.oss.driver.internal.core.channel.ClusterNameMismatchException; import com.datastax.oss.driver.internal.core.channel.DriverChannel; import com.datastax.oss.driver.internal.core.channel.DriverChannelOptions; +import com.datastax.oss.driver.internal.core.channel.EventCallback; import com.datastax.oss.driver.internal.core.config.ConfigChangeEvent; import com.datastax.oss.driver.internal.core.context.EventBus; import com.datastax.oss.driver.internal.core.context.InternalDriverContext; import com.datastax.oss.driver.internal.core.metadata.DefaultNode; +import com.datastax.oss.driver.internal.core.metadata.GracefulDisconnectEvent; import com.datastax.oss.driver.internal.core.metadata.TopologyEvent; +import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater; import com.datastax.oss.driver.internal.core.util.Loggers; import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures; import com.datastax.oss.driver.internal.core.util.concurrent.Reconnection; import com.datastax.oss.driver.internal.core.util.concurrent.RunOrSchedule; import com.datastax.oss.driver.internal.core.util.concurrent.UncaughtExceptions; import com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting; +import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList; import com.datastax.oss.driver.shaded.guava.common.collect.Sets; +import com.datastax.oss.protocol.internal.Message; +import com.datastax.oss.protocol.internal.ProtocolConstants; +import com.datastax.oss.protocol.internal.response.Event; import edu.umd.cs.findbugs.annotations.NonNull; import io.netty.util.concurrent.EventExecutor; import io.netty.util.concurrent.Future; @@ -228,11 +236,14 @@ private class SingleThreaded { private final DriverConfig config; private final ChannelFactory channelFactory; private final EventBus eventBus; + private final SessionMetricUpdater sessionMetricUpdater; + private final boolean gracefulDisconnectEnabled; // The channels that are currently connecting private final List> pendingChannels = new ArrayList<>(); private final Set closingChannels = new HashSet<>(); private final Reconnection reconnection; private final Object configListenerKey; + private final Object gracefulDisconnectListenerKey; private NodeDistance distance; private int wantedCount; @@ -252,6 +263,11 @@ private SingleThreaded( this.wantedCount = getConfiguredSize(distance); this.channelFactory = context.getChannelFactory(); this.eventBus = context.getEventBus(); + this.sessionMetricUpdater = context.getMetricsFactory().getSessionUpdater(); + this.gracefulDisconnectEnabled = + config + .getDefaultProfile() + .getBoolean(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, true); ReconnectionPolicy reconnectionPolicy = context.getReconnectionPolicy(); this.reconnection = new Reconnection( @@ -264,6 +280,10 @@ private SingleThreaded( this.configListenerKey = eventBus.register( ConfigChangeEvent.class, RunOrSchedule.on(adminExecutor, this::onConfigChanged)); + this.gracefulDisconnectListenerKey = + eventBus.register( + GracefulDisconnectEvent.class, + RunOrSchedule.on(adminExecutor, this::onGracefulDisconnect)); } private void connect() { @@ -291,12 +311,20 @@ private CompletionStage addMissingChannels() { int missing = wantedCount - channels.size(); LOG.debug("[{}] Trying to create {} missing channels", logPrefix, missing); - DriverChannelOptions options = - DriverChannelOptions.builder() - .withKeyspace(keyspaceName) - .withOwnerLogPrefix(sessionLogPrefix) - .build(); + for (int i = 0; i < missing; i++) { + DriverChannelOptions.Builder optionsBuilder = + DriverChannelOptions.builder() + .withKeyspace(keyspaceName) + .withOwnerLogPrefix(sessionLogPrefix); + + if (gracefulDisconnectEnabled) { + optionsBuilder.withEvents( + ImmutableList.of(ProtocolConstants.EventType.GRACEFUL_DISCONNECT), + new QueryConnectionEventCallback()); + } + + DriverChannelOptions options = optionsBuilder.build(); CompletionStage channelFuture = channelFactory.connect(node, options); pendingChannels.add(channelFuture); } @@ -474,6 +502,55 @@ private void onConfigChanged(@SuppressWarnings("unused") ConfigChangeEvent event resize(distance); } + private void onGracefulDisconnect(GracefulDisconnectEvent event) { + assert adminExecutor.inEventLoop(); + if (!event.node.equals(node)) { + return; + } + if (channels.size() == 0) { + return; + } + LOG.info( + "[{}] Received GRACEFUL_DISCONNECT for {}, closing all channels for this node gracefully", + logPrefix, + node); + // The graceful close allows in-flight requests to complete before channels are fully closed. + for (DriverChannel channel : channels) { + channel.close(); + } + } + + /** + * Event callback for query connections that handles GRACEFUL_DISCONNECT events. + * + *

This is called from the Netty I/O thread when an event is received on a query connection. + */ + private class QueryConnectionEventCallback implements EventCallback { + @Override + public void onEvent(Message eventMessage) { + if (!(eventMessage instanceof Event)) { + LOG.warn( + "[{}] Unsupported event class on query connection: {}", + logPrefix, + eventMessage.getClass().getName()); + return; + } + Event event = (Event) eventMessage; + if (ProtocolConstants.EventType.GRACEFUL_DISCONNECT.equals(event.type)) { + LOG.debug("[{}] Received GRACEFUL_DISCONNECT on query connection", logPrefix); + if (node instanceof DefaultNode) { + ((DefaultNode) node) + .getMetricUpdater() + .incrementCounter(DefaultNodeMetric.GRACEFUL_DISCONNECTS, null); + } + sessionMetricUpdater.incrementCounter(DefaultSessionMetric.GRACEFUL_DISCONNECTS, null); + eventBus.fire(new GracefulDisconnectEvent(node)); + } else { + LOG.warn("[{}] Unexpected event type on query connection: {}", logPrefix, event.type); + } + } + } + private CompletionStage setKeyspace(CqlIdentifier newKeyspaceName) { assert adminExecutor.inEventLoop(); if (setKeyspaceFuture != null && !setKeyspaceFuture.isDone()) { @@ -533,6 +610,7 @@ private void close() { reconnection.stop(); eventBus.unregister(configListenerKey, ConfigChangeEvent.class); + eventBus.unregister(gracefulDisconnectListenerKey, GracefulDisconnectEvent.class); // Close all channels, the pool future completes when all the channels futures have completed int toClose = closingChannels.size() + channels.size(); diff --git a/core/src/main/resources/reference.conf b/core/src/main/resources/reference.conf index 4ae83362e29..8f7a60141bf 100644 --- a/core/src/main/resources/reference.conf +++ b/core/src/main/resources/reference.conf @@ -535,6 +535,15 @@ datastax-java-driver { # change. # Overridable in a profile: no warn-on-init-error = true + + # Whether to subscribe to GRACEFUL_DISCONNECT events (CEP-59). When the server supports it, + # the driver will drain in-flight requests before closing connections during a node shutdown, + # instead of failing them with a timeout or connection error. + # + # Required: yes + # Modifiable at runtime: no + # Overridable in a profile: no + graceful-disconnect-enabled = true } # Advanced options for the built-in load-balancing policies. @@ -1557,6 +1566,10 @@ datastax-java-driver { # a Counter) // throttling.errors, + # The number of GRACEFUL_DISCONNECT events received from nodes that are shutting down + # gracefully (CEP-59), across all connections of the session (exposed as a Counter). + // graceful-disconnects, + # The throughput and latency percentiles of DSE continuous CQL requests (exposed as a # Timer). # @@ -1724,6 +1737,10 @@ datastax-java-driver { # See the description of the connection.max-orphan-requests option for more details. // pool.orphaned-streams, + # The number of GRACEFUL_DISCONNECT events received on this node's pooled connections, + # indicating that the node is shutting down gracefully (CEP-59) (exposed as a Counter). + // pool.graceful-disconnects, + # The number and rate of bytes sent to this node (exposed as a Meter if available, otherwise # as a Counter). // bytes-sent, diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/InFlightHandlerTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/InFlightHandlerTest.java index 35049e99af1..b776b41462e 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/InFlightHandlerTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/InFlightHandlerTest.java @@ -18,8 +18,10 @@ package com.datastax.oss.driver.internal.core.channel; import static com.datastax.oss.driver.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -33,6 +35,7 @@ import com.datastax.oss.protocol.internal.ProtocolConstants; import com.datastax.oss.protocol.internal.request.Query; import com.datastax.oss.protocol.internal.response.Error; +import com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent; import com.datastax.oss.protocol.internal.response.event.StatusChangeEvent; import com.datastax.oss.protocol.internal.response.result.SetKeyspace; import com.datastax.oss.protocol.internal.response.result.Void; @@ -644,6 +647,136 @@ private void addToPipeline() { addToPipelineWithEventCallback(null); } + @Test + public void should_initiate_graceful_drain_on_graceful_disconnect_event() { + // Given + EventCallback eventCallback = mock(EventCallback.class); + addToPipelineWithEventCallback(eventCallback); + when(streamIds.acquire()).thenReturn(42); + MockResponseCallback responseCallback = new MockResponseCallback(); + channel + .writeAndFlush( + new DriverChannel.RequestMessage(QUERY, false, Frame.NO_PAYLOAD, responseCallback)) + .awaitUninterruptibly(); + + // When + GracefulDisconnectEvent gracefulDisconnectEvent = + new com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent(); + Frame eventFrame = + Frame.forResponse( + DefaultProtocolVersion.V4.getCode(), + -1, + null, + Collections.emptyMap(), + Collections.emptyList(), + gracefulDisconnectEvent); + writeInboundFrame(eventFrame); + + // Then + // channel not closed yet because there is a pending request + assertThat(channel.closeFuture()).isNotDone(); + // callback was still notified + verify(eventCallback).onEvent(gracefulDisconnectEvent); + // new writes should be refused + ChannelFuture otherWriteFuture = + channel.writeAndFlush( + new DriverChannel.RequestMessage( + QUERY, false, Frame.NO_PAYLOAD, new MockResponseCallback())); + assertThat(otherWriteFuture) + .isFailed(e -> assertThat(e).isInstanceOf(IllegalStateException.class)); + + // When the pending request completes + Frame requestFrame = readOutboundFrame(); + writeInboundFrame(requestFrame, Void.INSTANCE); + + // Then the channel closes + assertThat(channel.closeFuture()).isSuccess(); + } + + @Test + public void should_close_immediately_on_graceful_disconnect_if_no_pending() { + // Given + EventCallback eventCallback = mock(EventCallback.class); + addToPipelineWithEventCallback(eventCallback); + + // When + GracefulDisconnectEvent gracefulDisconnectEvent = + new com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent(); + Frame eventFrame = + Frame.forResponse( + DefaultProtocolVersion.V4.getCode(), + -1, + null, + Collections.emptyMap(), + Collections.emptyList(), + gracefulDisconnectEvent); + writeInboundFrame(eventFrame); + + // Then + assertThat(channel.closeFuture()).isSuccess(); + verify(eventCallback).onEvent(gracefulDisconnectEvent); + } + + @Test + public void should_handle_duplicate_graceful_disconnect_events() { + // The server-side CEP-59 implementation is still evolving; a node might emit the event more + // than once (e.g. once per registered connection, or on a drain retry). The second event must + // not disrupt the drain already in progress. + // Given + EventCallback eventCallback = mock(EventCallback.class); + addToPipelineWithEventCallback(eventCallback); + when(streamIds.acquire()).thenReturn(42); + MockResponseCallback responseCallback = new MockResponseCallback(); + channel + .writeAndFlush( + new DriverChannel.RequestMessage(QUERY, false, Frame.NO_PAYLOAD, responseCallback)) + .awaitUninterruptibly(); + + // When: the same event is received twice while a request is still pending + for (int i = 0; i < 2; i++) { + Frame eventFrame = + Frame.forResponse( + DefaultProtocolVersion.V4.getCode(), + -1, + null, + Collections.emptyMap(), + Collections.emptyList(), + new com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent()); + writeInboundFrame(eventFrame); + } + + // Then: still draining, not closed abruptly + assertThat(channel.closeFuture()).isNotDone(); + verify(eventCallback, times(2)).onEvent(any()); + + // When the pending request completes, the drain finishes normally + Frame requestFrame = readOutboundFrame(); + writeInboundFrame(requestFrame, Void.INSTANCE); + assertThat(channel.closeFuture()).isSuccess(); + } + + @Test + public void should_handle_graceful_disconnect_without_event_callback() { + // Given + addToPipeline(); // no event callback + + // When + GracefulDisconnectEvent gracefulDisconnectEvent = + new com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent(); + Frame eventFrame = + Frame.forResponse( + DefaultProtocolVersion.V4.getCode(), + -1, + null, + Collections.emptyMap(), + Collections.emptyList(), + gracefulDisconnectEvent); + writeInboundFrame(eventFrame); + + // Then + assertThat(channel.closeFuture()).isSuccess(); + } + private void addToPipelineWithEventCallback(EventCallback eventCallback) { channel .pipeline() diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandlerGracefulDisconnectTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandlerGracefulDisconnectTest.java new file mode 100644 index 00000000000..858d970b7e7 --- /dev/null +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/channel/ProtocolInitHandlerGracefulDisconnectTest.java @@ -0,0 +1,231 @@ +/* + * 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.channel; + +import static com.datastax.oss.driver.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.datastax.oss.driver.api.core.DefaultProtocolVersion; +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.datastax.oss.driver.api.core.metadata.EndPoint; +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.InternalDriverContext; +import com.datastax.oss.driver.internal.core.metadata.TestNodeFactory; +import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList; +import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap; +import com.datastax.oss.protocol.internal.Frame; +import com.datastax.oss.protocol.internal.ProtocolConstants; +import com.datastax.oss.protocol.internal.request.Options; +import com.datastax.oss.protocol.internal.request.Register; +import com.datastax.oss.protocol.internal.request.Startup; +import com.datastax.oss.protocol.internal.response.Ready; +import com.datastax.oss.protocol.internal.response.Supported; +import io.netty.channel.ChannelFuture; +import java.net.InetSocketAddress; +import java.time.Duration; +import java.util.List; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +/** + * Coverage for GRACEFUL_DISCONNECT (CEP-59) registration during channel initialization: support is + * checked against each channel's own SUPPORTED response, and the event type is only included in + * REGISTER if the server advertises it. + */ +public class ProtocolInitHandlerGracefulDisconnectTest extends ChannelHandlerTestBase { + + private static final long QUERY_TIMEOUT_MILLIS = 100L; + private static final EndPoint END_POINT = TestNodeFactory.newEndPoint(1); + private static final Supported SUPPORTED_WITH_GRACEFUL_DISCONNECT = + new Supported( + ImmutableMap.of( + ProtocolConstants.EventType.GRACEFUL_DISCONNECT, + ImmutableList.of("true"), + "CQL_VERSION", + ImmutableList.of("3.4.7"))); + private static final Supported SUPPORTED_WITHOUT_GRACEFUL_DISCONNECT = + new Supported(ImmutableMap.of("CQL_VERSION", ImmutableList.of("3.4.7"))); + + @Mock private InternalDriverContext internalDriverContext; + @Mock private DriverConfig driverConfig; + @Mock private DriverExecutionProfile defaultProfile; + + private final ProtocolVersionRegistry protocolVersionRegistry = + new DefaultProtocolVersionRegistry("test"); + private HeartbeatHandler heartbeatHandler; + + @Before + @Override + public void setup() { + super.setup(); + MockitoAnnotations.initMocks(this); + when(internalDriverContext.getConfig()).thenReturn(driverConfig); + when(driverConfig.getDefaultProfile()).thenReturn(defaultProfile); + when(defaultProfile.getDuration(DefaultDriverOption.CONNECTION_INIT_QUERY_TIMEOUT)) + .thenReturn(Duration.ofMillis(QUERY_TIMEOUT_MILLIS)); + when(defaultProfile.getDuration(DefaultDriverOption.HEARTBEAT_INTERVAL)) + .thenReturn(Duration.ofSeconds(30)); + when(internalDriverContext.getProtocolVersionRegistry()).thenReturn(protocolVersionRegistry); + + channel + .pipeline() + .addLast( + ChannelFactory.INFLIGHT_HANDLER_NAME, + new InFlightHandler( + DefaultProtocolVersion.V4, + new StreamIdGenerator(100), + Integer.MAX_VALUE, + 100, + channel.newPromise(), + null, + "test")); + + heartbeatHandler = new HeartbeatHandler(defaultProfile); + } + + private ChannelFuture connectWithEvents(List eventTypes) { + DriverChannelOptions driverChannelOptions = + DriverChannelOptions.builder().withEvents(eventTypes, mock(EventCallback.class)).build(); + channel + .pipeline() + .addLast( + ChannelFactory.INIT_HANDLER_NAME, + new ProtocolInitHandler( + internalDriverContext, + DefaultProtocolVersion.V4, + null, + END_POINT, + driverChannelOptions, + heartbeatHandler, + true)); + return channel.connect(new InetSocketAddress("localhost", 9042)); + } + + /** Completes the OPTIONS and STARTUP steps. */ + private void initUntilAfterClusterName(Supported supportedResponse) { + Frame optionsFrame = readOutboundFrame(); + assertThat(optionsFrame.message).isInstanceOf(Options.class); + writeInboundFrame(optionsFrame, supportedResponse); + Frame startupFrame = readOutboundFrame(); + assertThat(startupFrame.message).isInstanceOf(Startup.class); + writeInboundFrame(startupFrame, new Ready()); + writeInboundFrame(readOutboundFrame(), TestResponses.clusterNameResponse("someClusterName")); + } + + @Test + public void should_register_graceful_disconnect_when_advertised_in_supported() { + ChannelFuture connectFuture = + connectWithEvents( + ImmutableList.of("STATUS_CHANGE", ProtocolConstants.EventType.GRACEFUL_DISCONNECT)); + + initUntilAfterClusterName(SUPPORTED_WITH_GRACEFUL_DISCONNECT); + Frame registerFrame = readOutboundFrame(); + assertThat(registerFrame.message).isInstanceOf(Register.class); + + assertThat(((Register) registerFrame.message).eventTypes) + .containsExactly("STATUS_CHANGE", ProtocolConstants.EventType.GRACEFUL_DISCONNECT); + writeInboundFrame(registerFrame, new Ready()); + assertThat(connectFuture).isSuccess(); + } + + @Test + public void should_not_register_graceful_disconnect_when_server_does_not_advertise_it() { + ChannelFuture connectFuture = + connectWithEvents( + ImmutableList.of("STATUS_CHANGE", ProtocolConstants.EventType.GRACEFUL_DISCONNECT)); + + initUntilAfterClusterName(SUPPORTED_WITHOUT_GRACEFUL_DISCONNECT); + Frame registerFrame = readOutboundFrame(); + assertThat(registerFrame.message).isInstanceOf(Register.class); + + assertThat(((Register) registerFrame.message).eventTypes).containsExactly("STATUS_CHANGE"); + writeInboundFrame(registerFrame, new Ready()); + assertThat(connectFuture).isSuccess(); + } + + @Test + public void should_not_register_graceful_disconnect_when_advertised_as_false() { + ChannelFuture connectFuture = + connectWithEvents( + ImmutableList.of("STATUS_CHANGE", ProtocolConstants.EventType.GRACEFUL_DISCONNECT)); + + initUntilAfterClusterName( + new Supported( + ImmutableMap.of( + ProtocolConstants.EventType.GRACEFUL_DISCONNECT, + ImmutableList.of("false"), + "CQL_VERSION", + ImmutableList.of("3.4.7")))); + Frame registerFrame = readOutboundFrame(); + assertThat(registerFrame.message).isInstanceOf(Register.class); + + assertThat(((Register) registerFrame.message).eventTypes).containsExactly("STATUS_CHANGE"); + writeInboundFrame(registerFrame, new Ready()); + assertThat(connectFuture).isSuccess(); + } + + @Test + public void should_skip_register_when_graceful_disconnect_was_the_only_event_type() { + // Pool channels only register for GRACEFUL_DISCONNECT; if the server does not support it, + // there is nothing left to register for. + ChannelFuture connectFuture = + connectWithEvents(ImmutableList.of(ProtocolConstants.EventType.GRACEFUL_DISCONNECT)); + + initUntilAfterClusterName(SUPPORTED_WITHOUT_GRACEFUL_DISCONNECT); + + assertThat(connectFuture).isSuccess(); + assertThat((Object) channel.readOutbound()).isNull(); + } + + @Test + public void should_detect_capability_from_supported_options_map() { + assertThat(ProtocolInitHandler.supportsGracefulDisconnect(null)).isFalse(); + assertThat(ProtocolInitHandler.supportsGracefulDisconnect(ImmutableMap.of())).isFalse(); + assertThat( + ProtocolInitHandler.supportsGracefulDisconnect( + ImmutableMap.of("CQL_VERSION", ImmutableList.of("3.4.7")))) + .isFalse(); + assertThat( + ProtocolInitHandler.supportsGracefulDisconnect( + ImmutableMap.of( + ProtocolConstants.EventType.GRACEFUL_DISCONNECT, ImmutableList.of()))) + .isTrue(); + assertThat( + ProtocolInitHandler.supportsGracefulDisconnect( + ImmutableMap.of( + ProtocolConstants.EventType.GRACEFUL_DISCONNECT, ImmutableList.of("true")))) + .isTrue(); + assertThat( + ProtocolInitHandler.supportsGracefulDisconnect( + ImmutableMap.of( + ProtocolConstants.EventType.GRACEFUL_DISCONNECT, ImmutableList.of("false")))) + .isFalse(); + assertThat( + ProtocolInitHandler.supportsGracefulDisconnect( + ImmutableMap.of( + ProtocolConstants.EventType.GRACEFUL_DISCONNECT, ImmutableList.of("FALSE")))) + .isFalse(); + } +} diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionEventsTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionEventsTest.java index cb83b523ebe..6c9311f3636 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionEventsTest.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionEventsTest.java @@ -20,18 +20,25 @@ import static com.datastax.oss.driver.Assertions.assertThat; import static org.awaitility.Awaitility.await; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.datastax.oss.driver.api.core.config.DefaultDriverOption; +import com.datastax.oss.driver.api.core.metadata.Metadata; +import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric; +import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric; import com.datastax.oss.driver.internal.core.channel.DriverChannel; import com.datastax.oss.driver.internal.core.channel.DriverChannelOptions; import com.datastax.oss.driver.internal.core.channel.EventCallback; +import com.datastax.oss.driver.internal.core.metadata.GracefulDisconnectEvent; import com.datastax.oss.driver.internal.core.metadata.TopologyEvent; import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList; import com.datastax.oss.protocol.internal.ProtocolConstants; import com.datastax.oss.protocol.internal.response.event.SchemaChangeEvent; import com.datastax.oss.protocol.internal.response.event.StatusChangeEvent; import com.datastax.oss.protocol.internal.response.event.TopologyChangeEvent; +import java.util.Optional; import java.util.concurrent.CompletableFuture; import org.junit.Test; import org.mockito.ArgumentCaptor; @@ -41,6 +48,8 @@ public class ControlConnectionEventsTest extends ControlConnectionTestBase { @Test public void should_register_for_all_events_if_topology_requested() { // Given + when(defaultProfile.getBoolean(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, true)) + .thenReturn(true); DriverChannel channel1 = newMockDriverChannel(1); ArgumentCaptor optionsCaptor = ArgumentCaptor.forClass(DriverChannelOptions.class); @@ -59,7 +68,8 @@ public void should_register_for_all_events_if_topology_requested() { .containsExactly( ProtocolConstants.EventType.SCHEMA_CHANGE, ProtocolConstants.EventType.STATUS_CHANGE, - ProtocolConstants.EventType.TOPOLOGY_CHANGE); + ProtocolConstants.EventType.TOPOLOGY_CHANGE, + ProtocolConstants.EventType.GRACEFUL_DISCONNECT); assertThat(channelOptions.eventCallback).isEqualTo(controlConnection); }); } @@ -67,6 +77,8 @@ public void should_register_for_all_events_if_topology_requested() { @Test public void should_register_for_schema_events_only_if_topology_not_requested() { // Given + when(defaultProfile.getBoolean(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, true)) + .thenReturn(false); DriverChannel channel1 = newMockDriverChannel(1); ArgumentCaptor optionsCaptor = ArgumentCaptor.forClass(DriverChannelOptions.class); @@ -87,6 +99,62 @@ public void should_register_for_schema_events_only_if_topology_not_requested() { }); } + @Test + public void should_not_register_for_graceful_disconnect_when_disabled() { + // Given + when(defaultProfile.getBoolean(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, true)) + .thenReturn(false); + DriverChannel channel1 = newMockDriverChannel(1); + ArgumentCaptor optionsCaptor = + ArgumentCaptor.forClass(DriverChannelOptions.class); + when(channelFactory.connect(eq(node1), optionsCaptor.capture())) + .thenReturn(CompletableFuture.completedFuture(channel1)); + + // When + controlConnection.init(true, false, false); + + // Then + await() + .untilAsserted( + () -> { + DriverChannelOptions channelOptions = optionsCaptor.getValue(); + assertThat(channelOptions.eventTypes) + .containsExactly( + ProtocolConstants.EventType.SCHEMA_CHANGE, + ProtocolConstants.EventType.STATUS_CHANGE, + ProtocolConstants.EventType.TOPOLOGY_CHANGE); + assertThat(channelOptions.eventCallback).isEqualTo(controlConnection); + }); + } + + @Test + public void should_process_graceful_disconnect_event() { + // Given + when(defaultProfile.getBoolean(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, true)) + .thenReturn(true); + DriverChannel channel1 = newMockDriverChannel(1); + Metadata metadata = mock(Metadata.class); + when(metadataManager.getMetadata()).thenReturn(metadata); + when(metadata.findNode(channel1.getEndPoint())).thenReturn(Optional.of(node1)); + ArgumentCaptor optionsCaptor = + ArgumentCaptor.forClass(DriverChannelOptions.class); + when(channelFactory.connect(eq(node1), optionsCaptor.capture())) + .thenReturn(CompletableFuture.completedFuture(channel1)); + controlConnection.init(true, false, false); + await().until(() -> optionsCaptor.getValue() != null); + EventCallback callback = optionsCaptor.getValue().eventCallback; + com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent event = + new com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent(); + + // When + callback.onEvent(event); + + // Then + verify(eventBus).fire(org.mockito.ArgumentMatchers.any(GracefulDisconnectEvent.class)); + verify(sessionMetricUpdater).incrementCounter(DefaultSessionMetric.GRACEFUL_DISCONNECTS, null); + verify(nodeMetricUpdater).incrementCounter(DefaultNodeMetric.GRACEFUL_DISCONNECTS, null); + } + @Test public void should_process_status_change_events() { // Given diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionTestBase.java b/core/src/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionTestBase.java index c52199465a8..015767793f9 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionTestBase.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/control/ControlConnectionTestBase.java @@ -42,6 +42,8 @@ import com.datastax.oss.driver.internal.core.metadata.MetadataManager; import com.datastax.oss.driver.internal.core.metadata.TestNodeFactory; import com.datastax.oss.driver.internal.core.metrics.MetricsFactory; +import com.datastax.oss.driver.internal.core.metrics.NodeMetricUpdater; +import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater; import io.netty.channel.Channel; import io.netty.channel.DefaultChannelPromise; import io.netty.channel.DefaultEventLoopGroup; @@ -77,6 +79,8 @@ abstract class ControlConnectionTestBase { @Mock protected LoadBalancingPolicyWrapper loadBalancingPolicyWrapper; @Mock protected MetadataManager metadataManager; @Mock protected MetricsFactory metricsFactory; + @Mock protected SessionMetricUpdater sessionMetricUpdater; + @Mock protected NodeMetricUpdater nodeMetricUpdater; protected DefaultNode node1; protected DefaultNode node2; @@ -118,6 +122,8 @@ public void setup() { when(context.getLoadBalancingPolicyWrapper()).thenReturn(loadBalancingPolicyWrapper); when(context.getMetricsFactory()).thenReturn(metricsFactory); + when(metricsFactory.getSessionUpdater()).thenReturn(sessionMetricUpdater); + when(metricsFactory.newNodeUpdater(any(Node.class))).thenReturn(nodeMetricUpdater); node1 = TestNodeFactory.newNode(1, context); node2 = TestNodeFactory.newNode(2, context); mockQueryPlan(node1, node2); diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/pool/ChannelPoolGracefulDisconnectTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/pool/ChannelPoolGracefulDisconnectTest.java new file mode 100644 index 00000000000..9d8eec3888b --- /dev/null +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/pool/ChannelPoolGracefulDisconnectTest.java @@ -0,0 +1,140 @@ +/* + * 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.pool; + +import static com.datastax.oss.driver.Assertions.assertThat; +import static com.datastax.oss.driver.Assertions.assertThatStage; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.datastax.oss.driver.api.core.config.DefaultDriverOption; +import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance; +import com.datastax.oss.driver.api.core.metrics.DefaultNodeMetric; +import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric; +import com.datastax.oss.driver.internal.core.channel.DriverChannel; +import com.datastax.oss.driver.internal.core.channel.DriverChannelOptions; +import com.datastax.oss.driver.internal.core.channel.MockChannelFactoryHelper; +import com.datastax.oss.driver.internal.core.metadata.DefaultNode; +import com.datastax.oss.driver.internal.core.metadata.GracefulDisconnectEvent; +import com.datastax.oss.driver.internal.core.metadata.TestNodeFactory; +import com.datastax.oss.protocol.internal.ProtocolConstants; +import java.util.concurrent.CompletionStage; +import org.junit.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.ArgumentMatchers; + +public class ChannelPoolGracefulDisconnectTest extends ChannelPoolTestBase { + + private ChannelPool initPool(boolean gracefulDisconnectEnabled, DriverChannel... channels) + throws Exception { + when(defaultProfile.getBoolean(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, true)) + .thenReturn(gracefulDisconnectEnabled); + when(defaultProfile.getInt(DefaultDriverOption.CONNECTION_POOL_LOCAL_SIZE)) + .thenReturn(channels.length); + + MockChannelFactoryHelper.Builder factoryHelperBuilder = + MockChannelFactoryHelper.builder(channelFactory); + for (DriverChannel channel : channels) { + factoryHelperBuilder.success(node, channel); + } + MockChannelFactoryHelper factoryHelper = factoryHelperBuilder.build(); + + CompletionStage poolFuture = + ChannelPool.init(node, null, NodeDistance.LOCAL, context, "test"); + factoryHelper.waitForCalls(node, channels.length); + assertThatStage(poolFuture).isSuccess(); + return poolFuture.toCompletableFuture().get(); + } + + @Test + public void should_request_graceful_disconnect_events_when_enabled() throws Exception { + DriverChannel channel1 = newMockDriverChannel(1); + initPool(true, channel1); + + ArgumentCaptor optionsCaptor = + ArgumentCaptor.forClass(DriverChannelOptions.class); + verify(channelFactory).connect(eq(node), optionsCaptor.capture()); + assertThat(optionsCaptor.getValue().eventTypes) + .containsExactly(ProtocolConstants.EventType.GRACEFUL_DISCONNECT); + assertThat(optionsCaptor.getValue().eventCallback).isNotNull(); + } + + @Test + public void should_not_request_graceful_disconnect_events_when_disabled() throws Exception { + DriverChannel channel1 = newMockDriverChannel(1); + initPool(false, channel1); + + ArgumentCaptor optionsCaptor = + ArgumentCaptor.forClass(DriverChannelOptions.class); + verify(channelFactory).connect(eq(node), optionsCaptor.capture()); + assertThat(optionsCaptor.getValue().eventTypes).isEmpty(); + } + + @Test + public void should_close_all_channels_when_graceful_disconnect_event_for_node() throws Exception { + DriverChannel channel1 = newMockDriverChannel(1); + DriverChannel channel2 = newMockDriverChannel(2); + initPool(true, channel1, channel2); + + // As fired by the control connection when it receives the event for this node: + eventBus.fire(new GracefulDisconnectEvent(node)); + + verify(channel1, VERIFY_TIMEOUT).close(); + verify(channel2, VERIFY_TIMEOUT).close(); + } + + @Test + public void should_ignore_graceful_disconnect_event_for_other_node() throws Exception { + DriverChannel channel1 = newMockDriverChannel(1); + initPool(true, channel1); + + DefaultNode otherNode = TestNodeFactory.newNode(2, context); + eventBus.fire(new GracefulDisconnectEvent(otherNode)); + + // Wait for the event to be processed on the admin executor, then check nothing was closed: + verify(eventBus, VERIFY_TIMEOUT).fire(ArgumentMatchers.any(GracefulDisconnectEvent.class)); + Thread.sleep(200); + verify(channel1, never()).close(); + } + + @Test + public void should_drain_and_increment_metrics_when_event_received_on_query_connection() + throws Exception { + DriverChannel channel1 = newMockDriverChannel(1); + initPool(true, channel1); + + ArgumentCaptor optionsCaptor = + ArgumentCaptor.forClass(DriverChannelOptions.class); + verify(channelFactory).connect(eq(node), optionsCaptor.capture()); + + // Simulate the server sending GRACEFUL_DISCONNECT on the pooled connection: + optionsCaptor + .getValue() + .eventCallback + .onEvent(new com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent()); + + verify(nodeMetricUpdater, VERIFY_TIMEOUT) + .incrementCounter(DefaultNodeMetric.GRACEFUL_DISCONNECTS, null); + verify(sessionMetricUpdater, VERIFY_TIMEOUT) + .incrementCounter(DefaultSessionMetric.GRACEFUL_DISCONNECTS, null); + verify(eventBus, VERIFY_TIMEOUT).fire(ArgumentMatchers.any(GracefulDisconnectEvent.class)); + verify(channel1, VERIFY_TIMEOUT).close(); + } +} diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/pool/ChannelPoolTestBase.java b/core/src/test/java/com/datastax/oss/driver/internal/core/pool/ChannelPoolTestBase.java index 2f8056e49e0..1e6734533b5 100644 --- a/core/src/test/java/com/datastax/oss/driver/internal/core/pool/ChannelPoolTestBase.java +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/pool/ChannelPoolTestBase.java @@ -37,6 +37,7 @@ import com.datastax.oss.driver.internal.core.metadata.TestNodeFactory; import com.datastax.oss.driver.internal.core.metrics.MetricsFactory; import com.datastax.oss.driver.internal.core.metrics.NodeMetricUpdater; +import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater; import io.netty.channel.Channel; import io.netty.channel.DefaultChannelPromise; import io.netty.channel.DefaultEventLoopGroup; @@ -63,6 +64,7 @@ abstract class ChannelPoolTestBase { @Mock protected ChannelFactory channelFactory; @Mock protected MetricsFactory metricsFactory; @Mock protected NodeMetricUpdater nodeMetricUpdater; + @Mock protected SessionMetricUpdater sessionMetricUpdater; protected DefaultNode node; protected EventBus eventBus; private DefaultEventLoopGroup adminEventLoopGroup; @@ -89,6 +91,7 @@ public void setup() { when(context.getMetricsFactory()).thenReturn(metricsFactory); when(metricsFactory.newNodeUpdater(any(Node.class))).thenReturn(nodeMetricUpdater); + when(metricsFactory.getSessionUpdater()).thenReturn(sessionMetricUpdater); node = TestNodeFactory.newNode(1, context); } diff --git a/core/src/test/java/com/datastax/oss/driver/internal/core/protocol/GracefulDisconnectWireCompatTest.java b/core/src/test/java/com/datastax/oss/driver/internal/core/protocol/GracefulDisconnectWireCompatTest.java new file mode 100644 index 00000000000..fa9effe0869 --- /dev/null +++ b/core/src/test/java/com/datastax/oss/driver/internal/core/protocol/GracefulDisconnectWireCompatTest.java @@ -0,0 +1,149 @@ +/* + * 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.protocol; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.catchThrowable; + +import com.datastax.oss.protocol.internal.Compressor; +import com.datastax.oss.protocol.internal.Frame; +import com.datastax.oss.protocol.internal.FrameCodec; +import com.datastax.oss.protocol.internal.ProtocolConstants; +import com.datastax.oss.protocol.internal.response.event.GracefulDisconnectEvent; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.UnpooledByteBufAllocator; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import org.junit.Test; + +/** + * Wire-level compatibility tests for the CEP-59 {@code GRACEFUL_DISCONNECT} event + * (CASSANDRA-21191). + * + *

The server-side implementation is still in flux, so these tests exercise the full decode path + * on raw bytes, as the server would send them, rather than on pre-built message objects. They pin + * down: + * + *

+ */ +public class GracefulDisconnectWireCompatTest { + + private final FrameCodec frameCodec = + FrameCodec.defaultClient( + new ByteBufPrimitiveCodec(UnpooledByteBufAllocator.DEFAULT), Compressor.none()); + + /** Builds a raw response envelope: version | flags | streamId | opcode | length | body. */ + private ByteBuf rawEventFrame(int protocolVersion, byte[] body) { + ByteBuf buffer = UnpooledByteBufAllocator.DEFAULT.buffer(); + buffer.writeByte(protocolVersion | 0x80); // response direction bit + buffer.writeByte(0); // flags + buffer.writeShort(-1); // stream id: events always use -1 + buffer.writeByte(ProtocolConstants.Opcode.EVENT); + buffer.writeInt(body.length); + buffer.writeBytes(body); + return buffer; + } + + private static byte[] eventBody(String eventType, byte[] extra) { + byte[] typeBytes = eventType.getBytes(StandardCharsets.UTF_8); + byte[] body = new byte[2 + typeBytes.length + extra.length]; + body[0] = (byte) (typeBytes.length >> 8); + body[1] = (byte) typeBytes.length; + System.arraycopy(typeBytes, 0, body, 2, typeBytes.length); + System.arraycopy(extra, 0, body, 2 + typeBytes.length, extra.length); + return body; + } + + @Test + public void should_decode_current_server_format() { + // Body is exactly the type string, as sent by the CASSANDRA-21191 baseline: + ByteBuf raw = + rawEventFrame(ProtocolConstants.Version.V4, eventBody("GRACEFUL_DISCONNECT", new byte[0])); + + Frame frame = frameCodec.decode(raw); + + assertThat(frame.streamId).isEqualTo(-1); + assertThat(frame.message).isInstanceOf(GracefulDisconnectEvent.class); + } + + @Test + public void should_decode_v5_envelope() { + ByteBuf raw = + rawEventFrame(ProtocolConstants.Version.V5, eventBody("GRACEFUL_DISCONNECT", new byte[0])); + + Frame frame = frameCodec.decode(raw); + + assertThat(frame.message).isInstanceOf(GracefulDisconnectEvent.class); + } + + @Test + public void should_tolerate_extra_body_bytes_from_future_server() { + // A plausible CEP-59 evolution: the server appends the grace period (an [int], here 5000ms) + // to the event body. An old driver must keep working, ignoring the extra payload. + byte[] extra = {0x00, 0x00, 0x13, (byte) 0x88}; + ByteBuf raw = + rawEventFrame(ProtocolConstants.Version.V4, eventBody("GRACEFUL_DISCONNECT", extra)); + + Frame frame = frameCodec.decode(raw); + + assertThat(frame.message).isInstanceOf(GracefulDisconnectEvent.class); + } + + @Test + public void should_round_trip_encoded_event() { + Frame outgoing = + Frame.forResponse( + ProtocolConstants.Version.V4, + -1, + null, + Collections.emptyMap(), + Collections.emptyList(), + new GracefulDisconnectEvent()); + FrameCodec serverCodec = + FrameCodec.defaultServer( + new ByteBufPrimitiveCodec(UnpooledByteBufAllocator.DEFAULT), Compressor.none()); + + Frame decoded = frameCodec.decode(serverCodec.encode(outgoing)); + + assertThat(decoded.message).isInstanceOf(GracefulDisconnectEvent.class); + } + + @Test + public void should_reject_unknown_event_type() { + // Documents the current failure mode if the server ever renames the event or pushes a type + // the driver does not know: decoding fails. This is why ProtocolInitHandler must only + // REGISTER for event types the server advertised (see + // ProtocolInitHandlerGracefulDisconnectTest) — a server never pushes events that were not + // registered. + ByteBuf raw = + rawEventFrame( + ProtocolConstants.Version.V4, eventBody("GRACEFUL_DISCONNECT_V2", new byte[0])); + + Throwable t = catchThrowable(() -> frameCodec.decode(raw)); + + assertThat(t) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Unsupported event type"); + } +} diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/connection/GracefulDisconnectIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/connection/GracefulDisconnectIT.java new file mode 100644 index 00000000000..98380e29072 --- /dev/null +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/connection/GracefulDisconnectIT.java @@ -0,0 +1,134 @@ +/* + * 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.connection; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +import com.codahale.metrics.Counter; +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.core.metrics.DefaultSessionMetric; +import com.datastax.oss.driver.api.testinfra.ccm.CustomCcmRule; +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.SessionUtils; +import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; +import org.junit.ClassRule; +import org.junit.Test; + +/** + * Exercises CEP-59 graceful disconnect (CASSANDRA-21191) against a real cluster: when a node is + * drained, it sends a GRACEFUL_DISCONNECT event on every registered connection before closing the + * transport, and the driver must drain its pool to that node and fail over without surfacing any + * exception to the application. + * + *

Requires a server that implements the GRACEFUL_DISCONNECT event; on older servers the test is + * skipped by the version requirement below. + */ +public class GracefulDisconnectIT { + + @ClassRule + public static final CustomCcmRule CCM_RULE = CustomCcmRule.builder().withNodes(2).build(); + + private static final String QUERY = "SELECT * FROM system.local"; + + @BackendRequirement( + type = BackendType.CASSANDRA, + minInclusive = "7.0", + description = "Graceful disconnect (CEP-59 / CASSANDRA-21191) requires server-side support") + @Test + public void should_fail_over_without_disruption_when_node_drains() throws Exception { + DriverConfigLoader loader = + SessionUtils.configLoaderBuilder() + .withStringList( + DefaultDriverOption.METRICS_SESSION_ENABLED, + ImmutableList.of(DefaultSessionMetric.GRACEFUL_DISCONNECTS.getPath())) + .build(); + + try (CqlSession session = SessionUtils.newSession(CCM_RULE, loader)) { + + // Sanity check before the drain: + session.execute(QUERY); + + // Steady query load for the whole duration of the test, collecting any exception that + // reaches the application: + AtomicLong successes = new AtomicLong(); + List failures = new CopyOnWriteArrayList<>(); + AtomicBoolean stopped = new AtomicBoolean(); + Thread load = + new Thread( + () -> { + while (!stopped.get()) { + try { + session.execute(QUERY); + successes.incrementAndGet(); + } catch (RuntimeException e) { + failures.add(e); + } + } + }, + "graceful-disconnect-load"); + load.start(); + + try { + // Drain node 2: the server stops accepting new requests and sends GRACEFUL_DISCONNECT on + // every connection registered for it, then closes the transport. + CCM_RULE.getCcmBridge().nodetool(2, "drain"); + + // The driver must have observed the event (this is also the end-to-end check for the + // session-level metric): + Counter gracefulDisconnects = + (Counter) + session + .getMetrics() + .orElseThrow(() -> new AssertionError("expected metrics to be enabled")) + .getSessionMetric(DefaultSessionMetric.GRACEFUL_DISCONNECTS) + .orElseThrow( + () -> new AssertionError("expected graceful-disconnects metric to exist")); + await() + .atMost(30, TimeUnit.SECONDS) + .untilAsserted(() -> assertThat(gracefulDisconnects.getCount()).isGreaterThan(0)); + + // Queries must keep succeeding after the drain (load fails over to the other node): + long successesAfterEvent = successes.get(); + await() + .atMost(30, TimeUnit.SECONDS) + .until(() -> successes.get() > successesAfterEvent + 100); + } finally { + stopped.set(true); + load.join(TimeUnit.SECONDS.toMillis(10)); + if (load.isAlive()) { + load.interrupt(); + load.join(TimeUnit.SECONDS.toMillis(5)); + } + } + assertThat(load.isAlive()).as("load thread should have terminated").isFalse(); + + // The whole point of graceful disconnect: the shutdown must be invisible to the + // application, no request may fail. + assertThat(failures).isEmpty(); + } + } +} diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/core/metrics/DropwizardMetricsIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/core/metrics/DropwizardMetricsIT.java index e0184516e21..0bbd85d0b4c 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/core/metrics/DropwizardMetricsIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/core/metrics/DropwizardMetricsIT.java @@ -111,6 +111,7 @@ protected void assertMetricsPresent(CqlSession session) { break; case CQL_CLIENT_TIMEOUTS: case THROTTLING_ERRORS: + case GRACEFUL_DISCONNECTS: assertThat(m).isInstanceOf(Counter.class); assertThat(((Counter) m).getCount()).isZero(); break; @@ -170,6 +171,7 @@ protected void assertMetricsPresent(CqlSession session) { case SPECULATIVE_EXECUTIONS: case CONNECTION_INIT_ERRORS: case AUTHENTICATION_ERRORS: + case GRACEFUL_DISCONNECTS: assertThat(m).isInstanceOf(Counter.class); assertThat(((Counter) m).getCount()).isZero(); break; diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/metrics/micrometer/MicrometerMetricsIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/metrics/micrometer/MicrometerMetricsIT.java index c38df1e2026..1111d250bf7 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/metrics/micrometer/MicrometerMetricsIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/metrics/micrometer/MicrometerMetricsIT.java @@ -102,6 +102,7 @@ protected void assertMetricsPresent(CqlSession session) { break; case CQL_CLIENT_TIMEOUTS: case THROTTLING_ERRORS: + case GRACEFUL_DISCONNECTS: assertThat(m).isInstanceOf(Counter.class); assertThat(((Counter) m).count()).isZero(); break; @@ -154,6 +155,7 @@ protected void assertMetricsPresent(CqlSession session) { case SPECULATIVE_EXECUTIONS: case CONNECTION_INIT_ERRORS: case AUTHENTICATION_ERRORS: + case GRACEFUL_DISCONNECTS: assertThat(m).isInstanceOf(Counter.class); assertThat(((Counter) m).count()).isZero(); break; diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/metrics/microprofile/MicroProfileMetricsIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/metrics/microprofile/MicroProfileMetricsIT.java index aa04c058a49..ce2fddd6baa 100644 --- a/integration-tests/src/test/java/com/datastax/oss/driver/metrics/microprofile/MicroProfileMetricsIT.java +++ b/integration-tests/src/test/java/com/datastax/oss/driver/metrics/microprofile/MicroProfileMetricsIT.java @@ -106,6 +106,7 @@ protected void assertMetricsPresent(CqlSession session) { assertThat(((Meter) m).getCount()).isGreaterThan(0); break; case CQL_CLIENT_TIMEOUTS: + case GRACEFUL_DISCONNECTS: case THROTTLING_ERRORS: assertThat(m).isInstanceOf(Counter.class); assertThat(((Counter) m).getCount()).isZero(); @@ -160,6 +161,7 @@ protected void assertMetricsPresent(CqlSession session) { case SPECULATIVE_EXECUTIONS: case CONNECTION_INIT_ERRORS: case AUTHENTICATION_ERRORS: + case GRACEFUL_DISCONNECTS: assertThat(m).isInstanceOf(Counter.class); assertThat(((Counter) m).getCount()).isZero(); break; diff --git a/metrics/micrometer/src/main/java/com/datastax/oss/driver/internal/metrics/micrometer/MicrometerNodeMetricUpdater.java b/metrics/micrometer/src/main/java/com/datastax/oss/driver/internal/metrics/micrometer/MicrometerNodeMetricUpdater.java index cb8303de965..46d5d77fc24 100644 --- a/metrics/micrometer/src/main/java/com/datastax/oss/driver/internal/metrics/micrometer/MicrometerNodeMetricUpdater.java +++ b/metrics/micrometer/src/main/java/com/datastax/oss/driver/internal/metrics/micrometer/MicrometerNodeMetricUpdater.java @@ -75,6 +75,7 @@ public MicrometerNodeMetricUpdater( initializeCounter(DefaultNodeMetric.SPECULATIVE_EXECUTIONS, profile); initializeCounter(DefaultNodeMetric.CONNECTION_INIT_ERRORS, profile); initializeCounter(DefaultNodeMetric.AUTHENTICATION_ERRORS, profile); + initializeCounter(DefaultNodeMetric.GRACEFUL_DISCONNECTS, profile); initializeTimer(DefaultNodeMetric.CQL_MESSAGES, profile); initializeTimer(DseNodeMetric.GRAPH_MESSAGES, profile); diff --git a/metrics/micrometer/src/main/java/com/datastax/oss/driver/internal/metrics/micrometer/MicrometerSessionMetricUpdater.java b/metrics/micrometer/src/main/java/com/datastax/oss/driver/internal/metrics/micrometer/MicrometerSessionMetricUpdater.java index 559054ab510..f4a302e881d 100644 --- a/metrics/micrometer/src/main/java/com/datastax/oss/driver/internal/metrics/micrometer/MicrometerSessionMetricUpdater.java +++ b/metrics/micrometer/src/main/java/com/datastax/oss/driver/internal/metrics/micrometer/MicrometerSessionMetricUpdater.java @@ -49,6 +49,7 @@ public MicrometerSessionMetricUpdater( initializeCounter(DefaultSessionMetric.CQL_CLIENT_TIMEOUTS, profile); initializeCounter(DefaultSessionMetric.THROTTLING_ERRORS, profile); + initializeCounter(DefaultSessionMetric.GRACEFUL_DISCONNECTS, profile); initializeCounter(DseSessionMetric.GRAPH_CLIENT_TIMEOUTS, profile); initializeTimer(DefaultSessionMetric.CQL_REQUESTS, profile); diff --git a/metrics/microprofile/src/main/java/com/datastax/oss/driver/internal/metrics/microprofile/MicroProfileNodeMetricUpdater.java b/metrics/microprofile/src/main/java/com/datastax/oss/driver/internal/metrics/microprofile/MicroProfileNodeMetricUpdater.java index 8a2d235b59e..d10940e554e 100644 --- a/metrics/microprofile/src/main/java/com/datastax/oss/driver/internal/metrics/microprofile/MicroProfileNodeMetricUpdater.java +++ b/metrics/microprofile/src/main/java/com/datastax/oss/driver/internal/metrics/microprofile/MicroProfileNodeMetricUpdater.java @@ -71,6 +71,7 @@ public MicroProfileNodeMetricUpdater( initializeCounter(DefaultNodeMetric.SPECULATIVE_EXECUTIONS, profile); initializeCounter(DefaultNodeMetric.CONNECTION_INIT_ERRORS, profile); initializeCounter(DefaultNodeMetric.AUTHENTICATION_ERRORS, profile); + initializeCounter(DefaultNodeMetric.GRACEFUL_DISCONNECTS, profile); initializeTimer(DefaultNodeMetric.CQL_MESSAGES, profile); initializeTimer(DseNodeMetric.GRAPH_MESSAGES, profile); diff --git a/metrics/microprofile/src/main/java/com/datastax/oss/driver/internal/metrics/microprofile/MicroProfileSessionMetricUpdater.java b/metrics/microprofile/src/main/java/com/datastax/oss/driver/internal/metrics/microprofile/MicroProfileSessionMetricUpdater.java index f3c906e4422..a06c708fc23 100644 --- a/metrics/microprofile/src/main/java/com/datastax/oss/driver/internal/metrics/microprofile/MicroProfileSessionMetricUpdater.java +++ b/metrics/microprofile/src/main/java/com/datastax/oss/driver/internal/metrics/microprofile/MicroProfileSessionMetricUpdater.java @@ -45,6 +45,7 @@ public MicroProfileSessionMetricUpdater( initializeCounter(DefaultSessionMetric.CQL_CLIENT_TIMEOUTS, profile); initializeCounter(DefaultSessionMetric.THROTTLING_ERRORS, profile); + initializeCounter(DefaultSessionMetric.GRACEFUL_DISCONNECTS, profile); initializeCounter(DseSessionMetric.GRAPH_CLIENT_TIMEOUTS, profile); initializeTimer(DefaultSessionMetric.CQL_REQUESTS, profile);