123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349 |
- // 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())
-
- atomic.AddInt32(&q.workerStartedCounter, 1) // Only increase counter, used for debugging
-
- 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 := <-wg.popItemErr:
- 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
- }
- }
- }
|