Add IsEmpty and begin just making the queues composed WorkerPools

Author: zeripath

modules/queue/queue.go

@@ -59,6 +59,7 @@ type Queue interface {
 	Run(atShutdown, atTerminate func(context.Context, func()))
 	Push(Data) error
 	Flush(time.Duration) error
+	IsEmpty() bool
 }
 
 // DummyQueueType is the type for the dummy queue
@@ -86,6 +87,11 @@ func (b *DummyQueue) Flush(time.Duration) error {
 	return nil
 }
 
+// IsEmpty asserts that the queue is empty
+func (b *DummyQueue) IsEmpty() bool {
+	return true
+}
+
 var queuesMap = map[Type]NewQueueFunc{DummyQueueType: NewDummyQueue}
 
 // RegisteredTypes provides the list of requested types of queues

modules/queue/queue_channel.go

@@ -7,7 +7,6 @@ package queue
 import (
 	"context"
 	"fmt"
-	"time"
 
 	"code.gitea.io/gitea/modules/log"
 )
@@ -27,7 +26,7 @@ type ChannelQueueConfiguration struct {
 // A channel queue is not persistable and does not shutdown or terminate cleanly
 // It is basically a very thin wrapper around a WorkerPool
 type ChannelQueue struct {
-	pool     *WorkerPool
+	*WorkerPool
 	exemplar interface{}
 	workers  int
 	name     string
@@ -44,12 +43,12 @@ func NewChannelQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue, erro
 		config.BatchLength = 1
 	}
 	queue := &ChannelQueue{
-		pool:     NewWorkerPool(handle, config.WorkerPoolConfiguration),
-		exemplar: exemplar,
-		workers:  config.Workers,
-		name:     config.Name,
+		WorkerPool: NewWorkerPool(handle, config.WorkerPoolConfiguration),
+		exemplar:   exemplar,
+		workers:    config.Workers,
+		name:       config.Name,
 	}
-	queue.pool.qid = GetManager().Add(config.Name, ChannelQueueType, config, exemplar, queue.pool)
+	queue.qid = GetManager().Add(config.Name, ChannelQueueType, config, exemplar, queue)
 	return queue, nil
 }
 
@@ -63,7 +62,7 @@ func (c *ChannelQueue) Run(atShutdown, atTerminate func(context.Context, func())
 	})
 	log.Debug("ChannelQueue: %s Starting", c.name)
 	go func() {
-		_ = c.pool.AddWorkers(c.workers, 0)
+		_ = c.AddWorkers(c.workers, 0)
 	}()
 }
 
@@ -72,7 +71,7 @@ func (c *ChannelQueue) Push(data Data) error {
 	if !assignableTo(data, c.exemplar) {
 		return fmt.Errorf("Unable to assign data: %v to same type as exemplar: %v in queue: %s", data, c.exemplar, c.name)
 	}
-	c.pool.Push(data)
+	c.WorkerPool.Push(data)
 	return nil
 }
 
@@ -81,11 +80,6 @@ func (c *ChannelQueue) Name() string {
 	return c.name
 }
 
-// Flush flushes the queue and blocks till the queue is empty
-func (c *ChannelQueue) Flush(timeout time.Duration) error {
-	return c.pool.Flush(timeout)
-}
-
 func init() {
 	queuesMap[ChannelQueueType] = NewChannelQueue
 }

modules/queue/queue_disk.go

@@ -145,6 +145,14 @@ func (l *LevelQueue) Flush(timeout time.Duration) error {
 	return l.pool.Flush(timeout)
 }
 
