You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
240 lines
4.8 KiB
Go
240 lines
4.8 KiB
Go
package queue
|
|
|
|
import (
|
|
"runtime"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"zero/core/logx"
|
|
"zero/core/rescue"
|
|
"zero/core/stat"
|
|
"zero/core/threading"
|
|
"zero/core/timex"
|
|
)
|
|
|
|
const queueName = "queue"
|
|
|
|
type (
|
|
Queue struct {
|
|
name string
|
|
metrics *stat.Metrics
|
|
producerFactory ProducerFactory
|
|
producerRoutineGroup *threading.RoutineGroup
|
|
consumerFactory ConsumerFactory
|
|
consumerRoutineGroup *threading.RoutineGroup
|
|
producerCount int
|
|
consumerCount int
|
|
active int32
|
|
channel chan string
|
|
quit chan struct{}
|
|
listeners []QueueListener
|
|
eventLock sync.Mutex
|
|
eventChannels []chan interface{}
|
|
}
|
|
|
|
QueueListener interface {
|
|
OnPause()
|
|
OnResume()
|
|
}
|
|
|
|
QueuePoller interface {
|
|
Name() string
|
|
Poll() string
|
|
}
|
|
|
|
QueuePusher interface {
|
|
Name() string
|
|
Push(string) error
|
|
}
|
|
)
|
|
|
|
func NewQueue(producerFactory ProducerFactory, consumerFactory ConsumerFactory) *Queue {
|
|
queue := &Queue{
|
|
metrics: stat.NewMetrics(queueName),
|
|
producerFactory: producerFactory,
|
|
producerRoutineGroup: threading.NewRoutineGroup(),
|
|
consumerFactory: consumerFactory,
|
|
consumerRoutineGroup: threading.NewRoutineGroup(),
|
|
producerCount: runtime.NumCPU(),
|
|
consumerCount: runtime.NumCPU() << 1,
|
|
channel: make(chan string),
|
|
quit: make(chan struct{}),
|
|
}
|
|
queue.SetName(queueName)
|
|
|
|
return queue
|
|
}
|
|
|
|
func (queue *Queue) AddListener(listener QueueListener) {
|
|
queue.listeners = append(queue.listeners, listener)
|
|
}
|
|
|
|
func (queue *Queue) Broadcast(message interface{}) {
|
|
go func() {
|
|
queue.eventLock.Lock()
|
|
defer queue.eventLock.Unlock()
|
|
|
|
for _, channel := range queue.eventChannels {
|
|
channel <- message
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (queue *Queue) SetName(name string) {
|
|
queue.name = name
|
|
queue.metrics.SetName(name)
|
|
}
|
|
|
|
func (queue *Queue) SetNumConsumer(count int) {
|
|
queue.consumerCount = count
|
|
}
|
|
|
|
func (queue *Queue) SetNumProducer(count int) {
|
|
queue.producerCount = count
|
|
}
|
|
|
|
func (queue *Queue) Start() {
|
|
queue.startProducers(queue.producerCount)
|
|
queue.startConsumers(queue.consumerCount)
|
|
|
|
queue.producerRoutineGroup.Wait()
|
|
close(queue.channel)
|
|
queue.consumerRoutineGroup.Wait()
|
|
}
|
|
|
|
func (queue *Queue) Stop() {
|
|
close(queue.quit)
|
|
}
|
|
|
|
func (queue *Queue) consume(eventChan chan interface{}) {
|
|
var consumer Consumer
|
|
|
|
for {
|
|
var err error
|
|
if consumer, err = queue.consumerFactory(); err != nil {
|
|
logx.Errorf("Error on creating consumer: %v", err)
|
|
time.Sleep(time.Second)
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
|
|
for {
|
|
select {
|
|
case message, ok := <-queue.channel:
|
|
if ok {
|
|
queue.consumeOne(consumer, message)
|
|
} else {
|
|
logx.Info("Task channel was closed, quitting consumer...")
|
|
return
|
|
}
|
|
case event := <-eventChan:
|
|
consumer.OnEvent(event)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (queue *Queue) consumeOne(consumer Consumer, message string) {
|
|
threading.RunSafe(func() {
|
|
startTime := timex.Now()
|
|
defer func() {
|
|
duration := timex.Since(startTime)
|
|
queue.metrics.Add(stat.Task{
|
|
Duration: duration,
|
|
})
|
|
logx.WithDuration(duration).Infof("%s", message)
|
|
}()
|
|
|
|
if err := consumer.Consume(message); err != nil {
|
|
logx.Errorf("Error occurred while consuming %v: %v", message, err)
|
|
}
|
|
})
|
|
}
|
|
|
|
func (queue *Queue) pause() {
|
|
for _, listener := range queue.listeners {
|
|
listener.OnPause()
|
|
}
|
|
}
|
|
|
|
func (queue *Queue) produce() {
|
|
var producer Producer
|
|
|
|
for {
|
|
var err error
|
|
if producer, err = queue.producerFactory(); err != nil {
|
|
logx.Errorf("Error on creating producer: %v", err)
|
|
time.Sleep(time.Second)
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
|
|
atomic.AddInt32(&queue.active, 1)
|
|
producer.AddListener(routineListener{
|
|
queue: queue,
|
|
})
|
|
|
|
for {
|
|
select {
|
|
case <-queue.quit:
|
|
logx.Info("Quitting producer")
|
|
return
|
|
default:
|
|
if v, ok := queue.produceOne(producer); ok {
|
|
queue.channel <- v
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (queue *Queue) produceOne(producer Producer) (string, bool) {
|
|
// avoid panic quit the producer, just log it and continue
|
|
defer rescue.Recover()
|
|
|
|
return producer.Produce()
|
|
}
|
|
|
|
func (queue *Queue) resume() {
|
|
for _, listener := range queue.listeners {
|
|
listener.OnResume()
|
|
}
|
|
}
|
|
|
|
func (queue *Queue) startConsumers(number int) {
|
|
for i := 0; i < number; i++ {
|
|
eventChan := make(chan interface{})
|
|
queue.eventLock.Lock()
|
|
queue.eventChannels = append(queue.eventChannels, eventChan)
|
|
queue.eventLock.Unlock()
|
|
queue.consumerRoutineGroup.Run(func() {
|
|
queue.consume(eventChan)
|
|
})
|
|
}
|
|
}
|
|
|
|
func (queue *Queue) startProducers(number int) {
|
|
for i := 0; i < number; i++ {
|
|
queue.producerRoutineGroup.Run(func() {
|
|
queue.produce()
|
|
})
|
|
}
|
|
}
|
|
|
|
type routineListener struct {
|
|
queue *Queue
|
|
}
|
|
|
|
func (rl routineListener) OnProducerPause() {
|
|
if atomic.AddInt32(&rl.queue.active, -1) <= 0 {
|
|
rl.queue.pause()
|
|
}
|
|
}
|
|
|
|
func (rl routineListener) OnProducerResume() {
|
|
if atomic.AddInt32(&rl.queue.active, 1) == 1 {
|
|
rl.queue.resume()
|
|
}
|
|
}
|