From e75ba3c598cb12eefebb617dd262cfbaac7cfc25 Mon Sep 17 00:00:00 2001 From: omniCoder77 Date: Sun, 8 Mar 2026 10:05:59 +0530 Subject: [PATCH 1/3] cep 59 --- 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 | 15 + .../cassandra/metrics/ClientMetrics.java | 23 + .../service/NativeTransportService.java | 6 + .../cassandra/service/StorageService.java | 150 +++++- .../service/StorageServiceMBean.java | 7 +- .../org/apache/cassandra/tools/NodeProbe.java | 2 +- .../cassandra/tools/nodetool/Drain.java | 3 +- .../org/apache/cassandra/transport/Event.java | 27 +- .../transport/InitialConnectionHandler.java | 2 + .../apache/cassandra/transport/Server.java | 19 + .../cassandra/transport/SimpleClient.java | 50 +- .../transport/messages/OptionsMessage.java | 3 + .../transport/messages/StartupMessage.java | 1 + .../test/GracefulDisconnectTest.java | 464 ++++++++++++++++++ ...ggableScheduledThreadPoolExecutorTest.java | 2 +- .../config/DatabaseDescriptorTest.java | 23 + .../cassandra/service/StorageServiceTest.java | 27 + 21 files changed, 931 insertions(+), 15 deletions(-) create mode 100644 doc/modules/cassandra/pages/managing/operating/graceful_disconnect.adoc create mode 100644 test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java 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..1aa4862e9e2c 100644 --- a/src/java/org/apache/cassandra/config/DatabaseDescriptor.java +++ b/src/java/org/apache/cassandra/config/DatabaseDescriptor.java @@ -2685,6 +2685,21 @@ 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) + { + conf.graceful_disconnect_grace_period = new DurationSpec.LongMillisecondsBound(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..e3cc38d16f40 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; + + 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..807fdf09df1b 100644 --- a/src/java/org/apache/cassandra/service/StorageService.java +++ b/src/java/org/apache/cassandra/service/StorageService.java @@ -40,13 +40,17 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -158,6 +162,7 @@ import org.apache.cassandra.locator.Replicas; import org.apache.cassandra.locator.SnitchAdapter; import org.apache.cassandra.locator.SystemReplicas; +import org.apache.cassandra.metrics.ClientMetrics; import org.apache.cassandra.metrics.Sampler; import org.apache.cassandra.metrics.SamplingManager; import org.apache.cassandra.metrics.StorageMetrics; @@ -219,7 +224,9 @@ import org.apache.cassandra.tcm.transformations.Startup; import org.apache.cassandra.tcm.transformations.Unregister; import org.apache.cassandra.transport.ClientResourceLimits; +import org.apache.cassandra.transport.Event; import org.apache.cassandra.transport.ProtocolVersion; +import org.apache.cassandra.transport.messages.EventMessage; import org.apache.cassandra.utils.ExecutorUtils; import org.apache.cassandra.utils.FBUtilities; import org.apache.cassandra.utils.JVMStabilityInspector; @@ -239,6 +246,9 @@ import org.apache.cassandra.utils.progress.jmx.JMXBroadcastExecutor; import org.apache.cassandra.utils.progress.jmx.JMXProgressSupport; +import io.netty.channel.group.ChannelGroup; +import io.netty.util.AttributeKey; + import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; @@ -292,9 +302,13 @@ public class StorageService extends NotificationBroadcasterSupport implements IE { private static final Logger logger = LoggerFactory.getLogger(StorageService.class); + private static ReentrantLock drainLock = new ReentrantLock(); + public static final int INDEFINITE = -1; public static final int RING_DELAY_MILLIS = getRingDelay(); // delay after which we assume ring has stablized + static final AttributeKey> EVENT_DISPATCHER = AttributeKey.valueOf("EVTDISP"); + { PathUtils.setDeletionListener(path -> { if (isDaemonSetupCompleted()) @@ -617,7 +631,9 @@ public void stopNativeTransport(boolean force) { throw new IllegalStateException("No configured daemon"); } - daemon.stopNativeTransport(force); + gracefulDisconnect(() -> { + daemon.stopNativeTransport(force); + }); } public boolean isNativeTransportRunning() @@ -715,7 +731,9 @@ public void stopDaemon() { if (daemon == null) throw new IllegalStateException("No configured daemon"); - daemon.deactivate(); + gracefulDisconnect(() -> { + daemon.deactivate(); + }); } // for testing only @@ -1199,6 +1217,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); @@ -3875,9 +3914,112 @@ public String getDrainProgress() /** * Shuts node off to writes, empties memtables and the commit log. */ - public synchronized void drain() throws IOException, InterruptedException, ExecutionException + public void drain() throws IOException, InterruptedException, ExecutionException, TimeoutException + { + if (!drainLock.tryLock()) + { + throw new IllegalStateException("Drain already in progress"); + } + try + { + CountDownLatch drainComplete = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + gracefulDisconnect(() -> { + try + { + drain(false); + } + catch (IOException | InterruptedException | ExecutionException e) + { + failure.set(e); + } + finally + { + drainComplete.countDown(); + } + }); + if (!drainComplete.await(DatabaseDescriptor.getGracefulDisconnectGracePeriod(), MILLISECONDS)) + throw new TimeoutException("Timed out waiting for drain to complete after graceful disconnect"); + Exception e = failure.get(); + if (e instanceof IOException) throw (IOException) e; + if (e instanceof InterruptedException) throw (InterruptedException) e; + if (e instanceof ExecutionException) throw (ExecutionException) e; + if (e != null) throw new RuntimeException(e); + } + finally + { + drainLock.unlock(); + } + } + + private void gracefulDisconnect(Runnable defaultAction) + { + NativeTransportService nts = CassandraDaemon.instance.nativeTransportService(); + if (nts == null || nts.getServer() == null) + { + defaultAction.run(); + return; + } + nts.getServer().stopAcceptingNewConnections(); + ChannelGroup channelGroup = nts.getServer().getChannelsSubscribedToGracefulDisconnect(); + if (channelGroup == null || channelGroup.isEmpty()) + { + defaultAction.run(); + return; + } + gracefulDisconnect(defaultAction, channelGroup); + } + + public void gracefulDisconnect(Runnable defaultAction, ChannelGroup channelGroup) { - drain(false); + if (channelGroup == null || channelGroup.isEmpty()) + { + defaultAction.run(); + return; + } + AtomicBoolean actionStarted = new AtomicBoolean(false); + AtomicInteger connectedChannels = new AtomicInteger(channelGroup.size()); + Runnable runOnceAction = () -> { + if (actionStarted.compareAndSet(false, true)) + defaultAction.run(); + }; + Runnable onTimeout = () -> { + if (actionStarted.get()) + return; + int remaining = connectedChannels.get(); + if (remaining > 0) + { + ClientMetrics.instance.markForcedDisconnect(remaining); + channelGroup.close(); + } + else + { + runOnceAction.run(); + } + }; + ScheduledFuture timeoutTask = ScheduledExecutors.nonPeriodicTasks + .schedule(onTimeout, DatabaseDescriptor.getGracefulDisconnectGracePeriod(), MILLISECONDS); + EventMessage gracefulDisconnectEventMessage = new EventMessage(new Event.GracefulDisconnect()); + channelGroup.forEach(channel -> { + Consumer dispatcher = channel.attr(EVENT_DISPATCHER).get(); + if (dispatcher != null) + dispatcher.accept(gracefulDisconnectEventMessage); + + channel.closeFuture().addListener(future -> { + ClientMetrics.instance.decrementConnectionsDraining(); + if (connectedChannels.decrementAndGet() == 0) + { + timeoutTask.cancel(false); + runOnceAction.run(); + } + }); + ClientMetrics.instance.incrementConnectionsDraining(); + }); + if (connectedChannels.get() == 0) + { + timeoutTask.cancel(false); + runOnceAction.run(); + } } protected synchronized void drain(boolean isFinalShutdown) throws IOException, InterruptedException, ExecutionException diff --git a/src/java/org/apache/cassandra/service/StorageServiceMBean.java b/src/java/org/apache/cassandra/service/StorageServiceMBean.java index abc868903127..13437e5ec896 100644 --- a/src/java/org/apache/cassandra/service/StorageServiceMBean.java +++ b/src/java/org/apache/cassandra/service/StorageServiceMBean.java @@ -640,7 +640,7 @@ default int upgradeSSTables(String keyspaceName, boolean excludeCurrentVersion, public String getDrainProgress(); /** makes node unavailable for writes, flushes memtables and replays commitlog. */ - public void drain() throws IOException, InterruptedException, ExecutionException; + public void drain() throws IOException, InterruptedException, ExecutionException, TimeoutException; /** * Truncates (deletes) the given table from the provided keyspace. @@ -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/NodeProbe.java b/src/java/org/apache/cassandra/tools/NodeProbe.java index c6e774f5aaea..41be861ba229 100644 --- a/src/java/org/apache/cassandra/tools/NodeProbe.java +++ b/src/java/org/apache/cassandra/tools/NodeProbe.java @@ -772,7 +772,7 @@ public AuthCacheMBean getAuthCacheMBean(String cacheName) } } - public void drain() throws IOException, InterruptedException, ExecutionException + public void drain() throws IOException, InterruptedException, ExecutionException, TimeoutException { ssProxy.drain(); } diff --git a/src/java/org/apache/cassandra/tools/nodetool/Drain.java b/src/java/org/apache/cassandra/tools/nodetool/Drain.java index fa61b4598960..f2c987e79a61 100644 --- a/src/java/org/apache/cassandra/tools/nodetool/Drain.java +++ b/src/java/org/apache/cassandra/tools/nodetool/Drain.java @@ -20,6 +20,7 @@ import java.io.IOException; import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeoutException; import org.apache.cassandra.tools.NodeProbe; @@ -34,7 +35,7 @@ public void execute(NodeProbe probe) try { probe.drain(); - } catch (IOException | InterruptedException | ExecutionException e) + } catch (IOException | InterruptedException | TimeoutException | 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/InitialConnectionHandler.java b/src/java/org/apache/cassandra/transport/InitialConnectionHandler.java index 28299363fb7e..7deb9770ff00 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,7 @@ 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()); + supportedOptions.put(StartupMessage.GRACEFUL_DISCONNECT, List.of(String.valueOf(DatabaseDescriptor.getGracefulDisconnectEnabled()))); 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..193123b63afb 100644 --- a/src/java/org/apache/cassandra/transport/Server.java +++ b/src/java/org/apache/cassandra/transport/Server.java @@ -80,6 +80,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 +127,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); @@ -149,6 +155,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 +164,16 @@ public synchronized void start() isRunning.set(true); } + 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 +347,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..e2784cf3daed 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; @@ -100,6 +100,8 @@ public class SimpleClient implements Closeable public static final int TIMEOUT_SECONDS = 10; + private final AtomicBoolean draining = new AtomicBoolean(false); + static { InternalLoggerFactory.setDefaultFactory(new Slf4JLoggerFactory()); @@ -112,7 +114,7 @@ public class SimpleClient implements Closeable 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; // We don't track connection really, so we don't need one Connection per channel @@ -238,8 +240,9 @@ public SimpleClient connect(boolean useCompression, boolean throwOnOverload) thr options.put(StartupMessage.COMPRESSION, "LZ4"); connection.setCompressor(Compressor.LZ4Compressor.instance); } + if (version.isGreaterOrEqualTo(ProtocolVersion.V5)) + options.put(StartupMessage.GRACEFUL_DISCONNECT, "GRACEFUL_DISCONNECT"); execute(new StartupMessage(options)); - return this; } @@ -324,6 +327,10 @@ 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); @@ -639,6 +646,18 @@ private void configureLegacyPipeline(ChannelHandlerContext ctx) } } + private void handleGracefulDisconnect() + { + draining.set(true); + ChannelFuture writeFuture = lastWriteFuture; + if (writeFuture != null) + writeFuture.addListener(f -> channel.close()); + else + channel.close(); + channel.closeFuture().addListener(f -> bootstrap.group().shutdownGracefully()); + } + + @ChannelHandler.Sharable static class MessageBatchEncoder extends MessageToMessageEncoder> { @@ -700,6 +719,14 @@ protected void initChannel(Channel channel) throws Exception @ChannelHandler.Sharable static class ResponseHandler extends SimpleChannelInboundHandler { + + private final SimpleClient client; + + public ResponseHandler(SimpleClient client) + { + this.client = client; + } + public final BlockingQueue responses = new SynchronousQueue<>(true); public EventHandler eventHandler; @@ -719,8 +746,19 @@ public void handleResponse(Channel channel, Message.Response r) if (r instanceof EventMessage) { + 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(((EventMessage) r).event); + eventHandler.onEvent(event); } else responses.put(r); @@ -899,4 +937,4 @@ private ChannelFuture[] writeLargeMessage(ChannelHandlerContext ctx, Envelope f) return futures.toArray(EMPTY_FUTURES_ARRAY); } } -} +} \ 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..6f6fe8717bd2 --- /dev/null +++ b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java @@ -0,0 +1,464 @@ +/* + * 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.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import org.awaitility.Awaitility; +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.metrics.ClientMetrics; +import org.apache.cassandra.service.CassandraDaemon; +import org.apache.cassandra.service.StorageService; +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 io.netty.channel.group.ChannelGroup; + +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 testGracefulDisconnectEventReceivedOnDrain() throws IOException, InterruptedException + { + 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()) + { + SimpleClient.SimpleEventHandler handler = new SimpleClient.SimpleEventHandler(); + client.setEventHandler(handler); + client.connect(false); + client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + cluster.get(1).runOnInstance(() -> { + ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect(); + StorageService.instance.gracefulDisconnect(() -> { + }, channelGroup); + }); + + Event event = handler.queue.poll(10, TimeUnit.SECONDS); + assertThat(event).isNotNull(); + assertThat(event.type).isEqualTo(Event.Type.GRACEFUL_DISCONNECT); + } + } + } + + @Test + public void testDefaultActionCalledAfterAllConnectionsClose() 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); + client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + client.close(); + Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> { + int drainingCount = cluster.get(1).callOnInstance(() -> + ClientMetrics.instance.connectionsDraining.get()); + assertThat(drainingCount).isEqualTo(0); + }); + + boolean actionCalled = cluster.get(1).callOnInstance(() -> { + AtomicBoolean called = new AtomicBoolean(false); + ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect(); + StorageService.instance.gracefulDisconnect(() -> called.set(true), channelGroup); + return called.get(); + }); + + assertThat(actionCalled).isTrue(); + } + } + } + + @Test + public void testDefaultActionCalledAfterMaxDrainMs() throws IOException, InterruptedException + { + 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); + client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + cluster.get(1).runOnInstance(() -> { + StorageService.instance.setGracefulDisconnectGracePeriod(3000); + }); + + cluster.get(1).runOnInstance(() -> { + ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect(); + StorageService.instance.gracefulDisconnect(() -> { + }, channelGroup); + }); + + Thread.sleep(5000); + + boolean finishedDraining = cluster.get(1).callOnInstance(() -> ClientMetrics.instance.connectionsDraining.get() == 0); + assertThat(finishedDraining).isTrue(); + } + finally + { + cluster.get(1).runOnInstance(() -> { + StorageService.instance.setGracefulDisconnectGracePeriod(5000); + }); + } + } + } + + @Test + public void testMultipleConnectionsAllReceiveEvent() throws IOException, InterruptedException + { + 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()) + { + SimpleClient.SimpleEventHandler handler1 = new SimpleClient.SimpleEventHandler(); + SimpleClient.SimpleEventHandler handler2 = new SimpleClient.SimpleEventHandler(); + client1.setEventHandler(handler1); + client2.setEventHandler(handler2); + + 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))); + + cluster.get(1).runOnInstance(() -> { + ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect(); + StorageService.instance.gracefulDisconnect(() -> { + }, channelGroup); + }); + + Event event1 = handler1.queue.poll(10, TimeUnit.SECONDS); + Event event2 = handler2.queue.poll(10, TimeUnit.SECONDS); + + assertThat(event1).isNotNull(); + assertThat(event2).isNotNull(); + assertThat(event1.type).isEqualTo(Event.Type.GRACEFUL_DISCONNECT); + assertThat(event2.type).isEqualTo(Event.Type.GRACEFUL_DISCONNECT); + } + } + } + + @Test + public void testMixedFleetOnlySubscribedReceivesEvent() throws IOException, InterruptedException + { + try (Cluster cluster = buildCluster(1, true)) + { + InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); + try (SimpleClient subscribedClient = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V5) + .build(); + SimpleClient nonSubscribedClient = SimpleClient.builder(nativeAddr.getHostString(), 9042) + .protocolVersion(ProtocolVersion.V4) + .build()) + { + SimpleClient.SimpleEventHandler subscribedHandler = new SimpleClient.SimpleEventHandler(); + SimpleClient.SimpleEventHandler nonSubscribedHandler = new SimpleClient.SimpleEventHandler(); + subscribedClient.setEventHandler(subscribedHandler); + nonSubscribedClient.setEventHandler(nonSubscribedHandler); + + subscribedClient.connect(false); + nonSubscribedClient.connect(false); + subscribedClient.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); + + cluster.get(1).runOnInstance(() -> { + ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() + .nativeTransportService() + .getChannelsSubscribedToGracefulDisconnect(); + StorageService.instance.gracefulDisconnect(() -> { + }, channelGroup); + }); + + Event subscribedEvent = subscribedHandler.queue.poll(10, TimeUnit.SECONDS); + Event nonSubscribedEvent = nonSubscribedHandler.queue.poll(3, TimeUnit.SECONDS); + + assertThat(subscribedEvent).isNotNull(); + assertThat(subscribedEvent.type).isEqualTo(Event.Type.GRACEFUL_DISCONNECT); + assertThat(nonSubscribedEvent).isNull(); + } + } + } +} \ No newline at end of file diff --git a/test/unit/org/apache/cassandra/concurrent/DebuggableScheduledThreadPoolExecutorTest.java b/test/unit/org/apache/cassandra/concurrent/DebuggableScheduledThreadPoolExecutorTest.java index ff3be281f383..0a05d283966d 100644 --- a/test/unit/org/apache/cassandra/concurrent/DebuggableScheduledThreadPoolExecutorTest.java +++ b/test/unit/org/apache/cassandra/concurrent/DebuggableScheduledThreadPoolExecutorTest.java @@ -56,7 +56,7 @@ public static void tearDown() } @Test - public void testShutdown() throws ExecutionException, InterruptedException, IOException + public void testShutdown() throws ExecutionException, InterruptedException, IOException, TimeoutException { ScheduledExecutorPlus testPool = executorFactory().scheduled("testpool"); diff --git a/test/unit/org/apache/cassandra/config/DatabaseDescriptorTest.java b/test/unit/org/apache/cassandra/config/DatabaseDescriptorTest.java index c3bdd00e21ae..a16f295ab636 100644 --- a/test/unit/org/apache/cassandra/config/DatabaseDescriptorTest.java +++ b/test/unit/org/apache/cassandra/config/DatabaseDescriptorTest.java @@ -817,6 +817,29 @@ public void testRowIndexSizeWarnEnabledAbortDisabled() DatabaseDescriptor.applyThresholdsValidations(conf); } + @Test + public void testGracefulDisconnectEnabled() + { + Config config = new Config(); + Assert.assertFalse("Default value of graceful_disconnect_enabled must be false", config.graceful_disconnect_enabled); + } + + @Test + public void testGracefulDisconnectGracePeriod() + { + long originalValue = DatabaseDescriptor.getGracefulDisconnectGracePeriod(); + Assert.assertEquals("Default value of graceful_disconnect_grace_period must be 5000", 5000, originalValue); + try + { + DatabaseDescriptor.setGracefulDisconnectGracePeriod(3000); + Assert.assertEquals("graceful_disconnect_grace_period should be updated to 3000", 3000, DatabaseDescriptor.getGracefulDisconnectGracePeriod()); + } + finally + { + DatabaseDescriptor.setGracefulDisconnectGracePeriod(originalValue); + } + } + @Test public void testRowIndexSizeAbortEnabledWarnDisabled() { 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() { From bcefd0dbe03638d1ac3c33896add35f2bda98613 Mon Sep 17 00:00:00 2001 From: omniCoder77 Date: Sun, 8 Mar 2026 10:05:59 +0530 Subject: [PATCH 2/3] cep 59 # Conflicts: # src/java/org/apache/cassandra/service/StorageService.java --- src/java/org/apache/cassandra/service/StorageService.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/java/org/apache/cassandra/service/StorageService.java b/src/java/org/apache/cassandra/service/StorageService.java index 807fdf09df1b..9d7caa96e6ec 100644 --- a/src/java/org/apache/cassandra/service/StorageService.java +++ b/src/java/org/apache/cassandra/service/StorageService.java @@ -40,7 +40,7 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.CountDownLatch; +import org.apache.cassandra.utils.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ThreadLocalRandom; @@ -3922,7 +3922,7 @@ public void drain() throws IOException, InterruptedException, ExecutionException } try { - CountDownLatch drainComplete = new CountDownLatch(1); + CountDownLatch drainComplete = CountDownLatch.newCountDownLatch(1); AtomicReference failure = new AtomicReference<>(); gracefulDisconnect(() -> { try @@ -3935,7 +3935,7 @@ public void drain() throws IOException, InterruptedException, ExecutionException } finally { - drainComplete.countDown(); + drainComplete.decrement(); } }); if (!drainComplete.await(DatabaseDescriptor.getGracefulDisconnectGracePeriod(), MILLISECONDS)) From c0641da72d2410767cd406cb7087f36674549b5c Mon Sep 17 00:00:00 2001 From: Rishabh Saraswat Date: Thu, 20 Aug 2026 00:08:05 +0530 Subject: [PATCH 3/3] creating a new FSM based model --- .../service/GracefulDisconnectLifecycle.java | 156 ++++++++++++++ .../cassandra/service/StorageService.java | 140 +++---------- .../apache/cassandra/transport/Server.java | 5 + .../cassandra/transport/SimpleClient.java | 10 +- .../test/GracefulDisconnectTest.java | 194 ------------------ .../GracefulDisconnectLifecycleTest.java | 152 ++++++++++++++ 6 files changed, 347 insertions(+), 310 deletions(-) create mode 100644 src/java/org/apache/cassandra/service/GracefulDisconnectLifecycle.java create mode 100644 test/unit/org/apache/cassandra/service/GracefulDisconnectLifecycleTest.java diff --git a/src/java/org/apache/cassandra/service/GracefulDisconnectLifecycle.java b/src/java/org/apache/cassandra/service/GracefulDisconnectLifecycle.java new file mode 100644 index 000000000000..75df44a5a529 --- /dev/null +++ b/src/java/org/apache/cassandra/service/GracefulDisconnectLifecycle.java @@ -0,0 +1,156 @@ +/* + * 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.service; + +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.Event; +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 io.netty.util.AttributeKey; + +final class GracefulDisconnectLifecycle +{ + static final AttributeKey> EVENT_DISPATCHER = + AttributeKey.valueOf("EVTDISP"); + + // 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/service/StorageService.java b/src/java/org/apache/cassandra/service/StorageService.java index 9d7caa96e6ec..f6bb7f155075 100644 --- a/src/java/org/apache/cassandra/service/StorageService.java +++ b/src/java/org/apache/cassandra/service/StorageService.java @@ -40,17 +40,13 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; -import org.apache.cassandra.utils.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; -import java.util.concurrent.ScheduledFuture; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; -import java.util.concurrent.locks.ReentrantLock; -import java.util.function.Consumer; import java.util.function.Supplier; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -224,9 +220,7 @@ import org.apache.cassandra.tcm.transformations.Startup; import org.apache.cassandra.tcm.transformations.Unregister; import org.apache.cassandra.transport.ClientResourceLimits; -import org.apache.cassandra.transport.Event; import org.apache.cassandra.transport.ProtocolVersion; -import org.apache.cassandra.transport.messages.EventMessage; import org.apache.cassandra.utils.ExecutorUtils; import org.apache.cassandra.utils.FBUtilities; import org.apache.cassandra.utils.JVMStabilityInspector; @@ -247,7 +241,6 @@ import org.apache.cassandra.utils.progress.jmx.JMXProgressSupport; import io.netty.channel.group.ChannelGroup; -import io.netty.util.AttributeKey; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; @@ -302,13 +295,9 @@ public class StorageService extends NotificationBroadcasterSupport implements IE { private static final Logger logger = LoggerFactory.getLogger(StorageService.class); - private static ReentrantLock drainLock = new ReentrantLock(); - public static final int INDEFINITE = -1; public static final int RING_DELAY_MILLIS = getRingDelay(); // delay after which we assume ring has stablized - static final AttributeKey> EVENT_DISPATCHER = AttributeKey.valueOf("EVTDISP"); - { PathUtils.setDeletionListener(path -> { if (isDaemonSetupCompleted()) @@ -631,9 +620,15 @@ public void stopNativeTransport(boolean force) { throw new IllegalStateException("No configured daemon"); } - gracefulDisconnect(() -> { - daemon.stopNativeTransport(force); - }); + try + { + gracefulDisconnect(); + } + catch (InterruptedException e) + { + throw new RuntimeException(e); + } + daemon.stopNativeTransport(force); } public boolean isNativeTransportRunning() @@ -731,9 +726,15 @@ public void stopDaemon() { if (daemon == null) throw new IllegalStateException("No configured daemon"); - gracefulDisconnect(() -> { - daemon.deactivate(); - }); + try + { + gracefulDisconnect(); + } + catch (InterruptedException e) + { + throw new RuntimeException(e); + } + daemon.deactivate(); } // for testing only @@ -3914,111 +3915,28 @@ public String getDrainProgress() /** * Shuts node off to writes, empties memtables and the commit log. */ - public void drain() throws IOException, InterruptedException, ExecutionException, TimeoutException + public synchronized void drain() throws IOException, InterruptedException, ExecutionException, TimeoutException { - if (!drainLock.tryLock()) - { - throw new IllegalStateException("Drain already in progress"); - } - try - { - CountDownLatch drainComplete = CountDownLatch.newCountDownLatch(1); - AtomicReference failure = new AtomicReference<>(); - gracefulDisconnect(() -> { - try - { - drain(false); - } - catch (IOException | InterruptedException | ExecutionException e) - { - failure.set(e); - } - finally - { - drainComplete.decrement(); - } - }); - if (!drainComplete.await(DatabaseDescriptor.getGracefulDisconnectGracePeriod(), MILLISECONDS)) - throw new TimeoutException("Timed out waiting for drain to complete after graceful disconnect"); - Exception e = failure.get(); - if (e instanceof IOException) throw (IOException) e; - if (e instanceof InterruptedException) throw (InterruptedException) e; - if (e instanceof ExecutionException) throw (ExecutionException) e; - if (e != null) throw new RuntimeException(e); - } - finally - { - drainLock.unlock(); - } + gracefulDisconnect(); + drain(false); } - private void gracefulDisconnect(Runnable defaultAction) + private void gracefulDisconnect() throws InterruptedException { NativeTransportService nts = CassandraDaemon.instance.nativeTransportService(); - if (nts == null || nts.getServer() == null) - { - defaultAction.run(); - return; - } nts.getServer().stopAcceptingNewConnections(); + ChannelGroup channelGroup = nts.getServer().getChannelsSubscribedToGracefulDisconnect(); - if (channelGroup == null || channelGroup.isEmpty()) - { - defaultAction.run(); - return; - } - gracefulDisconnect(defaultAction, channelGroup); - } + ClientMetrics.instance.connectionsDraining.set(channelGroup.size()); - public void gracefulDisconnect(Runnable defaultAction, ChannelGroup channelGroup) - { - if (channelGroup == null || channelGroup.isEmpty()) + try { - defaultAction.run(); - return; + int forcedDisconnects = new GracefulDisconnectLifecycle(channelGroup, channel -> ClientMetrics.instance.decrementConnectionsDraining()).run(); + ClientMetrics.instance.markForcedDisconnect(forcedDisconnects); } - AtomicBoolean actionStarted = new AtomicBoolean(false); - AtomicInteger connectedChannels = new AtomicInteger(channelGroup.size()); - Runnable runOnceAction = () -> { - if (actionStarted.compareAndSet(false, true)) - defaultAction.run(); - }; - Runnable onTimeout = () -> { - if (actionStarted.get()) - return; - int remaining = connectedChannels.get(); - if (remaining > 0) - { - ClientMetrics.instance.markForcedDisconnect(remaining); - channelGroup.close(); - } - else - { - runOnceAction.run(); - } - }; - ScheduledFuture timeoutTask = ScheduledExecutors.nonPeriodicTasks - .schedule(onTimeout, DatabaseDescriptor.getGracefulDisconnectGracePeriod(), MILLISECONDS); - EventMessage gracefulDisconnectEventMessage = new EventMessage(new Event.GracefulDisconnect()); - channelGroup.forEach(channel -> { - Consumer dispatcher = channel.attr(EVENT_DISPATCHER).get(); - if (dispatcher != null) - dispatcher.accept(gracefulDisconnectEventMessage); - - channel.closeFuture().addListener(future -> { - ClientMetrics.instance.decrementConnectionsDraining(); - if (connectedChannels.decrementAndGet() == 0) - { - timeoutTask.cancel(false); - runOnceAction.run(); - } - }); - ClientMetrics.instance.incrementConnectionsDraining(); - }); - if (connectedChannels.get() == 0) + catch (TimeoutException e) { - timeoutTask.cancel(false); - runOnceAction.run(); + throw new RuntimeException(e); } } diff --git a/src/java/org/apache/cassandra/transport/Server.java b/src/java/org/apache/cassandra/transport/Server.java index 193123b63afb..21ea89f5adc8 100644 --- a/src/java/org/apache/cassandra/transport/Server.java +++ b/src/java/org/apache/cassandra/transport/Server.java @@ -164,6 +164,11 @@ 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()) diff --git a/src/java/org/apache/cassandra/transport/SimpleClient.java b/src/java/org/apache/cassandra/transport/SimpleClient.java index e2784cf3daed..a1b6d92439c4 100644 --- a/src/java/org/apache/cassandra/transport/SimpleClient.java +++ b/src/java/org/apache/cassandra/transport/SimpleClient.java @@ -121,7 +121,7 @@ public class SimpleClient implements Closeable protected Connection connection; protected Bootstrap bootstrap; protected Channel channel; - protected ChannelFuture lastWriteFuture; + private volatile ChannelFuture lastWriteFuture; protected String compression; @@ -649,11 +649,11 @@ private void configureLegacyPipeline(ChannelHandlerContext ctx) private void handleGracefulDisconnect() { draining.set(true); + channel.eventLoop().execute(() -> { ChannelFuture writeFuture = lastWriteFuture; - if (writeFuture != null) - writeFuture.addListener(f -> channel.close()); - else - channel.close(); + ChannelFuture closeAfter = (writeFuture != null) ? writeFuture : channel.newSucceededFuture(); + closeAfter.addListener(f -> channel.close()); + }); channel.closeFuture().addListener(f -> bootstrap.group().shutdownGracefully()); } diff --git a/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java index 6f6fe8717bd2..33187e85da73 100644 --- a/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java +++ b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java @@ -21,19 +21,14 @@ import java.io.IOException; import java.net.InetSocketAddress; import java.util.Collections; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import org.awaitility.Awaitility; 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.metrics.ClientMetrics; import org.apache.cassandra.service.CassandraDaemon; -import org.apache.cassandra.service.StorageService; import org.apache.cassandra.transport.Event; import org.apache.cassandra.transport.Message; import org.apache.cassandra.transport.ProtocolVersion; @@ -44,8 +39,6 @@ import org.apache.cassandra.transport.messages.StartupMessage; import org.apache.cassandra.transport.messages.SupportedMessage; -import io.netty.channel.group.ChannelGroup; - import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; @@ -69,7 +62,6 @@ public Cluster buildCluster(int nodeCount, boolean gracefulDisconnectEnabled) th .start(); } - @Test public void testGracefulDisconnectAdvertisedWhenEnabled() throws IOException { @@ -275,190 +267,4 @@ public void testMultipleConnectionsCanSubscribe() throws IOException } } } - - @Test - public void testGracefulDisconnectEventReceivedOnDrain() throws IOException, InterruptedException - { - 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()) - { - SimpleClient.SimpleEventHandler handler = new SimpleClient.SimpleEventHandler(); - client.setEventHandler(handler); - client.connect(false); - client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); - - cluster.get(1).runOnInstance(() -> { - ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() - .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect(); - StorageService.instance.gracefulDisconnect(() -> { - }, channelGroup); - }); - - Event event = handler.queue.poll(10, TimeUnit.SECONDS); - assertThat(event).isNotNull(); - assertThat(event.type).isEqualTo(Event.Type.GRACEFUL_DISCONNECT); - } - } - } - - @Test - public void testDefaultActionCalledAfterAllConnectionsClose() 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); - client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); - - client.close(); - Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> { - int drainingCount = cluster.get(1).callOnInstance(() -> - ClientMetrics.instance.connectionsDraining.get()); - assertThat(drainingCount).isEqualTo(0); - }); - - boolean actionCalled = cluster.get(1).callOnInstance(() -> { - AtomicBoolean called = new AtomicBoolean(false); - ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() - .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect(); - StorageService.instance.gracefulDisconnect(() -> called.set(true), channelGroup); - return called.get(); - }); - - assertThat(actionCalled).isTrue(); - } - } - } - - @Test - public void testDefaultActionCalledAfterMaxDrainMs() throws IOException, InterruptedException - { - 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); - client.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); - - cluster.get(1).runOnInstance(() -> { - StorageService.instance.setGracefulDisconnectGracePeriod(3000); - }); - - cluster.get(1).runOnInstance(() -> { - ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() - .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect(); - StorageService.instance.gracefulDisconnect(() -> { - }, channelGroup); - }); - - Thread.sleep(5000); - - boolean finishedDraining = cluster.get(1).callOnInstance(() -> ClientMetrics.instance.connectionsDraining.get() == 0); - assertThat(finishedDraining).isTrue(); - } - finally - { - cluster.get(1).runOnInstance(() -> { - StorageService.instance.setGracefulDisconnectGracePeriod(5000); - }); - } - } - } - - @Test - public void testMultipleConnectionsAllReceiveEvent() throws IOException, InterruptedException - { - 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()) - { - SimpleClient.SimpleEventHandler handler1 = new SimpleClient.SimpleEventHandler(); - SimpleClient.SimpleEventHandler handler2 = new SimpleClient.SimpleEventHandler(); - client1.setEventHandler(handler1); - client2.setEventHandler(handler2); - - 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))); - - cluster.get(1).runOnInstance(() -> { - ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() - .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect(); - StorageService.instance.gracefulDisconnect(() -> { - }, channelGroup); - }); - - Event event1 = handler1.queue.poll(10, TimeUnit.SECONDS); - Event event2 = handler2.queue.poll(10, TimeUnit.SECONDS); - - assertThat(event1).isNotNull(); - assertThat(event2).isNotNull(); - assertThat(event1.type).isEqualTo(Event.Type.GRACEFUL_DISCONNECT); - assertThat(event2.type).isEqualTo(Event.Type.GRACEFUL_DISCONNECT); - } - } - } - - @Test - public void testMixedFleetOnlySubscribedReceivesEvent() throws IOException, InterruptedException - { - try (Cluster cluster = buildCluster(1, true)) - { - InetSocketAddress nativeAddr = cluster.get(1).config().broadcastAddress(); - try (SimpleClient subscribedClient = SimpleClient.builder(nativeAddr.getHostString(), 9042) - .protocolVersion(ProtocolVersion.V5) - .build(); - SimpleClient nonSubscribedClient = SimpleClient.builder(nativeAddr.getHostString(), 9042) - .protocolVersion(ProtocolVersion.V4) - .build()) - { - SimpleClient.SimpleEventHandler subscribedHandler = new SimpleClient.SimpleEventHandler(); - SimpleClient.SimpleEventHandler nonSubscribedHandler = new SimpleClient.SimpleEventHandler(); - subscribedClient.setEventHandler(subscribedHandler); - nonSubscribedClient.setEventHandler(nonSubscribedHandler); - - subscribedClient.connect(false); - nonSubscribedClient.connect(false); - subscribedClient.execute(new RegisterMessage(Collections.singletonList(Event.Type.GRACEFUL_DISCONNECT))); - - cluster.get(1).runOnInstance(() -> { - ChannelGroup channelGroup = CassandraDaemon.getInstanceForTesting() - .nativeTransportService() - .getChannelsSubscribedToGracefulDisconnect(); - StorageService.instance.gracefulDisconnect(() -> { - }, channelGroup); - }); - - Event subscribedEvent = subscribedHandler.queue.poll(10, TimeUnit.SECONDS); - Event nonSubscribedEvent = nonSubscribedHandler.queue.poll(3, TimeUnit.SECONDS); - - assertThat(subscribedEvent).isNotNull(); - assertThat(subscribedEvent.type).isEqualTo(Event.Type.GRACEFUL_DISCONNECT); - assertThat(nonSubscribedEvent).isNull(); - } - } - } } \ No newline at end of file diff --git a/test/unit/org/apache/cassandra/service/GracefulDisconnectLifecycleTest.java b/test/unit/org/apache/cassandra/service/GracefulDisconnectLifecycleTest.java new file mode 100644 index 000000000000..88b9fb2e5e39 --- /dev/null +++ b/test/unit/org/apache/cassandra/service/GracefulDisconnectLifecycleTest.java @@ -0,0 +1,152 @@ +/* + * 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.service; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +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.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 (InterruptedException | TimeoutException 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 (InterruptedException | TimeoutException 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 (InterruptedException | TimeoutException 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(); + } +} \ No newline at end of file