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() {