diff --git a/internal/rdb/rdb.go b/internal/rdb/rdb.go index 6848da9..3606fb1 100644 --- a/internal/rdb/rdb.go +++ b/internal/rdb/rdb.go @@ -194,12 +194,10 @@ func (r *RDB) dequeue(qkeys ...interface{}) (msgjson string, deadline int64, err // ARGV[3] -> task ID // Note: LREM count ZERO means "remove all elements equal to val" var doneCmd = redis.NewScript(` -local x = redis.call("LREM", KEYS[1], 0, ARGV[1]) -if x == 0 then +if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then return redis.error_reply("NOT FOUND") end -x = redis.call("ZREM", KEYS[2], ARGV[1]) -if x == 0 then +if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then return redis.error_reply("NOT FOUND") end local n = redis.call("INCR", KEYS[3]) @@ -228,12 +226,18 @@ func (r *RDB) Done(msg *base.TaskMessage) error { } // KEYS[1] -> asynq:in_progress -// KEYS[2] -> asynq:queues: +// KEYS[2] -> asynq:deadlines +// KEYS[3] -> asynq:queues: // ARGV[1] -> base.TaskMessage value // Note: Use RPUSH to push to the head of the queue. var requeueCmd = redis.NewScript(` -redis.call("LREM", KEYS[1], 0, ARGV[1]) -redis.call("RPUSH", KEYS[2], ARGV[1]) +if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then + return redis.error_reply("NOT FOUND") +end +if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then + return redis.error_reply("NOT FOUND") +end +redis.call("RPUSH", KEYS[3], ARGV[1]) return redis.status_reply("OK")`) // Requeue moves the task from in-progress queue to the specified queue. @@ -243,7 +247,7 @@ func (r *RDB) Requeue(msg *base.TaskMessage) error { return err } return requeueCmd.Run(r.client, - []string{base.InProgressQueue, base.QueueKey(msg.Queue)}, + []string{base.InProgressQueue, base.KeyDeadlines, base.QueueKey(msg.Queue)}, encoded).Err() } @@ -324,12 +328,10 @@ func (r *RDB) ScheduleUnique(msg *base.TaskMessage, processAt time.Time, ttl tim // ARGV[3] -> retry_at UNIX timestamp // ARGV[4] -> stats expiration timestamp var retryCmd = redis.NewScript(` -local x = redis.call("LREM", KEYS[1], 0, ARGV[1]) -if x == 0 then +if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then return redis.error_reply("NOT FOUND") end -x = redis.call("ZREM", KEYS[2], ARGV[1]) -if x == 0 then +if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then return redis.error_reply("NOT FOUND") end redis.call("ZADD", KEYS[3], ARGV[3], ARGV[2]) @@ -383,12 +385,10 @@ const ( // ARGV[5] -> max number of tasks in dead queue (e.g., 100) // ARGV[6] -> stats expiration timestamp var killCmd = redis.NewScript(` -local x = redis.call("LREM", KEYS[1], 0, ARGV[1]) -if x == 0 then +if redis.call("LREM", KEYS[1], 0, ARGV[1]) == 0 then return redis.error_reply("NOT FOUND") end -x = redis.call("ZREM", KEYS[2], ARGV[1]) -if x == 0 then +if redis.call("ZREM", KEYS[2], ARGV[1]) == 0 then return redis.error_reply("NOT FOUND") end redis.call("ZADD", KEYS[3], ARGV[3], ARGV[2]) diff --git a/internal/rdb/rdb_test.go b/internal/rdb/rdb_test.go index 34930e8..713b48c 100644 --- a/internal/rdb/rdb_test.go +++ b/internal/rdb/rdb_test.go @@ -533,38 +533,73 @@ func TestDone(t *testing.T) { func TestRequeue(t *testing.T) { r := setup(t) - t1 := h.NewTaskMessage("send_email", nil) - t2 := h.NewTaskMessage("export_csv", nil) - t3 := h.NewTaskMessageWithQueue("send_email", nil, "critical") + now := time.Now() + t1 := &base.TaskMessage{ + ID: xid.New(), + Type: "send_email", + Payload: nil, + Queue: "default", + Timeout: 1800, + } + t2 := &base.TaskMessage{ + ID: xid.New(), + Type: "export_csv", + Payload: nil, + Queue: "default", + Timeout: 3000, + } + t3 := &base.TaskMessage{ + ID: xid.New(), + Type: "send_email", + Payload: nil, + Queue: "critical", + Timeout: 80, + } + t1Deadline := int(now.Unix()) + t1.Timeout + t2Deadline := int(now.Unix()) + t2.Timeout + t3Deadline := int(now.Unix()) + t3.Timeout tests := []struct { enqueued map[string][]*base.TaskMessage // initial state of queues inProgress []*base.TaskMessage // initial state of the in-progress list + deadlines []h.ZSetEntry // initial state of the deadlines set target *base.TaskMessage // task to requeue wantEnqueued map[string][]*base.TaskMessage // final state of queues wantInProgress []*base.TaskMessage // final state of the in-progress list + wantDeadlines []h.ZSetEntry // final state of the deadlines set }{ { enqueued: map[string][]*base.TaskMessage{ base.DefaultQueueName: {}, }, inProgress: []*base.TaskMessage{t1, t2}, - target: t1, + deadlines: []h.ZSetEntry{ + {Msg: t1, Score: float64(t1Deadline)}, + {Msg: t2, Score: float64(t2Deadline)}, + }, + target: t1, wantEnqueued: map[string][]*base.TaskMessage{ base.DefaultQueueName: {t1}, }, wantInProgress: []*base.TaskMessage{t2}, + wantDeadlines: []h.ZSetEntry{ + {Msg: t2, Score: float64(t2Deadline)}, + }, }, { enqueued: map[string][]*base.TaskMessage{ base.DefaultQueueName: {t1}, }, inProgress: []*base.TaskMessage{t2}, - target: t2, + deadlines: []h.ZSetEntry{ + {Msg: t2, Score: float64(t2Deadline)}, + }, + target: t2, wantEnqueued: map[string][]*base.TaskMessage{ base.DefaultQueueName: {t1, t2}, }, wantInProgress: []*base.TaskMessage{}, + wantDeadlines: []h.ZSetEntry{}, }, { enqueued: map[string][]*base.TaskMessage{ @@ -572,12 +607,19 @@ func TestRequeue(t *testing.T) { "critical": {}, }, inProgress: []*base.TaskMessage{t2, t3}, - target: t3, + deadlines: []h.ZSetEntry{ + {Msg: t2, Score: float64(t2Deadline)}, + {Msg: t3, Score: float64(t3Deadline)}, + }, + target: t3, wantEnqueued: map[string][]*base.TaskMessage{ base.DefaultQueueName: {t1}, "critical": {t3}, }, wantInProgress: []*base.TaskMessage{t2}, + wantDeadlines: []h.ZSetEntry{ + {Msg: t2, Score: float64(t2Deadline)}, + }, }, } @@ -587,6 +629,7 @@ func TestRequeue(t *testing.T) { h.SeedEnqueuedQueue(t, r.client, msgs, qname) } h.SeedInProgressQueue(t, r.client, tc.inProgress) + h.SeedDeadlines(t, r.client, tc.deadlines) err := r.Requeue(tc.target) if err != nil { @@ -605,6 +648,10 @@ func TestRequeue(t *testing.T) { if diff := cmp.Diff(tc.wantInProgress, gotInProgress, h.SortMsgOpt); diff != "" { t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.InProgressQueue, diff) } + gotDeadlines := h.GetDeadlinesEntries(t, r.client) + if diff := cmp.Diff(tc.wantDeadlines, gotDeadlines, h.SortZSetEntryOpt); diff != "" { + t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.KeyDeadlines, diff) + } } }