package queue import ( "errors" "runtime" "sync" ) type ( // A Queue is a message queue. Queue struct { logger Logger workerCount int routineGroup *routineGroup quit chan struct{} worker Worker stopOnce sync.Once } ) // Option for queue system type Option func(*Queue) // ErrMissingWorker missing define worker var ErrMissingWorker = errors.New("missing worker module") // WithWorkerCount set worker count func WithWorkerCount(num int) Option { return func(q *Queue) { q.workerCount = num } } // WithLogger set custom logger func WithLogger(l Logger) Option { return func(q *Queue) { q.logger = l } } // WithWorker set custom worker func WithWorker(w Worker) Option { return func(q *Queue) { q.worker = w } } // NewQueue returns a Queue. func NewQueue(opts ...Option) (*Queue, error) { q := &Queue{ workerCount: runtime.NumCPU(), routineGroup: newRoutineGroup(), quit: make(chan struct{}), logger: new(defaultLogger), } // Loop through each option for _, opt := range opts { // Call the option giving the instantiated opt(q) } if q.worker == nil { return nil, ErrMissingWorker } return q, nil } // Capacity for queue max size func (q *Queue) Capacity() int { return q.worker.Capacity() } // Usage for count of queue usage func (q *Queue) Usage() int { return q.worker.Usage() } // Start to enable all worker func (q *Queue) Start() { q.startWorker() } // Shutdown stops all queues. func (q *Queue) Shutdown() { q.stopOnce.Do(func() { q.worker.Shutdown() close(q.quit) }) } // Wait all process func (q *Queue) Wait() { q.routineGroup.Wait() } // Queue to queue all job func (q *Queue) Queue(job QueuedMessage) error { return q.worker.Queue(job) } func (q *Queue) work(num int) { if err := q.worker.BeforeRun(); err != nil { q.logger.Fatal(err) } q.routineGroup.Run(func() { // to handle panic cases from inside the worker // in such case, we start a new goroutine defer func() { if err := recover(); err != nil { q.logger.Error(err) go q.work(num) } }() q.logger.Info("started the worker num ", num) q.worker.Run(q.quit) q.logger.Info("closed the worker num ", num) }) if err := q.worker.AfterRun(); err != nil { q.logger.Fatal(err) } } func (q *Queue) startWorker() { for i := 0; i < q.workerCount; i++ { go q.work(i) } }