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

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/querier/worker/processor_manager.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package worker
     7  
     8  import (
     9  	"context"
    10  	"sync"
    11  	"time"
    12  
    13  	"go.uber.org/atomic"
    14  	"google.golang.org/grpc"
    15  )
    16  
    17  const (
    18  	notifyShutdownTimeout = 5 * time.Second
    19  )
    20  
    21  // Manages processor goroutines for single grpc connection.
    22  type processorManager struct {
    23  	p       processor
    24  	conn    *grpc.ClientConn
    25  	address string
    26  
    27  	// Main context to control all goroutines.
    28  	ctx context.Context
    29  	wg  sync.WaitGroup
    30  
    31  	// Cancel functions for individual goroutines.
    32  	cancelsMu sync.Mutex
    33  	cancels   []context.CancelFunc
    34  
    35  	currentProcessors *atomic.Int32
    36  }
    37  
    38  func newProcessorManager(ctx context.Context, p processor, conn *grpc.ClientConn, address string) *processorManager {
    39  	return &processorManager{
    40  		p:                 p,
    41  		ctx:               ctx,
    42  		conn:              conn,
    43  		address:           address,
    44  		currentProcessors: atomic.NewInt32(0),
    45  	}
    46  }
    47  
    48  func (pm *processorManager) stop() {
    49  	// Notify the remote query-frontend or query-scheduler we're shutting down.
    50  	// We use a new context to make sure it's not cancelled.
    51  	notifyCtx, cancel := context.WithTimeout(context.Background(), notifyShutdownTimeout)
    52  	defer cancel()
    53  	pm.p.notifyShutdown(notifyCtx, pm.conn, pm.address)
    54  
    55  	// Stop all goroutines.
    56  	pm.concurrency(0)
    57  
    58  	// Wait until they finish.
    59  	pm.wg.Wait()
    60  
    61  	_ = pm.conn.Close()
    62  }
    63  
    64  func (pm *processorManager) concurrency(n int) {
    65  	pm.cancelsMu.Lock()
    66  	defer pm.cancelsMu.Unlock()
    67  
    68  	if n < 0 {
    69  		n = 0
    70  	}
    71  
    72  	for len(pm.cancels) < n {
    73  		ctx, cancel := context.WithCancel(pm.ctx)
    74  		pm.cancels = append(pm.cancels, cancel)
    75  
    76  		pm.wg.Add(1)
    77  		go func() {
    78  			defer pm.wg.Done()
    79  
    80  			pm.currentProcessors.Inc()
    81  			defer pm.currentProcessors.Dec()
    82  
    83  			pm.p.processQueriesOnSingleStream(ctx, pm.conn, pm.address)
    84  		}()
    85  	}
    86  
    87  	for len(pm.cancels) > n {
    88  		pm.cancels[0]()
    89  		pm.cancels = pm.cancels[1:]
    90  	}
    91  }