TWEAK: limit script args with small batch (#119)

Author: taylorchu

redis_queue.go

@@ -10,6 +10,19 @@ import (
 	"github.com/redis/go-redis/v9"
 )
 
+func batchSlice(n int) [][]int {
+	const size = 1000
+	var batches [][]int
+	for i := 0; i < n; i += size {
+		j := i + size
+		if j > n {
+			j = n
+		}
+		batches = append(batches, []int{i, j})
+	}
+	return batches
+}
+
 type redisQueue struct {
 	client redis.UniversalClient
 
@@ -170,6 +183,16 @@ func (q *redisQueue) Enqueue(job *Job, opt *EnqueueOptions) error {
 }
 
 func (q *redisQueue) BulkEnqueue(jobs []*Job, opt *EnqueueOptions) error {
+	for _, batch := range batchSlice(len(jobs)) {
+		err := q.bulkEnqueueSmallBatch(jobs[batch[0]:batch[1]], opt)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (q *redisQueue) bulkEnqueueSmallBatch(jobs []*Job, opt *EnqueueOptions) error {
 	err := opt.Validate()
 	if err != nil {
 		return err
@@ -201,6 +224,18 @@ func (q *redisQueue) Dequeue(opt *DequeueOptions) (*Job, error) {
 }
 
 func (q *redisQueue) BulkDequeue(count int64, opt *DequeueOptions) ([]*Job, error) {
+	var jobs []*Job
+	for _, batch := range batchSlice(int(count)) {
+		foundJobs, err := q.bulkDequeueSmallBatch(int64(batch[1]-batch[0]), opt)
+		if err != nil {
+			return nil, err
+		}
+		jobs = append(jobs, foundJobs...)
+	}
+	return jobs, nil
+}
+
+func (q *redisQueue) bulkDequeueSmallBatch(count int64, opt *DequeueOptions) ([]*Job, error) {
 	err := opt.Validate()
 	if err != nil {
 		return nil, err
@@ -236,6 +271,16 @@ func (q *redisQueue) Ack(job *Job, opt *AckOptions) error {
 }
 
 func (q *redisQueue) BulkAck(jobs []*Job, opt *AckOptions) error {
+	for _, batch := range batchSlice(len(jobs)) {
+		err := q.bulkAckSmallBatch(jobs[batch[0]:batch[1]], opt)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (q *redisQueue) bulkAckSmallBatch(jobs []*Job, opt *AckOptions) error {
 	err := opt.Validate()
 	if err != nil {
 		return err
@@ -253,6 +298,18 @@ func (q *redisQueue) BulkAck(jobs []*Job, opt *AckOptions) error {
 }
 
 func (q *redisQueue) BulkFind(jobIDs []string, opt *FindOptions) ([]*Job, error) {
+	var jobs []*Job
+	for _, batch := range batchSlice(len(jobIDs)) {
+		foundJobs, err := q.bulkFindSmallBatch(jobIDs[batch[0]:batch[1]], opt)
+		if err != nil {
+			return nil, err
+		}
+		jobs = append(jobs, foundJobs...)
+	}
+	return jobs, nil
+}
+
+func (q *redisQueue) bulkFindSmallBatch(jobIDs []string, opt *FindOptions) ([]*Job, error) {
 	err := opt.Validate()
 	if err != nil {
 		return nil, err

redis_queue_test.go

@@ -327,3 +327,27 @@ func TestRedisQueueGetQueueMetrics(t *testing.T) {
 	require.EqualValues(t, 1, m.ScheduledTotal)
 	require.True(t, m.Latency == 0)
 }
+
+func TestRedisQueueBulkEnqueue(t *testing.T) {
+	client := redistest.NewClient()
+	defer client.Close()
+	require.NoError(t, redistest.Reset(client))
+	q := NewRedisQueue(client)
+
+	const jobCount = 100000
+	var jobs []*Job
+	for i := 0; i < jobCount; i++ {
+		job := NewJob()
+		jobs = append(jobs, job)
+	}
+
+	err := q.BulkEnqueue(jobs, &EnqueueOptions{
+		Namespace: "{ns1}",
+		QueueID:   "q1",
+	})
+	require.NoError(t, err)
+
+	count, err := client.ZCard(context.Background(), "{ns1}:queue:q1").Result()
+	require.NoError(t, err)
+	require.Equal(t, int64(jobCount), count)
+}

sidekiq/external_queue.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
-	"time"
 
 	"github.com/taylorchu/work"
 )
@@ -82,22 +81,3 @@ func newSidekiqJob(job *work.Job, sqQueue, sqClass string) (*sidekiqJob, error)
 	}
 	return &sqJob, nil
 }
-
-func newJob(sqJob *sidekiqJob) (*work.Job, error) {
-	updatedAt := sqJob.CreatedAt
-	for _, ts := range []float64{sqJob.FailedAt, sqJob.RetriedAt} {
-		if ts > updatedAt {
-			updatedAt = ts
-		}
-	}
-	job := work.Job{
-		ID:         sqJob.ID,
-		Payload:    sqJob.Args,
-		CreatedAt:  time.Unix(int64(sqJob.CreatedAt), 0),
-		UpdatedAt:  time.Unix(int64(updatedAt), 0),
-		EnqueuedAt: time.Unix(int64(sqJob.EnqueuedAt), 0),
-		Retries:    sqJob.RetryCount,
-		LastError:  sqJob.ErrorMessage,
-	}
-	return &job, nil
-}