thriftbp: interrupt in-flight calls when their context is canceled

[rokob]

💸 TL;DR

Make thriftbp enforce the caller's context deadline while a pooled client is performing an RPC.

After leasing a client, clientPool.pooledCall now closes the leased connection when the request context is canceled. Closing the connection interrupts any blocked socket read or write. The call returns the context cancellation cause, and the interrupted connection is discarded instead of being returned to the pool for reuse.

Calls that complete before cancellation continue to reuse their connections normally.

📜 Details

Apache Thrift's Go header transport does not enforce the context deadline while reading a framed response body.

THeaderTransport.ReadFrame checks the context while reading the frame header, but then reads the complete body using io.CopyN. Meanwhile, TSocket.Read sets a new now + SocketTimeout deadline before every socket read. A response that continues making incremental progress can therefore run well beyond the RPC's absolute context deadline.

When a client is leased:

1. Register a callback for context cancellation.
2. Close the leased client if cancellation occurs, interrupting blocked I/O.
3. Synchronize with the callback before completing the call.
4. Return context.Cause(ctx) when cancellation wins.
5. Ensure the closed connection is discarded rather than reused.

The implementation also avoids racing a successful call against a late cancellation callback or closing the same client twice.

The target is enforcement of the RPC context’s absolute deadline, including stalled mid-frame reads. Although framed Thrift buffers the response before Flush, delivery over the socket remains a byte stream and the client can block after receiving only part of the frame. Closing the exclusively leased connection interrupts that read. The connection is then discarded because unread or late response bytes make protocol-safe reuse impossible.

🧪 Testing Steps / Validation

Added coverage for:

• A server that sends a frame body slowly enough to exceed the context deadline while continuing to make progress within the socket timeout.
• Cancellation interrupting the blocked call promptly.
• The interrupted connection being discarded.
• A call completing before its deadline and returning its connection to the pool.
• An already-canceled context.
• Cancellation racing with successful completion.
• No leaked goroutines under repeated cancellation.

✅ Checks

• CI tests (if present) are passing
• Adheres to code style for repo
• Contributor License Agreement (CLA) completed if not a Reddit employee

[fishy]

while what you described is correct, in reality, especially with framed transportation, the size in the frame envelope requires the full response to be already generated to calculate, so when you receive it on the client, that means the server already generated the full response and there's nothing really incremental there (for majority of the implementations the response are send in the same packet with the envelope). in a controlled environment this really is not something would happen.

[mathyourlife-reddit]

Reconnect under bursts: Release replaces the closed client synchronously via
opener() with no rate limit, so a cancel storm (upstream slowdown -> mass
timeouts) can kick off ~one reconnect per canceled lease at once, with new
calls hitting ErrExhausted and capacity shrinking if opens fail. Discard should
be right(?) since a canceled response can leave unread bytes.

On @fishy's point: framed transport means the server buffers the full response
before Flush, but the client still doesn't get the frame atomically. Looks like this can
block mid-frame if the peer stalls on write or delivery slows. Worth saying in
the description whether the target is a stalled mid-frame read or just
enforcing the absolute deadline, since it read as "can't happen."

[rokob]

This is intended to ensure the invariant that an RPC must stop when its absolute context deadline expires. It is exposed by what looks like a stalled mid-frame read where io.CopyN is performing multiple socket reads and blocking between them even though the entire body has been constructed on the server.

I have traces taking >2s to complete when there is a fixed deadline of 1.05s. The current rolling per read socket timeout is not compatible with an absolute deadline.