[SPARK-26604][CORE] Clean up channel registration for StreamManager

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

@@ -90,8 +90,8 @@ 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);
+      buf = streamManager.getChunk(msg.streamChunkId.streamId, msg.streamChunkId.chunkIndex,
+        channel);
     } catch (Exception e) {
       logger.error(String.format("Error opening block %s for request from %s",
         msg.streamChunkId, getRemoteAddress(channel)), e);

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

@@ -71,15 +71,22 @@ public OneForOneStreamManager() {
     streams = new ConcurrentHashMap<>();
   }
 
-  @Override
+
+  /**
+   * Associates a stream with a single client connection, which is guaranteed to be the only reader
+   * of the stream. Once the connection is closed, the stream will never be used again, enabling
+   * cleanup by `connectionTerminated`.
+   */
   public void registerChannel(Channel channel, long streamId) {
     if (streams.containsKey(streamId)) {
       streams.get(streamId).associatedChannel = channel;
     }
   }
 
   @Override
-  public ManagedBuffer getChunk(long streamId, int chunkIndex) {
+  public ManagedBuffer getChunk(long streamId, int chunkIndex, Channel channel) {
+    registerChannel(channel, streamId);
+
     StreamState state = streams.get(streamId);
     if (chunkIndex != state.curChunk) {
       throw new IllegalStateException(String.format(
@@ -100,9 +107,9 @@ public ManagedBuffer getChunk(long streamId, int chunkIndex) {
   }
 
   @Override
-  public ManagedBuffer openStream(String streamChunkId) {
+  public ManagedBuffer openStream(String streamChunkId, Channel channel) {
     Pair<Long, Integer> streamChunkIdPair = parseStreamChunkId(streamChunkId);
-    return getChunk(streamChunkIdPair.getLeft(), streamChunkIdPair.getRight());
+    return getChunk(streamChunkIdPair.getLeft(), streamChunkIdPair.getRight(), channel);
   }
 
   public static String genStreamChunkId(long streamId, int chunkId) {

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

@@ -42,9 +42,10 @@ public abstract class StreamManager {
    * The returned ManagedBuffer will be release()'d after being written to the network.
    *
    * @param streamId id of a stream that has been previously registered with the StreamManager.
+   * @param channel The connection used to serve chunk request.
    * @param chunkIndex 0-indexed chunk of the stream that's requested
    */
-  public abstract ManagedBuffer getChunk(long streamId, int chunkIndex);
+  public abstract ManagedBuffer getChunk(long streamId, int chunkIndex, Channel channel);
 
   /**
    * Called in response to a stream() request. The returned data is streamed to the client
@@ -54,22 +55,13 @@ public abstract class StreamManager {
    * {@link #getChunk(long, int)} method.
    *
    * @param streamId id of a stream that has been previously registered with the StreamManager.
+   * @param channel The connection used to serve stream request.
    * @return A managed buffer for the stream, or null if the stream was not found.
    */
-  public ManagedBuffer openStream(String streamId) {
+  public ManagedBuffer openStream(String streamId, Channel channel) {
     throw new UnsupportedOperationException();
   }
 
-  /**
-   * Associates a stream with a single client connection, which is guaranteed to be the only reader
-   * of the stream. The getChunk() method will be called serially on this connection and once the
-   * connection is closed, the stream will never be used again, enabling cleanup.
-   *
-   * This must be called before the first getChunk() on the stream, but it may be invoked multiple
-   * times with the same channel and stream id.
-   */
-  public void registerChannel(Channel channel, long streamId) { }
-
   /**
    * Indicates that the given channel has been terminated. After this occurs, we are guaranteed not
    * to read from the associated streams again, so any state can be cleaned up.

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

@@ -126,7 +126,7 @@ private void processStreamRequest(final StreamRequest req) {
     }
     ManagedBuffer buf;
     try {
-      buf = streamManager.openStream(req.streamId);
+      buf = streamManager.openStream(req.streamId, channel);
     } catch (Exception e) {
       logger.error(String.format(
         "Error opening stream %s for request from %s", req.streamId, getRemoteAddress(channel)), e);

common/network-common/src/test/java/org/apache/spark/network/ChunkFetchIntegrationSuite.java

@@ -32,6 +32,7 @@
 
 import com.google.common.collect.Sets;
 import com.google.common.io.Closeables;
+import io.netty.channel.Channel;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -92,7 +93,7 @@ public static void setUp() throws Exception {
 
     streamManager = new StreamManager() {
       @Override
-      public ManagedBuffer getChunk(long streamId, int chunkIndex) {
+      public ManagedBuffer getChunk(long streamId, int chunkIndex, Channel channel) {
         assertEquals(STREAM_ID, streamId);
         if (chunkIndex == BUFFER_CHUNK_INDEX) {
           return new NioManagedBuffer(buf);

common/network-common/src/test/java/org/apache/spark/network/ChunkFetchRequestHandlerSuite.java

@@ -65,7 +65,6 @@ public void handleChunkFetchRequest() throws Exception {
     managedBuffers.add(new TestManagedBuffer(30));
     managedBuffers.add(new TestManagedBuffer(40));
     long streamId = streamManager.registerStream("test-app", managedBuffers.iterator());
-    streamManager.registerChannel(channel, streamId);
     TransportClient reverseClient = mock(TransportClient.class);
     ChunkFetchRequestHandler requestHandler = new ChunkFetchRequestHandler(reverseClient,
       rpcHandler.getStreamManager(), 2L);

common/network-common/src/test/java/org/apache/spark/network/RequestTimeoutIntegrationSuite.java

@@ -18,6 +18,7 @@
 package org.apache.spark.network;
 
 import com.google.common.util.concurrent.Uninterruptibles;
+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.ChunkReceivedCallback;
@@ -65,7 +66,7 @@ public void setUp() throws Exception {
 
     defaultManager = new StreamManager() {
       @Override
-      public ManagedBuffer getChunk(long streamId, int chunkIndex) {
+      public ManagedBuffer getChunk(long streamId, int chunkIndex, Channel channel) {
         throw new UnsupportedOperationException();
       }
     };
@@ -184,7 +185,7 @@ public void furtherRequestsDelay() throws Exception {
     final byte[] response = new byte[16];
     final StreamManager manager = new StreamManager() {
       @Override
-      public ManagedBuffer getChunk(long streamId, int chunkIndex) {
+      public ManagedBuffer getChunk(long streamId, int chunkIndex, Channel channel) {
         Uninterruptibles.sleepUninterruptibly(FOREVER, TimeUnit.MILLISECONDS);
         return new NioManagedBuffer(ByteBuffer.wrap(response));
       }

common/network-common/src/test/java/org/apache/spark/network/StreamSuite.java

@@ -31,6 +31,7 @@
 import java.util.concurrent.TimeUnit;
 
 import com.google.common.io.Files;
+import io.netty.channel.Channel;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -70,12 +71,12 @@ public static void setUp() throws Exception {
     final TransportConf conf = new TransportConf("shuffle", MapConfigProvider.EMPTY);
     final StreamManager streamManager = new StreamManager() {
       @Override
-      public ManagedBuffer getChunk(long streamId, int chunkIndex) {
+      public ManagedBuffer getChunk(long streamId, int chunkIndex, Channel channel) {
         throw new UnsupportedOperationException();
       }
 
       @Override
-      public ManagedBuffer openStream(String streamId) {
+      public ManagedBuffer openStream(String streamId, Channel channel) {
         return testData.openStream(conf, streamId);
       }
     };

common/network-common/src/test/java/org/apache/spark/network/TransportRequestHandlerSuite.java

@@ -59,7 +59,6 @@ public void handleStreamRequest() throws Exception {
     managedBuffers.add(new TestManagedBuffer(30));
     managedBuffers.add(new TestManagedBuffer(40));
     long streamId = streamManager.registerStream("test-app", managedBuffers.iterator());
-    streamManager.registerChannel(channel, streamId);
     TransportClient reverseClient = mock(TransportClient.class);
     TransportRequestHandler requestHandler = new TransportRequestHandler(channel, reverseClient,
       rpcHandler, 2L);

common/network-common/src/test/java/org/apache/spark/network/sasl/SparkSaslSuite.java

@@ -259,7 +259,7 @@ public void testFileRegionEncryption() throws Exception {
     try {
       TransportConf conf = new TransportConf("shuffle", new MapConfigProvider(testConf));
       StreamManager sm = mock(StreamManager.class);
-      when(sm.getChunk(anyLong(), anyInt())).thenAnswer(invocation ->
+      when(sm.getChunk(anyLong(), anyInt(), any(Channel.class))).thenAnswer(invocation ->
           new FileSegmentManagedBuffer(conf, file, 0, file.length()));
 
       RpcHandler rpcHandler = mock(RpcHandler.class);

core/src/main/scala/org/apache/spark/rpc/netty/NettyStreamManager.scala

@@ -19,6 +19,8 @@ package org.apache.spark.rpc.netty
 import java.io.File
 import java.util.concurrent.ConcurrentHashMap
 
+import io.netty.channel.Channel
+
 import org.apache.spark.network.buffer.{FileSegmentManagedBuffer, ManagedBuffer}
 import org.apache.spark.network.server.StreamManager
 import org.apache.spark.rpc.RpcEnvFileServer
@@ -43,11 +45,11 @@ private[netty] class NettyStreamManager(rpcEnv: NettyRpcEnv)
   private val jars = new ConcurrentHashMap[String, File]()
   private val dirs = new ConcurrentHashMap[String, File]()
 
-  override def getChunk(streamId: Long, chunkIndex: Int): ManagedBuffer = {
+  override def getChunk(streamId: Long, chunkIndex: Int, channel: Channel): ManagedBuffer = {
     throw new UnsupportedOperationException()
   }
 
-  override def openStream(streamId: String): ManagedBuffer = {
+  override def openStream(streamId: String, channel: Channel): ManagedBuffer = {
     val Array(ftype, fname) = streamId.stripPrefix("/").split("/", 2)
     val file = ftype match {
       case "files" => files.get(fname)