CASSPYTHON-16: Add GRACEFUL_DISCONNECT event support (CEP-59)

CHANGELOG.rst

@@ -1,3 +1,11 @@
+Unreleased
+==========
+
+Features
+--------
+* Support CEP-59 graceful disconnect: register for GRACEFUL_DISCONNECT events and drain
+  connections without disrupting in-flight requests when a node shuts down (CASSPYTHON-16)
+
 3.30.1
 ======
 June 19, 2026

cassandra/cluster.py

@@ -758,6 +758,18 @@ def default_retry_policy(self, policy):
     will be an instance of :class:`~cassandra.metrics.Metrics`.
     """
 
+    graceful_disconnect_enabled = True
+    """
+    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: in-flight requests
+    are allowed to complete before the connections are closed, and new
+    requests fail over to other nodes.
+
+    Defaults to :const:`True`. The capability is negotiated per connection;
+    on servers that do not advertise support this option has no effect.
+    """
+
     metrics = None
     """
     An instance of :class:`cassandra.metrics.Metrics` if :attr:`.metrics_enabled` is
@@ -1069,7 +1081,8 @@ def __init__(self,
                  ssl_context=None,
                  endpoint_factory=None,
                  cloud=None,
-                 column_encryption_policy=None):
+                 column_encryption_policy=None,
+                 graceful_disconnect_enabled=True):
         """
         ``executor_threads`` defines the number of threads in a pool for handling asynchronous tasks such as
         establishing connection pools or refreshing metadata.
@@ -1266,6 +1279,7 @@ def __init__(self,
         self.connect_timeout = connect_timeout
         self.prepare_on_all_hosts = prepare_on_all_hosts
         self.reprepare_on_up = reprepare_on_up
+        self.graceful_disconnect_enabled = graceful_disconnect_enabled
 
         self._listeners = set()
         self._listener_lock = Lock()
@@ -1971,6 +1985,26 @@ def signal_connection_failure(self, host, connection_exc, is_host_addition, expe
             self.on_down(host, is_host_addition, expect_host_to_be_down)
         return is_down
 
+    def on_graceful_disconnect(self, host):
+        """
+        Called when a GRACEFUL_DISCONNECT event (CEP-59) is received on any
+        connection to `host`: the node announced that it is shutting down
+        gracefully. All pools for that host are drained so that in-flight
+        requests complete before their connections are closed.
+
+        Intended for internal use only.
+        """
+        log.info("Received GRACEFUL_DISCONNECT for host %s, "
+                 "the node is shutting down gracefully", host)
+        if self.metrics_enabled and self.metrics:
+            self.metrics.on_graceful_disconnect()
+        if host is None:
+            return
+        for session in tuple(self.sessions):
+            pool = session._pools.get(host)
+            if pool:
+                pool.on_graceful_disconnect()
+
     def add_host(self, endpoint, datacenter=None, rack=None, signal=True, refresh_nodes=True):
         """
         Called when adding initial contact points and when the control
@@ -3567,11 +3601,19 @@ def _try_connect(self, host):
         # this object (after a dereferencing a weakref)
         self_weakref = weakref.ref(self, partial(_clear_watcher, weakref.proxy(connection)))
         try:
-            connection.register_watchers({
+            watchers = {
                 "TOPOLOGY_CHANGE": partial(_watch_callback, self_weakref, '_handle_topology_change'),
                 "STATUS_CHANGE": partial(_watch_callback, self_weakref, '_handle_status_change'),
                 "SCHEMA_CHANGE": partial(_watch_callback, self_weakref, '_handle_schema_change')
-            }, register_timeout=self._timeout)
+            }
+            # GRACEFUL_DISCONNECT (CEP-59) support is negotiated per
+            # connection: only register if this connection's SUPPORTED
+            # response advertised it.
+            if (self._cluster.graceful_disconnect_enabled
+                    and connection.supports_graceful_disconnect):
+                watchers["GRACEFUL_DISCONNECT"] = partial(
+                    _watch_callback, self_weakref, '_handle_graceful_disconnect')
+            connection.register_watchers(watchers, register_timeout=self._timeout)
 
             sel_peers = self._get_peers_query(self.PeersQueryType.PEERS, connection)
             sel_local = self._SELECT_LOCAL if self._token_meta_enabled else self._SELECT_LOCAL_NO_TOKENS
