diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 47b8ee3f6e..271863ab6d 100644 --- a/CHANGELOG.rst +++ b/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 diff --git a/cassandra/cluster.py b/cassandra/cluster.py index f2e111894b..ee0b17611f 100644 --- a/cassandra/cluster.py +++ b/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 diff --git a/cassandra/connection.py b/cassandra/connection.py index d9bca62718..7525c698e3 100644 --- a/cassandra/connection.py +++ b/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)) diff --git a/cassandra/metrics.py b/cassandra/metrics.py index a585000266..a734bf2f8c 100644 --- a/cassandra/metrics.py +++ b/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. diff --git a/cassandra/pool.py b/cassandra/pool.py index d060eb23e4..3abe4421d8 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -409,6 +409,7 @@ def __init__(self, host, host_distance, session): self._keyspace = session.keyspace if self._keyspace: self._connection.set_keyspace_blocking(self._keyspace) + self._maybe_register_graceful_disconnect(self._connection) log.debug("Finished initializing connection for host %s", self.host) def _get_connection(self): @@ -419,6 +420,9 @@ def _get_connection(self): conn = self._connection if not conn: raise NoConnectionsAvailable() + if conn.is_draining: + raise NoConnectionsAvailable( + "Connection to %s is draining (graceful disconnect)" % (self.host,)) return conn def borrow_connection(self, timeout): @@ -503,6 +507,39 @@ def on_orphaned_stream_released(self): with self._stream_available_condition: self._stream_available_condition.notify() + def _maybe_register_graceful_disconnect(self, connection): + """ + Register for GRACEFUL_DISCONNECT events (CEP-59) on a pooled + connection, if the feature is enabled and this connection's server + advertised support for it. + """ + if (self._session.cluster.graceful_disconnect_enabled + and connection.supports_graceful_disconnect): + connection.register_watcher( + "GRACEFUL_DISCONNECT", self._on_graceful_disconnect_event) + + def _on_graceful_disconnect_event(self, event): + """ + Called from the connection's event thread when a GRACEFUL_DISCONNECT + event is received on one of this pool's connections. + """ + log.debug("Received GRACEFUL_DISCONNECT on connection to %s", self.host) + self._session.cluster.on_graceful_disconnect(self.host) + + def on_graceful_disconnect(self): + """ + Gracefully drain all connections to this host (CEP-59): in-flight + requests are allowed to complete before the connections are closed. + """ + if self.is_shutdown: + return + connections = self.get_connections() + if not connections: + return + log.info("Draining all connections to %s gracefully", self.host) + for connection in connections: + connection.start_graceful_drain() + def _replace(self, connection): with self._lock: if self.is_shutdown: @@ -513,6 +550,7 @@ def _replace(self, connection): conn = self._session.cluster.connection_factory(self.host.endpoint, on_orphaned_stream_released=self.on_orphaned_stream_released) if self._keyspace: conn.set_keyspace_blocking(self._keyspace) + self._maybe_register_graceful_disconnect(conn) self._connection = conn except Exception: log.warning("Failed reconnecting %s. Retrying." % (self.host.endpoint,)) @@ -616,6 +654,9 @@ def __init__(self, host, host_distance, session): for conn in self._connections: conn.set_keyspace_blocking(self._keyspace) + for conn in self._connections: + self._maybe_register_graceful_disconnect(conn) + self._trash = set() self._next_trash_allowed_at = time.time() self.open_count = core_conns @@ -715,6 +756,7 @@ def _add_conn_if_under_max(self): conn = self._session.cluster.connection_factory(self.host.endpoint, on_orphaned_stream_released=self.on_orphaned_stream_released) if self._keyspace: conn.set_keyspace_blocking(self._session.keyspace) + self._maybe_register_graceful_disconnect(conn) self._next_trash_allowed_at = time.time() + _MIN_TRASH_INTERVAL with self._lock: new_connections = self._connections[:] + [conn] @@ -818,6 +860,39 @@ def on_orphaned_stream_released(self): """ self._signal_available_conn() + def _maybe_register_graceful_disconnect(self, connection): + """ + Register for GRACEFUL_DISCONNECT events (CEP-59) on a pooled + connection, if the feature is enabled and this connection's server + advertised support for it. + """ + if (self._session.cluster.graceful_disconnect_enabled + and connection.supports_graceful_disconnect): + connection.register_watcher( + "GRACEFUL_DISCONNECT", self._on_graceful_disconnect_event) + + def _on_graceful_disconnect_event(self, event): + """ + Called from the connection's event thread when a GRACEFUL_DISCONNECT + event is received on one of this pool's connections. + """ + log.debug("Received GRACEFUL_DISCONNECT on connection to %s", self.host) + self._session.cluster.on_graceful_disconnect(self.host) + + def on_graceful_disconnect(self): + """ + Gracefully drain all connections to this host (CEP-59): in-flight + requests are allowed to complete before the connections are closed. + """ + if self.is_shutdown: + return + connections = self.get_connections() + if not connections: + return + log.info("Draining all connections to %s gracefully", self.host) + for connection in connections: + connection.start_graceful_drain() + def _maybe_trash_connection(self, connection): core_conns = self._session.cluster.get_core_connections_per_host(self.host_distance) did_trash = False diff --git a/cassandra/protocol.py b/cassandra/protocol.py index b1c4183cf4..767933f06d 100644 --- a/cassandra/protocol.py +++ b/cassandra/protocol.py @@ -987,7 +987,8 @@ def send_body(self, f, protocol_version): known_event_types = frozenset(( 'TOPOLOGY_CHANGE', 'STATUS_CHANGE', - 'SCHEMA_CHANGE' + 'SCHEMA_CHANGE', + 'GRACEFUL_DISCONNECT' )) @@ -1032,6 +1033,12 @@ def recv_status_change(cls, f, protocol_version): address = read_inet(f) return dict(change_type=change_type, address=address) + @classmethod + def recv_graceful_disconnect(cls, f, protocol_version): + # The GRACEFUL_DISCONNECT event (CEP-59) has no body; the type + # string is enough. + return {} + @classmethod def recv_schema_change(cls, f, protocol_version): # "CREATED", "DROPPED", or "UPDATED" diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py index a9da91009a..ca28f14c3a 100644 --- a/tests/integration/__init__.py +++ b/tests/integration/__init__.py @@ -349,6 +349,7 @@ def _id_and_mark(f): greaterthanorequalcass3_11 = unittest.skipUnless(CASSANDRA_VERSION >= Version('3.11'), 'Cassandra version 3.11 or greater required') greaterthanorequalcass40 = unittest.skipUnless(CASSANDRA_VERSION >= Version('4.0'), 'Cassandra version 4.0 or greater required') greaterthanorequalcass50 = unittest.skipUnless(CASSANDRA_VERSION >= Version('5.0-beta'), 'Cassandra version 5.0 or greater required') +greaterthanorequalcass70 = unittest.skipUnless(CASSANDRA_VERSION >= Version('7.0'), 'Cassandra version 7.0 or greater required') lessthanorequalcass40 = unittest.skipUnless(CASSANDRA_VERSION <= Version('4.0'), 'Cassandra version less or equal to 4.0 required') lessthancass40 = unittest.skipUnless(CASSANDRA_VERSION < Version('4.0'), 'Cassandra version less than 4.0 required') lessthancass30 = unittest.skipUnless(CASSANDRA_VERSION < Version('3.0'), 'Cassandra version less then 3.0 required') diff --git a/tests/integration/standard/test_graceful_disconnect.py b/tests/integration/standard/test_graceful_disconnect.py new file mode 100644 index 0000000000..10ccbd8b5b --- /dev/null +++ b/tests/integration/standard/test_graceful_disconnect.py @@ -0,0 +1,114 @@ +# 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. + +""" +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 connections 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 whole module is skipped by the version requirement below. +""" + +import threading +import time +import unittest + +from tests.integration import (get_node, greaterthanorequalcass70, + remove_cluster, requirecassandra, + use_cluster, TestCluster) +from tests.util import wait_until + +GRACEFUL_DISCONNECT_CLUSTER_NAME = 'graceful_disconnect_cluster' +QUERY = "SELECT * FROM system.local" + + +def setup_module(): + # The server-side feature is disabled by default: + use_cluster(GRACEFUL_DISCONNECT_CLUSTER_NAME, [2], + configuration_options={'graceful_disconnect_enabled': True}) + + +def teardown_module(): + remove_cluster() + + +@requirecassandra +@greaterthanorequalcass70 +class GracefulDisconnectTests(unittest.TestCase): + + def test_fail_over_without_disruption_when_node_drains(self): + cluster = TestCluster(metrics_enabled=True) + session = cluster.connect(wait_for_all_pools=True) + try: + # 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: + successes = [0] + failures = [] + stopped = threading.Event() + + def load(): + while not stopped.is_set(): + try: + session.execute(QUERY) + successes[0] += 1 + time.sleep(0.005) + except Exception as exc: + failures.append(exc) + + load_thread = threading.Thread(target=load, + name='graceful-disconnect-load') + load_thread.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. + get_node(2).nodetool('drain') + + # The driver must have observed the event (this is also the + # end-to-end check for the metrics counter): + wait_until( + lambda: cluster.metrics.stats.graceful_disconnects > 0, + 0.5, 60) + self.assertGreater(cluster.metrics.stats.graceful_disconnects, 0) + + # Queries must keep succeeding after the drain (load fails + # over to the other node): + successes_after_event = successes[0] + wait_until( + lambda: successes[0] > successes_after_event + 100, + 0.5, 60) + self.assertGreater(successes[0], successes_after_event + 100) + finally: + stopped.set() + load_thread.join(timeout=10) + + self.assertFalse(load_thread.is_alive(), + "load thread should have terminated") + + # The whole point of graceful disconnect: the shutdown must be + # invisible to the application, no request may fail. + self.assertEqual( + failures, [], + "expected no disruptive exceptions, got: %r" % (failures,)) + finally: + cluster.shutdown() diff --git a/tests/unit/test_connection.py b/tests/unit/test_connection.py index 3bca654c55..3b87adaaab 100644 --- a/tests/unit/test_connection.py +++ b/tests/unit/test_connection.py @@ -311,10 +311,10 @@ def send_msg(msg, req_id, msg_callback): max_request_id=127, lock=Lock(), in_flight=0, is_idle=True, - is_defunct=False, is_closed=False, + is_draining=False, is_defunct=False, is_closed=False, get_request_id=lambda: request_id, send_msg=Mock(side_effect=send_msg)) - non_idle_connection = Mock(spec=Connection, in_flight=0, is_idle=False, is_defunct=False, is_closed=False) + non_idle_connection = Mock(spec=Connection, in_flight=0, is_idle=False, is_draining=False, is_defunct=False, is_closed=False) get_holders = self.make_get_holders(1) holder = get_holders.return_value[0] @@ -353,7 +353,7 @@ def test_no_req_ids(self, *args): max_connection = Mock(spec=Connection, host='localhost', lock=Lock(), max_request_id=in_flight - 1, in_flight=in_flight, - is_idle=True, is_defunct=False, is_closed=False) + is_idle=True, is_draining=False, is_defunct=False, is_closed=False) holder = get_holders.return_value[0] holder.get_connections.return_value.append(max_connection) @@ -379,7 +379,7 @@ def send_msg(msg, req_id, msg_callback): max_request_id=127, lock=Lock(), in_flight=0, is_idle=True, - is_defunct=False, is_closed=False, + is_draining=False, is_defunct=False, is_closed=False, get_request_id=lambda: request_id, send_msg=Mock(side_effect=send_msg)) holder = get_holders.return_value[0] @@ -409,7 +409,7 @@ def send_msg(msg, req_id, msg_callback): max_request_id=127, lock=Lock(), in_flight=0, is_idle=True, - is_defunct=False, is_closed=False, + is_draining=False, is_defunct=False, is_closed=False, get_request_id=lambda: request_id, send_msg=Mock(side_effect=send_msg)) holder = get_holders.return_value[0] diff --git a/tests/unit/test_graceful_disconnect.py b/tests/unit/test_graceful_disconnect.py new file mode 100644 index 0000000000..5c871e4b4a --- /dev/null +++ b/tests/unit/test_graceful_disconnect.py @@ -0,0 +1,446 @@ +# 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. + +""" +Tests for graceful disconnect support (CEP-59): decoding the +GRACEFUL_DISCONNECT event, per-connection capability negotiation against the +SUPPORTED response, draining connections without disrupting in-flight +requests, and the control connection / pool / cluster plumbing. +""" + +import unittest +from io import BytesIO +from threading import Lock +from unittest.mock import Mock, NonCallableMagicMock, patch + +from cassandra.cluster import Cluster, ControlConnection, Session +from cassandra.connection import (Connection, ConnectionShutdown, + DefaultEndPoint) +from cassandra.metrics import Metrics +from cassandra.pool import Host, HostConnection, NoConnectionsAvailable +from cassandra.policies import HostDistance, SimpleConvictionPolicy +from cassandra.protocol import (EventMessage, NotSupportedError, + ProtocolHandler, SupportedMessage, + write_string, write_stringmultimap) +from cassandra.connection import _Frame + +from tests.unit.test_control_connection import (MockCluster, MockConnection, + FakeTime) + + +class GracefulDisconnectEventDecodeTest(unittest.TestCase): + """ + Decodes the CEP-59 GRACEFUL_DISCONNECT event body as the server would + send it (the event is body-less: only the type string is present). + """ + + @staticmethod + def _event_body(event_type): + buf = BytesIO() + write_string(buf, event_type) + return buf.getvalue() + + def test_decode_graceful_disconnect_event(self): + for protocol_version in (3, 4, 5, 6): + msg = EventMessage.recv_body( + BytesIO(self._event_body('GRACEFUL_DISCONNECT')), protocol_version) + self.assertIsInstance(msg, EventMessage) + self.assertEqual(msg.event_type, 'GRACEFUL_DISCONNECT') + self.assertEqual(msg.event_args, {}) + + def test_decode_is_case_insensitive(self): + msg = EventMessage.recv_body( + BytesIO(self._event_body('graceful_disconnect')), 4) + self.assertEqual(msg.event_type, 'GRACEFUL_DISCONNECT') + + def test_unknown_event_type_raises(self): + self.assertRaises( + NotSupportedError, + EventMessage.recv_body, + BytesIO(self._event_body('GRACEFUL_DISCONNECT_V2')), 4) + + +class SupportsGracefulDisconnectTest(unittest.TestCase): + """ + The capability is negotiated per connection from the SUPPORTED options + string-multimap; the server may send the key with an explicit "false" + value when the feature is disabled. + """ + + def test_detect_capability_from_supported_options(self): + supports = Connection._supports_graceful_disconnect + self.assertFalse(supports(None)) + self.assertFalse(supports({})) + self.assertFalse(supports({'CQL_VERSION': ['3.4.7']})) + self.assertTrue(supports({'GRACEFUL_DISCONNECT': []})) + self.assertTrue(supports({'GRACEFUL_DISCONNECT': ['true']})) + self.assertFalse(supports({'GRACEFUL_DISCONNECT': ['false']})) + self.assertFalse(supports({'GRACEFUL_DISCONNECT': ['FALSE']})) + + +class ConnectionGracefulDrainTest(unittest.TestCase): + """ + Draining behavior of a single connection when a GRACEFUL_DISCONNECT + event is received on it. + """ + + def make_connection(self): + c = Connection(DefaultEndPoint('1.2.3.4')) + c._socket = Mock() + c._socket.send.side_effect = lambda x: len(x) + c.close = Mock() + return c + + def _options_response_requests(self, connection): + return {0: (connection._handle_options_response, + ProtocolHandler.decode_message, [])} + + def _process_supported_message(self, connection, options): + options_buf = BytesIO() + write_stringmultimap(options_buf, options) + body = options_buf.getvalue() + connection.process_msg( + _Frame(version=4, flags=0, stream=0, opcode=SupportedMessage.opcode, + body_offset=9, end_pos=9 + len(body)), + body) + + def test_options_response_sets_supports_graceful_disconnect(self): + c = self.make_connection() + c._requests = self._options_response_requests(c) + self.assertFalse(c.supports_graceful_disconnect) + self._process_supported_message(c, { + 'CQL_VERSION': ['3.4.7'], + 'COMPRESSION': [], + 'GRACEFUL_DISCONNECT': ['true'] + }) + self.assertTrue(c.supports_graceful_disconnect) + + def test_options_response_without_capability(self): + c = self.make_connection() + c._requests = self._options_response_requests(c) + self._process_supported_message(c, { + 'CQL_VERSION': ['3.4.7'], + 'COMPRESSION': [] + }) + self.assertFalse(c.supports_graceful_disconnect) + + def test_options_response_with_capability_disabled(self): + c = self.make_connection() + c._requests = self._options_response_requests(c) + self._process_supported_message(c, { + 'CQL_VERSION': ['3.4.7'], + 'COMPRESSION': [], + 'GRACEFUL_DISCONNECT': ['false'] + }) + self.assertFalse(c.supports_graceful_disconnect) + + def test_drain_waits_for_pending_requests(self): + c = self.make_connection() + # a pending request on stream 5 + response_callback = Mock() + c._requests = {5: (response_callback, ProtocolHandler.decode_message, [])} + + c.start_graceful_drain() + + # connection not closed yet because there is a pending request + self.assertTrue(c.is_draining) + self.assertEqual(c.close.call_count, 0) + + # new writes are refused + self.assertRaises(ConnectionShutdown, c.send_msg, + Mock(), 6, Mock()) + + # when the pending request completes, the connection closes + options_buf = BytesIO() + write_stringmultimap(options_buf, {'CQL_VERSION': ['3.4.7'], + 'COMPRESSION': []}) + body = options_buf.getvalue() + c.process_msg( + _Frame(version=4, flags=0, stream=5, opcode=SupportedMessage.opcode, + body_offset=9, end_pos=9 + len(body)), + body) + + response_callback.assert_called_once() + c.close.assert_called_once_with() + + def test_drain_closes_immediately_when_no_pending(self): + c = self.make_connection() + c.start_graceful_drain() + self.assertTrue(c.is_draining) + c.close.assert_called_once_with() + + def test_drain_is_idempotent(self): + c = self.make_connection() + c.start_graceful_drain() + c.start_graceful_drain() + c.close.assert_called_once_with() + + def test_pushed_event_starts_drain_and_notifies_watchers(self): + c = self.make_connection() + watcher = Mock() + c._push_watchers['GRACEFUL_DISCONNECT'].add(watcher) + + c.handle_pushed(EventMessage('GRACEFUL_DISCONNECT', {})) + + # the drain started (and completed, since nothing was pending) + # before the watcher was notified + self.assertTrue(c.is_draining) + c.close.assert_called_once_with() + watcher.assert_called_once_with({}) + + def test_pushed_event_drains_even_if_watcher_errors(self): + c = self.make_connection() + watcher = Mock(side_effect=RuntimeError("misbehaving callback")) + c._push_watchers['GRACEFUL_DISCONNECT'].add(watcher) + + c.handle_pushed(EventMessage('GRACEFUL_DISCONNECT', {})) + + self.assertTrue(c.is_draining) + c.close.assert_called_once_with() + + def test_other_events_do_not_drain(self): + c = self.make_connection() + c.handle_pushed(EventMessage('STATUS_CHANGE', + {'change_type': 'UP', + 'address': ('1.2.3.4', 9042)})) + self.assertFalse(c.is_draining) + self.assertEqual(c.close.call_count, 0) + + +class ControlConnectionGracefulDisconnectTest(unittest.TestCase): + + def setUp(self): + self.cluster = MockCluster() + self.cluster.graceful_disconnect_enabled = True + self.cluster.on_graceful_disconnect = Mock() + self.connection = MockConnection() + self.time = FakeTime() + + self.control_connection = ControlConnection(self.cluster, 1, 0, 0, 0) + self.control_connection._connection = self.connection + self.control_connection._time = self.time + + def _registered_watchers(self, connection): + host = self.cluster.metadata.get_host(DefaultEndPoint("192.168.1.0")) + self.cluster.connection_factory = Mock(return_value=connection) + connection.register_watchers = Mock() + connection.close = Mock() + # _try_connect issues the peers/local queries with fail_on_error=False, + # which returns (success, result) pairs + peers_result, local_result = connection.wait_for_responses.return_value + connection.wait_for_responses = Mock( + return_value=((True, peers_result), (True, local_result))) + cc = self.control_connection + cc._refresh_node_list_and_token_map = Mock() + cc._refresh_schema = Mock() + cc._try_connect(host) + args, kwargs = connection.register_watchers.call_args + return args[0] + + def test_registers_graceful_disconnect_when_supported(self): + connection = MockConnection() + connection.supports_graceful_disconnect = True + + watchers = self._registered_watchers(connection) + + self.assertEqual( + sorted(watchers.keys()), + ['GRACEFUL_DISCONNECT', 'SCHEMA_CHANGE', 'STATUS_CHANGE', + 'TOPOLOGY_CHANGE']) + + def test_does_not_register_when_server_does_not_advertise(self): + connection = MockConnection() + connection.supports_graceful_disconnect = False + + watchers = self._registered_watchers(connection) + + self.assertEqual( + sorted(watchers.keys()), + ['SCHEMA_CHANGE', 'STATUS_CHANGE', 'TOPOLOGY_CHANGE']) + + def test_does_not_register_when_disabled(self): + self.cluster.graceful_disconnect_enabled = False + connection = MockConnection() + connection.supports_graceful_disconnect = True + + watchers = self._registered_watchers(connection) + + self.assertEqual( + sorted(watchers.keys()), + ['SCHEMA_CHANGE', 'STATUS_CHANGE', 'TOPOLOGY_CHANGE']) + + def test_handle_graceful_disconnect(self): + self.control_connection.reconnect = Mock() + + self.control_connection._handle_graceful_disconnect({}) + + host = self.cluster.metadata.get_host(DefaultEndPoint("192.168.1.0")) + self.cluster.on_graceful_disconnect.assert_called_once_with(host) + # the control connection proactively moves to another host + self.control_connection.reconnect.assert_called_once_with() + + def test_handle_graceful_disconnect_without_connection(self): + self.control_connection.reconnect = Mock() + self.control_connection._connection = None + + self.control_connection._handle_graceful_disconnect({}) + + self.assertEqual(self.cluster.on_graceful_disconnect.call_count, 0) + self.control_connection.reconnect.assert_called_once_with() + + +class HostConnectionGracefulDisconnectTest(unittest.TestCase): + + def make_session(self, graceful_disconnect_enabled=True): + session = NonCallableMagicMock(spec=Session, keyspace='foobarkeyspace') + session.cluster.get_core_connections_per_host.return_value = 1 + session.cluster.get_max_requests_per_connection.return_value = 1 + session.cluster.get_max_connections_per_host.return_value = 1 + session.cluster.graceful_disconnect_enabled = graceful_disconnect_enabled + return session + + def make_connection_mock(self, supports_graceful_disconnect=True): + return NonCallableMagicMock( + spec=Connection, in_flight=0, is_draining=False, is_defunct=False, + is_closed=False, max_request_id=100, lock=Lock(), + supports_graceful_disconnect=supports_graceful_disconnect) + + def make_pool(self, session, conn): + host = Mock(spec=Host, address='ip1') + session.cluster.connection_factory.return_value = conn + return HostConnection(host, HostDistance.LOCAL, session) + + def test_registers_watcher_when_enabled_and_supported(self): + session = self.make_session() + conn = self.make_connection_mock() + pool = self.make_pool(session, conn) + + conn.register_watcher.assert_called_once_with( + "GRACEFUL_DISCONNECT", pool._on_graceful_disconnect_event) + + def test_does_not_register_watcher_when_disabled(self): + session = self.make_session(graceful_disconnect_enabled=False) + conn = self.make_connection_mock() + self.make_pool(session, conn) + + self.assertEqual(conn.register_watcher.call_count, 0) + + def test_does_not_register_watcher_when_not_supported(self): + session = self.make_session() + conn = self.make_connection_mock(supports_graceful_disconnect=False) + self.make_pool(session, conn) + + self.assertEqual(conn.register_watcher.call_count, 0) + + def test_event_on_pooled_connection_notifies_cluster(self): + session = self.make_session() + conn = self.make_connection_mock() + pool = self.make_pool(session, conn) + + pool._on_graceful_disconnect_event({}) + + session.cluster.on_graceful_disconnect.assert_called_once_with(pool.host) + + def test_on_graceful_disconnect_drains_connections(self): + session = self.make_session() + conn = self.make_connection_mock() + pool = self.make_pool(session, conn) + + pool.on_graceful_disconnect() + + conn.start_graceful_drain.assert_called_once_with() + + def test_on_graceful_disconnect_noop_when_shutdown(self): + session = self.make_session() + conn = self.make_connection_mock() + pool = self.make_pool(session, conn) + pool.shutdown() + conn.start_graceful_drain.reset_mock() + + pool.on_graceful_disconnect() + + self.assertEqual(conn.start_graceful_drain.call_count, 0) + + def test_borrow_fails_over_while_draining(self): + session = self.make_session() + conn = self.make_connection_mock() + pool = self.make_pool(session, conn) + + conn.is_draining = True + + self.assertRaises(NoConnectionsAvailable, + pool.borrow_connection, 0) + + +class ClusterGracefulDisconnectTest(unittest.TestCase): + + def test_enabled_by_default(self): + cluster = Cluster(contact_points=['127.0.0.1']) + self.assertTrue(cluster.graceful_disconnect_enabled) + + def test_can_be_disabled(self): + cluster = Cluster(contact_points=['127.0.0.1'], + graceful_disconnect_enabled=False) + self.assertFalse(cluster.graceful_disconnect_enabled) + + def _cluster_with_mock_session(self): + cluster = Cluster(contact_points=['127.0.0.1']) + session = Mock(spec=Session) + host = Mock(spec=Host, address='ip1') + other_host = Mock(spec=Host, address='ip2') + pool = Mock() + session._pools = {host: pool} + cluster.sessions.add(session) + return cluster, session, host, other_host, pool + + def test_on_graceful_disconnect_drains_pools_for_host(self): + cluster, session, host, _, pool = self._cluster_with_mock_session() + + cluster.on_graceful_disconnect(host) + + pool.on_graceful_disconnect.assert_called_once_with() + + def test_on_graceful_disconnect_ignores_other_hosts(self): + cluster, session, host, other_host, pool = self._cluster_with_mock_session() + + cluster.on_graceful_disconnect(other_host) + + self.assertEqual(pool.on_graceful_disconnect.call_count, 0) + + def test_on_graceful_disconnect_ignores_unknown_host(self): + cluster, _, _, _, pool = self._cluster_with_mock_session() + + cluster.on_graceful_disconnect(None) + + self.assertEqual(pool.on_graceful_disconnect.call_count, 0) + + def test_on_graceful_disconnect_increments_metrics(self): + cluster, _, host, _, _ = self._cluster_with_mock_session() + cluster.metrics_enabled = True + cluster.metrics = Mock(spec=Metrics) + + cluster.on_graceful_disconnect(host) + + cluster.metrics.on_graceful_disconnect.assert_called_once_with() + + +class MetricsGracefulDisconnectTest(unittest.TestCase): + + def test_graceful_disconnects_counter(self): + metrics = Metrics(Mock()) + before = metrics.stats.graceful_disconnects + metrics.on_graceful_disconnect() + self.assertEqual(metrics.stats.graceful_disconnects, before + 1) diff --git a/tests/unit/test_host_connection_pool.py b/tests/unit/test_host_connection_pool.py index d8b5ca976e..0aab1919c8 100644 --- a/tests/unit/test_host_connection_pool.py +++ b/tests/unit/test_host_connection_pool.py @@ -39,7 +39,7 @@ def make_session(self): def test_borrow_and_return(self): host = Mock(spec=Host, address='ip1') session = self.make_session() - conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, max_request_id=100) + conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_draining=False, is_defunct=False, is_closed=False, max_request_id=100) session.cluster.connection_factory.return_value = conn pool = self.PoolImpl(host, HostDistance.LOCAL, session) @@ -58,7 +58,7 @@ def test_borrow_and_return(self): def test_failed_wait_for_connection(self): host = Mock(spec=Host, address='ip1') session = self.make_session() - conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, max_request_id=100) + conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_draining=False, is_defunct=False, is_closed=False, max_request_id=100) session.cluster.connection_factory.return_value = conn pool = self.PoolImpl(host, HostDistance.LOCAL, session) @@ -76,7 +76,7 @@ def test_failed_wait_for_connection(self): def test_successful_wait_for_connection(self): host = Mock(spec=Host, address='ip1') session = self.make_session() - conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, max_request_id=100, lock=Lock()) + conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_draining=False, is_defunct=False, is_closed=False, max_request_id=100, lock=Lock()) session.cluster.connection_factory.return_value = conn pool = self.PoolImpl(host, HostDistance.LOCAL, session) @@ -100,7 +100,7 @@ def get_second_conn(): def test_spawn_when_at_max(self): host = Mock(spec=Host, address='ip1') session = self.make_session() - conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, max_request_id=100) + conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_draining=False, is_defunct=False, is_closed=False, max_request_id=100) conn.max_request_id = 100 session.cluster.connection_factory.return_value = conn @@ -126,7 +126,7 @@ def test_spawn_when_at_max(self): def test_return_defunct_connection(self): host = Mock(spec=Host, address='ip1') session = self.make_session() - conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, + conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_draining=False, is_defunct=False, is_closed=False, max_request_id=100, signaled_error=False) session.cluster.connection_factory.return_value = conn @@ -145,7 +145,7 @@ def test_return_defunct_connection(self): def test_return_defunct_connection_on_down_host(self): host = Mock(spec=Host, address='ip1') session = self.make_session() - conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, + conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_draining=False, is_defunct=False, is_closed=False, max_request_id=100, signaled_error=False, orphaned_threshold_reached=False) session.cluster.connection_factory.return_value = conn @@ -167,7 +167,7 @@ def test_return_defunct_connection_on_down_host(self): def test_return_closed_connection(self): host = Mock(spec=Host, address='ip1') session = self.make_session() - conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=True, max_request_id=100, + conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_draining=False, is_defunct=False, is_closed=True, max_request_id=100, signaled_error=False, orphaned_threshold_reached=False) session.cluster.connection_factory.return_value = conn @@ -214,7 +214,7 @@ class HostConnectionPoolTests(_PoolTests): def test_all_connections_trashed(self): host = Mock(spec=Host, address='ip1') session = self.make_session() - conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_defunct=False, is_closed=False, max_request_id=100, + conn = NonCallableMagicMock(spec=Connection, in_flight=0, is_draining=False, is_defunct=False, is_closed=False, max_request_id=100, lock=Lock()) session.cluster.connection_factory.return_value = conn session.cluster.get_core_connections_per_host.return_value = 1