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

     1  package selfprofiling
     2  
     3  import (
     4  	"context"
     5  	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
     6  	"runtime"
     7  	"runtime/debug"
     8  	"time"
     9  
    10  	"github.com/pyroscope-io/client/pyroscope"
    11  	"github.com/pyroscope-io/client/upstream"
    12  
    13  	"github.com/pyroscope-io/pyroscope/pkg/convert/pprof"
    14  	"github.com/pyroscope-io/pyroscope/pkg/ingestion"
    15  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    16  	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
    17  )
    18  
    19  func NewSession(logger pyroscope.Logger, ingester ingestion.Ingester, appName string, tags map[string]string) *pyroscope.Session {
    20  	runtime.SetMutexProfileFraction(5)
    21  	runtime.SetBlockProfileRate(5)
    22  
    23  	session, _ := pyroscope.NewSession(pyroscope.SessionConfig{
    24  		Upstream: NewUpstream(logger, ingester),
    25  		AppName:  appName,
    26  		ProfilingTypes: []pyroscope.ProfileType{
    27  			pyroscope.ProfileCPU,
    28  			pyroscope.ProfileInuseObjects,
    29  			pyroscope.ProfileAllocObjects,
    30  			pyroscope.ProfileInuseSpace,
    31  			pyroscope.ProfileAllocSpace,
    32  			pyroscope.ProfileGoroutines,
    33  			pyroscope.ProfileMutexCount,
    34  			pyroscope.ProfileMutexDuration,
    35  			pyroscope.ProfileBlockCount,
    36  			pyroscope.ProfileBlockDuration,
    37  		},
    38  		SampleRate: 100,
    39  		UploadRate: 10 * time.Second,
    40  		Logger:     logger,
    41  		Tags:       tags,
    42  	})
    43  	return session
    44  }
    45  
    46  type Upstream struct {
    47  	logger   pyroscope.Logger
    48  	ingester ingestion.Ingester
    49  }
    50  
    51  func NewUpstream(logger pyroscope.Logger, ingester ingestion.Ingester) *Upstream {
    52  	return &Upstream{
    53  		logger:   logger,
    54  		ingester: ingester,
    55  	}
    56  }
    57  
    58  func (u *Upstream) Upload(j *upstream.UploadJob) {
    59  	defer func() {
    60  		if r := recover(); r != nil {
    61  			u.logger.Errorf("panic recovered: %v; %v", r, string(debug.Stack()))
    62  		}
    63  	}()
    64  
    65  	key, err := segment.ParseKey(j.Name)
    66  	if err != nil {
    67  		u.logger.Errorf("invalid key %q: %v", j.Name, err)
    68  		return
    69  	}
    70  
    71  	if len(j.Profile) == 0 {
    72  		return
    73  	}
    74  
    75  	profile := &pprof.RawProfile{
    76  		Profile:          j.Profile,
    77  		SampleTypeConfig: tree.DefaultSampleTypeMapping,
    78  	}
    79  	if j.SampleTypeConfig != nil {
    80  		profile.SampleTypeConfig = sampleTypeConfigFromUpstream(j.SampleTypeConfig)
    81  	}
    82  	if len(j.PrevProfile) > 0 {
    83  		profile.PreviousProfile = j.PrevProfile
    84  	}
    85  
    86  	err = u.ingester.Ingest(context.TODO(), &ingestion.IngestInput{
    87  		Profile: profile,
    88  		Metadata: ingestion.Metadata{
    89  			SpyName:   j.SpyName,
    90  			StartTime: j.StartTime,
    91  			EndTime:   j.EndTime,
    92  			Key:       key,
    93  		},
    94  	})
    95  
    96  	if err != nil {
    97  		u.logger.Errorf("failed to store a local profile: %v", err)
    98  	}
    99  }
   100  
   101  func (*Upstream) Flush() {
   102  
   103  }
   104  
   105  func sampleTypeConfigFromUpstream(config map[string]*upstream.SampleType) map[string]*tree.SampleTypeConfig {
   106  	res := make(map[string]*tree.SampleTypeConfig)
   107  	for k, v := range config {
   108  		vv := &tree.SampleTypeConfig{
   109  			Units:       metadata.Units(v.Units),
   110  			DisplayName: v.DisplayName,
   111  			Aggregation: metadata.AggregationType(v.Aggregation),
   112  			Cumulative:  v.Cumulative,
   113  			Sampled:     v.Sampled,
   114  		}
   115  		res[k] = vv
   116  	}
   117  	return res
   118  }