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

[ChuckLin2025]

What changes were proposed in this pull request?

TransportResponseHandler.handle() matches an incoming StreamResponse / StreamFailure to its callback by FIFO streamCallbacks.poll() order. This relies on the assumption (from SPARK-11265 / #9206) that the server answers StreamRequests in the same order the client sent them, so a simple queue rather than a streamId -> callback map suffices. The response carries a streamId and the client stored the expected streamId when it registered the callback, but the two were never compared.

This PR adds a streamId equality check after poll(), factored into a single helper verifyStreamCallbackMatches(...) that both the StreamResponse and StreamFailure branches call. On a mismatch it:

1. Fails the polled callback (under its own registered streamId) with an IOException, so its caller does not hang waiting for a response it will never correctly receive (poll() has already removed that callback from the queue, so it would otherwise be orphaned).
2. Throws IllegalStateException, which propagates to Netty's exceptionCaught, so the connection is torn down and its remaining outstanding requests are re-fetched in order on a fresh channel.

It also emits a dedicated logger.error at the detection site whose message contains desynced and both streamIds and the remote address, so this otherwise-silent condition is greppable.

Why are the changes needed?

This is a defensive check. Under correct operation the streamId equality always holds and the check is a no-op: responses to StreamRequests arrive on a single connection in the order the client sent them, and each callback is registered under the exact streamId it requested, so the head of the FIFO queue always corresponds to the next response. A mismatch is not reachable by any normal client/server interaction -- it could only be produced by memory or hardware corruption (for example a bit flip in the streamId, or a corrupted queue).

The concern is what happens if such corruption ever does occur. Because the StreamResponse wire message carries only (streamId, byteCount) -- no blockId, no content check -- and StreamInterceptor enforces only the server-declared byte count, a wrong-but-self-consistent block would pass every existing check and be delivered silently to the reader as shuffle data corruption on the fetch-to-disk path (OneForOneBlockFetcher -> client.stream()), or leave a reader's task hung waiting for a response it never correctly receives. This check converts that silent, undetectable outcome into a loud, retriable fetch failure. It cannot false-fire: since the client registers each callback under the exact streamId it requested, registered streamId == response streamId is an invariant of every correct delivery, so tearing down the (evidently corrupted) connection is the safe response.

Does this PR introduce any user-facing change?

No. On correct executions the equality always holds, so behavior is unchanged. The check only affects the (normally unreachable) corruption case, turning silent wrong results / a hang into an existing, retriable fetch-failure code path; no API or result-schema change.

How was this patch tested?

Added four unit tests to TransportResponseHandlerSuite (each simulates a desync by registering one streamId and handling a response for another):

• streamResponseWithMismatchedStreamIdThrows -- a StreamResponse whose streamId differs from the head-of-queue callback throws IllegalStateException (message contains desynced), does not deliver success to the wrong callback, and fails the polled callback under its own streamId.
• streamFailureWithMismatchedStreamIdThrows -- same for the StreamFailure branch.
• desyncTearsDownConnectionAndFailsAllOutstandingRequestsRetriably -- on a channel with an innocent concurrent chunk fetch in flight, a desync fails the polled stream callback inline and, via the exceptionCaught teardown path, fails the remaining chunk fetch -- all retriably, none receiving data.
• streamResponseWithMatchingStreamIdIsDelivered -- regression guard: a matching streamId is still delivered normally.

Ran build/sbt 'network-common/testOnly org.apache.spark.network.TransportResponseHandlerSuite': Passed, 15 total, 0 failed.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 4.8)

[ChuckLin2025]

@cloud-fan and @Ngone51 could you take a look at this PR thanks !

[uros-b]

The new method builds its log message with String.format(...) and passes the result to logger.error(msg) as a plain pre-formatted string. Every other logger.error/logger.warn call with dynamic arguments in this file (migrated under SPARK-48209) uses the structured MDC form: logger.error("... {} ...", MDC.of(LogKeys.STREAM_ID, ...), MDC.of(LogKeys.HOST_PORT, ...)). Both LogKeys.STREAM_ID and LogKeys.HOST_PORT already exist and are already imported in this file; the fix is mechanical.

[ChuckLin2025]

Good catch, thanks. Fixed in 53ba42f: the desync logger.error now uses the structured MDC form (MDC.of(LogKeys.STREAM_ID, responseStreamId), MDC.of(LogKeys.HOST_PORT, getRemoteAddress(channel))), matching the other calls in this file. Since the Logger API takes only MDC args for the {} slots and the file uses distinct keys per value (e.g. HOST_PORT/HOST_PORT2) rather than repeating one, I kept the structured line to the arriving STREAM_ID + HOST_PORT and left the full detail (both streamIds and the response/failure kind) on the IOException/IllegalStateException, which is logged with full context on the exceptionCaught teardown path.

[tdcmeehan]

This detects callback-queue desynchronization, not general transport or data corruption. A payload bit flip leaves the stream ID unchanged and is therefore undetected. If SPARK-59142’s goal is corruption detection, should we instead verify an end-to-end block checksum (or rely on an authenticated transport such as TLS)? If the goal is only asserting the FIFO protocol invariant, the PR description, implementation, and tests should be narrowed accordingly.