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>

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

@@ -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

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 {

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<String, List<String>> supportedOptions) {
+    if (supportedOptions == null) {
+      return false;
+    }
+    List<String> 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<String> 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<String, List<String>> 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,15 +337,15 @@ 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 {
               setConnectSuccess();
             }
           }
         } else if (step == Step.SET_KEYSPACE && response instanceof SetKeyspace) {
-          if (!options.eventTypes.isEmpty()) {
+          if (!eventTypes.isEmpty()) {
             step = Step.REGISTER;
             send();
           } else {

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<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 +647,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(ProtocolConstants.EventType.GRACEFUL_DISCONNECT);
+    }
     return builder.build();
   }
 }

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 + ")";
+  }
+}

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,

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(