From becd26479b34c96e491304d049fa30e571c56df3 Mon Sep 17 00:00:00 2001 From: Ken Hibino Date: Mon, 10 Aug 2020 21:49:12 -0700 Subject: [PATCH] Update WriteServerState and ClearServerState in RDB --- internal/rdb/rdb.go | 66 ++++++++++++++++++++-------------------- internal/rdb/rdb_test.go | 44 +++++++++++++++------------ 2 files changed, 57 insertions(+), 53 deletions(-) diff --git a/internal/rdb/rdb.go b/internal/rdb/rdb.go index 017fe9a..c78d58d 100644 --- a/internal/rdb/rdb.go +++ b/internal/rdb/rdb.go @@ -309,7 +309,7 @@ func (r *RDB) ScheduleUnique(msg *base.TaskMessage, processAt time.Time, ttl tim // KEYS[2] -> asynq:{}:deadlines // KEYS[3] -> asynq:{}:retry // KEYS[4] -> asynq:{}:processed: -// KEYS[5] -> asynq:{}:failure: +// KEYS[5] -> asynq:{}:failed: // ARGV[1] -> base.TaskMessage value to remove from base.InProgressQueue queue // ARGV[2] -> base.TaskMessage value to add to Retry queue // ARGV[3] -> retry_at UNIX timestamp @@ -364,7 +364,7 @@ const ( // KEYS[2] -> asynq:{}:deadlines // KEYS[3] -> asynq:{}:dead // KEYS[4] -> asynq:{}:processed: -// KEYS[5] -> asynq:{}:failure: +// KEYS[5] -> asynq:{}:failed: // ARGV[1] -> base.TaskMessage value to remove from base.InProgressQueue queue // ARGV[2] -> base.TaskMessage value to add to Dead queue // ARGV[3] -> died_at UNIX timestamp @@ -408,10 +408,10 @@ func (r *RDB) Kill(msg *base.TaskMessage, errMsg string) error { now := time.Now() limit := now.AddDate(0, 0, -deadExpirationInDays).Unix() // 90 days ago processedKey := base.ProcessedKey(msg.Queue, now) - failureKey := base.FailureKey(msg.Queue, now) + failedKey := base.FailedKey(msg.Queue, now) expireAt := now.Add(statsTTL) return killCmd.Run(r.client, - []string{base.InProgressKey(msg.Queue), base.DeadlinesKey(msg.Queue), base.DeadKey(msg.Queue), processedKey, failureKey}, + []string{base.InProgressKey(msg.Queue), base.DeadlinesKey(msg.Queue), base.DeadKey(msg.Queue), processedKey, failedKey}, msgToRemove, msgToAdd, now.Unix(), limit, maxDeadTasks, expireAt.Unix()).Err() } @@ -454,7 +454,7 @@ func (r *RDB) forward(src, dst string) (int, error) { // forwardAll moves tasks with a score less than the current unix time from the src zset, // until there's no more tasks. -func (r *RDB) forwardAll(src, dst string) error { +func (r *RDB) forwardAll(src, dst string) (err error) { n := 1 for n != 0 { n, err = r.forward(src, dst) @@ -488,25 +488,20 @@ func (r *RDB) ListDeadlineExceeded(deadline time.Time, qnames ...string) ([]*bas return msgs, nil } -// KEYS[1] -> asynq:servers: -// KEYS[2] -> asynq:servers -// KEYS[3] -> asynq:workers -// KEYS[4] -> asynq:workers -// ARGV[1] -> expiration time -// ARGV[2] -> TTL in seconds -// ARGV[3] -> server info -// ARGV[4:] -> alternate key-value pair of (worker id, worker data) +// KEYS[1] -> asynq:servers:{} +// KEYS[2] -> asynq:workers:{} +// ARGV[1] -> TTL in seconds +// ARGV[2] -> server info +// ARGV[3:] -> alternate key-value pair of (worker id, worker data) // Note: Add key to ZSET with expiration time as score. // ref: https://github.com/antirez/redis/issues/135#issuecomment-2361996 var writeServerStateCmd = redis.NewScript(` -redis.call("SETEX", KEYS[1], ARGV[2], ARGV[3]) -redis.call("ZADD", KEYS[2], ARGV[1], KEYS[1]) -redis.call("DEL", KEYS[3]) -for i = 4, table.getn(ARGV)-1, 2 do - redis.call("HSET", KEYS[3], ARGV[i], ARGV[i+1]) +redis.call("SETEX", KEYS[1], ARGV[1], ARGV[2]) +redis.call("DEL", KEYS[2]) +for i = 3, table.getn(ARGV)-1, 2 do + redis.call("HSET", KEYS[2], ARGV[i], ARGV[i+1]) end -redis.call("EXPIRE", KEYS[3], ARGV[2]) -redis.call("ZADD", KEYS[4], ARGV[1], KEYS[3]) +redis.call("EXPIRE", KEYS[2], ARGV[1]) return redis.status_reply("OK")`) // WriteServerState writes server state data to redis with expiration set to the value ttl. @@ -516,7 +511,7 @@ func (r *RDB) WriteServerState(info *base.ServerInfo, workers []*base.WorkerInfo return err } exp := time.Now().Add(ttl).UTC() - args := []interface{}{float64(exp.Unix()), ttl.Seconds(), bytes} // args to the lua script + args := []interface{}{ttl.Seconds(), bytes} // args to the lua script for _, w := range workers { bytes, err := json.Marshal(w) if err != nil { @@ -526,28 +521,33 @@ func (r *RDB) WriteServerState(info *base.ServerInfo, workers []*base.WorkerInfo } skey := base.ServerInfoKey(info.Host, info.PID, info.ServerID) wkey := base.WorkersKey(info.Host, info.PID, info.ServerID) - return writeServerStateCmd.Run(r.client, - []string{skey, base.AllServers, wkey, base.AllWorkers}, - args...).Err() + if err := r.client.ZAdd(base.AllServers, &redis.Z{Score: float64(exp.Unix()), Member: skey}).Err(); err != nil { + return err + } + if err := r.client.ZAdd(base.AllWorkers, &redis.Z{Score: float64(exp.Unix()), Member: wkey}).Err(); err != nil { + return err + } + return writeServerStateCmd.Run(r.client, []string{skey, wkey}, args...).Err() } -// KEYS[1] -> asynq:servers -// KEYS[2] -> asynq:servers: -// KEYS[3] -> asynq:workers -// KEYS[4] -> asynq:workers +// KEYS[1] -> asynq:servers:{} +// KEYS[2] -> asynq:workers:{} var clearServerStateCmd = redis.NewScript(` -redis.call("ZREM", KEYS[1], KEYS[2]) +redis.call("DEL", KEYS[1]) redis.call("DEL", KEYS[2]) -redis.call("ZREM", KEYS[3], KEYS[4]) -redis.call("DEL", KEYS[4]) return redis.status_reply("OK")`) // ClearServerState deletes server state data from redis. func (r *RDB) ClearServerState(host string, pid int, serverID string) error { skey := base.ServerInfoKey(host, pid, serverID) wkey := base.WorkersKey(host, pid, serverID) - return clearServerStateCmd.Run(r.client, - []string{base.AllServers, skey, base.AllWorkers, wkey}).Err() + if err := r.client.ZRem(base.AllServers, skey).Err(); err != nil { + return err + } + if err := r.client.ZRem(base.AllWorkers, wkey).Err(); err != nil { + return err + } + return clearServerStateCmd.Run(r.client, []string{skey, wkey}).Err() } // CancelationPubSub returns a pubsub for cancelation messages. diff --git a/internal/rdb/rdb_test.go b/internal/rdb/rdb_test.go index a668bb1..f3da3cd 100644 --- a/internal/rdb/rdb_test.go +++ b/internal/rdb/rdb_test.go @@ -267,7 +267,7 @@ func TestDequeue(t *testing.T) { for _, tc := range tests { h.FlushDB(t, r.client) // clean up db before each test case - h.SeedAllEnqueuedQueues(t, r.client, msgs, queue, tc.enqueued) + h.SeedAllEnqueuedQueues(t, r.client, tc.enqueued) gotMsg, gotDeadline, err := r.Dequeue(tc.args...) if err != tc.err { @@ -438,6 +438,7 @@ func TestDone(t *testing.T) { t3Deadline := now.Unix() + t3.Deadline tests := []struct { + desc string inProgress map[string][]*base.TaskMessage // initial state of the in-progress list deadlines map[string][]base.Z // initial state of deadlines set target *base.TaskMessage // task to remove @@ -445,6 +446,7 @@ func TestDone(t *testing.T) { wantDeadlines map[string][]base.Z // final state of the deadline set }{ { + desc: "removes message from the correct queue", inProgress: map[string][]*base.TaskMessage{ "default": {t1}, "custom": {t2}, @@ -455,8 +457,8 @@ func TestDone(t *testing.T) { }, target: t1, wantInProgress: map[string][]*base.TaskMessage{ - "default": {t2}, - "custom": {}, + "default": {}, + "custom": {t2}, }, wantDeadlines: map[string][]base.Z{ "default": {}, @@ -464,6 +466,7 @@ func TestDone(t *testing.T) { }, }, { + desc: "with one queue", inProgress: map[string][]*base.TaskMessage{ "default": {t1}, }, @@ -479,6 +482,7 @@ func TestDone(t *testing.T) { }, }, { + desc: "with multiple messages in a queue", inProgress: map[string][]*base.TaskMessage{ "default": {t1, t3}, "custom": {t2}, @@ -489,7 +493,7 @@ func TestDone(t *testing.T) { }, target: t3, wantInProgress: map[string][]*base.TaskMessage{ - "defualt": {t1}, + "default": {t1}, "custom": {t2}, }, wantDeadlines: map[string][]base.Z{ @@ -517,21 +521,21 @@ func TestDone(t *testing.T) { err := r.Done(tc.target) if err != nil { - t.Errorf("(*RDB).Done(task) = %v, want nil", err) + t.Errorf("%s; (*RDB).Done(task) = %v, want nil", tc.desc, err) continue } for queue, want := range tc.wantInProgress { gotInProgress := h.GetInProgressMessages(t, r.client, queue) if diff := cmp.Diff(want, gotInProgress, h.SortMsgOpt); diff != "" { - t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.InProgressKey(queue), diff) + t.Errorf("%s; mismatch found in %q: (-want, +got):\n%s", tc.desc, base.InProgressKey(queue), diff) continue } } for queue, want := range tc.wantDeadlines { gotDeadlines := h.GetDeadlinesEntries(t, r.client, queue) if diff := cmp.Diff(want, gotDeadlines, h.SortZSetEntryOpt); diff != "" { - t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.DeadlinesKey(queue), diff) + t.Errorf("%s; mismatch found in %q: (-want, +got):\n%s", tc.desc, base.DeadlinesKey(queue), diff) continue } } @@ -539,16 +543,16 @@ func TestDone(t *testing.T) { processedKey := base.ProcessedKey(tc.target.Queue, time.Now()) gotProcessed := r.client.Get(processedKey).Val() if gotProcessed != "1" { - t.Errorf("GET %q = %q, want 1", processedKey, gotProcessed) + t.Errorf("%s; GET %q = %q, want 1", tc.desc, processedKey, gotProcessed) } gotTTL := r.client.TTL(processedKey).Val() if gotTTL > statsTTL { - t.Errorf("TTL %q = %v, want less than or equal to %v", processedKey, gotTTL, statsTTL) + t.Errorf("%s; TTL %q = %v, want less than or equal to %v", tc.desc, processedKey, gotTTL, statsTTL) } if len(tc.target.UniqueKey) > 0 && r.client.Exists(tc.target.UniqueKey).Val() != 0 { - t.Errorf("Uniqueness lock %q still exists", tc.target.UniqueKey) + t.Errorf("%s; Uniqueness lock %q still exists", tc.desc, tc.target.UniqueKey) } } } @@ -597,7 +601,7 @@ func TestRequeue(t *testing.T) { inProgress: map[string][]*base.TaskMessage{ "default": {t1, t2}, }, - deadlines: []base.Z{ + deadlines: map[string][]base.Z{ "default": { {Message: t1, Score: t1Deadline}, {Message: t2, Score: t2Deadline}, @@ -610,8 +614,8 @@ func TestRequeue(t *testing.T) { wantInProgress: map[string][]*base.TaskMessage{ "default": {t2}, }, - wantDeadlines: []base.Z{ - "defult": { + wantDeadlines: map[string][]base.Z{ + "default": { {Message: t2, Score: t2Deadline}, }, }, @@ -649,8 +653,8 @@ func TestRequeue(t *testing.T) { "critical": {t3}, }, deadlines: map[string][]base.Z{ - "defualt": {{Message: t2, Score: t2Deadline}}, - "critial": {{Message: t3, Score: t3Deadline}}, + "default": {{Message: t2, Score: t2Deadline}}, + "critical": {{Message: t3, Score: t3Deadline}}, }, target: t3, wantEnqueued: map[string][]*base.TaskMessage{ @@ -670,7 +674,7 @@ func TestRequeue(t *testing.T) { for _, tc := range tests { h.FlushDB(t, r.client) // clean up db before each test case - h.SeedAllEnqueuedQueues(t, r.client, msgs, tc.enqueued) + h.SeedAllEnqueuedQueues(t, r.client, tc.enqueued) h.SeedAllInProgressQueues(t, r.client, tc.inProgress) h.SeedAllDeadlines(t, r.client, tc.deadlines) @@ -694,7 +698,7 @@ func TestRequeue(t *testing.T) { } for qname, want := range tc.wantDeadlines { gotDeadlines := h.GetDeadlinesEntries(t, r.client, qname) - if diff := cmp.Diff(wnt, gotDeadlines, h.SortZSetEntryOpt); diff != "" { + if diff := cmp.Diff(want, gotDeadlines, h.SortZSetEntryOpt); diff != "" { t.Errorf("mismatch found in %q: (-want, +got):\n%s", base.DeadlinesKey(qname), diff) } } @@ -1028,7 +1032,7 @@ func TestKill(t *testing.T) { wantInProgress: map[string][]*base.TaskMessage{ "default": {t2}, }, - wantDeadlines: map[string]base.Z{ + wantDeadlines: map[string][]base.Z{ "default": {{Message: t2, Score: t2Deadline}}, }, wantDead: map[string][]base.Z{ @@ -1090,7 +1094,7 @@ func TestKill(t *testing.T) { "default": {t1}, "custom": {}, }, - wantDeadlines: map[string]base.Z{ + wantDeadlines: map[string][]base.Z{ "default": {{Message: t1, Score: t1Deadline}}, "custom": {}, }, @@ -1354,7 +1358,7 @@ func TestListDeadlineExceeded(t *testing.T) { }, qnames: []string{"default", "critical"}, t: time.Now(), - want: []*base.TaskMessage{t1, t2}, + want: []*base.TaskMessage{t1, t3}, }, { desc: "with empty in-progress queue",