Skip to content

[TPU Raiden] Fix [Failure/ReceiveDisconnect]: propagate broken inbound streams to fail active receive sessions and release staging slots immediately. - #1041

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_987666213
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_987666213

Conversation

@copybara-service

Copy link
Copy Markdown

[TPU Raiden] Fix [Failure/ReceiveDisconnect]: propagate broken inbound streams to fail active receive sessions and release staging slots immediately.

When an inbound TCP push connection fails, resets, or suffers an I/O read error,
the transport server previously logged the error but did not notify the owning delegate.
Without this notification, decode receive sessions remained oblivious to the disconnect
and stalled until per-transfer timeouts (30s) expired, keeping host staging buffers locked
and causing Head-of-Line blocking / slot starvation for subsequent healthy transfers.

This change:

  1. Adds OnReceiveFailed to RawBufferTransportDelegate and BlockTransportDelegate.
  2. In BlockTransport and RawBufferTransport, invokes OnReceiveFailed(uuid, status)
    whenever an incoming push stream fails.
  3. Wires OnReceiveFailed through KVCacheManagerBase's TransferEventHooks to
    KVCacheManagerWithTransfer::OnReceiveFailed.
  4. In KVCacheManagerWithTransfer::OnReceiveFailed, immediately finishes the active
    receive session with the error status, setting draining=true and releasing
    host staging slots back to the free pool immediately on socket disconnect.
  5. Adds unit tests in block_transport_test and kv_cache_manager_with_transfer_test
    verifying immediate failure notification and receive session settlement.

@google-cla

google-cla Bot commented Sep 24, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

…d streams to fail active receive sessions and release staging slots immediately.

When an inbound TCP push connection fails, resets, or suffers an I/O read error,
the transport server previously logged the error but did not notify the owning delegate.
Without this notification, decode receive sessions remained oblivious to the disconnect
and stalled until per-transfer timeouts (30s) expired, keeping host staging buffers locked
and causing Head-of-Line blocking / slot starvation for subsequent healthy transfers.

This change:
1. Adds OnReceiveFailed to RawBufferTransportDelegate and BlockTransportDelegate.
2. In BlockTransport and RawBufferTransport, invokes OnReceiveFailed(uuid, status)
   whenever an incoming push stream fails.
3. Wires OnReceiveFailed through KVCacheManagerBase's TransferEventHooks to
   KVCacheManagerWithTransfer::OnReceiveFailed.
4. In KVCacheManagerWithTransfer::OnReceiveFailed, immediately finishes the active
   receive session with the error status, setting draining=true and releasing
   host staging slots back to the free pool immediately on socket disconnect.
5. Adds unit tests in block_transport_test and kv_cache_manager_with_transfer_test
   verifying immediate failure notification and receive session settlement.

PiperOrigin-RevId: 987666213
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

0 participants