diff --git a/internal/asynqtest/asynqtest.go b/internal/asynqtest/asynqtest.go index 9fc527e..ab7a961 100644 --- a/internal/asynqtest/asynqtest.go +++ b/internal/asynqtest/asynqtest.go @@ -283,14 +283,21 @@ func seedRedisList(tb testing.TB, c redis.UniversalClient, key string, } key := base.TaskKey(msg.Queue, msg.ID.String()) data := map[string]interface{}{ - "msg": encoded, - "state": state.String(), - "timeout": msg.Timeout, - "deadline": msg.Deadline, + "msg": encoded, + "state": state.String(), + "timeout": msg.Timeout, + "deadline": msg.Deadline, + "unique_key": msg.UniqueKey, } if err := c.HSet(key, data).Err(); err != nil { tb.Fatal(err) } + if len(msg.UniqueKey) > 0 { + err := c.SetNX(msg.UniqueKey, msg.ID.String(), 1*time.Minute).Err() + if err != nil { + tb.Fatalf("Failed to set unique lock in redis: %v", err) + } + } } } @@ -306,14 +313,21 @@ func seedRedisZSet(tb testing.TB, c redis.UniversalClient, key string, } key := base.TaskKey(msg.Queue, msg.ID.String()) data := map[string]interface{}{ - "msg": encoded, - "state": state.String(), - "timeout": msg.Timeout, - "deadline": msg.Deadline, + "msg": encoded, + "state": state.String(), + "timeout": msg.Timeout, + "deadline": msg.Deadline, + "unique_key": msg.UniqueKey, } if err := c.HSet(key, data).Err(); err != nil { tb.Fatal(err) } + if len(msg.UniqueKey) > 0 { + err := c.SetNX(msg.UniqueKey, msg.ID.String(), 1*time.Minute).Err() + if err != nil { + tb.Fatalf("Failed to set unique lock in redis: %v", err) + } + } } } diff --git a/internal/base/base.go b/internal/base/base.go index 0c9f34d..6123dd1 100644 --- a/internal/base/base.go +++ b/internal/base/base.go @@ -170,6 +170,7 @@ func SchedulerHistoryKey(entryID string) string { } // UniqueKey returns a redis key with the given type, payload, and queue name. +// FIXME: We probably need to generate a hash of payload to make this key unique func UniqueKey(qname, tasktype string, payload []byte) string { return fmt.Sprintf("%sunique:%s:%s", QueueKeyPrefix(qname), tasktype, string(payload)) } diff --git a/internal/rdb/inspect.go b/internal/rdb/inspect.go index c84bd4d..2417563 100644 --- a/internal/rdb/inspect.go +++ b/internal/rdb/inspect.go @@ -509,7 +509,7 @@ func (r *RDB) ListArchived(qname string, pgn Pagination) ([]base.Z, error) { if !r.client.SIsMember(base.AllQueues, qname).Val() { return nil, errors.E(op, errors.NotFound, &errors.QueueNotFoundError{Queue: qname}) } - res, err := r.listZSetEntries(base.ArchivedKey(qname), qname, pgn) + zs, err := r.listZSetEntries(base.ArchivedKey(qname), qname, pgn) if err != nil { return nil, errors.E(op, errors.CanonicalCode(err), err) } @@ -719,7 +719,7 @@ func (r *RDB) runAll(zset, qname string) (int64, error) { } res, err := runAllCmd.Run(r.client, keys, argv...).Result() if err != nil { - return err + return 0, err } n, ok := res.(int64) if !ok { @@ -994,6 +994,10 @@ else return redis.error_reply("task is not found in zset: " .. tostring(state)) end end +local unique_key = redis.call("HGET", KEYS[1], "unique_key") +if unique_key ~= nil and unique_key ~= "" and redis.call("GET", unique_key) == ARGV[1] then + redis.call("DEL", unique_key) +end return redis.call("DEL", KEYS[1]) `) @@ -1089,7 +1093,12 @@ func (r *RDB) DeleteAllScheduledTasks(qname string) (int64, error) { var deleteAllCmd = redis.NewScript(` local ids = redis.call("ZRANGE", KEYS[1], 0, -1) for _, id in ipairs(ids) do - redis.call("DEL", ARGV[1] .. id) + local task_key = ARGV[1] .. id + local unique_key = redis.call("HGET", task_key, "unique_key") + if unique_key ~= nil and unique_key ~= "" and redis.call("GET", unique_key) == id then + redis.call("DEL", unique_key) + end + redis.call("DEL", task_key) end redis.call("DEL", KEYS[1]) return table.getn(ids)`) diff --git a/internal/rdb/inspect_test.go b/internal/rdb/inspect_test.go index 00500ad..c09356f 100644 --- a/internal/rdb/inspect_test.go +++ b/internal/rdb/inspect_test.go @@ -5,6 +5,7 @@ package rdb import ( + "encoding/json" "fmt" "testing" "time" @@ -3117,6 +3118,62 @@ func TestDeletePendingTask(t *testing.T) { } } +func TestDeleteTaskWithUniqueLock(t *testing.T) { + r := setup(t) + defer r.Close() + m1 := &base.TaskMessage{ + ID: uuid.New(), + Type: "email", + Payload: h.JSON(map[string]interface{}{"user_id": json.Number("123")}), + Queue: base.DefaultQueueName, + UniqueKey: base.UniqueKey(base.DefaultQueueName, "email", h.JSON(map[string]interface{}{"user_id": 123})), + } + t1 := time.Now().Add(3 * time.Hour) + + tests := []struct { + scheduled map[string][]base.Z + qname string + id uuid.UUID + uniqueKey string + wantScheduled map[string][]*base.TaskMessage + }{ + { + scheduled: map[string][]base.Z{ + "default": { + {Message: m1, Score: t1.Unix()}, + }, + }, + qname: "default", + id: m1.ID, + uniqueKey: m1.UniqueKey, + wantScheduled: map[string][]*base.TaskMessage{ + "default": {}, + }, + }, + } + + for _, tc := range tests { + h.FlushDB(t, r.client) // clean up db before each test case + h.SeedAllScheduledQueues(t, r.client, tc.scheduled) + + if got := r.DeleteTask(tc.qname, tc.id); got != nil { + t.Errorf("r.DeleteTask(%q, %v) returned error: %v", tc.qname, tc.id, got) + continue + } + + for qname, want := range tc.wantScheduled { + gotScheduled := h.GetScheduledMessages(t, r.client, qname) + if diff := cmp.Diff(want, gotScheduled, h.SortMsgOpt); diff != "" { + t.Errorf("mismatch found in %q; (-want, +got)\n%s", base.ScheduledKey(qname), diff) + } + } + + if r.client.Exists(tc.uniqueKey).Val() != 0 { + t.Errorf("Uniqueness lock %q still exists", tc.uniqueKey) + } + } +} + func TestDeleteTaskError(t *testing.T) { r := setup(t) defer r.Close() @@ -3305,6 +3362,7 @@ func TestDeleteAllArchivedTasksWithUniqueKey(t *testing.T) { archived map[string][]base.Z qname string want int64 + uniqueKeys []string // list of unique keys that should be cleared wantArchived map[string][]*base.TaskMessage }{ { @@ -3315,8 +3373,9 @@ func TestDeleteAllArchivedTasksWithUniqueKey(t *testing.T) { {Message: m3, Score: time.Now().Unix()}, }, }, - qname: "default", - want: 3, + qname: "default", + want: 3, + uniqueKeys: []string{m1.UniqueKey, m2.UniqueKey}, wantArchived: map[string][]*base.TaskMessage{ "default": {}, }, @@ -3326,18 +3385,6 @@ func TestDeleteAllArchivedTasksWithUniqueKey(t *testing.T) { for _, tc := range tests { h.FlushDB(t, r.client) // clean up db before each test case h.SeedAllArchivedQueues(t, r.client, tc.archived) - var uniqueKeys []string // list of unique keys set in redis - for _, zs := range tc.archived { - for _, z := range zs { - if len(z.Message.UniqueKey) > 0 { - err := r.client.SetNX(z.Message.UniqueKey, z.Message.ID.String(), time.Minute).Err() - if err != nil { - t.Fatalf("Failed to set unique lock in redis: %v", err) - } - uniqueKeys = append(uniqueKeys, z.Message.UniqueKey) - } - } - } got, err := r.DeleteAllArchivedTasks(tc.qname) if err != nil { @@ -3353,7 +3400,7 @@ func TestDeleteAllArchivedTasksWithUniqueKey(t *testing.T) { } } - for _, uniqueKey := range uniqueKeys { + for _, uniqueKey := range tc.uniqueKeys { if r.client.Exists(uniqueKey).Val() != 0 { t.Errorf("Uniqueness lock %q still exists", uniqueKey) } diff --git a/internal/rdb/rdb.go b/internal/rdb/rdb.go index dc94194..36a8343 100644 --- a/internal/rdb/rdb.go +++ b/internal/rdb/rdb.go @@ -114,7 +114,8 @@ redis.call("HSET", KEYS[2], "msg", ARGV[3], "state", "pending", "timeout", ARGV[4], - "deadline", ARGV[5]) + "deadline", ARGV[5], + "unique_key", KEYS[1]) redis.call("LPUSH", KEYS[3], ARGV[1]) return 1 `) @@ -407,7 +408,8 @@ redis.call("HSET", KEYS[2], "msg", ARGV[4], "state", "scheduled", "timeout", ARGV[5], - "deadline", ARGV[6]) + "deadline", ARGV[6], + "unique_key", KEYS[1]) redis.call("ZADD", KEYS[3], ARGV[3], ARGV[1]) return 1 `) diff --git a/internal/rdb/rdb_test.go b/internal/rdb/rdb_test.go index fc6508a..4e2b7d7 100644 --- a/internal/rdb/rdb_test.go +++ b/internal/rdb/rdb_test.go @@ -191,7 +191,11 @@ func TestEnqueueUnique(t *testing.T) { } deadline := r.client.HGet(taskKey, "deadline").Val() // "deadline" field if want := strconv.Itoa(int(tc.msg.Deadline)); deadline != want { - t.Errorf("deadline field under task-ke is set to %v, want %v", deadline, want) + t.Errorf("deadline field under task-key is set to %v, want %v", deadline, want) + } + uniqueKey := r.client.HGet(taskKey, "unique_key").Val() // "unique_key" field + if uniqueKey != tc.msg.UniqueKey { + t.Errorf("uniqueue_key field under task key is set to %q, want %q", uniqueKey, tc.msg.UniqueKey) } // Check queue is in the AllQueues set. @@ -1008,7 +1012,11 @@ func TestScheduleUnique(t *testing.T) { } deadline := r.client.HGet(taskKey, "deadline").Val() // "deadline" field if want := strconv.Itoa(int(tc.msg.Deadline)); deadline != want { - t.Errorf("deadline field under task-ke is set to %v, want %v", deadline, want) + t.Errorf("deadline field under task-key is set to %v, want %v", deadline, want) + } + uniqueKey := r.client.HGet(taskKey, "unique_key").Val() // "unique_key" field + if uniqueKey != tc.msg.UniqueKey { + t.Errorf("uniqueue_key field under task key is set to %q, want %q", uniqueKey, tc.msg.UniqueKey) } // Check queue is in the AllQueues set.