Implement Pausable and Pushback for ChannelQueues and ChannelUniqueQueues

zeripath · zeripath · commit a73c034fa83e · 2021-05-19T22:03:53.000+01:00
Signed-off-by: Andrew Thornton &lt;art27@cantab.net&gt;
diff --git a/modules/queue/queue_channel.go b/modules/queue/queue_channel.go
@@ -7,6 +7,8 @@ package queue
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
+	"time"
 
 	"code.gitea.io/gitea/modules/log"
 )
@@ -51,7 +53,6 @@ func NewChannelQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue, erro
 	shutdownCtx, shutdownCtxCancel := context.WithCancel(terminateCtx)
 
 	queue := &ChannelQueue{
-		WorkerPool:         NewWorkerPool(handle, config.WorkerPoolConfiguration),
 		shutdownCtx:        shutdownCtx,
 		shutdownCtxCancel:  shutdownCtxCancel,
 		terminateCtx:       terminateCtx,
@@ -60,6 +61,23 @@ func NewChannelQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue, erro
 		workers:            config.Workers,
 		name:               config.Name,
 	}
+	queue.WorkerPool = NewWorkerPool(func(data ...Data) []Data {
+		unhandled := handle(data...)
+		if len(unhandled) > 0 {
+			// We can only pushback to the channel if we're paused.
+			if queue.IsPaused() {
+				atomic.AddInt64(&queue.numInQueue, int64(len(unhandled)))
+				go func() {
+					for _, datum := range data {
+						queue.dataChan <- datum
+					}
+				}()
+				return nil
+			}
+		}
+		return unhandled
+	}, config.WorkerPoolConfiguration)
+
 	queue.qid = GetManager().Add(queue, ChannelQueueType, config, exemplar)
 	return queue, nil
 }
@@ -81,6 +99,52 @@ func (q *ChannelQueue) Push(data Data) error {
 	return nil
 }
 
