Add DrainOnContextTimeout to return error on context expiration - #3946
Add DrainOnContextTimeout to return error on context expiration#3946Ananyaas wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 4 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.
| internal.Logger.Printf(ctx, "redis: context timeout drain failed for conn[%d]: %v", cn.GetID(), readErr) | ||
| return false | ||
| } | ||
| return true |
There was a problem hiding this comment.
Pipeline drain reads one reply
High Severity
When DrainOnContextTimeout is enabled, a context error during multi-reply operations (like pipelines) causes drainConnOnContextTimeout to consume only one reply. This leaves unread replies on the socket, returning a desynchronized connection to the pool and potentially corrupting subsequent command parsing.
Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.
| if c.shouldDrainOnContextTimeout(err) { | ||
| if c.drainConnOnContextTimeout(ctx, cn) { | ||
| p.Put(ctx, cn) | ||
| return |
There was a problem hiding this comment.
Drain path skips CSC probe
Medium Severity
The releaseConnToPool function's new context-timeout drain path returns early after a successful drain. This bypasses essential post-command processing, such as marking client-side cache read pending and updating HIMPORT hooks, potentially leading to client/server state divergence and incorrect cache invalidation.
Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.
| drainCtx, cancel := context.WithTimeout(context.Background(), c.opt.ContextTimeoutDrainTimeout) | ||
| defer cancel() | ||
|
|
||
| readErr := cn.WithReader(drainCtx, c.opt.ContextTimeoutDrainTimeout, func(rd *proto.Reader) error { |
There was a problem hiding this comment.
Negative drain timeout hangs release
High Severity
If ContextTimeoutDrainTimeout is negative, drainConnOnContextTimeout calls WithReader with that duration, which skips setting a socket read deadline while proto.Reader does not honor context cancellation. The release path can block indefinitely on ReadReply, holding the connection and stalling pool return.
Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.
| } | ||
| } | ||
| _, err := rd.ReadReply() | ||
| if err == nil || err == proto.Nil { |
There was a problem hiding this comment.
Drain misreads client-handled pushes
Medium Severity
On RESP3, processPendingPushNotificationWithReader leaves client-handled pub/sub pushes (e.g. message) on the stream, but drain then performs a single ReadReply and treats success as a fully aligned connection. That read often consumes the push frame instead of the timed-out command reply, so the real reply can remain and the connection may be re-pooled out of sync.
Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5939f7b86d
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| // Only activate this path for explicit opt-in behavior and for context | ||
| // deadline or cancellation errors, which are the cases where a bounded drain | ||
| // is still safe to attempt. | ||
| return c.opt != nil && c.opt.DrainOnContextTimeout && isContextError(err) |
There was a problem hiding this comment.
Drain socket timeouts caused by context deadlines
When ContextTimeoutEnabled makes a command respect a caller deadline, internal/pool.Conn.WithReader enforces that deadline via SetReadDeadline, so a blocked read normally returns a net.Error timeout such as i/o timeout rather than context.DeadlineExceeded. Because this predicate only accepts errors.Is(..., context.Canceled/DeadlineExceeded), the new opt-in drain path is skipped for the main context-deadline scenario and the connection is still removed as a bad conn, so DrainOnContextTimeout has no effect for typical timed-out commands.
Useful? React with 👍 / 👎.
| return err | ||
| } | ||
| } | ||
| _, err := rd.ReadReply() |
There was a problem hiding this comment.
Drain every pending pipeline reply before re-pooling
For Pipeline/TxPipeline calls, withPipelineConn also releases through this drain path, but a context error can occur after a batch with multiple commands has already been written. Draining only one RESP reply can put the connection back after consuming the first response while later responses are still in the socket; if they have not reached the bufio buffer yet, Put will not reject the connection and the next borrower can read a stale pipeline reply as its command result.
Useful? React with 👍 / 👎.
|
Hello @Ananyaas and thank you for this contribution! I will review it shortly, for now feel free to check the bots feedback. |


Issue : #3808 (comment)
Summary
This change adds an opt-in path for handling commands that finish with a context deadline or cancellation error. When enabled, the client will make a bounded attempt to drain any remaining reply stream before returning the connection to the pool. The original context error is still returned to the caller, but the connection can often be safely reused instead of being discarded.
Why
This helps reduce unnecessary connection churn and avoids protocol desynchronization after timed-out or canceled commands. It is especially useful in workloads with frequent context expiration, where preserving connection reuse improves efficiency and stability.
Test Case Run
Note
Medium Risk
Changes shared connection release behavior on a critical path when enabled; mis-draining could desync protocol, though failed drains remove the connection and the feature is opt-in.
Overview
Adds opt-in handling when a command ends with a context deadline or cancellation error. With
DrainOnContextTimeout,releaseConnToPooltries a bounded read of the outstanding reply (defaultContextTimeoutDrainTimeout50ms) so the connection can bePutback instead of discarded; the caller still gets the original context error. On RESP3, pending push frames are processed before reading the command reply.New options live on
Optionsand are forwarded through cluster, ring, sentinel, and universal client option structs. Unit tests cover successful drain/re-pool and RESP3 push processing ahead of the reply.Reviewed by Cursor Bugbot for commit 5939f7b. Bugbot is set up for automated code reviews on this repo. Configure here.