feat: support cancellable sleep

batch.go

@@ -192,10 +192,25 @@ func (c *Client) weakFetchBatch(ctx context.Context, keys []string, expire time.
 			go func(i int) {
 				defer wg.Done()
 				r, err := c.luaGet(ctx, keys[i], owner)
+				ticker := time.NewTimer(c.Options.LockSleep)
+				defer ticker.Stop()
 				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 <-ticker.C:
+						// equal to time.Sleep(c.Options.LockSleep) but can be canceled
+					}
 					r, err = c.luaGet(ctx, keys[i], owner)
+					// Reset ticker after luaGet
+					// If we reset ticker before luaGet, since luaGet takes a period of time,
+					// the actual sleep time will be shorter than expected
+					if !ticker.Stop() && len(ticker.C) > 0 {
+						<-ticker.C
+					}
+					ticker.Reset(c.Options.LockSleep)
 				}
 				if err != nil {
 					ch <- pair{idx: i, data: "", err: err}
@@ -302,10 +317,24 @@ func (c *Client) strongFetchBatch(ctx context.Context, keys []string, expire tim
 			go func(i int) {
 				defer wg.Done()
 				r, err := c.luaGet(ctx, keys[i], owner)
+				ticker := time.NewTimer(c.Options.LockSleep)
+				defer ticker.Stop()
 				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()}
+					case <-ticker.C:
+						// equal to time.Sleep(c.Options.LockSleep) but can be canceled
+					}
 					r, err = c.luaGet(ctx, keys[i], owner)
+					// Reset ticker after luaGet
+					// If we reset ticker before luaGet, since luaGet takes a period of time,
+					// the actual sleep time will be shorter than expected
+					if !ticker.Stop() && len(ticker.C) > 0 {
+						<-ticker.C
+					}
+					ticker.Reset(c.Options.LockSleep)
 				}
 				if err != nil {
 					ch <- pair{idx: i, data: "", err: err}

client.go

@@ -3,6 +3,7 @@ package rockscache
 import (
 	"context"
 	"fmt"
+	"log"
 	"math"
 	"math/rand"
 	"time"
@@ -193,10 +194,25 @@ func (c *Client) weakFetch(ctx context.Context, key string, expire time.Duration
 	debugf("weakFetch: key=%s", key)
 	owner := shortuuid.New()
 	r, err := c.luaGet(ctx, key, owner)
+	ticker := time.NewTimer(c.Options.LockSleep)
+	defer ticker.Stop()
 	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 <-ticker.C:
+			log.Printf("ticker")
+			// equal to time.Sleep(c.Options.LockSleep) but can be canceled
+		}
 		r, err = c.luaGet(ctx, key, owner)
+		// Reset ticker after luaGet
+		// If we reset ticker before luaGet, since luaGet takes a period of time,
+		// the actual sleep time will be shorter than expected
+		if !ticker.Stop() && len(ticker.C) > 0 {
+			<-ticker.C
+		}
+		ticker.Reset(c.Options.LockSleep)
 	}
 	if err != nil {
 		return "", err
@@ -217,10 +233,25 @@ func (c *Client) strongFetch(ctx context.Context, key string, expire time.Durati
 	debugf("strongFetch: key=%s", key)
 	owner := shortuuid.New()
 	r, err := c.luaGet(ctx, key, owner)
+	ticker := time.NewTimer(c.Options.LockSleep)
+	defer ticker.Stop()
 	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 <-ticker.C:
+			log.Printf("ticker")
+			// equal to time.Sleep(c.Options.LockSleep) but can be canceled
+		}
 		r, err = c.luaGet(ctx, key, owner)
+		// Reset ticker after luaGet
+		// If we reset ticker before luaGet, since luaGet takes a period of time,
+		// the actual sleep time will be shorter than expected
+		if !ticker.Stop() && len(ticker.C) > 0 {
+			<-ticker.C
+		}
+		ticker.Reset(c.Options.LockSleep)
 	}
 	if err != nil {
 		return "", err

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=