Skip to content

Add DrainOnContextTimeout to return error on context expiration #3946

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Ananyaas wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Add DrainOnContextTimeout to return error on context expiration #3946

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
027b32d
avoid closing pooled connections on context deadline when response ca…
Ananyaas Aug 3, 2026
a991124
Merge branch 'master' into master
Ananyaas Aug 4, 2026
5939f7b
Proposal: avoid closing pooled connections on context deadline when r…
Ananyaas Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 112 additions & 0 deletions internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,55 @@ import (
"github.com/redis/go-redis/v9/auth"
"github.com/redis/go-redis/v9/internal/pool"
"github.com/redis/go-redis/v9/internal/proto"
"github.com/redis/go-redis/v9/push"

. "github.com/bsm/ginkgo/v2"
. "github.com/bsm/gomega"
)

var ctx = context.TODO()

type stubPooler struct {
putCalls int
removeCalls int
putConns []*pool.Conn
removeConns []*pool.Conn
}

type trackingPushProcessor struct {
calls int
}

func (p *trackingPushProcessor) GetHandler(string) push.NotificationHandler { return nil }
func (p *trackingPushProcessor) ProcessPendingNotifications(context.Context, push.NotificationHandlerContext, *proto.Reader) error {
p.calls++
return nil
}
func (p *trackingPushProcessor) RegisterHandler(string, push.NotificationHandler, bool) error {
return nil
}
func (p *trackingPushProcessor) UnregisterHandler(string) error { return nil }

func (s *stubPooler) NewConn(context.Context) (*pool.Conn, error) { return nil, nil }
func (s *stubPooler) CloseConn(context.Context, *pool.Conn, string, string) error { return nil }
func (s *stubPooler) Get(context.Context) (*pool.Conn, error) { return nil, nil }
func (s *stubPooler) Put(context.Context, *pool.Conn) {
s.putCalls++
s.putConns = append(s.putConns, nil)
}
func (s *stubPooler) Remove(context.Context, *pool.Conn, error) {
s.removeCalls++
s.removeConns = append(s.removeConns, nil)
}
func (s *stubPooler) Len() int { return 0 }
func (s *stubPooler) IdleLen() int { return 0 }
func (s *stubPooler) Stats() *pool.Stats { return &pool.Stats{} }
func (s *stubPooler) Size() int { return 0 }
func (s *stubPooler) AddPoolHook(pool.PoolHook) {}
func (s *stubPooler) RemovePoolHook(pool.PoolHook) {}
func (s *stubPooler) RemoveWithoutTurn(context.Context, *pool.Conn, error) {}
func (s *stubPooler) Close() error { return nil }

type capturingLogger struct {
mu sync.Mutex
logs []string
Expand Down Expand Up @@ -766,6 +808,76 @@ var _ = Describe("withConn", Label("NonRedisEnterprise"), func() {
})
})

// TestReleaseConnDrainsContextTimeoutAndPoolsConnection verifies that a
// context-timeout error can still be surfaced to the caller while allowing the
// connection to be re-pooled when the remaining reply can be drained safely.
func TestReleaseConnDrainsContextTimeoutAndPoolsConnection(t *testing.T) {
server, clientConn := net.Pipe()
defer func() { _ = server.Close(); _ = clientConn.Close() }()

go func() {
_, _ = server.Write([]byte("+OK\r\n"))
_ = server.Close()
}()

cn := pool.NewConn(clientConn)
pooler := &stubPooler{}
client := &baseClient{
opt: &Options{
DrainOnContextTimeout: true,
ContextTimeoutDrainTimeout: 50 * time.Millisecond,
},
connPool: pooler,
}

client.releaseConn(context.Background(), cn, context.DeadlineExceeded)

if pooler.putCalls != 1 {
t.Fatalf("expected connection to be put back into pool, got %d puts and %d removes", pooler.putCalls, pooler.removeCalls)
}
if pooler.removeCalls != 0 {
t.Fatalf("expected connection not to be removed, got %d removes", pooler.removeCalls)
}
}

