CASSJAVA-124: Add GRACEFUL_DISCONNECT support (CEP-59)

bom/pom.xml

@@ -78,7 +78,7 @@
       <dependency>
         <groupId>com.datastax.oss</groupId>
         <artifactId>native-protocol</artifactId>
-        <version>1.5.2</version>
+        <version>1.5.3-SNAPSHOT</version>
       </dependency>
     </dependencies>
   </dependencyManagement>

ci/run-tests.sh

@@ -4,7 +4,11 @@
 . ~/env.txt
 cd $(dirname "$(readlink -f "$0")")/..
 printenv | sort
-mvn -B -V install -DskipTests -Dmaven.javadoc.skip=true
+# Install snapshot dependencies (e.g. native-protocol) that are not published to a public repo yet.
+# Fail the build immediately if this does not succeed, and pass -U below so that any stale
+# resolution-failure markers in the local repository do not block the freshly installed snapshot.
+./install-snapshots.sh || exit 1
+mvn -U -B -V install -DskipTests -Dmaven.javadoc.skip=true
 jabba use ${TEST_JAVA_VERSION}
 # Find out the latest patch version of Cassandra
 PATCH_SERVER_VERSION=$(curl -s https://downloads.apache.org/cassandra/ | grep -oP '(?<=href=\")[0-9]+\.[0-9]+\.[0-9]+(?=)' | sort -rV | uniq -w 3 | grep $SERVER_VERSION)

core/src/main/java/com/datastax/oss/driver/api/core/config/DefaultDriverOption.java

@@ -1041,7 +1041,15 @@ public enum DefaultDriverOption implements DriverOption {
    *
    * <p>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.
+   *
+   * <p>Value-type: boolean
+   */
+  GRACEFUL_DISCONNECT_ENABLED("advanced.connection.graceful-disconnect-enabled");
 
   private final String path;
 

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));

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<Boolean> 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.

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<String, DefaultNodeMetric> BY_PATH = sortByPath();

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<String, DefaultSessionMetric> BY_PATH = sortByPath();

core/src/main/java/com/datastax/oss/driver/internal/core/channel/ChannelFactory.java

@@ -339,6 +339,12 @@ protected void initChannel(Channel channel) {
                 options.eventCallback,
                 options.ownerLogPrefix);
         HeartbeatHandler heartbeatHandler = new HeartbeatHandler(defaultConfig);
+        // Always query OPTIONS on the first channel (to discover the product type), and on any
+        // channel that intends to register for GRACEFUL_DISCONNECT: capabilities can differ from
+        // node to node (e.g. mixed-version clusters), so each channel must filter its REGISTER
+        // against its own SUPPORTED response.
+        boolean querySupportedOptions =
+            productType == null || options.eventTypes.contains(GracefulDisconnectEvent.EVENT_TYPE);
         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

core/src/main/java/com/datastax/oss/driver/internal/core/channel/GracefulDisconnectEvent.java

@@ -0,0 +1,51 @@
+/*
+ * 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 com.datastax.oss.driver.api.core.metadata.Node;
+import net.jcip.annotations.Immutable;
+
+/**
+ * This event indicates that the server is shutting down gracefully and the driver should:
+ *
+ * <ul>
+ *   <li>Stop sending new requests on the affected connection
+ *   <li>Allow in-flight requests to complete
+ *   <li>Begin reconnection attempts with exponential backoff
+ * </ul>
+ *
+ * <p>This is part of CEP-59: Graceful Disconnect – In-Band Connection Draining for Node Shutdown.
+ */
+@Immutable
+public class GracefulDisconnectEvent {
+
+  /** The event type string as defined in the native protocol. */
+  public static final String EVENT_TYPE = "GRACEFUL_DISCONNECT";
+
+  /** The node that sent the graceful disconnect event. */
+  public final Node node;
+
+  public GracefulDisconnectEvent(Node node) {
+    this.node = node;
+  }
+
+  @Override
+  public String toString() {
+    return "GracefulDisconnectEvent{node=" + node + '}';
+  }
+}

core/src/main/java/com/datastax/oss/driver/internal/core/channel/InFlightHandler.java

