driver/redis: a cluster or sentinel PubSub follows its subscription to another node (v1.6.0, Refs #887)

driver/redis/cluster.go

@@ -490,6 +490,11 @@ func (c *ClusterClient) refreshLoop() {
 	for {
 		select {
 		case <-ticker.C:
+			// closeCh wins a tie with the ticker: select picks at random
+			// when both are ready, and a refresh after Close dials nodes.
+			if c.closed.Load() {
+				return
+			}
 			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 			_ = c.refreshTopology(ctx)
 			cancel()
@@ -956,16 +961,41 @@ func (c *ClusterClient) Publish(ctx context.Context, channel, message string) (i
 }
 
 // Subscribe opens a pub/sub connection on any cluster node and subscribes to
-// channels (SUBSCRIBE is cluster-wide). See [Client.Subscribe]: when the node is
-// found down on the first SUBSCRIBE the PubSub and a nil error are returned and
-// the reconnect retries that one node address.
+// channels (SUBSCRIBE is node-local, and the cluster delivers every PUBLISH to
+// every node). See [Client.Subscribe]: when the node is found down on the first
+// SUBSCRIBE the PubSub and a nil error are returned and the reconnect replays
+// the channels.
+//
+// The PubSub follows its subscription: when the node that holds it goes away,
+// it subscribes again on another node and keeps delivering to the same
+// [PubSub.Channel]. A message published while it switches is lost (a gap), and
+// the PubSub itself delivers from one conn at a time: it stops listening to the
+// old conn before the new one subscribes. The cluster, though, sends a PUBLISH
+// to every node over its bus asynchronously, so a message in flight around the
+// switch can in rare cases reach the old and the new node and be delivered
+// twice. While the
+// PubSub has no conn its control calls ([PubSub.Subscribe] and the others)
+// return [ErrPubSubConnDown]; nothing else is surfaced as long as a node can be
+// reached, and the PubSub retries with backoff, with no limit, until
+// [PubSub.Close]. When the ClusterClient is closed, a PubSub whose conn is down
+// (or goes down afterwards) gives up: its Channel() is closed and its control
+// calls return [ErrClosed]. A PubSub whose conn stays up after the close (the
+// event loop is shared with another Client, or WithEngine) keeps it until
+// [PubSub.Close].
 func (c *ClusterClient) Subscribe(ctx context.Context, channels ...string) (*PubSub, error) {
-	// Regular SUBSCRIBE is cluster-wide: pick any node.
+	// Regular SUBSCRIBE is node-local: pick any node.
 	node := c.anyNode()
 	if node == nil {
 		return nil, errors.New("celeris-redis: no cluster nodes available")
 	}
-	return node.client.Subscribe(ctx, channels...)
+	ps, err := c.openPubSub(ctx, node.client, "")
+	if err != nil {
+		return nil, err
+	}
+	if len(channels) == 0 {
+		return ps, nil
+	}
+	return ps.firstControl(ps.Subscribe(ctx, channels...))
 }
 
 // SSubscribe opens a pub/sub connection and subscribes to shard channels
@@ -974,9 +1004,25 @@ func (c *ClusterClient) Subscribe(ctx context.Context, channels ...string) (*Pub
 // SHOULD hash to the same slot (use hash tags); mixing slots in one PubSub is
 // undefined behavior in cluster mode. Like [Client.Subscribe], it returns the
 // PubSub and a nil error when the conn is found down on the first SSUBSCRIBE:
-// the reconnect replays the channels, to the same node address. It does not
-// follow a slot move, so a caller that wants the new owner closes the PubSub
-// and calls SSubscribe again.
+// the reconnect replays the channels.
+//
+// The PubSub follows its slot: when the owner goes away, or the slot moves (the
+// server drops the subscription with an unsolicited sunsubscribe, or answers
+// the SSUBSCRIBE with -MOVED), it refreshes the cluster topology, subscribes
+// again on the new owner and keeps delivering to the same [PubSub.Channel],
+// with the PubSub's reconnect backoff between attempts and no limit. A message
+// published while it switches is lost (a gap); one owner serves a slot, so the
+// switch cannot deliver a message twice. One PubSub serves one slot, the slot of
+// its first channel; a channel of another slot added to it later is not served.
+// If another node owns that slot, the SSUBSCRIBE is answered -MOVED, which is not
+// reported, and the channel is not subscribed. If the same node owns it, the node
+// accepts it and delivers it, until the next reroute of any kind, whose replay
+// sends only the channels of the served slot and leaves it out; the move of that
+// slot does not reroute the PubSub. If the first channel is removed while a
+// channel of another node's slot remains, that channel is not subscribed until
+// the next reroute of any kind, whose replay serves the re-selected slot. See
+// [ClusterClient.Subscribe] for what the caller sees while it is down and when
+// the ClusterClient is closed.
 func (c *ClusterClient) SSubscribe(ctx context.Context, channels ...string) (*PubSub, error) {
 	if len(channels) == 0 {
 		return nil, errors.New("celeris-redis: SSubscribe requires at least one channel")
@@ -989,7 +1035,7 @@ func (c *ClusterClient) SSubscribe(ctx context.Context, channels ...string) (*Pu
 	if node == nil {
 		return nil, errors.New("celeris-redis: no cluster nodes available")
 	}
-	ps, err := node.client.newPubSub(ctx)
+	ps, err := c.openPubSub(ctx, node.client, channels[0])
 	if err != nil {
 		return nil, err
 	}

driver/redis/cluster_node.go

@@ -10,13 +10,55 @@ type clusterNode struct {
 	client  *Client
 	primary bool
 	// failing is set when the node is unreachable; cleared after a successful
-	// command or topology refresh confirms it is alive.
-	failing atomic.Bool
+	// command or topology refresh confirms it is alive. A PubSub dial also
+	// sets it, and clears only what it set (see failFlag).
+	failing failFlag
 	// latency stores the last-measured round-trip time in nanoseconds, updated
 	// during topology refresh. Used by RouteByLatency to pick the fastest node.
 	latency atomic.Int64
 }
 
+// failFlag is clusterNode.failing: a flag that also remembers whether it was last
+// set by a failed PubSub dial. One atomic word holds both bits, so every change
+// is one atomic operation and none can be lost to another:
+//
+//	0           healthy
+//	failingBit  failing, set by the command path or found by a topology refresh
+//	failingBit|byPubSubBit  failing, set by a failed PubSub dial
+//
+// Store, which the command path and the topology refresh use, takes the node
+// over: it replaces the word, so a later successful PubSub dial leaves it alone.
+// markPubSub replaces it with the PubSub's own verdict. clearPubSub clears the
+// flag only if the word is exactly the PubSub's own mark (one compare and
+// swap), so it never erases a mark that changed after it looked.
+type failFlag struct {
+	w atomic.Uint32
+}
+
+const (
+	failingBit  uint32 = 1
+	byPubSubBit uint32 = 2
+)
+
+// Load reports whether the node is marked failing.
+func (f *failFlag) Load() bool { return f.w.Load()&failingBit != 0 }
+
+// Store sets the flag as the command path or the topology refresh found it.
+func (f *failFlag) Store(v bool) {
+	if v {
+		f.w.Store(failingBit)
+		return
+	}
+	f.w.Store(0)
+}
+
+// markPubSub sets the flag because a PubSub dial failed.
+func (f *failFlag) markPubSub() { f.w.Store(failingBit | byPubSubBit) }
+
+// clearPubSub clears the flag after a successful PubSub dial, but only if a
+// failed PubSub dial set it and nothing has changed it since.
+func (f *failFlag) clearPubSub() { f.w.CompareAndSwap(failingBit|byPubSubBit, 0) }
+
 // Close tears down the node's client.
 func (n *clusterNode) Close() error {
 	if n.client != nil {

driver/redis/conn.go

@@ -761,6 +761,21 @@ func (c *redisConn) releaseResult(req *redisRequest) {
 	putRequest(req)
 }
 
+// testHookHandshakeRelease, when non-nil, runs for every request handshake gives
+// back to the pool. It is nil in production (handshake runs once per conn, so
+// the check is off every hot path); a test sets it to see that a handshake that
+// failed for any reason but a server error reply gives nothing back (the reader
+// may still be filling the request in), and resets it to nil when it is done.
+var testHookHandshakeRelease atomic.Pointer[func(req *redisRequest)]
+
+// releaseHandshakeReq is releaseResult for handshake's requests.
+func (c *redisConn) releaseHandshakeReq(req *redisRequest) {
+	if h := testHookHandshakeRelease.Load(); h != nil {
+		(*h)(req)
+	}
+	c.releaseResult(req)
+}
+
 // handshake runs HELLO (with optional fallback to AUTH+SELECT) to negotiate
 // protocol version and authentication.
 func (c *redisConn) handshake(ctx context.Context) error {
@@ -784,7 +799,7 @@ func (c *redisConn) handshake(ctx context.Context) error {
 	}
 	req, err := c.exec(ctx, helloArgs...)
 	if err == nil {
-		defer c.releaseResult(req)
+		defer c.releaseHandshakeReq(req)
 		c.state.proto.Store(int32(targetProto))
 		// After HELLO OK, select db.
 		if c.cfg.DB != 0 {
@@ -794,12 +809,18 @@ func (c *redisConn) handshake(ctx context.Context) error {
 		}
 		return nil
 	}
-	// Release the request from the failed exec — it may hold pooled Reader
-	// slices that would otherwise leak.
-	c.releaseResult(req)
 	// If HELLO was rejected, fall back to legacy auth + select.
 	var rerr *RedisError
 	if errors.As(err, &rerr) {
+		// The server's error reply was dispatched: dispatch is done with req.
+		// Any other error (the caller's context ended, the conn closed) leaves
+		// the request to the garbage collector, as every other caller of exec
+		// does: the reader goroutine may still be filling it in, and a request
+		// given back to the pool now could be reused by another conn and
+		// woken by that late finish with a reply that is not its own. Nothing
+		// leaks: req.owned is never set in production code, so releaseResult
+		// only skips a sync.Pool reuse here.
+		c.releaseHandshakeReq(req)
 		c.state.proto.Store(2)
 		return c.legacyAuth(ctx)
 	}

driver/redis/handshake_ctx_ended_887_unix_test.go

@@ -0,0 +1,126 @@
+//go:build linux || darwin
+
+package redis
+
+// A Do whose context ends while a fresh conn's HELLO reply is arriving: exec
+// returns the context's error, but the reader goroutine already popped the
+// request and is filling it in (state.go dispatch). handshake gave that request
+// back to the pool on every failed exec, so another Do could take the same
+// object, and the reader's late finish woke it before its own reply came: an
+// empty reply with a nil error on a healthy conn, and -race names the two
+// writers (found by the review of celeris#887 / PR #1154; the code is from
+// April, 5cb0f9b).
+//
+// The window is the microseconds around the context's end, so the cell is a
+// stress test: 1500 rounds with random delays and timeouts, and it is only
+// meaningful under -race (-count=20 is the bar; without -race it counts wrong
+// replies, which is probabilistic). It is left out of the Windows job, whose
+// timer granularity would make it pass without testing anything.
+//
+// It is meant to run once per CI job. Every round opens a connection, and a
+// local -count=20 uses about half of the macOS ephemeral port range (TIME_WAIT);
+// run it with a modest -count, and leave a minute between big runs.
+
+import (
+	"bufio"
+	"context"
+	"math/rand/v2"
+	"net"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestHandshakeWhoseContextEndedDoesNotRecycleItsRequestWhileTheReaderHoldsIt887(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = ln.Close() })
+	go func() {
+		for {
+			c, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer func() { _ = c.Close() }()
+				r := bufio.NewReader(c)
+				w := bufio.NewWriter(c)
+				for {
+					cmd, err := readCommandInto(r, nil)
+					if err != nil {
+						return
+					}
+					switch strings.ToUpper(cmd[0]) {
+					case "HELLO":
+						time.Sleep(time.Duration(rand.IntN(300)) * time.Microsecond)
+						handleHELLO(w, 3)
+					case "PING":
+						writeSimple(w, "PONG")
+					default:
+						writeSimple(w, "OK")
+					}
+					if w.Flush() != nil {
+						return
+					}
+				}
+			}(c)
+		}
+	}()
+	addr := ln.Addr().String()
+
+	// A healthy client whose PING replies are checked while the others'
+	// handshakes are cut short: a recycled request wakes one of them with a
+	// reply that is not its own.
+	checker, err := NewClient(addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = checker.Close() })
+	var bad atomic.Int64
+	stop := make(chan struct{})
+	var wg sync.WaitGroup
+	for g := 0; g < 2; g++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				v, err := checker.Do(context.Background(), "PING")
+				if err != nil {
+					// (A dial error: the loopback runs out of ports under this
+					// churn on macOS. Only a reply that is not the conn's own
+					// counts.)
+					continue
+				}
+				if v == nil || string(v.Str) != "PONG" {
+					if bad.Add(1) <= 3 {
+						t.Logf("the checker's PING got v=%+v, want PONG", v)
+					}
+				}
+			}
+		}()
+	}
+	for i := 0; i < 1000; i++ {
+		c, err := NewClient(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(50+rand.IntN(300))*time.Microsecond)
+		_, _ = c.Do(ctx, "PING")
+		cancel()
+		_ = c.Close()
+	}
+	close(stop)
+	wg.Wait()
+	if n := bad.Load(); n != 0 {
+		t.Errorf("%d PINGs on a healthy conn got a reply that was not theirs: a request of a failed handshake was recycled while its reader still held it", n)
+	}
+}