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

     1  package storage
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"runtime/debug"
     7  	"sync"
     8  
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/prometheus/client_golang/prometheus/promauto"
    11  	"github.com/sirupsen/logrus"
    12  )
    13  
    14  type IngestionQueue struct {
    15  	logger logrus.FieldLogger
    16  	putter Putter
    17  
    18  	wg    sync.WaitGroup
    19  	queue chan *PutInput
    20  	stop  chan struct{}
    21  
    22  	discardedTotal prometheus.Counter
    23  }
    24  
    25  const (
    26  	defaultQueueSize = 100
    27  	defaultWorkers   = 1
    28  )
    29  
    30  func NewIngestionQueue(logger logrus.FieldLogger, putter Putter, r prometheus.Registerer, c *Config) *IngestionQueue {
    31  	queueSize := c.queueSize
    32  	if queueSize == 0 {
    33  		queueSize = defaultQueueSize
    34  	}
    35  	queueWorkers := c.queueWorkers
    36  	if queueWorkers == 0 {
    37  		queueWorkers = defaultWorkers
    38  	}
    39  
    40  	q := IngestionQueue{
    41  		logger: logger,
    42  		putter: putter,
    43  		queue:  make(chan *PutInput, queueSize),
    44  		stop:   make(chan struct{}),
    45  
    46  		discardedTotal: promauto.With(r).NewCounter(prometheus.CounterOpts{
    47  			Name: "pyroscope_ingestion_queue_discarded_total",
    48  			Help: "number of ingestion requests discarded",
    49  		}),
    50  	}
    51  
    52  	q.wg.Add(queueWorkers)
    53  	for i := 0; i < queueWorkers; i++ {
    54  		go q.runQueueWorker()
    55  	}
    56  
    57  	return &q
    58  }
    59  
    60  func (s *IngestionQueue) Stop() {
    61  	close(s.stop)
    62  	s.wg.Wait()
    63  }
    64  
    65  func (s *IngestionQueue) Put(ctx context.Context, input *PutInput) error {
    66  	select {
    67  	case <-ctx.Done():
    68  	case <-s.stop:
    69  	case s.queue <- input:
    70  		// Once input is queued, context cancellation is ignored.
    71  		return nil
    72  	default:
    73  		// Drop data if the queue is full.
    74  	}
    75  	s.discardedTotal.Inc()
    76  	return nil
    77  }
    78  
    79  func (s *IngestionQueue) runQueueWorker() {
    80  	defer s.wg.Done()
    81  	for {
    82  		select {
    83  		case input, ok := <-s.queue:
    84  			if ok {
    85  				if err := s.safePut(input); err != nil {
    86  					s.logger.WithField("key", input.Key.Normalized()).WithError(err).Error("error happened while ingesting data")
    87  				}
    88  			}
    89  		case <-s.stop:
    90  			return
    91  		}
    92  	}
    93  }
    94  
    95  func (s *IngestionQueue) safePut(input *PutInput) (err error) {
    96  	defer func() {
    97  		if r := recover(); r != nil {
    98  			err = fmt.Errorf("panic recovered: %v; %v", r, string(debug.Stack()))
    99  		}
   100  	}()
   101  	// TODO(kolesnikovae): It's better to derive a context that is cancelled on Stop.
   102  	return s.putter.Put(context.TODO(), input)
   103  }