github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/worker/processor_manager.go (about)

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