use sorted-set to release stale locks

x/rate/semaphore.go

@@ -4,9 +4,11 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/go-redis/redis/v8"
 	"github.com/hibiken/asynq"
+	asynqcontext "github.com/hibiken/asynq/internal/context"
 )
 
 // NewSemaphore creates a new counting Semaphore.
@@ -38,28 +40,73 @@ type Semaphore struct {
 	name           string
 }
 
+// KEYS[1] -> asynq:sema:<scope>
+// ARGV[1] -> max concurrency
+// ARGV[2] -> current time in unix time
+// ARGV[3] -> deadline in unix time
+// ARGV[4] -> task ID
 var acquireCmd = redis.NewScript(`
-local lockCount = tonumber(redis.call('GET', KEYS[1]))
+redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", tonumber(ARGV[2])-1)
+local lockCount = redis.call("ZCARD", KEYS[1])
 
-if (not lockCount or lockCount < tonumber(ARGV[1])) then
-  redis.call('INCR', KEYS[1])
-  return true
+if (lockCount < tonumber(ARGV[1])) then
+     redis.call("ZADD", KEYS[1], ARGV[3], ARGV[4])
+     return true
 else
-  return false
+    return false
 end
 `)
 
 // Acquire will try to acquire lock on the counting semaphore.
-// - Returns true, iff semaphore key exists and current value is less than maxConcurrency
-// - Returns false otherwise
-func (s *Semaphore) Acquire(ctx context.Context) bool {
-	val, _ := acquireCmd.Run(ctx, s.rc, []string{semaphoreKey(s.name)}, []interface{}{s.maxConcurrency}...).Bool()
-	return val
+// - Returns (true, nil), iff semaphore key exists and current value is less than maxConcurrency
+// - Returns (false, nil) when lock cannot be acquired
+// - Returns (false, error) otherwise
+//
+// The context.Context passed to Acquire must have a deadline set,
+// this ensures that lock is released if the job goroutine crashes and does not call Release.
+func (s *Semaphore) Acquire(ctx context.Context) (bool, error) {
+	d, ok := ctx.Deadline()
+	if !ok {
+		return false, fmt.Errorf("provided context must have a deadline")
+	}
+
+	taskID, ok := asynqcontext.GetTaskID(ctx)
+	if !ok {
+		return false, fmt.Errorf("provided context is missing task ID value")
+	}
+
+	b, err := acquireCmd.Run(ctx, s.rc,
+		[]string{semaphoreKey(s.name)},
+		[]interface{}{
+			s.maxConcurrency,
+			time.Now().Unix(),
+			d.Unix(),
+			taskID,
+		}...).Bool()
+	if err == redis.Nil {
+		return b, nil
+	}
+
+	return b, err
 }
 
 // Release will release the lock on the counting semaphore.
-func (s *Semaphore) Release(ctx context.Context) {
-	s.rc.Decr(ctx, semaphoreKey(s.name))
+func (s *Semaphore) Release(ctx context.Context) error {
+	taskID, ok := asynqcontext.GetTaskID(ctx)
+	if !ok {
+		return fmt.Errorf("provided context is missing task ID value")
+	}
+
+	n, err := s.rc.ZRem(ctx, semaphoreKey(s.name), taskID).Result()
+	if err != nil {
+		return fmt.Errorf("redis command failed: %v", err)
+	}
+
+	if n == 0 {
+		return fmt.Errorf("no lock found for task %q", taskID)
+	}
+
+	return nil
 }
 
 // Close closes the connection with redis.
@@ -68,5 +115,5 @@ func (s *Semaphore) Close() error {
 }
 
 func semaphoreKey(name string) string {
-	return fmt.Sprintf("asynq:sema:{%s}", name)
+	return fmt.Sprintf("asynq:sema:%s", name)
 }