From 35e6cf16971b9750ad57d0881b1db68cb060dd56 Mon Sep 17 00:00:00 2001 From: omniCoder77 Date: Sun, 8 Mar 2026 10:05:59 +0530 Subject: [PATCH 1/6] cep 59 --- CHANGES.txt | 1 + conf/cassandra.yaml | 13 + conf/cassandra_latest.yaml | 14 + .../operating/graceful_disconnect.adoc | 91 ++++ .../org/apache/cassandra/config/Config.java | 4 + .../cassandra/config/DatabaseDescriptor.java | 18 + .../cassandra/metrics/ClientMetrics.java | 23 + .../service/NativeTransportService.java | 6 + .../cassandra/service/StorageService.java | 21 + .../service/StorageServiceMBean.java | 5 + .../cassandra/tools/nodetool/Drain.java | 3 +- .../org/apache/cassandra/transport/Event.java | 27 +- .../GracefulDisconnectLifecycle.java | 153 +++++++ .../transport/InitialConnectionHandler.java | 3 + .../apache/cassandra/transport/Server.java | 60 ++- .../cassandra/transport/SimpleClient.java | 414 +++++++++--------- .../transport/messages/OptionsMessage.java | 3 + .../transport/messages/StartupMessage.java | 1 + .../test/GracefulDisconnectTest.java | 334 ++++++++++++++ .../config/DatabaseDescriptorTest.java | 24 + .../cassandra/metrics/ClientMetricsTest.java | 16 + .../cassandra/service/StorageServiceTest.java | 27 ++ .../GracefulDisconnectLifecycleTest.java | 239 ++++++++++ 23 files changed, 1287 insertions(+), 213 deletions(-) create mode 100644 doc/modules/cassandra/pages/managing/operating/graceful_disconnect.adoc create mode 100644 src/java/org/apache/cassandra/transport/GracefulDisconnectLifecycle.java create mode 100644 test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java create mode 100644 test/unit/org/apache/cassandra/transport/GracefulDisconnectLifecycleTest.java diff --git a/CHANGES.txt b/CHANGES.txt index 5bbd8cba36fc..124c735e5a0c 100644 --- a/CHANGES.txt +++ b/CHANGES.txt @@ -6,6 +6,7 @@ * Avoid using ObjectUtils.getFirstNonNull in Schema (CASSANDRA-21394) * Allow nodetool garbagecollect to take a user defined list of SSTables (CASSANDRA-16767) * Add a guardrail for misprepared statements (CASSANDRA-21139) + * [CEP-59] Implementation of In-Band Connection Draining (Graceful Disconnect) (CASSANDRA-21191) Merged from 6.0: * Make cqlsh prompt to reset to no keyspace set by USE after dropping that keyspace (CASSANDRA-21548) * Implement CMS rediscovery and recovery protocol (CASSANDRA-20476) diff --git a/conf/cassandra.yaml b/conf/cassandra.yaml index 7201fc8c6efc..5ef5585a3ad2 100644 --- a/conf/cassandra.yaml +++ b/conf/cassandra.yaml @@ -1104,6 +1104,19 @@ native_transport_allow_older_protocols: true # native_transport_rate_limiting_enabled: false # native_transport_max_requests_per_second: 1000000 +# When enabled, nodes will signal connected clients before shutting down, +# allowing in-flight requests to complete without client-visible timeouts. +# This applies to intentional shutdowns (nodetool drain, rolling restarts, +# controlled JVM shutdown). Clients must subscribe to the GRACEFUL_DISCONNECT +# event via REGISTER to benefit from this behavior. +# Requires driver support for the GRACEFUL_DISCONNECT event type. +# See: doc/modules/cassandra/pages/managing/operating/graceful_disconnect.adoc +# Defaults to false. +# graceful_disconnect_enabled: false + +# Time given to clients to stop sending new requests after the GRACEFUL_DISCONNECT event is emitted. +# graceful_disconnect_grace_period: 5s + # The address or interface to bind the native transport server to. # # Set rpc_address OR rpc_interface, not both. diff --git a/conf/cassandra_latest.yaml b/conf/cassandra_latest.yaml index 63da0583276e..16a6bc1d8ae2 100644 --- a/conf/cassandra_latest.yaml +++ b/conf/cassandra_latest.yaml @@ -1087,6 +1087,20 @@ native_transport_allow_older_protocols: true # native_transport_rate_limiting_enabled: false # native_transport_max_requests_per_second: 1000000 +# When enabled, nodes will signal connected clients before shutting down, +# allowing in-flight requests to complete without client-visible timeouts. +# This applies to intentional shutdowns (nodetool drain, rolling restarts, +# controlled JVM shutdown). Clients must subscribe to the GRACEFUL_DISCONNECT +# event via REGISTER to benefit from this behavior. +# +# Requires driver support for the GRACEFUL_DISCONNECT event type. +# See: doc/modules/cassandra/pages/managing/operating/graceful_disconnect.adoc + +graceful_disconnect_enabled: true + +# Time given to clients to stop sending new requests after the GRACEFUL_DISCONNECT event is emitted. +# graceful_disconnect_grace_period: 5s + # The address or interface to bind the native transport server to. # # Set rpc_address OR rpc_interface, not both. diff --git a/doc/modules/cassandra/pages/managing/operating/graceful_disconnect.adoc b/doc/modules/cassandra/pages/managing/operating/graceful_disconnect.adoc new file mode 100644 index 000000000000..a4e33faf64ad --- /dev/null +++ b/doc/modules/cassandra/pages/managing/operating/graceful_disconnect.adoc @@ -0,0 +1,91 @@ += Graceful Disconnect — In-Band Connection Draining for Cassandra Node Shutdown + +== Vocabulary + +in-band:: same connection +In-flight requests:: requests that have been sent but not yet completed + +== Introduction + +When a Cassandra node has to be taken offline client drivers have no reliable, in-band signal that the node is going away. + +* Drivers keep sending requests to a shutting-down node until they hit a socket close or timeout. +* In-flight requests are abandoned, producing `ReadTimeoutException` errors on the client side. +* Retry storms emerge as clients simultaneously rediscover the topology. + +== Solution + +Introduce an in-band signal — `GRACEFUL_DISCONNECT` — so that the server can notify clients (that have subscribed) connection before closing it. This gives drivers time to: + +. Stop sending new requests on that connection/node. +. Let all in-flight requests complete. +. Close that socket connection (this notifies server that there are no pending queries on client side). +. Try reconnecting with exponential backoff. + +== New Configurations Introduced + +`graceful_disconnect_enabled`:: A configuration that enables the server to perform graceful disconnect. +`graceful_disconnect_grace_period`:: A configuration that forces shutdown of any active driver connection to the node after the grace period expires. + +[cols="1,1,3,1", options="header"] +|=== +| Parameter | Type | Description | Default +| `graceful_disconnect_enabled` | Boolean | A configuration that enables server to perform graceful disconnect. | false +| `graceful_disconnect_grace_period` | Duration (ms) | A configuration that determines after how much time to force close a socket connection, if there is any pending connection from driver side. | 5s +|=== + +== Client Compatibility & Upgrade Strategy + +=== Legacy Drivers — No Change Required + +Drivers that do not support graceful disconnect are not affected by any value of `graceful_disconnect_enabled` or `graceful_disconnect_grace_period`. + +=== Required Driver Versions + +To benefit from graceful draining, a compatible driver must be used. + +[cols="1,1", options="header"] +|=== +| Driver | Status +| Java | https://issues.apache.org/jira/browse/CASSJAVA-124[In progress] +| Python | https://issues.apache.org/jira/browse/CASSPYTHON-16[In progress] +| Node.js | https://issues.apache.org/jira/browse/CASSNODEJS-5[In progress] +| Go | https://issues.apache.org/jira/browse/CASSGO-117[In progress] +| C++ | https://issues.apache.org/jira/browse/CASSCPP-7[In progress] +|=== + +=== Mixed-Fleet Rollout + +During a rolling upgrade where some nodes support the feature and others do not: + +* Nodes *with* `graceful_disconnect_enabled: true` advertise the capability in `SUPPORTED` and drivers subscribe. +* Nodes *without* the feature omit the key from `SUPPORTED`. Compatible drivers silently fall back to the legacy TCP teardown path for those nodes. + +No special coordination is needed. Drivers handle both peers transparently within the same session. + +== Operational Visibility + +=== Metrics + +[cols="1,3", options="header"] +|=== +| Metric | Description +| `ConnectionsDraining` | Current count of connections in the Draining state (subscribed and awaiting in-flight completion). Should rise at drain start and fall to zero before shutdown completes. +| `ForcedDisconnects` | Cumulative count of connections force-closed after `graceful_disconnect_grace_period` expired without the driver closing cleanly. Persistent non-zero values here indicate driver-side issues or an undersized grace period. +|=== + +== Failure Modes & Edge Cases + +=== What if the driver does not support `GRACEFUL_DISCONNECT`? + +If the driver never sent `REGISTER` for this event type (legacy driver, or a compatible driver connected to a legacy node), the server applies *standard TCP teardown* — the same behavior as before this feature existed. There is no error, no retry storm specific to this feature, and no action required from operators. + +This is the designed fallback and not an error condition. + +=== What will happen if two different driver, one supporting graceful disconnect while other not, are talking to a node supporting graceful disconnect? + +Driver supporting graceful disconnect will disconnect gracefully, while the one not supporting graceful disconnect will behave as a legacy driver. + +=== What if the server crashes instead of draining cleanly? + +`GRACEFUL_DISCONNECT` is *not* a crash-safety mechanism. If the JVM is killed with `SIGKILL`, the node loses power, or an OOM kill occurs, no `GRACEFUL_DISCONNECT` event is emitted. Drivers fall back to socket-close detection and gossip `DOWN` events — the current behavior. \ No newline at end of file diff --git a/src/java/org/apache/cassandra/config/Config.java b/src/java/org/apache/cassandra/config/Config.java index 011d149333c5..3b3abda52ca2 100644 --- a/src/java/org/apache/cassandra/config/Config.java +++ b/src/java/org/apache/cassandra/config/Config.java @@ -143,6 +143,10 @@ public static Set splitCommaDelimited(String src) /** Triggers automatic allocation of tokens if set, based on the provided replica count for a datacenter */ public Integer allocate_tokens_for_local_replication_factor = null; + public boolean graceful_disconnect_enabled = false; + + public volatile DurationSpec.LongMillisecondsBound graceful_disconnect_grace_period = new DurationSpec.LongMillisecondsBound(5000); + @Replaces(oldName = "native_transport_idle_timeout_in_ms", converter = Converters.MILLIS_DURATION_LONG, deprecated = true) public DurationSpec.LongMillisecondsBound native_transport_idle_timeout = new DurationSpec.LongMillisecondsBound("0ms"); diff --git a/src/java/org/apache/cassandra/config/DatabaseDescriptor.java b/src/java/org/apache/cassandra/config/DatabaseDescriptor.java index 6554b6ff05b3..065abaec1fae 100644 --- a/src/java/org/apache/cassandra/config/DatabaseDescriptor.java +++ b/src/java/org/apache/cassandra/config/DatabaseDescriptor.java @@ -2685,6 +2685,24 @@ public static void setRpcTimeout(long timeOutInMillis) conf.request_timeout = new DurationSpec.LongMillisecondsBound(timeOutInMillis); } + public static long getGracefulDisconnectGracePeriod() + { + return conf.graceful_disconnect_grace_period.toMilliseconds(); + } + + public static void setGracefulDisconnectGracePeriod(long gracefulDisconnectGracePeriod) + { + if (gracefulDisconnectGracePeriod > 0) + conf.graceful_disconnect_grace_period = new DurationSpec.LongMillisecondsBound(gracefulDisconnectGracePeriod); + else + throw new IllegalArgumentException(String.format("{} <= 0, not allowed, non positive values not allowed", gracefulDisconnectGracePeriod)); + } + + public static boolean getGracefulDisconnectEnabled() + { + return conf.graceful_disconnect_enabled; + } + public static long getReadRpcTimeout(TimeUnit unit) { return conf.read_request_timeout.to(unit); diff --git a/src/java/org/apache/cassandra/metrics/ClientMetrics.java b/src/java/org/apache/cassandra/metrics/ClientMetrics.java index 3ca4885a5d71..370f9f8e93b3 100644 --- a/src/java/org/apache/cassandra/metrics/ClientMetrics.java +++ b/src/java/org/apache/cassandra/metrics/ClientMetrics.java @@ -71,6 +71,10 @@ public final class ClientMetrics @VisibleForTesting Gauge connectedNativeClients; + public AtomicInteger connectionsDraining = new AtomicInteger(); + + public Meter forcedDisconnects; + @VisibleForTesting Gauge encryptedConnectedNativeClients; @@ -173,6 +177,21 @@ public void markProtocolException() protocolException.mark(); } + public void incrementConnectionsDraining() + { + connectionsDraining.incrementAndGet(); + } + + public void decrementConnectionsDraining() + { + connectionsDraining.decrementAndGet(); + } + + public void markForcedDisconnect(int forceDisconnectedClients) + { + forcedDisconnects.mark(forceDisconnectedClients); + } + public void markSSLHandshakeException() { sslHandshakeException.mark(); @@ -197,6 +216,10 @@ public synchronized void init(Server servers) registerGauge("ClientsByProtocolVersion", "clientsByProtocolVersion", this::recentClientStats); registerGauge("RequestsSize", ClientResourceLimits::getCurrentGlobalUsage); + connectionsDraining = new AtomicInteger(); + registerGauge("ConnectionsDraining", connectionsDraining::get); + forcedDisconnects = registerMeter("ForcedDisconnects"); + CassandraReservoir ipUsageReservoir = ClientResourceLimits.ipUsageReservoir(); Metrics.register(factory.createMetricName("RequestsSizeByIpDistribution"), new OverrideHistogram(ipUsageReservoir) diff --git a/src/java/org/apache/cassandra/service/NativeTransportService.java b/src/java/org/apache/cassandra/service/NativeTransportService.java index 28817df396fc..894f8f782800 100644 --- a/src/java/org/apache/cassandra/service/NativeTransportService.java +++ b/src/java/org/apache/cassandra/service/NativeTransportService.java @@ -37,6 +37,7 @@ import io.netty.channel.EventLoopGroup; import io.netty.channel.epoll.Epoll; import io.netty.channel.epoll.EpollEventLoopGroup; +import io.netty.channel.group.ChannelGroup; import io.netty.channel.nio.NioEventLoopGroup; import io.netty.util.Version; @@ -163,6 +164,11 @@ Server getServer() return server; } + public ChannelGroup getChannelsSubscribedToGracefulDisconnect() + { + return server.getChannelsSubscribedToGracefulDisconnect(); + } + public void clearConnectionHistory() { server.clearConnectionHistory(); diff --git a/src/java/org/apache/cassandra/service/StorageService.java b/src/java/org/apache/cassandra/service/StorageService.java index b4c834355841..bb16af8d1da8 100644 --- a/src/java/org/apache/cassandra/service/StorageService.java +++ b/src/java/org/apache/cassandra/service/StorageService.java @@ -1199,6 +1199,27 @@ public long getRpcTimeout() return DatabaseDescriptor.getRpcTimeout(MILLISECONDS); } + + @Override + public boolean getGracefulDisconnectEnabled() + { + return DatabaseDescriptor.getGracefulDisconnectEnabled(); + } + + @Override + public void setGracefulDisconnectGracePeriod(long value) + { + if (value <= 0 && DatabaseDescriptor.getGracefulDisconnectEnabled()) + throw new IllegalArgumentException("Graceful disconnect grace period must be positive when graceful disconnect is enabled. Got " + value); + DatabaseDescriptor.setGracefulDisconnectGracePeriod(value); + } + + @Override + public long getGracefulDisconnectGracePeriod() + { + return DatabaseDescriptor.getGracefulDisconnectGracePeriod(); + } + public void setReadRpcTimeout(long value) { DatabaseDescriptor.setReadRpcTimeout(value); diff --git a/src/java/org/apache/cassandra/service/StorageServiceMBean.java b/src/java/org/apache/cassandra/service/StorageServiceMBean.java index abc868903127..2b63ecab1b39 100644 --- a/src/java/org/apache/cassandra/service/StorageServiceMBean.java +++ b/src/java/org/apache/cassandra/service/StorageServiceMBean.java @@ -795,6 +795,11 @@ default int upgradeSSTables(String keyspaceName, boolean excludeCurrentVersion, public void setRpcTimeout(long value); public long getRpcTimeout(); + public void setGracefulDisconnectGracePeriod(long value); + public long getGracefulDisconnectGracePeriod(); + + public boolean getGracefulDisconnectEnabled(); + public void setReadRpcTimeout(long value); public long getReadRpcTimeout(); diff --git a/src/java/org/apache/cassandra/tools/nodetool/Drain.java b/src/java/org/apache/cassandra/tools/nodetool/Drain.java index fa61b4598960..a71e7ffebfc1 100644 --- a/src/java/org/apache/cassandra/tools/nodetool/Drain.java +++ b/src/java/org/apache/cassandra/tools/nodetool/Drain.java @@ -34,7 +34,8 @@ public void execute(NodeProbe probe) try { probe.drain(); - } catch (IOException | InterruptedException | ExecutionException e) + } + catch (IOException | InterruptedException | ExecutionException e) { throw new RuntimeException("Error occurred during flushing", e); } diff --git a/src/java/org/apache/cassandra/transport/Event.java b/src/java/org/apache/cassandra/transport/Event.java index a1209a13eaf5..2d5a564f40b4 100644 --- a/src/java/org/apache/cassandra/transport/Event.java +++ b/src/java/org/apache/cassandra/transport/Event.java @@ -36,7 +36,8 @@ public enum Type TOPOLOGY_CHANGE(ProtocolVersion.V3), STATUS_CHANGE(ProtocolVersion.V3), SCHEMA_CHANGE(ProtocolVersion.V3), - TRACE_COMPLETE(ProtocolVersion.V4); + TRACE_COMPLETE(ProtocolVersion.V4), + GRACEFUL_DISCONNECT(ProtocolVersion.V5); public final ProtocolVersion minimumVersion; @@ -66,6 +67,8 @@ public static Event deserialize(ByteBuf cb, ProtocolVersion version) return StatusChange.deserializeEvent(cb, version); case SCHEMA_CHANGE: return SchemaChange.deserializeEvent(cb, version); + case GRACEFUL_DISCONNECT: + return GracefulDisconnect.deserializeEvent(cb, version); } throw new AssertionError(); } @@ -442,4 +445,26 @@ public boolean equals(Object other) && Objects.equal(argTypes, scc.argTypes); } } + + public static class GracefulDisconnect extends Event + { + public GracefulDisconnect() + { + super(Type.GRACEFUL_DISCONNECT); + } + + @Override + protected int eventSerializedSize(ProtocolVersion version) + { + return 0; + } + + @Override + protected void serializeEvent(ByteBuf dest, ProtocolVersion version) {} + + public static GracefulDisconnect deserializeEvent(ByteBuf cb, ProtocolVersion version) + { + return new GracefulDisconnect(); + } + } } diff --git a/src/java/org/apache/cassandra/transport/GracefulDisconnectLifecycle.java b/src/java/org/apache/cassandra/transport/GracefulDisconnectLifecycle.java new file mode 100644 index 000000000000..5479bb60c6e5 --- /dev/null +++ b/src/java/org/apache/cassandra/transport/GracefulDisconnectLifecycle.java @@ -0,0 +1,153 @@ +/* + * 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 org.apache.cassandra.transport; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; + +import org.apache.cassandra.config.DatabaseDescriptor; +import org.apache.cassandra.transport.messages.EventMessage; +import org.apache.cassandra.utils.concurrent.CountDownLatch; + +import io.netty.channel.Channel; +import io.netty.channel.group.ChannelGroup; + +import static org.apache.cassandra.transport.Dispatcher.EVENT_DISPATCHER; + +final class GracefulDisconnectLifecycle +{ + // Stable snapshot of subscribed channels. + private final List channels; + + // Number of channels still awaiting closure. + private final AtomicInteger remainingChannels; + + // Signals that the graceful disconnect phase is complete. + private final CountDownLatch completion; + + // Tracks the current lifecycle state. + private final AtomicReference state; + + private final long gracefulDisconnectGracePeriod; + + // Called when each channel is closed. + private final Consumer onChannelClosed; + + // Number of clients forcefully disconnected after the grace period. + private final AtomicInteger forcedDisconnects = new AtomicInteger(); + + // Best-effort upper bound on how long forced Channel#close() calls should + // take to actually complete once triggered. Not user-configurable — + // this is a safety margin against the drain hanging. + private long hardDeadlineBufferMillis = 5000; + + public GracefulDisconnectLifecycle(ChannelGroup channelGroup, + Consumer onChannelClosed) + { + this(channelGroup, + onChannelClosed, + DatabaseDescriptor.getGracefulDisconnectGracePeriod(), + 5000); + } + + GracefulDisconnectLifecycle(ChannelGroup channelGroup, + Consumer onChannelClosed, + long gracefulDisconnectGracePeriod, + long hardDeadlineBufferMillis) + { + channels = new ArrayList<>(channelGroup); + remainingChannels = new AtomicInteger(channels.size()); + completion = CountDownLatch.newCountDownLatch(1); + state = new AtomicReference<>(State.WAITING_FOR_CLIENTS); + this.gracefulDisconnectGracePeriod = gracefulDisconnectGracePeriod; + this.onChannelClosed = onChannelClosed; + this.hardDeadlineBufferMillis = hardDeadlineBufferMillis; + } + + int run() throws InterruptedException, TimeoutException + { + startGracefulDisconnect(); + if (completion.await(gracefulDisconnectGracePeriod, TimeUnit.MILLISECONDS)) return forcedDisconnects.get(); + onGracePeriodExpired(); + if (!completion.await(hardDeadlineBufferMillis, TimeUnit.MILLISECONDS)) + throw new TimeoutException("Graceful disconnect did not complete even after forced close, " + remainingChannels.get() + " channel(s) still open"); + return forcedDisconnects.get(); + } + + private void startGracefulDisconnect() + { + if (remainingChannels.get() == 0) + { + complete(); + return; + } + + EventMessage eventMessage = new EventMessage(new Event.GracefulDisconnect()); + channels.forEach(channel -> { + channel.closeFuture().addListener(future -> onChannelClosed(channel)); + Consumer dispatcher = channel.attr(EVENT_DISPATCHER).get(); + if (dispatcher != null) + dispatcher.accept(eventMessage); + }); + } + + private void complete() + { + State previous = state.getAndSet(State.COMPLETE); + if (previous != State.COMPLETE) + completion.decrement(); + } + + private void onChannelClosed(Channel channel) + { + onChannelClosed.accept(channel); + + if (remainingChannels.decrementAndGet() == 0) + complete(); + } + + private void onGracePeriodExpired() + { + if (!state.compareAndSet(State.WAITING_FOR_CLIENTS, State.FORCE_CLOSING)) + return; + + int actuallyClosed = 0; + for (Channel channel : channels) + { + if (channel.isOpen()) + { + channel.close(); + actuallyClosed++; + } + } + forcedDisconnects.set(actuallyClosed); + } + + private enum State + { + WAITING_FOR_CLIENTS, + FORCE_CLOSING, + COMPLETE + } +} \ No newline at end of file diff --git a/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java b/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java index 28299363fb7e..b3bb479a0b66 100644 --- a/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java +++ b/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java @@ -28,6 +28,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.apache.cassandra.config.DatabaseDescriptor; import org.apache.cassandra.cql3.QueryProcessor; import org.apache.cassandra.net.AsyncChannelPromise; import org.apache.cassandra.transport.ClientResourceLimits.Overload; @@ -91,6 +92,8 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List li supportedOptions.put(StartupMessage.CQL_VERSION, cqlVersions); supportedOptions.put(StartupMessage.COMPRESSION, compressions); supportedOptions.put(StartupMessage.PROTOCOL_VERSIONS, ProtocolVersion.supportedVersions()); + if (DatabaseDescriptor.getGracefulDisconnectEnabled()) + supportedOptions.put(StartupMessage.GRACEFUL_DISCONNECT, List.of("true")); SupportedMessage supported = new SupportedMessage(supportedOptions); outbound = supported.encode(inbound.header.version, inbound.header.streamId); ctx.writeAndFlush(outbound); diff --git a/src/java/org/apache/cassandra/transport/Server.java b/src/java/org/apache/cassandra/transport/Server.java index d7af40eb96f7..a8fcafb41f98 100644 --- a/src/java/org/apache/cassandra/transport/Server.java +++ b/src/java/org/apache/cassandra/transport/Server.java @@ -28,6 +28,7 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.LockSupport; import java.util.function.BooleanSupplier; @@ -44,6 +45,7 @@ import org.apache.cassandra.db.marshal.AbstractType; import org.apache.cassandra.db.marshal.UserType; import org.apache.cassandra.locator.InetAddressAndPort; +import org.apache.cassandra.metrics.ClientMetrics; import org.apache.cassandra.schema.KeyspaceMetadata; import org.apache.cassandra.schema.Schema; import org.apache.cassandra.schema.SchemaChangeListener; @@ -80,6 +82,7 @@ public class Server implements CassandraDaemon.Server private static final boolean useEpoll = NativeTransportService.useEpoll(); private final ConnectionTracker connectionTracker; + private Channel bindChannel; private final Connection.Factory connectionFactory = new Connection.Factory() { @@ -126,6 +129,11 @@ private Server (Builder builder) Schema.instance.registerListener(notifier); } + public ChannelGroup getChannelsSubscribedToGracefulDisconnect() + { + return connectionTracker.groups.get(Event.Type.GRACEFUL_DISCONNECT); + } + public void stop() { stop(false); @@ -133,8 +141,38 @@ public void stop() public void stop(boolean force) { - if (isRunning.compareAndSet(true, false)) - close(force); + if (isRunning.compareAndSet(true, false)) + { + if (!force && DatabaseDescriptor.getGracefulDisconnectEnabled()) + { + gracefulDisconnect(); + } + close(force); + } + } + + private void gracefulDisconnect() + { + stopAcceptingNewConnections(); + + ChannelGroup channelGroup = getChannelsSubscribedToGracefulDisconnect(); + ClientMetrics.instance.connectionsDraining.set(channelGroup.size()); + + try + { + int forcedDisconnects = new GracefulDisconnectLifecycle(channelGroup, + channel -> ClientMetrics.instance.decrementConnectionsDraining()).run(); + ClientMetrics.instance.markForcedDisconnect(forcedDisconnects); + } + catch (InterruptedException e) + { + logger.warn("Graceful disconnect interrupted", e); + Thread.currentThread().interrupt(); + } + catch (TimeoutException e) + { + logger.warn("Graceful disconnect timed out waiting for channels to close; proceeding with shutdown", e); + } } public boolean isRunning() @@ -149,6 +187,7 @@ public synchronized void start() // Configure the server. ChannelFuture bindFuture = pipelineConfigurator.initializeChannel(workerGroup, socket, connectionFactory); + bindChannel = bindFuture.channel(); if (!bindFuture.awaitUninterruptibly().isSuccess()) throw new IllegalStateException(String.format("Failed to bind port %d on %s.", socket.getPort(), socket.getAddress().getHostAddress()), bindFuture.cause()); @@ -157,6 +196,21 @@ public synchronized void start() isRunning.set(true); } + /** + * Permanently stops the native transport acceptor from accepting new connections. + * This does NOT reopen on its own — a full native transport restart is required + * to accept new connections again after this is called. + */ + public void stopAcceptingNewConnections() + { + if (bindChannel != null && bindChannel.isOpen()) + { + logger.info("Stopping native transport acceptor on {}", bindChannel.localAddress()); + // syncUninterruptibly ensures we wait for the port to actually close + bindChannel.close().syncUninterruptibly(); + } + } + public int countConnectedClients() { return connectionTracker.countConnectedClients(); @@ -330,6 +384,8 @@ public boolean isRunning() public void register(Event.Type type, Channel ch) { + if (type == Event.Type.GRACEFUL_DISCONNECT && !DatabaseDescriptor.getGracefulDisconnectEnabled()) + return; groups.get(type).add(ch); } diff --git a/src/java/org/apache/cassandra/transport/SimpleClient.java b/src/java/org/apache/cassandra/transport/SimpleClient.java index aee18ac227cb..65a5c9e08085 100644 --- a/src/java/org/apache/cassandra/transport/SimpleClient.java +++ b/src/java/org/apache/cassandra/transport/SimpleClient.java @@ -29,7 +29,7 @@ import java.util.Queue; import java.util.concurrent.BlockingQueue; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.concurrent.SynchronousQueue; // checkstyle: permit this import +import java.util.concurrent.SynchronousQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; @@ -82,7 +82,7 @@ import io.netty.handler.codec.MessageToMessageDecoder; import io.netty.handler.codec.MessageToMessageEncoder; import io.netty.handler.ssl.SslContext; -import io.netty.util.concurrent.Promise; // checkstyle: permit this import +import io.netty.util.concurrent.Promise; import io.netty.util.concurrent.PromiseCombiner; import io.netty.util.internal.logging.InternalLoggerFactory; import io.netty.util.internal.logging.Slf4JLoggerFactory; @@ -99,81 +99,21 @@ public class SimpleClient implements Closeable { public static final int TIMEOUT_SECONDS = 10; - - static - { - InternalLoggerFactory.setDefaultFactory(new Slf4JLoggerFactory()); - } - private static final Logger logger = LoggerFactory.getLogger(SimpleClient.class); - public final String host; public final int port; - private final EncryptionOptions.ClientEncryptionOptions encryptionOptions; - private final int largeMessageThreshold; - - protected final ResponseHandler responseHandler = new ResponseHandler(); + protected final ResponseHandler responseHandler = new ResponseHandler(this); protected final Connection.Tracker tracker = new ConnectionTracker(); protected final ProtocolVersion version; + private final AtomicBoolean draining = new AtomicBoolean(false); + private final EncryptionOptions.ClientEncryptionOptions encryptionOptions; + private final int largeMessageThreshold; // We don't track connection really, so we don't need one Connection per channel protected Connection connection; protected Bootstrap bootstrap; protected Channel channel; - protected ChannelFuture lastWriteFuture; - protected String compression; - - public static class Builder - { - private final String host; - private final int port; - private EncryptionOptions.ClientEncryptionOptions encryptionOptions = new EncryptionOptions.ClientEncryptionOptions(); - private ProtocolVersion version = ProtocolVersion.CURRENT; - private boolean useBeta = false; - private int largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE; - - private Builder(String host, int port) - { - this.host = host; - this.port = port; - } - - public Builder encryption(EncryptionOptions.ClientEncryptionOptions options) - { - this.encryptionOptions = options; - return this; - } - - public Builder useBeta() - { - this.useBeta = true; - return this; - } - - public Builder protocolVersion(ProtocolVersion version) - { - this.version = version; - return this; - } - - public Builder largeMessageThreshold(int bytes) - { - largeMessageThreshold = bytes; - return this; - } - - public SimpleClient build() - { - if (version.isBeta() && !useBeta) - throw new IllegalArgumentException(String.format("Beta version of server used (%s), but USE_BETA flag is not set", version)); - return new SimpleClient(this); - } - } - - public static Builder builder(String host, int port) - { - return new Builder(host, port); - } + private volatile ChannelFuture lastWriteFuture; private SimpleClient(Builder builder) { @@ -208,9 +148,7 @@ public SimpleClient(String host, int port, ProtocolVersion version, boolean useB this.version = version; this.encryptionOptions = encryptionOptions.applyConfig(); - this.largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE - - Math.max(FrameEncoderCrc.HEADER_AND_TRAILER_LENGTH, - FrameEncoderLZ4.HEADER_AND_TRAILER_LENGTH); + this.largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE - Math.max(FrameEncoderCrc.HEADER_AND_TRAILER_LENGTH, FrameEncoderLZ4.HEADER_AND_TRAILER_LENGTH); } public SimpleClient(String host, int port) @@ -218,6 +156,23 @@ public SimpleClient(String host, int port) this(host, port, new EncryptionOptions.ClientEncryptionOptions()); } + public static Builder builder(String host, int port) + { + return new Builder(host, port); + } + + /** + * The stream id to frame an outbound client request with. SimpleClient carries the intended id on the + * request's (dummy) source envelope (see {@link #execute(List)} and callers that pipeline requests). + * When no source has been assigned, we fall back to 0, which is sufficient for the non-pipelined path + * where only a single request is ever in flight. + */ + private static int outboundStreamId(Message message) + { + Envelope source = message.getSource(); + return source == null ? 0 : source.header.streamId; + } + public SimpleClient connect(boolean useCompression) throws IOException { return connect(useCompression, false); @@ -229,8 +184,7 @@ public SimpleClient connect(boolean useCompression, boolean throwOnOverload) thr Map options = new HashMap<>(); options.put(StartupMessage.CQL_VERSION, "3.0.0"); - if (throwOnOverload) - options.put(StartupMessage.THROW_ON_OVERLOAD, "1"); + if (throwOnOverload) options.put(StartupMessage.THROW_ON_OVERLOAD, "1"); connection.setThrowOnOverload(throwOnOverload); if (useCompression) @@ -239,7 +193,6 @@ public SimpleClient connect(boolean useCompression, boolean throwOnOverload) thr connection.setCompressor(Compressor.LZ4Compressor.instance); } execute(new StartupMessage(options)); - return this; } @@ -252,13 +205,10 @@ public void setEventHandler(EventHandler eventHandler) void establishConnection() throws IOException { // Configure the client. - bootstrap = new Bootstrap() - .group(new NioEventLoopGroup(new NamedThreadFactory("SimpleClient-nioEventLoopGroup"))) - .channel(io.netty.channel.socket.nio.NioSocketChannel.class) - .option(ChannelOption.TCP_NODELAY, true); + bootstrap = new Bootstrap().group(new NioEventLoopGroup(new NamedThreadFactory("SimpleClient-nioEventLoopGroup"))).channel(io.netty.channel.socket.nio.NioSocketChannel.class).option(ChannelOption.TCP_NODELAY, true); // Configure the pipeline factory. - if(encryptionOptions.getEnabled()) + if (encryptionOptions.getEnabled()) { bootstrap.handler(new SecureInitializer(largeMessageThreshold)); } @@ -279,35 +229,34 @@ void establishConnection() throws IOException public ResultMessage execute(String query, ConsistencyLevel consistency) { - return execute(query, Collections.emptyList(), consistency); + return execute(query, Collections.emptyList(), consistency); } public ResultMessage execute(String query, List values, ConsistencyLevel consistencyLevel) { Message.Response msg = execute(new QueryMessage(query, QueryOptions.forInternalCalls(consistencyLevel, values))); assert msg instanceof ResultMessage; - return (ResultMessage)msg; + return (ResultMessage) msg; } public ResultMessage.Prepared prepare(String query) { Message.Response msg = execute(new PrepareMessage(query, null)); assert msg instanceof ResultMessage.Prepared; - return (ResultMessage.Prepared)msg; + return (ResultMessage.Prepared) msg; } public ResultMessage executePrepared(ResultMessage.Prepared prepared, List values, ConsistencyLevel consistency) { Message.Response msg = execute(new ExecuteMessage(prepared.statementId, prepared.resultMetadataId, QueryOptions.forInternalCalls(consistency, values))); assert msg instanceof ResultMessage; - return (ResultMessage)msg; + return (ResultMessage) msg; } public void close() { // Wait until all messages are flushed before closing the channel. - if (lastWriteFuture != null) - lastWriteFuture.awaitUninterruptibly(); + if (lastWriteFuture != null) lastWriteFuture.awaitUninterruptibly(); // Close the connection. Make sure the close operation ends because // all I/O operations are asynchronous in Netty. @@ -324,15 +273,18 @@ public Message.Response execute(Message.Request request) public Message.Response execute(Message.Request request, boolean throwOnErrorResponse) { + if (draining.get()) + { + throw new RuntimeException("Connection is draining (GRACEFUL_DISCONNECT received)"); + } try { request.attach(connection); lastWriteFuture = channel.writeAndFlush(Collections.singletonList(request)); Message.Response msg = responseHandler.responses.poll(TIMEOUT_SECONDS, TimeUnit.SECONDS); - if (msg == null) - throw new RuntimeException("timeout"); + if (msg == null) throw new RuntimeException("timeout"); if (throwOnErrorResponse && msg instanceof ErrorMessage) - throw new RuntimeException((Throwable)((ErrorMessage)msg).error); + throw new RuntimeException((Throwable) ((ErrorMessage) msg).error); return msg; } catch (InterruptedException e) @@ -361,10 +313,8 @@ public Map execute(List requ for (int i = 0; i < requests.size(); i++) { Message.Response msg = responseHandler.responses.poll(deadline - currentTimeMillis(), TimeUnit.MILLISECONDS); - if (msg == null) - throw new RuntimeException("timeout"); - if (msg instanceof ErrorMessage) - throw new RuntimeException((Throwable) ((ErrorMessage) msg).error); + if (msg == null) throw new RuntimeException("timeout"); + if (msg instanceof ErrorMessage) throw new RuntimeException((Throwable) ((ErrorMessage) msg).error); rrMap.put(requests.get(msg.getSource().header.streamId), msg); } } @@ -383,16 +333,15 @@ public Map execute(List requ } } - /** - * The stream id to frame an outbound client request with. SimpleClient carries the intended id on the - * request's (dummy) source envelope (see {@link #execute(List)} and callers that pipeline requests). - * When no source has been assigned, we fall back to 0, which is sufficient for the non-pipelined path - * where only a single request is ever in flight. - */ - private static int outboundStreamId(Message message) + private void handleGracefulDisconnect() { - Envelope source = message.getSource(); - return source == null ? 0 : source.header.streamId; + draining.set(true); + channel.eventLoop().execute(() -> { + ChannelFuture writeFuture = lastWriteFuture; + ChannelFuture closeAfter = (writeFuture != null) ? writeFuture : channel.newSucceededFuture(); + closeAfter.addListener(f -> channel.close()); + }); + channel.closeFuture().addListener(f -> bootstrap.group().shutdownGracefully()); } public interface EventHandler @@ -400,6 +349,53 @@ public interface EventHandler void onEvent(Event event); } + public static class Builder + { + private final String host; + private final int port; + private EncryptionOptions.ClientEncryptionOptions encryptionOptions = new EncryptionOptions.ClientEncryptionOptions(); + private ProtocolVersion version = ProtocolVersion.CURRENT; + private boolean useBeta = false; + private int largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE; + + private Builder(String host, int port) + { + this.host = host; + this.port = port; + } + + public Builder encryption(EncryptionOptions.ClientEncryptionOptions options) + { + this.encryptionOptions = options; + return this; + } + + public Builder useBeta() + { + this.useBeta = true; + return this; + } + + public Builder protocolVersion(ProtocolVersion version) + { + this.version = version; + return this; + } + + public Builder largeMessageThreshold(int bytes) + { + largeMessageThreshold = bytes; + return this; + } + + public SimpleClient build() + { + if (version.isBeta() && !useBeta) + throw new IllegalArgumentException(String.format("Beta version of server used (%s), but USE_BETA flag is not set", version)); + return new SimpleClient(this); + } + } + public static class SimpleEventHandler implements EventHandler { public final BlockingQueue queue = newBlockingQueue(); @@ -412,7 +408,9 @@ public void onEvent(Event event) private static class ConnectionTracker implements Connection.Tracker { - public void addConnection(Channel ch, Connection connection) {} + public void addConnection(Channel ch, Connection connection) + { + } @Override public boolean isRunning() @@ -423,19 +421,19 @@ public boolean isRunning() private static class HandlerNames { - private static final String ENVELOPE_DECODER = "envelopeDecoder"; - private static final String ENVELOPE_ENCODER = "envelopeEncoder"; - private static final String COMPRESSOR = "compressor"; - private static final String DECOMPRESSOR = "decompressor"; - private static final String MESSAGE_DECODER = "messageDecoder"; - private static final String MESSAGE_ENCODER = "messageEncoder"; + private static final String ENVELOPE_DECODER = "envelopeDecoder"; + private static final String ENVELOPE_ENCODER = "envelopeEncoder"; + private static final String COMPRESSOR = "compressor"; + private static final String DECOMPRESSOR = "decompressor"; + private static final String MESSAGE_DECODER = "messageDecoder"; + private static final String MESSAGE_ENCODER = "messageEncoder"; - private static final String INITIAL_HANDLER = "intitialHandler"; - private static final String RESPONSE_HANDLER = "responseHandler"; + private static final String INITIAL_HANDLER = "intitialHandler"; + private static final String RESPONSE_HANDLER = "responseHandler"; - private static final String FRAME_DECODER = "frameDecoder"; - private static final String FRAME_ENCODER = "frameEncoder"; - private static final String PROCESSOR = "processor"; + private static final String FRAME_DECODER = "frameDecoder"; + private static final String FRAME_ENCODER = "frameEncoder"; + private static final String PROCESSOR = "processor"; } private static class InitialHandler extends MessageToMessageDecoder @@ -443,6 +441,7 @@ private static class InitialHandler extends MessageToMessageDecoder final ProtocolVersion version; final ResponseHandler responseHandler; final int largeMessageThreshold; + InitialHandler(ProtocolVersion version, ResponseHandler responseHandler, int largeMessageThreshold) { this.version = version; @@ -453,7 +452,7 @@ private static class InitialHandler extends MessageToMessageDecoder @Override protected void decode(ChannelHandlerContext ctx, Envelope request, List results) { - switch(request.header.type) + switch (request.header.type) { case READY: case AUTHENTICATE: @@ -475,9 +474,7 @@ protected void decode(ChannelHandlerContext ctx, Envelope request, List results.add(request); break; default: - throw new ProtocolException(String.format("Unexpected %s request expecting " + - "READY, AUTHENTICATE, ERROR or SUPPORTED", - request.header.type)); + throw new ProtocolException(String.format("Unexpected %s request expecting " + "READY, AUTHENTICATE, ERROR or SUPPORTED", request.header.type)); } } @@ -557,34 +554,22 @@ public void release() } }; - CQLMessageHandler processor = - new CQLMessageHandler(ctx.channel(), - null, - version, - frameDecoder, - envelopeDecoder, - messageDecoder, - responseConsumer, - payloadAllocator, - queueCapacity, - QueueBackpressure.NO_OP, - resources, - handler -> {}, - errorHandler, - ctx.channel().attr(Connection.attributeKey).get().isThrowOnOverload()) + CQLMessageHandler processor = new CQLMessageHandler(ctx.channel(), null, version, frameDecoder, envelopeDecoder, messageDecoder, responseConsumer, payloadAllocator, queueCapacity, QueueBackpressure.NO_OP, resources, handler -> { + }, errorHandler, ctx.channel().attr(Connection.attributeKey).get().isThrowOnOverload()) + { + protected boolean processRequest(Envelope request, Overload overload) { - protected boolean processRequest(Envelope request, Overload overload) - { - boolean continueProcessing = super.processRequest(request, overload); - releaseCapacity(Ints.checkedCast(request.header.bodySizeInBytes)); - return continueProcessing; - } - }; + boolean continueProcessing = super.processRequest(request, overload); + releaseCapacity(Ints.checkedCast(request.header.bodySizeInBytes)); + return continueProcessing; + } + }; pipeline.addLast(HandlerNames.FRAME_DECODER, frameDecoder); pipeline.addLast(HandlerNames.FRAME_ENCODER, frameEncoder); pipeline.addLast(HandlerNames.PROCESSOR, processor); - pipeline.addLast(HandlerNames.MESSAGE_ENCODER, new ChannelOutboundHandlerAdapter() { + pipeline.addLast(HandlerNames.MESSAGE_ENCODER, new ChannelOutboundHandlerAdapter() + { public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception { @@ -612,20 +597,16 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) private FrameDecoder frameDecoder(ChannelHandlerContext ctx, BufferPoolAllocator allocator) { Connection conn = ctx.channel().attr(Connection.attributeKey).get(); - if (conn.getCompressor() == null) - return FrameDecoderCrc.create(allocator); - if (conn.getCompressor() instanceof Compressor.LZ4Compressor) - return FrameDecoderLZ4.fast(allocator); + if (conn.getCompressor() == null) return FrameDecoderCrc.create(allocator); + if (conn.getCompressor() instanceof Compressor.LZ4Compressor) return FrameDecoderLZ4.fast(allocator); throw new ProtocolException("Unsupported compressor: " + conn.getCompressor().getClass().getCanonicalName()); } private FrameEncoder frameEncoder(ChannelHandlerContext ctx) { Connection conn = ctx.channel().attr(Connection.attributeKey).get(); - if (conn.getCompressor() == null) - return FrameEncoderCrc.instance; - if (conn.getCompressor() instanceof Compressor.LZ4Compressor) - return FrameEncoderLZ4.fastInstance; + if (conn.getCompressor() == null) return FrameEncoderCrc.instance; + if (conn.getCompressor() instanceof Compressor.LZ4Compressor) return FrameEncoderLZ4.fastInstance; throw new ProtocolException("Unsupported compressor: " + conn.getCompressor().getClass().getCanonicalName()); } @@ -640,10 +621,13 @@ private void configureLegacyPipeline(ChannelHandlerContext ctx) } @ChannelHandler.Sharable - static class MessageBatchEncoder extends MessageToMessageEncoder> + static class MessageBatchEncoder extends MessageToMessageEncoder> { public static final MessageBatchEncoder instance = new MessageBatchEncoder(); - private MessageBatchEncoder(){} + + private MessageBatchEncoder() + { + } public void encode(ChannelHandlerContext ctx, List messages, List results) { @@ -656,52 +640,17 @@ public void encode(ChannelHandlerContext ctx, List messages, List - { - private int largeMessageThreshold; - Initializer(int largeMessageThreshold) - { - this.largeMessageThreshold = largeMessageThreshold; - } - - protected void initChannel(Channel channel) throws Exception - { - connection = new Connection(channel, version, tracker); - channel.attr(Connection.attributeKey).set(connection); - - ChannelPipeline pipeline = channel.pipeline(); -// pipeline.addLast("debug", new LoggingHandler(LogLevel.INFO)); - pipeline.addLast(HandlerNames.ENVELOPE_DECODER, new Envelope.Decoder()); - pipeline.addLast(HandlerNames.ENVELOPE_ENCODER, Envelope.Encoder.instance); - pipeline.addLast(HandlerNames.INITIAL_HANDLER, new InitialHandler(version, responseHandler, largeMessageThreshold)); - pipeline.addLast(HandlerNames.MESSAGE_DECODER, PreV5Handlers.ProtocolDecoder.instance); - pipeline.addLast(HandlerNames.MESSAGE_ENCODER, MessageBatchEncoder.instance); - pipeline.addLast(HandlerNames.RESPONSE_HANDLER, responseHandler); - } - } - - private class SecureInitializer extends Initializer - { - SecureInitializer(int largeMessageThreshold) - { - super(largeMessageThreshold); - } - - protected void initChannel(Channel channel) throws Exception - { - super.initChannel(channel); - SslContext sslContext = SSLFactory.getOrCreateSslContext(encryptionOptions, encryptionOptions.getClientAuth(), - ISslContextFactory.SocketType.CLIENT, SSL_FACTORY_CONTEXT_DESCRIPTION); - InetSocketAddress peer = encryptionOptions.require_endpoint_verification ? new InetSocketAddress(host, port) : null; - channel.pipeline().addFirst("ssl", newSslHandler(channel, sslContext, peer)); - } - } - @ChannelHandler.Sharable static class ResponseHandler extends SimpleChannelInboundHandler { + public final BlockingQueue responses = new SynchronousQueue<>(true); + private final SimpleClient client; public EventHandler eventHandler; + public ResponseHandler(SimpleClient client) + { + this.client = client; + } @Override public void channelRead0(ChannelHandlerContext ctx, Message.Response r) @@ -719,11 +668,19 @@ public void handleResponse(Channel channel, Message.Response r) if (r instanceof EventMessage) { - if (eventHandler != null) - eventHandler.onEvent(((EventMessage) r).event); + Event event = ((EventMessage) r).event; + + if (event.type == Event.Type.GRACEFUL_DISCONNECT) + { + logger.info("Received GRACEFUL_DISCONNECT. Entering draining mode."); + if (eventHandler != null) eventHandler.onEvent(event); + client.handleGracefulDisconnect(); + return; + } + + if (eventHandler != null) eventHandler.onEvent(event); } - else - responses.put(r); + else responses.put(r); } catch (InterruptedException e) { @@ -776,15 +733,13 @@ public void enqueue(Envelope message) public void releaseAll() { Envelope e; - while ((e = outbound.poll()) != null) - e.release(); + while ((e = outbound.poll()) != null) e.release(); } public void schedule(ChannelHandlerContext ctx) { if (scheduled.compareAndSet(false, true)) - ctx.executor().scheduleAtFixedRate(() -> maybeWrite(ctx, ctx.voidPromise()), - 10, 10, TimeUnit.MILLISECONDS); + ctx.executor().scheduleAtFixedRate(() -> maybeWrite(ctx, ctx.voidPromise()), 10, 10, TimeUnit.MILLISECONDS); } public void maybeWrite(ChannelHandlerContext ctx, Promise promise) @@ -848,10 +803,8 @@ private ChannelFuture flushBuffer(ChannelHandlerContext ctx, List mess private FrameEncoder.Payload allocate(int size, boolean selfContained) { - FrameEncoder.Payload payload = frameEncoder.allocator() - .allocate(selfContained, Math.min(size, largeMessageThreshold)); - if (size >= largeMessageThreshold) - payload.buffer.limit(largeMessageThreshold); + FrameEncoder.Payload payload = frameEncoder.allocator().allocate(selfContained, Math.min(size, largeMessageThreshold)); + if (size >= largeMessageThreshold) payload.buffer.limit(largeMessageThreshold); return payload; } @@ -870,8 +823,7 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) buf = payload.buffer; // BufferPool may give us a buffer larger than we asked for. // FrameEncoder may object if buffer.remaining is >= MAX_SIZE. - if (payloadSize >= largeMessageThreshold) - buf.limit(largeMessageThreshold); + if (payloadSize >= largeMessageThreshold) buf.limit(largeMessageThreshold); if (firstFrame) { @@ -880,8 +832,7 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) } int remaining = Math.min(buf.remaining(), f.body.readableBytes()); - if (remaining > 0) - buf.put(f.body.slice(f.body.readerIndex(), remaining).nioBuffer()); + if (remaining > 0) buf.put(f.body.slice(f.body.readerIndex(), remaining).nioBuffer()); f.body.readerIndex(f.body.readerIndex() + remaining); payload.finish(); @@ -891,12 +842,57 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) promise.addListener(result -> { if (!result.isSuccess()) logger.warn("Failed to send frame of large message, size: " + remaining, result.cause()); - else - logger.trace("Sent frame of large message, size: {}", remaining); + else logger.trace("Sent frame of large message, size: {}", remaining); }); } f.release(); return futures.toArray(EMPTY_FUTURES_ARRAY); } } -} + + private class Initializer extends ChannelInitializer + { + private final int largeMessageThreshold; + + Initializer(int largeMessageThreshold) + { + this.largeMessageThreshold = largeMessageThreshold; + } + + protected void initChannel(Channel channel) throws Exception + { + connection = new Connection(channel, version, tracker); + channel.attr(Connection.attributeKey).set(connection); + + ChannelPipeline pipeline = channel.pipeline(); +// pipeline.addLast("debug", new LoggingHandler(LogLevel.INFO)); + pipeline.addLast(HandlerNames.ENVELOPE_DECODER, new Envelope.Decoder()); + pipeline.addLast(HandlerNames.ENVELOPE_ENCODER, Envelope.Encoder.instance); + pipeline.addLast(HandlerNames.INITIAL_HANDLER, new InitialHandler(version, responseHandler, largeMessageThreshold)); + pipeline.addLast(HandlerNames.MESSAGE_DECODER, PreV5Handlers.ProtocolDecoder.instance); + pipeline.addLast(HandlerNames.MESSAGE_ENCODER, MessageBatchEncoder.instance); + pipeline.addLast(HandlerNames.RESPONSE_HANDLER, responseHandler); + } + } + + private class SecureInitializer extends Initializer + { + SecureInitializer(int largeMessageThreshold) + { + super(largeMessageThreshold); + } + + protected void initChannel(Channel channel) throws Exception + { + super.initChannel(channel); + SslContext sslContext = SSLFactory.getOrCreateSslContext(encryptionOptions, encryptionOptions.getClientAuth(), ISslContextFactory.SocketType.CLIENT, SSL_FACTORY_CONTEXT_DESCRIPTION); + InetSocketAddress peer = encryptionOptions.require_endpoint_verification ? new InetSocketAddress(host, port) : null; + channel.pipeline().addFirst("ssl", newSslHandler(channel, sslContext, peer)); + } + } + + static + { + InternalLoggerFactory.setDefaultFactory(new Slf4JLoggerFactory()); + } +} \ No newline at end of file diff --git a/src/java/org/apache/cassandra/transport/messages/OptionsMessage.java b/src/java/org/apache/cassandra/transport/messages/OptionsMessage.java index 77a76e336897..ee598afe1e9a 100644 --- a/src/java/org/apache/cassandra/transport/messages/OptionsMessage.java +++ b/src/java/org/apache/cassandra/transport/messages/OptionsMessage.java @@ -22,6 +22,7 @@ import java.util.List; import java.util.Map; +import org.apache.cassandra.config.DatabaseDescriptor; import org.apache.cassandra.cql3.QueryProcessor; import org.apache.cassandra.service.QueryState; import org.apache.cassandra.transport.Compressor; @@ -73,6 +74,8 @@ protected Message.Response execute(QueryState state, Dispatcher.RequestTime requ Map> supported = new HashMap>(); supported.put(StartupMessage.CQL_VERSION, cqlVersions); supported.put(StartupMessage.COMPRESSION, compressions); + if (DatabaseDescriptor.getGracefulDisconnectEnabled()) + supported.put(StartupMessage.GRACEFUL_DISCONNECT, List.of("true")); supported.put(StartupMessage.PROTOCOL_VERSIONS, ProtocolVersion.supportedVersions()); return new SupportedMessage(supported); diff --git a/src/java/org/apache/cassandra/transport/messages/StartupMessage.java b/src/java/org/apache/cassandra/transport/messages/StartupMessage.java index c76db6826e30..acdb1cf6f92e 100644 --- a/src/java/org/apache/cassandra/transport/messages/StartupMessage.java +++ b/src/java/org/apache/cassandra/transport/messages/StartupMessage.java @@ -51,6 +51,7 @@ public class StartupMessage extends Message.Request public static final String DRIVER_NAME = "DRIVER_NAME"; public static final String DRIVER_VERSION = "DRIVER_VERSION"; public static final String THROW_ON_OVERLOAD = "THROW_ON_OVERLOAD"; + public static final String GRACEFUL_DISCONNECT = "GRACEFUL_DISCONNECT"; public static final Message.Codec codec = new Message.Codec() { diff --git a/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java new file mode 100644 index 000000000000..792a7db28b69 --- /dev/null +++ b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java @@ -0,0 +1,334 @@ +/* + * 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 org.apache.cassandra.distributed.test; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.util.Collections; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +import org.junit.BeforeClass; +import org.junit.Test; + +import org.apache.cassandra.config.DatabaseDescriptor; +import org.apache.cassandra.distributed.Cluster; +import org.apache.cassandra.distributed.api.Feature; +import org.apache.cassandra.service.CassandraDaemon; +import org.apache.cassandra.transport.Event; +import org.apache.cassandra.transport.Message; +import org.apache.cassandra.transport.ProtocolVersion; +import org.apache.cassandra.transport.SimpleClient; +import org.apache.cassandra.transport.messages.OptionsMessage; +import org.apache.cassandra.transport.messages.ReadyMessage; +import org.apache.cassandra.transport.messages.RegisterMessage; +import org.apache.cassandra.transport.messages.StartupMessage; +import org.apache.cassandra.transport.messages.SupportedMessage; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +public class GracefulDisconnectTest +{ + + @BeforeClass + public static void setUp() throws IOException + { + DatabaseDescriptor.daemonInitialization(); + } + + public Cluster buildCluster(int nodeCount, boolean gracefulDisconnectEnabled) throws IOException + { + return Cluster + .build(nodeCount) + .withConfig(config -> + config + .with(Feature.NATIVE_PROTOCOL, Feature.GOSSIP) + .set("graceful_disconnect_enabled", gracefulDisconnectEnabled)) + .start(); + } + + @Test + public void testGracefulDisconnectAdvertisedWhenEnabled() throws IOException + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042).build()) + { + client.connect(false); + Message.Response response = client.execute(new OptionsMessage()); + SupportedMessage supported = (SupportedMessage) response; + + assertThat(supported.supported.containsKey(StartupMessage.GRACEFUL_DISCONNECT)) + .as("GRACEFUL_DISCONNECT should be advertised in SUPPORTED when enabled") + .isTrue(); + } + } + } + + @Test + public void testGracefulDisconnectDoesNotAdvertisedWhenNotEnabled() throws IOException + { + try (Cluster cluster = buildCluster(1, false)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042).build()) + { + client.connect(false); + SupportedMessage supported = (SupportedMessage) client.execute(new OptionsMessage()); + + assertThat(supported.supported.containsKey(StartupMessage.GRACEFUL_DISCONNECT)) + .as("GRACEFUL_DISCONNECT should NOT be advertised in SUPPORTED when disabled") + .isFalse(); + } + } + } + + @Test + public void testSubscriptionViaREGISTER() throws IOException + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V5) + .build()) + { + client.connect(false); + Message.Response response = client.execute( + new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + assertThat(response).isInstanceOf(ReadyMessage.class); + + int subscribedCount = cluster.get(1).callOnInstance(() -> + CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect() + .size()); + + assertThat(subscribedCount) + .as("One channel should be subscribed to GRACEFUL_DISCONNECT") + .isEqualTo(1); + } + } + } + + @Test + public void testRegisterGracefulDisconnectRejectedOnV4() throws IOException + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042).protocolVersion(ProtocolVersion.V4).build()) + { + client.connect(false); + + assertThatThrownBy(() -> client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT)))) + .hasCauseInstanceOf(org.apache.cassandra.transport.ProtocolException.class); + + int subscribedCount = cluster.get(1).callOnInstance(() -> + CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect() + .size()); + + assertThat(subscribedCount) + .as("V4 client should not be able to subscribe") + .isEqualTo(0); + } + } + } + + @Test + public void testNonSubscribedClientDoesNotReceiveEvent() throws IOException + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V4) + .build()) + { + client.connect(false); + int subscribedCount = cluster.get(1).callOnInstance(() -> + CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect() + .size()); + assertThat(subscribedCount).isEqualTo(0); + } + } + } + + @Test + public void testServerStopsAcceptingNewConnectionsOnDrain() throws IOException + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042).build()) + { + client.connect(false); + cluster.get(1).nodetool("drain"); + + assertThatThrownBy(() -> { + SimpleClient newClient = SimpleClient.builder(nativeAddr.getHostString(), 9042).build(); + newClient.connect(false); + }).as("Server should reject new connections after stopAcceptingNewConnections") + .isNotNull(); + cluster.get(1).shutdown(); + } + } + } + + @Test + public void testNoEventEmittedWhenDisabled() throws IOException + { + try (Cluster cluster = buildCluster(1, false)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042).build()) + { + client.connect(false); + client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + int subscribedCount = cluster.get(1).callOnInstance(() -> + CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect() + .size()); + + assertThat(subscribedCount).isEqualTo(0); + + boolean disabled = cluster.get(1).callOnInstance(() -> !DatabaseDescriptor.getGracefulDisconnectEnabled()); + assertThat(disabled).isTrue(); + } + } + } + + @Test + public void testDrainProceedsImmediatelyWithNoSubscribedConnections() throws IOException + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V4) + .build()) + { + client.connect(false); + cluster.get(1).nodetoolResult("drain").asserts().success(); + cluster.get(1).shutdown(); + } + } + } + + @Test + public void testMultipleConnectionsCanSubscribe() throws IOException + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client1 = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V5) + .build(); + SimpleClient client2 = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V5) + .build()) + { + client1.connect(false); + client2.connect(false); + + client1.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + client2.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + int subscribedCount = cluster.get(1).callOnInstance(() -> + CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect() + .size()); + + assertThat(subscribedCount).isEqualTo(2); + } + } + } + + @Test + public void testInFlightQueryCompletesDuringDrain() throws Exception + { + try (Cluster cluster = buildCluster(1, true)) + { + cluster.schemaChange("CREATE KEYSPACE ks WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}"); + cluster.schemaChange("CREATE TABLE ks.tbl (id int PRIMARY KEY, val text)"); + cluster.get(1).executeInternal("INSERT INTO ks.tbl (id, val) VALUES (1, 'test_val')"); + + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V5) + .build()) + { + client.connect(false); + client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + CompletableFuture drainFuture = CompletableFuture.runAsync(() -> { + cluster.get(1).nodetool("drain"); + }); + + Message.Response response = client.execute(new org.apache.cassandra.transport.messages.QueryMessage( + "SELECT val FROM ks.tbl WHERE id = 1", + org.apache.cassandra.cql3.QueryOptions.DEFAULT + )); + + assertThat(response) + .as("Query must succeed during drain without connection drop or timeout") + .isInstanceOf(org.apache.cassandra.transport.messages.ResultMessage.Rows.class); + + drainFuture.get(10, TimeUnit.SECONDS); + + assertThatThrownBy(() -> { + SimpleClient newClient = SimpleClient.builder(nativeAddr.getHostString(), 9042).build(); + newClient.connect(false); + }).isNotNull(); + } + } + } + + @Test + public void testCooperativeClientDisconnectsWithinGracePeriod() throws Exception + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + + SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V5) + .build(); + client.connect(false); + client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + cluster.get(1).nodetool("drain"); + + long forcedCount = cluster.get(1).callOnInstance(() -> + org.apache.cassandra.metrics.ClientMetrics.instance.forcedDisconnects.getCount() + ); + assertThat(forcedCount).as("ForcedDisconnects metric should be 0 for a cooperative client").isEqualTo(0); + } + } +} \ No newline at end of file diff --git a/test/unit/org/apache/cassandra/config/DatabaseDescriptorTest.java b/test/unit/org/apache/cassandra/config/DatabaseDescriptorTest.java index c3bdd00e21ae..c1b2ad8fbd88 100644 --- a/test/unit/org/apache/cassandra/config/DatabaseDescriptorTest.java +++ b/test/unit/org/apache/cassandra/config/DatabaseDescriptorTest.java @@ -817,6 +817,30 @@ public void testRowIndexSizeWarnEnabledAbortDisabled() DatabaseDescriptor.applyThresholdsValidations(conf); } + @Test + public void testGracefulDisconnectEnabled() + { + Assertions.assertThat(DatabaseDescriptor.getGracefulDisconnectEnabled()).as("Graceful disconnect should be disabled by default").isFalse(); + } + + @Test + public void testGracefulDisconnectGracePeriod() + { + long originalValue = DatabaseDescriptor.getGracefulDisconnectGracePeriod(); + Assertions.assertThat(originalValue).as("Default value of graceful_disconnect_grace_period must be 5000").isEqualTo(5000); + try + { + DatabaseDescriptor.setGracefulDisconnectGracePeriod(3000); + Assertions.assertThat(DatabaseDescriptor.getGracefulDisconnectGracePeriod()).as("graceful_disconnect_grace_period should be updated to 3000").isEqualTo(3000); + Assertions.assertThatThrownBy(() -> DatabaseDescriptor.setGracefulDisconnectGracePeriod(0)).isInstanceOf(IllegalArgumentException.class); + Assertions.assertThatThrownBy(() -> DatabaseDescriptor.setGracefulDisconnectGracePeriod(-10)).isInstanceOf(IllegalArgumentException.class); + } + finally + { + DatabaseDescriptor.setGracefulDisconnectGracePeriod(originalValue); + } + } + @Test public void testRowIndexSizeAbortEnabledWarnDisabled() { diff --git a/test/unit/org/apache/cassandra/metrics/ClientMetricsTest.java b/test/unit/org/apache/cassandra/metrics/ClientMetricsTest.java index 1ecae7f67cd9..2f8131a05ecd 100644 --- a/test/unit/org/apache/cassandra/metrics/ClientMetricsTest.java +++ b/test/unit/org/apache/cassandra/metrics/ClientMetricsTest.java @@ -239,4 +239,20 @@ public void testConnectedClientsAndAuthMetrics() throws SSLException assertEquals(0, clientMetrics.encryptedConnectedNativeClients.getValue().intValue()); assertEquals(0, passwordConnections.getValue().intValue()); } + + @Test + public void testGracefulDisconnectMetrics() + { + assertEquals(0, clientMetrics.connectionsDraining.get()); + long initialForcedDisconnects = clientMetrics.forcedDisconnects.getCount(); + + clientMetrics.incrementConnectionsDraining(); + assertEquals(1, clientMetrics.connectionsDraining.get()); + + clientMetrics.decrementConnectionsDraining(); + assertEquals(0, clientMetrics.connectionsDraining.get()); + + clientMetrics.markForcedDisconnect(3); + assertEquals(initialForcedDisconnects + 3, clientMetrics.forcedDisconnects.getCount()); + } } diff --git a/test/unit/org/apache/cassandra/service/StorageServiceTest.java b/test/unit/org/apache/cassandra/service/StorageServiceTest.java index 8742cdac2395..08ee30839f77 100644 --- a/test/unit/org/apache/cassandra/service/StorageServiceTest.java +++ b/test/unit/org/apache/cassandra/service/StorageServiceTest.java @@ -21,6 +21,7 @@ import java.util.concurrent.ScheduledFuture; import java.util.concurrent.atomic.AtomicInteger; +import org.assertj.core.api.Assertions; import org.junit.Before; import org.junit.BeforeClass; import org.junit.Test; @@ -201,6 +202,32 @@ public void testColumnIndexSizeInKiB() } } + @Test + public void testGracefulDisconnectGracePeriod() + { + StorageService storageService = StorageService.instance; + long originalGracePeriod = storageService.getGracefulDisconnectGracePeriod(); + try + { + storageService.setGracefulDisconnectGracePeriod(3000); + Assertions.assertThat(3000).isEqualTo(storageService.getGracefulDisconnectGracePeriod()); + + Assertions.assertThatThrownBy(() -> storageService.setGracefulDisconnectGracePeriod(-1)).isInstanceOf(IllegalArgumentException.class); + + Assertions.assertThat(3000).isEqualTo(storageService.getGracefulDisconnectGracePeriod()); + } + finally + { + storageService.setGracefulDisconnectGracePeriod(originalGracePeriod); + } + } + + @Test + public void testGracefulDisconnectEnabled() + { + Assertions.assertThat(StorageService.instance.getGracefulDisconnectEnabled()).isFalse(); + } + @Test public void testColumnIndexCacheSizeInKiB() { diff --git a/test/unit/org/apache/cassandra/transport/GracefulDisconnectLifecycleTest.java b/test/unit/org/apache/cassandra/transport/GracefulDisconnectLifecycleTest.java new file mode 100644 index 000000000000..7a704f91967a --- /dev/null +++ b/test/unit/org/apache/cassandra/transport/GracefulDisconnectLifecycleTest.java @@ -0,0 +1,239 @@ +/* + * 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 org.apache.cassandra.transport; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.BeforeClass; +import org.junit.Test; + +import org.apache.cassandra.config.DatabaseDescriptor; + +import io.netty.channel.Channel; +import io.netty.channel.DefaultChannelId; +import io.netty.channel.embedded.EmbeddedChannel; +import io.netty.channel.group.DefaultChannelGroup; +import io.netty.util.concurrent.GlobalEventExecutor; + +import static org.assertj.core.api.Assertions.assertThat; + +public class GracefulDisconnectLifecycleTest +{ + @BeforeClass + public static void setup() + { + DatabaseDescriptor.daemonInitialization(); + } + + @Test + public void testCompletesImmediatelyWithNoChannels() throws Exception + { + DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + + GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, channel -> {}); + + assertThat(lifecycle.run()).isEqualTo(0); + } + + @Test + public void testCompletesWhenChannelCloses() throws Exception + { + DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + + EmbeddedChannel channel = new EmbeddedChannel(); + channelGroup.add(channel); + + AtomicReference closedChannel = new AtomicReference<>(); + + GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, closedChannel::set); + + Future result = CompletableFuture.supplyAsync(() -> { + try + { + return lifecycle.run(); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + }); + + assertThat(result.isDone()).isFalse(); + + channel.close(); + + assertThat(result.get(5, TimeUnit.SECONDS)).isEqualTo(0); + assertThat(closedChannel.get()).isSameAs(channel); + } + + @Test + public void testForceClosesChannelAfterGracePeriod() throws Exception + { + DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + + EmbeddedChannel channel = new EmbeddedChannel(); + channelGroup.add(channel); + + GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ignored -> {}, 100, 100); + + Future result = CompletableFuture.supplyAsync(() -> { + try + { + return lifecycle.run(); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + }); + + int forcedDisconnects = result.get(5, TimeUnit.SECONDS); + + assertThat(forcedDisconnects).isEqualTo(1); + assertThat(channel.isOpen()).isFalse(); + } + + @Test + public void testWaitsForAllChannelsToClose() throws Exception + { + DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + + EmbeddedChannel channel1 = new EmbeddedChannel(); + EmbeddedChannel channel2 = new EmbeddedChannel(); + + channelGroup.add(channel1); + channelGroup.add(channel2); + + GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ignored -> {}, 1000, 100); + + Future result = CompletableFuture.supplyAsync(() -> { + try + { + return lifecycle.run(); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + }); + + channel1.close(); + + assertThat(result.isDone()).isFalse(); + + channel2.close(); + + assertThat(result.get(5, TimeUnit.SECONDS)).isEqualTo(0); + assertThat(channel1.isOpen()).isFalse(); + assertThat(channel2.isOpen()).isFalse(); + } + + @Test + public void testMixedCooperativeAndUncooperativeChannels() throws Exception + { + DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + + EmbeddedChannel cooperativeChannel = new EmbeddedChannel(DefaultChannelId.newInstance()); + EmbeddedChannel uncooperativeChannel = new EmbeddedChannel(DefaultChannelId.newInstance()); + + channelGroup.add(cooperativeChannel); + channelGroup.add(uncooperativeChannel); + + GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ignored -> {}, 200, 100); + + Future result = CompletableFuture.supplyAsync(() -> { + try + { + return lifecycle.run(); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + }); + + // Close only the cooperative channel immediately + cooperativeChannel.close(); + + // The lifecycle should wait for the grace period to expire for the uncooperative channel + int forcedDisconnects = result.get(5, TimeUnit.SECONDS); + + assertThat(forcedDisconnects).as("Only the uncooperative channel should be force closed").isEqualTo(1); + assertThat(cooperativeChannel.isOpen()).isFalse(); + assertThat(uncooperativeChannel.isOpen()).isFalse(); + } + + @Test + public void testChannelAlreadyClosedBeforeLifecycleStart() throws Exception + { + DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + + EmbeddedChannel alreadyClosedChannel = new EmbeddedChannel(); + alreadyClosedChannel.close(); // Closed before lifecycle starts + + channelGroup.add(alreadyClosedChannel); + + GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ignored -> {}, 1000, 100); + + int forcedDisconnects = lifecycle.run(); + + assertThat(forcedDisconnects).isEqualTo(0); + assertThat(alreadyClosedChannel.isOpen()).isFalse(); + } + + @Test + public void testCallbackInvokedForEachChannel() throws Exception + { + DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + + EmbeddedChannel ch1 = new EmbeddedChannel(DefaultChannelId.newInstance()); + EmbeddedChannel ch2 = new EmbeddedChannel(DefaultChannelId.newInstance()); + + channelGroup.add(ch1); + channelGroup.add(ch2); + + AtomicInteger closedCallbackCount = new AtomicInteger(0); + + // Grace period 100ms: ch1 will close cooperatively, ch2 will be force closed + GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ch -> closedCallbackCount.incrementAndGet(), 100, 100); + + Future result = CompletableFuture.supplyAsync(() -> { + try + { + return lifecycle.run(); + } + catch (Exception e) + { + throw new RuntimeException(e); + } + }); + + ch1.close(); + + result.get(5, TimeUnit.SECONDS); + + assertThat(closedCallbackCount.get()) + .as("Callback must be invoked exactly once per channel regardless of how it closed") + .isEqualTo(2); + } +} \ No newline at end of file From 635137a69f1c26e26954c8ad188115b3d91b2e0c Mon Sep 17 00:00:00 2001 From: Rishabh Saraswat Date: Thu, 3 Sep 2026 10:57:21 +0530 Subject: [PATCH 2/6] reformated SimpleClient.java --- .../GracefulDisconnectLifecycle.java | 153 -------- .../cassandra/transport/SimpleClient.java | 335 ++++++++++-------- ...eTest.java => GracefulDisconnectTest.java} | 0 3 files changed, 185 insertions(+), 303 deletions(-) delete mode 100644 src/java/org/apache/cassandra/transport/GracefulDisconnectLifecycle.java rename test/unit/org/apache/cassandra/transport/{GracefulDisconnectLifecycleTest.java => GracefulDisconnectTest.java} (100%) diff --git a/src/java/org/apache/cassandra/transport/GracefulDisconnectLifecycle.java b/src/java/org/apache/cassandra/transport/GracefulDisconnectLifecycle.java deleted file mode 100644 index 5479bb60c6e5..000000000000 --- a/src/java/org/apache/cassandra/transport/GracefulDisconnectLifecycle.java +++ /dev/null @@ -1,153 +0,0 @@ -/* - * 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 org.apache.cassandra.transport; - -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Consumer; - -import org.apache.cassandra.config.DatabaseDescriptor; -import org.apache.cassandra.transport.messages.EventMessage; -import org.apache.cassandra.utils.concurrent.CountDownLatch; - -import io.netty.channel.Channel; -import io.netty.channel.group.ChannelGroup; - -import static org.apache.cassandra.transport.Dispatcher.EVENT_DISPATCHER; - -final class GracefulDisconnectLifecycle -{ - // Stable snapshot of subscribed channels. - private final List channels; - - // Number of channels still awaiting closure. - private final AtomicInteger remainingChannels; - - // Signals that the graceful disconnect phase is complete. - private final CountDownLatch completion; - - // Tracks the current lifecycle state. - private final AtomicReference state; - - private final long gracefulDisconnectGracePeriod; - - // Called when each channel is closed. - private final Consumer onChannelClosed; - - // Number of clients forcefully disconnected after the grace period. - private final AtomicInteger forcedDisconnects = new AtomicInteger(); - - // Best-effort upper bound on how long forced Channel#close() calls should - // take to actually complete once triggered. Not user-configurable — - // this is a safety margin against the drain hanging. - private long hardDeadlineBufferMillis = 5000; - - public GracefulDisconnectLifecycle(ChannelGroup channelGroup, - Consumer onChannelClosed) - { - this(channelGroup, - onChannelClosed, - DatabaseDescriptor.getGracefulDisconnectGracePeriod(), - 5000); - } - - GracefulDisconnectLifecycle(ChannelGroup channelGroup, - Consumer onChannelClosed, - long gracefulDisconnectGracePeriod, - long hardDeadlineBufferMillis) - { - channels = new ArrayList<>(channelGroup); - remainingChannels = new AtomicInteger(channels.size()); - completion = CountDownLatch.newCountDownLatch(1); - state = new AtomicReference<>(State.WAITING_FOR_CLIENTS); - this.gracefulDisconnectGracePeriod = gracefulDisconnectGracePeriod; - this.onChannelClosed = onChannelClosed; - this.hardDeadlineBufferMillis = hardDeadlineBufferMillis; - } - - int run() throws InterruptedException, TimeoutException - { - startGracefulDisconnect(); - if (completion.await(gracefulDisconnectGracePeriod, TimeUnit.MILLISECONDS)) return forcedDisconnects.get(); - onGracePeriodExpired(); - if (!completion.await(hardDeadlineBufferMillis, TimeUnit.MILLISECONDS)) - throw new TimeoutException("Graceful disconnect did not complete even after forced close, " + remainingChannels.get() + " channel(s) still open"); - return forcedDisconnects.get(); - } - - private void startGracefulDisconnect() - { - if (remainingChannels.get() == 0) - { - complete(); - return; - } - - EventMessage eventMessage = new EventMessage(new Event.GracefulDisconnect()); - channels.forEach(channel -> { - channel.closeFuture().addListener(future -> onChannelClosed(channel)); - Consumer dispatcher = channel.attr(EVENT_DISPATCHER).get(); - if (dispatcher != null) - dispatcher.accept(eventMessage); - }); - } - - private void complete() - { - State previous = state.getAndSet(State.COMPLETE); - if (previous != State.COMPLETE) - completion.decrement(); - } - - private void onChannelClosed(Channel channel) - { - onChannelClosed.accept(channel); - - if (remainingChannels.decrementAndGet() == 0) - complete(); - } - - private void onGracePeriodExpired() - { - if (!state.compareAndSet(State.WAITING_FOR_CLIENTS, State.FORCE_CLOSING)) - return; - - int actuallyClosed = 0; - for (Channel channel : channels) - { - if (channel.isOpen()) - { - channel.close(); - actuallyClosed++; - } - } - forcedDisconnects.set(actuallyClosed); - } - - private enum State - { - WAITING_FOR_CLIENTS, - FORCE_CLOSING, - COMPLETE - } -} \ No newline at end of file diff --git a/src/java/org/apache/cassandra/transport/SimpleClient.java b/src/java/org/apache/cassandra/transport/SimpleClient.java index 65a5c9e08085..81254e40bb40 100644 --- a/src/java/org/apache/cassandra/transport/SimpleClient.java +++ b/src/java/org/apache/cassandra/transport/SimpleClient.java @@ -29,7 +29,7 @@ import java.util.Queue; import java.util.concurrent.BlockingQueue; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.SynchronousQueue; // checkstyle: permit this import import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; @@ -99,21 +99,76 @@ public class SimpleClient implements Closeable { public static final int TIMEOUT_SECONDS = 10; + + static + { + InternalLoggerFactory.setDefaultFactory(new Slf4JLoggerFactory()); + } + private static final Logger logger = LoggerFactory.getLogger(SimpleClient.class); + public final String host; public final int port; + private final EncryptionOptions.ClientEncryptionOptions encryptionOptions; + private final int largeMessageThreshold; + protected final ResponseHandler responseHandler = new ResponseHandler(this); protected final Connection.Tracker tracker = new ConnectionTracker(); protected final ProtocolVersion version; - private final AtomicBoolean draining = new AtomicBoolean(false); - private final EncryptionOptions.ClientEncryptionOptions encryptionOptions; - private final int largeMessageThreshold; // We don't track connection really, so we don't need one Connection per channel protected Connection connection; protected Bootstrap bootstrap; protected Channel channel; + protected ChannelFuture lastWriteFuture; + private final AtomicBoolean draining = new AtomicBoolean(false); protected String compression; - private volatile ChannelFuture lastWriteFuture; + + public static class Builder + { + private final String host; + private final int port; + private EncryptionOptions.ClientEncryptionOptions encryptionOptions = new EncryptionOptions.ClientEncryptionOptions(); + private ProtocolVersion version = ProtocolVersion.CURRENT; + private boolean useBeta = false; + private int largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE; + + private Builder(String host, int port) + { + this.host = host; + this.port = port; + } + + public Builder encryption(EncryptionOptions.ClientEncryptionOptions options) + { + this.encryptionOptions = options; + return this; + } + + public Builder useBeta() + { + this.useBeta = true; + return this; + } + + public Builder protocolVersion(ProtocolVersion version) + { + this.version = version; + return this; + } + + public Builder largeMessageThreshold(int bytes) + { + largeMessageThreshold = bytes; + return this; + } + + public SimpleClient build() + { + if (version.isBeta() && !useBeta) + throw new IllegalArgumentException(String.format("Beta version of server used (%s), but USE_BETA flag is not set", version)); + return new SimpleClient(this); + } + } private SimpleClient(Builder builder) { @@ -148,7 +203,9 @@ public SimpleClient(String host, int port, ProtocolVersion version, boolean useB this.version = version; this.encryptionOptions = encryptionOptions.applyConfig(); - this.largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE - Math.max(FrameEncoderCrc.HEADER_AND_TRAILER_LENGTH, FrameEncoderLZ4.HEADER_AND_TRAILER_LENGTH); + this.largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE - + Math.max(FrameEncoderCrc.HEADER_AND_TRAILER_LENGTH, + FrameEncoderLZ4.HEADER_AND_TRAILER_LENGTH); } public SimpleClient(String host, int port) @@ -161,18 +218,6 @@ public static Builder builder(String host, int port) return new Builder(host, port); } - /** - * The stream id to frame an outbound client request with. SimpleClient carries the intended id on the - * request's (dummy) source envelope (see {@link #execute(List)} and callers that pipeline requests). - * When no source has been assigned, we fall back to 0, which is sufficient for the non-pipelined path - * where only a single request is ever in flight. - */ - private static int outboundStreamId(Message message) - { - Envelope source = message.getSource(); - return source == null ? 0 : source.header.streamId; - } - public SimpleClient connect(boolean useCompression) throws IOException { return connect(useCompression, false); @@ -184,7 +229,8 @@ public SimpleClient connect(boolean useCompression, boolean throwOnOverload) thr Map options = new HashMap<>(); options.put(StartupMessage.CQL_VERSION, "3.0.0"); - if (throwOnOverload) options.put(StartupMessage.THROW_ON_OVERLOAD, "1"); + if (throwOnOverload) + options.put(StartupMessage.THROW_ON_OVERLOAD, "1"); connection.setThrowOnOverload(throwOnOverload); if (useCompression) @@ -193,6 +239,7 @@ public SimpleClient connect(boolean useCompression, boolean throwOnOverload) thr connection.setCompressor(Compressor.LZ4Compressor.instance); } execute(new StartupMessage(options)); + return this; } @@ -205,10 +252,12 @@ public void setEventHandler(EventHandler eventHandler) void establishConnection() throws IOException { // Configure the client. - bootstrap = new Bootstrap().group(new NioEventLoopGroup(new NamedThreadFactory("SimpleClient-nioEventLoopGroup"))).channel(io.netty.channel.socket.nio.NioSocketChannel.class).option(ChannelOption.TCP_NODELAY, true); + bootstrap = new Bootstrap().group(new NioEventLoopGroup(new NamedThreadFactory("SimpleClient-nioEventLoopGroup"))) + .channel(io.netty.channel.socket.nio.NioSocketChannel.class) + .option(ChannelOption.TCP_NODELAY, true); // Configure the pipeline factory. - if (encryptionOptions.getEnabled()) + if(encryptionOptions.getEnabled()) { bootstrap.handler(new SecureInitializer(largeMessageThreshold)); } @@ -229,21 +278,21 @@ void establishConnection() throws IOException public ResultMessage execute(String query, ConsistencyLevel consistency) { - return execute(query, Collections.emptyList(), consistency); + return execute(query, Collections.emptyList(), consistency); } public ResultMessage execute(String query, List values, ConsistencyLevel consistencyLevel) { Message.Response msg = execute(new QueryMessage(query, QueryOptions.forInternalCalls(consistencyLevel, values))); assert msg instanceof ResultMessage; - return (ResultMessage) msg; + return (ResultMessage)msg; } public ResultMessage.Prepared prepare(String query) { Message.Response msg = execute(new PrepareMessage(query, null)); assert msg instanceof ResultMessage.Prepared; - return (ResultMessage.Prepared) msg; + return (ResultMessage.Prepared)msg; } public ResultMessage executePrepared(ResultMessage.Prepared prepared, List values, ConsistencyLevel consistency) @@ -256,7 +305,8 @@ public ResultMessage executePrepared(ResultMessage.Prepared prepared, List execute(List requ for (int i = 0; i < requests.size(); i++) { Message.Response msg = responseHandler.responses.poll(deadline - currentTimeMillis(), TimeUnit.MILLISECONDS); - if (msg == null) throw new RuntimeException("timeout"); - if (msg instanceof ErrorMessage) throw new RuntimeException((Throwable) ((ErrorMessage) msg).error); + if (msg == null) + throw new RuntimeException("timeout"); + if (msg instanceof ErrorMessage) + throw new RuntimeException((Throwable) ((ErrorMessage) msg).error); rrMap.put(requests.get(msg.getSource().header.streamId), msg); } } @@ -333,15 +386,16 @@ public Map execute(List requ } } - private void handleGracefulDisconnect() + /** + * The stream id to frame an outbound client request with. SimpleClient carries the intended id on the + * request's (dummy) source envelope (see {@link #execute(List)} and callers that pipeline requests). + * When no source has been assigned, we fall back to 0, which is sufficient for the non-pipelined path + * where only a single request is ever in flight. + */ + private static int outboundStreamId(Message message) { - draining.set(true); - channel.eventLoop().execute(() -> { - ChannelFuture writeFuture = lastWriteFuture; - ChannelFuture closeAfter = (writeFuture != null) ? writeFuture : channel.newSucceededFuture(); - closeAfter.addListener(f -> channel.close()); - }); - channel.closeFuture().addListener(f -> bootstrap.group().shutdownGracefully()); + Envelope source = message.getSource(); + return source == null ? 0 : source.header.streamId; } public interface EventHandler @@ -349,51 +403,15 @@ public interface EventHandler void onEvent(Event event); } - public static class Builder + private void handleGracefulDisconnect() { - private final String host; - private final int port; - private EncryptionOptions.ClientEncryptionOptions encryptionOptions = new EncryptionOptions.ClientEncryptionOptions(); - private ProtocolVersion version = ProtocolVersion.CURRENT; - private boolean useBeta = false; - private int largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE; - - private Builder(String host, int port) - { - this.host = host; - this.port = port; - } - - public Builder encryption(EncryptionOptions.ClientEncryptionOptions options) - { - this.encryptionOptions = options; - return this; - } - - public Builder useBeta() - { - this.useBeta = true; - return this; - } - - public Builder protocolVersion(ProtocolVersion version) - { - this.version = version; - return this; - } - - public Builder largeMessageThreshold(int bytes) - { - largeMessageThreshold = bytes; - return this; - } - - public SimpleClient build() - { - if (version.isBeta() && !useBeta) - throw new IllegalArgumentException(String.format("Beta version of server used (%s), but USE_BETA flag is not set", version)); - return new SimpleClient(this); - } + draining.set(true); + channel.eventLoop().execute(() -> { + ChannelFuture writeFuture = lastWriteFuture; + ChannelFuture closeAfter = (writeFuture != null) ? writeFuture : channel.newSucceededFuture(); + closeAfter.addListener(f -> channel.close()); + }); + channel.closeFuture().addListener(f -> bootstrap.group().shutdownGracefully()); } public static class SimpleEventHandler implements EventHandler @@ -421,19 +439,19 @@ public boolean isRunning() private static class HandlerNames { - private static final String ENVELOPE_DECODER = "envelopeDecoder"; - private static final String ENVELOPE_ENCODER = "envelopeEncoder"; - private static final String COMPRESSOR = "compressor"; - private static final String DECOMPRESSOR = "decompressor"; - private static final String MESSAGE_DECODER = "messageDecoder"; - private static final String MESSAGE_ENCODER = "messageEncoder"; + private static final String ENVELOPE_DECODER = "envelopeDecoder"; + private static final String ENVELOPE_ENCODER = "envelopeEncoder"; + private static final String COMPRESSOR = "compressor"; + private static final String DECOMPRESSOR = "decompressor"; + private static final String MESSAGE_DECODER = "messageDecoder"; + private static final String MESSAGE_ENCODER = "messageEncoder"; - private static final String INITIAL_HANDLER = "intitialHandler"; - private static final String RESPONSE_HANDLER = "responseHandler"; + private static final String INITIAL_HANDLER = "intitialHandler"; + private static final String RESPONSE_HANDLER = "responseHandler"; - private static final String FRAME_DECODER = "frameDecoder"; - private static final String FRAME_ENCODER = "frameEncoder"; - private static final String PROCESSOR = "processor"; + private static final String FRAME_DECODER = "frameDecoder"; + private static final String FRAME_ENCODER = "frameEncoder"; + private static final String PROCESSOR = "processor"; } private static class InitialHandler extends MessageToMessageDecoder @@ -441,7 +459,6 @@ private static class InitialHandler extends MessageToMessageDecoder final ProtocolVersion version; final ResponseHandler responseHandler; final int largeMessageThreshold; - InitialHandler(ProtocolVersion version, ResponseHandler responseHandler, int largeMessageThreshold) { this.version = version; @@ -452,7 +469,7 @@ private static class InitialHandler extends MessageToMessageDecoder @Override protected void decode(ChannelHandlerContext ctx, Envelope request, List results) { - switch (request.header.type) + switch(request.header.type) { case READY: case AUTHENTICATE: @@ -474,7 +491,9 @@ protected void decode(ChannelHandlerContext ctx, Envelope request, List results.add(request); break; default: - throw new ProtocolException(String.format("Unexpected %s request expecting " + "READY, AUTHENTICATE, ERROR or SUPPORTED", request.header.type)); + throw new ProtocolException(String.format("Unexpected %s request expecting " + + "READY, AUTHENTICATE, ERROR or SUPPORTED", + request.header.type)); } } @@ -554,8 +573,21 @@ public void release() } }; - CQLMessageHandler processor = new CQLMessageHandler(ctx.channel(), null, version, frameDecoder, envelopeDecoder, messageDecoder, responseConsumer, payloadAllocator, queueCapacity, QueueBackpressure.NO_OP, resources, handler -> { - }, errorHandler, ctx.channel().attr(Connection.attributeKey).get().isThrowOnOverload()) + CQLMessageHandler processor = + new CQLMessageHandler(ctx.channel(), + null, + version, + frameDecoder, + envelopeDecoder, + messageDecoder, + responseConsumer, + payloadAllocator, + queueCapacity, + QueueBackpressure.NO_OP, + resources, + handler -> {}, + errorHandler, + ctx.channel().attr(Connection.attributeKey).get().isThrowOnOverload()) { protected boolean processRequest(Envelope request, Overload overload) { @@ -568,8 +600,7 @@ protected boolean processRequest(Envelope request, Overload overload) pipeline.addLast(HandlerNames.FRAME_DECODER, frameDecoder); pipeline.addLast(HandlerNames.FRAME_ENCODER, frameEncoder); pipeline.addLast(HandlerNames.PROCESSOR, processor); - pipeline.addLast(HandlerNames.MESSAGE_ENCODER, new ChannelOutboundHandlerAdapter() - { + pipeline.addLast(HandlerNames.MESSAGE_ENCODER, new ChannelOutboundHandlerAdapter() { public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception { @@ -597,16 +628,20 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) private FrameDecoder frameDecoder(ChannelHandlerContext ctx, BufferPoolAllocator allocator) { Connection conn = ctx.channel().attr(Connection.attributeKey).get(); - if (conn.getCompressor() == null) return FrameDecoderCrc.create(allocator); - if (conn.getCompressor() instanceof Compressor.LZ4Compressor) return FrameDecoderLZ4.fast(allocator); + if (conn.getCompressor() == null) + return FrameDecoderCrc.create(allocator); + if (conn.getCompressor() instanceof Compressor.LZ4Compressor) + return FrameDecoderLZ4.fast(allocator); throw new ProtocolException("Unsupported compressor: " + conn.getCompressor().getClass().getCanonicalName()); } private FrameEncoder frameEncoder(ChannelHandlerContext ctx) { Connection conn = ctx.channel().attr(Connection.attributeKey).get(); - if (conn.getCompressor() == null) return FrameEncoderCrc.instance; - if (conn.getCompressor() instanceof Compressor.LZ4Compressor) return FrameEncoderLZ4.fastInstance; + if (conn.getCompressor() == null) + return FrameEncoderCrc.instance; + if (conn.getCompressor() instanceof Compressor.LZ4Compressor) + return FrameEncoderLZ4.fastInstance; throw new ProtocolException("Unsupported compressor: " + conn.getCompressor().getClass().getCanonicalName()); } @@ -621,14 +656,10 @@ private void configureLegacyPipeline(ChannelHandlerContext ctx) } @ChannelHandler.Sharable - static class MessageBatchEncoder extends MessageToMessageEncoder> + static class MessageBatchEncoder extends MessageToMessageEncoder> { public static final MessageBatchEncoder instance = new MessageBatchEncoder(); - private MessageBatchEncoder() - { - } - public void encode(ChannelHandlerContext ctx, List messages, List results) { Connection connection = ctx.channel().attr(Connection.attributeKey).get(); @@ -640,10 +671,34 @@ public void encode(ChannelHandlerContext ctx, List messages, List + { + private final int largeMessageThreshold; + + Initializer(int largeMessageThreshold) + { + this.largeMessageThreshold = largeMessageThreshold; + } + + protected void initChannel(Channel channel) throws Exception + { + connection = new Connection(channel, version, tracker); + channel.attr(Connection.attributeKey).set(connection); + + ChannelPipeline pipeline = channel.pipeline(); +// pipeline.addLast("debug", new LoggingHandler(LogLevel.INFO)); + pipeline.addLast(HandlerNames.ENVELOPE_DECODER, new Envelope.Decoder()); + pipeline.addLast(HandlerNames.ENVELOPE_ENCODER, Envelope.Encoder.instance); + pipeline.addLast(HandlerNames.INITIAL_HANDLER, new InitialHandler(version, responseHandler, largeMessageThreshold)); + pipeline.addLast(HandlerNames.MESSAGE_DECODER, PreV5Handlers.ProtocolDecoder.instance); + pipeline.addLast(HandlerNames.MESSAGE_ENCODER, MessageBatchEncoder.instance); + pipeline.addLast(HandlerNames.RESPONSE_HANDLER, responseHandler); + } + } + @ChannelHandler.Sharable static class ResponseHandler extends SimpleChannelInboundHandler { - public final BlockingQueue responses = new SynchronousQueue<>(true); private final SimpleClient client; public EventHandler eventHandler; @@ -668,19 +723,22 @@ public void handleResponse(Channel channel, Message.Response r) if (r instanceof EventMessage) { - Event event = ((EventMessage) r).event; + Event event = ((EventMessage)r).event; if (event.type == Event.Type.GRACEFUL_DISCONNECT) { logger.info("Received GRACEFUL_DISCONNECT. Entering draining mode."); - if (eventHandler != null) eventHandler.onEvent(event); + if (eventHandler != null) + eventHandler.onEvent(event); client.handleGracefulDisconnect(); return; } - if (eventHandler != null) eventHandler.onEvent(event); + if (eventHandler != null) + eventHandler.onEvent(event); } - else responses.put(r); + else + responses.put(r); } catch (InterruptedException e) { @@ -733,13 +791,15 @@ public void enqueue(Envelope message) public void releaseAll() { Envelope e; - while ((e = outbound.poll()) != null) e.release(); + while ((e = outbound.poll()) != null) + e.release(); } public void schedule(ChannelHandlerContext ctx) { if (scheduled.compareAndSet(false, true)) - ctx.executor().scheduleAtFixedRate(() -> maybeWrite(ctx, ctx.voidPromise()), 10, 10, TimeUnit.MILLISECONDS); + ctx.executor().scheduleAtFixedRate(() -> maybeWrite(ctx, ctx.voidPromise()), + 10, 10, TimeUnit.MILLISECONDS); } public void maybeWrite(ChannelHandlerContext ctx, Promise promise) @@ -803,8 +863,10 @@ private ChannelFuture flushBuffer(ChannelHandlerContext ctx, List mess private FrameEncoder.Payload allocate(int size, boolean selfContained) { - FrameEncoder.Payload payload = frameEncoder.allocator().allocate(selfContained, Math.min(size, largeMessageThreshold)); - if (size >= largeMessageThreshold) payload.buffer.limit(largeMessageThreshold); + FrameEncoder.Payload payload = frameEncoder.allocator() + .allocate(selfContained, Math.min(size, largeMessageThreshold)); + if (size >= largeMessageThreshold) + payload.buffer.limit(largeMessageThreshold); return payload; } @@ -823,7 +885,8 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) buf = payload.buffer; // BufferPool may give us a buffer larger than we asked for. // FrameEncoder may object if buffer.remaining is >= MAX_SIZE. - if (payloadSize >= largeMessageThreshold) buf.limit(largeMessageThreshold); + if (payloadSize >= largeMessageThreshold) + buf.limit(largeMessageThreshold); if (firstFrame) { @@ -832,7 +895,8 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) } int remaining = Math.min(buf.remaining(), f.body.readableBytes()); - if (remaining > 0) buf.put(f.body.slice(f.body.readerIndex(), remaining).nioBuffer()); + if (remaining > 0) + buf.put(f.body.slice(f.body.readerIndex(), remaining).nioBuffer()); f.body.readerIndex(f.body.readerIndex() + remaining); payload.finish(); @@ -842,7 +906,8 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) promise.addListener(result -> { if (!result.isSuccess()) logger.warn("Failed to send frame of large message, size: " + remaining, result.cause()); - else logger.trace("Sent frame of large message, size: {}", remaining); + else + logger.trace("Sent frame of large message, size: {}", remaining); }); } f.release(); @@ -850,31 +915,6 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) } } - private class Initializer extends ChannelInitializer - { - private final int largeMessageThreshold; - - Initializer(int largeMessageThreshold) - { - this.largeMessageThreshold = largeMessageThreshold; - } - - protected void initChannel(Channel channel) throws Exception - { - connection = new Connection(channel, version, tracker); - channel.attr(Connection.attributeKey).set(connection); - - ChannelPipeline pipeline = channel.pipeline(); -// pipeline.addLast("debug", new LoggingHandler(LogLevel.INFO)); - pipeline.addLast(HandlerNames.ENVELOPE_DECODER, new Envelope.Decoder()); - pipeline.addLast(HandlerNames.ENVELOPE_ENCODER, Envelope.Encoder.instance); - pipeline.addLast(HandlerNames.INITIAL_HANDLER, new InitialHandler(version, responseHandler, largeMessageThreshold)); - pipeline.addLast(HandlerNames.MESSAGE_DECODER, PreV5Handlers.ProtocolDecoder.instance); - pipeline.addLast(HandlerNames.MESSAGE_ENCODER, MessageBatchEncoder.instance); - pipeline.addLast(HandlerNames.RESPONSE_HANDLER, responseHandler); - } - } - private class SecureInitializer extends Initializer { SecureInitializer(int largeMessageThreshold) @@ -890,9 +930,4 @@ protected void initChannel(Channel channel) throws Exception channel.pipeline().addFirst("ssl", newSslHandler(channel, sslContext, peer)); } } - - static - { - InternalLoggerFactory.setDefaultFactory(new Slf4JLoggerFactory()); - } } \ No newline at end of file diff --git a/test/unit/org/apache/cassandra/transport/GracefulDisconnectLifecycleTest.java b/test/unit/org/apache/cassandra/transport/GracefulDisconnectTest.java similarity index 100% rename from test/unit/org/apache/cassandra/transport/GracefulDisconnectLifecycleTest.java rename to test/unit/org/apache/cassandra/transport/GracefulDisconnectTest.java From 3ee100e3f086870f0a3d1110f650ccb79983f635 Mon Sep 17 00:00:00 2001 From: Rishabh Saraswat Date: Thu, 3 Sep 2026 11:22:51 +0530 Subject: [PATCH 3/6] removed Lifecycle way of Draining --- .../service/NativeTransportService.java | 1 + .../apache/cassandra/transport/Server.java | 37 ++-- .../cassandra/transport/SimpleClient.java | 109 ++++----- .../transport/GracefulDisconnectTest.java | 208 ++++-------------- 4 files changed, 117 insertions(+), 238 deletions(-) diff --git a/src/java/org/apache/cassandra/service/NativeTransportService.java b/src/java/org/apache/cassandra/service/NativeTransportService.java index 894f8f782800..1df9a68e13c2 100644 --- a/src/java/org/apache/cassandra/service/NativeTransportService.java +++ b/src/java/org/apache/cassandra/service/NativeTransportService.java @@ -164,6 +164,7 @@ Server getServer() return server; } + @VisibleForTesting public ChannelGroup getChannelsSubscribedToGracefulDisconnect() { return server.getChannelsSubscribedToGracefulDisconnect(); diff --git a/src/java/org/apache/cassandra/transport/Server.java b/src/java/org/apache/cassandra/transport/Server.java index a8fcafb41f98..3e74c261513c 100644 --- a/src/java/org/apache/cassandra/transport/Server.java +++ b/src/java/org/apache/cassandra/transport/Server.java @@ -28,7 +28,6 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.LockSupport; import java.util.function.BooleanSupplier; @@ -156,23 +155,33 @@ private void gracefulDisconnect() stopAcceptingNewConnections(); ChannelGroup channelGroup = getChannelsSubscribedToGracefulDisconnect(); + if (channelGroup.isEmpty()) + return; + ClientMetrics.instance.connectionsDraining.set(channelGroup.size()); + channelGroup.forEach(channel -> + channel.closeFuture().addListener(future -> ClientMetrics.instance.decrementConnectionsDraining()) + ); - try - { - int forcedDisconnects = new GracefulDisconnectLifecycle(channelGroup, - channel -> ClientMetrics.instance.decrementConnectionsDraining()).run(); - ClientMetrics.instance.markForcedDisconnect(forcedDisconnects); - } - catch (InterruptedException e) - { - logger.warn("Graceful disconnect interrupted", e); - Thread.currentThread().interrupt(); - } - catch (TimeoutException e) + connectionTracker.send(new Event.GracefulDisconnect()); + + long gracePeriod = DatabaseDescriptor.getGracefulDisconnectGracePeriod(); + int forcedDisconnects = 0; + + boolean completedCleanly = channelGroup.newCloseFuture().awaitUninterruptibly(gracePeriod, TimeUnit.MILLISECONDS); + + if (!completedCleanly) { - logger.warn("Graceful disconnect timed out waiting for channels to close; proceeding with shutdown", e); + forcedDisconnects = channelGroup.size(); + if (forcedDisconnects > 0) + { + logger.warn("Draining grace period of {}ms elapsed; {} active client connection(s) failed to close cleanly and will be forcefully terminated.", + gracePeriod, forcedDisconnects); + channelGroup.close(); + } } + + ClientMetrics.instance.markForcedDisconnect(forcedDisconnects); } public boolean isRunning() diff --git a/src/java/org/apache/cassandra/transport/SimpleClient.java b/src/java/org/apache/cassandra/transport/SimpleClient.java index 81254e40bb40..0a96294fe361 100644 --- a/src/java/org/apache/cassandra/transport/SimpleClient.java +++ b/src/java/org/apache/cassandra/transport/SimpleClient.java @@ -82,7 +82,7 @@ import io.netty.handler.codec.MessageToMessageDecoder; import io.netty.handler.codec.MessageToMessageEncoder; import io.netty.handler.ssl.SslContext; -import io.netty.util.concurrent.Promise; +import io.netty.util.concurrent.Promise; // checkstyle: permit this import import io.netty.util.concurrent.PromiseCombiner; import io.netty.util.internal.logging.InternalLoggerFactory; import io.netty.util.internal.logging.Slf4JLoggerFactory; @@ -170,6 +170,12 @@ public SimpleClient build() } } + + public static Builder builder(String host, int port) + { + return new Builder(host, port); + } + private SimpleClient(Builder builder) { this.host = builder.host; @@ -204,8 +210,8 @@ public SimpleClient(String host, int port, ProtocolVersion version, boolean useB this.version = version; this.encryptionOptions = encryptionOptions.applyConfig(); this.largeMessageThreshold = FrameEncoder.Payload.MAX_SIZE - - Math.max(FrameEncoderCrc.HEADER_AND_TRAILER_LENGTH, - FrameEncoderLZ4.HEADER_AND_TRAILER_LENGTH); + Math.max(FrameEncoderCrc.HEADER_AND_TRAILER_LENGTH, + FrameEncoderLZ4.HEADER_AND_TRAILER_LENGTH); } public SimpleClient(String host, int port) @@ -213,11 +219,6 @@ public SimpleClient(String host, int port) this(host, port, new EncryptionOptions.ClientEncryptionOptions()); } - public static Builder builder(String host, int port) - { - return new Builder(host, port); - } - public SimpleClient connect(boolean useCompression) throws IOException { return connect(useCompression, false); @@ -252,9 +253,10 @@ public void setEventHandler(EventHandler eventHandler) void establishConnection() throws IOException { // Configure the client. - bootstrap = new Bootstrap().group(new NioEventLoopGroup(new NamedThreadFactory("SimpleClient-nioEventLoopGroup"))) - .channel(io.netty.channel.socket.nio.NioSocketChannel.class) - .option(ChannelOption.TCP_NODELAY, true); + bootstrap = new Bootstrap() + .group(new NioEventLoopGroup(new NamedThreadFactory("SimpleClient-nioEventLoopGroup"))) + .channel(io.netty.channel.socket.nio.NioSocketChannel.class) + .option(ChannelOption.TCP_NODELAY, true); // Configure the pipeline factory. if(encryptionOptions.getEnabled()) @@ -299,7 +301,7 @@ public ResultMessage executePrepared(ResultMessage.Prepared prepared, List processor = - new CQLMessageHandler(ctx.channel(), - null, - version, - frameDecoder, - envelopeDecoder, - messageDecoder, - responseConsumer, - payloadAllocator, - queueCapacity, - QueueBackpressure.NO_OP, - resources, - handler -> {}, - errorHandler, - ctx.channel().attr(Connection.attributeKey).get().isThrowOnOverload()) - { - protected boolean processRequest(Envelope request, Overload overload) + new CQLMessageHandler(ctx.channel(), + null, + version, + frameDecoder, + envelopeDecoder, + messageDecoder, + responseConsumer, + payloadAllocator, + queueCapacity, + QueueBackpressure.NO_OP, + resources, + handler -> {}, + errorHandler, + ctx.channel().attr(Connection.attributeKey).get().isThrowOnOverload()) { - boolean continueProcessing = super.processRequest(request, overload); - releaseCapacity(Ints.checkedCast(request.header.bodySizeInBytes)); - return continueProcessing; - } - }; + protected boolean processRequest(Envelope request, Overload overload) + { + boolean continueProcessing = super.processRequest(request, overload); + releaseCapacity(Ints.checkedCast(request.header.bodySizeInBytes)); + return continueProcessing; + } + }; pipeline.addLast(HandlerNames.FRAME_DECODER, frameDecoder); pipeline.addLast(HandlerNames.FRAME_ENCODER, frameEncoder); @@ -659,6 +659,7 @@ private void configureLegacyPipeline(ChannelHandlerContext ctx) static class MessageBatchEncoder extends MessageToMessageEncoder> { public static final MessageBatchEncoder instance = new MessageBatchEncoder(); + private MessageBatchEncoder(){} public void encode(ChannelHandlerContext ctx, List messages, List results) { @@ -673,8 +674,7 @@ public void encode(ChannelHandlerContext ctx, List messages, List { - private final int largeMessageThreshold; - + private int largeMessageThreshold; Initializer(int largeMessageThreshold) { this.largeMessageThreshold = largeMessageThreshold; @@ -696,6 +696,23 @@ protected void initChannel(Channel channel) throws Exception } } + private class SecureInitializer extends Initializer + { + SecureInitializer(int largeMessageThreshold) + { + super(largeMessageThreshold); + } + + protected void initChannel(Channel channel) throws Exception + { + super.initChannel(channel); + SslContext sslContext = SSLFactory.getOrCreateSslContext(encryptionOptions, encryptionOptions.getClientAuth(), + ISslContextFactory.SocketType.CLIENT, SSL_FACTORY_CONTEXT_DESCRIPTION); + InetSocketAddress peer = encryptionOptions.require_endpoint_verification ? new InetSocketAddress(host, port) : null; + channel.pipeline().addFirst("ssl", newSslHandler(channel, sslContext, peer)); + } + } + @ChannelHandler.Sharable static class ResponseHandler extends SimpleChannelInboundHandler { @@ -914,20 +931,4 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) return futures.toArray(EMPTY_FUTURES_ARRAY); } } - - private class SecureInitializer extends Initializer - { - SecureInitializer(int largeMessageThreshold) - { - super(largeMessageThreshold); - } - - protected void initChannel(Channel channel) throws Exception - { - super.initChannel(channel); - SslContext sslContext = SSLFactory.getOrCreateSslContext(encryptionOptions, encryptionOptions.getClientAuth(), ISslContextFactory.SocketType.CLIENT, SSL_FACTORY_CONTEXT_DESCRIPTION); - InetSocketAddress peer = encryptionOptions.require_endpoint_verification ? new InetSocketAddress(host, port) : null; - channel.pipeline().addFirst("ssl", newSslHandler(channel, sslContext, peer)); - } - } } \ No newline at end of file diff --git a/test/unit/org/apache/cassandra/transport/GracefulDisconnectTest.java b/test/unit/org/apache/cassandra/transport/GracefulDisconnectTest.java index 7a704f91967a..24931e6c0e06 100644 --- a/test/unit/org/apache/cassandra/transport/GracefulDisconnectTest.java +++ b/test/unit/org/apache/cassandra/transport/GracefulDisconnectTest.java @@ -21,23 +21,21 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; import org.junit.BeforeClass; import org.junit.Test; import org.apache.cassandra.config.DatabaseDescriptor; -import io.netty.channel.Channel; import io.netty.channel.DefaultChannelId; import io.netty.channel.embedded.EmbeddedChannel; +import io.netty.channel.group.ChannelGroup; import io.netty.channel.group.DefaultChannelGroup; import io.netty.util.concurrent.GlobalEventExecutor; import static org.assertj.core.api.Assertions.assertThat; -public class GracefulDisconnectLifecycleTest +public class GracefulDisconnectTest { @BeforeClass public static void setup() @@ -46,194 +44,64 @@ public static void setup() } @Test - public void testCompletesImmediatelyWithNoChannels() throws Exception + public void testEmptyChannelGroupGracefulHandling() { - DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); - - GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, channel -> {}); - - assertThat(lifecycle.run()).isEqualTo(0); + ChannelGroup group = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + boolean completed = group.newCloseFuture().awaitUninterruptibly(100, TimeUnit.MILLISECONDS); + assertThat(completed).isTrue(); + assertThat(group.size()).isEqualTo(0); } @Test - public void testCompletesWhenChannelCloses() throws Exception + public void testChannelGroupAutomaticRemovalOnClose() { - DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); - - EmbeddedChannel channel = new EmbeddedChannel(); - channelGroup.add(channel); - - AtomicReference closedChannel = new AtomicReference<>(); - - GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, closedChannel::set); - - Future result = CompletableFuture.supplyAsync(() -> { - try - { - return lifecycle.run(); - } - catch (Exception e) - { - throw new RuntimeException(e); - } - }); - - assertThat(result.isDone()).isFalse(); + ChannelGroup group = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + EmbeddedChannel channel = new EmbeddedChannel(DefaultChannelId.newInstance()); + group.add(channel); + assertThat(group.size()).isEqualTo(1); channel.close(); - - assertThat(result.get(5, TimeUnit.SECONDS)).isEqualTo(0); - assertThat(closedChannel.get()).isSameAs(channel); - } - - @Test - public void testForceClosesChannelAfterGracePeriod() throws Exception - { - DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); - - EmbeddedChannel channel = new EmbeddedChannel(); - channelGroup.add(channel); - - GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ignored -> {}, 100, 100); - - Future result = CompletableFuture.supplyAsync(() -> { - try - { - return lifecycle.run(); - } - catch (Exception e) - { - throw new RuntimeException(e); - } - }); - - int forcedDisconnects = result.get(5, TimeUnit.SECONDS); - - assertThat(forcedDisconnects).isEqualTo(1); - assertThat(channel.isOpen()).isFalse(); - } - - @Test - public void testWaitsForAllChannelsToClose() throws Exception - { - DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); - - EmbeddedChannel channel1 = new EmbeddedChannel(); - EmbeddedChannel channel2 = new EmbeddedChannel(); - - channelGroup.add(channel1); - channelGroup.add(channel2); - - GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ignored -> {}, 1000, 100); - - Future result = CompletableFuture.supplyAsync(() -> { - try - { - return lifecycle.run(); - } - catch (Exception e) - { - throw new RuntimeException(e); - } - }); - - channel1.close(); - - assertThat(result.isDone()).isFalse(); - - channel2.close(); - - assertThat(result.get(5, TimeUnit.SECONDS)).isEqualTo(0); - assertThat(channel1.isOpen()).isFalse(); - assertThat(channel2.isOpen()).isFalse(); + assertThat(group.size()).isEqualTo(0); } @Test - public void testMixedCooperativeAndUncooperativeChannels() throws Exception + public void testCloseFutureTriggersWhenAllChannelsTerminate() throws Exception { - DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); - - EmbeddedChannel cooperativeChannel = new EmbeddedChannel(DefaultChannelId.newInstance()); - EmbeddedChannel uncooperativeChannel = new EmbeddedChannel(DefaultChannelId.newInstance()); - - channelGroup.add(cooperativeChannel); - channelGroup.add(uncooperativeChannel); - - GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ignored -> {}, 200, 100); - - Future result = CompletableFuture.supplyAsync(() -> { - try - { - return lifecycle.run(); - } - catch (Exception e) - { - throw new RuntimeException(e); - } - }); - - // Close only the cooperative channel immediately - cooperativeChannel.close(); - - // The lifecycle should wait for the grace period to expire for the uncooperative channel - int forcedDisconnects = result.get(5, TimeUnit.SECONDS); - - assertThat(forcedDisconnects).as("Only the uncooperative channel should be force closed").isEqualTo(1); - assertThat(cooperativeChannel.isOpen()).isFalse(); - assertThat(uncooperativeChannel.isOpen()).isFalse(); - } - - @Test - public void testChannelAlreadyClosedBeforeLifecycleStart() throws Exception - { - DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); - - EmbeddedChannel alreadyClosedChannel = new EmbeddedChannel(); - alreadyClosedChannel.close(); // Closed before lifecycle starts + ChannelGroup group = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + EmbeddedChannel ch1 = new EmbeddedChannel(DefaultChannelId.newInstance()); + EmbeddedChannel ch2 = new EmbeddedChannel(DefaultChannelId.newInstance()); + group.add(ch1); + group.add(ch2); - channelGroup.add(alreadyClosedChannel); + Future closeFutureCompletion = CompletableFuture.supplyAsync(() -> group.newCloseFuture().awaitUninterruptibly(3000, TimeUnit.MILLISECONDS)); - GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ignored -> {}, 1000, 100); + assertThat(closeFutureCompletion.isDone()).isFalse(); - int forcedDisconnects = lifecycle.run(); + ch1.close(); + ch2.close(); - assertThat(forcedDisconnects).isEqualTo(0); - assertThat(alreadyClosedChannel.isOpen()).isFalse(); + assertThat(closeFutureCompletion.get(2, TimeUnit.SECONDS)).isTrue(); } @Test - public void testCallbackInvokedForEachChannel() throws Exception + public void testGracePeriodTimeoutTriggersBulkForceClose() { - DefaultChannelGroup channelGroup = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + ChannelGroup group = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + EmbeddedChannel cooperative = new EmbeddedChannel(DefaultChannelId.newInstance()); + EmbeddedChannel uncooperative = new EmbeddedChannel(DefaultChannelId.newInstance()); + group.add(cooperative); + group.add(uncooperative); - EmbeddedChannel ch1 = new EmbeddedChannel(DefaultChannelId.newInstance()); - EmbeddedChannel ch2 = new EmbeddedChannel(DefaultChannelId.newInstance()); - - channelGroup.add(ch1); - channelGroup.add(ch2); + cooperative.close(); - AtomicInteger closedCallbackCount = new AtomicInteger(0); + boolean completedCleanly = group.newCloseFuture().awaitUninterruptibly(50, TimeUnit.MILLISECONDS); - // Grace period 100ms: ch1 will close cooperatively, ch2 will be force closed - GracefulDisconnectLifecycle lifecycle = new GracefulDisconnectLifecycle(channelGroup, ch -> closedCallbackCount.incrementAndGet(), 100, 100); - - Future result = CompletableFuture.supplyAsync(() -> { - try - { - return lifecycle.run(); - } - catch (Exception e) - { - throw new RuntimeException(e); - } - }); - - ch1.close(); + assertThat(completedCleanly).isFalse(); - result.get(5, TimeUnit.SECONDS); + int unclosedCount = group.size(); + assertThat(unclosedCount).isEqualTo(1); - assertThat(closedCallbackCount.get()) - .as("Callback must be invoked exactly once per channel regardless of how it closed") - .isEqualTo(2); + group.close(); + assertThat(uncooperative.isOpen()).isFalse(); } -} \ No newline at end of file +} From 0694e8436721f4691c20945e608876b09cdf268c Mon Sep 17 00:00:00 2001 From: Rishabh Saraswat Date: Thu, 3 Sep 2026 11:54:43 +0530 Subject: [PATCH 4/6] added client metrics tests --- .../cassandra/metrics/ClientMetricsTest.java | 32 +++++++++++++------ 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/test/unit/org/apache/cassandra/metrics/ClientMetricsTest.java b/test/unit/org/apache/cassandra/metrics/ClientMetricsTest.java index 2f8131a05ecd..85b5d3040a4e 100644 --- a/test/unit/org/apache/cassandra/metrics/ClientMetricsTest.java +++ b/test/unit/org/apache/cassandra/metrics/ClientMetricsTest.java @@ -34,6 +34,7 @@ import org.apache.cassandra.auth.IAuthenticator.AuthenticationMode; import org.apache.cassandra.config.DatabaseDescriptor; import org.apache.cassandra.config.OverrideConfigurationLoader; +import org.apache.cassandra.service.CassandraDaemon; import org.apache.cassandra.service.EmbeddedCassandraService; import org.apache.cassandra.transport.TlsTestUtils; @@ -241,18 +242,29 @@ public void testConnectedClientsAndAuthMetrics() throws SSLException } @Test - public void testGracefulDisconnectMetrics() + public void testGracefulDisconnectMetricsViaServerLifecycle() throws Exception { - assertEquals(0, clientMetrics.connectionsDraining.get()); - long initialForcedDisconnects = clientMetrics.forcedDisconnects.getCount(); + boolean original = DatabaseDescriptor.getGracefulDisconnectEnabled(); + DatabaseDescriptor.getRawConfig().graceful_disconnect_enabled = true; - clientMetrics.incrementConnectionsDraining(); - assertEquals(1, clientMetrics.connectionsDraining.get()); - - clientMetrics.decrementConnectionsDraining(); - assertEquals(0, clientMetrics.connectionsDraining.get()); + try + { + assertEquals(0, clientMetrics.connectionsDraining.get()); + long initialForcedDisconnects = clientMetrics.forcedDisconnects.getCount(); - clientMetrics.markForcedDisconnect(3); - assertEquals(initialForcedDisconnects + 3, clientMetrics.forcedDisconnects.getCount()); + try (Cluster cluster = clusterBuilder().withCredentials("cassandra", "cassandra").build(); + Session ignored = cluster.connect()) + { + assertEquals(2, clientMetrics.connectedNativeClients.getValue().intValue()); + CassandraDaemon.getInstanceForTesting().nativeTransportService().stop(); + assertEquals(0, clientMetrics.connectionsDraining.get()); + assertEquals(initialForcedDisconnects, clientMetrics.forcedDisconnects.getCount()); + } + } + finally + { + DatabaseDescriptor.getRawConfig().graceful_disconnect_enabled = original; + CassandraDaemon.getInstanceForTesting().nativeTransportService().start(); + } } } From a325bad881f3848f91a4dd60827b412dc23ff0c6 Mon Sep 17 00:00:00 2001 From: Rishabh Saraswat Date: Thu, 3 Sep 2026 12:46:13 +0530 Subject: [PATCH 5/6] removed ambuguity from functions names and a race condition --- .../cassandra/config/DatabaseDescriptor.java | 7 +++---- .../cassandra/metrics/ClientMetrics.java | 6 ------ .../service/NativeTransportService.java | 5 ++--- .../cassandra/service/StorageService.java | 2 -- .../apache/cassandra/tools/nodetool/Drain.java | 3 +-- .../transport/InitialConnectionHandler.java | 3 ++- .../org/apache/cassandra/transport/Server.java | 8 ++++---- .../cassandra/transport/SimpleClient.java | 18 +++++++++++------- .../transport/messages/OptionsMessage.java | 3 ++- .../test/GracefulDisconnectTest.java | 17 ++++++----------- 10 files changed, 31 insertions(+), 41 deletions(-) diff --git a/src/java/org/apache/cassandra/config/DatabaseDescriptor.java b/src/java/org/apache/cassandra/config/DatabaseDescriptor.java index 065abaec1fae..5961942b1de9 100644 --- a/src/java/org/apache/cassandra/config/DatabaseDescriptor.java +++ b/src/java/org/apache/cassandra/config/DatabaseDescriptor.java @@ -2692,10 +2692,9 @@ public static long getGracefulDisconnectGracePeriod() public static void setGracefulDisconnectGracePeriod(long gracefulDisconnectGracePeriod) { - if (gracefulDisconnectGracePeriod > 0) - conf.graceful_disconnect_grace_period = new DurationSpec.LongMillisecondsBound(gracefulDisconnectGracePeriod); - else - throw new IllegalArgumentException(String.format("{} <= 0, not allowed, non positive values not allowed", gracefulDisconnectGracePeriod)); + if (gracefulDisconnectGracePeriod <= 0) + throw new IllegalArgumentException(String.format("graceful_disconnect_grace_period must be positive, got %d", gracefulDisconnectGracePeriod)); + conf.graceful_disconnect_grace_period = new DurationSpec.LongMillisecondsBound(gracefulDisconnectGracePeriod); } public static boolean getGracefulDisconnectEnabled() diff --git a/src/java/org/apache/cassandra/metrics/ClientMetrics.java b/src/java/org/apache/cassandra/metrics/ClientMetrics.java index 370f9f8e93b3..1d7ff0f467ed 100644 --- a/src/java/org/apache/cassandra/metrics/ClientMetrics.java +++ b/src/java/org/apache/cassandra/metrics/ClientMetrics.java @@ -177,11 +177,6 @@ public void markProtocolException() protocolException.mark(); } - public void incrementConnectionsDraining() - { - connectionsDraining.incrementAndGet(); - } - public void decrementConnectionsDraining() { connectionsDraining.decrementAndGet(); @@ -216,7 +211,6 @@ public synchronized void init(Server servers) registerGauge("ClientsByProtocolVersion", "clientsByProtocolVersion", this::recentClientStats); registerGauge("RequestsSize", ClientResourceLimits::getCurrentGlobalUsage); - connectionsDraining = new AtomicInteger(); registerGauge("ConnectionsDraining", connectionsDraining::get); forcedDisconnects = registerMeter("ForcedDisconnects"); diff --git a/src/java/org/apache/cassandra/service/NativeTransportService.java b/src/java/org/apache/cassandra/service/NativeTransportService.java index 1df9a68e13c2..eaccc6360ce2 100644 --- a/src/java/org/apache/cassandra/service/NativeTransportService.java +++ b/src/java/org/apache/cassandra/service/NativeTransportService.java @@ -37,7 +37,6 @@ import io.netty.channel.EventLoopGroup; import io.netty.channel.epoll.Epoll; import io.netty.channel.epoll.EpollEventLoopGroup; -import io.netty.channel.group.ChannelGroup; import io.netty.channel.nio.NioEventLoopGroup; import io.netty.util.Version; @@ -165,9 +164,9 @@ Server getServer() } @VisibleForTesting - public ChannelGroup getChannelsSubscribedToGracefulDisconnect() + public int getChannelsSubscribedToGracefulDisconnectCount() { - return server.getChannelsSubscribedToGracefulDisconnect(); + return server.getChannelsSubscribedToGracefulDisconnect().size(); } public void clearConnectionHistory() diff --git a/src/java/org/apache/cassandra/service/StorageService.java b/src/java/org/apache/cassandra/service/StorageService.java index bb16af8d1da8..6d937862904b 100644 --- a/src/java/org/apache/cassandra/service/StorageService.java +++ b/src/java/org/apache/cassandra/service/StorageService.java @@ -1209,8 +1209,6 @@ public boolean getGracefulDisconnectEnabled() @Override public void setGracefulDisconnectGracePeriod(long value) { - if (value <= 0 && DatabaseDescriptor.getGracefulDisconnectEnabled()) - throw new IllegalArgumentException("Graceful disconnect grace period must be positive when graceful disconnect is enabled. Got " + value); DatabaseDescriptor.setGracefulDisconnectGracePeriod(value); } diff --git a/src/java/org/apache/cassandra/tools/nodetool/Drain.java b/src/java/org/apache/cassandra/tools/nodetool/Drain.java index a71e7ffebfc1..fa61b4598960 100644 --- a/src/java/org/apache/cassandra/tools/nodetool/Drain.java +++ b/src/java/org/apache/cassandra/tools/nodetool/Drain.java @@ -34,8 +34,7 @@ public void execute(NodeProbe probe) try { probe.drain(); - } - catch (IOException | InterruptedException | ExecutionException e) + } catch (IOException | InterruptedException | ExecutionException e) { throw new RuntimeException("Error occurred during flushing", e); } diff --git a/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java b/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java index b3bb479a0b66..0e9e4771a44f 100644 --- a/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java +++ b/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java @@ -83,6 +83,7 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List li cqlVersions.add(QueryProcessor.CQL_VERSION.toString()); List compressions = new ArrayList<>(); + final List gracefulDisconnect = List.of("true"); if (Compressor.SnappyCompressor.instance != null) compressions.add("snappy"); // LZ4 is always available since worst case scenario it default to a pure JAVA implem. @@ -93,7 +94,7 @@ protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List li supportedOptions.put(StartupMessage.COMPRESSION, compressions); supportedOptions.put(StartupMessage.PROTOCOL_VERSIONS, ProtocolVersion.supportedVersions()); if (DatabaseDescriptor.getGracefulDisconnectEnabled()) - supportedOptions.put(StartupMessage.GRACEFUL_DISCONNECT, List.of("true")); + supportedOptions.put(StartupMessage.GRACEFUL_DISCONNECT, gracefulDisconnect); SupportedMessage supported = new SupportedMessage(supportedOptions); outbound = supported.encode(inbound.header.version, inbound.header.streamId); ctx.writeAndFlush(outbound); diff --git a/src/java/org/apache/cassandra/transport/Server.java b/src/java/org/apache/cassandra/transport/Server.java index 3e74c261513c..35091032faac 100644 --- a/src/java/org/apache/cassandra/transport/Server.java +++ b/src/java/org/apache/cassandra/transport/Server.java @@ -143,9 +143,7 @@ public void stop(boolean force) if (isRunning.compareAndSet(true, false)) { if (!force && DatabaseDescriptor.getGracefulDisconnectEnabled()) - { gracefulDisconnect(); - } close(force); } } @@ -158,9 +156,11 @@ private void gracefulDisconnect() if (channelGroup.isEmpty()) return; - ClientMetrics.instance.connectionsDraining.set(channelGroup.size()); channelGroup.forEach(channel -> - channel.closeFuture().addListener(future -> ClientMetrics.instance.decrementConnectionsDraining()) + { + ClientMetrics.instance.connectionsDraining.set(channelGroup.size()); + channel.closeFuture().addListener(future -> ClientMetrics.instance.decrementConnectionsDraining()); + } ); connectionTracker.send(new Event.GracefulDisconnect()); diff --git a/src/java/org/apache/cassandra/transport/SimpleClient.java b/src/java/org/apache/cassandra/transport/SimpleClient.java index 0a96294fe361..325d7c12ca88 100644 --- a/src/java/org/apache/cassandra/transport/SimpleClient.java +++ b/src/java/org/apache/cassandra/transport/SimpleClient.java @@ -28,6 +28,7 @@ import java.util.Map; import java.util.Queue; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.SynchronousQueue; // checkstyle: permit this import import java.util.concurrent.TimeUnit; @@ -121,6 +122,7 @@ public class SimpleClient implements Closeable protected Channel channel; protected ChannelFuture lastWriteFuture; private final AtomicBoolean draining = new AtomicBoolean(false); + private volatile CompletableFuture inFlight = CompletableFuture.completedFuture(null); protected String compression; public static class Builder @@ -323,12 +325,14 @@ public Message.Response execute(Message.Request request) return execute(request, true); } + public Message.Response execute(Message.Request request, boolean throwOnErrorResponse) { if (draining.get()) - { throw new RuntimeException("Connection is draining (GRACEFUL_DISCONNECT received)"); - } + + CompletableFuture requestCompletion = new CompletableFuture<>(); + inFlight = requestCompletion; try { request.attach(connection); @@ -344,6 +348,10 @@ public Message.Response execute(Message.Request request, boolean throwOnErrorRes { throw new UncheckedInterruptedException(e); } + finally + { + requestCompletion.complete(null); + } } public Map execute(List requests) @@ -408,11 +416,7 @@ public interface EventHandler private void handleGracefulDisconnect() { draining.set(true); - channel.eventLoop().execute(() -> { - ChannelFuture writeFuture = lastWriteFuture; - ChannelFuture closeAfter = (writeFuture != null) ? writeFuture : channel.newSucceededFuture(); - closeAfter.addListener(f -> channel.close()); - }); + inFlight.thenRun(() -> channel.eventLoop().execute(channel::close)); channel.closeFuture().addListener(f -> bootstrap.group().shutdownGracefully()); } diff --git a/src/java/org/apache/cassandra/transport/messages/OptionsMessage.java b/src/java/org/apache/cassandra/transport/messages/OptionsMessage.java index ee598afe1e9a..01de08bb7be1 100644 --- a/src/java/org/apache/cassandra/transport/messages/OptionsMessage.java +++ b/src/java/org/apache/cassandra/transport/messages/OptionsMessage.java @@ -66,6 +66,7 @@ protected Message.Response execute(QueryState state, Dispatcher.RequestTime requ cqlVersions.add(QueryProcessor.CQL_VERSION.toString()); List compressions = new ArrayList(); + final List gracefulDisconnect = List.of("true"); if (Compressor.SnappyCompressor.instance != null) compressions.add("snappy"); // LZ4 is always available since worst case scenario it default to a pure JAVA implem. @@ -75,7 +76,7 @@ protected Message.Response execute(QueryState state, Dispatcher.RequestTime requ supported.put(StartupMessage.CQL_VERSION, cqlVersions); supported.put(StartupMessage.COMPRESSION, compressions); if (DatabaseDescriptor.getGracefulDisconnectEnabled()) - supported.put(StartupMessage.GRACEFUL_DISCONNECT, List.of("true")); + supported.put(StartupMessage.GRACEFUL_DISCONNECT, gracefulDisconnect); supported.put(StartupMessage.PROTOCOL_VERSIONS, ProtocolVersion.supportedVersions()); return new SupportedMessage(supported); diff --git a/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java index 792a7db28b69..a544d6edfe4b 100644 --- a/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java +++ b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java @@ -120,8 +120,7 @@ public void testSubscriptionViaREGISTER() throws IOException int subscribedCount = cluster.get(1).callOnInstance(() -> CassandraDaemon.getInstanceForTesting() .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect() - .size()); + .getChannelsSubscribedToGracefulDisconnectCount()); assertThat(subscribedCount) .as("One channel should be subscribed to GRACEFUL_DISCONNECT") @@ -146,8 +145,7 @@ public void testRegisterGracefulDisconnectRejectedOnV4() throws IOException int subscribedCount = cluster.get(1).callOnInstance(() -> CassandraDaemon.getInstanceForTesting() .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect() - .size()); + .getChannelsSubscribedToGracefulDisconnectCount()); assertThat(subscribedCount) .as("V4 client should not be able to subscribe") @@ -170,8 +168,7 @@ public void testNonSubscribedClientDoesNotReceiveEvent() throws IOException int subscribedCount = cluster.get(1).callOnInstance(() -> CassandraDaemon.getInstanceForTesting() .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect() - .size()); + .getChannelsSubscribedToGracefulDisconnectCount()); assertThat(subscribedCount).isEqualTo(0); } } @@ -212,8 +209,7 @@ public void testNoEventEmittedWhenDisabled() throws IOException int subscribedCount = cluster.get(1).callOnInstance(() -> CassandraDaemon.getInstanceForTesting() .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect() - .size()); + .getChannelsSubscribedToGracefulDisconnectCount()); assertThat(subscribedCount).isEqualTo(0); @@ -262,8 +258,7 @@ public void testMultipleConnectionsCanSubscribe() throws IOException int subscribedCount = cluster.get(1).callOnInstance(() -> CassandraDaemon.getInstanceForTesting() .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect() - .size()); + .getChannelsSubscribedToGracefulDisconnectCount()); assertThat(subscribedCount).isEqualTo(2); } @@ -331,4 +326,4 @@ public void testCooperativeClientDisconnectsWithinGracePeriod() throws Exception assertThat(forcedCount).as("ForcedDisconnects metric should be 0 for a cooperative client").isEqualTo(0); } } -} \ No newline at end of file +} From e0adab2887aad48f250d124a943453dd36a3303f Mon Sep 17 00:00:00 2001 From: Rishabh Saraswat Date: Thu, 3 Sep 2026 13:10:18 +0530 Subject: [PATCH 6/6] changed Protocol Version of GRACEFUL_DISCONNECT from V5 to V4 and updated relevent files --- doc/native_protocol_v5.spec | 19 +++++++++++++++++++ .../org/apache/cassandra/transport/Event.java | 2 +- .../cassandra/transport/SimpleClient.java | 2 +- .../test/GracefulDisconnectTest.java | 12 ++++-------- 4 files changed, 25 insertions(+), 10 deletions(-) diff --git a/doc/native_protocol_v5.spec b/doc/native_protocol_v5.spec index b5543fa6f4cd..539903b5cd1a 100644 --- a/doc/native_protocol_v5.spec +++ b/doc/native_protocol_v5.spec @@ -698,6 +698,12 @@ Table of Contents for events on all connections, as this would only result in receiving multiple times the same event messages, wasting bandwidth. + Note: Unlike cluster-wide events (such as TOPOLOGY_CHANGE, STATUS_CHANGE, + and SCHEMA_CHANGE) which are distributed to any registered connection, the + "GRACEFUL_DISCONNECT" event is strictly connection-specific. To benefit + from graceful connection draining, clients must register for the "GRACEFUL_DISCONNECT" + event type on all connections they establish. + 4.2. Responses @@ -757,6 +763,11 @@ Table of Contents version description. For example: 3/v3, 4/v4, 5/v5-beta. If a version is in beta, it will have the word "beta" in its description. + - "GRACEFUL_DISCONNECT": if graceful disconnect is supported and enabled + by the server, this key will be present with the list of supported values + set to ["true"]. If the key is absent, the server does not support or has + disabled the graceful disconnect feature. + 4.2.5. RESULT @@ -987,6 +998,8 @@ Table of Contents consists of a [string] and an [inet], corresponding respectively to the type of change ("NEW_NODE" or "REMOVED_NODE") followed by the address of the new/removed node. + + - "STATUS_CHANGE": events related to change of node status. Currently, up/down events are sent. The body of the message (after the event type) consists of a [string] and an [inet], corresponding respectively to the @@ -1013,6 +1026,12 @@ Table of Contents - [string] keyspace containing the user defined function / aggregate - [string] the function/aggregate name - [string list] one string for each argument type (as CQL type) + - "GRACEFUL_DISCONNECT": events related to an impending clean shutdown of + the node. This event signals to the client that the server is shutting down + and will close the connection after the configured grace period. The body of + the message (after the event type) is empty. Upon receiving this event, + the client must stop sending new queries on this connection, let in-flight + queries finish, and cleanly close the socket. All EVENT messages have a streamId of -1 (Section 2.4.1.3). diff --git a/src/java/org/apache/cassandra/transport/Event.java b/src/java/org/apache/cassandra/transport/Event.java index 2d5a564f40b4..f56b9b98c4b4 100644 --- a/src/java/org/apache/cassandra/transport/Event.java +++ b/src/java/org/apache/cassandra/transport/Event.java @@ -37,7 +37,7 @@ public enum Type STATUS_CHANGE(ProtocolVersion.V3), SCHEMA_CHANGE(ProtocolVersion.V3), TRACE_COMPLETE(ProtocolVersion.V4), - GRACEFUL_DISCONNECT(ProtocolVersion.V5); + GRACEFUL_DISCONNECT(ProtocolVersion.V4); public final ProtocolVersion minimumVersion; diff --git a/src/java/org/apache/cassandra/transport/SimpleClient.java b/src/java/org/apache/cassandra/transport/SimpleClient.java index 325d7c12ca88..24951343274d 100644 --- a/src/java/org/apache/cassandra/transport/SimpleClient.java +++ b/src/java/org/apache/cassandra/transport/SimpleClient.java @@ -28,7 +28,7 @@ import java.util.Map; import java.util.Queue; import java.util.concurrent.BlockingQueue; -import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletableFuture; // checkstyle: permit this import import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.SynchronousQueue; // checkstyle: permit this import import java.util.concurrent.TimeUnit; diff --git a/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java index a544d6edfe4b..ee79e36d1c39 100644 --- a/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java +++ b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java @@ -130,7 +130,7 @@ public void testSubscriptionViaREGISTER() throws IOException } @Test - public void testRegisterGracefulDisconnectRejectedOnV4() throws IOException + public void testRegisterGracefulDisconnectAcceptedOnV4() throws IOException { try (Cluster cluster = buildCluster(1, true)) { @@ -138,18 +138,14 @@ public void testRegisterGracefulDisconnectRejectedOnV4() throws IOException try (SimpleClient client = SimpleClient.builder(nativeAddr.getHostString(), 9042).protocolVersion(ProtocolVersion.V4).build()) { client.connect(false); - - assertThatThrownBy(() -> client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT)))) - .hasCauseInstanceOf(org.apache.cassandra.transport.ProtocolException.class); - + client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); int subscribedCount = cluster.get(1).callOnInstance(() -> CassandraDaemon.getInstanceForTesting() .nativeTransportService() .getChannelsSubscribedToGracefulDisconnectCount()); - assertThat(subscribedCount) - .as("V4 client should not be able to subscribe") - .isEqualTo(0); + .as("V4 client should be able to subscribe") + .isEqualTo(1); } } }