// TestReleaseConnDrainsContextTimeoutAndProcessesPushNotifications exercises
// the same drain path for RESP3 connections and confirms that pending push
// notifications are processed before the reply is consumed.
func TestReleaseConnDrainsContextTimeoutAndProcessesPushNotifications(t *testing.T) {
server, clientConn := net.Pipe()
defer func() { _ = server.Close(); _ = clientConn.Close() }()

go func() {
_, _ = server.Write([]byte(">2\r\n$7\r\nmessage\r\n$5\r\nhello\r\n+OK\r\n"))
_ = server.Close()
}()

cn := pool.NewConn(clientConn)
pooler := &stubPooler{}
processor := &trackingPushProcessor{}
client := &baseClient{
opt: &Options{
Protocol: 3,
DrainOnContextTimeout: true,
ContextTimeoutDrainTimeout: 50 * time.Millisecond,
},
connPool: pooler,
pushProcessor: processor,
}

client.releaseConn(context.Background(), cn, context.DeadlineExceeded)

if processor.calls != 1 {
t.Fatalf("expected push processor to be invoked once, got %d", processor.calls)
}
if pooler.putCalls != 1 {
t.Fatalf("expected connection to be put back into pool, got %d puts and %d removes", pooler.putCalls, pooler.removeCalls)
}
if pooler.removeCalls != 0 {
t.Fatalf("expected connection not to be removed, got %d removes", pooler.removeCalls)
}
}

