From 8ba7676873c36afb42d125a639fbf3788a942c86 Mon Sep 17 00:00:00 2001 From: Andrii Makarets <60238228+Makarechi@users.noreply.github.com> Date: Wed, 22 Jul 2026 21:36:46 +0200 Subject: [PATCH 1/6] fix(rueidisaside): clean up lock after canceled acquisition --- rueidisaside/aside.go | 19 +++++++- rueidisaside/aside_test.go | 94 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 2 deletions(-) diff --git a/rueidisaside/aside.go b/rueidisaside/aside.go index b40290ab..8a836845 100644 --- a/rueidisaside/aside.go +++ b/rueidisaside/aside.go @@ -157,14 +157,29 @@ retry: val, err := resp.ToString() if rueidis.IsRedisNil(err) && fn != nil { // cache miss, prepare to populate the value by fn() + if err = ctx.Err(); err != nil { + return val, err + } var id string if id, err = c.keepalive(); err == nil { // acquire client id + if err = ctx.Err(); err != nil { + return val, err + } + // Wait for the lock response even if the caller is canceled. A canceled + // write can still reach Redis, so ownership must be known before cleanup. + acquireCtx := context.WithoutCancel(ctx) if c.useLuaLock { - val, err = acquireLock.Exec(ctx, c.client, []string{key}, []string{id, strconv.FormatInt(ttl.Milliseconds(), 10)}).ToString() + val, err = acquireLock.Exec(acquireCtx, 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() + val, err = c.client.Do(acquireCtx, c.client.B().Set().Key(key).Value(id).Nx().Get().Px(ttl).Build()).ToString() } + if ctxErr := ctx.Err(); ctxErr != nil { + if rueidis.IsRedisNil(err) { + delkey.Exec(context.Background(), c.client, []string{key}, []string{id}) + } + return "", ctxErr + } 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) diff --git a/rueidisaside/aside_test.go b/rueidisaside/aside_test.go index 0ef43a34..260d0125 100644 --- a/rueidisaside/aside_test.go +++ b/rueidisaside/aside_test.go @@ -2,6 +2,7 @@ package rueidisaside import ( "context" + "errors" "math/rand" "strconv" "sync" @@ -14,6 +15,40 @@ import ( var addr = []string{"127.0.0.1:6379"} +type cancelAfterAcquireClient struct { + rueidis.Client + key string + cancel context.CancelFunc + accepted bool +} + +func (c *cancelAfterAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { + commands := cmd.Commands() + if isAcquireCommand(commands, c.key) { + resp := c.Client.Do(ctx, cmd) + if !rueidis.IsRedisNil(resp.Error()) { + return resp + } + c.accepted = true + c.cancel() + if err := ctx.Err(); err != nil { + return rueidis.NewErrorResult(err) + } + return resp + } + return c.Client.Do(ctx, cmd) +} + +func isAcquireCommand(commands []string, key string) bool { + if len(commands) == 7 && commands[0] == "SET" { + return commands[1] == key && commands[3] == "NX" && commands[4] == "GET" + } + if len(commands) == 6 && (commands[0] == "EVALSHA" || commands[0] == "EVAL") { + return commands[2] == "1" && commands[3] == key + } + return false +} + func makeClient(t *testing.T, addr []string) CacheAsideClient { client, err := NewClient(ClientOption{ ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, @@ -303,6 +338,65 @@ func TestWriteCancelLL(t *testing.T) { } } +func TestAcquireCancelCleanup(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "set" + if useLuaLock { + name = "lua" + } + t.Run(name, func(t *testing.T) { + key := strconv.Itoa(rand.Int()) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var wrapped *cancelAfterAcquireClient + client, err := NewClient(ClientOption{ + ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, + ClientTTL: time.Second, + UseLuaLock: useLuaLock, + ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { + client, err := rueidis.NewClient(option) + if err != nil { + return nil, err + } + wrapped = &cancelAfterAcquireClient{Client: client, key: key, cancel: cancel} + return wrapped, nil + }, + }) + if err != nil { + t.Fatal(err) + } + if err = client.Client().Do(context.Background(), client.Client().B().Del().Key(key).Build()).Error(); err != nil { + t.Fatal(err) + } + defer func() { + client.Client().Do(context.Background(), client.Client().B().Del().Key(key).Build()) + client.Close() + }() + + loaderCalled := false + _, err = client.Get(ctx, time.Minute, key, func(context.Context, string) (string, error) { + loaderCalled = true + return "value", nil + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + if !wrapped.accepted { + t.Fatal("Redis did not accept the lock command") + } + if loaderCalled { + t.Fatal("loader should not be called when lock acquisition is canceled") + } + + val, err := client.Client().Do(context.Background(), client.Client().B().Get().Key(key).Build()).ToString() + if !rueidis.IsRedisNil(err) { + t.Fatalf("expected canceled lock acquisition to release the placeholder, got value %q and error %v", val, err) + } + }) + } +} + func TestTimeout(t *testing.T) { client := makeClient(t, addr).(*Client) defer client.Close() From c9aa4692139ef0b2067107b9b2d1e2dfdd4309e6 Mon Sep 17 00:00:00 2001 From: Andrii Makarets <60238228+Makarechi@users.noreply.github.com> Date: Sat, 25 Jul 2026 20:53:25 +0200 Subject: [PATCH 2/6] refactor(rueidisaside): remove redundant context checks --- rueidisaside/aside.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/rueidisaside/aside.go b/rueidisaside/aside.go index 8a836845..9e314fb4 100644 --- a/rueidisaside/aside.go +++ b/rueidisaside/aside.go @@ -157,14 +157,8 @@ retry: val, err := resp.ToString() if rueidis.IsRedisNil(err) && fn != nil { // cache miss, prepare to populate the value by fn() - if err = ctx.Err(); err != nil { - return val, err - } var id string if id, err = c.keepalive(); err == nil { // acquire client id - if err = ctx.Err(); err != nil { - return val, err - } // Wait for the lock response even if the caller is canceled. A canceled // write can still reach Redis, so ownership must be known before cleanup. acquireCtx := context.WithoutCancel(ctx) From 51940d3153578e0428a202661fbee9728ef6471c Mon Sep 17 00:00:00 2001 From: Andrii Makarets <60238228+Makarechi@users.noreply.github.com> Date: Sun, 26 Jul 2026 19:55:03 +0200 Subject: [PATCH 3/6] fix(rueidisaside): guard canceled cache fills --- rueidisaside/README.md | 6 +- rueidisaside/aside.go | 52 ++- rueidisaside/aside_test.go | 94 ------ rueidisaside/attempt.go | 240 +++++++++++++ rueidisaside/attempt_test.go | 632 +++++++++++++++++++++++++++++++++++ rueidisaside/flight.go | 44 +++ rueidisaside/flight_test.go | 16 + 7 files changed, 956 insertions(+), 128 deletions(-) create mode 100644 rueidisaside/attempt.go create mode 100644 rueidisaside/attempt_test.go create mode 100644 rueidisaside/flight.go create mode 100644 rueidisaside/flight_test.go diff --git a/rueidisaside/README.md b/rueidisaside/README.md index db44bfe4..c6b8ed5d 100644 --- a/rueidisaside/README.md +++ b/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 inside a Lua script. +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) diff --git a/rueidisaside/aside.go b/rueidisaside/aside.go index 9e314fb4..647cccec 100644 --- a/rueidisaside/aside.go +++ b/rueidisaside/aside.go @@ -5,9 +5,9 @@ import ( "encoding/binary" "encoding/hex" "math/rand" - "strconv" "strings" "sync" + "sync/atomic" "time" "unsafe" @@ -19,7 +19,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,6 +35,7 @@ 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, } @@ -55,8 +56,10 @@ type Client struct { client rueidis.Client ctx context.Context waits map[string]chan struct{} + flights map[flightKey]*flight cancel context.CancelFunc id string + attempts atomic.Uint64 ttl time.Duration mu sync.Mutex useLuaLock bool @@ -72,6 +75,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,32 +162,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 - // Wait for the lock response even if the caller is canceled. A canceled - // write can still reach Redis, so ownership must be known before cleanup. - acquireCtx := context.WithoutCancel(ctx) - if c.useLuaLock { - val, err = acquireLock.Exec(acquireCtx, c.client, []string{key}, []string{id, strconv.FormatInt(ttl.Milliseconds(), 10)}).ToString() - } else { - val, err = c.client.Do(acquireCtx, 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 ctxErr := ctx.Err(); ctxErr != nil { - if rueidis.IsRedisNil(err) { - delkey.Exec(context.Background(), c.client, []string{key}, []string{id}) - } - return "", ctxErr - } - 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) } } @@ -229,7 +223,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() } @@ -249,8 +243,4 @@ func OverrideCacheTTL(ctx context.Context, ttl time.Duration) { } } -var ( - delkey = rueidis.NewLuaScript(`if redis.call("GET",KEYS[1]) == ARGV[1] then return redis.call("DEL",KEYS[1]) else return 0 end`) - setkey = rueidis.NewLuaScript(`if redis.call("GET",KEYS[1]) == ARGV[1] then return redis.call("SET",KEYS[1],ARGV[2],"PX",ARGV[3]) else return 0 end`) - acquireLock = rueidis.NewLuaScript(`if redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) then return nil else return redis.call("GET", KEYS[1]) end`) -) +var delkey = rueidis.NewLuaScript(`if redis.call("GET",KEYS[1]) == ARGV[1] then return redis.call("DEL",KEYS[1]) else return 0 end`) diff --git a/rueidisaside/aside_test.go b/rueidisaside/aside_test.go index 260d0125..0ef43a34 100644 --- a/rueidisaside/aside_test.go +++ b/rueidisaside/aside_test.go @@ -2,7 +2,6 @@ package rueidisaside import ( "context" - "errors" "math/rand" "strconv" "sync" @@ -15,40 +14,6 @@ import ( var addr = []string{"127.0.0.1:6379"} -type cancelAfterAcquireClient struct { - rueidis.Client - key string - cancel context.CancelFunc - accepted bool -} - -func (c *cancelAfterAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { - commands := cmd.Commands() - if isAcquireCommand(commands, c.key) { - resp := c.Client.Do(ctx, cmd) - if !rueidis.IsRedisNil(resp.Error()) { - return resp - } - c.accepted = true - c.cancel() - if err := ctx.Err(); err != nil { - return rueidis.NewErrorResult(err) - } - return resp - } - return c.Client.Do(ctx, cmd) -} - -func isAcquireCommand(commands []string, key string) bool { - if len(commands) == 7 && commands[0] == "SET" { - return commands[1] == key && commands[3] == "NX" && commands[4] == "GET" - } - if len(commands) == 6 && (commands[0] == "EVALSHA" || commands[0] == "EVAL") { - return commands[2] == "1" && commands[3] == key - } - return false -} - func makeClient(t *testing.T, addr []string) CacheAsideClient { client, err := NewClient(ClientOption{ ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, @@ -338,65 +303,6 @@ func TestWriteCancelLL(t *testing.T) { } } -func TestAcquireCancelCleanup(t *testing.T) { - for _, useLuaLock := range []bool{false, true} { - name := "set" - if useLuaLock { - name = "lua" - } - t.Run(name, func(t *testing.T) { - key := strconv.Itoa(rand.Int()) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - var wrapped *cancelAfterAcquireClient - client, err := NewClient(ClientOption{ - ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, - ClientTTL: time.Second, - UseLuaLock: useLuaLock, - ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { - client, err := rueidis.NewClient(option) - if err != nil { - return nil, err - } - wrapped = &cancelAfterAcquireClient{Client: client, key: key, cancel: cancel} - return wrapped, nil - }, - }) - if err != nil { - t.Fatal(err) - } - if err = client.Client().Do(context.Background(), client.Client().B().Del().Key(key).Build()).Error(); err != nil { - t.Fatal(err) - } - defer func() { - client.Client().Do(context.Background(), client.Client().B().Del().Key(key).Build()) - client.Close() - }() - - loaderCalled := false - _, err = client.Get(ctx, time.Minute, key, func(context.Context, string) (string, error) { - loaderCalled = true - return "value", nil - }) - if !errors.Is(err, context.Canceled) { - t.Fatalf("expected context.Canceled, got %v", err) - } - if !wrapped.accepted { - t.Fatal("Redis did not accept the lock command") - } - if loaderCalled { - t.Fatal("loader should not be called when lock acquisition is canceled") - } - - val, err := client.Client().Do(context.Background(), client.Client().B().Get().Key(key).Build()).ToString() - if !rueidis.IsRedisNil(err) { - t.Fatalf("expected canceled lock acquisition to release the placeholder, got value %q and error %v", val, err) - } - }) - } -} - func TestTimeout(t *testing.T) { client := makeClient(t, addr).(*Client) defer client.Close() diff --git a/rueidisaside/attempt.go b/rueidisaside/attempt.go new file mode 100644 index 00000000..775d5f23 --- /dev/null +++ b/rueidisaside/attempt.go @@ -0,0 +1,240 @@ +package rueidisaside + +import ( + "context" + "fmt" + "strconv" + "sync" + "time" + + "github.com/redis/rueidis" + "github.com/redis/rueidis/internal/cmds" +) + +const attemptGuardPrefix = "rueidisaside:attempt:" + +const ( + acquireRevoked int64 = iota + acquireSucceeded + acquireOccupied +) + +type attempt struct { + guard string + token string +} + +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) { + var cleanup bool + a := c.newAttempt(key, id) + defer c.finishFlight(key, id, f) + defer func() { + if cleanup { + // Cleanup must finish before the flight is released. If Redis cannot + // confirm it, rotate the client id so a late command cannot match a + // subsequent flight from this client. + cleanupCtx, cancel := context.WithTimeout(context.Background(), c.ttl) + cleanupErr := revokeAttempt.Exec(cleanupCtx, c.client, []string{key, a.guard}, []string{id, a.token}).Error() + cancel() + if cleanupErr != nil { + c.abandonGeneration(id) + if err == nil { + err = cleanupErr + } + } + } + }() + + deadline, _ := ctx.Deadline() + guardTTL := time.Until(deadline) + if guardTTL <= 0 { + if err = ctx.Err(); err == nil { + err = context.DeadlineExceeded + } + return "", err + } + if guardTTL < time.Millisecond { + guardTTL = time.Millisecond + } + if err = c.client.Do(ctx, c.client.B().Set().Key(a.guard).Value(a.token).Nx().Px(guardTTL).Build()).Error(); err != nil { + return "", err + } + cleanup = true + + var acquired bool + if val, acquired, err = c.acquire(ctx, key, id, a); err != nil { + return val, err + } + if !acquired { + cleanup = false + return val, nil + } + + ctx = context.WithValue(ctx, ttlKey, &ttl) + if val, err = fn(ctx, key); err != nil { + return val, err + } + var status int64 + status, err = guardedSet.Exec( + ctx, + c.client, + []string{key, a.guard}, + []string{id, a.token, val, strconv.FormatInt(ttl.Milliseconds(), 10)}, + ).AsInt64() + if err == nil && status == 1 { + cleanup = false + } + return val, err +} + +func (c *Client) acquire(ctx context.Context, key, id string, a attempt) (val string, acquired bool, err error) { + script := guardedAcquire + if c.useLuaLock { + script = guardedAcquireLegacy + } + resp, err := script.Exec( + ctx, + c.client, + []string{key, a.guard}, + []string{id, a.token}, + ).ToArray() + if err != nil { + return "", false, err + } + if len(resp) == 0 { + return "", false, fmt.Errorf("rueidisaside: empty acquire response") + } + status, err := resp[0].AsInt64() + if err != nil { + return "", false, fmt.Errorf("rueidisaside: invalid acquire response: %w", err) + } + switch status { + case acquireRevoked: + if err = ctx.Err(); err == nil { + err = context.DeadlineExceeded + } + return "", false, err + case acquireSucceeded: + return "", true, nil + case acquireOccupied: + if len(resp) != 2 { + return "", false, fmt.Errorf("rueidisaside: occupied acquire response has %d elements", len(resp)) + } + val, err = resp[1].ToString() + return val, false, err + default: + return "", false, fmt.Errorf("rueidisaside: unknown acquire status %d", status) + } +} + +func (c *Client) newAttempt(key, id string) attempt { + n := c.attempts.Add(1) + token := id + ":" + strconv.FormatUint(n, 10) + return attempt{ + guard: attemptGuardKey(key, token), + token: token, + } +} + +func (c *Client) abandonGeneration(id string) { + c.mu.Lock() + if c.id == id { + c.id = "" + } + c.mu.Unlock() +} + +var ( + guardedAcquire = rueidis.NewLuaScript(` +if redis.call("GET", KEYS[2]) ~= ARGV[2] then + return {0} +end +local ttl = redis.call("PTTL", KEYS[2]) +if ttl <= 0 then + return {0} +end +local current = redis.call("SET", KEYS[1], ARGV[1], "NX", "GET", "PX", ttl) +if current then + redis.call("DEL", KEYS[2]) + return {2, current} +end +return {1} +`) + guardedAcquireLegacy = rueidis.NewLuaScript(` +if redis.call("GET", KEYS[2]) ~= ARGV[2] then + return {0} +end +local ttl = redis.call("PTTL", KEYS[2]) +if ttl <= 0 then + return {0} +end +local current = redis.call("GET", KEYS[1]) +if current then + redis.call("DEL", KEYS[2]) + return {2, current} +end +redis.call("SET", KEYS[1], ARGV[1], "PX", ttl) +return {1} +`) + guardedSet = rueidis.NewLuaScript(` +if redis.call("GET", KEYS[2]) ~= ARGV[2] then + return 0 +end +if redis.call("GET", KEYS[1]) == ARGV[1] then + redis.call("SET", KEYS[1], ARGV[3], "PX", ARGV[4]) + redis.call("DEL", KEYS[2]) + return 1 +end +redis.call("DEL", KEYS[2]) +return 0 +`) + revokeAttempt = rueidis.NewLuaScriptRetryable(` +if redis.call("GET", KEYS[2]) == ARGV[2] then + redis.call("DEL", KEYS[2]) +end +if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("DEL", KEYS[1]) +end +return 0 +`) +) + +const slotTagAlphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_" + +var ( + slotTags [16384][3]byte + slotTagsOnce sync.Once +) + +func attemptGuardKey(key, token string) string { + slotTagsOnce.Do(initSlotTags) + tag := slotTags[cmds.Slot(key)] + return attemptGuardPrefix + "{" + string(tag[:]) + "}:" + token +} + +func initSlotTags() { + // Lua requires the cache key and its attempt guard to share a cluster slot. + // Build one printable hash tag for every possible slot. + remaining := len(slotTags) + for i := 0; i < len(slotTagAlphabet) && remaining != 0; i++ { + for j := 0; j < len(slotTagAlphabet) && remaining != 0; j++ { + for k := 0; k < len(slotTagAlphabet) && remaining != 0; k++ { + tag := [3]byte{slotTagAlphabet[i], slotTagAlphabet[j], slotTagAlphabet[k]} + slot := cmds.Slot(string(tag[:])) + if slotTags[slot][0] == 0 { + slotTags[slot] = tag + remaining-- + } + } + } + } + if remaining != 0 { + panic("rueidisaside: unable to generate Redis slot tags") + } +} diff --git a/rueidisaside/attempt_test.go b/rueidisaside/attempt_test.go new file mode 100644 index 00000000..013030c2 --- /dev/null +++ b/rueidisaside/attempt_test.go @@ -0,0 +1,632 @@ +package rueidisaside + +import ( + "context" + "errors" + "math/rand" + "strconv" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/redis/rueidis" + "github.com/redis/rueidis/internal/cmds" +) + +type delayedAcquireClient struct { + rueidis.Client + key string + once sync.Once + blocked chan struct{} + pending rueidis.Completed + guard string +} + +func (c *delayedAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { + delayed := false + if isGuardedScript(cmd.Commands(), c.key, 2) { + c.once.Do(func() { + delayed = true + c.pending = cmd + c.guard = cmd.Commands()[4] + close(c.blocked) + }) + } + if delayed { + <-ctx.Done() + return rueidis.NewErrorResult(ctx.Err()) + } + return c.Client.Do(ctx, cmd) +} + +func (c *delayedAcquireClient) executePending() rueidis.RedisResult { + return c.Client.Do(context.Background(), c.pending) +} + +type blockingAcquireClient struct { + rueidis.Client + key string + started chan struct{} + release chan struct{} + calls atomic.Int64 +} + +type blockingCleanupClient struct { + rueidis.Client + key string + cancel context.CancelFunc + cleanupStarted chan struct{} + cleanupOnce sync.Once + calls atomic.Int64 +} + +func (c *blockingCleanupClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { + if isGuardedScript(cmd.Commands(), c.key, 2) { + if c.calls.Add(1) == 1 { + resp := c.Client.Do(ctx, cmd) + if resp.Error() == nil { + c.cancel() + return rueidis.NewErrorResult(ctx.Err()) + } + return resp + } + c.cleanupOnce.Do(func() { close(c.cleanupStarted) }) + <-ctx.Done() + return rueidis.NewErrorResult(ctx.Err()) + } + return c.Client.Do(ctx, cmd) +} + +func (c *blockingAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { + if isGuardedScript(cmd.Commands(), c.key, 2) { + if c.calls.Add(1) == 1 { + close(c.started) + select { + case <-c.release: + case <-ctx.Done(): + return rueidis.NewErrorResult(ctx.Err()) + } + } + } + return c.Client.Do(ctx, cmd) +} + +func isGuardedScript(commands []string, key string, argCount int) bool { + if len(commands) != 5+argCount || (commands[0] != "EVALSHA" && commands[0] != "EVAL") { + return false + } + return commands[2] == "2" && commands[3] == key +} + +func preloadAcquire(client rueidis.Client, useLuaLock bool, key string) error { + script := guardedAcquire + if useLuaLock { + script = guardedAcquireLegacy + } + guard := attemptGuardKey(key, "preload") + _, err := script.Exec( + context.Background(), + client, + []string{key, guard}, + []string{"preload-id", "preload"}, + ).ToArray() + return err +} + +func newDelayedAcquireCache(t *testing.T, key string, useLuaLock bool) (*Client, *delayedAcquireClient) { + t.Helper() + var wrapped *delayedAcquireClient + client, err := NewClient(ClientOption{ + ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, + ClientTTL: time.Second, + UseLuaLock: useLuaLock, + ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { + client, err := rueidis.NewClient(option) + if err != nil { + return nil, err + } + if err = preloadAcquire(client, useLuaLock, key); err != nil { + client.Close() + return nil, err + } + wrapped = &delayedAcquireClient{ + Client: client, + key: key, + blocked: make(chan struct{}), + } + return wrapped, nil + }, + }) + if err != nil { + t.Fatal(err) + } + return client.(*Client), wrapped +} + +func newBlockingAcquireCache(t *testing.T, key string, useLuaLock bool) (*Client, *blockingAcquireClient) { + t.Helper() + var wrapped *blockingAcquireClient + client, err := NewClient(ClientOption{ + ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, + ClientTTL: time.Second, + UseLuaLock: useLuaLock, + ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { + client, err := rueidis.NewClient(option) + if err != nil { + return nil, err + } + if err = preloadAcquire(client, useLuaLock, key); err != nil { + client.Close() + return nil, err + } + wrapped = &blockingAcquireClient{ + Client: client, + key: key, + started: make(chan struct{}), + release: make(chan struct{}), + } + return wrapped, nil + }, + }) + if err != nil { + t.Fatal(err) + } + return client.(*Client), wrapped +} + +func newBlockingCleanupCache( + t *testing.T, + key string, + useLuaLock bool, + cancel context.CancelFunc, +) (*Client, *blockingCleanupClient) { + t.Helper() + var wrapped *blockingCleanupClient + client, err := NewClient(ClientOption{ + ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, + ClientTTL: 100 * time.Millisecond, + UseLuaLock: useLuaLock, + ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { + client, err := rueidis.NewClient(option) + if err != nil { + return nil, err + } + if err = preloadAcquire(client, useLuaLock, key); err != nil { + client.Close() + return nil, err + } + wrapped = &blockingCleanupClient{ + Client: client, + key: key, + cancel: cancel, + cleanupStarted: make(chan struct{}), + } + return wrapped, nil + }, + }) + if err != nil { + t.Fatal(err) + } + return client.(*Client), wrapped +} + +func TestCanceledAcquireCannotArriveLate(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "redis-7" + if useLuaLock { + name = "legacy" + } + t.Run(name, func(t *testing.T) { + key := "late-acquire-" + strconv.Itoa(rand.Int()) + client, wrapped := newDelayedAcquireCache(t, key, useLuaLock) + defer client.Close() + defer client.client.Do(context.Background(), client.client.B().Del().Key(key).Build()) + + ctx, cancel := context.WithCancel(context.Background()) + first := make(chan getResult, 1) + var firstLoaderCalled atomic.Bool + go func() { + val, err := client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { + firstLoaderCalled.Store(true) + return "old", nil + }) + first <- getResult{val: val, err: err} + }() + + select { + case <-wrapped.blocked: + case <-time.After(time.Second): + t.Fatal("acquire did not reach the delay point") + } + + second := make(chan getResult, 1) + go func() { + val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + return "new", nil + }) + second <- getResult{val: val, err: err} + }() + + cancel() + select { + case result := <-first: + if !errors.Is(result.err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", result.err) + } + case <-time.After(time.Second): + t.Fatal("canceled leader did not return") + } + if firstLoaderCalled.Load() { + t.Fatal("canceled acquire called the loader") + } + + select { + case result := <-second: + if result.err != nil || result.val != "new" { + t.Fatalf("follower returned %q, %v", result.val, result.err) + } + case <-time.After(2 * time.Second): + t.Fatal("follower did not retry after cleanup") + } + + resp, err := wrapped.executePending().ToArray() + if err != nil { + t.Fatal(err) + } + status, err := resp[0].AsInt64() + if err != nil || status != acquireRevoked { + t.Fatalf("late acquire returned status %d, %v", status, err) + } + + val, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() + if err != nil || val != "new" { + t.Fatalf("late acquire changed the cached value to %q, %v", val, err) + } + if err = client.client.Do(context.Background(), client.client.B().Get().Key(wrapped.guard).Build()).Error(); !rueidis.IsRedisNil(err) { + t.Fatalf("attempt guard was not revoked: %v", err) + } + }) + } +} + +func TestSingleflightSerializesAcquire(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "redis-7" + if useLuaLock { + name = "legacy" + } + t.Run(name, func(t *testing.T) { + key := "singleflight-" + strconv.Itoa(rand.Int()) + client, wrapped := newBlockingAcquireCache(t, key, useLuaLock) + defer client.Close() + defer client.client.Do(context.Background(), client.client.B().Del().Key(key).Build()) + + first := make(chan getResult, 1) + go func() { + val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + return "value", nil + }) + first <- getResult{val: val, err: err} + }() + select { + case <-wrapped.started: + case <-time.After(time.Second): + t.Fatal("first acquire did not start") + } + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + var secondLoaderCalled atomic.Bool + _, secondErr := client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { + secondLoaderCalled.Store(true) + return "second", nil + }) + close(wrapped.release) + + select { + case result := <-first: + if result.err != nil || result.val != "value" { + t.Fatalf("leader returned %q, %v", result.val, result.err) + } + case <-time.After(time.Second): + t.Fatal("leader did not finish") + } + if !errors.Is(secondErr, context.DeadlineExceeded) { + t.Fatalf("expected follower deadline, got %v", secondErr) + } + if secondLoaderCalled.Load() { + t.Fatal("follower called the loader") + } + if calls := wrapped.calls.Load(); calls != 1 { + t.Fatalf("expected one acquire, got %d", calls) + } + }) + } +} + +func TestCleanupTimeoutRotatesClientGeneration(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "redis-7" + if useLuaLock { + name = "legacy" + } + t.Run(name, func(t *testing.T) { + key := "cleanup-timeout-" + strconv.Itoa(rand.Int()) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + client, wrapped := newBlockingCleanupCache(t, key, useLuaLock, cancel) + defer client.Close() + defer client.client.Do(context.Background(), client.client.B().Del().Key(key).Build()) + + started := time.Now() + loaderCalled := false + _, err := client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { + loaderCalled = true + return "value", nil + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + if elapsed := time.Since(started); elapsed > time.Second { + t.Fatalf("cleanup exceeded its bound: %s", elapsed) + } + if loaderCalled { + t.Fatal("loader ran after the canceled acquire") + } + select { + case <-wrapped.cleanupStarted: + default: + t.Fatal("cleanup was not attempted") + } + + client.mu.Lock() + id := client.id + client.mu.Unlock() + if id != "" { + t.Fatalf("client generation was not rotated: %q", id) + } + }) + } +} + +type getResult struct { + val string + err error +} + +func TestSingleflightKeepsDifferentKeysConcurrent(t *testing.T) { + client := makeClient(t, addr) + defer client.Close() + keys := []string{ + "parallel-a-" + strconv.Itoa(rand.Int()), + "parallel-b-" + strconv.Itoa(rand.Int()), + } + started := make(chan string, len(keys)) + release := make(chan struct{}) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(release) }) + + results := make(chan getResult, len(keys)) + for _, key := range keys { + key := key + go func() { + val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + started <- key + <-release + return key, nil + }) + results <- getResult{val: val, err: err} + }() + } + for range keys { + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("different keys were serialized") + } + } + releaseOnce.Do(func() { close(release) }) + for range keys { + select { + case result := <-results: + if result.err != nil { + t.Fatal(result.err) + } + case <-time.After(time.Second): + t.Fatal("parallel cache fill did not finish") + } + } +} + +func TestLoaderPanicReleasesFlight(t *testing.T) { + client := makeClient(t, addr) + defer client.Close() + key := "panic-" + strconv.Itoa(rand.Int()) + + func() { + defer func() { + if recover() == nil { + t.Fatal("loader did not panic") + } + }() + _, _ = client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + panic("loader panic") + }) + }() + + val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + return "recovered", nil + }) + if err != nil || val != "recovered" { + t.Fatalf("flight was not released after panic: %q, %v", val, err) + } +} + +func TestDelayedSetCannotOverwriteNewAttempt(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "redis-7" + if useLuaLock { + name = "legacy" + } + t.Run(name, func(t *testing.T) { + client := makeClient(t, addr).(*Client) + client.useLuaLock = useLuaLock + defer client.Close() + + key := "late-set-" + strconv.Itoa(rand.Int()) + id := PlaceholderPrefix + "late-set-owner" + first := attempt{token: "first-token"} + first.guard = attemptGuardKey(key, first.token) + second := attempt{token: "second-token"} + second.guard = attemptGuardKey(key, second.token) + defer client.client.Do( + context.Background(), + client.client.B().Del().Key(key, first.guard, second.guard).Build(), + ) + + if err := client.client.Do(context.Background(), client.client.B().Set().Key(first.guard).Value(first.token).Px(time.Second).Build()).Error(); err != nil { + t.Fatal(err) + } + if _, acquired, err := client.acquire(context.Background(), key, id, first); err != nil || !acquired { + t.Fatalf("first acquire failed: acquired=%v err=%v", acquired, err) + } + if err := revokeAttempt.Exec(context.Background(), client.client, []string{key, first.guard}, []string{id, first.token}).Error(); err != nil { + t.Fatal(err) + } + + if err := client.client.Do(context.Background(), client.client.B().Set().Key(second.guard).Value(second.token).Px(time.Second).Build()).Error(); err != nil { + t.Fatal(err) + } + if _, acquired, err := client.acquire(context.Background(), key, id, second); err != nil || !acquired { + t.Fatalf("second acquire failed: acquired=%v err=%v", acquired, err) + } + if status, err := guardedSet.Exec( + context.Background(), + client.client, + []string{key, second.guard}, + []string{id, second.token, "new", "1000"}, + ).AsInt64(); err != nil || status != 1 { + t.Fatalf("second set failed: status=%d err=%v", status, err) + } + if status, err := guardedSet.Exec( + context.Background(), + client.client, + []string{key, first.guard}, + []string{id, first.token, "old", "1000"}, + ).AsInt64(); err != nil || status != 0 { + t.Fatalf("delayed set returned status=%d err=%v", status, err) + } + + val, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() + if err != nil || val != "new" { + t.Fatalf("delayed set changed the value to %q, %v", val, err) + } + }) + } +} + +func TestCleanupRemovesAcceptedAcquire(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "redis-7" + if useLuaLock { + name = "legacy" + } + t.Run(name, func(t *testing.T) { + client := makeClient(t, addr).(*Client) + client.useLuaLock = useLuaLock + defer client.Close() + + key := "accepted-acquire-" + strconv.Itoa(rand.Int()) + id := PlaceholderPrefix + "accepted-owner" + a := attempt{token: "accepted-token"} + a.guard = attemptGuardKey(key, a.token) + defer client.client.Do(context.Background(), client.client.B().Del().Key(key, a.guard).Build()) + + if err := client.client.Do(context.Background(), client.client.B().Set().Key(a.guard).Value(a.token).Px(time.Second).Build()).Error(); err != nil { + t.Fatal(err) + } + if _, acquired, err := client.acquire(context.Background(), key, id, a); err != nil || !acquired { + t.Fatalf("acquire failed: acquired=%v err=%v", acquired, err) + } + if err := revokeAttempt.Exec(context.Background(), client.client, []string{key, a.guard}, []string{id, a.token}).Error(); err != nil { + t.Fatal(err) + } + if err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).Error(); !rueidis.IsRedisNil(err) { + t.Fatalf("placeholder was not removed: %v", err) + } + if err := client.client.Do(context.Background(), client.client.B().Get().Key(a.guard).Build()).Error(); !rueidis.IsRedisNil(err) { + t.Fatalf("guard was not removed: %v", err) + } + }) + } +} + +func TestAcquireTTLDoesNotOutliveGuard(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "redis-7" + if useLuaLock { + name = "legacy" + } + t.Run(name, func(t *testing.T) { + client := makeClient(t, addr).(*Client) + client.useLuaLock = useLuaLock + defer client.Close() + + key := "guard-ttl-" + strconv.Itoa(rand.Int()) + id := PlaceholderPrefix + "ttl-owner" + a := attempt{token: "ttl-token"} + a.guard = attemptGuardKey(key, a.token) + defer client.client.Do(context.Background(), client.client.B().Del().Key(key, a.guard).Build()) + + if err := client.client.Do(context.Background(), client.client.B().Set().Key(a.guard).Value(a.token).Px(time.Second).Build()).Error(); err != nil { + t.Fatal(err) + } + if _, acquired, err := client.acquire(context.Background(), key, id, a); err != nil || !acquired { + t.Fatalf("acquire failed: acquired=%v err=%v", acquired, err) + } + pttls := client.client.DoMulti( + context.Background(), + client.client.B().Pttl().Key(key).Build(), + client.client.B().Pttl().Key(a.guard).Build(), + ) + keyTTL, err := pttls[0].AsInt64() + if err != nil { + t.Fatal(err) + } + guardTTL, err := pttls[1].AsInt64() + if err != nil { + t.Fatal(err) + } + if keyTTL <= 0 || guardTTL <= 0 || keyTTL > guardTTL+20 { + t.Fatalf("unexpected TTLs: key=%dms guard=%dms", keyTTL, guardTTL) + } + }) + } +} + +func TestAttemptGuardKeyUsesSameSlot(t *testing.T) { + slotTagsOnce.Do(initSlotTags) + for slot, tag := range slotTags { + if got := cmds.Slot(string(tag[:])); got != uint16(slot) { + t.Fatalf("slot tag %q maps to %d, expected %d", tag, got, slot) + } + } + + keys := []string{ + "", + "plain", + "key{tag}suffix", + "{}", + "unclosed{tag", + "empty{}tag", + "\x00binary\xff", + "nested{{tag}}", + } + for _, key := range keys { + guard := attemptGuardKey(key, "token") + if got, want := cmds.Slot(guard), cmds.Slot(key); got != want { + t.Fatalf("guard for %q maps to slot %d, expected %d", key, got, want) + } + } +} diff --git a/rueidisaside/flight.go b/rueidisaside/flight.go new file mode 100644 index 00000000..edca82f3 --- /dev/null +++ b/rueidisaside/flight.go @@ -0,0 +1,44 @@ +package rueidisaside + +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) +} diff --git a/rueidisaside/flight_test.go b/rueidisaside/flight_test.go new file mode 100644 index 00000000..11185b37 --- /dev/null +++ b/rueidisaside/flight_test.go @@ -0,0 +1,16 @@ +package rueidisaside + +import "testing" + +func TestBeginFlightRejectsStaleGeneration(t *testing.T) { + c := &Client{ + id: "new-generation", + flights: make(map[flightKey]*flight), + } + if f, _ := c.beginFlight("key", "old-generation"); f != nil { + t.Fatal("stale generation created a flight") + } + if f, leader := c.beginFlight("key", "new-generation"); f == nil || !leader { + t.Fatal("current generation did not create a flight") + } +} From 1c592a6117f3b05f81df90ecb437d64938b4e96f Mon Sep 17 00:00:00 2001 From: Andrii Makarets <60238228+Makarechi@users.noreply.github.com> Date: Mon, 27 Jul 2026 00:51:14 +0200 Subject: [PATCH 4/6] fix(rueidisaside): order canceled fills on one pipeline --- rueidisaside/README.md | 2 +- rueidisaside/aside.go | 9 +- rueidisaside/aside_test.go | 7 +- rueidisaside/attempt.go | 240 ------------- rueidisaside/attempt_test.go | 632 ----------------------------------- rueidisaside/flight.go | 46 +++ rueidisaside/flight_test.go | 293 +++++++++++++++- 7 files changed, 350 insertions(+), 879 deletions(-) delete mode 100644 rueidisaside/attempt.go delete mode 100644 rueidisaside/attempt_test.go diff --git a/rueidisaside/README.md b/rueidisaside/README.md index c6b8ed5d..ce2dda63 100644 --- a/rueidisaside/README.md +++ b/rueidisaside/README.md @@ -110,7 +110,7 @@ func main() { ## Limitation -By default, lock acquisition uses the Redis 7.0 `SET NX GET` behavior inside a Lua script. +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: diff --git a/rueidisaside/aside.go b/rueidisaside/aside.go index 647cccec..2ec51605 100644 --- a/rueidisaside/aside.go +++ b/rueidisaside/aside.go @@ -7,7 +7,6 @@ import ( "math/rand" "strings" "sync" - "sync/atomic" "time" "unsafe" @@ -39,6 +38,7 @@ func NewClient(option ClientOption) (cc CacheAsideClient, err error) { 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) @@ -59,7 +59,6 @@ type Client struct { flights map[flightKey]*flight cancel context.CancelFunc id string - attempts atomic.Uint64 ttl time.Duration mu sync.Mutex useLuaLock bool @@ -243,4 +242,8 @@ func OverrideCacheTTL(ctx context.Context, ttl time.Duration) { } } -var delkey = rueidis.NewLuaScript(`if redis.call("GET",KEYS[1]) == ARGV[1] then return redis.call("DEL",KEYS[1]) else return 0 end`) +var ( + delkey = rueidis.NewLuaScript(`if redis.call("GET",KEYS[1]) == ARGV[1] then return redis.call("DEL",KEYS[1]) else return 0 end`) + setkey = rueidis.NewLuaScript(`if redis.call("GET",KEYS[1]) == ARGV[1] then return redis.call("SET",KEYS[1],ARGV[2],"PX",ARGV[3]) else return 0 end`) + acquireLock = rueidis.NewLuaScript(`if redis.call("SET", KEYS[1], ARGV[1], "NX", "PX", ARGV[2]) then return nil else return redis.call("GET", KEYS[1]) end`) +) diff --git a/rueidisaside/aside_test.go b/rueidisaside/aside_test.go index 0ef43a34..a617fcaf 100644 --- a/rueidisaside/aside_test.go +++ b/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) { diff --git a/rueidisaside/attempt.go b/rueidisaside/attempt.go deleted file mode 100644 index 775d5f23..00000000 --- a/rueidisaside/attempt.go +++ /dev/null @@ -1,240 +0,0 @@ -package rueidisaside - -import ( - "context" - "fmt" - "strconv" - "sync" - "time" - - "github.com/redis/rueidis" - "github.com/redis/rueidis/internal/cmds" -) - -const attemptGuardPrefix = "rueidisaside:attempt:" - -const ( - acquireRevoked int64 = iota - acquireSucceeded - acquireOccupied -) - -type attempt struct { - guard string - token string -} - -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) { - var cleanup bool - a := c.newAttempt(key, id) - defer c.finishFlight(key, id, f) - defer func() { - if cleanup { - // Cleanup must finish before the flight is released. If Redis cannot - // confirm it, rotate the client id so a late command cannot match a - // subsequent flight from this client. - cleanupCtx, cancel := context.WithTimeout(context.Background(), c.ttl) - cleanupErr := revokeAttempt.Exec(cleanupCtx, c.client, []string{key, a.guard}, []string{id, a.token}).Error() - cancel() - if cleanupErr != nil { - c.abandonGeneration(id) - if err == nil { - err = cleanupErr - } - } - } - }() - - deadline, _ := ctx.Deadline() - guardTTL := time.Until(deadline) - if guardTTL <= 0 { - if err = ctx.Err(); err == nil { - err = context.DeadlineExceeded - } - return "", err - } - if guardTTL < time.Millisecond { - guardTTL = time.Millisecond - } - if err = c.client.Do(ctx, c.client.B().Set().Key(a.guard).Value(a.token).Nx().Px(guardTTL).Build()).Error(); err != nil { - return "", err - } - cleanup = true - - var acquired bool - if val, acquired, err = c.acquire(ctx, key, id, a); err != nil { - return val, err - } - if !acquired { - cleanup = false - return val, nil - } - - ctx = context.WithValue(ctx, ttlKey, &ttl) - if val, err = fn(ctx, key); err != nil { - return val, err - } - var status int64 - status, err = guardedSet.Exec( - ctx, - c.client, - []string{key, a.guard}, - []string{id, a.token, val, strconv.FormatInt(ttl.Milliseconds(), 10)}, - ).AsInt64() - if err == nil && status == 1 { - cleanup = false - } - return val, err -} - -func (c *Client) acquire(ctx context.Context, key, id string, a attempt) (val string, acquired bool, err error) { - script := guardedAcquire - if c.useLuaLock { - script = guardedAcquireLegacy - } - resp, err := script.Exec( - ctx, - c.client, - []string{key, a.guard}, - []string{id, a.token}, - ).ToArray() - if err != nil { - return "", false, err - } - if len(resp) == 0 { - return "", false, fmt.Errorf("rueidisaside: empty acquire response") - } - status, err := resp[0].AsInt64() - if err != nil { - return "", false, fmt.Errorf("rueidisaside: invalid acquire response: %w", err) - } - switch status { - case acquireRevoked: - if err = ctx.Err(); err == nil { - err = context.DeadlineExceeded - } - return "", false, err - case acquireSucceeded: - return "", true, nil - case acquireOccupied: - if len(resp) != 2 { - return "", false, fmt.Errorf("rueidisaside: occupied acquire response has %d elements", len(resp)) - } - val, err = resp[1].ToString() - return val, false, err - default: - return "", false, fmt.Errorf("rueidisaside: unknown acquire status %d", status) - } -} - -func (c *Client) newAttempt(key, id string) attempt { - n := c.attempts.Add(1) - token := id + ":" + strconv.FormatUint(n, 10) - return attempt{ - guard: attemptGuardKey(key, token), - token: token, - } -} - -func (c *Client) abandonGeneration(id string) { - c.mu.Lock() - if c.id == id { - c.id = "" - } - c.mu.Unlock() -} - -var ( - guardedAcquire = rueidis.NewLuaScript(` -if redis.call("GET", KEYS[2]) ~= ARGV[2] then - return {0} -end -local ttl = redis.call("PTTL", KEYS[2]) -if ttl <= 0 then - return {0} -end -local current = redis.call("SET", KEYS[1], ARGV[1], "NX", "GET", "PX", ttl) -if current then - redis.call("DEL", KEYS[2]) - return {2, current} -end -return {1} -`) - guardedAcquireLegacy = rueidis.NewLuaScript(` -if redis.call("GET", KEYS[2]) ~= ARGV[2] then - return {0} -end -local ttl = redis.call("PTTL", KEYS[2]) -if ttl <= 0 then - return {0} -end -local current = redis.call("GET", KEYS[1]) -if current then - redis.call("DEL", KEYS[2]) - return {2, current} -end -redis.call("SET", KEYS[1], ARGV[1], "PX", ttl) -return {1} -`) - guardedSet = rueidis.NewLuaScript(` -if redis.call("GET", KEYS[2]) ~= ARGV[2] then - return 0 -end -if redis.call("GET", KEYS[1]) == ARGV[1] then - redis.call("SET", KEYS[1], ARGV[3], "PX", ARGV[4]) - redis.call("DEL", KEYS[2]) - return 1 -end -redis.call("DEL", KEYS[2]) -return 0 -`) - revokeAttempt = rueidis.NewLuaScriptRetryable(` -if redis.call("GET", KEYS[2]) == ARGV[2] then - redis.call("DEL", KEYS[2]) -end -if redis.call("GET", KEYS[1]) == ARGV[1] then - return redis.call("DEL", KEYS[1]) -end -return 0 -`) -) - -const slotTagAlphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_" - -var ( - slotTags [16384][3]byte - slotTagsOnce sync.Once -) - -func attemptGuardKey(key, token string) string { - slotTagsOnce.Do(initSlotTags) - tag := slotTags[cmds.Slot(key)] - return attemptGuardPrefix + "{" + string(tag[:]) + "}:" + token -} - -func initSlotTags() { - // Lua requires the cache key and its attempt guard to share a cluster slot. - // Build one printable hash tag for every possible slot. - remaining := len(slotTags) - for i := 0; i < len(slotTagAlphabet) && remaining != 0; i++ { - for j := 0; j < len(slotTagAlphabet) && remaining != 0; j++ { - for k := 0; k < len(slotTagAlphabet) && remaining != 0; k++ { - tag := [3]byte{slotTagAlphabet[i], slotTagAlphabet[j], slotTagAlphabet[k]} - slot := cmds.Slot(string(tag[:])) - if slotTags[slot][0] == 0 { - slotTags[slot] = tag - remaining-- - } - } - } - } - if remaining != 0 { - panic("rueidisaside: unable to generate Redis slot tags") - } -} diff --git a/rueidisaside/attempt_test.go b/rueidisaside/attempt_test.go deleted file mode 100644 index 013030c2..00000000 --- a/rueidisaside/attempt_test.go +++ /dev/null @@ -1,632 +0,0 @@ -package rueidisaside - -import ( - "context" - "errors" - "math/rand" - "strconv" - "sync" - "sync/atomic" - "testing" - "time" - - "github.com/redis/rueidis" - "github.com/redis/rueidis/internal/cmds" -) - -type delayedAcquireClient struct { - rueidis.Client - key string - once sync.Once - blocked chan struct{} - pending rueidis.Completed - guard string -} - -func (c *delayedAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { - delayed := false - if isGuardedScript(cmd.Commands(), c.key, 2) { - c.once.Do(func() { - delayed = true - c.pending = cmd - c.guard = cmd.Commands()[4] - close(c.blocked) - }) - } - if delayed { - <-ctx.Done() - return rueidis.NewErrorResult(ctx.Err()) - } - return c.Client.Do(ctx, cmd) -} - -func (c *delayedAcquireClient) executePending() rueidis.RedisResult { - return c.Client.Do(context.Background(), c.pending) -} - -type blockingAcquireClient struct { - rueidis.Client - key string - started chan struct{} - release chan struct{} - calls atomic.Int64 -} - -type blockingCleanupClient struct { - rueidis.Client - key string - cancel context.CancelFunc - cleanupStarted chan struct{} - cleanupOnce sync.Once - calls atomic.Int64 -} - -func (c *blockingCleanupClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { - if isGuardedScript(cmd.Commands(), c.key, 2) { - if c.calls.Add(1) == 1 { - resp := c.Client.Do(ctx, cmd) - if resp.Error() == nil { - c.cancel() - return rueidis.NewErrorResult(ctx.Err()) - } - return resp - } - c.cleanupOnce.Do(func() { close(c.cleanupStarted) }) - <-ctx.Done() - return rueidis.NewErrorResult(ctx.Err()) - } - return c.Client.Do(ctx, cmd) -} - -func (c *blockingAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { - if isGuardedScript(cmd.Commands(), c.key, 2) { - if c.calls.Add(1) == 1 { - close(c.started) - select { - case <-c.release: - case <-ctx.Done(): - return rueidis.NewErrorResult(ctx.Err()) - } - } - } - return c.Client.Do(ctx, cmd) -} - -func isGuardedScript(commands []string, key string, argCount int) bool { - if len(commands) != 5+argCount || (commands[0] != "EVALSHA" && commands[0] != "EVAL") { - return false - } - return commands[2] == "2" && commands[3] == key -} - -func preloadAcquire(client rueidis.Client, useLuaLock bool, key string) error { - script := guardedAcquire - if useLuaLock { - script = guardedAcquireLegacy - } - guard := attemptGuardKey(key, "preload") - _, err := script.Exec( - context.Background(), - client, - []string{key, guard}, - []string{"preload-id", "preload"}, - ).ToArray() - return err -} - -func newDelayedAcquireCache(t *testing.T, key string, useLuaLock bool) (*Client, *delayedAcquireClient) { - t.Helper() - var wrapped *delayedAcquireClient - client, err := NewClient(ClientOption{ - ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, - ClientTTL: time.Second, - UseLuaLock: useLuaLock, - ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { - client, err := rueidis.NewClient(option) - if err != nil { - return nil, err - } - if err = preloadAcquire(client, useLuaLock, key); err != nil { - client.Close() - return nil, err - } - wrapped = &delayedAcquireClient{ - Client: client, - key: key, - blocked: make(chan struct{}), - } - return wrapped, nil - }, - }) - if err != nil { - t.Fatal(err) - } - return client.(*Client), wrapped -} - -func newBlockingAcquireCache(t *testing.T, key string, useLuaLock bool) (*Client, *blockingAcquireClient) { - t.Helper() - var wrapped *blockingAcquireClient - client, err := NewClient(ClientOption{ - ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, - ClientTTL: time.Second, - UseLuaLock: useLuaLock, - ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { - client, err := rueidis.NewClient(option) - if err != nil { - return nil, err - } - if err = preloadAcquire(client, useLuaLock, key); err != nil { - client.Close() - return nil, err - } - wrapped = &blockingAcquireClient{ - Client: client, - key: key, - started: make(chan struct{}), - release: make(chan struct{}), - } - return wrapped, nil - }, - }) - if err != nil { - t.Fatal(err) - } - return client.(*Client), wrapped -} - -func newBlockingCleanupCache( - t *testing.T, - key string, - useLuaLock bool, - cancel context.CancelFunc, -) (*Client, *blockingCleanupClient) { - t.Helper() - var wrapped *blockingCleanupClient - client, err := NewClient(ClientOption{ - ClientOption: rueidis.ClientOption{InitAddress: addr, PipelineMultiplex: -1, SelectDB: 5}, - ClientTTL: 100 * time.Millisecond, - UseLuaLock: useLuaLock, - ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { - client, err := rueidis.NewClient(option) - if err != nil { - return nil, err - } - if err = preloadAcquire(client, useLuaLock, key); err != nil { - client.Close() - return nil, err - } - wrapped = &blockingCleanupClient{ - Client: client, - key: key, - cancel: cancel, - cleanupStarted: make(chan struct{}), - } - return wrapped, nil - }, - }) - if err != nil { - t.Fatal(err) - } - return client.(*Client), wrapped -} - -func TestCanceledAcquireCannotArriveLate(t *testing.T) { - for _, useLuaLock := range []bool{false, true} { - name := "redis-7" - if useLuaLock { - name = "legacy" - } - t.Run(name, func(t *testing.T) { - key := "late-acquire-" + strconv.Itoa(rand.Int()) - client, wrapped := newDelayedAcquireCache(t, key, useLuaLock) - defer client.Close() - defer client.client.Do(context.Background(), client.client.B().Del().Key(key).Build()) - - ctx, cancel := context.WithCancel(context.Background()) - first := make(chan getResult, 1) - var firstLoaderCalled atomic.Bool - go func() { - val, err := client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { - firstLoaderCalled.Store(true) - return "old", nil - }) - first <- getResult{val: val, err: err} - }() - - select { - case <-wrapped.blocked: - case <-time.After(time.Second): - t.Fatal("acquire did not reach the delay point") - } - - second := make(chan getResult, 1) - go func() { - val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { - return "new", nil - }) - second <- getResult{val: val, err: err} - }() - - cancel() - select { - case result := <-first: - if !errors.Is(result.err, context.Canceled) { - t.Fatalf("expected context.Canceled, got %v", result.err) - } - case <-time.After(time.Second): - t.Fatal("canceled leader did not return") - } - if firstLoaderCalled.Load() { - t.Fatal("canceled acquire called the loader") - } - - select { - case result := <-second: - if result.err != nil || result.val != "new" { - t.Fatalf("follower returned %q, %v", result.val, result.err) - } - case <-time.After(2 * time.Second): - t.Fatal("follower did not retry after cleanup") - } - - resp, err := wrapped.executePending().ToArray() - if err != nil { - t.Fatal(err) - } - status, err := resp[0].AsInt64() - if err != nil || status != acquireRevoked { - t.Fatalf("late acquire returned status %d, %v", status, err) - } - - val, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() - if err != nil || val != "new" { - t.Fatalf("late acquire changed the cached value to %q, %v", val, err) - } - if err = client.client.Do(context.Background(), client.client.B().Get().Key(wrapped.guard).Build()).Error(); !rueidis.IsRedisNil(err) { - t.Fatalf("attempt guard was not revoked: %v", err) - } - }) - } -} - -func TestSingleflightSerializesAcquire(t *testing.T) { - for _, useLuaLock := range []bool{false, true} { - name := "redis-7" - if useLuaLock { - name = "legacy" - } - t.Run(name, func(t *testing.T) { - key := "singleflight-" + strconv.Itoa(rand.Int()) - client, wrapped := newBlockingAcquireCache(t, key, useLuaLock) - defer client.Close() - defer client.client.Do(context.Background(), client.client.B().Del().Key(key).Build()) - - first := make(chan getResult, 1) - go func() { - val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { - return "value", nil - }) - first <- getResult{val: val, err: err} - }() - select { - case <-wrapped.started: - case <-time.After(time.Second): - t.Fatal("first acquire did not start") - } - - ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) - defer cancel() - var secondLoaderCalled atomic.Bool - _, secondErr := client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { - secondLoaderCalled.Store(true) - return "second", nil - }) - close(wrapped.release) - - select { - case result := <-first: - if result.err != nil || result.val != "value" { - t.Fatalf("leader returned %q, %v", result.val, result.err) - } - case <-time.After(time.Second): - t.Fatal("leader did not finish") - } - if !errors.Is(secondErr, context.DeadlineExceeded) { - t.Fatalf("expected follower deadline, got %v", secondErr) - } - if secondLoaderCalled.Load() { - t.Fatal("follower called the loader") - } - if calls := wrapped.calls.Load(); calls != 1 { - t.Fatalf("expected one acquire, got %d", calls) - } - }) - } -} - -func TestCleanupTimeoutRotatesClientGeneration(t *testing.T) { - for _, useLuaLock := range []bool{false, true} { - name := "redis-7" - if useLuaLock { - name = "legacy" - } - t.Run(name, func(t *testing.T) { - key := "cleanup-timeout-" + strconv.Itoa(rand.Int()) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - client, wrapped := newBlockingCleanupCache(t, key, useLuaLock, cancel) - defer client.Close() - defer client.client.Do(context.Background(), client.client.B().Del().Key(key).Build()) - - started := time.Now() - loaderCalled := false - _, err := client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { - loaderCalled = true - return "value", nil - }) - if !errors.Is(err, context.Canceled) { - t.Fatalf("expected context.Canceled, got %v", err) - } - if elapsed := time.Since(started); elapsed > time.Second { - t.Fatalf("cleanup exceeded its bound: %s", elapsed) - } - if loaderCalled { - t.Fatal("loader ran after the canceled acquire") - } - select { - case <-wrapped.cleanupStarted: - default: - t.Fatal("cleanup was not attempted") - } - - client.mu.Lock() - id := client.id - client.mu.Unlock() - if id != "" { - t.Fatalf("client generation was not rotated: %q", id) - } - }) - } -} - -type getResult struct { - val string - err error -} - -func TestSingleflightKeepsDifferentKeysConcurrent(t *testing.T) { - client := makeClient(t, addr) - defer client.Close() - keys := []string{ - "parallel-a-" + strconv.Itoa(rand.Int()), - "parallel-b-" + strconv.Itoa(rand.Int()), - } - started := make(chan string, len(keys)) - release := make(chan struct{}) - var releaseOnce sync.Once - defer releaseOnce.Do(func() { close(release) }) - - results := make(chan getResult, len(keys)) - for _, key := range keys { - key := key - go func() { - val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { - started <- key - <-release - return key, nil - }) - results <- getResult{val: val, err: err} - }() - } - for range keys { - select { - case <-started: - case <-time.After(time.Second): - t.Fatal("different keys were serialized") - } - } - releaseOnce.Do(func() { close(release) }) - for range keys { - select { - case result := <-results: - if result.err != nil { - t.Fatal(result.err) - } - case <-time.After(time.Second): - t.Fatal("parallel cache fill did not finish") - } - } -} - -func TestLoaderPanicReleasesFlight(t *testing.T) { - client := makeClient(t, addr) - defer client.Close() - key := "panic-" + strconv.Itoa(rand.Int()) - - func() { - defer func() { - if recover() == nil { - t.Fatal("loader did not panic") - } - }() - _, _ = client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { - panic("loader panic") - }) - }() - - val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { - return "recovered", nil - }) - if err != nil || val != "recovered" { - t.Fatalf("flight was not released after panic: %q, %v", val, err) - } -} - -func TestDelayedSetCannotOverwriteNewAttempt(t *testing.T) { - for _, useLuaLock := range []bool{false, true} { - name := "redis-7" - if useLuaLock { - name = "legacy" - } - t.Run(name, func(t *testing.T) { - client := makeClient(t, addr).(*Client) - client.useLuaLock = useLuaLock - defer client.Close() - - key := "late-set-" + strconv.Itoa(rand.Int()) - id := PlaceholderPrefix + "late-set-owner" - first := attempt{token: "first-token"} - first.guard = attemptGuardKey(key, first.token) - second := attempt{token: "second-token"} - second.guard = attemptGuardKey(key, second.token) - defer client.client.Do( - context.Background(), - client.client.B().Del().Key(key, first.guard, second.guard).Build(), - ) - - if err := client.client.Do(context.Background(), client.client.B().Set().Key(first.guard).Value(first.token).Px(time.Second).Build()).Error(); err != nil { - t.Fatal(err) - } - if _, acquired, err := client.acquire(context.Background(), key, id, first); err != nil || !acquired { - t.Fatalf("first acquire failed: acquired=%v err=%v", acquired, err) - } - if err := revokeAttempt.Exec(context.Background(), client.client, []string{key, first.guard}, []string{id, first.token}).Error(); err != nil { - t.Fatal(err) - } - - if err := client.client.Do(context.Background(), client.client.B().Set().Key(second.guard).Value(second.token).Px(time.Second).Build()).Error(); err != nil { - t.Fatal(err) - } - if _, acquired, err := client.acquire(context.Background(), key, id, second); err != nil || !acquired { - t.Fatalf("second acquire failed: acquired=%v err=%v", acquired, err) - } - if status, err := guardedSet.Exec( - context.Background(), - client.client, - []string{key, second.guard}, - []string{id, second.token, "new", "1000"}, - ).AsInt64(); err != nil || status != 1 { - t.Fatalf("second set failed: status=%d err=%v", status, err) - } - if status, err := guardedSet.Exec( - context.Background(), - client.client, - []string{key, first.guard}, - []string{id, first.token, "old", "1000"}, - ).AsInt64(); err != nil || status != 0 { - t.Fatalf("delayed set returned status=%d err=%v", status, err) - } - - val, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() - if err != nil || val != "new" { - t.Fatalf("delayed set changed the value to %q, %v", val, err) - } - }) - } -} - -func TestCleanupRemovesAcceptedAcquire(t *testing.T) { - for _, useLuaLock := range []bool{false, true} { - name := "redis-7" - if useLuaLock { - name = "legacy" - } - t.Run(name, func(t *testing.T) { - client := makeClient(t, addr).(*Client) - client.useLuaLock = useLuaLock - defer client.Close() - - key := "accepted-acquire-" + strconv.Itoa(rand.Int()) - id := PlaceholderPrefix + "accepted-owner" - a := attempt{token: "accepted-token"} - a.guard = attemptGuardKey(key, a.token) - defer client.client.Do(context.Background(), client.client.B().Del().Key(key, a.guard).Build()) - - if err := client.client.Do(context.Background(), client.client.B().Set().Key(a.guard).Value(a.token).Px(time.Second).Build()).Error(); err != nil { - t.Fatal(err) - } - if _, acquired, err := client.acquire(context.Background(), key, id, a); err != nil || !acquired { - t.Fatalf("acquire failed: acquired=%v err=%v", acquired, err) - } - if err := revokeAttempt.Exec(context.Background(), client.client, []string{key, a.guard}, []string{id, a.token}).Error(); err != nil { - t.Fatal(err) - } - if err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).Error(); !rueidis.IsRedisNil(err) { - t.Fatalf("placeholder was not removed: %v", err) - } - if err := client.client.Do(context.Background(), client.client.B().Get().Key(a.guard).Build()).Error(); !rueidis.IsRedisNil(err) { - t.Fatalf("guard was not removed: %v", err) - } - }) - } -} - -func TestAcquireTTLDoesNotOutliveGuard(t *testing.T) { - for _, useLuaLock := range []bool{false, true} { - name := "redis-7" - if useLuaLock { - name = "legacy" - } - t.Run(name, func(t *testing.T) { - client := makeClient(t, addr).(*Client) - client.useLuaLock = useLuaLock - defer client.Close() - - key := "guard-ttl-" + strconv.Itoa(rand.Int()) - id := PlaceholderPrefix + "ttl-owner" - a := attempt{token: "ttl-token"} - a.guard = attemptGuardKey(key, a.token) - defer client.client.Do(context.Background(), client.client.B().Del().Key(key, a.guard).Build()) - - if err := client.client.Do(context.Background(), client.client.B().Set().Key(a.guard).Value(a.token).Px(time.Second).Build()).Error(); err != nil { - t.Fatal(err) - } - if _, acquired, err := client.acquire(context.Background(), key, id, a); err != nil || !acquired { - t.Fatalf("acquire failed: acquired=%v err=%v", acquired, err) - } - pttls := client.client.DoMulti( - context.Background(), - client.client.B().Pttl().Key(key).Build(), - client.client.B().Pttl().Key(a.guard).Build(), - ) - keyTTL, err := pttls[0].AsInt64() - if err != nil { - t.Fatal(err) - } - guardTTL, err := pttls[1].AsInt64() - if err != nil { - t.Fatal(err) - } - if keyTTL <= 0 || guardTTL <= 0 || keyTTL > guardTTL+20 { - t.Fatalf("unexpected TTLs: key=%dms guard=%dms", keyTTL, guardTTL) - } - }) - } -} - -func TestAttemptGuardKeyUsesSameSlot(t *testing.T) { - slotTagsOnce.Do(initSlotTags) - for slot, tag := range slotTags { - if got := cmds.Slot(string(tag[:])); got != uint16(slot) { - t.Fatalf("slot tag %q maps to %d, expected %d", tag, got, slot) - } - } - - keys := []string{ - "", - "plain", - "key{tag}suffix", - "{}", - "unclosed{tag", - "empty{}tag", - "\x00binary\xff", - "nested{{tag}}", - } - for _, key := range keys { - guard := attemptGuardKey(key, "token") - if got, want := cmds.Slot(guard), cmds.Slot(key); got != want { - t.Fatalf("guard for %q maps to slot %d, expected %d", key, got, want) - } - } -} diff --git a/rueidisaside/flight.go b/rueidisaside/flight.go index edca82f3..787fd7ce 100644 --- a/rueidisaside/flight.go +++ b/rueidisaside/flight.go @@ -1,5 +1,13 @@ package rueidisaside +import ( + "context" + "strconv" + "time" + + "github.com/redis/rueidis" +) + type flightKey struct { key string id string // client id generation @@ -42,3 +50,41 @@ func (c *Client) resetFlightsLocked() { } 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 +} diff --git a/rueidisaside/flight_test.go b/rueidisaside/flight_test.go index 11185b37..0f3c3249 100644 --- a/rueidisaside/flight_test.go +++ b/rueidisaside/flight_test.go @@ -1,6 +1,73 @@ package rueidisaside -import "testing" +import ( + "context" + "errors" + "math/rand" + "strconv" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/redis/rueidis" +) + +type cancelAfterAcquireClient struct { + rueidis.Client + key string + cancel context.CancelFunc + accepted atomic.Bool +} + +func (c *cancelAfterAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { + if isAcquireCommand(cmd.Commands(), c.key) { + resp := c.Client.Do(ctx, cmd) + if rueidis.IsRedisNil(resp.Error()) && c.accepted.CompareAndSwap(false, true) { + c.cancel() + return rueidis.NewErrorResult(ctx.Err()) + } + return resp + } + return c.Client.Do(ctx, cmd) +} + +type blockingAcquireClient struct { + rueidis.Client + key string + started chan struct{} + release chan struct{} + calls atomic.Int64 +} + +func (c *blockingAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { + if isAcquireCommand(cmd.Commands(), c.key) { + if c.calls.Add(1) == 1 { + close(c.started) + } + select { + case <-c.release: + case <-ctx.Done(): + return rueidis.NewErrorResult(ctx.Err()) + } + } + return c.Client.Do(ctx, cmd) +} + +func isAcquireCommand(commands []string, key string) bool { + if len(commands) == 7 && commands[0] == "SET" { + return commands[1] == key && commands[3] == "NX" && commands[4] == "GET" + } + if len(commands) == 6 && (commands[0] == "EVALSHA" || commands[0] == "EVAL") { + return commands[2] == "1" && commands[3] == key + } + return false +} + +type getResult struct { + val string + err error +} func TestBeginFlightRejectsStaleGeneration(t *testing.T) { c := &Client{ @@ -10,7 +77,229 @@ func TestBeginFlightRejectsStaleGeneration(t *testing.T) { if f, _ := c.beginFlight("key", "old-generation"); f != nil { t.Fatal("stale generation created a flight") } - if f, leader := c.beginFlight("key", "new-generation"); f == nil || !leader { + f, leader := c.beginFlight("key", "new-generation") + if f == nil || !leader { t.Fatal("current generation did not create a flight") } + follower, leader := c.beginFlight("key", "new-generation") + if follower != f || leader { + t.Fatal("same key did not join the active flight") + } + if other, leader := c.beginFlight("other-key", "new-generation"); other == nil || !leader { + t.Fatal("different key did not create an independent flight") + } + select { + case <-f.done: + t.Fatal("flight completed before the leader finished") + default: + } + c.finishFlight("key", "new-generation", f) + select { + case <-f.done: + default: + t.Fatal("flight did not wake its follower") + } +} + +func TestAcquireCancellationCleansPlaceholder(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "redis-7" + if useLuaLock { + name = "legacy" + } + t.Run(name, func(t *testing.T) { + key := "canceled-acquire-" + strconv.Itoa(rand.Int()) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + var wrapped *cancelAfterAcquireClient + client, err := NewClient(ClientOption{ + ClientOption: rueidis.ClientOption{InitAddress: addr, SelectDB: 5}, + ClientTTL: time.Second, + UseLuaLock: useLuaLock, + ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { + client, err := rueidis.NewClient(option) + if err != nil { + return nil, err + } + wrapped = &cancelAfterAcquireClient{ + Client: client, + key: key, + cancel: cancel, + } + return wrapped, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer client.Close() + defer client.Client().Do(context.Background(), client.Client().B().Del().Key(key).Build()) + + loaderCalled := false + _, err = client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { + loaderCalled = true + return "value", nil + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + if !wrapped.accepted.Load() { + t.Fatal("Redis did not accept the lock command") + } + if loaderCalled { + t.Fatal("loader ran after the canceled acquire") + } + + val, err := client.Client().Do(context.Background(), client.Client().B().Get().Key(key).Build()).ToString() + if !rueidis.IsRedisNil(err) { + t.Fatalf("expected cleanup to remove the placeholder, got %q, %v", val, err) + } + }) + } +} + +func TestSingleflightSerializesAcquire(t *testing.T) { + for _, useLuaLock := range []bool{false, true} { + name := "redis-7" + if useLuaLock { + name = "legacy" + } + t.Run(name, func(t *testing.T) { + key := "singleflight-" + strconv.Itoa(rand.Int()) + + var wrapped *blockingAcquireClient + client, err := NewClient(ClientOption{ + ClientOption: rueidis.ClientOption{InitAddress: addr, SelectDB: 5}, + ClientTTL: time.Second, + UseLuaLock: useLuaLock, + ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { + client, err := rueidis.NewClient(option) + if err != nil { + return nil, err + } + wrapped = &blockingAcquireClient{ + Client: client, + key: key, + started: make(chan struct{}), + release: make(chan struct{}), + } + return wrapped, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer client.Close() + defer client.Client().Do(context.Background(), client.Client().B().Del().Key(key).Build()) + + first := make(chan getResult, 1) + go func() { + val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + return "value", nil + }) + first <- getResult{val: val, err: err} + }() + select { + case <-wrapped.started: + case <-time.After(time.Second): + t.Fatal("first acquire did not start") + } + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + var secondLoaderCalled atomic.Bool + _, secondErr := client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { + secondLoaderCalled.Store(true) + return "second", nil + }) + if !errors.Is(secondErr, context.DeadlineExceeded) { + t.Fatalf("expected follower deadline, got %v", secondErr) + } + if secondLoaderCalled.Load() { + t.Fatal("follower called the loader") + } + if calls := wrapped.calls.Load(); calls != 1 { + t.Fatalf("expected one acquire while the flight was active, got %d", calls) + } + + close(wrapped.release) + select { + case result := <-first: + if result.err != nil || result.val != "value" { + t.Fatalf("leader returned %q, %v", result.val, result.err) + } + case <-time.After(time.Second): + t.Fatal("leader did not finish") + } + }) + } +} + +func TestSingleflightKeepsDifferentKeysConcurrent(t *testing.T) { + client := makeClient(t, addr) + defer client.Close() + keys := []string{ + "parallel-a-" + strconv.Itoa(rand.Int()), + "parallel-b-" + strconv.Itoa(rand.Int()), + } + started := make(chan string, len(keys)) + release := make(chan struct{}) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(release) }) + + results := make(chan getResult, len(keys)) + for _, key := range keys { + key := key + go func() { + val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + started <- key + <-release + return key, nil + }) + results <- getResult{val: val, err: err} + }() + } + for range keys { + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("different keys were serialized") + } + } + releaseOnce.Do(func() { close(release) }) + for range keys { + select { + case result := <-results: + if result.err != nil { + t.Fatal(result.err) + } + case <-time.After(time.Second): + t.Fatal("parallel cache fill did not finish") + } + } +} + +func TestLoaderPanicReleasesFlight(t *testing.T) { + client := makeClient(t, addr) + defer client.Close() + key := "panic-" + strconv.Itoa(rand.Int()) + + func() { + defer func() { + if recover() == nil { + t.Fatal("loader did not panic") + } + }() + _, _ = client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + panic("loader panic") + }) + }() + + val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + return "recovered", nil + }) + if err != nil || val != "recovered" { + t.Fatalf("flight was not released after panic: %q, %v", val, err) + } } From 81fe858d20782fa9eb0461d49c0e8993706e1aa9 Mon Sep 17 00:00:00 2001 From: Andrii Makarets <60238228+Makarechi@users.noreply.github.com> Date: Tue, 28 Jul 2026 09:52:39 +0200 Subject: [PATCH 5/6] fix(rueidisaside): keep flights across client id changes --- rueidisaside/aside.go | 5 +- rueidisaside/aside_test.go | 161 ++++++++++++++++++------------------ rueidisaside/flight.go | 30 ++----- rueidisaside/flight_test.go | 18 +++- 4 files changed, 104 insertions(+), 110 deletions(-) diff --git a/rueidisaside/aside.go b/rueidisaside/aside.go index 2ec51605..abaf9fad 100644 --- a/rueidisaside/aside.go +++ b/rueidisaside/aside.go @@ -34,7 +34,7 @@ func NewClient(option ClientOption) (cc CacheAsideClient, err error) { } ca := &Client{ waits: make(map[string]chan struct{}), - flights: make(map[flightKey]*flight), + flights: make(map[string]*flight), ttl: option.ClientTTL, useLuaLock: option.UseLuaLock, } @@ -56,7 +56,7 @@ type Client struct { client rueidis.Client ctx context.Context waits map[string]chan struct{} - flights map[flightKey]*flight + flights map[string]*flight cancel context.CancelFunc id string ttl time.Duration @@ -74,7 +74,6 @@ 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() diff --git a/rueidisaside/aside_test.go b/rueidisaside/aside_test.go index a617fcaf..a211d551 100644 --- a/rueidisaside/aside_test.go +++ b/rueidisaside/aside_test.go @@ -2,6 +2,7 @@ package rueidisaside import ( "context" + "errors" "math/rand" "strconv" "sync" @@ -346,102 +347,100 @@ func TestTimeoutLL(t *testing.T) { func TestDisconnect(t *testing.T) { client := makeClient(t, addr).(*Client) + testDisconnectWaitsForFlight(t, client) +} + +func TestDisconnectLL(t *testing.T) { + client := makeClientWithLuaLock(t, addr).(*Client) + testDisconnectWaitsForFlight(t, client) +} + +func testDisconnectWaitsForFlight(t *testing.T, client *Client) { + t.Helper() defer client.Close() key := strconv.Itoa(rand.Int()) - ch := make(chan string, 2) - val, err := client.Get(context.Background(), time.Second*5, key, func(ctx context.Context, key string) (val string, err error) { - id1, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() - if err != nil { - t.Error(err) - } - go func() { - val, err := client.Get(context.Background(), time.Second*5, key, func(ctx context.Context, key string) (val string, err error) { - id2, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() - if err != nil { - t.Error(err) - } - ch <- id2 - return "2", nil - }) - if val != "2" { - t.Error(err) - } - }() - client.onInvalidation(nil) // simulate disconnection - id2 := <-ch - if id1 == id2 { - t.Error("id not changed") - } - ch <- id1 - ch <- id2 - return "1", nil - }) - if val != "1" { - t.Fatal(err) + defer client.client.Do(context.Background(), client.client.B().Del().Key(key).Build()) + + leaderErr := errors.New("leader stopped after disconnect") + leaderReady := make(chan string, 1) + leaderRelease := make(chan struct{}) + var releaseOnce sync.Once + releaseLeader := func() { + releaseOnce.Do(func() { close(leaderRelease) }) } - val, err = client.Get(context.Background(), time.Millisecond*500, key, nil) - if val != "2" { - t.Error(err) + defer releaseLeader() + + leaderResult := make(chan getResult, 1) + go func() { + val, err := client.Get(context.Background(), time.Second*5, key, func(ctx context.Context, key string) (string, error) { + id1, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() + if err != nil { + return "", err + } + client.onInvalidation(nil) // simulate disconnection + leaderReady <- id1 + <-leaderRelease + return "", leaderErr + }) + leaderResult <- getResult{val: val, err: err} + }() + + var id1 string + select { + case id1 = <-leaderReady: + case <-time.After(time.Second): + t.Fatal("leader did not reach the loader") } - err = client.client.Do(context.Background(), client.client.B().Get().Key(<-ch).Build()).Error() // id1 - if !rueidis.IsRedisNil(err) { - t.Error(err) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + var followerLoaderCalled atomic.Bool + _, err := client.Get(ctx, time.Second*5, key, func(context.Context, string) (string, error) { + followerLoaderCalled.Store(true) + return "unexpected", nil + }) + cancel() + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected follower to wait for the active flight, got %v", err) } - err = client.client.Do(context.Background(), client.client.B().Get().Key(<-ch).Build()).Error() // id2 - if err != nil { - t.Error(err) + if followerLoaderCalled.Load() { + t.Fatal("follower started another loader while the old generation was active") } - time.Sleep(client.ttl) // wait old refresh goroutine exit -} -func TestDisconnectLL(t *testing.T) { - client := makeClientWithLuaLock(t, addr).(*Client) - defer client.Close() - key := strconv.Itoa(rand.Int()) - ch := make(chan string, 2) - val, err := client.Get(context.Background(), time.Second*5, key, func(ctx context.Context, key string) (val string, err error) { - id1, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() - if err != nil { - t.Error(err) + releaseLeader() + select { + case result := <-leaderResult: + if !errors.Is(result.err, leaderErr) { + t.Fatalf("expected leader error, got %q, %v", result.val, result.err) } - go func() { - val, err := client.Get(context.Background(), time.Second*5, key, func(ctx context.Context, key string) (val string, err error) { - id2, err := client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() - if err != nil { - t.Error(err) - } - ch <- id2 - return "2", nil - }) - if val != "2" { - t.Error(err) - } - }() - client.onInvalidation(nil) // simulate disconnection - id2 := <-ch - if id1 == id2 { - t.Error("id not changed") - } - ch <- id1 - ch <- id2 - return "1", nil + case <-time.After(time.Second): + t.Fatal("leader did not finish") + } + + var id2 string + val, err := client.Get(context.Background(), time.Second*5, key, func(context.Context, string) (string, error) { + var err error + id2, err = client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() + return "2", err }) - if val != "1" { - t.Fatal(err) + if err != nil || val != "2" { + t.Fatalf("new generation did not populate the cache: %q, %v", val, err) } - val, err = client.Get(context.Background(), time.Millisecond*500, key, nil) - if val != "2" { - t.Error(err) + if id1 == id2 { + t.Fatal("client id did not change") + } + + val, err = client.client.Do(context.Background(), client.client.B().Get().Key(key).Build()).ToString() + if err != nil || val != "2" { + t.Fatalf("unexpected cache value: %q, %v", val, err) } - err = client.client.Do(context.Background(), client.client.B().Get().Key(<-ch).Build()).Error() // id1 + err = client.client.Do(context.Background(), client.client.B().Get().Key(id1).Build()).Error() if !rueidis.IsRedisNil(err) { - t.Error(err) + t.Fatalf("old client marker still exists: %v", err) } - err = client.client.Do(context.Background(), client.client.B().Get().Key(<-ch).Build()).Error() // id2 + err = client.client.Do(context.Background(), client.client.B().Get().Key(id2).Build()).Error() if err != nil { - t.Error(err) + t.Fatalf("new client marker is missing: %v", err) } - time.Sleep(client.ttl) // wait old refresh goroutine exit } func TestMultipleClient(t *testing.T) { diff --git a/rueidisaside/flight.go b/rueidisaside/flight.go index 787fd7ce..e6776a60 100644 --- a/rueidisaside/flight.go +++ b/rueidisaside/flight.go @@ -8,11 +8,6 @@ import ( "github.com/redis/rueidis" ) -type flightKey struct { - key string - id string // client id generation -} - type flight struct { done chan struct{} } @@ -21,34 +16,25 @@ func (c *Client) beginFlight(key, id string) (f *flight, leader bool) { c.mu.Lock() defer c.mu.Unlock() + if f = c.flights[key]; f != nil { + return f, false + } 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 + c.flights[key] = f return f, true } -func (c *Client) finishFlight(key, id string, f *flight) { +func (c *Client) finishFlight(key 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 { + if c.flights[key] == f { + delete(c.flights, key) close(f.done) } - c.flights = make(map[flightKey]*flight) } func (c *Client) populate( @@ -59,7 +45,7 @@ func (c *Client) populate( f *flight, ) (val string, err error) { cleanup := true - defer c.finishFlight(key, id, f) + defer c.finishFlight(key, f) defer func() { if cleanup { delkey.Exec(context.Background(), c.client, []string{key}, []string{id}) diff --git a/rueidisaside/flight_test.go b/rueidisaside/flight_test.go index 0f3c3249..78c14118 100644 --- a/rueidisaside/flight_test.go +++ b/rueidisaside/flight_test.go @@ -69,22 +69,32 @@ type getResult struct { err error } -func TestBeginFlightRejectsStaleGeneration(t *testing.T) { +func TestBeginFlightWaitsAcrossGenerations(t *testing.T) { c := &Client{ id: "new-generation", - flights: make(map[flightKey]*flight), + flights: make(map[string]*flight), } - if f, _ := c.beginFlight("key", "old-generation"); f != nil { + if f, leader := c.beginFlight("key", "old-generation"); f != nil || leader { t.Fatal("stale generation created a flight") } f, leader := c.beginFlight("key", "new-generation") if f == nil || !leader { t.Fatal("current generation did not create a flight") } + staleFollower, leader := c.beginFlight("key", "old-generation") + if staleFollower != f || leader { + t.Fatal("stale generation did not join the current flight") + } follower, leader := c.beginFlight("key", "new-generation") if follower != f || leader { t.Fatal("same key did not join the active flight") } + if staleOther, leader := c.beginFlight("other-key", "old-generation"); staleOther != nil || leader { + t.Fatal("stale generation joined a flight for another key") + } + if len(c.flights) != 1 { + t.Fatalf("stale generation changed the flights map: %d entries", len(c.flights)) + } if other, leader := c.beginFlight("other-key", "new-generation"); other == nil || !leader { t.Fatal("different key did not create an independent flight") } @@ -93,7 +103,7 @@ func TestBeginFlightRejectsStaleGeneration(t *testing.T) { t.Fatal("flight completed before the leader finished") default: } - c.finishFlight("key", "new-generation", f) + c.finishFlight("key", f) select { case <-f.done: default: From 1854a98fba38610d38f559f96058990eb7469740 Mon Sep 17 00:00:00 2001 From: Andrii Makarets <60238228+Makarechi@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:10:09 +0200 Subject: [PATCH 6/6] fix(rueidisaside): start flights before keepalive --- rueidisaside/aside.go | 24 +++---- rueidisaside/flight.go | 15 +++-- rueidisaside/flight_test.go | 123 ++++++++++++++++++++++++++++++------ 3 files changed, 123 insertions(+), 39 deletions(-) diff --git a/rueidisaside/aside.go b/rueidisaside/aside.go index abaf9fad..76817a1a 100644 --- a/rueidisaside/aside.go +++ b/rueidisaside/aside.go @@ -159,24 +159,18 @@ retry: val, err := resp.ToString() 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 { - f, leader := c.beginFlight(key, id) - if f == nil { // the client id changed while preparing the flight + f, leader := c.beginFlight(key) + if !leader { + select { + case <-f.done: goto retry + case <-ctx.Done(): + return "", ctx.Err() + case <-c.ctx.Done(): + return "", c.ctx.Err() } - 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) } + val, err = c.populate(ctx, ttl, key, fn, f) } if err != nil { diff --git a/rueidisaside/flight.go b/rueidisaside/flight.go index e6776a60..49726b66 100644 --- a/rueidisaside/flight.go +++ b/rueidisaside/flight.go @@ -12,16 +12,13 @@ type flight struct { done chan struct{} } -func (c *Client) beginFlight(key, id string) (f *flight, leader bool) { +func (c *Client) beginFlight(key string) (f *flight, leader bool) { c.mu.Lock() defer c.mu.Unlock() if f = c.flights[key]; f != nil { return f, false } - if c.id != id { - return nil, false - } f = &flight{done: make(chan struct{})} c.flights[key] = f return f, true @@ -40,12 +37,18 @@ func (c *Client) finishFlight(key string, f *flight) { func (c *Client) populate( ctx context.Context, ttl time.Duration, - key, id string, + key string, fn func(ctx context.Context, key string) (val string, err error), f *flight, ) (val string, err error) { - cleanup := true defer c.finishFlight(key, f) + + id, err := c.keepalive() + if err != nil { + return "", err + } + + cleanup := true defer func() { if cleanup { delkey.Exec(context.Background(), c.client, []string{key}, []string{id}) diff --git a/rueidisaside/flight_test.go b/rueidisaside/flight_test.go index 78c14118..07750fbe 100644 --- a/rueidisaside/flight_test.go +++ b/rueidisaside/flight_test.go @@ -5,6 +5,7 @@ import ( "errors" "math/rand" "strconv" + "strings" "sync" "sync/atomic" "testing" @@ -54,6 +55,28 @@ func (c *blockingAcquireClient) Do(ctx context.Context, cmd rueidis.Completed) r return c.Client.Do(ctx, cmd) } +type blockingKeepaliveClient struct { + rueidis.Client + started chan struct{} + release chan struct{} + calls atomic.Int64 +} + +func (c *blockingKeepaliveClient) Do(ctx context.Context, cmd rueidis.Completed) rueidis.RedisResult { + commands := cmd.Commands() + if len(commands) >= 2 && commands[0] == "SET" && strings.HasPrefix(commands[1], PlaceholderPrefix) { + if c.calls.Add(1) == 1 { + close(c.started) + select { + case <-c.release: + case <-ctx.Done(): + return rueidis.NewErrorResult(ctx.Err()) + } + } + } + return c.Client.Do(ctx, cmd) +} + func isAcquireCommand(commands []string, key string) bool { if len(commands) == 7 && commands[0] == "SET" { return commands[1] == key && commands[3] == "NX" && commands[4] == "GET" @@ -69,33 +92,23 @@ type getResult struct { err error } -func TestBeginFlightWaitsAcrossGenerations(t *testing.T) { +func TestBeginFlightIsPerKey(t *testing.T) { c := &Client{ - id: "new-generation", flights: make(map[string]*flight), } - if f, leader := c.beginFlight("key", "old-generation"); f != nil || leader { - t.Fatal("stale generation created a flight") - } - f, leader := c.beginFlight("key", "new-generation") + f, leader := c.beginFlight("key") if f == nil || !leader { - t.Fatal("current generation did not create a flight") - } - staleFollower, leader := c.beginFlight("key", "old-generation") - if staleFollower != f || leader { - t.Fatal("stale generation did not join the current flight") + t.Fatal("first caller did not create a flight") } - follower, leader := c.beginFlight("key", "new-generation") + c.id = "new-generation" + follower, leader := c.beginFlight("key") if follower != f || leader { - t.Fatal("same key did not join the active flight") - } - if staleOther, leader := c.beginFlight("other-key", "old-generation"); staleOther != nil || leader { - t.Fatal("stale generation joined a flight for another key") + t.Fatal("same key did not join the active flight after the client id changed") } if len(c.flights) != 1 { - t.Fatalf("stale generation changed the flights map: %d entries", len(c.flights)) + t.Fatalf("follower changed the flights map: %d entries", len(c.flights)) } - if other, leader := c.beginFlight("other-key", "new-generation"); other == nil || !leader { + if other, leader := c.beginFlight("other-key"); other == nil || !leader { t.Fatal("different key did not create an independent flight") } select { @@ -169,6 +182,80 @@ func TestAcquireCancellationCleansPlaceholder(t *testing.T) { } } +func TestSingleflightStartsBeforeKeepalive(t *testing.T) { + key := "singleflight-keepalive-" + strconv.Itoa(rand.Int()) + + var wrapped *blockingKeepaliveClient + client, err := NewClient(ClientOption{ + ClientOption: rueidis.ClientOption{InitAddress: addr, SelectDB: 5}, + ClientTTL: 5 * time.Second, + ClientBuilder: func(option rueidis.ClientOption) (rueidis.Client, error) { + client, err := rueidis.NewClient(option) + if err != nil { + return nil, err + } + wrapped = &blockingKeepaliveClient{ + Client: client, + started: make(chan struct{}), + release: make(chan struct{}), + } + return wrapped, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer client.Close() + defer client.Client().Do(context.Background(), client.Client().B().Del().Key(key).Build()) + defer func() { + select { + case <-wrapped.release: + default: + close(wrapped.release) + } + }() + + first := make(chan getResult, 1) + go func() { + val, err := client.Get(context.Background(), time.Second, key, func(context.Context, string) (string, error) { + return "value", nil + }) + first <- getResult{val: val, err: err} + }() + select { + case <-wrapped.started: + case <-time.After(time.Second): + t.Fatal("first keepalive did not start") + } + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + var secondLoaderCalled atomic.Bool + _, secondErr := client.Get(ctx, time.Second, key, func(context.Context, string) (string, error) { + secondLoaderCalled.Store(true) + return "second", nil + }) + if !errors.Is(secondErr, context.DeadlineExceeded) { + t.Fatalf("expected follower deadline, got %v", secondErr) + } + if secondLoaderCalled.Load() { + t.Fatal("follower called the loader") + } + if calls := wrapped.calls.Load(); calls != 1 { + t.Fatalf("expected one keepalive while the flight was active, got %d", calls) + } + + close(wrapped.release) + select { + case result := <-first: + if result.err != nil || result.val != "value" { + t.Fatalf("leader returned %q, %v", result.val, result.err) + } + case <-time.After(time.Second): + t.Fatal("leader did not finish") + } +} + func TestSingleflightSerializesAcquire(t *testing.T) { for _, useLuaLock := range []bool{false, true} { name := "redis-7"