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:
+ *
+ *
+ * - the current contract: an EVENT envelope whose body is just the type string;
+ *
- forward compatibility: a future server that appends a payload (e.g. the grace period) to
+ * the event body must not break decoding;
+ *
- the failure mode for an unknown event type, which is why the driver must never REGISTER for
+ * types the server did not advertise.
+ *
+ */
+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);