131 lines
2.3 KiB
Go
131 lines
2.3 KiB
Go
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)
|
|
}
|
|
}
|