From 3ad62127df5c2692e9ab62ec15f8a9d529bb1a10 Mon Sep 17 00:00:00 2001 From: zeripath Date: Fri, 13 Jan 2023 20:42:42 +0000 Subject: [PATCH] Correctly handle select on multiple channels in Queues (#22146) (#22428) Backport #22146 There are a few places in FlushQueueWithContext which make an incorrect assumption about how `select` on multiple channels works. The problem is best expressed by looking at the following example: ```go package main import "fmt" func main() { closedChan := make(chan struct{}) close(closedChan) toClose := make(chan struct{}) count := 0 for { select { case <-closedChan: count++ fmt.Println(count) if count == 2 { close(toClose) } case <-toClose: return } } } ``` This PR double-checks that the contexts are closed outside of checking if there is data in the dataChan. It also rationalises the WorkerPool FlushWithContext because the previous implementation failed to handle pausing correctly. This will probably fix the underlying problem in #22145 Fix #22145 Signed-off-by: Andrew Thornton --- modules/queue/queue_channel.go | 26 ---------------- modules/queue/unique_queue_channel.go | 30 ------------------ modules/queue/workerpool.go | 44 ++++++++++++++++++++++++++- 3 files changed, 43 insertions(+), 57 deletions(-) diff --git a/modules/queue/queue_channel.go b/modules/queue/queue_channel.go index 028023d50..3a375fdb4 100644 --- a/modules/queue/queue_channel.go +++ b/modules/queue/queue_channel.go @@ -110,32 +110,6 @@ func (q *ChannelQueue) Flush(timeout time.Duration) error { 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) - paused, _ := q.IsPausedIsResumed() - for { - select { - case <-paused: - return nil - case data, ok := <-q.dataChan: - if !ok { - return nil - } - 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() diff --git a/modules/queue/unique_queue_channel.go b/modules/queue/unique_queue_channel.go index d1bf7239e..8458e8c52 100644 --- a/modules/queue/unique_queue_channel.go +++ b/modules/queue/unique_queue_channel.go @@ -9,7 +9,6 @@ import ( "fmt" "runtime/pprof" "sync" - "sync/atomic" "time" "code.gitea.io/gitea/modules/container" @@ -168,35 +167,6 @@ func (q *ChannelUniqueQueue) Flush(timeout time.Duration) error { 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) - paused, _ := q.IsPausedIsResumed() - for { - select { - case <-paused: - return nil - default: - } - select { - case data, ok := <-q.dataChan: - if !ok { - return nil - } - 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/workerpool.go b/modules/queue/workerpool.go index bdf04a363..9403bfdde 100644 --- a/modules/queue/workerpool.go +++ b/modules/queue/workerpool.go @@ -464,13 +464,43 @@ func (p *WorkerPool) IsEmpty() bool { return atomic.LoadInt64(&p.numInQueue) == 0 } +// contextError returns either ctx.Done(), the base context's error or nil +func (p *WorkerPool) contextError(ctx context.Context) error { + select { + case <-p.baseCtx.Done(): + return p.baseCtx.Err() + case <-ctx.Done(): + return ctx.Err() + default: + return nil + } +} + // 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 { log.Trace("WorkerPool: %d Flush", p.qid) + paused, _ := p.IsPausedIsResumed() for { + // Because select will return any case that is satisified at random we precheck here before looking at dataChan. select { - case data := <-p.dataChan: + case <-paused: + // Ensure that even if paused that the cancelled error is still sent + return p.contextError(ctx) + case <-p.baseCtx.Done(): + return p.baseCtx.Err() + case <-ctx.Done(): + return ctx.Err() + default: + } + + select { + case <-paused: + return p.contextError(ctx) + case data, ok := <-p.dataChan: + if !ok { + return nil + } if unhandled := p.handle(data); unhandled != nil { log.Error("Unhandled Data whilst flushing queue %d", p.qid) } @@ -496,6 +526,7 @@ func (p *WorkerPool) doWork(ctx context.Context) { paused, _ := p.IsPausedIsResumed() data := make([]Data, 0, p.batchLength) for { + // Because select will return any case that is satisified at random we precheck here before looking at dataChan. select { case <-paused: log.Trace("Worker for Queue %d Pausing", p.qid) @@ -516,8 +547,19 @@ func (p *WorkerPool) doWork(ctx context.Context) { log.Trace("Worker shutting down") return } + case <-ctx.Done(): + if len(data) > 0 { + log.Trace("Handling: %d data, %v", len(data), data) + if unhandled := p.handle(data...); unhandled != nil { + log.Error("Unhandled Data in queue %d", p.qid) + } + atomic.AddInt64(&p.numInQueue, -1*int64(len(data))) + } + log.Trace("Worker shutting down") + return default: } + select { case <-paused: // go back around