From 67a9e8aa001a7e8b504c681cd5ed00b423d25336 Mon Sep 17 00:00:00 2001 From: Ken Hibino Date: Thu, 21 Nov 2019 21:45:27 -0800 Subject: [PATCH] Track in-progress tasks with redis SET --- asynq.go | 2 ++ processor.go | 15 +++++++++++---- rdb.go | 37 +++++++++++++++++++++++++++---------- rdb_test.go | 20 ++++++++++++++++---- 4 files changed, 56 insertions(+), 18 deletions(-) diff --git a/asynq.go b/asynq.go index a27c14d..53b1297 100644 --- a/asynq.go +++ b/asynq.go @@ -3,6 +3,8 @@ package asynq /* TODOs: - [P0] Write tests +- [P0] Add heartbeats and rescuer +- [P0] Make all moving operations atomic - [P1] Add Support for multiple queues - [P1] User defined max-retry count - [P2] Web UI diff --git a/processor.go b/processor.go index e2330a0..eb91ad8 100644 --- a/processor.go +++ b/processor.go @@ -56,12 +56,14 @@ func (p *processor) start() { }() } +// exec pulls a task out of the queue and starts a worker goroutine to +// process the task. func (p *processor) exec() { + // NOTE: BLPOP needs to timeout to avoid blocking forever + // in case of a program shutdown or additon of a new queue. const timeout = 5 * time.Second - // pull a task out of the queue and process it // TODO(hibiken): sort the list of queues in order of priority - // NOTE: BLPOP needs to timeout in case a new queue is added. - msg, err := p.rdb.bpop(timeout, p.rdb.listQueues()...) + msg, err := p.rdb.dequeue(timeout, p.rdb.listQueues()...) if err != nil { switch err { case errQueuePopTimeout: @@ -79,7 +81,12 @@ func (p *processor) exec() { t := &Task{Type: msg.Type, Payload: msg.Payload} p.sema <- struct{}{} // acquire token go func(task *Task) { - defer func() { <-p.sema }() // release token + defer func() { + if err := p.rdb.srem(inProgress, msg); err != nil { + log.Printf("[SERVER ERROR] SREM failed: %v\n", err) + } + <-p.sema // release token + }() err := p.handler(task) if err != nil { retryTask(p.rdb, msg, err) diff --git a/rdb.go b/rdb.go index 33136b0..4aa7da5 100644 --- a/rdb.go +++ b/rdb.go @@ -13,11 +13,12 @@ import ( // Redis keys const ( - queuePrefix = "asynq:queues:" // LIST - allQueues = "asynq:queues" // SET - scheduled = "asynq:scheduled" // ZSET - retry = "asynq:retry" // ZSET - dead = "asynq:dead" // ZSET + queuePrefix = "asynq:queues:" // LIST + allQueues = "asynq:queues" // SET + scheduled = "asynq:scheduled" // ZSET + retry = "asynq:retry" // ZSET + dead = "asynq:dead" // ZSET + inProgress = "asynq:in_progress" // SET ) var ( @@ -55,18 +56,22 @@ func (r *rdb) push(msg *taskMessage) error { return nil } -// bpop blocks until there is a taskMessage available to be processed, -// returns immediately if there are already tasks waiting to be processed. -func (r *rdb) bpop(timeout time.Duration, keys ...string) (*taskMessage, error) { +// dequeue blocks until there is a taskMessage available to be processed, +// once available, it adds the task to "in progress" set and returns the task. +func (r *rdb) dequeue(timeout time.Duration, keys ...string) (*taskMessage, error) { + // TODO(hibiken): Make BLPOP & SADD atomic. res, err := r.client.BLPop(timeout, keys...).Result() if err != nil { if err != redis.Nil { - return nil, fmt.Errorf("command BLPOP %v %v failed: %v", - timeout, keys, err) + return nil, fmt.Errorf("command BLPOP %v %v failed: %v", timeout, keys, err) } return nil, errQueuePopTimeout } q, data := res[0], res[1] + err = r.client.SAdd(inProgress, data).Err() + if err != nil { + return nil, fmt.Errorf("command SADD %q %v failed: %v", inProgress, data, err) + } var msg taskMessage err = json.Unmarshal([]byte(data), &msg) if err != nil { @@ -76,6 +81,18 @@ func (r *rdb) bpop(timeout time.Duration, keys ...string) (*taskMessage, error) return &msg, nil } +func (r *rdb) srem(key string, msg *taskMessage) error { + bytes, err := json.Marshal(msg) + if err != nil { + return fmt.Errorf("could not encode task into JSON: %v", err) + } + err = r.client.SRem(key, string(bytes)).Err() + if err != nil { + return fmt.Errorf("command SREM %s %s failed: %v", key, string(bytes), err) + } + return nil +} + // zadd adds the taskMessage to the specified zset (sorted set) with the given score. func (r *rdb) zadd(zset string, zscore float64, msg *taskMessage) error { bytes, err := json.Marshal(msg) diff --git a/rdb_test.go b/rdb_test.go index 410e5a8..0eb5541 100644 --- a/rdb_test.go +++ b/rdb_test.go @@ -51,7 +51,7 @@ func TestPush(t *testing.T) { } } -func TestBPopImmediateReturn(t *testing.T) { +func TestDequeueImmediateReturn(t *testing.T) { r := setup() msg := &taskMessage{ Type: "GenerateCSVExport", @@ -60,19 +60,31 @@ func TestBPopImmediateReturn(t *testing.T) { } r.push(msg) - res, err := r.bpop(time.Second, "asynq:queues:csv") + res, err := r.dequeue(time.Second, "asynq:queues:csv") if err != nil { t.Fatalf("r.bpop() failed: %v", err) } + if !cmp.Equal(res, msg) { t.Errorf("cmp.Equal(res, msg) = %t, want %t", false, true) } + jobs := client.SMembers(inProgress).Val() + if len(jobs) != 1 { + t.Fatalf("len(jobs) = %d, want %d", len(jobs), 1) + } + var tm taskMessage + if err := json.Unmarshal([]byte(jobs[0]), &tm); err != nil { + t.Fatalf("json.Marshal() failed: %v", err) + } + if diff := cmp.Diff(res, &tm); diff != "" { + t.Errorf("cmp.Diff(res, tm) = %s", diff) + } } -func TestBPopTimeout(t *testing.T) { +func TestDequeueTimeout(t *testing.T) { r := setup() - _, err := r.bpop(time.Second, "asynq:queues:default") + _, err := r.dequeue(time.Second, "asynq:queues:default") if err != errQueuePopTimeout { t.Errorf("err = %v, want %v", err, errQueuePopTimeout) }