Use zset for aggregation set to preserve score

internal/asynqtest/asynqtest.go

@@ -254,10 +254,10 @@ func SeedGroup(tb testing.TB, r redis.UniversalClient, entries []base.Z, qname,
 	seedRedisZSet(tb, r, base.GroupKey(qname, gname), entries, base.TaskStateAggregating)
 }
 
-func SeedAggregationSet(tb testing.TB, r redis.UniversalClient, msgs []*base.TaskMessage, qname, gname, setID string) {
+func SeedAggregationSet(tb testing.TB, r redis.UniversalClient, entries []base.Z, qname, gname, setID string) {
 	tb.Helper()
 	r.SAdd(context.Background(), base.AllQueues, qname)
-	seedRedisSet(tb, r, base.AggregationSetKey(qname, gname, setID), msgs, base.TaskStateAggregating)
+	seedRedisZSet(tb, r, base.AggregationSetKey(qname, gname, setID), entries, base.TaskStateAggregating)
 }
 
 // SeedAllPendingQueues initializes all of the specified queues with the given messages.
@@ -386,33 +386,6 @@ func seedRedisZSet(tb testing.TB, c redis.UniversalClient, key string,
 	}
 }
 
-func seedRedisSet(tb testing.TB, c redis.UniversalClient, key string,
-	msgs []*base.TaskMessage, state base.TaskState) {
-	tb.Helper()
-	for _, msg := range msgs {
-		encoded := MustMarshal(tb, msg)
-		if err := c.SAdd(context.Background(), key, msg.ID).Err(); err != nil {
-			tb.Fatal(err)
-		}
-		taskKey := base.TaskKey(msg.Queue, msg.ID)
-		data := map[string]interface{}{
-			"msg":        encoded,
-			"state":      state.String(),
-			"unique_key": msg.UniqueKey,
-			"group":      msg.GroupKey,
-		}
-		if err := c.HSet(context.Background(), taskKey, data).Err(); err != nil {
-			tb.Fatal(err)
-		}
-		if len(msg.UniqueKey) > 0 {
-			err := c.SetNX(context.Background(), msg.UniqueKey, msg.ID, 1*time.Minute).Err()
-			if err != nil {
-				tb.Fatalf("Failed to set unique lock in redis: %v", err)
-			}
-		}
-	}
-}
-
 // GetPendingMessages returns all pending messages in the given queue.
 // It also asserts the state field of the task.
 func GetPendingMessages(tb testing.TB, r redis.UniversalClient, qname string) []*base.TaskMessage {

internal/rdb/rdb.go

@@ -1015,9 +1015,9 @@ if size == 0 then
 end
 local maxSize = tonumber(ARGV[1])
 if maxSize ~= 0 and size >= maxSize then
-	local msgs = redis.call("ZRANGE", KEYS[1], 0, maxSize-1)
+	local res = redis.call("ZRANGE", KEYS[1], 0, maxSize-1, "WITHSCORES")
-	for _, msg in ipairs(msgs) do
+	for i=1, table.getn(res)-1, 2 do
-		redis.call("SADD", KEYS[2], msg)
+		redis.call("ZADD", KEYS[2], tonumber(res[i+1]), res[i])
 	end
 	redis.call("ZREMRANGEBYRANK", KEYS[1], 0, maxSize-1)
 	redis.call("ZADD", KEYS[3], ARGV[5], ARGV[4])
@@ -1030,9 +1030,9 @@ if maxDelay ~= 0 then
 	local oldestEntryScore = tonumber(oldestEntry[2])
 	local maxDelayTime = currentTime - maxDelay
 	if oldestEntryScore <= maxDelayTime then
-		local msgs = redis.call("ZRANGE", KEYS[1], 0, maxSize-1)
+		local res = redis.call("ZRANGE", KEYS[1], 0, maxSize-1, "WITHSCORES")
-		for _, msg in ipairs(msgs) do
+		for i=1, table.getn(res)-1, 2 do
-			redis.call("SADD", KEYS[2], msg)
+			redis.call("ZADD", KEYS[2], tonumber(res[i+1]), res[i])
 		end
 		redis.call("ZREMRANGEBYRANK", KEYS[1], 0, maxSize-1)
 		redis.call("ZADD", KEYS[3], ARGV[5], ARGV[4])
@@ -1043,9 +1043,9 @@ local latestEntry = redis.call("ZREVRANGE", KEYS[1], 0, 0, "WITHSCORES")
 local latestEntryScore = tonumber(latestEntry[2])
 local gracePeriodStartTime = currentTime - tonumber(ARGV[3])
 if latestEntryScore <= gracePeriodStartTime then
-	local msgs = redis.call("ZRANGE", KEYS[1], 0, maxSize-1)
+	local res = redis.call("ZRANGE", KEYS[1], 0, maxSize-1, "WITHSCORES")
-	for _, msg in ipairs(msgs) do
+	for i=1, table.getn(res)-1, 2 do
-		redis.call("SADD", KEYS[2], msg)
+		redis.call("ZADD", KEYS[2], tonumber(res[i+1]), res[i])
 	end
 	redis.call("ZREMRANGEBYRANK", KEYS[1], 0, maxSize-1)
 	redis.call("ZADD", KEYS[3], ARGV[5], ARGV[4])
@@ -1101,7 +1101,7 @@ func (r *RDB) AggregationCheck(qname, gname string, t time.Time, gracePeriod, ma
 // ARGV[1] -> task key prefix
 var readAggregationSetCmd = redis.NewScript(`
 local msgs = {}
-local ids = redis.call("SMEMBERS", KEYS[1])
+local ids = redis.call("ZRANGE", KEYS[1], 0, -1)
 for _, id in ipairs(ids) do
 	local key = ARGV[1] .. id
 	table.insert(msgs, redis.call("HGET", key, "msg"))
@@ -1142,7 +1142,7 @@ func (r *RDB) ReadAggregationSet(qname, gname, setID string) ([]*base.TaskMessag
 // -------
 // ARGV[1] -> task key prefix
 var deleteAggregationSetCmd = redis.NewScript(`
-local ids = redis.call("SMEMBERS", KEYS[1])
+local ids = redis.call("ZRANGE", KEYS[1], 0, -1)
 for _, id in ipairs(ids)  do
 	redis.call("DEL", ARGV[1] .. id)
 end

internal/rdb/rdb_test.go

@@ -3359,6 +3359,7 @@ func TestDeleteAggregationSet(t *testing.T) {
 	r := setup(t)
 	defer r.Close()
 
+	now := time.Now()
 	ctx := context.Background()
 	setID := uuid.NewString()
 	msg1 := h.NewTaskMessageBuilder().SetType("foo").SetQueue("default").SetGroup("mygroup").Build()
@@ -3366,16 +3367,20 @@ func TestDeleteAggregationSet(t *testing.T) {
 	msg3 := h.NewTaskMessageBuilder().SetType("baz").SetQueue("default").SetGroup("mygroup").Build()
 
 	tests := []struct {
-		aggregationSet []*base.TaskMessage
+		aggregationSet []base.Z
 		qname          string
 		gname          string
 		setID          string
 	}{
 		{
-			aggregationSet: []*base.TaskMessage{msg1, msg2, msg3},
+			aggregationSet: []base.Z{
-			qname:          "default",
+				{msg1, now.Add(-3 * time.Minute).Unix()},
-			gname:          "mygroup",
+				{msg2, now.Add(-2 * time.Minute).Unix()},
-			setID:          setID,
+				{msg3, now.Add(-1 * time.Minute).Unix()},
+			},
+			qname: "default",
+			gname: "mygroup",
+			setID: setID,
 		},
 	}
 
@@ -3393,8 +3398,8 @@ func TestDeleteAggregationSet(t *testing.T) {
 		}
 
 		// Check all tasks in the set are deleted.
-		for _, m := range tc.aggregationSet {
+		for _, z := range tc.aggregationSet {
-			taskKey := base.TaskKey(m.Queue, m.ID)
+			taskKey := base.TaskKey(z.Message.Queue, z.Message.ID)
 			if r.client.Exists(ctx, taskKey).Val() != 0 {
 				t.Errorf("task key %q still exists", taskKey)
 			}