chore(queue): support option setting (#593)

main.go

@@ -316,7 +316,7 @@ func main() {
 		logx.LogError.Fatal(err)
 	}
 
-	w := simple.NewWorker(int(cfg.Core.QueueNum))
+	w := simple.NewWorker(simple.WithQueueNum(int(cfg.Core.QueueNum)))
 	q := queue.NewQueue(w, int(cfg.Core.WorkerNum))
 	q.Start()
 

queue/simple/simple.go

@@ -2,6 +2,7 @@ package simple
 
 import (
 	"errors"
+	"runtime"
 
 	"github.com/appleboy/gorush/gorush"
 	"github.com/appleboy/gorush/queue"
@@ -9,6 +10,9 @@ import (
 
 var _ queue.Worker = (*Worker)(nil)
 
+// Option for queue system
+type Option func(*Worker)
+
 // Worker for simple queue using channel
 type Worker struct {
 	queueNotification chan gorush.PushNotification
@@ -49,9 +53,24 @@ func (s *Worker) Queue(job interface{}) error {
 	}
 }
 
-// NewWorker for struct
-func NewWorker(num int) *Worker {
-	return &Worker{
-		queueNotification: make(chan gorush.PushNotification, num),
+// WithQueueNum setup the capcity of queue
+func WithQueueNum(num int) Option {
+	return func(w *Worker) {
+		w.queueNotification = make(chan gorush.PushNotification, num)
 	}
 }
+
+// NewWorker for struc
+func NewWorker(opts ...Option) *Worker {
+	w := &Worker{
+		queueNotification: make(chan gorush.PushNotification, runtime.NumCPU()<<1),
+	}
+
+	// Loop through each option
+	for _, opt := range opts {
+		// Call the option giving the instantiated
+		opt(w)
+	}
+
+	return w
+}

router/server_test.go

@@ -46,7 +46,7 @@ func TestMain(m *testing.M) {
 		log.Fatal(err)
 	}
 
-	w = simple.NewWorker(int(cfg.Core.QueueNum))
+	w = simple.NewWorker()
 	q = queue.NewQueue(w, 4)
 	q.Start()
 	defer func() {