@@ -33,6 +33,7 @@
 import com.datastax.oss.protocol.internal.Frame;
 import com.datastax.oss.protocol.internal.Message;
 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 +219,13 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
 
     if (streamId < 0) {
       Message event = responseFrame.message;
+      if (event instanceof Event
+          && GracefulDisconnectEvent.EVENT_TYPE.equals(((Event) event).type)) {
+        // Start draining this channel first, so that the drain is not compromised if the
+        // callback below misbehaves.
+        LOG.debug("[{}] Received GRACEFUL_DISCONNECT, initiating graceful drain", logPrefix);
+        startGracefulShutdown(ctx);
+      }
       if (eventCallback == null) {
         LOG.debug("[{}] Received event {} but no callback was registered", logPrefix, event);
       } else {

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,30 @@ protected boolean setConnectSuccess() {
     return result;
   }
 
+  /**
+   * Whether a SUPPORTED response advertises the CEP-59 graceful disconnect capability.
+   *
+   * <p>The server signals support with a {@code GRACEFUL_DISCONNECT} key; depending on the code
+   * path it may also send the key with an explicit {@code "false"} value when the feature is
+   * disabled, so any value other than {@code false} (in any case) is treated as supported.
+   */
+  @VisibleForTesting
+  static boolean supportsGracefulDisconnect(Map<String, List<String>> supportedOptions) {
+    if (supportedOptions == null) {
+      return false;
+    }
+    List<String> values = supportedOptions.get(GracefulDisconnectEvent.EVENT_TYPE);
+    if (values == null) {
+      return false;
+    }
+    for (String value : values) {
+      if ("false".equalsIgnoreCase(value)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   private enum Step {
     OPTIONS,
     STARTUP,
@@ -157,6 +184,8 @@ private class InitRequest extends ChannelHandlerRequest {
     private Message request;
     private Authenticator authenticator;
     private ByteBuffer authResponseToken;
+    private List<String> lastRegisterEventTypes;
+    private boolean retriedRegisterWithoutGracefulDisconnect;
 
     InitRequest(ChannelHandlerContext ctx) {
       super(ctx, timeoutMillis);
@@ -183,12 +212,36 @@ Message getRequest() {
         case AUTH_RESPONSE:
           return request = new AuthResponse(authResponseToken);
         case REGISTER:
-          return request = new Register(options.eventTypes);
+          return request = new Register(lastRegisterEventTypes = filterSupportedEventTypes());
         default:
           throw new AssertionError("unhandled step: " + step);
       }
     }
 
+    /**
+     * Filters the requested event types to only include those supported by the server.
+     *
+     * <p>Specifically, GRACEFUL_DISCONNECT is only included if this channel's own SUPPORTED
+     * response advertises the capability. Capabilities are tracked per connection: in a
+     * mixed-version cluster (e.g. during a rolling upgrade), some nodes may support graceful
+     * disconnect while others don't, so a channel never relies on what another channel negotiated.
+     * Channels that request this event type always run the OPTIONS step (see {@link
+     * ChannelFactory}), so the channel attribute is populated by the time REGISTER is sent.
+     */
+    private List<String> filterSupportedEventTypes() {
+      List<String> filteredEventTypes = new ArrayList<>(options.eventTypes);
+
+      if (filteredEventTypes.contains(GracefulDisconnectEvent.EVENT_TYPE)) {
+        Map<String, List<String>> supportedOptions = channel.attr(DriverChannel.OPTIONS_KEY).get();
+        if (!supportsGracefulDisconnect(supportedOptions)
+            || retriedRegisterWithoutGracefulDisconnect) {
+          filteredEventTypes.remove(GracefulDisconnectEvent.EVENT_TYPE);
+        }
+      }
+
+      return filteredEventTypes;
+    }
+
     @Override
     void send() {
       stepNumber++;
@@ -339,6 +392,23 @@ void onResponse(Message response) {
           } else if (step == Step.SET_KEYSPACE
               && error.code == ProtocolConstants.ErrorCode.INVALID) {
             fail(new InvalidKeyspaceException(error.message));
+          } else if (step == Step.REGISTER
+              && !retriedRegisterWithoutGracefulDisconnect
+              && lastRegisterEventTypes != null
+              && lastRegisterEventTypes.contains(GracefulDisconnectEvent.EVENT_TYPE)
+              && (serverOrProtocolError || error.code == ProtocolConstants.ErrorCode.INVALID)) {
+            // The server rejected our REGISTER, most likely because it does not recognize the
+            // GRACEFUL_DISCONNECT event type: a pre-CEP-59 node in a mixed-version cluster, or a
+            // node where the (still evolving) CEP-59 wire contract has changed. Losing graceful
+            // disconnect on this connection is benign; failing channel init is disruptive. So
+            // retry once without the event type instead of failing.
+            LOG.warn(
+                "[{}] Server rejected REGISTER including {} ({}), retrying without it",
+                logPrefix,
+                GracefulDisconnectEvent.EVENT_TYPE,
+                error.message);
+            retriedRegisterWithoutGracefulDisconnect = true;
+            send();
           } else {
             failOnUnexpected(error);
           }

core/src/main/java/com/datastax/oss/driver/internal/core/control/ControlConnection.java

@@ -26,10 +26,12 @@
 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.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.channel.GracefulDisconnectEvent;
 import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
 import com.datastax.oss.driver.internal.core.metadata.DefaultTopologyMonitor;
 import com.datastax.oss.driver.internal.core.metadata.DistanceEvent;
@@ -178,8 +180,8 @@ public void onEvent(Message eventMessage) {
     if (!(eventMessage instanceof Event)) {
       LOG.warn("[{}] Unsupported event class: {}", logPrefix, eventMessage.getClass().getName());
     } else {
-      LOG.debug("[{}] Processing incoming event {}", logPrefix, eventMessage);
       Event event = (Event) eventMessage;
+      LOG.debug("[{}] Processing incoming event {}", logPrefix, eventMessage);
       switch (event.type) {
         case ProtocolConstants.EventType.TOPOLOGY_CHANGE:
           processTopologyChange(event);
@@ -190,6 +192,9 @@ public void onEvent(Message eventMessage) {
         case ProtocolConstants.EventType.SCHEMA_CHANGE:
           processSchemaChange(event);
           break;
+        case GracefulDisconnectEvent.EVENT_TYPE:
+          processGracefulDisconnect();
+          break;
         default:
           LOG.warn("[{}] Unsupported event type: {}", logPrefix, event.type);
       }
@@ -242,6 +247,30 @@ 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 -> context.getEventBus().fire(new GracefulDisconnectEvent(node)));
+    }
+    // The control connection will handle reconnection automatically when the channel closes.
+    // The ChannelPool will close all its channels when it receives the GracefulDisconnectEvent,
+    // which will cause the NodeStateManager to set the node to DOWN state and trigger the
+    // LoadBalancingPolicy to remove it from the live set.
+  }
+
   private class SingleThreaded {
     private final InternalDriverContext context;
     private final DriverConfig config;
@@ -292,7 +321,13 @@ private void init(
       }
       initWasCalled = true;
       try {
-        ImmutableList<String> eventTypes = buildEventTypes(listenToClusterEvents);
+        boolean gracefulDisconnectEnabled =
+            context
+                .getConfig()
+                .getDefaultProfile()
+                .getBoolean(DefaultDriverOption.GRACEFUL_DISCONNECT_ENABLED, true);
+        ImmutableList<String> eventTypes =
+            buildEventTypes(listenToClusterEvents, gracefulDisconnectEnabled);
         LOG.debug("[{}] Initializing with event types {}", logPrefix, eventTypes);
         channelOptions =
             DriverChannelOptions.builder()
@@ -606,14 +641,18 @@ private boolean isAuthFailure(Throwable error) {
     return true;
   }
 
-  private static ImmutableList<String> buildEventTypes(boolean listenClusterEvents) {
+  private static ImmutableList<String> buildEventTypes(
+      boolean listenClusterEvents, boolean gracefulDisconnectEnabled) {
     ImmutableList.Builder<String> builder = ImmutableList.builder();
     builder.add(ProtocolConstants.EventType.SCHEMA_CHANGE);
     if (listenClusterEvents) {
       builder
           .add(ProtocolConstants.EventType.STATUS_CHANGE)
           .add(ProtocolConstants.EventType.TOPOLOGY_CHANGE);
     }
+    if (gracefulDisconnectEnabled) {
+      builder.add(GracefulDisconnectEvent.EVENT_TYPE);
+    }
     return builder.build();
   }
 }

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,