@@ -3922,6 +3964,17 @@ def _handle_schema_change(self, event):
         delay = self._delay_for_event_type('schema_change', self._schema_event_refresh_window)
         self._cluster.scheduler.schedule_unique(delay, self.refresh_schema, **event)
 
+    def _handle_graceful_disconnect(self, event):
+        log.info("[control connection] Received GRACEFUL_DISCONNECT event, "
+                 "the server is shutting down gracefully")
+        connection = self._connection
+        if connection:
+            host = self._cluster.metadata.get_host(connection.endpoint)
+            self._cluster.on_graceful_disconnect(host)
+        # The control connection is draining and will be closed by the
+        # server; proactively move it to another host.
+        self.reconnect()
+
     def wait_for_schema_agreement(self, connection=None, preloaded_results=None, wait_time=None):
 
         total_timeout = wait_time if wait_time is not None else self._cluster.max_schema_agreement_wait

cassandra/connection.py

@@ -724,6 +724,16 @@ class Connection(object):
 
     is_defunct = False
     is_closed = False
+
+    # Set to true when a GRACEFUL_DISCONNECT event (CEP-59) was received on
+    # this connection: no new requests are accepted, and the connection is
+    # closed once all in-flight requests have completed.
+    is_draining = False
+
+    # Whether this connection's SUPPORTED response advertised the CEP-59
+    # graceful disconnect capability (negotiated per connection).
+    supports_graceful_disconnect = False
+
     lock = None
     user_type_map = None
 
@@ -1058,17 +1068,53 @@ def get_request_id(self):
 
     def handle_pushed(self, response):
         log.debug("Message pushed from server: %r", response)
+        if getattr(response, 'event_type', None) == 'GRACEFUL_DISCONNECT':
+            # Start draining this connection first, so that the drain is not
+            # compromised if a callback below misbehaves.
+            self.start_graceful_drain()
         for cb in self._push_watchers.get(response.event_type, []):
             try:
                 cb(response.event_args)
             except Exception:
                 log.exception("Pushed event handler errored, ignoring:")
 
+    def start_graceful_drain(self):
+        """
+        Start a graceful drain of this connection (CEP-59): new requests are
+        refused, in-flight requests are allowed to complete, and the
+        connection is closed once the last in-flight request has completed.
+        """
+        with self.lock:
+            if self.is_defunct or self.is_closed or self.is_draining:
+                return
+            self.is_draining = True
+            has_pending = bool(self._requests) or bool(self._continuous_paging_sessions)
+        log.debug("Draining connection (%s) to %s gracefully; pending requests: %s",
+                  id(self), self.endpoint, has_pending)
+        if not has_pending:
+            self.close()
+
+    def _maybe_finish_graceful_drain(self):
+        """
+        Complete a graceful drain by closing the connection once nothing is
+        in flight anymore.
+        """
+        if not self.is_draining or self.is_closed or self.is_defunct:
+            return
+        with self.lock:
+            has_pending = bool(self._requests) or bool(self._continuous_paging_sessions)
+        if not has_pending:
+            log.debug("Graceful drain of connection (%s) to %s complete, closing",
+                      id(self), self.endpoint)
+            self.close()
+
     def send_msg(self, msg, request_id, cb, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=None):
         if self.is_defunct:
             raise ConnectionShutdown("Connection to %s is defunct" % self.endpoint)
         elif self.is_closed:
             raise ConnectionShutdown("Connection to %s is closed" % self.endpoint)
