forgejo/modules/queue/workergroup.go

350 lines
10 KiB
Go

// Copyright 2023 The Gitea Authors. All rights reserved.
// SPDX-License-Identifier: MIT
package queue
import (
"context"
"runtime/pprof"
"sync"
"sync/atomic"
"time"
"code.gitea.io/gitea/modules/log"
)
var (
infiniteTimerC = make(chan time.Time)
batchDebounceDuration = 100 * time.Millisecond
workerIdleDuration = 1 * time.Second
shutdownDefaultTimeout = 2 * time.Second
unhandledItemRequeueDuration atomic.Int64 // to avoid data race during test
)
func init() {
unhandledItemRequeueDuration.Store(int64(5 * time.Second))
}
// workerGroup is a group of workers to work with a WorkerPoolQueue
type workerGroup[T any] struct {
q *WorkerPoolQueue[T]
wg sync.WaitGroup
ctxWorker context.Context
ctxWorkerCancel context.CancelFunc
batchBuffer []T
popItemChan chan []byte
popItemErr chan error
}
func (wg *workerGroup[T]) doPrepareWorkerContext() {
wg.ctxWorker, wg.ctxWorkerCancel = context.WithCancel(wg.q.ctxRun)
}
// doDispatchBatchToWorker dispatches a batch of items to worker's channel.
// If the channel is full, it tries to start a new worker if possible.
func (q *WorkerPoolQueue[T]) doDispatchBatchToWorker(wg *workerGroup[T], flushChan chan flushType) {
batch := wg.batchBuffer
wg.batchBuffer = nil
if len(batch) == 0 {
return
}
full := false
select {
case q.batchChan <- batch:
default:
full = true
}
// TODO: the logic could be improved in the future, to avoid a data-race between "doStartNewWorker" and "workerNum"
// The root problem is that if we skip "doStartNewWorker" here, the "workerNum" might be decreased by other workers later
// So ideally, it should check whether there are enough workers by some approaches, and start new workers if necessary.
q.workerNumMu.Lock()
noWorker := q.workerNum == 0
if full || noWorker {
if q.workerNum < q.workerMaxNum || noWorker && q.workerMaxNum <= 0 {
q.workerNum++
q.doStartNewWorker(wg)
}
}
q.workerNumMu.Unlock()
if full {
select {
case q.batchChan <- batch:
case flush := <-flushChan:
q.doWorkerHandle(batch)
q.doFlush(wg, flush)
case <-q.ctxRun.Done():
wg.batchBuffer = batch // return the batch to buffer, the "doRun" function will handle it
}
}
}
// doWorkerHandle calls the safeHandler to handle a batch of items, and it increases/decreases the active worker number.
// If the context has been canceled, it should not be caller because the "Push" still needs the context, in such case, call q.safeHandler directly
func (q *WorkerPoolQueue[T]) doWorkerHandle(batch []T) {
q.workerNumMu.Lock()
q.workerActiveNum++
q.workerNumMu.Unlock()
defer func() {
q.workerNumMu.Lock()
q.workerActiveNum--
q.workerNumMu.Unlock()
}()
unhandled := q.safeHandler(batch...)
// if none of the items were handled, it should back-off for a few seconds
// in this case the handler (eg: document indexer) may have encountered some errors/failures
if len(unhandled) == len(batch) && unhandledItemRequeueDuration.Load() != 0 {
log.Error("Queue %q failed to handle batch of %d items, backoff for a few seconds", q.GetName(), len(batch))
select {
case <-q.ctxRun.Done():
case <-time.After(time.Duration(unhandledItemRequeueDuration.Load())):
}
}
for _, item := range unhandled {
if err := q.Push(item); err != nil {
if !q.basePushForShutdown(item) {
log.Error("Failed to requeue item for queue %q when calling handler: %v", q.GetName(), err)
}
}
}
}
// basePushForShutdown tries to requeue items into the base queue when the WorkerPoolQueue is shutting down.
// If the queue is shutting down, it returns true and try to push the items
// Otherwise it does nothing and returns false
func (q *WorkerPoolQueue[T]) basePushForShutdown(items ...T) bool {
shutdownTimeout := time.Duration(q.shutdownTimeout.Load())
if shutdownTimeout == 0 {
return false
}
ctxShutdown, ctxShutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout)
defer ctxShutdownCancel()
for _, item := range items {
// if there is still any error, the queue can do nothing instead of losing the items
if err := q.baseQueue.PushItem(ctxShutdown, q.marshal(item)); err != nil {
log.Error("Failed to requeue item for queue %q when shutting down: %v", q.GetName(), err)
}
}
return true
}
// doStartNewWorker starts a new worker for the queue, the worker reads from worker's channel and handles the items.
func (q *WorkerPoolQueue[T]) doStartNewWorker(wp *workerGroup[T]) {
wp.wg.Add(1)
go func() {
defer wp.wg.Done()
log.Debug("Queue %q starts new worker", q.GetName())
defer log.Debug("Queue %q stops idle worker", q.GetName())
t := time.NewTicker(workerIdleDuration)
defer t.Stop()
keepWorking := true
stopWorking := func() {
q.workerNumMu.Lock()
keepWorking = false
q.workerNum--
q.workerNumMu.Unlock()
}
for keepWorking {
select {
case <-wp.ctxWorker.Done():
stopWorking()
case batch, ok := <-q.batchChan:
if !ok {
stopWorking()
continue
}
q.doWorkerHandle(batch)
// reset the idle ticker, and drain the tick after reset in case a tick is already triggered
t.Reset(workerIdleDuration)
select {
case <-t.C:
default:
}
case <-t.C:
q.workerNumMu.Lock()
keepWorking = q.workerNum <= 1 // keep the last worker running
if !keepWorking {
q.workerNum--
}
q.workerNumMu.Unlock()
}
}
}()
}
// doFlush flushes the queue: it tries to read all items from the queue and handles them.
// It is for testing purpose only. It's not designed to work for a cluster.
func (q *WorkerPoolQueue[T]) doFlush(wg *workerGroup[T], flush flushType) {
log.Debug("Queue %q starts flushing", q.GetName())
defer log.Debug("Queue %q finishes flushing", q.GetName())
// stop all workers, and prepare a new worker context to start new workers
wg.ctxWorkerCancel()
wg.wg.Wait()
defer func() {
close(flush)
wg.doPrepareWorkerContext()
}()
// drain the batch channel first
loop:
for {
select {
case batch := <-q.batchChan:
q.doWorkerHandle(batch)
default:
break loop
}
}
// drain the popItem channel
emptyCounter := 0
for {
select {
case data, dataOk := <-wg.popItemChan:
if !dataOk {
return
}
emptyCounter = 0
if v, jsonOk := q.unmarshal(data); !jsonOk {
continue
} else {
q.doWorkerHandle([]T{v})
}
case err := <-wg.popItemErr:
if !q.isCtxRunCanceled() {
log.Error("Failed to pop item from queue %q (doFlush): %v", q.GetName(), err)
}
return
case <-q.ctxRun.Done():
log.Debug("Queue %q is shutting down", q.GetName())
return
case <-time.After(20 * time.Millisecond):
// There is no reliable way to make sure all queue items are consumed by the Flush, there always might be some items stored in some buffers/temp variables.
// If we run Gitea in a cluster, we can even not guarantee all items are consumed in a deterministic instance.
// Luckily, the "Flush" trick is only used in tests, so far so good.
if cnt, _ := q.baseQueue.Len(q.ctxRun); cnt == 0 && len(wg.popItemChan) == 0 {
emptyCounter++
}
if emptyCounter >= 2 {
return
}
}
}
}
func (q *WorkerPoolQueue[T]) isCtxRunCanceled() bool {
select {
case <-q.ctxRun.Done():
return true
default:
return false
}
}
var skipFlushChan = make(chan flushType) // an empty flush chan, used to skip reading other flush requests
// doRun is the main loop of the queue. All related "doXxx" functions are executed in its context.
func (q *WorkerPoolQueue[T]) doRun() {
pprof.SetGoroutineLabels(q.ctxRun)
log.Debug("Queue %q starts running", q.GetName())
defer log.Debug("Queue %q stops running", q.GetName())
wg := &workerGroup[T]{q: q}
wg.doPrepareWorkerContext()
wg.popItemChan, wg.popItemErr = popItemByChan(q.ctxRun, q.baseQueue.PopItem)
defer func() {
q.ctxRunCancel()
// drain all data on the fly
// since the queue is shutting down, the items can't be dispatched to workers because the context is canceled
// it can't call doWorkerHandle either, because there is no chance to push unhandled items back to the queue
var unhandled []T
close(q.batchChan)
for batch := range q.batchChan {
unhandled = append(unhandled, batch...)
}
unhandled = append(unhandled, wg.batchBuffer...)
for data := range wg.popItemChan {
if v, ok := q.unmarshal(data); ok {
unhandled = append(unhandled, v)
}
}
shutdownTimeout := time.Duration(q.shutdownTimeout.Load())
if shutdownTimeout != 0 {
// if there is a shutdown context, try to push the items back to the base queue
q.basePushForShutdown(unhandled...)
workerDone := make(chan struct{})
// the only way to wait for the workers, because the handlers do not have context to wait for
go func() { wg.wg.Wait(); close(workerDone) }()
select {
case <-workerDone:
case <-time.After(shutdownTimeout):
log.Error("Queue %q is shutting down, but workers are still running after timeout", q.GetName())
}
} else {
// if there is no shutdown context, just call the handler to try to handle the items. if the handler fails again, the items are lost
q.safeHandler(unhandled...)
}
close(q.shutdownDone)
}()
var batchDispatchC <-chan time.Time = infiniteTimerC
for {
select {
case data, dataOk := <-wg.popItemChan:
if !dataOk {
return
}
if v, jsonOk := q.unmarshal(data); !jsonOk {
testRecorder.Record("pop:corrupted:%s", data) // in rare cases the levelqueue(leveldb) might be corrupted
continue
} else {
wg.batchBuffer = append(wg.batchBuffer, v)
}
if len(wg.batchBuffer) >= q.batchLength {
q.doDispatchBatchToWorker(wg, q.flushChan)
} else if batchDispatchC == infiniteTimerC {
batchDispatchC = time.After(batchDebounceDuration)
} // else: batchDispatchC is already a debounce timer, it will be triggered soon
case <-batchDispatchC:
batchDispatchC = infiniteTimerC
q.doDispatchBatchToWorker(wg, q.flushChan)
case flush := <-q.flushChan:
// before flushing, it needs to try to dispatch the batch to worker first, in case there is no worker running
// after the flushing, there is at least one worker running, so "doFlush" could wait for workers to finish
// since we are already in a "flush" operation, so the dispatching function shouldn't read the flush chan.
q.doDispatchBatchToWorker(wg, skipFlushChan)
q.doFlush(wg, flush)
case err, errOk := <-wg.popItemErr:
if !errOk {
return
}
if !q.isCtxRunCanceled() {
log.Error("Failed to pop item from queue %q (doRun): %v", q.GetName(), err)
}
return
case <-q.ctxRun.Done():
log.Debug("Queue %q is shutting down", q.GetName())
return
}
}
}