[SPARK-59142][CORE] Verify streamId when matching stream fetch responses to callbacks in TransportResponseHandler

common/network-common/src/main/java/org/apache/spark/network/client/TransportResponseHandler.java

@@ -235,6 +235,9 @@ public void handle(ResponseMessage message) throws Exception {
     } else if (message instanceof StreamResponse resp) {
       Pair<String, StreamCallback> entry = streamCallbacks.poll();
       if (entry != null) {
+        // Guard against a desynced callback queue before using the polled callback. Under correct
+        // operation this is always a no-op; see verifyStreamCallbackMatches.
+        verifyStreamCallbackMatches(entry, resp.streamId, "response");
         StreamCallback callback = entry.getRight();
         if (resp.byteCount > 0) {
           StreamInterceptor<ResponseMessage> interceptor = new StreamInterceptor<>(
@@ -269,6 +272,9 @@ public void handle(ResponseMessage message) throws Exception {
     } else if (message instanceof StreamFailure resp) {
       Pair<String, StreamCallback> entry = streamCallbacks.poll();
       if (entry != null) {
+        // Same guard as the StreamResponse branch: verify the polled callback is the one this
+        // failure is for before routing it. Under correct operation this is always a no-op.
+        verifyStreamCallbackMatches(entry, resp.streamId, "failure");
         StreamCallback callback = entry.getRight();
         try {
           callback.onFailure(resp.streamId, new RuntimeException(resp.error));
@@ -284,6 +290,55 @@ public void handle(ResponseMessage message) throws Exception {
     }
   }
 
+  /**
+   * Verifies that the callback polled from the head of the FIFO {@link #streamCallbacks} queue is
+   * the one this stream response/failure is for, by comparing the callback's registered streamId
+   * with the streamId carried by the response.
+   *
+   * <p>Under correct operation this equality always holds and the method is a no-op: responses to
+   * {@code StreamRequest}s arrive on a single connection in the order the client sent them (see
+   * SPARK-11265), and the client registers each callback under the exact streamId it requested, so
+   * the head of the queue always corresponds to the next response. A mismatch is therefore not
+   * reachable by any normal client/server interaction; it could only be produced by memory or
+   * hardware corruption (e.g. a bit flip in the streamId or a corrupted queue). This is a defensive
+   * check that turns such corruption -- which would otherwise silently deliver the wrong block's
+   * bytes to a reader -- into a loud, retriable failure.
+   *
+   * <p>On a mismatch it fails the polled callback under its own streamId (so its caller does not
+   * hang waiting for a response it will never correctly receive; {@code poll()} has already removed
+   * it from the queue) and throws {@link IllegalStateException}, which propagates to Netty's
+   * {@code exceptionCaught} so the connection is torn down and its remaining outstanding requests
+   * are re-fetched in order on a fresh channel.
+   */
+  private void verifyStreamCallbackMatches(
+      Pair<String, StreamCallback> entry, String responseStreamId, String kind) {
+    if (entry.getLeft().equals(responseStreamId)) {
+      return;
+    }
+    // Log at the detection site, in the structured MDC form used throughout this file, so this
+    // otherwise-silent guard is directly greppable ("desynced") for incidence analysis, independent
+    // of the generic connection-exception log the thrown IllegalStateException produces downstream.
+    logger.error("Stream callback queue desynced: received streamId {} does not match the head of "
+        + "the callback queue from {}. This is unreachable under correct operation and may "
+        + "indicate memory or hardware corruption; failing the connection to avoid delivering the "
+        + "wrong block.",
+      MDC.of(LogKeys.STREAM_ID, responseStreamId),
+      MDC.of(LogKeys.HOST_PORT, getRemoteAddress(channel)));
+    // Full detail (both streamIds) goes on the exception that fails the callback and tears down the
+    // connection.
+    String msg = String.format(
+      "Stream callback queue desynced: %s streamId %s does not match the head of the callback "
+        + "queue (streamId %s) from %s. This is unreachable under correct operation and may "
+        + "indicate memory or hardware corruption; failing the connection to avoid delivering the "
+        + "wrong block.", kind, responseStreamId, entry.getLeft(), getRemoteAddress(channel));
+    try {
+      entry.getRight().onFailure(entry.getLeft(), new IOException(msg));
+    } catch (IOException ioe) {
+      logger.warn("Error in stream failure handler.", ioe);
+    }
+    throw new IllegalStateException(msg);
+  }
+
   /** Returns total number of outstanding requests (fetch requests + rpcs) */
   public int numOutstandingRequests() {
     return outstandingFetches.size() + outstandingRpcs.size() + streamCallbacks.size() +

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

@@ -27,6 +27,8 @@
 import org.mockito.ArgumentCaptor;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.*;
 
 import org.apache.spark.network.buffer.NioManagedBuffer;
@@ -196,6 +198,114 @@ public void failStreamCallbackWhenInstallingInterceptorFails() throws Exception
     assertEquals(0, handler.numOutstandingRequests());
   }
 
+  @Test
+  public void streamResponseWithMismatchedStreamIdThrows() throws Exception {
+    // The FIFO streamCallbacks queue matches responses to callbacks by poll() order. A response
+    // whose streamId does not match the head-of-queue callback's registered streamId is unreachable
+    // under correct operation (responses arrive in order on one connection and each callback is
+    // registered under its own streamId) -- it would only arise from memory/hardware corruption,
+    // and delivering the response would feed the wrong block's bytes to the callback. The defensive
+    // guard must throw instead (SPARK-59142); throwing propagates to Netty's exceptionCaught -> the
+    // connection is torn down and its outstanding requests re-fetched in order on a fresh channel.
+    // A mismatch is simulated by registering "stream-A" and handling a response for "stream-B".
+    Channel c = new LocalChannel();
+    c.pipeline().addLast(TransportFrameDecoder.HANDLER_NAME, new TransportFrameDecoder());
+    TransportResponseHandler handler = new TransportResponseHandler(c);
+
+    StreamCallback cb = mock(StreamCallback.class);
+    handler.addStreamCallback("stream-A", cb);
+
+    // A response for a different streamId arrives at the head of the FIFO queue.
+    StreamResponse mismatched = new StreamResponse("stream-B", 1234L, null);
+    IllegalStateException e = assertThrows(IllegalStateException.class,
+      () -> handler.handle(mismatched));
+    assertTrue(e.getMessage().contains("desynced"),
+      "expected a desync error, got: " + e.getMessage());
+    // The mismatched response must NOT have been delivered to the wrong callback as success...
+    verify(cb, never()).onComplete(any());
+    // ...and the polled callback is failed (with its OWN streamId) so its caller does not hang.
+    verify(cb, times(1)).onFailure(eq("stream-A"), isA(IOException.class));
+  }
+
+  @Test
+  public void streamFailureWithMismatchedStreamIdThrows() throws Exception {
+    Channel c = new LocalChannel();
+    c.pipeline().addLast(TransportFrameDecoder.HANDLER_NAME, new TransportFrameDecoder());
+    TransportResponseHandler handler = new TransportResponseHandler(c);
+
+    StreamCallback cb = mock(StreamCallback.class);
+    handler.addStreamCallback("stream-A", cb);
+
+    StreamFailure mismatched = new StreamFailure("stream-B", "uh-oh");
+    IllegalStateException e = assertThrows(IllegalStateException.class,
+      () -> handler.handle(mismatched));
+    assertTrue(e.getMessage().contains("desynced"),
+      "expected a desync error, got: " + e.getMessage());
+    // The failure must NOT be routed under the wrong (response) streamId; the polled callback is
+    // failed under its OWN streamId instead, so its caller does not hang.
+    verify(cb, never()).onFailure(eq("stream-B"), any());
+    verify(cb, times(1)).onFailure(eq("stream-A"), isA(IOException.class));
+  }
+
+  @Test
+  public void desyncTearsDownConnectionAndFailsAllOutstandingRequestsRetriably() throws Exception {
+    // Upstream impact of the streamId assert. In production, throwing from handle() propagates to
+    // TransportChannelHandler.exceptionCaught, which calls responseHandler.exceptionCaught (failing
+    // EVERY outstanding request on the channel) and then ctx.close(). This test simulates that
+    // sequence and shows the meaning for callers: when a stream-callback desync is detected, the
+    // whole (poisoned) connection is torn down and ALL its in-flight requests -- the mismatched
+    // stream AND any innocent concurrent chunk-fetch sharing the channel -- fail with a retriable
+    // error. None receive data. Upstream, each onFailure becomes a FetchFailedException -> stage
+    // retry on a fresh, in-order connection. The cost of a detected desync is a retry, never
+    // corrupt bytes.
+    Channel c = new LocalChannel();
+    c.pipeline().addLast(TransportFrameDecoder.HANDLER_NAME, new TransportFrameDecoder());
+    TransportResponseHandler handler = new TransportResponseHandler(c);
+
+    // An innocent chunk fetch is in flight on the same connection.
+    StreamChunkId chunkId = new StreamChunkId(1, 0);
+    ChunkReceivedCallback chunkCb = mock(ChunkReceivedCallback.class);
+    handler.addFetchRequest(chunkId, chunkCb);
+    // ...and a stream fetch for "stream-A".
+    StreamCallback streamCb = mock(StreamCallback.class);
+    handler.addStreamCallback("stream-A", streamCb);
+    assertEquals(2, handler.numOutstandingRequests());
+
+    // A StreamResponse for the wrong streamId arrives -> handle() throws (desync detected).
+    // The desynced (polled) stream callback is failed inline so its caller does not hang.
+    IllegalStateException thrown = assertThrows(IllegalStateException.class,
+      () -> handler.handle(new StreamResponse("stream-B", 1234L, null)));
+    assertTrue(thrown.getMessage().contains("desynced"));
+    verify(streamCb, times(1)).onFailure(eq("stream-A"), any());
+    verify(streamCb, never()).onComplete(any());
+
+    // Netty then invokes exceptionCaught with the thrown cause; this is the teardown path that
+    // fails the connection's REMAINING outstanding requests (the innocent concurrent chunk fetch).
+    handler.exceptionCaught(thrown);
+    verify(chunkCb, times(1)).onFailure(eq(0), any());
+
+    // Net result: no request on the poisoned connection received data; all failed retriably.
+    assertEquals(0, handler.numOutstandingRequests());
+  }
+
+  @Test
+  public void streamResponseWithMatchingStreamIdIsDelivered() throws Exception {
+    // Regression guard: the streamId check must not disturb the normal in-order case. A response
+    // whose streamId matches the head-of-queue callback is handled exactly as before.
+    Channel c = new LocalChannel();
+    c.pipeline().addLast(TransportFrameDecoder.HANDLER_NAME, new TransportFrameDecoder());
+    TransportResponseHandler handler = new TransportResponseHandler(c);
+
+    StreamCallback cb = mock(StreamCallback.class);
+    handler.addStreamCallback("stream", cb);
+    assertEquals(1, handler.numOutstandingRequests());
+
+    // byteCount == 0 -> the handler calls onComplete inline (no interceptor install needed).
+    handler.handle(new StreamResponse("stream", 0L, null));
+    verify(cb, times(1)).onComplete(eq("stream"));
+    assertEquals(0, handler.numOutstandingRequests());
+  }
+
   @Test
   public void handleSuccessfulMergedBlockMeta() throws Exception {
     TransportResponseHandler handler = new TransportResponseHandler(new LocalChannel());