+// IsEmpty checks whether the queue is empty
+func (l *LevelQueue) IsEmpty() bool {
+	if !l.pool.IsEmpty() {
+		return false
+	}
+	return l.queue.Len() == 0
+}
+
 // Shutdown this queue and stop processing
 func (l *LevelQueue) Shutdown() {
 	l.lock.Lock()

modules/queue/queue_disk_channel.go

@@ -134,7 +134,7 @@ func (p *PersistableChannelQueue) Run(atShutdown, atTerminate func(context.Conte
 
 	p.lock.Lock()
 	if p.internal == nil {
-		err := p.setInternal(atShutdown, p.ChannelQueue.pool.handle, p.exemplar)
+		err := p.setInternal(atShutdown, p.ChannelQueue.handle, p.exemplar)
 		p.lock.Unlock()
 		if err != nil {
 			log.Fatal("Unable to create internal queue for %s Error: %v", p.Name(), err)
@@ -150,21 +150,21 @@ func (p *PersistableChannelQueue) Run(atShutdown, atTerminate func(context.Conte
 	go p.internal.Run(func(_ context.Context, _ func()) {}, func(_ context.Context, _ func()) {})
 
 	go func() {
-		_ = p.ChannelQueue.pool.AddWorkers(p.workers, 0)
+		_ = p.ChannelQueue.AddWorkers(p.workers, 0)
 	}()
 
 	log.Trace("PersistableChannelQueue: %s Waiting til closed", p.delayedStarter.name)
 	<-p.closed
 	log.Trace("PersistableChannelQueue: %s Cancelling pools", p.delayedStarter.name)
-	p.ChannelQueue.pool.cancel()
+	p.ChannelQueue.cancel()
 	p.internal.(*LevelQueue).pool.cancel()
 	log.Trace("PersistableChannelQueue: %s Waiting til done", p.delayedStarter.name)
-	p.ChannelQueue.pool.Wait()
+	p.ChannelQueue.Wait()
 	p.internal.(*LevelQueue).pool.Wait()
 	// Redirect all remaining data in the chan to the internal channel
 	go func() {
 		log.Trace("PersistableChannelQueue: %s Redirecting remaining data", p.delayedStarter.name)
-		for data := range p.ChannelQueue.pool.dataChan {
+		for data := range p.ChannelQueue.dataChan {
 			_ = p.internal.Push(data)
 		}
 		log.Trace("PersistableChannelQueue: %s Done Redirecting remaining data", p.delayedStarter.name)
@@ -174,7 +174,20 @@ func (p *PersistableChannelQueue) Run(atShutdown, atTerminate func(context.Conte
 
 // Flush flushes the queue and blocks till the queue is empty
 func (p *PersistableChannelQueue) Flush(timeout time.Duration) error {
-	return p.ChannelQueue.pool.Flush(timeout)
+	return p.ChannelQueue.Flush(timeout)
+}
+
+// IsEmpty checks if a queue is empty
+func (p *PersistableChannelQueue) IsEmpty() bool {
+	if !p.ChannelQueue.IsEmpty() {
+		return false
+	}
+	p.lock.Lock()
+	defer p.lock.Unlock()
+	if p.internal == nil {
+		return false
+	}
+	return p.internal.IsEmpty()
 }
 
 // Shutdown processing this queue

modules/queue/queue_redis.go

@@ -24,6 +24,7 @@ const RedisQueueType Type = "redis"
 type redisClient interface {
 	RPush(key string, args ...interface{}) *redis.IntCmd
 	LPop(key string) *redis.StringCmd
+	LLen(key string) *redis.IntCmd
 	Ping() *redis.StatusCmd
 	Close() error
 }
@@ -168,6 +169,19 @@ func (r *RedisQueue) Flush(timeout time.Duration) error {
 	return r.pool.Flush(timeout)
 }
 
+// IsEmpty checks if the queue is empty
+func (r *RedisQueue) IsEmpty() bool {
+	if !r.pool.IsEmpty() {
+		return false
+	}
+	length, err := r.client.LLen(r.queueName).Result()
+	if err != nil {
+		log.Error("Error whilst getting queue length for %s: Error: %v", r.name, err)
+		return false
+	}
+	return length == 0
+}
+
 // Shutdown processing from this queue
 func (r *RedisQueue) Shutdown() {
 	log.Trace("RedisQueue: %s Shutting down", r.name)

modules/queue/queue_wrapped.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"code.gitea.io/gitea/modules/log"
@@ -87,10 +88,11 @@ func (q *delayedStarter) setInternal(atShutdown func(context.Context, func()), h
 // WrappedQueue wraps a delayed starting queue
 type WrappedQueue struct {
 	delayedStarter
-	lock     sync.Mutex
-	handle   HandlerFunc
-	exemplar interface{}
-	channel  chan Data
+	lock       sync.Mutex
+	handle     HandlerFunc
+	exemplar   interface{}
+	channel    chan Data
+	numInQueue int64
 }
 
 // NewWrappedQueue will attempt to create a queue of the provided type,
@@ -140,6 +142,7 @@ func (q *WrappedQueue) Push(data Data) error {
 	if !assignableTo(data, q.exemplar) {
 		return fmt.Errorf("unable to assign data: %v to same type as exemplar: %v in %s", data, q.exemplar, q.name)
 	}
+	atomic.AddInt64(&q.numInQueue, 1)
 	q.channel <- data
 	return nil
 }
@@ -181,6 +184,20 @@ func (q *WrappedQueue) Flush(timeout time.Duration) error {
 	}
 }
 
+// IsEmpty checks whether the queue is empty
+func (q *WrappedQueue) IsEmpty() bool {
+	val := atomic.LoadInt64(&q.numInQueue)
+	if val != 0 {
+		return false
+	}
+	q.lock.Lock()
+	defer q.lock.Unlock()
+	if q.internal == nil {
+		return false
+	}
+	return q.internal.IsEmpty()
+}
+
 // Run starts to run the queue and attempts to create the internal queue
 func (q *WrappedQueue) Run(atShutdown, atTerminate func(context.Context, func())) {
 	log.Debug("WrappedQueue: %s Starting", q.name)
@@ -195,6 +212,7 @@ func (q *WrappedQueue) Run(atShutdown, atTerminate func(context.Context, func())
 		go func() {
 			for data := range q.channel {
 				_ = q.internal.Push(data)
+				atomic.AddInt64(&q.numInQueue, -1)
 			}
 		}()
 	} else {

modules/queue/workerpool.go

@@ -7,6 +7,7 @@ package queue
 import (
 	"context"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"code.gitea.io/gitea/modules/log"
@@ -30,6 +31,7 @@ type WorkerPool struct {
 	blockTimeout       time.Duration
 	boostTimeout       time.Duration
 	boostWorkers       int
+	numInQueue         int64
 }
 
 // WorkerPoolConfiguration is the basic configuration for a WorkerPool
@@ -65,6 +67,7 @@ func NewWorkerPool(handle HandlerFunc, config WorkerPoolConfiguration) *WorkerPo
 // Push pushes the data to the internal channel
 func (p *WorkerPool) Push(data Data) {
 	p.lock.Lock()
+	atomic.AddInt64(&p.numInQueue, 1)
 	if p.blockTimeout > 0 && p.boostTimeout > 0 && (p.numberOfWorkers <= p.maxNumberOfWorkers || p.maxNumberOfWorkers < 0) {
 		p.lock.Unlock()
 		p.pushBoost(data)
@@ -273,6 +276,7 @@ func (p *WorkerPool) CleanUp(ctx context.Context) {
 	close(p.dataChan)
 	for data := range p.dataChan {
 		p.handle(data)
+		atomic.AddInt64(&p.numInQueue, -1)
 		select {
 		case <-ctx.Done():
 			log.Warn("WorkerPool: %d Cleanup context closed before finishing clean-up", p.qid)
@@ -290,6 +294,11 @@ func (p *WorkerPool) Flush(timeout time.Duration) error {
 	return p.FlushWithContext(ctx)
 }
 
+// IsEmpty returns if true if the worker queue is empty
+func (p *WorkerPool) IsEmpty() bool {
+	return atomic.LoadInt64(&p.numInQueue) == 0
+}
+
 // FlushWithContext is very similar to CleanUp but it will return as soon as the dataChan is empty
 // NB: The worker will not be registered with the manager.
 func (p *WorkerPool) FlushWithContext(ctx context.Context) error {
@@ -298,6 +307,7 @@ func (p *WorkerPool) FlushWithContext(ctx context.Context) error {
 		select {
 		case data := <-p.dataChan:
 			p.handle(data)
+			atomic.AddInt64(&p.numInQueue, -1)
 		case <-p.baseCtx.Done():
 			return p.baseCtx.Err()
 		case <-ctx.Done():
@@ -317,6 +327,7 @@ func (p *WorkerPool) doWork(ctx context.Context) {
 			if len(data) > 0 {
 				log.Trace("Handling: %d data, %v", len(data), data)
 				p.handle(data...)
+				atomic.AddInt64(&p.numInQueue, -1*int64(len(data)))
 			}
 			log.Trace("Worker shutting down")
 			return
@@ -326,6 +337,7 @@ func (p *WorkerPool) doWork(ctx context.Context) {
 				if len(data) > 0 {
 					log.Trace("Handling: %d data, %v", len(data), data)
 					p.handle(data...)
+					atomic.AddInt64(&p.numInQueue, -1*int64(len(data)))
 				}
 				log.Trace("Worker shutting down")
 				return
@@ -334,6 +346,7 @@ func (p *WorkerPool) doWork(ctx context.Context) {
 			if len(data) >= p.batchLength {
 				log.Trace("Handling: %d data, %v", len(data), data)
 				p.handle(data...)
+				atomic.AddInt64(&p.numInQueue, -1*int64(len(data)))
 				data = make([]Data, 0, p.batchLength)
 			}
 		default:
@@ -349,6 +362,7 @@ func (p *WorkerPool) doWork(ctx context.Context) {
 				if len(data) > 0 {
 					log.Trace("Handling: %d data, %v", len(data), data)
 					p.handle(data...)
+					atomic.AddInt64(&p.numInQueue, -1*int64(len(data)))
 				}
 				log.Trace("Worker shutting down")
 				return
@@ -364,6 +378,7 @@ func (p *WorkerPool) doWork(ctx context.Context) {
 					if len(data) > 0 {
 						log.Trace("Handling: %d data, %v", len(data), data)
 						p.handle(data...)
+						atomic.AddInt64(&p.numInQueue, -1*int64(len(data)))
 					}
 					log.Trace("Worker shutting down")
 					return
@@ -372,13 +387,15 @@ func (p *WorkerPool) doWork(ctx context.Context) {
 				if len(data) >= p.batchLength {
 					log.Trace("Handling: %d data, %v", len(data), data)
 					p.handle(data...)
+					atomic.AddInt64(&p.numInQueue, -1*int64(len(data)))
 					data = make([]Data, 0, p.batchLength)
 				}
 			case <-timer.C:
 				delay = time.Millisecond * 100
 				if len(data) > 0 {
 					log.Trace("Handling: %d data, %v", len(data), data)
 					p.handle(data...)
+					atomic.AddInt64(&p.numInQueue, -1*int64(len(data)))
 					data = make([]Data, 0, p.batchLength)
 				}