From 857b1a1060cd5112b1ee93bfb2c7f6a7a7309ab5 Mon Sep 17 00:00:00 2001 From: zeripath Date: Fri, 30 Dec 2022 00:06:47 +0000 Subject: [PATCH] Correctly handle select on multiple channels in Queues (#22146) 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 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 028023d50032c..3a375fdb4124c 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 d1bf7239eb945..8458e8c52e9cf 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 bdf04a363b79f..9403bfddeca65 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 <-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 data := <-p.dataChan: + 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