Track in-progress tasks with redis SET

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

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)

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)

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)
 	}