+// Flush flushes the channel with a timeout - the Flush worker will be registered as a flush worker with the manager
+func (q *ChannelQueue) Flush(timeout time.Duration) error {
+	q.lock.Lock()
+	paused := q.paused
+	q.lock.Unlock()
+	if paused {
+		return nil
+	}
+	ctx, cancel := q.commonRegisterWorkers(1, timeout, true)
+	defer cancel()
+	return q.FlushWithContext(ctx)
+}
+
+// FlushWithContext is very similar to CleanUp but it will return as soon as the dataChan is empty
+func (q *ChannelQueue) FlushWithContext(ctx context.Context) error {
+	log.Trace("ChannelQueue: %d Flush", q.qid)
+	for {
+		q.lock.Lock()
+		paused := q.paused
+		q.lock.Unlock()
+		if paused {
+			return nil
+		}
+		select {
+		case data := <-q.dataChan:
+			if q.IsPaused() {
+				// we're paused so we should push this back and stop
+				// (whilst handle will check this too we need to stop the flusher for this to work.)
+				go func() {
+					q.dataChan <- data
+				}()
+				return nil
+			} else if unhandled := q.handle(data); unhandled != nil {
+				log.Error("Unhandled Data whilst flushing queue %d", q.qid)
+			}
+			atomic.AddInt64(&q.numInQueue, -1)
+		case <-q.baseCtx.Done():
+			return q.baseCtx.Err()
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			return nil
+		}
+	}
+}
+
 // Shutdown processing from this queue
 func (q *ChannelQueue) Shutdown() {
 	q.lock.Lock()
@@ -94,6 +158,7 @@ func (q *ChannelQueue) Shutdown() {
 	log.Trace("ChannelQueue: %s Shutting down", q.name)
 	go func() {
 		log.Trace("ChannelQueue: %s Flushing", q.name)
+		// We can't use Cleanup here because that will close the channel
 		if err := q.FlushWithContext(q.terminateCtx); err != nil {
 			log.Warn("ChannelQueue: %s Terminated before completed flushing", q.name)
 			return
diff --git a/modules/queue/queue_disk_channel.go b/modules/queue/queue_disk_channel.go
@@ -51,7 +51,20 @@ func NewPersistableChannelQueue(handle HandlerFunc, cfg, exemplar interface{}) (
 	}
 	config := configInterface.(PersistableChannelQueueConfiguration)
 
-	channelQueue, err := NewChannelQueue(handle, ChannelQueueConfiguration{
+	queue := &PersistableChannelQueue{
+		closed: make(chan struct{}),
+	}
+
+	wrappedHandle := func(data ...Data) (failed []Data) {
+		for _, unhandled := range handle(data...) {
+			if fail := queue.PushBack(unhandled); fail != nil {
+				failed = append(failed, fail)
+			}
+		}
+		return
+	}
+
+	channelQueue, err := NewChannelQueue(wrappedHandle, ChannelQueueConfiguration{
 		WorkerPoolConfiguration: WorkerPoolConfiguration{
 			QueueLength:  config.QueueLength,
 			BatchLength:  config.BatchLength,
@@ -84,15 +97,12 @@ func NewPersistableChannelQueue(handle HandlerFunc, cfg, exemplar interface{}) (
 		DataDir: config.DataDir,
 	}
 
-	levelQueue, err := NewLevelQueue(handle, levelCfg, exemplar)
+	levelQueue, err := NewLevelQueue(wrappedHandle, levelCfg, exemplar)
 	if err == nil {
-		queue := &PersistableChannelQueue{
-			channelQueue: channelQueue.(*ChannelQueue),
-			delayedStarter: delayedStarter{
-				internal: levelQueue.(*LevelQueue),
-				name:     config.Name,
-			},
-			closed: make(chan struct{}),
+		queue.channelQueue = channelQueue.(*ChannelQueue)
+		queue.delayedStarter = delayedStarter{
+			internal: levelQueue.(*LevelQueue),
+			name:     config.Name,
 		}
 		_ = GetManager().Add(queue, PersistableChannelQueueType, config, exemplar)
 		return queue, nil
@@ -102,16 +112,13 @@ func NewPersistableChannelQueue(handle HandlerFunc, cfg, exemplar interface{}) (
 		return nil, ErrInvalidConfiguration{cfg: cfg}
 	}
 
-	queue := &PersistableChannelQueue{
-		channelQueue: channelQueue.(*ChannelQueue),
-		delayedStarter: delayedStarter{
-			cfg:         levelCfg,
-			underlying:  LevelQueueType,
-			timeout:     config.Timeout,
-			maxAttempts: config.MaxAttempts,
-			name:        config.Name,
-		},
-		closed: make(chan struct{}),
+	queue.channelQueue = channelQueue.(*ChannelQueue)
+	queue.delayedStarter = delayedStarter{
+		cfg:         levelCfg,
+		underlying:  LevelQueueType,
+		timeout:     config.Timeout,
+		maxAttempts: config.MaxAttempts,
+		name:        config.Name,
 	}
 	_ = GetManager().Add(queue, PersistableChannelQueueType, config, exemplar)
 	return queue, nil
@@ -132,6 +139,19 @@ func (q *PersistableChannelQueue) Push(data Data) error {
 	}
 }
 
+// PushBack will push the indexer data to queue
+func (q *PersistableChannelQueue) PushBack(data Data) error {
+	select {
+	case <-q.closed:
+		if pbr, ok := q.internal.(PushBackable); ok {
+			return pbr.PushBack(data)
+		}
+		return q.internal.Push(data)
+	default:
+		return q.channelQueue.Push(data)
+	}
+}
+
 // Run starts to run the queue
 func (q *PersistableChannelQueue) Run(atShutdown, atTerminate func(func())) {
 	log.Debug("PersistableChannelQueue: %s Starting", q.delayedStarter.name)
diff --git a/modules/queue/unique_queue_channel.go b/modules/queue/unique_queue_channel.go
@@ -8,6 +8,8 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"sync/atomic"
+	"time"
 
 	"code.gitea.io/gitea/modules/log"
 )
@@ -69,7 +71,16 @@ func NewChannelUniqueQueue(handle HandlerFunc, cfg, exemplar interface{}) (Queue
 			delete(queue.table, datum)
 			queue.lock.Unlock()
 			if u := handle(datum); u != nil {
-				unhandled = append(unhandled, u...)
+				if queue.IsPaused() {
+					// We can only pushback to the channel if we're paused.
+					go func() {
+						if err := queue.Push(u); err != nil {
+							log.Error("Unable to push back to queue %d", queue.qid)
+						}
+					}()
+				} else {
+					unhandled = append(unhandled, u...)
+				}
 			}
 		}
 		return unhandled
@@ -131,6 +142,52 @@ func (q *ChannelUniqueQueue) Has(data Data) (bool, error) {
 	return has, nil
 }
 
+// Flush flushes the channel with a timeout - the Flush worker will be registered as a flush worker with the manager
+func (q *ChannelUniqueQueue) Flush(timeout time.Duration) error {
+	q.lock.Lock()
+	paused := q.paused
+	q.lock.Unlock()
+	if paused {
+		return nil
+	}
+	ctx, cancel := q.commonRegisterWorkers(1, timeout, true)
+	defer cancel()
+	return q.FlushWithContext(ctx)
+}
+
+// FlushWithContext is very similar to CleanUp but it will return as soon as the dataChan is empty
+func (q *ChannelUniqueQueue) FlushWithContext(ctx context.Context) error {
+	log.Trace("ChannelUniqueQueue: %d Flush", q.qid)
+	for {
+		q.lock.Lock()
+		paused := q.paused
+		q.lock.Unlock()
+		if paused {
+			return nil
+		}
+		select {
+		case data := <-q.dataChan:
+			if q.IsPaused() {
+				// we're paused so we should push this back and stop
+				// (whilst handle will check this too we need to stop the flusher for this to work.)
+				go func() {
+					q.dataChan <- data
+				}()
+				return nil
+			} else if unhandled := q.handle(data); unhandled != nil {
+				log.Error("Unhandled Data whilst flushing queue %d", q.qid)
+			}
+			atomic.AddInt64(&q.numInQueue, -1)
+		case <-q.baseCtx.Done():
+			return q.baseCtx.Err()
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			return nil
+		}
+	}
+}
+
 // Shutdown processing from this queue
 func (q *ChannelUniqueQueue) Shutdown() {
 	log.Trace("ChannelUniqueQueue: %s Shutting down", q.name)
diff --git a/modules/queue/unique_queue_disk_channel.go b/modules/queue/unique_queue_disk_channel.go
@@ -51,7 +51,20 @@ func NewPersistableChannelUniqueQueue(handle HandlerFunc, cfg, exemplar interfac
 	}
 	config := configInterface.(PersistableChannelUniqueQueueConfiguration)
 
-	channelUniqueQueue, err := NewChannelUniqueQueue(handle, ChannelUniqueQueueConfiguration{
+	queue := &PersistableChannelUniqueQueue{
+		closed: make(chan struct{}),
+	}
+
+	wrappedHandle := func(data ...Data) (failed []Data) {
+		for _, unhandled := range handle(data...) {
+			if fail := queue.PushBack(unhandled); fail != nil {
+				failed = append(failed, fail)
+			}
+		}
+		return
+	}
+
+	channelUniqueQueue, err := NewChannelUniqueQueue(wrappedHandle, ChannelUniqueQueueConfiguration{
 		WorkerPoolConfiguration: WorkerPoolConfiguration{
 			QueueLength:  config.QueueLength,
 			BatchLength:  config.BatchLength,
@@ -84,10 +97,7 @@ func NewPersistableChannelUniqueQueue(handle HandlerFunc, cfg, exemplar interfac
 		DataDir: config.DataDir,
 	}
 
-	queue := &PersistableChannelUniqueQueue{
-		channelQueue: channelUniqueQueue.(*ChannelUniqueQueue),
-		closed:       make(chan struct{}),
-	}
+	queue.channelQueue = channelUniqueQueue.(*ChannelUniqueQueue)
 
 	levelQueue, err := NewLevelUniqueQueue(func(data ...Data) []Data {
 		for _, datum := range data {
@@ -143,6 +153,19 @@ func (q *PersistableChannelUniqueQueue) PushFunc(data Data, fn func() error) err
 	}
 }
 
+// PushBack will push the indexer data to queue
+func (q *PersistableChannelUniqueQueue) PushBack(data Data) error {
+	select {
+	case <-q.closed:
+		if pbr, ok := q.internal.(PushBackable); ok {
+			return pbr.PushBack(data)
+		}
+		return q.internal.Push(data)
+	default:
+		return q.channelQueue.Push(data)
+	}
+}
+
 // Has will test if the queue has the data
 func (q *PersistableChannelUniqueQueue) Has(data Data) (bool, error) {
 	// This is more difficult...
