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

     1  package storage
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"math/big"
     7  	"time"
     8  
     9  	"github.com/sirupsen/logrus"
    10  
    11  	"github.com/pyroscope-io/pyroscope/pkg/model/appmetadata"
    12  	"github.com/pyroscope-io/pyroscope/pkg/storage/dimension"
    13  	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
    14  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    15  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
    16  )
    17  
    18  type PutInput struct {
    19  	StartTime       time.Time
    20  	EndTime         time.Time
    21  	Key             *segment.Key
    22  	Val             *tree.Tree
    23  	SpyName         string
    24  	SampleRate      uint32
    25  	Units           metadata.Units
    26  	AggregationType metadata.AggregationType
    27  }
    28  
    29  func (s *Storage) Put(ctx context.Context, pi *PutInput) error {
    30  	if s.hc.IsOutOfDiskSpace() {
    31  		return errOutOfSpace
    32  	}
    33  	if pi.StartTime.Before(s.retentionPolicy().LowerTimeBoundary()) {
    34  		return errRetention
    35  	}
    36  
    37  	if err := segment.ValidateKey(pi.Key); err != nil {
    38  		return err
    39  	}
    40  
    41  	if err := s.appSvc.CreateOrUpdate(ctx, appmetadata.ApplicationMetadata{
    42  		FQName:          pi.Key.AppName(),
    43  		SpyName:         pi.SpyName,
    44  		SampleRate:      pi.SampleRate,
    45  		Units:           pi.Units,
    46  		AggregationType: pi.AggregationType,
    47  	}); err != nil {
    48  		s.logger.Error("error saving metadata", err)
    49  	}
    50  
    51  	s.putTotal.Inc()
    52  	if pi.Key.HasProfileID() {
    53  		if err := s.ensureAppSegmentExists(pi); err != nil {
    54  			return err
    55  		}
    56  		return s.exemplars.insert(ctx, pi)
    57  	}
    58  
    59  	s.logger.WithFields(logrus.Fields{
    60  		"startTime":       pi.StartTime.String(),
    61  		"endTime":         pi.EndTime.String(),
    62  		"key":             pi.Key.Normalized(),
    63  		"samples":         pi.Val.Samples(),
    64  		"units":           pi.Units,
    65  		"aggregationType": pi.AggregationType,
    66  	}).Debug("storage.Put")
    67  
    68  	if err := s.labels.PutLabels(pi.Key.Labels()); err != nil {
    69  		return fmt.Errorf("unable to write labels: %w", err)
    70  	}
    71  
    72  	sk := pi.Key.SegmentKey()
    73  	for k, v := range pi.Key.Labels() {
    74  		key := k + ":" + v
    75  		r, err := s.dimensions.GetOrCreate(key)
    76  		if err != nil {
    77  			s.logger.Errorf("dimensions cache for %v: %v", key, err)
    78  			continue
    79  		}
    80  		r.(*dimension.Dimension).Insert([]byte(sk))
    81  		s.dimensions.Put(key, r)
    82  	}
    83  
    84  	r, err := s.segments.GetOrCreate(sk)
    85  	if err != nil {
    86  		return fmt.Errorf("segments cache for %v: %v", sk, err)
    87  	}
    88  
    89  	st := r.(*segment.Segment)
    90  	st.SetMetadata(metadata.Metadata{
    91  		SpyName:         pi.SpyName,
    92  		SampleRate:      pi.SampleRate,
    93  		Units:           pi.Units,
    94  		AggregationType: pi.AggregationType,
    95  	})
    96  
    97  	samples := pi.Val.Samples()
    98  	err = st.Put(pi.StartTime, pi.EndTime, samples, func(depth int, t time.Time, r *big.Rat, addons []segment.Addon) {
    99  		tk := pi.Key.TreeKey(depth, t)
   100  		res, err := s.trees.GetOrCreate(tk)
   101  		if err != nil {
   102  			s.logger.Errorf("trees cache for %v: %v", tk, err)
   103  			return
   104  		}
   105  		cachedTree := res.(*tree.Tree)
   106  		treeClone := pi.Val.Clone(r)
   107  		for _, addon := range addons {
   108  			if res, ok := s.trees.Lookup(pi.Key.TreeKey(addon.Depth, addon.T)); ok {
   109  				ta := res.(*tree.Tree)
   110  				ta.RLock()
   111  				treeClone.Merge(ta)
   112  				ta.RUnlock()
   113  			}
   114  		}
   115  		cachedTree.Lock()
   116  		cachedTree.Merge(treeClone)
   117  		cachedTree.Unlock()
   118  		s.trees.Put(tk, cachedTree)
   119  	})
   120  	if err != nil {
   121  		return err
   122  	}
   123  
   124  	s.segments.Put(sk, st)
   125  	return nil
   126  }