github.com/grafana/pyroscope@v1.18.0/pkg/querier/worker/util.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  
     3  package worker
     4  
     5  import (
     6  	"context"
     7  	"time"
     8  
     9  	"github.com/go-kit/log"
    10  	"github.com/go-kit/log/level"
    11  	"go.uber.org/atomic"
    12  )
    13  
    14  // newExecutionContext returns a new execution context (execCtx) that wraps the input workerCtx and
    15  // it used to run the querier's worker loop and execute queries.
    16  // The purpose of the execution context is to gracefully shutdown queriers, waiting
    17  // until inflight queries are terminated before the querier process exits.
    18  //
    19  // The caller must call execCancel() once done.
    20  //
    21  // How it's used:
    22  //
    23  // - The querier worker's loop run in a dedicated context, called the "execution context".
    24  //
    25  // - The execution context is canceled when the worker context gets cancelled (ie. querier is shutting down)
    26  // and there's no inflight query execution. In case there's an inflight query, the execution context is canceled
    27  // once the inflight query terminates and the response has been sent.
    28  func newExecutionContext(workerCtx context.Context, logger log.Logger) (execCtx context.Context, execCancel context.CancelFunc, inflightQuery *atomic.Bool) {
    29  	execCtx, execCancel = context.WithCancel(context.Background())
    30  	inflightQuery = atomic.NewBool(false)
    31  
    32  	go func() {
    33  		// Wait until it's safe to cancel the execution context, which is when one of the following conditions happen:
    34  		// - The worker context has been canceled and there's no inflight query
    35  		// - The execution context itself has been explicitly canceled
    36  		select {
    37  		case <-workerCtx.Done():
    38  			level.Debug(logger).Log("msg", "querier worker context has been canceled, waiting until there's no inflight query")
    39  
    40  			for inflightQuery.Load() {
    41  				select {
    42  				case <-execCtx.Done():
    43  					// In the meanwhile, the execution context has been explicitly canceled, so we should just terminate.
    44  					return
    45  				case <-time.After(100 * time.Millisecond):
    46  					// Going to check it again.
    47  				}
    48  			}
    49  
    50  			level.Debug(logger).Log("msg", "querier worker context has been canceled and there's no inflight query, canceling the execution context too")
    51  			execCancel()
    52  		case <-execCtx.Done():
    53  			level.Debug(logger).Log("msg", "querier worker context has been canceled and there's no inflight query, canceling the execution context too")
    54  			// Nothing to do. The execution context has been explicitly canceled.
    55  		}
    56  	}()
    57  
    58  	return
    59  }