feat(conn): add DisposableConn to discard on Close (#4039)

* feat(conn): add DisposableConn to discard on Close

A Conn carries session state (AUTH, SELECT, CLIENT SETNAME, CLIENT
TRACKING, RESET, ...) that cannot be undone on release, so returning its
connection to the parent pool on Close hands that state to unrelated
callers: with PoolSize=1, CLIENT INFO on the pooled client reported the
Conn's ACL user and DB after Close.

Conn keeps returning its connection to the pool. DisposableConn is the
same thing except Close removes the connection from the parent pool, via
a new StickyConnPool.DiscardOnClose honored in freeConn. Get still
reuses an idle parent connection, so the cost is one dial per
DisposableConn lifetime. Tx sticky pools are unchanged.

* refactor: rename DisposableConn to EphemeralConn

---------

Co-authored-by: vlady-kotsev <vladygk8@gmail.com>

README.md

@@ -641,7 +641,10 @@ them, while a raw `Do(ctx, "himport", "prepare", ...)` bypasses that
 entirely, with no replay, recovery, or discard propagation.
 
 For session-scoped work without a typed API, hold a dedicated connection
-(`client.Conn()`) for its whole lifetime and close it afterwards.
+(`client.EphemeralConn()`) for its whole lifetime and close it afterwards.
+`client.Conn()` is also a dedicated connection, but `Close` returns it to the
+pool with whatever session state is left on it, so it only suits commands that
+leave the connection as they found it.
 
 ## Typed Errors
 

internal/pool/pool_sticky.go

@@ -45,6 +45,13 @@ type StickyConnPool struct {
 	// the parent's background drainer.
 	onFirstConn func(*Conn)
 
+	// discardOnClose makes Close remove the claimed connection from the parent
+	// pool instead of returning it. Client.EphemeralConn sets it: the session
+	// state a dedicated connection accumulates (AUTH, SELECT, CLIENT SETNAME,
+	// CLIENT TRACKING, RESET, ...) cannot be undone, and a pooled caller must
+	// never inherit it.
+	discardOnClose bool
+
 	_badConnError atomic.Value
 }
 
@@ -110,6 +117,13 @@ func (p *StickyConnPool) SetOnFirstConn(fn func(*Conn)) {
 	p.onFirstConn = fn
 }
 
+// DiscardOnClose configures the pool to remove the claimed connection from
+// the parent pool on Close instead of returning it. It must be called before
+// the pool is used.
+func (p *StickyConnPool) DiscardOnClose() {
+	p.discardOnClose = true
+}
+
 func (p *StickyConnPool) Put(ctx context.Context, cn *Conn) {
 	defer func() {
 		if recover() != nil {
@@ -130,6 +144,8 @@ func (p *StickyConnPool) Put(ctx context.Context, cn *Conn) {
 func (p *StickyConnPool) freeConn(ctx context.Context, cn *Conn) {
 	if err := p.badConnError(); err != nil {
 		p.pool.Remove(ctx, cn, err)
+	} else if p.discardOnClose {
+		p.pool.Remove(ctx, cn, ErrClosed)
 	} else {
 		p.pool.Put(ctx, cn)
 	}

internal/pool/pool_sticky_close_on_put_test.go

@@ -7,9 +7,10 @@ import (
 	"testing"
 )
 
-// stubPooler records removals; Get hands out a fixed conn.
+// stubPooler records puts and removals; Get hands out a fixed conn.
 type stubPooler struct {
 	cn      *Conn
+	put     bool
 	removed bool
 }
 
@@ -18,7 +19,7 @@ func (s *stubPooler) CloseConn(context.Context, *Conn, string, string) error {
 	return nil
 }
 func (s *stubPooler) Get(context.Context) (*Conn, error) { return s.cn, nil }
-func (s *stubPooler) Put(context.Context, *Conn)         {}
+func (s *stubPooler) Put(context.Context, *Conn)         { s.put = true }
 func (s *stubPooler) Remove(_ context.Context, _ *Conn, _ error) {
 	s.removed = true
 }
@@ -70,3 +71,35 @@ func TestStickyConnPoolPutHonorsCloseOnPut(t *testing.T) {
 		t.Error("underlying connection must be removed from the parent pool")
 	}
 }
+
+// A sticky pool configured with DiscardOnClose must remove its connection from
+// the parent on Close instead of returning it: a dedicated Conn carries session
+// state (AUTH, SELECT, ...) that the next pooled caller must not inherit.
+func TestStickyConnPoolDiscardOnClose(t *testing.T) {
+	ctx := context.Background()
+	for _, discard := range []bool{false, true} {
+		client, server := net.Pipe()
+
+		stub := &stubPooler{cn: NewConn(client)}
+		p := NewStickyConnPool(stub)
+		if discard {
+			p.DiscardOnClose()
+		}
+
+		cn, err := p.Get(ctx)
+		if err != nil {
+			t.Fatalf("discard=%v: get: %v", discard, err)
+		}
+		p.Put(ctx, cn)
+		if err := p.Close(); err != nil {
+			t.Fatalf("discard=%v: close: %v", discard, err)
+		}
+
+		if stub.removed != discard || stub.put == discard {
+			t.Errorf("discard=%v: removed=%v put=%v", discard, stub.removed, stub.put)
+		}
+
+		client.Close()
+		server.Close()
+	}
+}

redis.go

@@ -2676,11 +2676,34 @@ func (c *Client) Close() error {
 	return firstErr
 }
 
+// Conn returns a Conn that runs every command on a single connection borrowed
+// from this client's pool. Close returns that connection to the pool, so any
+// session state left on it (AUTH, SELECT, CLIENT SETNAME, CLIENT TRACKING, ...)
+// is still there for the next pooled caller. Use EphemeralConn when the Conn
+// changes session state.
 func (c *Client) Conn() *Conn {
+	return c.conn(false)
+}
+
+// EphemeralConn is Conn, except Close discards the underlying connection
+// instead of returning it to the pool. Session state a Conn sets (AUTH,
+// SELECT, CLIENT SETNAME, CLIENT TRACKING, RESET, ...) cannot be undone
+// generically on release, so a connection that carried it must not go on to
+// serve unrelated callers. The cost is one dial per EphemeralConn lifetime;
+// Get still reuses an idle pooled connection.
+func (c *Client) EphemeralConn() *Conn {
+	return c.conn(true)
+}
+
+func (c *Client) conn(discardOnClose bool) *Conn {
+	sticky := c.baseClient.newStickyConnPool()
+	if discardOnClose {
+		sticky.DiscardOnClose()
+	}
 	// Share the HIMPORT fieldset registry: the sticky pool borrows
 	// connections from this client's pool, so fieldsets prepared on them
 	// stay valid after the connections are returned.
-	conn := newConn(c.opt, c.baseClient.newStickyConnPool(), &c.hooksMixin, c.himport)
+	conn := newConn(c.opt, sticky, &c.hooksMixin, c.himport)
 	// A sticky client does not serve cache hits, but a new pool connection first
 	// initialized through it may later be reused by the parent. Share the
 	// successful-attachment signal so that connection is tracked exactly when

redis_test.go

@@ -245,6 +245,42 @@ var _ = Describe("Client", func() {
 		Expect(err).NotTo(HaveOccurred())
 	})
 
+	It("should not return an EphemeralConn's session state to the pool", Label("NonRedisEnterprise"), func() {
+		opt := redisOptions()
+		opt.PoolSize = 1
+		db := redis.NewClient(opt)
+
+		defer func() {
+			Expect(db.Close()).NotTo(HaveOccurred())
+		}()
+
+		conn := db.EphemeralConn()
+		Expect(conn.Select(ctx, 1).Err()).NotTo(HaveOccurred())
+		Expect(conn.Close()).NotTo(HaveOccurred())
+
+		// With a single pool slot, a connection handed back by Close would
+		// serve this command and report db=1.
+		info, err := db.ClientInfo(ctx).Result()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(info.DB).To(Equal(0))
+	})
+
+	It("should return a Conn to the pool on Close", func() {
+		opt := redisOptions()
+		opt.PoolSize = 1
+		db := redis.NewClient(opt)
+
+		defer func() {
+			Expect(db.Close()).NotTo(HaveOccurred())
+		}()
+
+		conn := db.Conn()
+		Expect(conn.Ping(ctx).Err()).NotTo(HaveOccurred())
+		Expect(conn.Close()).NotTo(HaveOccurred())
+
+		Expect(db.PoolStats().IdleConns).To(Equal(uint32(1)))
+	})
+
 	It("should client PROTO 2", func() {
 		opt := redisOptions()
 		opt.Protocol = 2