github.com/grafana/pyroscope@v1.18.0/cmd/profilecli/upload.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"os"
     6  	"time"
     7  
     8  	"connectrpc.com/connect"
     9  	"github.com/go-kit/log/level"
    10  	"github.com/google/uuid"
    11  
    12  	pushv1 "github.com/grafana/pyroscope/api/gen/proto/go/push/v1"
    13  	"github.com/grafana/pyroscope/api/gen/proto/go/push/v1/pushv1connect"
    14  	connectapi "github.com/grafana/pyroscope/pkg/api/connect"
    15  	"github.com/grafana/pyroscope/pkg/model"
    16  	"github.com/grafana/pyroscope/pkg/pprof"
    17  )
    18  
    19  func (c *phlareClient) pusherClient() pushv1connect.PusherServiceClient {
    20  	return pushv1connect.NewPusherServiceClient(
    21  		c.httpClient(),
    22  		c.URL,
    23  		append(
    24  			connectapi.DefaultClientOptions(),
    25  			c.protocolOption(),
    26  		)...,
    27  	)
    28  }
    29  
    30  type uploadParams struct {
    31  	*phlareClient
    32  	paths             []string
    33  	extraLabels       map[string]string
    34  	overrideTimestamp bool
    35  }
    36  
    37  func addUploadParams(cmd commander) *uploadParams {
    38  	var (
    39  		params = &uploadParams{
    40  			extraLabels: map[string]string{},
    41  		}
    42  	)
    43  	params.phlareClient = addPhlareClient(cmd)
    44  
    45  	cmd.Arg("path", "Path(s) to profile(s) to upload").Required().ExistingFilesVar(&params.paths)
    46  	cmd.Flag("extra-labels", "Add additional labels to the profile(s)").StringMapVar(&params.extraLabels)
    47  	cmd.Flag("override-timestamp", "Set the profile timestamp to now").BoolVar(&params.overrideTimestamp)
    48  	return params
    49  }
    50  
    51  func upload(ctx context.Context, params *uploadParams) (err error) {
    52  	pc := params.pusherClient()
    53  
    54  	lblStrings := make([]string, 0, len(params.extraLabels)*2)
    55  	for key, value := range params.extraLabels {
    56  		lblStrings = append(lblStrings, key, value)
    57  	}
    58  
    59  	var (
    60  		lbl        = model.LabelsFromStrings(lblStrings...)
    61  		series     = make([]*pushv1.RawProfileSeries, len(params.paths))
    62  		lblBuilder = model.NewLabelsBuilder(lbl)
    63  	)
    64  	for idx, path := range params.paths {
    65  		lblBuilder.Reset(lbl)
    66  
    67  		data, err := os.ReadFile(path)
    68  		if err != nil {
    69  			return nil
    70  		}
    71  
    72  		profile, err := pprof.RawFromBytes(data)
    73  		if err != nil {
    74  			return err
    75  		}
    76  
    77  		if params.overrideTimestamp {
    78  			profile.TimeNanos = time.Now().UnixNano()
    79  			data, err = pprof.Marshal(profile.Profile, true)
    80  			if err != nil {
    81  				return err
    82  			}
    83  		}
    84  
    85  		// detect name if no name has been set
    86  		if lbl.Get(model.LabelNameProfileName) == "" {
    87  			name := "unknown"
    88  			for _, t := range profile.SampleType {
    89  				if sid := int(t.Type); sid < len(profile.StringTable) {
    90  					if s := profile.StringTable[sid]; s == "cpu" {
    91  						name = "process_cpu"
    92  						break
    93  					} else if s == "alloc_space" || s == "inuse_space" {
    94  						name = "memory"
    95  						break
    96  					} else {
    97  						level.Debug(logger).Log("msg", "unspecific/unknown profile sample type", "profile", s)
    98  					}
    99  				}
   100  			}
   101  			lblBuilder.Set(model.LabelNameProfileName, name)
   102  		}
   103  
   104  		// set a default service_name label if one is not provided
   105  		if lbl.Get(model.LabelNameServiceName) == "" {
   106  			lblBuilder.Set(model.LabelNameServiceName, "profilecli-upload")
   107  		}
   108  
   109  		series[idx] = &pushv1.RawProfileSeries{
   110  			Labels: lblBuilder.Labels(),
   111  			Samples: []*pushv1.RawSample{{
   112  				ID:         uuid.New().String(),
   113  				RawProfile: data,
   114  			}},
   115  		}
   116  	}
   117  
   118  	_, err = pc.Push(ctx, connect.NewRequest(&pushv1.PushRequest{
   119  		Series: series,
   120  	}))
   121  
   122  	if err != nil {
   123  		return err
   124  	}
   125  
   126  	for idx := range series {
   127  		level.Info(logger).Log("msg", "successfully uploaded profile", "id", series[idx].Samples[0].ID, "labels", model.Labels(series[idx].Labels).ToPrometheusLabels().String(), "path", params.paths[idx])
   128  	}
   129  
   130  	return nil
   131  }