Unblock processor shutdown process if processor is waiting for semaphore

token

internal/rdb/rdb.go

@@ -116,6 +116,25 @@ func (r *RDB) Done(msg *TaskMessage) error {
 	return nil
 }
 
+// Requeue moves the task from in-progress queue to the default
+// queue.
+func (r *RDB) Requeue(msg *TaskMessage) error {
+	bytes, err := json.Marshal(msg)
+	if err != nil {
+		return fmt.Errorf("could not marshal %+v to json: %v", msg, err)
+	}
+	// KEYS[1] -> asynq:in_progress
+	// KEYS[2] -> asynq:queues:default
+	// ARGV[1] -> taskMessage value
+	script := redis.NewScript(`
+	redis.call("LREM", KEYS[1], 0, ARGV[1])
+	redis.call("RPUSH", KEYS[2], ARGV[1])
+	return redis.status_reply("OK")
+	`)
+	_, err = script.Run(r.client, []string{inProgressQ, defaultQ}, string(bytes)).Result()
+	return err
+}
+
 // Schedule adds the task to the backlog queue to be processed in the future.
 func (r *RDB) Schedule(msg *TaskMessage, processAt time.Time) error {
 	bytes, err := json.Marshal(msg)

internal/rdb/rdb_test.go

@@ -112,6 +112,59 @@ func TestDone(t *testing.T) {
 	}
 }
 
+func TestRequeue(t *testing.T) {
+	r := setup(t)
+	t1 := newTaskMessage("send_email", nil)
+	t2 := newTaskMessage("export_csv", nil)
+
+	tests := []struct {
+		enqueued       []*TaskMessage // initial state of the default queue
+		inProgress     []*TaskMessage // initial state of the in-progress list
+		target         *TaskMessage   // task to requeue
+		wantEnqueued   []*TaskMessage // final state of the default queue
+		wantInProgress []*TaskMessage // final state of the in-progress list
+	}{
+		{
+			enqueued:       []*TaskMessage{},
+			inProgress:     []*TaskMessage{t1, t2},
+			target:         t1,
+			wantEnqueued:   []*TaskMessage{t1},
+			wantInProgress: []*TaskMessage{t2},
+		},
+		{
+			enqueued:       []*TaskMessage{t1},
+			inProgress:     []*TaskMessage{t2},
+			target:         t2,
+			wantEnqueued:   []*TaskMessage{t1, t2},
+			wantInProgress: []*TaskMessage{},
+		},
+	}
+
+	for _, tc := range tests {
+		flushDB(t, r) // clean up db before each test case
+		seedDefaultQueue(t, r, tc.enqueued)
+		seedInProgressQueue(t, r, tc.inProgress)
+
+		err := r.Requeue(tc.target)
+		if err != nil {
+			t.Errorf("(*RDB).Requeue(task) = %v, want nil", err)
+			continue
+		}
+
+		gotEnqueuedRaw := r.client.LRange(defaultQ, 0, -1).Val()
+		gotEnqueued := mustUnmarshalSlice(t, gotEnqueuedRaw)
+		if diff := cmp.Diff(tc.wantEnqueued, gotEnqueued, sortMsgOpt); diff != "" {
+			t.Errorf("mismatch found in %q: (-want, +got):\n%s", defaultQ, diff)
+		}
+
+		gotInProgressRaw := r.client.LRange(inProgressQ, 0, -1).Val()
+		gotInProgress := mustUnmarshalSlice(t, gotInProgressRaw)
+		if diff := cmp.Diff(tc.wantInProgress, gotInProgress, sortMsgOpt); diff != "" {
+			t.Errorf("mismatch found in %q: (-want, +got):\n%s", inProgressQ, diff)
+		}
+	}
+}
+
 func TestKill(t *testing.T) {
 	r := setup(t)
 	t1 := newTaskMessage("send_email", nil)