From 58181ecb1214cd9d2e1a18b21bd495d8b4771505 Mon Sep 17 00:00:00 2001 From: Abhishek Pal Date: Sun, 2 Aug 2026 15:14:35 +0530 Subject: [PATCH] RATIS-2630. Netty data-stream shutdown leaks in-flight streams --- .../netty/server/DataStreamManagement.java | 2 + .../apache/ratis/netty/server/StreamMap.java | 7 + .../server/TestDataStreamManagement.java | 128 ++++++++++++++++++ 3 files changed, 137 insertions(+) diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/server/DataStreamManagement.java b/ratis-netty/src/main/java/org/apache/ratis/netty/server/DataStreamManagement.java index 005e0b8e61..aaae801733 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/server/DataStreamManagement.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/server/DataStreamManagement.java @@ -254,6 +254,8 @@ void shutdown() { timeout -> LOG.warn("{}: requestExecutor shutdown timeout in {}", this, timeout)); ConcurrentUtils.shutdownAndWait(TimeDuration.ONE_SECOND, writeExecutor, timeout -> LOG.warn("{}: writeExecutor shutdown timeout in {}", this, timeout)); + // The executors are stopped, so no more reads will add to the stream map. + cleanUp(streams.keySet()); } private CompletableFuture stream(RaftClientRequest request, StateMachine stateMachine) { diff --git a/ratis-netty/src/main/java/org/apache/ratis/netty/server/StreamMap.java b/ratis-netty/src/main/java/org/apache/ratis/netty/server/StreamMap.java index 073698cb80..949496d328 100644 --- a/ratis-netty/src/main/java/org/apache/ratis/netty/server/StreamMap.java +++ b/ratis-netty/src/main/java/org/apache/ratis/netty/server/StreamMap.java @@ -21,6 +21,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.HashSet; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.function.Function; @@ -52,4 +54,9 @@ STREAM remove(ClientInvocationId key) { LOG.debug("remove({}) returns {}", key, info); return info; } + + /** @return a snapshot of the keys currently in this map. */ + Set keySet() { + return new HashSet<>(map.keySet()); + } } diff --git a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java index fd7fae6bc2..b795785736 100644 --- a/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java +++ b/ratis-test/src/test/java/org/apache/ratis/netty/server/TestDataStreamManagement.java @@ -27,6 +27,7 @@ import org.apache.ratis.netty.metrics.NettyServerStreamRpcMetrics; import org.apache.ratis.proto.RaftProtos.DataStreamPacketHeaderProto.Type; import org.apache.ratis.protocol.ClientId; +import org.apache.ratis.protocol.ClientInvocationId; import org.apache.ratis.protocol.DataStreamReply; import org.apache.ratis.protocol.Message; import org.apache.ratis.protocol.RaftClientReply; @@ -35,9 +36,12 @@ import org.apache.ratis.protocol.RaftPeer; import org.apache.ratis.protocol.RaftPeerId; import org.apache.ratis.protocol.exceptions.ReadIndexException; +import org.apache.ratis.server.DataStreamMap; import org.apache.ratis.server.RaftServer; import org.apache.ratis.statemachine.StateMachine; import org.apache.ratis.statemachine.StateMachine.DataApi; +import org.apache.ratis.statemachine.StateMachine.DataChannel; +import org.apache.ratis.statemachine.StateMachine.DataStream; import org.apache.ratis.statemachine.impl.BaseStateMachine; import org.apache.ratis.thirdparty.com.google.protobuf.ByteString; import org.apache.ratis.thirdparty.io.netty.buffer.ByteBuf; @@ -63,6 +67,8 @@ import java.util.List; import java.util.Set; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; @@ -494,6 +500,128 @@ void readCleansChannelMapOnEarlyException() throws Exception { } } + @Test + void shutdownReleasesInFlightStreams() throws Exception { + final RaftPeerId serverId = RaftPeerId.valueOf("s1"); + final RaftGroupId groupId = RaftGroupId.randomId(); + final ClientId clientId = ClientId.randomId(); + + final AtomicBoolean cleanedUp = new AtomicBoolean(); + final DataChannel dataChannel = new DataChannel() { + private volatile boolean open = true; + @Override public void force(boolean metadata) { } + @Override public int write(ByteBuffer src) { + final int remaining = src.remaining(); + src.position(src.limit()); + return remaining; + } + @Override public boolean isOpen() { return open; } + @Override public void close() { open = false; } + }; + final DataStream dataStream = new DataStream() { + @Override public DataChannel getDataChannel() { return dataChannel; } + @Override public CompletableFuture cleanUp() { + cleanedUp.set(true); + return CompletableFuture.completedFuture(null); + } + }; + final DataApi dataApi = new DataApi() { + @Override public CompletableFuture stream(RaftClientRequest request) { + return CompletableFuture.completedFuture(dataStream); + } + }; + final StateMachine stateMachine = new BaseStateMachine() { + @Override public DataApi data() { return dataApi; } + }; + + // A real DataStreamMap so we can observe the stream being registered and later released. + final ConcurrentMap> entries = new ConcurrentHashMap<>(); + final DataStreamMap dataStreamMap = new DataStreamMap() { + @Override public CompletableFuture computeIfAbsent(ClientInvocationId id, + Function> f) { + return entries.computeIfAbsent(id, f); + } + @Override public CompletableFuture remove(ClientInvocationId id) { + return entries.remove(id); + } + }; + + final RaftServer.Division division = newStreamingDivision(serverId, stateMachine, dataStreamMap); + final RaftServer server = newRaftServer(serverId, new RaftProperties(), groupId, division); + final NettyServerStreamRpcMetrics metrics = new NettyServerStreamRpcMetrics("s1"); + final DataStreamManagement management = new DataStreamManagement(server, metrics); + + final EmbeddedChannel embeddedChannel = new EmbeddedChannel(new ChannelInboundHandlerAdapter()); + final ChannelHandlerContext ctx = embeddedChannel.pipeline().firstContext(); + assertNotNull(ctx, "ChannelHandlerContext should be initialized"); + + // Non-primary header (server id != request server id) so there are no remote successors. + final RaftClientRequest raftClientRequest = RaftClientRequest.newBuilder() + .setClientId(clientId) + .setServerId(RaftPeerId.valueOf("s2")) + .setGroupId(groupId) + .setCallId(1L) + .setMessage(Message.valueOf(ByteString.copyFromUtf8("header"))) + .setType(RaftClientRequest.writeRequestType()) + .build(); + final ByteBuffer header = ClientProtoUtils.toRaftClientRequestProtoByteBuffer(raftClientRequest); + final DataStreamRequestByteBuf request = new DataStreamRequestByteBuf( + clientId, + Type.STREAM_HEADER, + raftClientRequest.getCallId(), + 0L, + Collections.singletonList(StandardWriteOption.FLUSH), + Unpooled.wrappedBuffer(header)); + + final CheckedBiFunction, Set, IOException> getStreams = + (r, p) -> Collections.emptySet(); + + try { + management.read(request, ctx, getStreams); + // The header registers an in-flight stream that is NOT removed (no CLOSE was sent). + JavaUtils.attempt(() -> assertEquals(1, entries.size()), 10, + TimeDuration.valueOf(100, TimeUnit.MILLISECONDS), "stream registered", null); + + management.shutdown(); + + assertTrue(entries.isEmpty(), "in-flight stream should be unregistered on shutdown"); + assertTrue(cleanedUp.get(), "DataStream.cleanUp should be invoked on shutdown"); + } finally { + embeddedChannel.finishAndReleaseAll(); + } + } + + private static RaftServer.Division newStreamingDivision( + RaftPeerId serverId, StateMachine stateMachine, DataStreamMap dataStreamMap) { + return (RaftServer.Division) Proxy.newProxyInstance(RaftServer.Division.class.getClassLoader(), + new Class[]{RaftServer.Division.class}, + (proxy, method, args) -> { + switch (method.getName()) { + case "getStateMachine": + return stateMachine; + case "getDataStreamMap": + return dataStreamMap; + case "getId": + return serverId; + case "getCommitInfos": + return Collections.emptyList(); + case "getRaftConf": + // Not dereferenced for a non-primary request without a routing table. + return null; + case "close": + return null; + case "toString": + return stateMachine.toString(); + case "hashCode": + return System.identityHashCode(proxy); + case "equals": + return proxy == args[0]; + default: + throw new UnsupportedOperationException(method.toString()); + } + }); + } + private static class ReadOnlyRequest { private final DataStreamRequestByteBuf request; private final ByteBuf headerBuf;