TWEAK: speed up sidekiq pull ack (#113)

Author: taylorchu

sidekiq/pull.go

@@ -120,16 +120,26 @@ func (q *sidekiqQueue) Pull(opt *PullOptions) error {
 		}
 	}()
 
-	pull := func() error {
-		res, err := q.dequeueScript.Run(ctx, q.client, []string{queueNamespace},
-			queueNamespace,
-			queueID,
-		).Result()
-		if err != nil {
-			return err
+	res, err := q.dequeueScript.Run(ctx, q.client, []string{queueNamespace},
+		queueNamespace,
+		queueID,
+	).Result()
+	if err != nil {
+		if errors.Is(err, redis.Nil) {
+			return nil
 		}
+		return err
+	}
+	queue := opt.Queue
+	if queue == nil {
+		queue = q.RedisQueue
+	}
+	jobm := res.([]interface{})
+	jobs := make([]*work.Job, len(jobm))
+	queueIDs := make([]string, len(jobm))
+	for i, iface := range jobm {
 		var sqJob sidekiqJob
-		err = json.NewDecoder(strings.NewReader(res.(string))).Decode(&sqJob)
+		err := json.NewDecoder(strings.NewReader(iface.(string))).Decode(&sqJob)
 		if err != nil {
 			return err
 		}
@@ -141,48 +151,62 @@ func (q *sidekiqQueue) Pull(opt *PullOptions) error {
 		if err != nil {
 			return err
 		}
-		queue := opt.Queue
-		if queue == nil {
-			queue = q.RedisQueue
+		jobs[i] = job
+		queueIDs[i] = FormatQueueID(sqJob.Queue, sqJob.Class)
+	}
+	found := make([]*work.Job, len(jobs))
+	if finder, ok := queue.(work.BulkJobFinder); ok {
+		jobIDs := make([]string, len(jobs))
+		for i, job := range jobs {
+			jobIDs[i] = job.ID
+		}
+		// best effort to check for duplicates
+		foundJobs, err := finder.BulkFind(jobIDs, &work.FindOptions{
+			Namespace: opt.Namespace,
+		})
+		if err != nil {
+			return err
 		}
-		var found bool
-		if finder, ok := queue.(work.BulkJobFinder); ok {
-			// best effort to check for duplicates
-			jobs, err := finder.BulkFind([]string{job.ID}, &work.FindOptions{
+		found = foundJobs
+	}
+	if bulkEnqueuer, ok := queue.(work.BulkEnqueuer); ok {
+		m := make(map[string][]*work.Job)
+		for i, job := range jobs {
+			if found[i] != nil {
+				continue
+			}
+			queueID := queueIDs[i]
+			m[queueID] = append(m[queueID], job)
+		}
+		for queueID, jobs := range m {
+			err := bulkEnqueuer.BulkEnqueue(jobs, &work.EnqueueOptions{
 				Namespace: opt.Namespace,
+				QueueID:   queueID,
 			})
 			if err != nil {
 				return err
 			}
-			found = len(jobs) == 1 && jobs[0] != nil
 		}
-		if !found {
+	} else {
+		for i, job := range jobs {
+			if found[i] != nil {
+				continue
+			}
 			err := queue.Enqueue(job, &work.EnqueueOptions{
 				Namespace: opt.Namespace,
-				QueueID:   FormatQueueID(sqJob.Queue, sqJob.Class),
+				QueueID:   queueIDs[i],
 			})
 			if err != nil {
 				return err
 			}
 		}
-		err = q.ackScript.Run(ctx, q.client, []string{queueNamespace},
-			queueNamespace,
-			queueID,
-			res.(string),
-		).Err()
-		if err != nil {
-			return err
-		}
-		return nil
 	}
-
-	for {
-		err := pull()
-		if err != nil {
-			if errors.Is(err, redis.Nil) {
-				return nil
-			}
-			return err
-		}
+	err = q.ackScript.Run(ctx, q.client, []string{queueNamespace},
+		queueNamespace,
+		queueID,
+	).Err()
+	if err != nil {
+		return err
 	}
+	return nil
 }

sidekiq/pull_test.go

@@ -520,3 +520,43 @@ func TestPullDequeueHeartbeat(t *testing.T) {
 	require.Equal(t, "123", z[0].Member)
 	require.EqualValues(t, now.Unix()+100, z[0].Score)
 }
+
+func TestPullDequeueAck(t *testing.T) {
+	client := redistest.NewClient()
+	defer client.Close()
+	require.NoError(t, redistest.Reset(client))
+
+	jobIDs := []string{"83b27ea26dd65821239ca6aa", "ebb3186ec09c42642f980a20", "7e59de95478191698aa69b22"}
+	for _, jobID := range jobIDs {
+		err := client.LPush(context.Background(), pullTestSidekiqQueueKey, fmt.Sprintf(`{"class":"TestWorker","args":[],"retry":3,"queue":"default","backtrace":true,"jid":%q,"created_at":1567788641.0875323,"enqueued_at":1567788642.0879307,"retry_count":2,"error_message":"error: test","error_class":"StandardError","failed_at":1567791043,"retried_at":1567791046}"`, jobID)).Err()
+		require.NoError(t, err)
+	}
+
+	q := NewQueue(client)
+
+	res, err := q.(*sidekiqQueue).dequeueScript.Run(context.Background(), client, []string{pullTestNamespace},
+		pullTestSidekiqNamespace,
+		fmt.Sprintf("queue:%s", pullTestSidekiqQueue),
+	).Result()
+	require.NoError(t, err)
+
+	jobm := res.([]interface{})
+	require.Len(t, jobm, 3)
+	require.Contains(t, jobm[0], "7e59de95478191698aa69b22")
+	require.Contains(t, jobm[1], "ebb3186ec09c42642f980a20")
+	require.Contains(t, jobm[2], "83b27ea26dd65821239ca6aa")
+
+	count, err := client.LLen(context.Background(), pullTestSidekiqQueueKey).Result()
+	require.NoError(t, err)
+	require.Equal(t, int64(3), count)
+
+	err = q.(*sidekiqQueue).ackScript.Run(context.Background(), client, []string{pullTestNamespace},
+		pullTestSidekiqNamespace,
+		fmt.Sprintf("queue:%s", pullTestSidekiqQueue),
+	).Err()
+	require.NoError(t, err)
+
+	count, err = client.LLen(context.Background(), pullTestSidekiqQueueKey).Result()
+	require.NoError(t, err)
+	require.Equal(t, int64(0), count)
+}

sidekiq/queue.go

@@ -98,16 +98,15 @@ func NewQueue(client redis.UniversalClient) Queue {
 	local queue_id = ARGV[2]
 
 	local queue_key = table.concat({queue_ns, queue_id}, ":")
-	return redis.call("lindex", queue_key, -1)
+	return redis.call("lrange", queue_key, 0, -1)
 	`)
 
 	ackScript := redis.NewScript(`
 	local queue_ns = ARGV[1]
 	local queue_id = ARGV[2]
-	local jobm = ARGV[3]
 
 	local queue_key = table.concat({queue_ns, queue_id}, ":")
-	return redis.call("lrem", queue_key, -1, jobm)
+	return redis.call("del", queue_key)
 	`)
 
 	dequeueStartScript := redis.NewScript(`
@@ -156,7 +155,6 @@ func NewQueue(client redis.UniversalClient) Queue {
 	local puller_queue_key = table.concat({queue_ns, queue_id}, ":")
 	local pullers_key = table.concat({queue_ns, "pullers"}, ":")
 	if redis.call("llen", puller_queue_key) == 0 then
-		redis.call("del", puller_queue_key)
 		return redis.call("zrem", pullers_key, queue_id)
 	end
 	return 0