[SPARK-26418][SHUFFLE] Only OpenBlocks without any ChunkFetch for one stream will cause memory leak in ExternalShuffleService

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

@@ -90,7 +90,6 @@ protected void channelRead0(
     ManagedBuffer buf;
     try {
       streamManager.checkAuthorization(client, msg.streamChunkId.streamId);
-      streamManager.registerChannel(channel, msg.streamChunkId.streamId);
       buf = streamManager.getChunk(msg.streamChunkId.streamId, msg.streamChunkId.chunkIndex);
     } catch (Exception e) {
       logger.error(String.format("Error opening block %s for request from %s",

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

@@ -23,6 +23,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicLong;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import io.netty.channel.Channel;
 import org.apache.commons.lang3.tuple.ImmutablePair;
@@ -202,4 +203,8 @@ public long registerStream(String appId, Iterator<ManagedBuffer> buffers) {
     return myStreamId;
   }
 
+  @VisibleForTesting
+  public long getStreamsSize() {
+    return streams.size();
+  }
 }

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalShuffleBlockHandler.java

@@ -92,6 +92,7 @@ protected void handleMessage(
         checkAuth(client, msg.appId);
         long streamId = streamManager.registerStream(client.getClientId(),
           new ManagedBufferIterator(msg.appId, msg.execId, msg.blockIds));
+        streamManager.registerChannel(client.getChannel(), streamId);
         if (logger.isTraceEnabled()) {
           logger.trace("Registered streamId {} with {} buffers for client {} from host {}",
                        streamId,

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/StreamStatesCleanupSuite.java

@@ -0,0 +1,66 @@
+/*
+ * 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.spark.network.shuffle;
+
+import io.netty.channel.Channel;
+import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NioManagedBuffer;
+import org.apache.spark.network.client.RpcResponseCallback;
+import org.apache.spark.network.client.TransportClient;
+import org.apache.spark.network.client.TransportResponseHandler;
+import org.apache.spark.network.server.OneForOneStreamManager;
+import org.apache.spark.network.server.RpcHandler;
+import org.apache.spark.network.shuffle.protocol.OpenBlocks;
+import org.junit.Test;
+
+import java.nio.ByteBuffer;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class StreamStatesCleanupSuite {
+
+  @Test
+  public void testStreamsAreRemovedCorrectly() {
+    OneForOneStreamManager streamManager = new OneForOneStreamManager();
+    ExternalShuffleBlockResolver blockResolver = mock(ExternalShuffleBlockResolver.class);
+    TransportClient reverseClient
+      = new TransportClient(mock(Channel.class), mock(TransportResponseHandler.class));
+    RpcHandler handler = new ExternalShuffleBlockHandler(streamManager, blockResolver);
+
+    ManagedBuffer block0Marker = new NioManagedBuffer(ByteBuffer.wrap(new byte[3]));
+    ManagedBuffer block1Marker = new NioManagedBuffer(ByteBuffer.wrap(new byte[7]));
+    when(blockResolver.getBlockData("app0", "exec1", 0, 0, 0))
+      .thenReturn(block0Marker);
+    when(blockResolver.getBlockData("app0", "exec1", 0, 0, 1))
+      .thenReturn(block1Marker);
+    ByteBuffer openBlocks = new OpenBlocks("app0", "exec1",
+      new String[]{"shuffle_0_0_0", "shuffle_0_0_1"})
+      .toByteBuffer();
+
+    RpcResponseCallback callback = mock(RpcResponseCallback.class);
+
+    // Open blocks
+    handler.receive(reverseClient, openBlocks, callback);
+    // Connection closed before any FetchChunk request received
+    streamManager.connectionTerminated(reverseClient.getChannel());
+
+    assert streamManager.getStreamsSize() == 0;
+  }
+
+}