chore(queue): Support task function (#28)

* chore(queue): Support task function

Signed-off-by: Bo-Yi Wu <[email protected]>

* chore: update

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

queue.go

@@ -1,6 +1,7 @@
 package queue
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"runtime"
@@ -12,6 +13,9 @@ import (
 // ErrQueueShutdown close the queue.
 var ErrQueueShutdown = errors.New("queue has been closed")
 
+// TaskFunc is the task function
+type TaskFunc func(context.Context) error
+
 type (
 	// A Queue is a message queue.
 	Queue struct {
@@ -28,6 +32,7 @@ type (
 
 	// Job with Timeout
 	Job struct {
+		Task    TaskFunc      `json:"-"`
 		Timeout time.Duration `json:"timeout"`
 		Body    []byte        `json:"body"`
 	}
@@ -157,6 +162,27 @@ func (q *Queue) QueueWithTimeout(timeout time.Duration, job QueuedMessage) error
 	return q.handleQueue(timeout, job)
 }
 
+func (q *Queue) handleQueueTask(timeout time.Duration, task TaskFunc) error {
+	if atomic.LoadInt32(&q.stopFlag) == 1 {
+		return ErrQueueShutdown
+	}
+
+	return q.worker.Queue(Job{
+		Task:    task,
+		Timeout: timeout,
+	})
+}
+
+// QueueTask to queue job task
+func (q *Queue) QueueTask(task TaskFunc) error {
+	return q.handleQueueTask(q.timeout, task)
+}
+
+// QueueTaskWithTimeout to queue job task with timeout
+func (q *Queue) QueueTaskWithTimeout(timeout time.Duration, task TaskFunc) error {
+	return q.handleQueueTask(timeout, task)
+}
+
 func (q *Queue) work() {
 	if atomic.LoadInt32(&q.stopFlag) == 1 {
 		return

queue_test.go

@@ -1,6 +1,7 @@
 package queue
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -33,7 +34,7 @@ func TestNewQueue(t *testing.T) {
 }
 
 func TestWorkerNum(t *testing.T) {
-	w := &queueWorker{
+	w := &messageWorker{
 		messages: make(chan QueuedMessage, 100),
 	}
 	q, err := NewQueue(
@@ -52,7 +53,7 @@ func TestWorkerNum(t *testing.T) {
 }
 
 func TestShtdonwOnce(t *testing.T) {
-	w := &queueWorker{
+	w := &messageWorker{
 		messages: make(chan QueuedMessage, 100),
 	}
 	q, err := NewQueue(
@@ -76,7 +77,7 @@ func TestWorkerStatus(t *testing.T) {
 	m := mockMessage{
 		message: "foobar",
 	}
-	w := &queueWorker{
+	w := &messageWorker{
 		messages: make(chan QueuedMessage, 100),
 	}
 	q, err := NewQueue(
@@ -99,7 +100,7 @@ func TestWorkerStatus(t *testing.T) {
 }
 
 func TestWorkerPanic(t *testing.T) {
-	w := &queueWorker{
+	w := &messageWorker{
 		messages: make(chan QueuedMessage, 10),
 	}
 	q, err := NewQueue(
@@ -127,7 +128,7 @@ func TestWorkerPanic(t *testing.T) {
 }
 
 func TestCapacityReached(t *testing.T) {
-	w := &queueWorker{
+	w := &messageWorker{
 		messages: make(chan QueuedMessage, 1),
 	}
 	q, err := NewQueue(
@@ -148,7 +149,7 @@ func TestCapacityReached(t *testing.T) {
 }
 
 func TestCloseQueueAfterShutdown(t *testing.T) {
-	w := &queueWorker{
+	w := &messageWorker{
 		messages: make(chan QueuedMessage, 10),
 	}
 	q, err := NewQueue(
@@ -174,3 +175,37 @@ func TestCloseQueueAfterShutdown(t *testing.T) {
 	assert.Error(t, err)
 	assert.Equal(t, ErrQueueShutdown, err)
 }
+
+func TestQueueTaskJob(t *testing.T) {
+	w := &taskWorker{
+		messages: make(chan QueuedMessage, 10),
+	}
+	q, err := NewQueue(
+		WithWorker(w),
+		WithWorkerCount(5),
+		WithLogger(NewLogger()),
+	)
+	assert.NoError(t, err)
+	assert.NotNil(t, q)
+	q.Start()
+	assert.NoError(t, q.QueueTask(func(ctx context.Context) error {
+		time.Sleep(120 * time.Millisecond)
+		q.logger.Info("Add new task 1")
+		return nil
+	}))
+	assert.NoError(t, q.QueueTask(func(ctx context.Context) error {
+		time.Sleep(100 * time.Millisecond)
+		q.logger.Info("Add new task 2")
+		return nil
+	}))
+	assert.NoError(t, q.QueueTaskWithTimeout(50*time.Millisecond, func(ctx context.Context) error {
+		time.Sleep(80 * time.Millisecond)
+		return nil
+	}))
+	time.Sleep(50 * time.Millisecond)
+	q.Shutdown()
+	assert.Equal(t, ErrQueueShutdown, q.QueueTask(func(ctx context.Context) error {
+		return nil
+	}))
+	q.Wait()
+}

simple/simple.go

@@ -58,7 +58,11 @@ func (s *Worker) handle(job queue.Job) error {
 		}()
 
 		// run custom process function
-		done <- s.runFunc(ctx, job)
+		if job.Task != nil {
+			done <- job.Task(ctx)
+		} else {
+			done <- s.runFunc(ctx, job)
+		}
 	}()
 
 	select {
@@ -97,6 +101,11 @@ func (s *Worker) Run() error {
 	for task := range s.taskQueue {
 		var data queue.Job
 		_ = json.Unmarshal(task.Bytes(), &data)
+		if v, ok := task.(queue.Job); ok {
+			if v.Task != nil {
+				data.Task = v.Task
+			}
+		}
 		if err := s.handle(data); err != nil {
 			s.logger.Error(err.Error())
 		}

simple/simple_test.go

@@ -327,3 +327,43 @@ func TestJobComplete(t *testing.T) {
 	assert.Error(t, err)
 	assert.Equal(t, errors.New("job completed"), err)
 }
+
+func TestTaskJobComplete(t *testing.T) {
+	job := queue.Job{
+		Timeout: 100 * time.Millisecond,
+		Task: func(ctx context.Context) error {
+			return errors.New("job completed")
+		},
+	}
+	w := NewWorker()
+
+	err := w.handle(job)
+	assert.Error(t, err)
+	assert.Equal(t, errors.New("job completed"), err)
+
+	job = queue.Job{
+		Timeout: 250 * time.Millisecond,
+		Task: func(ctx context.Context) error {
+			return nil
+		},
+	}
+
+	w = NewWorker()
+	done := make(chan error)
+	go func() {
+		done <- w.handle(job)
+	}()
+	assert.NoError(t, w.Shutdown())
+	err = <-done
+	assert.NoError(t, err)
+
+	// job timeout
+	job = queue.Job{
+		Timeout: 50 * time.Millisecond,
+		Task: func(ctx context.Context) error {
+			time.Sleep(60 * time.Millisecond)
+			return nil
+		},
+	}
+	assert.Equal(t, context.DeadlineExceeded, w.handle(job))
+}

worker.go

@@ -1,6 +1,7 @@
 package queue
 
 import (
+	"context"
 	"errors"
 	"time"
 )
@@ -30,7 +31,7 @@ type QueuedMessage interface {
 
 var (
 	_ Worker = (*emptyWorker)(nil)
-	_ Worker = (*queueWorker)(nil)
+	_ Worker = (*messageWorker)(nil)
 )
 
 type emptyWorker struct{}
@@ -43,13 +44,13 @@ func (w *emptyWorker) Queue(job QueuedMessage) error { return nil }
 func (w *emptyWorker) Capacity() int                 { return 0 }
 func (w *emptyWorker) Usage() int                    { return 0 }
 
-type queueWorker struct {
+type messageWorker struct {
 	messages chan QueuedMessage
 }
 
-func (w *queueWorker) BeforeRun() error { return nil }
-func (w *queueWorker) AfterRun() error  { return nil }
-func (w *queueWorker) Run() error {
+func (w *messageWorker) BeforeRun() error { return nil }
+func (w *messageWorker) AfterRun() error  { return nil }
+func (w *messageWorker) Run() error {
 	for msg := range w.messages {
 		if string(msg.Bytes()) == "panic" {
 			panic("show panic")
@@ -59,18 +60,51 @@ func (w *queueWorker) Run() error {
 	return nil
 }
 
-func (w *queueWorker) Shutdown() error {
+func (w *messageWorker) Shutdown() error {
 	close(w.messages)
 	return nil
 }
 
-func (w *queueWorker) Queue(job QueuedMessage) error {
+func (w *messageWorker) Queue(job QueuedMessage) error {
 	select {
 	case w.messages <- job:
 		return nil
 	default:
 		return errors.New("max capacity reached")
 	}
 }
-func (w *queueWorker) Capacity() int { return cap(w.messages) }
-func (w *queueWorker) Usage() int    { return len(w.messages) }
+func (w *messageWorker) Capacity() int { return cap(w.messages) }
+func (w *messageWorker) Usage() int    { return len(w.messages) }
+
+type taskWorker struct {
+	messages chan QueuedMessage
+}
+
+func (w *taskWorker) BeforeRun() error { return nil }
+func (w *taskWorker) AfterRun() error  { return nil }
+func (w *taskWorker) Run() error {
+	for msg := range w.messages {
+		if v, ok := msg.(Job); ok {
+			if v.Task != nil {
+				_ = v.Task(context.Background())
+			}
+		}
+	}
+	return nil
+}
+
+func (w *taskWorker) Shutdown() error {
+	close(w.messages)
+	return nil
+}
+
+func (w *taskWorker) Queue(job QueuedMessage) error {
+	select {
+	case w.messages <- job:
+		return nil
+	default:
+		return errors.New("max capacity reached")
+	}
+}
+func (w *taskWorker) Capacity() int { return cap(w.messages) }
+func (w *taskWorker) Usage() int    { return len(w.messages) }