var _ = Describe("ClusterClient", func() {
var client *ClusterClient

Expand Down
18 changes: 18 additions & 0 deletions options.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,17 @@ type Options struct {
//
// Experimental: this API may change in a minor release.
ClientSideCacheStrategy CSCStrategy

// DrainOnContextTimeout enables a bounded drain of the reply stream when a
// command finishes with a context deadline or cancellation error. The caller
// still receives the original context error, but the connection may be
// re-pooled if the outstanding reply can be safely consumed.
DrainOnContextTimeout bool

// ContextTimeoutDrainTimeout bounds how long the client will spend draining
// a reply stream after a context timeout or cancellation error. If the drain
// does not complete within this duration, the connection is removed.
ContextTimeoutDrainTimeout time.Duration
}

// CSCStrategy selects the client-side caching invalidation architecture. Set via
Expand Down Expand Up @@ -531,6 +542,13 @@ func (opt *Options) init() {
case 0:
opt.WriteTimeout = opt.ReadTimeout
}

// Default to a short drain window so context-timeout recovery remains
// bounded and does not introduce long stalls on the release path.
if opt.ContextTimeoutDrainTimeout == 0 {
opt.ContextTimeoutDrainTimeout = 50 * time.Millisecond
}

if opt.PoolTimeout == 0 {
if opt.ReadTimeout > 0 {
opt.PoolTimeout = opt.ReadTimeout + time.Second
Expand Down
12 changes: 8 additions & 4 deletions osscluster.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,9 +110,11 @@ type ClusterOptions struct {
// See Options.DialerRetryBackoff for details.
DialerRetryBackoff func(attempt int) time.Duration

ReadTimeout time.Duration
WriteTimeout time.Duration
ContextTimeoutEnabled bool
ReadTimeout time.Duration
WriteTimeout time.Duration
ContextTimeoutEnabled bool
DrainOnContextTimeout bool
ContextTimeoutDrainTimeout time.Duration

// MaxConcurrentDials is the maximum number of concurrent connection creation goroutines.
// If <= 0, defaults to PoolSize. If > PoolSize, it will be capped at PoolSize.
Expand Down Expand Up @@ -456,7 +458,9 @@ func (opt *ClusterOptions) clientOptions() *Options {
ReadTimeout: opt.ReadTimeout,
WriteTimeout: opt.WriteTimeout,

ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
DrainOnContextTimeout: opt.DrainOnContextTimeout,
ContextTimeoutDrainTimeout: opt.ContextTimeoutDrainTimeout,

PoolFIFO: opt.PoolFIFO,
PoolSize: opt.PoolSize,
Expand Down
51 changes: 51 additions & 0 deletions redis.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1046,6 +1046,19 @@ func (c *baseClient) releaseConn(ctx context.Context, cn *pool.Conn, err error)
// tracking is on. Limiter accounting stays with the callers, whose shapes
// differ. Shared by releaseConn and withPipelineConn so the two cannot drift.
func (c *baseClient) releaseConnToPool(ctx context.Context, p pool.Pooler, cn *pool.Conn, err error) {
// If the command finished with a context error and drain-on-timeout is
// enabled, try to consume the outstanding reply and restore protocol
// alignment before re-pooling the connection. The original context error is
// still preserved for the caller.
if c.shouldDrainOnContextTimeout(err) {
if c.drainConnOnContextTimeout(ctx, cn) {
p.Put(ctx, cn)
return

@cursor cursor Bot Aug 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.

}
p.Remove(ctx, cn, err)
return
}

if isBadConn(err, false, c.opt.Addr) {
p.Remove(ctx, cn, err)
return
Expand All @@ -1070,6 +1083,44 @@ func (c *baseClient) releaseConnToPool(ctx context.Context, p pool.Pooler, cn *p
p.Put(ctx, cn)
}

func (c *baseClient) shouldDrainOnContextTimeout(err error) bool {
// 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)

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

}

func (c *baseClient) drainConnOnContextTimeout(ctx context.Context, cn *pool.Conn) bool {
if c.opt == nil || cn == nil || cn.IsClosed() {
return false
}

drainCtx, cancel := context.WithTimeout(context.Background(), c.opt.ContextTimeoutDrainTimeout)
defer cancel()

readErr := cn.WithReader(drainCtx, c.opt.ContextTimeoutDrainTimeout, func(rd *proto.Reader) error {

@cursor cursor Bot Aug 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.

// RESP3 connections may already have pending push frames buffered ahead
// of the reply. Process them with the configured push processor so the
// stream remains aligned before we read the command reply.
if c.opt.Protocol == 3 && c.pushProcessor != nil {
if err := c.processPendingPushNotificationWithReader(drainCtx, cn, rd); err != nil {
return err
}
}
_, err := rd.ReadReply()

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

if err == nil || err == proto.Nil {

@cursor cursor Bot Aug 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.

return nil
}
return err
})

if readErr != nil {
internal.Logger.Printf(ctx, "redis: context timeout drain failed for conn[%d]: %v", cn.GetID(), readErr)
return false
}
return true

@cursor cursor Bot Aug 4, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 5939f7b. Configure here.

}

func (c *baseClient) withConn(
ctx context.Context, fn func(context.Context, *pool.Conn) error,
) error {
Expand Down
24 changes: 14 additions & 10 deletions ring.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,9 +117,11 @@ type RingOptions struct {
// See Options.DialerRetryBackoff for details.
DialerRetryBackoff func(attempt int) time.Duration

ReadTimeout time.Duration
WriteTimeout time.Duration
ContextTimeoutEnabled bool
ReadTimeout time.Duration
WriteTimeout time.Duration
ContextTimeoutEnabled bool
DrainOnContextTimeout bool
ContextTimeoutDrainTimeout time.Duration

// PoolFIFO uses FIFO mode for each node connection pool GET/PUT (default LIFO).
PoolFIFO bool
Expand Down Expand Up @@ -239,13 +241,15 @@ func (opt *RingOptions) clientOptions() *Options {

MaxRetries: -1,

DialTimeout: opt.DialTimeout,
DialerRetries: opt.DialerRetries,
DialerRetryTimeout: opt.DialerRetryTimeout,
DialerRetryBackoff: opt.DialerRetryBackoff,
ReadTimeout: opt.ReadTimeout,
WriteTimeout: opt.WriteTimeout,
ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
DialTimeout: opt.DialTimeout,
DialerRetries: opt.DialerRetries,
DialerRetryTimeout: opt.DialerRetryTimeout,
DialerRetryBackoff: opt.DialerRetryBackoff,
ReadTimeout: opt.ReadTimeout,
WriteTimeout: opt.WriteTimeout,
ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
DrainOnContextTimeout: opt.DrainOnContextTimeout,
ContextTimeoutDrainTimeout: opt.ContextTimeoutDrainTimeout,

PoolFIFO: opt.PoolFIFO,
PoolSize: opt.PoolSize,
Expand Down
20 changes: 14 additions & 6 deletions sentinel.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,9 +108,11 @@ type FailoverOptions struct {
// See Options.DialerRetryBackoff for details.
DialerRetryBackoff func(attempt int) time.Duration

ReadTimeout time.Duration
WriteTimeout time.Duration
ContextTimeoutEnabled bool
ReadTimeout time.Duration
WriteTimeout time.Duration
ContextTimeoutEnabled bool
DrainOnContextTimeout bool
ContextTimeoutDrainTimeout time.Duration

// ReadBufferSize is the size of the bufio.Reader buffer for each connection.
// Larger buffers can improve performance for commands that return large responses.
Expand Down Expand Up @@ -229,7 +231,9 @@ func (opt *FailoverOptions) clientOptions() *Options {
ReadTimeout: opt.ReadTimeout,
WriteTimeout: opt.WriteTimeout,

ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
DrainOnContextTimeout: opt.DrainOnContextTimeout,
ContextTimeoutDrainTimeout: opt.ContextTimeoutDrainTimeout,

PoolFIFO: opt.PoolFIFO,
PoolSize: opt.PoolSize,
Expand Down Expand Up @@ -284,7 +288,9 @@ func (opt *FailoverOptions) sentinelOptions(addr string) *Options {
ReadTimeout: opt.ReadTimeout,
WriteTimeout: opt.WriteTimeout,

ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
DrainOnContextTimeout: opt.DrainOnContextTimeout,
ContextTimeoutDrainTimeout: opt.ContextTimeoutDrainTimeout,

PoolFIFO: opt.PoolFIFO,
PoolSize: opt.PoolSize,
Expand Down Expand Up @@ -350,7 +356,9 @@ func (opt *FailoverOptions) clusterOptions() *ClusterOptions {
ReadTimeout: opt.ReadTimeout,
WriteTimeout: opt.WriteTimeout,

ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
ContextTimeoutEnabled: opt.ContextTimeoutEnabled,
DrainOnContextTimeout: opt.DrainOnContextTimeout,
ContextTimeoutDrainTimeout: opt.ContextTimeoutDrainTimeout,

PoolFIFO: opt.PoolFIFO,
PoolSize: opt.PoolSize,
Expand Down
20 changes: 14 additions & 6 deletions universal.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,9 +70,11 @@ type UniversalOptions struct {
// default: 100 milliseconds
DialerRetryTimeout time.Duration

ReadTimeout time.Duration
WriteTimeout time.Duration
ContextTimeoutEnabled bool
ReadTimeout time.Duration
WriteTimeout time.Duration
ContextTimeoutEnabled bool
DrainOnContextTimeout bool
ContextTimeoutDrainTimeout time.Duration

// ReadBufferSize is the size of the bufio.Reader buffer for each connection.
// Larger buffers can improve performance for commands that return large responses.
Expand Down Expand Up @@ -210,7 +212,9 @@ func (o *UniversalOptions) Cluster() *ClusterOptions {
ReadTimeout: o.ReadTimeout,
WriteTimeout: o.WriteTimeout,

ContextTimeoutEnabled: o.ContextTimeoutEnabled,
ContextTimeoutEnabled: o.ContextTimeoutEnabled,
DrainOnContextTimeout: o.DrainOnContextTimeout,
ContextTimeoutDrainTimeout: o.ContextTimeoutDrainTimeout,

ReadBufferSize: o.ReadBufferSize,
WriteBufferSize: o.WriteBufferSize,
Expand Down Expand Up @@ -277,7 +281,9 @@ func (o *UniversalOptions) Failover() *FailoverOptions {
ReadTimeout: o.ReadTimeout,
WriteTimeout: o.WriteTimeout,

ContextTimeoutEnabled: o.ContextTimeoutEnabled,
ContextTimeoutEnabled: o.ContextTimeoutEnabled,
DrainOnContextTimeout: o.DrainOnContextTimeout,
ContextTimeoutDrainTimeout: o.ContextTimeoutDrainTimeout,

ReadBufferSize: o.ReadBufferSize,
WriteBufferSize: o.WriteBufferSize,
Expand Down Expand Up @@ -338,7 +344,9 @@ func (o *UniversalOptions) Simple() *Options {
ReadTimeout: o.ReadTimeout,
WriteTimeout: o.WriteTimeout,

ContextTimeoutEnabled: o.ContextTimeoutEnabled,
ContextTimeoutEnabled: o.ContextTimeoutEnabled,
DrainOnContextTimeout: o.DrainOnContextTimeout,
ContextTimeoutDrainTimeout: o.ContextTimeoutDrainTimeout,

ReadBufferSize: o.ReadBufferSize,
WriteBufferSize: o.WriteBufferSize,
Expand Down
Loading