Fix memory leak of TransportRequestHandler.streamIds.

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

@@ -58,6 +58,11 @@ public OneForOneStreamManager() {
     streams = new ConcurrentHashMap<Long, StreamState>();
   }
 
+  @Override
+  public boolean streamHasNext(long streamId) {
+    return streams.containsKey(streamId);
+  }
+
   @Override
   public ManagedBuffer getChunk(long streamId, int chunkIndex) {
     StreamState state = streams.get(streamId);

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

@@ -43,6 +43,11 @@ public abstract class StreamManager {
    */
   public abstract ManagedBuffer getChunk(long streamId, int chunkIndex);
 
+  /**
+   * Indicates that if the specified stream has next chunks to read further.
+   */
+  public boolean streamHasNext(long streamId) { return true; }
+
   /**
    * Indicates that the TCP connection that was tied to the given stream has been terminated. After
    * this occurs, we are guaranteed not to read from the stream again, so any state can be cleaned

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

@@ -109,6 +109,9 @@ private void processFetchRequest(final ChunkFetchRequest req) {
     ManagedBuffer buf;
     try {
       buf = streamManager.getChunk(req.streamChunkId.streamId, req.streamChunkId.chunkIndex);
+      if (!streamManager.streamHasNext(req.streamChunkId.streamId)) {
+        streamIds.remove(req.streamChunkId.streamId);
+      }
     } catch (Exception e) {
       logger.error(String.format(
         "Error opening block %s for request from %s", req.streamChunkId, client), e);