[SPARK-30246][CORE] OneForOneStreamManager might leak memory in connectionTerminated

common/network-common/src/main/java/org/apache/spark/network/server/OneForOneStreamManager.java

@@ -18,6 +18,7 @@
 package org.apache.spark.network.server;
 
 import java.util.Iterator;
+import java.util.LinkedList;
 import java.util.Map;
 import java.util.Random;
 import java.util.concurrent.ConcurrentHashMap;
@@ -117,18 +118,21 @@ public static Pair<Long, Integer> parseStreamChunkId(String streamChunkId) {
 
   @Override
   public void connectionTerminated(Channel channel) {
+    LinkedList<StreamState> removedStates = new LinkedList<>();
     // Close all streams which have been associated with the channel.
     for (Map.Entry<Long, StreamState> entry: streams.entrySet()) {
       StreamState state = entry.getValue();
       if (state.associatedChannel == channel) {
-        streams.remove(entry.getKey());
-
-        // Release all remaining buffers.
-        while (state.buffers.hasNext()) {
-          ManagedBuffer buffer = state.buffers.next();
-          if (buffer != null) {
-            buffer.release();
-          }
+        removedStates.add(streams.remove(entry.getKey()));
+      }
+    }
+
+    for (StreamState state: removedStates) {
+      // Release all remaining buffers.
+      while (state.buffers.hasNext()) {
+        ManagedBuffer buffer = state.buffers.next();
+        if (buffer != null) {
+          buffer.release();
         }
       }
     }

common/network-common/src/test/java/org/apache/spark/network/server/OneForOneStreamManagerSuite.java

@@ -18,6 +18,7 @@
 package org.apache.spark.network.server;
 
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 
 import io.netty.channel.Channel;
@@ -96,4 +97,30 @@ public void managedBuffersAreFreedWhenConnectionIsClosed() {
     Mockito.verify(buffer2, Mockito.times(1)).release();
     Assert.assertEquals(0, manager.numStreamStates());
   }
+
+  @Test
+  public void streamStatesAreFreedWhenConnectionIsClosedEvenIfBufferIteratorThrowsException() {
+    OneForOneStreamManager manager = new OneForOneStreamManager();
+
+    Iterator<ManagedBuffer> buffers = Mockito.mock(Iterator.class);
+    Mockito.when(buffers.hasNext()).thenReturn(true);
+    Mockito.when(buffers.next()).thenThrow(RuntimeException.class);
+
+    Channel dummyChannel = Mockito.mock(Channel.class, Mockito.RETURNS_SMART_NULLS);
+    manager.registerStream("appId", buffers, dummyChannel);
+    manager.registerStream("appId", buffers, dummyChannel);
+
+    Assert.assertEquals(2, manager.numStreamStates());
+
+    try {
+      manager.connectionTerminated(dummyChannel);
+      Assert.fail("connectionTerminated should throws RuntimeException in buffers.next()");
+    } catch (RuntimeException e) {
+      Mockito.verify(buffers, Mockito.times(1)).hasNext();
+      Mockito.verify(buffers, Mockito.times(1)).next();
+
+      Assert.assertEquals(0, manager.numStreamStates());
+    }
+  }
+
 }