github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/remotewrite/queue.go (about)

     1  package remotewrite
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"runtime/debug"
     7  	"sync"
     8  
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/pyroscope-io/pyroscope/pkg/config"
    11  	"github.com/pyroscope-io/pyroscope/pkg/ingestion"
    12  	"github.com/sirupsen/logrus"
    13  )
    14  
    15  // this queue was based off storage/queue.go
    16  // TODO(eh-am): merge with that one
    17  type IngestionQueue struct {
    18  	logger   logrus.FieldLogger
    19  	ingester ingestion.Ingester
    20  
    21  	wg    sync.WaitGroup
    22  	queue chan *ingestion.IngestInput
    23  	stop  chan struct{}
    24  
    25  	metrics *queueMetrics
    26  }
    27  
    28  // NewIngestionQueue creates an IngestionQueue
    29  // Notice how a config.RemoteWriteTarget is taken as argument, even though
    30  // not all fields are used. This is done to simplify the API, as the alternative
    31  // is to take multiple arguments
    32  func NewIngestionQueue(logger logrus.FieldLogger, reg prometheus.Registerer, ingester ingestion.Ingester, targetName string, cfg config.RemoteWriteTarget) *IngestionQueue {
    33  	// Setup defaults
    34  	if cfg.QueueWorkers == 0 {
    35  		// This may be a very conservative value
    36  		// Since it's IO bounded work
    37  		cfg.QueueWorkers = numWorkers()
    38  	}
    39  
    40  	if cfg.QueueSize == 0 {
    41  		cfg.QueueSize = 100
    42  	}
    43  
    44  	q := IngestionQueue{
    45  		logger:   logger,
    46  		ingester: ingester,
    47  		queue:    make(chan *ingestion.IngestInput, cfg.QueueSize),
    48  		stop:     make(chan struct{}),
    49  		metrics:  newQueueMetrics(reg, targetName, cfg.Address),
    50  	}
    51  
    52  	q.wg.Add(cfg.QueueWorkers)
    53  	for i := 0; i < cfg.QueueWorkers; i++ {
    54  		go q.runQueueWorker()
    55  	}
    56  
    57  	q.metrics.mustRegister()
    58  	q.initMetrics(cfg.QueueSize, cfg.QueueWorkers)
    59  
    60  	return &q
    61  }
    62  
    63  func (q *IngestionQueue) initMetrics(queueSize int, queueWorkers int) {
    64  	q.metrics.capacity.Add(float64(queueSize))
    65  	q.metrics.numWorkers.Add(float64(queueWorkers))
    66  }
    67  
    68  func (q *IngestionQueue) Stop() {
    69  	close(q.stop)
    70  	q.wg.Wait()
    71  }
    72  
    73  func (q *IngestionQueue) Ingest(ctx context.Context, input *ingestion.IngestInput) error {
    74  	select {
    75  	case <-ctx.Done():
    76  	case <-q.stop:
    77  	case q.queue <- input:
    78  		q.metrics.pendingItems.Inc()
    79  		// Once input is queued, context cancellation is ignored.
    80  		return nil
    81  	default:
    82  		// Drop data if the queue is full.
    83  	}
    84  
    85  	q.logger.WithField("key", input.Metadata.Key.Normalized()).Debugf("dropping since there's not enough space in the queue")
    86  	q.metrics.droppedItems.Inc()
    87  	return nil
    88  }
    89  
    90  func (q *IngestionQueue) runQueueWorker() {
    91  	defer q.wg.Done()
    92  	for {
    93  		select {
    94  		case input := <-q.queue:
    95  			if err := q.safePut(input); err != nil {
    96  				q.logger.WithField("key", input.Metadata.Key.Normalized()).WithError(err).Error("error happened while ingesting data")
    97  			}
    98  			q.metrics.pendingItems.Dec()
    99  		case <-q.stop:
   100  			return
   101  		}
   102  	}
   103  }
   104  
   105  func (q *IngestionQueue) safePut(input *ingestion.IngestInput) (err error) {
   106  	defer func() {
   107  		if r := recover(); r != nil {
   108  			err = fmt.Errorf("panic recovered: %v; %v", r, string(debug.Stack()))
   109  		}
   110  	}()
   111  	// TODO(kolesnikovae): It's better to derive a context that is cancelled on Stop.
   112  	return q.ingester.Ingest(context.TODO(), input)
   113  }