RATIS-2197. Clean remote stream to resolve direct memory leak

symious · Nov 21, 2024 · d2b563e · d2b563e
1 parent 06affdf
commit d2b563e
Showing 1 changed file with 8 additions and 2 deletions.
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
@@ -390,8 +390,13 @@ static void sendDataStreamException(Throwable throwable, DataStreamRequestByteBu
   void cleanUp(Set<ClientInvocationId> ids) {
     for (ClientInvocationId clientInvocationId : ids) {
       Optional.ofNullable(streams.remove(clientInvocationId))
-          .map(StreamInfo::getLocal)
-          .ifPresent(LocalStream::cleanUp);
+          .ifPresent(streamInfo -> {
+            streamInfo.getDivision()
+                .getDataStreamMap()
+                .remove(clientInvocationId);
+            streamInfo.getLocal().cleanUp();
+            streamInfo.applyToRemotes(out -> out.out.closeAsync());
+          });
     }
   }
 
@@ -423,6 +428,7 @@ private void removeDataStream(ClientInvocationId invocationId, StreamInfo info)
     if (info != null) {
       info.getDivision().getDataStreamMap().remove(invocationId);
       info.getLocal().cleanUp();
+      info.applyToRemotes(out -> out.out.closeAsync());
     }
   }