chore: support using channel in queue (#3)

Author: appleboy

simple/simple.go

@@ -0,0 +1,97 @@
+package simple
+
+import (
+	"errors"
+	"runtime"
+
+	"github.com/appleboy/queue"
+)
+
+var _ queue.Worker = (*Worker)(nil)
+
+// Option for queue system
+type Option func(*Worker)
+
+var errMaxCapacity = errors.New("max capacity reached")
+
+// Worker for simple queue using channel
+type Worker struct {
+	queueNotification chan queue.QueuedMessage
+	runFunc           func(queue.QueuedMessage) error
+}
+
+// BeforeRun run script before start worker
+func (s *Worker) BeforeRun() error {
+	return nil
+}
+
+// AfterRun run script after start worker
+func (s *Worker) AfterRun() error {
+	return nil
+}
+
+// Run start the worker
+func (s *Worker) Run(_ chan struct{}) error {
+	for notification := range s.queueNotification {
+		_ = s.runFunc(notification)
+	}
+	return nil
+}
+
+// Shutdown worker
+func (s *Worker) Shutdown() error {
+	close(s.queueNotification)
+	return nil
+}
+
+// Capacity for channel
+func (s *Worker) Capacity() int {
+	return cap(s.queueNotification)
+}
+
+// Usage for count of channel usage
+func (s *Worker) Usage() int {
+	return len(s.queueNotification)
+}
+
+// Queue send notification to queue
+func (s *Worker) Queue(job queue.QueuedMessage) error {
+	select {
+	case s.queueNotification <- job:
+		return nil
+	default:
+		return errMaxCapacity
+	}
+}
+
+// WithQueueNum setup the capcity of queue
+func WithQueueNum(num int) Option {
+	return func(w *Worker) {
+		w.queueNotification = make(chan queue.QueuedMessage, num)
+	}
+}
+
+// WithRunFunc setup the run func of queue
+func WithRunFunc(fn func(queue.QueuedMessage) error) Option {
+	return func(w *Worker) {
+		w.runFunc = fn
+	}
+}
+
+// NewWorker for struc
+func NewWorker(opts ...Option) *Worker {
+	w := &Worker{
+		queueNotification: make(chan queue.QueuedMessage, runtime.NumCPU()<<1),
+		runFunc: func(msg queue.QueuedMessage) error {
+			return nil
+		},
+	}
+
+	// Loop through each option
+	for _, opt := range opts {
+		// Call the option giving the instantiated
+		opt(w)
+	}
+
+	return w
+}

simple/simple_test.go

@@ -0,0 +1,71 @@
+package simple
+
+import (
+	"runtime"
+	"testing"
+	"time"
+
+	"github.com/appleboy/queue"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type mockMessage struct {
+	msg string
+}
+
+func (m mockMessage) Bytes() []byte {
+	return []byte(m.msg)
+}
+
+func TestQueueUsage(t *testing.T) {
+	w := NewWorker()
+	assert.Equal(t, runtime.NumCPU()<<1, w.Capacity())
+	assert.Equal(t, 0, w.Usage())
+
+	assert.NoError(t, w.Queue(&mockMessage{}))
+	assert.Equal(t, 1, w.Usage())
+}
+
+func TestMaxCapacity(t *testing.T) {
+	w := NewWorker(WithQueueNum(2))
+	assert.Equal(t, 2, w.Capacity())
+	assert.Equal(t, 0, w.Usage())
+
+	assert.NoError(t, w.Queue(&mockMessage{}))
+	assert.Equal(t, 1, w.Usage())
+	assert.NoError(t, w.Queue(&mockMessage{}))
+	assert.Equal(t, 2, w.Usage())
+	assert.Error(t, w.Queue(&mockMessage{}))
+	assert.Equal(t, 2, w.Usage())
+
+	err := w.Queue(&mockMessage{})
+	assert.Equal(t, errMaxCapacity, err)
+}
+
+func TestCustomFuncAndWait(t *testing.T) {
+	m := mockMessage{
+		msg: "foo",
+	}
+	w := NewWorker(
+		WithRunFunc(func(msg queue.QueuedMessage) error {
+			time.Sleep(500 * time.Millisecond)
+			return nil
+		}),
+	)
+	q, err := queue.NewQueue(
+		queue.WithWorker(w),
+		queue.WithWorkerCount(2),
+	)
+	assert.NoError(t, err)
+	q.Start()
+	time.Sleep(100 * time.Millisecond)
+	assert.NoError(t, w.Queue(m))
+	assert.NoError(t, w.Queue(m))
+	assert.NoError(t, w.Queue(m))
+	assert.NoError(t, w.Queue(m))
+	time.Sleep(600 * time.Millisecond)
+	q.Shutdown()
+	q.Wait()
+	// you will see the execute time > 1000ms
+}

worker.go

@@ -2,7 +2,6 @@ package queue
 
 import (
 	"errors"
-	"log"
 	"time"
 )
 
@@ -52,7 +51,6 @@ func (w *queueWorker) BeforeRun() error { return nil }
 func (w *queueWorker) AfterRun() error  { return nil }
 func (w *queueWorker) Run(chan struct{}) error {
 	for msg := range w.messages {
-		log.Println("got message", msg)
 		if string(msg.Bytes()) == "panic" {
 			panic("show panic")
 		}