fix(rueidisaside): prevent canceled cache fills from leaking locks

rueidisaside/README.md

@@ -110,8 +110,8 @@ func main() {
 
 ## Limitation
 
-Currently, requires Redis >= 7.0.
-However, the `UseLuaLock` option is available and allows you to use the `rueidisaside` with older Redis versions < 7.0 as well.
+By default, lock acquisition uses the Redis 7.0 `SET NX GET` behavior.
+The `UseLuaLock` option switches the acquisition step to a Lua implementation that is compatible with older Redis versions < 7.0.
 
 To configure the Lua fallback option:
 
@@ -120,7 +120,7 @@ client, err := rueidisaside.NewClient(rueidisaside.ClientOption{
     ClientOption: rueidis.ClientOption{
         InitAddress: []string{"127.0.0.1:6379"},
     },
-    UseLuaLock: true, // Enable Lua script for older Redis versions
+    UseLuaLock: true, // Enable the Redis < 7 compatible lock implementation
 })
 if err != nil {
     panic(err)

rueidisaside/aside.go

@@ -5,7 +5,6 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"math/rand"
-	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -19,7 +18,7 @@ type ClientOption struct {
 	ClientBuilder func(option rueidis.ClientOption) (rueidis.Client, error)
 	ClientOption  rueidis.ClientOption
 	ClientTTL     time.Duration // TTL for the client marker, refreshed every 1/2 TTL. Defaults to 10s. The marker allows other clients to know if this client is still alive.
-	UseLuaLock    bool
+	UseLuaLock    bool          // Use the Redis < 7 compatible lock implementation.
 }
 
 type CacheAsideClient interface {
@@ -35,9 +34,11 @@ func NewClient(option ClientOption) (cc CacheAsideClient, err error) {
 	}
 	ca := &Client{
 		waits:      make(map[string]chan struct{}),
+		flights:    make(map[flightKey]*flight),
 		ttl:        option.ClientTTL,
 		useLuaLock: option.UseLuaLock,
 	}
+	option.ClientOption.PipelineMultiplex = -1 // ensure lock and cleanup commands use the same connection.
 	option.ClientOption.OnInvalidations = ca.onInvalidation
 	if option.ClientBuilder != nil {
 		ca.client, err = option.ClientBuilder(option.ClientOption)
@@ -55,6 +56,7 @@ type Client struct {
 	client     rueidis.Client
 	ctx        context.Context
 	waits      map[string]chan struct{}
+	flights    map[flightKey]*flight
 	cancel     context.CancelFunc
 	id         string
 	ttl        time.Duration
@@ -72,6 +74,7 @@ func (c *Client) onInvalidation(messages []rueidis.RedisMessage) {
 			close(ch)
 		}
 		c.waits = make(map[string]chan struct{})
+		c.resetFlightsLocked()
 	} else {
 		for _, m := range messages {
 			key, _ := m.ToString()
@@ -158,23 +161,22 @@ retry:
 
 	if rueidis.IsRedisNil(err) && fn != nil { // cache miss, prepare to populate the value by fn()
 		var id string
-		if id, err = c.keepalive(); err == nil { // acquire client id
-			if c.useLuaLock {
-				val, err = acquireLock.Exec(ctx, c.client, []string{key}, []string{id, strconv.FormatInt(ttl.Milliseconds(), 10)}).ToString()
-			} else {
-				val, err = c.client.Do(ctx, c.client.B().Set().Key(key).Value(id).Nx().Get().Px(ttl).Build()).ToString()
+		if id, err = c.keepalive(); err == nil {
+			f, leader := c.beginFlight(key, id)
+			if f == nil { // the client id changed while preparing the flight
+				goto retry
 			}
-
-			if rueidis.IsRedisNil(err) { // successfully set client id on the key as a lock
-				// attach TTL pointer to context for potential modification via OverrideCacheTTL
-				ctx = context.WithValue(ctx, ttlKey, &ttl)
-				if val, err = fn(ctx, key); err == nil {
-					err = setkey.Exec(ctx, c.client, []string{key}, []string{id, val, strconv.FormatInt(ttl.Milliseconds(), 10)}).Error()
-				}
-				if err != nil { // failed to populate the value, release the lock.
-					delkey.Exec(context.Background(), c.client, []string{key}, []string{id})
+			if !leader {
+				select {
+				case <-f.done:
+					goto retry
+				case <-ctx.Done():
+					return "", ctx.Err()
+				case <-c.ctx.Done():
+					return "", c.ctx.Err()
 				}
 			}
+			val, err = c.populate(ctx, ttl, key, id, fn, f)
 		}
 	}
 
@@ -220,7 +222,7 @@ func (c *Client) Close() {
 	id := c.id
 	c.mu.Unlock()
 	if id != "" {
-		c.client.Do(context.Background(), c.client.B().Del().Key(c.id).Build())
+		c.client.Do(context.Background(), c.client.B().Del().Key(id).Build())
 	}
 	c.client.Close()
 }

rueidisaside/aside_test.go

@@ -45,9 +45,11 @@ func TestClientErr(t *testing.T) {
 
 func TestWithClientBuilder(t *testing.T) {
 	var client rueidis.Client
+	var pipelineMultiplex int
 	c, err := NewClient(ClientOption{
-		ClientOption: rueidis.ClientOption{InitAddress: addr, SelectDB: 5},
+		ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: 3, SelectDB: 5},
 		ClientBuilder: func(option rueidis.ClientOption) (_ rueidis.Client, err error) {
+			pipelineMultiplex = option.PipelineMultiplex
 			client, err = rueidis.NewClient(option)
 			return client, err
 		},
@@ -59,6 +61,9 @@ func TestWithClientBuilder(t *testing.T) {
 	if c.Client() != client {
 		t.Fatal("client mismatched")
 	}
+	if pipelineMultiplex != -1 {
+		t.Fatalf("expected PipelineMultiplex -1, got %d", pipelineMultiplex)
+	}
 }
 
 func TestCacheFilled(t *testing.T) {

rueidisaside/flight.go

@@ -0,0 +1,90 @@
+package rueidisaside
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	"github.com/redis/rueidis"
+)
+
+type flightKey struct {
+	key string
+	id  string // client id generation
+}
+
+type flight struct {
+	done chan struct{}
+}
+
+func (c *Client) beginFlight(key, id string) (f *flight, leader bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.id != id {
+		return nil, false
+	}
+	fk := flightKey{key: key, id: id}
+	if f = c.flights[fk]; f != nil {
+		return f, false
+	}
+	f = &flight{done: make(chan struct{})}
+	c.flights[fk] = f
+	return f, true
+}
+
+func (c *Client) finishFlight(key, id string, f *flight) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	fk := flightKey{key: key, id: id}
+	if c.flights[fk] == f {
+		delete(c.flights, fk)
+		close(f.done)
+	}
+}
+
+func (c *Client) resetFlightsLocked() {
+	for _, f := range c.flights {
+		close(f.done)
+	}
+	c.flights = make(map[flightKey]*flight)
+}
+
+func (c *Client) populate(
+	ctx context.Context,
+	ttl time.Duration,
+	key, id string,
+	fn func(ctx context.Context, key string) (val string, err error),
+	f *flight,
+) (val string, err error) {
+	cleanup := true
+	defer c.finishFlight(key, id, f)
+	defer func() {
+		if cleanup {
+			delkey.Exec(context.Background(), c.client, []string{key}, []string{id})
+		}
+	}()
+
+	if c.useLuaLock {
+		val, err = acquireLock.Exec(ctx, c.client, []string{key}, []string{id, strconv.FormatInt(ttl.Milliseconds(), 10)}).ToString()
+	} else {
+		val, err = c.client.Do(ctx, c.client.B().Set().Key(key).Value(id).Nx().Get().Px(ttl).Build()).ToString()
+	}
+	if err == nil {
+		cleanup = false
+		return val, nil
+	}
+	if !rueidis.IsRedisNil(err) {
+		return val, err
+	}
+
+	ctx = context.WithValue(ctx, ttlKey, &ttl)
+	if val, err = fn(ctx, key); err == nil {
+		err = setkey.Exec(ctx, c.client, []string{key}, []string{id, val, strconv.FormatInt(ttl.Milliseconds(), 10)}).Error()
+	}
+	if err == nil {
+		cleanup = false
+	}
+	return val, err
+}