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

     1  package adhoc
     2  
     3  import (
     4  	"fmt"
     5  	"path"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/fatih/color"
    10  	"github.com/prometheus/client_golang/prometheus"
    11  	"github.com/pyroscope-io/pyroscope/pkg/adhoc/writer"
    12  	"github.com/pyroscope-io/pyroscope/pkg/analytics"
    13  	"github.com/pyroscope-io/pyroscope/pkg/health"
    14  	"github.com/sirupsen/logrus"
    15  
    16  	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
    17  	"github.com/pyroscope-io/pyroscope/pkg/config"
    18  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    19  )
    20  
    21  type runner interface {
    22  	Run() error
    23  }
    24  
    25  type mode int
    26  
    27  const (
    28  	modeExec mode = iota + 1
    29  	modeConnect
    30  	modePush
    31  	modePull
    32  )
    33  
    34  func (m mode) String() string {
    35  	switch m {
    36  	case modeExec:
    37  		return "exec"
    38  	case modeConnect:
    39  		return "connect"
    40  	case modePush:
    41  		return "push"
    42  	case modePull:
    43  		return "pull"
    44  	default:
    45  		return "unknown"
    46  	}
    47  }
    48  
    49  func Cli(cfg *config.Adhoc, args []string) error {
    50  	// Determine the mode to use to gather profiling data
    51  	var m mode
    52  	if cfg.Push {
    53  		if cfg.SpyName != "auto" {
    54  			return fmt.Errorf("'--push' and '--spy-name' can not be set together")
    55  		}
    56  		if cfg.Pid != 0 {
    57  			return fmt.Errorf("'--push' and '--pid' can not be set together")
    58  		}
    59  		if cfg.URL != "" {
    60  			return fmt.Errorf("'--push' and '--url' can not be set together")
    61  		}
    62  		m = modePush
    63  	} else if cfg.URL != "" {
    64  		if cfg.SpyName != "auto" {
    65  			return fmt.Errorf("'--url' and '--spy-name' can not be set together")
    66  		}
    67  		if cfg.Pid != 0 {
    68  			return fmt.Errorf("'--url' and '--pid' can not be set together")
    69  		}
    70  		m = modePull
    71  	} else if cfg.Pid != 0 {
    72  		m = modeConnect
    73  	} else if cfg.SpyName != "auto" {
    74  		m = modeExec
    75  	} else {
    76  		if len(args) == 0 {
    77  			return fmt.Errorf("could not detect the preferred profiling mode. Either pass the proper flag or some argument")
    78  		}
    79  		baseName := path.Base(args[0])
    80  		if spy.ResolveAutoName(baseName) == "" {
    81  			m = modePush
    82  		} else {
    83  			m = modeExec
    84  		}
    85  	}
    86  
    87  	level := logrus.PanicLevel
    88  	if l, err := logrus.ParseLevel(cfg.LogLevel); err == nil && !cfg.NoLogging {
    89  		level = l
    90  	}
    91  	logger := logrus.StandardLogger()
    92  	logger.SetLevel(level)
    93  	// an adhoc run shouldn't be a long-running process, make the output less verbose and more human-friendly.
    94  	logger.Formatter = &logrus.TextFormatter{DisableTimestamp: true}
    95  	logger.SetReportCaller(false)
    96  
    97  	switch cfg.OutputFormat {
    98  	case "html", "pprof", "collapsed", "none":
    99  	default:
   100  		return fmt.Errorf("invalid output format '%s', the only supported output formats are 'html', 'pprof' and 'collapsed'", cfg.OutputFormat)
   101  	}
   102  
   103  	st, err := storage.New(newStorageConfig(cfg), logger, prometheus.DefaultRegisterer, new(health.Controller), storage.NoopApplicationMetadataService{})
   104  	if err != nil {
   105  		return fmt.Errorf("could not initialize storage: %w", err)
   106  	}
   107  
   108  	var r runner
   109  	switch m {
   110  	case modeExec:
   111  		r, err = newExec(cfg, args, st, logger)
   112  	case modeConnect:
   113  		r, err = newConnect(cfg, st, logger)
   114  	case modePush:
   115  		r, err = newPush(cfg, args, st, logger)
   116  	case modePull:
   117  		r, err = newPull(cfg, args, st, logger)
   118  	default:
   119  		err = fmt.Errorf("could not determine the profiling mode. This shouldn't happen, please reporte the bug at %s",
   120  			color.GreenString("https://github.com/pyroscope-io/pyroscope/issues"),
   121  		)
   122  	}
   123  	if err != nil {
   124  		return err
   125  	}
   126  
   127  	t0 := time.Now()
   128  	status := "success"
   129  	if err := r.Run(); err != nil {
   130  		status = "failure"
   131  		logger.WithError(err).Error("running profiler")
   132  	}
   133  
   134  	wg := sync.WaitGroup{}
   135  	if !cfg.AnalyticsOptOut {
   136  		wg.Add(1)
   137  		go analytics.AdhocReport(m.String()+"-"+status, &wg)
   138  	}
   139  
   140  	if err := writer.NewWriter(cfg, st, logger).Write(t0, time.Now()); err != nil {
   141  		logger.WithError(err).Error("writing profiling data")
   142  	}
   143  
   144  	logger.Debug("stopping storage")
   145  	if err := st.Close(); err != nil {
   146  		logger.WithError(err).Error("storage close")
   147  	}
   148  	wg.Wait()
   149  	return err
   150  }
   151  
   152  func newStorageConfig(cfg *config.Adhoc) *storage.Config {
   153  	return storage.NewConfig(&config.Server{MaxNodesSerialization: cfg.MaxNodesSerialization}).WithInMemory()
   154  }