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/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/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..5961942b1de9 100644 --- a/src/java/org/apache/cassandra/config/DatabaseDescriptor.java +++ b/src/java/org/apache/cassandra/config/DatabaseDescriptor.java @@ -2685,6 +2685,23 @@ 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) + 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() + { + 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..1d7ff0f467ed 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,16 @@ public void markProtocolException() protocolException.mark(); } + public void decrementConnectionsDraining() + { + connectionsDraining.decrementAndGet(); + } + + public void markForcedDisconnect(int forceDisconnectedClients) + { + forcedDisconnects.mark(forceDisconnectedClients); + } + public void markSSLHandshakeException() { sslHandshakeException.mark(); @@ -197,6 +211,9 @@ public synchronized void init(Server servers) registerGauge("ClientsByProtocolVersion", "clientsByProtocolVersion", this::recentClientStats); registerGauge("RequestsSize", ClientResourceLimits::getCurrentGlobalUsage); + 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..eaccc6360ce2 100644 --- a/src/java/org/apache/cassandra/service/NativeTransportService.java +++ b/src/java/org/apache/cassandra/service/NativeTransportService.java @@ -163,6 +163,12 @@ Server getServer() return server; } + @VisibleForTesting + public int getChannelsSubscribedToGracefulDisconnectCount() + { + return server.getChannelsSubscribedToGracefulDisconnect().size(); + } + 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..6d937862904b 100644 --- a/src/java/org/apache/cassandra/service/StorageService.java +++ b/src/java/org/apache/cassandra/service/StorageService.java @@ -1199,6 +1199,25 @@ public long getRpcTimeout() return DatabaseDescriptor.getRpcTimeout(MILLISECONDS); } + + @Override + public boolean getGracefulDisconnectEnabled() + { + return DatabaseDescriptor.getGracefulDisconnectEnabled(); + } + + @Override + public void setGracefulDisconnectGracePeriod(long 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/transport/Event.java b/src/java/org/apache/cassandra/transport/Event.java index a1209a13eaf5..f56b9b98c4b4 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.V4); 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..0e9e4771a44f 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; @@ -82,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. @@ -91,6 +93,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, 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 d7af40eb96f7..35091032faac 100644 --- a/src/java/org/apache/cassandra/transport/Server.java +++ b/src/java/org/apache/cassandra/transport/Server.java @@ -44,6 +44,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 +81,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 +128,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 +140,48 @@ 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(); + if (channelGroup.isEmpty()) + return; + + channelGroup.forEach(channel -> + { + ClientMetrics.instance.connectionsDraining.set(channelGroup.size()); + channel.closeFuture().addListener(future -> ClientMetrics.instance.decrementConnectionsDraining()); + } + ); + + connectionTracker.send(new Event.GracefulDisconnect()); + + long gracePeriod = DatabaseDescriptor.getGracefulDisconnectGracePeriod(); + int forcedDisconnects = 0; + + boolean completedCleanly = channelGroup.newCloseFuture().awaitUninterruptibly(gracePeriod, TimeUnit.MILLISECONDS); + + if (!completedCleanly) + { + 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() @@ -149,6 +196,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 +205,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 +393,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..24951343274d 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; // checkstyle: permit this import import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.SynchronousQueue; // checkstyle: permit this import import java.util.concurrent.TimeUnit; @@ -112,7 +113,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 @@ -120,7 +121,8 @@ public class SimpleClient implements Closeable protected Bootstrap bootstrap; 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 @@ -170,6 +172,7 @@ public SimpleClient build() } } + public static Builder builder(String host, int port) { return new Builder(host, port); @@ -322,8 +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); @@ -339,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) @@ -400,6 +413,13 @@ public interface EventHandler void onEvent(Event event); } + private void handleGracefulDisconnect() + { + draining.set(true); + inFlight.thenRun(() -> channel.eventLoop().execute(channel::close)); + channel.closeFuture().addListener(f -> bootstrap.group().shutdownGracefully()); + } + public static class SimpleEventHandler implements EventHandler { public final BlockingQueue queue = newBlockingQueue(); @@ -701,7 +721,12 @@ protected void initChannel(Channel channel) throws Exception 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,8 +744,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 +935,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..01de08bb7be1 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; @@ -65,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. @@ -73,6 +75,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, gracefulDisconnect); 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..ee79e36d1c39 --- /dev/null +++ b/test/distributed/org/apache/cassandra/distributed/test/GracefulDisconnectTest.java @@ -0,0 +1,325 @@ +/* + * 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() + .getChannelsSubscribedToGracefulDisconnectCount()); + + assertThat(subscribedCount) + .as("One channel should be subscribed to GRACEFUL_DISCONNECT") + .isEqualTo(1); + } + } + } + + @Test + public void testRegisterGracefulDisconnectAcceptedOnV4() 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); + 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 be able to subscribe") + .isEqualTo(1); + } + } + } + + @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() + .getChannelsSubscribedToGracefulDisconnectCount()); + 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() + .getChannelsSubscribedToGracefulDisconnectCount()); + + 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() + .getChannelsSubscribedToGracefulDisconnectCount()); + + 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); + } + } +} 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..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; @@ -239,4 +240,31 @@ public void testConnectedClientsAndAuthMetrics() throws SSLException assertEquals(0, clientMetrics.encryptedConnectedNativeClients.getValue().intValue()); assertEquals(0, passwordConnections.getValue().intValue()); } + + @Test + public void testGracefulDisconnectMetricsViaServerLifecycle() throws Exception + { + boolean original = DatabaseDescriptor.getGracefulDisconnectEnabled(); + DatabaseDescriptor.getRawConfig().graceful_disconnect_enabled = true; + + try + { + assertEquals(0, clientMetrics.connectionsDraining.get()); + long initialForcedDisconnects = 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(); + } + } } 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/GracefulDisconnectTest.java b/test/unit/org/apache/cassandra/transport/GracefulDisconnectTest.java new file mode 100644 index 000000000000..24931e6c0e06 --- /dev/null +++ b/test/unit/org/apache/cassandra/transport/GracefulDisconnectTest.java @@ -0,0 +1,107 @@ +/* + * 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 org.junit.BeforeClass; +import org.junit.Test; + +import org.apache.cassandra.config.DatabaseDescriptor; + +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 GracefulDisconnectTest +{ + @BeforeClass + public static void setup() + { + DatabaseDescriptor.daemonInitialization(); + } + + @Test + public void testEmptyChannelGroupGracefulHandling() + { + 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 testChannelGroupAutomaticRemovalOnClose() + { + ChannelGroup group = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + EmbeddedChannel channel = new EmbeddedChannel(DefaultChannelId.newInstance()); + group.add(channel); + + assertThat(group.size()).isEqualTo(1); + channel.close(); + assertThat(group.size()).isEqualTo(0); + } + + @Test + public void testCloseFutureTriggersWhenAllChannelsTerminate() throws Exception + { + ChannelGroup group = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + EmbeddedChannel ch1 = new EmbeddedChannel(DefaultChannelId.newInstance()); + EmbeddedChannel ch2 = new EmbeddedChannel(DefaultChannelId.newInstance()); + group.add(ch1); + group.add(ch2); + + Future closeFutureCompletion = CompletableFuture.supplyAsync(() -> group.newCloseFuture().awaitUninterruptibly(3000, TimeUnit.MILLISECONDS)); + + assertThat(closeFutureCompletion.isDone()).isFalse(); + + ch1.close(); + ch2.close(); + + assertThat(closeFutureCompletion.get(2, TimeUnit.SECONDS)).isTrue(); + } + + @Test + public void testGracePeriodTimeoutTriggersBulkForceClose() + { + ChannelGroup group = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE); + EmbeddedChannel cooperative = new EmbeddedChannel(DefaultChannelId.newInstance()); + EmbeddedChannel uncooperative = new EmbeddedChannel(DefaultChannelId.newInstance()); + group.add(cooperative); + group.add(uncooperative); + + cooperative.close(); + + boolean completedCleanly = group.newCloseFuture().awaitUninterruptibly(50, TimeUnit.MILLISECONDS); + + assertThat(completedCleanly).isFalse(); + + int unclosedCount = group.size(); + assertThat(unclosedCount).isEqualTo(1); + + group.close(); + assertThat(uncooperative.isOpen()).isFalse(); + } +}