chore(queue): support custom run func (#603)
Signed-off-by: Bo-Yi Wu <appleboy.tw@gmail.com>
This commit is contained in:
parent
380162a38a
commit
ed25b0f42a
4
main.go
4
main.go
|
@ -320,7 +320,9 @@ func main() {
|
||||||
var w queue.Worker
|
var w queue.Worker
|
||||||
switch core.Queue(cfg.Queue.Engine) {
|
switch core.Queue(cfg.Queue.Engine) {
|
||||||
case core.LocalQueue:
|
case core.LocalQueue:
|
||||||
w = simple.NewWorker(simple.WithQueueNum(int(cfg.Core.QueueNum)))
|
w = simple.NewWorker(
|
||||||
|
simple.WithQueueNum(int(cfg.Core.QueueNum)),
|
||||||
|
)
|
||||||
case core.NSQ:
|
case core.NSQ:
|
||||||
w = nsq.NewWorker()
|
w = nsq.NewWorker()
|
||||||
default:
|
default:
|
||||||
|
|
|
@ -17,7 +17,8 @@ var errMaxCapacity = errors.New("max capacity reached")
|
||||||
|
|
||||||
// Worker for simple queue using channel
|
// Worker for simple queue using channel
|
||||||
type Worker struct {
|
type Worker struct {
|
||||||
queueNotification chan queue.QueuedMessage
|
QueueNotification chan queue.QueuedMessage
|
||||||
|
runFunc func(*Worker) error
|
||||||
}
|
}
|
||||||
|
|
||||||
// BeforeRun run script before start worker
|
// BeforeRun run script before start worker
|
||||||
|
@ -32,33 +33,29 @@ func (s *Worker) AfterRun() error {
|
||||||
|
|
||||||
// Run start the worker
|
// Run start the worker
|
||||||
func (s *Worker) Run(_ chan struct{}) error {
|
func (s *Worker) Run(_ chan struct{}) error {
|
||||||
for notification := range s.queueNotification {
|
return s.runFunc(s)
|
||||||
gorush.SendNotification(notification)
|
|
||||||
}
|
|
||||||
|
|
||||||
return nil
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Shutdown worker
|
// Shutdown worker
|
||||||
func (s *Worker) Shutdown() error {
|
func (s *Worker) Shutdown() error {
|
||||||
close(s.queueNotification)
|
close(s.QueueNotification)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Capacity for channel
|
// Capacity for channel
|
||||||
func (s *Worker) Capacity() int {
|
func (s *Worker) Capacity() int {
|
||||||
return cap(s.queueNotification)
|
return cap(s.QueueNotification)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Usage for count of channel usage
|
// Usage for count of channel usage
|
||||||
func (s *Worker) Usage() int {
|
func (s *Worker) Usage() int {
|
||||||
return len(s.queueNotification)
|
return len(s.QueueNotification)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Queue send notification to queue
|
// Queue send notification to queue
|
||||||
func (s *Worker) Queue(job queue.QueuedMessage) error {
|
func (s *Worker) Queue(job queue.QueuedMessage) error {
|
||||||
select {
|
select {
|
||||||
case s.queueNotification <- job:
|
case s.QueueNotification <- job:
|
||||||
return nil
|
return nil
|
||||||
default:
|
default:
|
||||||
return errMaxCapacity
|
return errMaxCapacity
|
||||||
|
@ -68,14 +65,27 @@ func (s *Worker) Queue(job queue.QueuedMessage) error {
|
||||||
// WithQueueNum setup the capcity of queue
|
// WithQueueNum setup the capcity of queue
|
||||||
func WithQueueNum(num int) Option {
|
func WithQueueNum(num int) Option {
|
||||||
return func(w *Worker) {
|
return func(w *Worker) {
|
||||||
w.queueNotification = make(chan queue.QueuedMessage, num)
|
w.QueueNotification = make(chan queue.QueuedMessage, num)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithRunFunc setup the run func of queue
|
||||||
|
func WithRunFunc(fn func(w *Worker) error) Option {
|
||||||
|
return func(w *Worker) {
|
||||||
|
w.runFunc = fn
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewWorker for struc
|
// NewWorker for struc
|
||||||
func NewWorker(opts ...Option) *Worker {
|
func NewWorker(opts ...Option) *Worker {
|
||||||
w := &Worker{
|
w := &Worker{
|
||||||
queueNotification: make(chan queue.QueuedMessage, runtime.NumCPU()<<1),
|
QueueNotification: make(chan queue.QueuedMessage, runtime.NumCPU()<<1),
|
||||||
|
runFunc: func(w *Worker) error {
|
||||||
|
for notification := range w.QueueNotification {
|
||||||
|
gorush.SendNotification(notification)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
// Loop through each option
|
// Loop through each option
|
||||||
|
|
Loading…
Reference in New Issue