From 6494db05009b362a2627c8a3beb5a5b9d2bc65fd Mon Sep 17 00:00:00 2001
From: afterincomparableyum
 <224495379+afterincomparableyum@users.noreply.github.com>
Date: Sat, 13 Jun 2026 09:46:13 -0500
Subject: [PATCH] [CELEBORN-2356] Fix NPE before null-check in
 handleChunkFetchRequest

getStreamState returns null for an unknown id. The streamState.buffers pattern match dereferences it before the null guard runs, throwing NullPointerException. The intended ChunkFetchFailure reply is unreachable. Here we should reorder null-check first, then compute storageMetrics.
---
 .../service/deploy/worker/FetchHandler.scala  | 16 ++++-----
 .../deploy/worker/FetchHandlerSuiteJ.java     | 33 +++++++++++++++++++
 2 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
index 7ad990e2bf4..880fee00381 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
+++ b/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
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java
index b54a0b99a56..76fb857e487 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java
+++ b/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 =