Support EVALSHA

batch.go

@@ -19,22 +19,7 @@ var (
 )
 
 func (c *Client) luaGetBatch(ctx context.Context, keys []string, owner string) ([]interface{}, error) {
-	res, err := callLua(ctx, c.rdb, `-- luaGetBatch
-    local rets = {}
-    for i, key in ipairs(KEYS)
-    do
-        local v = redis.call('HGET', key, 'value')
-        local lu = redis.call('HGET', key, 'lockUntil')
-        if lu ~= false and tonumber(lu) < tonumber(ARGV[1]) or lu == false and v == false then
-            redis.call('HSET', key, 'lockUntil', ARGV[2])
-            redis.call('HSET', key, 'lockOwner', ARGV[3])
-            table.insert(rets, { v, 'LOCKED' })
-        else
-            table.insert(rets, {v, lu})
-        end
-    end
-    return rets
-	`, keys, []interface{}{now(), now() + int64(c.Options.LockExpire/time.Second), owner})
+	res, err := callLua(ctx, c.rdb, getBatchScript, keys, []interface{}{now(), now() + int64(c.Options.LockExpire/time.Second), owner})
 	debugf("luaGetBatch return: %v, %v", res, err)
 	if err != nil {
 		return nil, err
@@ -51,20 +36,7 @@ func (c *Client) luaSetBatch(ctx context.Context, keys []string, values []string
 	for _, ex := range expires {
 		vals = append(vals, ex)
 	}
-	_, err := callLua(ctx, c.rdb, `-- luaSetBatch
-    local n = #KEYS
-    for i, key in ipairs(KEYS)
-    do
-        local o = redis.call('HGET', key, 'lockOwner')
-        if o ~= ARGV[1] then
-                return
-        end
-        redis.call('HSET', key, 'value', ARGV[i+1])
-        redis.call('HDEL', key, 'lockUntil')
-        redis.call('HDEL', key, 'lockOwner')
-        redis.call('EXPIRE', key, ARGV[i+1+n])
-    end
-	`, keys, vals)
+	_, err := callLua(ctx, c.rdb, setBatchScript, keys, vals)
 	return err
 }
 
@@ -194,7 +166,13 @@ func (c *Client) weakFetchBatch(ctx context.Context, keys []string, expire time.
 				r, err := c.luaGet(ctx, keys[i], owner)
 				for err == nil && r[0] == nil && r[1].(string) != locked {
 					debugf("batch weak: empty result for %s locked by other, so sleep %s", keys[i], c.Options.LockSleep.String())
-					time.Sleep(c.Options.LockSleep)
+					select {
+					case <-ctx.Done():
+						ch <- pair{idx: i, err: ctx.Err()}
+						return
+					case <-time.After(c.Options.LockSleep):
+						// equal to time.Sleep(c.Options.LockSleep) but can be canceled
+					}
 					r, err = c.luaGet(ctx, keys[i], owner)
 				}
 				if err != nil {
@@ -304,7 +282,13 @@ func (c *Client) strongFetchBatch(ctx context.Context, keys []string, expire tim
 				r, err := c.luaGet(ctx, keys[i], owner)
 				for err == nil && r[1] != nil && r[1] != locked { // locked by other
 					debugf("batch: locked by other, so sleep %s", c.Options.LockSleep)
-					time.Sleep(c.Options.LockSleep)
+					select {
+					case <-ctx.Done():
+						ch <- pair{idx: i, err: ctx.Err()}
+						return
+					case <-time.After(c.Options.LockSleep):
+						// equal to time.Sleep(c.Options.LockSleep) but can be canceled
+					}
 					r, err = c.luaGet(ctx, keys[i], owner)
 				}
 				if err != nil {
@@ -380,14 +364,8 @@ func (c *Client) TagAsDeletedBatch2(ctx context.Context, keys []string) error {
 		return nil
 	}
 	debugf("batch deleting: keys=%v", keys)
-	luaFn := func(con redisConn) error {
-		_, err := callLua(ctx, con, ` -- luaDeleteBatch
-		for i, key in ipairs(KEYS) do
-			redis.call('HSET', key, 'lockUntil', 0)
-			redis.call('HDEL', key, 'lockOwner')
-			redis.call('EXPIRE', key, ARGV[1])
-		end
-		`, keys, []interface{}{int64(c.Options.Delay / time.Second)})
+	luaFn := func(con redis.Scripter) error {
+		_, err := callLua(ctx, con, deleteBatchScript, keys, []interface{}{int64(c.Options.Delay / time.Second)})
 		return err
 	}
 	if c.Options.WaitReplicas > 0 {

batch_cover_test.go

@@ -123,3 +123,70 @@ func TestTagAsDeletedBatchWait(t *testing.T) {
 		assert.Error(t, err, fmt.Errorf("wait replicas 1 failed. result replicas: 0"))
 	}
 }
+
+func TestWeakFetchBatchCanceled(t *testing.T) {
+	clearCache()
+	rc := NewClient(rdb, NewDefaultOptions())
+	n := int(rand.Int31n(20) + 10)
+	idxs := genIdxs(n)
+	keys, values1, values2 := genKeys(idxs), genValues(n, "value_"), genValues(n, "eulav_")
+	values3 := genValues(n, "vvvv_")
+	go func() {
+		dc2 := NewClient(rdb, NewDefaultOptions())
+		v, err := dc2.FetchBatch(keys, 60*time.Second, genBatchDataFunc(values1, 450))
+		assert.Nil(t, err)
+		assert.Equal(t, values1, v)
+	}()
+	time.Sleep(20 * time.Millisecond)
+
+	began := time.Now()
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+	_, err := rc.FetchBatch2(ctx, keys, 60*time.Second, genBatchDataFunc(values2, 200))
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+	assertEqualDuration(t, time.Duration(200)*time.Millisecond, time.Since(began))
+
+	ctx, cancel = context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(200 * time.Millisecond)
+		cancel()
+	}()
+	began = time.Now()
+	_, err = rc.FetchBatch2(ctx, keys, 60*time.Second, genBatchDataFunc(values3, 200))
+	assert.ErrorIs(t, err, context.Canceled)
+	assertEqualDuration(t, time.Duration(200)*time.Millisecond, time.Since(began))
+}
+
+func TestStrongFetchBatchCanceled(t *testing.T) {
+	clearCache()
+	rc := NewClient(rdb, NewDefaultOptions())
+	rc.Options.StrongConsistency = true
+	n := int(rand.Int31n(20) + 10)
+	idxs := genIdxs(n)
+	keys, values1, values2 := genKeys(idxs), genValues(n, "value_"), genValues(n, "eulav_")
+	values3 := genValues(n, "vvvv_")
+	go func() {
+		dc2 := NewClient(rdb, NewDefaultOptions())
+		v, err := dc2.FetchBatch(keys, 60*time.Second, genBatchDataFunc(values1, 450))
+		assert.Nil(t, err)
+		assert.Equal(t, values1, v)
+	}()
+	time.Sleep(20 * time.Millisecond)
+
+	began := time.Now()
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+	_, err := rc.FetchBatch2(ctx, keys, 60*time.Second, genBatchDataFunc(values2, 200))
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+	assertEqualDuration(t, time.Duration(200)*time.Millisecond, time.Since(began))
+
+	ctx, cancel = context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(200 * time.Millisecond)
+		cancel()
+	}()
+	began = time.Now()
+	_, err = rc.FetchBatch2(ctx, keys, 60*time.Second, genBatchDataFunc(values3, 200))
+	assert.ErrorIs(t, err, context.Canceled)
+	assertEqualDuration(t, time.Duration(200)*time.Millisecond, time.Since(began))
+}

client.go

@@ -93,12 +93,8 @@ func (c *Client) TagAsDeleted2(ctx context.Context, key string) error {
 		return nil
 	}
 	debugf("deleting: key=%s", key)
-	luaFn := func(con redisConn) error {
-		_, err := callLua(ctx, con, ` --  delete
-		redis.call('HSET', KEYS[1], 'lockUntil', 0)
-		redis.call('HDEL', KEYS[1], 'lockOwner')
-		redis.call('EXPIRE', KEYS[1], ARGV[1])
-			`, []string{key}, []interface{}{int64(c.Options.Delay / time.Second)})
+	luaFn := func(con redis.Scripter) error {
+		_, err := callLua(ctx, con, deleteScript, []string{key}, []interface{}{int64(c.Options.Delay / time.Second)})
 		return err
 	}
 	if c.Options.WaitReplicas > 0 {
@@ -141,16 +137,7 @@ func (c *Client) Fetch2(ctx context.Context, key string, expire time.Duration, f
 }
 
 func (c *Client) luaGet(ctx context.Context, key string, owner string) ([]interface{}, error) {
-	res, err := callLua(ctx, c.rdb, ` -- luaGet
-	local v = redis.call('HGET', KEYS[1], 'value')
-	local lu = redis.call('HGET', KEYS[1], 'lockUntil')
-	if lu ~= false and tonumber(lu) < tonumber(ARGV[1]) or lu == false and v == false then
-		redis.call('HSET', KEYS[1], 'lockUntil', ARGV[2])
-		redis.call('HSET', KEYS[1], 'lockOwner', ARGV[3])
-		return { v, 'LOCKED' }
-	end
-	return {v, lu}
-	`, []string{key}, []interface{}{now(), now() + int64(c.Options.LockExpire/time.Second), owner})
+	res, err := callLua(ctx, c.rdb, getScript, []string{key}, []interface{}{now(), now() + int64(c.Options.LockExpire/time.Second), owner})
 	debugf("luaGet return: %v, %v", res, err)
 	if err != nil {
 		return nil, err
@@ -159,16 +146,7 @@ func (c *Client) luaGet(ctx context.Context, key string, owner string) ([]interf
 }
 
 func (c *Client) luaSet(ctx context.Context, key string, value string, expire int, owner string) error {
-	_, err := callLua(ctx, c.rdb, `-- luaSet
-	local o = redis.call('HGET', KEYS[1], 'lockOwner')
-	if o ~= ARGV[2] then
-			return
-	end
-	redis.call('HSET', KEYS[1], 'value', ARGV[1])
-	redis.call('HDEL', KEYS[1], 'lockUntil')
-	redis.call('HDEL', KEYS[1], 'lockOwner')
-	redis.call('EXPIRE', KEYS[1], ARGV[3])
-	`, []string{key}, []interface{}{value, owner, expire})
+	_, err := callLua(ctx, c.rdb, setScript, []string{key}, []interface{}{value, owner, expire})
 	return err
 }
 
@@ -195,7 +173,12 @@ func (c *Client) weakFetch(ctx context.Context, key string, expire time.Duration
 	r, err := c.luaGet(ctx, key, owner)
 	for err == nil && r[0] == nil && r[1].(string) != locked {
 		debugf("empty result for %s locked by other, so sleep %s", key, c.Options.LockSleep.String())
-		time.Sleep(c.Options.LockSleep)
+		select {
+		case <-ctx.Done():
+			return "", ctx.Err()
+		case <-time.After(c.Options.LockSleep):
+			// equal to time.Sleep(c.Options.LockSleep) but can be canceled
+		}
 		r, err = c.luaGet(ctx, key, owner)
 	}
 	if err != nil {
@@ -219,7 +202,12 @@ func (c *Client) strongFetch(ctx context.Context, key string, expire time.Durati
 	r, err := c.luaGet(ctx, key, owner)
 	for err == nil && r[1] != nil && r[1] != locked { // locked by other
 		debugf("locked by other, so sleep %s", c.Options.LockSleep)
-		time.Sleep(c.Options.LockSleep)
+		select {
+		case <-ctx.Done():
+			return "", ctx.Err()
+		case <-time.After(c.Options.LockSleep):
+			// equal to time.Sleep(c.Options.LockSleep) but can be canceled
+		}
 		r, err = c.luaGet(ctx, key, owner)
 	}
 	if err != nil {
@@ -248,16 +236,7 @@ func (c *Client) RawSet(ctx context.Context, key string, value string, expire ti
 // LockForUpdate locks the key, used in very strict strong consistency mode
 func (c *Client) LockForUpdate(ctx context.Context, key string, owner string) error {
 	lockUntil := math.Pow10(10)
-	res, err := callLua(ctx, c.rdb, ` -- luaLock
-	local lu = redis.call('HGET', KEYS[1], 'lockUntil')
-	local lo = redis.call('HGET', KEYS[1], 'lockOwner')
-	if lu == false or tonumber(lu) < tonumber(ARGV[2]) or lo == ARGV[1] then
-		redis.call('HSET', KEYS[1], 'lockUntil', ARGV[2])
-		redis.call('HSET', KEYS[1], 'lockOwner', ARGV[1])
-		return 'LOCKED'
-	end
-	return lo
-	`, []string{key}, []interface{}{owner, lockUntil})
+	res, err := callLua(ctx, c.rdb, lockScript, []string{key}, []interface{}{owner, lockUntil})
 	if err == nil && res != "LOCKED" {
 		return fmt.Errorf("%s has been locked by %s", key, res)
 	}
@@ -266,13 +245,6 @@ func (c *Client) LockForUpdate(ctx context.Context, key string, owner string) er
 
 // UnlockForUpdate unlocks the key, used in very strict strong consistency mode
 func (c *Client) UnlockForUpdate(ctx context.Context, key string, owner string) error {
-	_, err := callLua(ctx, c.rdb, ` -- luaUnlock
-	local lo = redis.call('HGET', KEYS[1], 'lockOwner')
-	if lo == ARGV[1] then
-		redis.call('HSET', KEYS[1], 'lockUntil', 0)
-		redis.call('HDEL', KEYS[1], 'lockOwner')
-		redis.call('EXPIRE', KEYS[1], ARGV[2])
-	end
-	`, []string{key}, []interface{}{owner, c.Options.LockExpire / time.Second})
+	_, err := callLua(ctx, c.rdb, unlockScript, []string{key}, []interface{}{owner, c.Options.LockExpire / time.Second})
 	return err
 }

client_test.go

@@ -65,6 +65,7 @@ func TestWeakFetch(t *testing.T) {
 
 	clearCache()
 	began := time.Now()
+	// nolint: goconst
 	expected := "value1"
 	go func() {
 		dc2 := NewClient(rdb, NewDefaultOptions())
@@ -147,6 +148,48 @@ func TestStrongErrorFetch(t *testing.T) {
 	assert.True(t, time.Since(began) < time.Duration(150)*time.Millisecond)
 }
 
+// nolint: unparam
+func assertEqualDuration(t *testing.T, expected, actual time.Duration) {
+	t.Helper()
+	delta := expected - actual
+	if delta < 0 {
+		delta = -delta
+	}
+	t.Logf("expected=%s, actual=%s, delta=%s", expected, actual, delta)
+	assert.Less(t, delta, time.Duration(2)*time.Millisecond)
+}
+
+func TestStrongFetchCanceled(t *testing.T) {
+	clearCache()
+	rc := NewClient(rdb, NewDefaultOptions())
+	rc.Options.StrongConsistency = true
+	expected := "value1"
+	go func() {
+		dc2 := NewClient(rdb, NewDefaultOptions())
+		v, err := dc2.Fetch(rdbKey, 60*time.Second, genDataFunc(expected, 450))
+		assert.Nil(t, err)
+		assert.Equal(t, expected, v)
+	}()
+	time.Sleep(20 * time.Millisecond)
+
+	began := time.Now()
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+	_, err := rc.Fetch2(ctx, rdbKey, 60*time.Second, genDataFunc(expected, 200))
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+	assertEqualDuration(t, time.Duration(200)*time.Millisecond, time.Since(began))
+
+	ctx, cancel = context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(200 * time.Millisecond)
+		cancel()
+	}()
+	began = time.Now()
+	_, err = rc.Fetch2(ctx, rdbKey, 60*time.Second, genDataFunc(expected, 200))
+	assert.ErrorIs(t, err, context.Canceled)
+	assertEqualDuration(t, time.Duration(200)*time.Millisecond, time.Since(began))
+}
+
 func TestWeakErrorFetch(t *testing.T) {
 	rc := NewClient(rdb, NewDefaultOptions())
 
@@ -165,6 +208,37 @@ func TestWeakErrorFetch(t *testing.T) {
 	assert.True(t, time.Since(began) < time.Duration(150)*time.Millisecond)
 }
 
+func TestWeakFetchCanceled(t *testing.T) {
+	rc := NewClient(rdb, NewDefaultOptions())
+
+	clearCache()
+	expected := "value1"
+	go func() {
+		dc2 := NewClient(rdb, NewDefaultOptions())
+		v, err := dc2.Fetch(rdbKey, 60*time.Second, genDataFunc(expected, 450))
+		assert.Nil(t, err)
+		assert.Equal(t, expected, v)
+	}()
+	time.Sleep(20 * time.Millisecond)
+
+	began := time.Now()
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+	_, err := rc.Fetch2(ctx, rdbKey, 60*time.Second, genDataFunc(expected, 200))
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+	assertEqualDuration(t, time.Duration(200)*time.Millisecond, time.Since(began))
+
+	ctx, cancel = context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(200 * time.Millisecond)
+		cancel()
+	}()
+	began = time.Now()
+	_, err = rc.Fetch2(ctx, rdbKey, 60*time.Second, genDataFunc(expected, 200))
+	assert.ErrorIs(t, err, context.Canceled)
+	assertEqualDuration(t, time.Duration(200)*time.Millisecond, time.Since(began))
+}
+
 func TestRawGet(t *testing.T) {
 	rc := NewClient(rdb, NewDefaultOptions())
 	_, err := rc.RawGet(ctx, "not-exists")

go.mod

@@ -5,7 +5,7 @@ go 1.18
 require (
 	github.com/lithammer/shortuuid v3.0.0+incompatible
 	github.com/redis/go-redis/v9 v9.0.3
-	github.com/stretchr/testify v1.8.1
+	github.com/stretchr/testify v1.8.4
 	golang.org/x/sync v0.1.0
 )
 

go.sum

@@ -22,6 +22,8 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
 golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=