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

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"encoding/json"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  	"strings"
    11  
    12  	gprofile "github.com/google/pprof/profile"
    13  	"github.com/grafana/dskit/runutil"
    14  	"github.com/k0kubun/pp/v3"
    15  	"github.com/klauspost/compress/gzip"
    16  	"github.com/mattn/go-isatty"
    17  	"github.com/pkg/errors"
    18  
    19  	googlev1 "github.com/grafana/pyroscope/api/gen/proto/go/google/v1"
    20  	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
    21  )
    22  
    23  const (
    24  	outputConsole = "console"
    25  	outputRaw     = "raw"
    26  	outputPprof   = "pprof="
    27  )
    28  
    29  func outputSeries(result []*typesv1.Labels) error {
    30  	enc := json.NewEncoder(os.Stdout)
    31  	m := make(map[string]interface{})
    32  	for _, s := range result {
    33  		clear(m)
    34  		for _, l := range s.Labels {
    35  			m[l.Name] = l.Value
    36  		}
    37  		if err := enc.Encode(m); err != nil {
    38  			return err
    39  		}
    40  	}
    41  	return nil
    42  }
    43  
    44  func outputMergeProfile(ctx context.Context, outputFlag string, profile *googlev1.Profile) error {
    45  	mypp := pp.New()
    46  	mypp.SetColoringEnabled(isatty.IsTerminal(os.Stdout.Fd()))
    47  	mypp.SetExportedOnly(true)
    48  
    49  	if outputFlag == outputConsole {
    50  		buf, err := profile.MarshalVT()
    51  		if err != nil {
    52  			return errors.Wrap(err, "failed to marshal protobuf")
    53  		}
    54  
    55  		p, err := gprofile.Parse(bytes.NewReader(buf))
    56  		if err != nil {
    57  			return errors.Wrap(err, "failed to parse profile")
    58  		}
    59  
    60  		fmt.Fprintln(output(ctx), p.String())
    61  		return nil
    62  
    63  	}
    64  
    65  	if outputFlag == outputRaw {
    66  		mypp.Print(profile)
    67  		return nil
    68  	}
    69  
    70  	if strings.HasPrefix(outputFlag, outputPprof) {
    71  		filePath := strings.TrimPrefix(outputFlag, outputPprof)
    72  		if filePath == "" {
    73  			return errors.New("no file path specified after pprof=")
    74  		}
    75  		buf, err := profile.MarshalVT()
    76  		if err != nil {
    77  			return errors.Wrap(err, "failed to marshal protobuf")
    78  		}
    79  
    80  		// open new file, fail when the file already exists
    81  		f, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0644)
    82  		if err != nil {
    83  			return errors.Wrap(err, "failed to create pprof file")
    84  		}
    85  		defer runutil.CloseWithErrCapture(&err, f, "failed to close pprof file")
    86  
    87  		gzipWriter := gzip.NewWriter(f)
    88  		defer runutil.CloseWithErrCapture(&err, gzipWriter, "failed to close pprof gzip writer")
    89  
    90  		if _, err := io.Copy(gzipWriter, bytes.NewReader(buf)); err != nil {
    91  			return errors.Wrap(err, "failed to write pprof")
    92  		}
    93  
    94  		return nil
    95  	}
    96  
    97  	return errors.Errorf("unknown output %s", outputFlag)
    98  }