[CELEBORN-2356] Fix NPE before null-check in handleChunkFetchRequest

worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala

@@ -557,6 +557,14 @@ class FetchHandler(
       s" to fetch block $streamChunkSlice")
 
     val streamState = chunkStreamManager.getStreamState(streamChunkSlice.streamId)
+    if (streamState == null) {
+      val message = s"Stream ${streamChunkSlice.streamId} is not registered with worker. " +
+        "This can happen if the worker was restarted recently."
+      logError(message)
+      workerSource.incCounter(storageMetrics._3)
+      client.getChannel.writeAndFlush(new ChunkFetchFailure(streamChunkSlice, message))
+      return
+    }
     val storageMetrics = streamState.buffers match {
       case _: FileChunkBuffers => (
           WorkerSource.FETCH_LOCAL_CHUNK_TIME,
@@ -567,14 +575,6 @@ class FetchHandler(
           WorkerSource.FETCH_MEMORY_CHUNK_SUCCESS_COUNT,
           WorkerSource.FETCH_MEMORY_CHUNK_FAIL_COUNT)
     }
-    if (streamState == null) {
-      val message = s"Stream ${streamChunkSlice.streamId} is not registered with worker. " +
-        "This can happen if the worker was restart recently."
-      logError(message)
-      workerSource.incCounter(storageMetrics._3)
-      client.getChannel.writeAndFlush(new ChunkFetchFailure(streamChunkSlice, message))
-      return
-    }
 
     maxChunkBeingTransferred.foreach { threshold =>
       val chunksBeingTransferred = chunkStreamManager.chunksBeingTransferred // take high cpu usage

worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java

@@ -50,6 +50,7 @@
 import org.apache.celeborn.common.network.client.RpcResponseCallback;
 import org.apache.celeborn.common.network.client.TransportClient;
 import org.apache.celeborn.common.network.client.TransportResponseHandler;
+import org.apache.celeborn.common.network.protocol.ChunkFetchFailure;
 import org.apache.celeborn.common.network.protocol.ChunkFetchSuccess;
 import org.apache.celeborn.common.network.protocol.Message;
 import org.apache.celeborn.common.network.protocol.OpenStream;
@@ -302,6 +303,38 @@ public void testDoNotDeleteOriginalFileWhenNonRangeLocalReadWorkInProgress() thr
     }
   }
 
+  @Test
+  public void testFetchChunkForUnregisteredStream() {
+    EmbeddedChannel channel = new EmbeddedChannel();
+    TransportClient client = new TransportClient(channel, mock(TransportResponseHandler.class));
+    FetchHandler fetchHandler = mockFetchHandler(null);
+
+    // the stream is never opened, so it is unknown to the ChunkStreamManager
+    long unregisteredStreamId = 12345;
+    fetchHandler.receive(
+        client,
+        new RpcRequest(
+            TransportClient.requestId(),
+            new NioManagedBuffer(
+                new TransportMessage(
+                        MessageType.CHUNK_FETCH_REQUEST,
+                        PbChunkFetchRequest.newBuilder()
+                            .setStreamChunkSlice(
+                                PbStreamChunkSlice.newBuilder()
+                                    .setStreamId(unregisteredStreamId)
+                                    .setChunkIndex(0)
+                                    .setOffset(0)
+                                    .setLen(Integer.MAX_VALUE))
+                            .build()
+                            .toByteArray())
+                    .toByteBuffer())),
+        createRpcResponseCallback(channel));
+
+    ChunkFetchFailure chunkFetchFailure = channel.readOutbound();
+    assertEquals(unregisteredStreamId, chunkFetchFailure.streamChunkSlice.streamId);
+    assertTrue(chunkFetchFailure.errorString.contains("not registered"));
+  }
+
   private FetchHandler mockFetchHandler(FileInfo fileInfo) {
     WorkerSource workerSource = mock(WorkerSource.class);
     TransportConf transportConf =