github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/agent/upstream/direct/direct.go (about)

     1  package direct
     2  
     3  import (
     4  	"context"
     5  	"runtime/debug"
     6  	"sync"
     7  
     8  	"github.com/sirupsen/logrus"
     9  
    10  	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
    11  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    12  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    13  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
    14  )
    15  
    16  const upstreamThreads = 1
    17  
    18  type Direct struct {
    19  	storage  storage.Putter
    20  	exporter storage.MetricsExporter
    21  	queue    chan *upstream.UploadJob
    22  	stop     chan struct{}
    23  	wg       sync.WaitGroup
    24  }
    25  
    26  func New(s storage.Putter, e storage.MetricsExporter) *Direct {
    27  	return &Direct{
    28  		storage:  s,
    29  		exporter: e,
    30  		queue:    make(chan *upstream.UploadJob, 100),
    31  		stop:     make(chan struct{}),
    32  	}
    33  }
    34  
    35  func (u *Direct) Start() {
    36  	u.wg.Add(upstreamThreads)
    37  	for i := 0; i < upstreamThreads; i++ {
    38  		go u.uploadLoop()
    39  	}
    40  }
    41  
    42  func (u *Direct) Stop() {
    43  	close(u.stop)
    44  	u.wg.Wait()
    45  }
    46  
    47  func (u *Direct) Upload(j *upstream.UploadJob) {
    48  	select {
    49  	case u.queue <- j:
    50  	case <-u.stop:
    51  		return
    52  	default:
    53  		logrus.Error("Direct upload queue is full, dropping a profile")
    54  	}
    55  }
    56  
    57  func (u *Direct) uploadProfile(j *upstream.UploadJob) {
    58  	key, err := segment.ParseKey(j.Name)
    59  	if err != nil {
    60  		logrus.WithField("key", key).Error("invalid key:")
    61  		return
    62  	}
    63  
    64  	pi := &storage.PutInput{
    65  		StartTime:       j.StartTime,
    66  		EndTime:         j.EndTime,
    67  		Key:             key,
    68  		Val:             tree.New(),
    69  		SpyName:         j.SpyName,
    70  		SampleRate:      j.SampleRate,
    71  		Units:           j.Units,
    72  		AggregationType: j.AggregationType,
    73  	}
    74  
    75  	cb := pi.Val.Insert
    76  	if o, ok := u.exporter.Evaluate(pi); ok {
    77  		cb = func(k []byte, v uint64) {
    78  			o.Observe(k, int(v))
    79  			cb(k, v)
    80  		}
    81  	}
    82  
    83  	j.Trie.Iterate(cb)
    84  	if err = u.storage.Put(context.TODO(), pi); err != nil {
    85  		logrus.WithError(err).Error("failed to store a local profile")
    86  	}
    87  }
    88  
    89  func (u *Direct) uploadLoop() {
    90  	defer u.wg.Done()
    91  	for {
    92  		select {
    93  		case j := <-u.queue:
    94  			u.safeUpload(j)
    95  		case <-u.stop:
    96  			return
    97  		}
    98  	}
    99  }
   100  
   101  // do safe upload
   102  func (u *Direct) safeUpload(j *upstream.UploadJob) {
   103  	defer func() {
   104  		if r := recover(); r != nil {
   105  			logrus.Errorf("panic recovered: %v; %v", r, string(debug.Stack()))
   106  		}
   107  	}()
   108  	u.uploadProfile(j)
   109  }