From 4349a170174ddc65ee8466fe1d079090eb12f236 Mon Sep 17 00:00:00 2001
From: Bo-Yi Wu <appleboy.tw@gmail.com>
Date: Sat, 17 Jul 2021 06:56:42 +0800
Subject: [PATCH] chore(queue): support option setting (#593)

---
 main.go                |  2 +-
 queue/simple/simple.go | 27 +++++++++++++++++++++++----
 router/server_test.go  |  2 +-
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/main.go b/main.go
index 15bb7ff..a2c5a94 100644
--- a/main.go
+++ b/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()
 
diff --git a/queue/simple/simple.go b/queue/simple/simple.go
index c86b370..6603ea5 100644
--- a/queue/simple/simple.go
+++ b/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
+}
diff --git a/router/server_test.go b/router/server_test.go
index d4a0e18..4f604ce 100644
--- a/router/server_test.go
+++ b/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() {