+        elif self.is_draining:
+            raise ConnectionShutdown("Connection to %s is draining (graceful disconnect)" % self.endpoint)
         elif not self._socket_writable:
             raise ConnectionBusy("Connection %s is overloaded" % self.endpoint)
 
@@ -1302,6 +1348,8 @@ def process_msg(self, header, body):
                 with self.lock:
                     self.request_ids.append(stream_id)
 
+        self._maybe_finish_graceful_drain()
+
     def new_continuous_paging_session(self, stream_id, decoder, row_factory, state):
         session = ContinuousPagingSession(stream_id, decoder, row_factory, self, state)
         self._continuous_paging_sessions[stream_id] = session
@@ -1341,6 +1389,8 @@ def _handle_options_response(self, options_response):
         supported_cql_versions = options_response.cql_versions
         remote_supported_compressions = options_response.options['COMPRESSION']
         self._product_type = options_response.options.get('PRODUCT_TYPE', [None])[0]
+        self.supports_graceful_disconnect = \
+            self._supports_graceful_disconnect(options_response.options)
 
         if self.cql_version:
             if self.cql_version not in supported_cql_versions:
@@ -1394,6 +1444,20 @@ def _handle_options_response(self, options_response):
 
         self._send_startup_message(compression_type, no_compact=self.no_compact)
 
+    @staticmethod
+    def _supports_graceful_disconnect(supported_options):
+        """
+        Whether a SUPPORTED response advertises the CEP-59 graceful disconnect
+        capability. The server may send the key with an explicit ``"false"``
+        value when the feature is disabled.
+        """
+        if supported_options is None:
+            return False
+        values = supported_options.get('GRACEFUL_DISCONNECT')
+        if values is None:
+            return False
+        return not any(value.lower() == 'false' for value in values)
+
     @defunct_on_error
     def _send_startup_message(self, compression=None, no_compact=False):
         log.debug("Sending StartupMessage on %s", self)
@@ -1709,6 +1773,11 @@ def run(self):
                     for connection in connections:
                         self._raise_if_stopped()
                         if not (connection.is_defunct or connection.is_closed):
+                            if connection.is_draining:
+                                # A graceful drain (CEP-59) is in progress: do not
+                                # send a heartbeat, the connection will close once
+                                # its in-flight requests have completed.
+                                continue
                             if connection.is_idle:
                                 try:
                                     futures.append(HeartbeatFuture(connection, owner))

cassandra/metrics.py

@@ -96,6 +96,13 @@ class Metrics(object):
     failed request was ignored based on the :class:`.RetryPolicy` decision.
     """
 
+    graceful_disconnects = None
+    """
+    A :class:`greplin.scales.IntStat` count of the number of
+    ``GRACEFUL_DISCONNECT`` events (CEP-59) received from nodes that are
+    shutting down gracefully, across all connections of the cluster.
+    """
+
     known_hosts = None
     """
     A :class:`greplin.scales.IntStat` count of the number of nodes in
@@ -131,6 +138,7 @@ def __init__(self, cluster_proxy):
                                        scales.IntStat('other_errors'),
                                        scales.IntStat('retries'),
                                        scales.IntStat('ignores'),
+                                       scales.IntStat('graceful_disconnects'),
 
                                        # gauges
                                        scales.Stat('known_hosts',
@@ -155,6 +163,7 @@ def __init__(self, cluster_proxy):
         self.other_errors = self.stats.other_errors
         self.retries = self.stats.retries
         self.ignores = self.stats.ignores
+        self.graceful_disconnects = self.stats.graceful_disconnects
         self.known_hosts = self.stats.known_hosts
         self.connected_to = self.stats.connected_to
         self.open_connections = self.stats.open_connections
@@ -180,6 +189,9 @@ def on_ignore(self):
     def on_retry(self):
         self.stats.retries += 1
 
+    def on_graceful_disconnect(self):
+        self.stats.graceful_disconnects += 1
+
     def get_stats(self):
         """
         Returns the metrics for the registered cluster instance.