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

     1  package exec
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	goexec "os/exec"
     7  	"os/signal"
     8  	"path"
     9  	"time"
    10  
    11  	"github.com/sirupsen/logrus"
    12  
    13  	"github.com/pyroscope-io/pyroscope/pkg/agent"
    14  	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
    15  	"github.com/pyroscope-io/pyroscope/pkg/agent/types"
    16  	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
    17  	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
    18  	"github.com/pyroscope-io/pyroscope/pkg/config"
    19  	"github.com/pyroscope-io/pyroscope/pkg/util/process"
    20  )
    21  
    22  type Exec struct {
    23  	Args               []string
    24  	Logger             *logrus.Logger
    25  	Upstream           upstream.Upstream
    26  	SpyName            string
    27  	ApplicationName    string
    28  	SampleRate         uint32
    29  	DetectSubprocesses bool
    30  	Tags               map[string]string
    31  	NoRootDrop         bool
    32  	UserName           string
    33  	GroupName          string
    34  	PHPSpyArgs         string
    35  }
    36  
    37  func NewExec(cfg *config.Exec, args []string) (*Exec, error) {
    38  	if len(args) == 0 {
    39  		return nil, fmt.Errorf("no arguments passed")
    40  	}
    41  
    42  	spyName := cfg.SpyName
    43  	if spyName == "auto" {
    44  		baseName := path.Base(args[0])
    45  		spyName = spy.ResolveAutoName(baseName)
    46  		if spyName == "" {
    47  			return nil, UnsupportedSpyError{Subcommand: "exec", Args: args}
    48  		}
    49  	}
    50  	if err := PerformChecks(spyName); err != nil {
    51  		return nil, err
    52  	}
    53  
    54  	logger := NewLogger(cfg.LogLevel, cfg.NoLogging)
    55  
    56  	rc := remote.RemoteConfig{
    57  		AuthToken:              cfg.AuthToken,
    58  		TenantID:               cfg.TenantID,
    59  		BasicAuthUser:          cfg.BasicAuthUser,
    60  		BasicAuthPassword:      cfg.BasicAuthPassword,
    61  		HTTPHeaders:            cfg.Headers,
    62  		UpstreamThreads:        cfg.UpstreamThreads,
    63  		UpstreamAddress:        cfg.ServerAddress,
    64  		UpstreamRequestTimeout: cfg.UpstreamRequestTimeout,
    65  	}
    66  	up, err := remote.New(rc, logger)
    67  	if err != nil {
    68  		return nil, fmt.Errorf("new remote upstream: %v", err)
    69  	}
    70  
    71  	// if the sample rate is zero, use the default value
    72  	sampleRate := uint32(types.DefaultSampleRate)
    73  	if cfg.SampleRate != 0 {
    74  		sampleRate = uint32(cfg.SampleRate)
    75  	}
    76  
    77  	return &Exec{
    78  		Args:               args,
    79  		Logger:             logger,
    80  		Upstream:           up,
    81  		SpyName:            spyName,
    82  		ApplicationName:    CheckApplicationName(logger, cfg.ApplicationName, spyName, args),
    83  		SampleRate:         sampleRate,
    84  		DetectSubprocesses: cfg.DetectSubprocesses,
    85  		Tags:               cfg.Tags,
    86  		NoRootDrop:         cfg.NoRootDrop,
    87  		UserName:           cfg.UserName,
    88  		GroupName:          cfg.GroupName,
    89  		PHPSpyArgs:         cfg.PHPSpyArgs,
    90  	}, nil
    91  }
    92  
    93  func (e *Exec) Run() error {
    94  	e.Logger.WithFields(logrus.Fields{
    95  		"args": fmt.Sprintf("%q", e.Args),
    96  	}).Debug("starting command")
    97  
    98  	// The channel buffer capacity should be sufficient to be keep up with
    99  	// the expected signal rate (in case of Exec all the signals to be relayed
   100  	// to the child process)
   101  	c := make(chan os.Signal, 10)
   102  	var cmd *goexec.Cmd
   103  	// Note that we don't specify which signals to be sent: any signal to be
   104  	// relayed to the child process (including SIGINT and SIGTERM).
   105  	signal.Notify(c)
   106  	cmd = goexec.Command(e.Args[0], e.Args[1:]...)
   107  	cmd.Stderr = os.Stderr
   108  	cmd.Stdout = os.Stdout
   109  	cmd.Stdin = os.Stdin
   110  	if err := adjustCmd(cmd, e.NoRootDrop, e.UserName, e.GroupName); err != nil {
   111  		return err
   112  	}
   113  	if err := cmd.Start(); err != nil {
   114  		return err
   115  	}
   116  	defer func() {
   117  		signal.Stop(c)
   118  		close(c)
   119  	}()
   120  
   121  	sc := agent.SessionConfig{
   122  		Upstream:         e.Upstream,
   123  		AppName:          e.ApplicationName,
   124  		Tags:             e.Tags,
   125  		ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU},
   126  		SpyName:          e.SpyName,
   127  		SampleRate:       e.SampleRate,
   128  		UploadRate:       10 * time.Second,
   129  		Pid:              cmd.Process.Pid,
   130  		WithSubprocesses: e.DetectSubprocesses,
   131  		Logger:           e.Logger,
   132  		PHPSpyArgs:       e.PHPSpyArgs,
   133  	}
   134  	session, err := agent.NewSession(sc)
   135  	if err != nil {
   136  		return fmt.Errorf("new session: %w", err)
   137  	}
   138  
   139  	logrus.WithFields(logrus.Fields{
   140  		"app-name":            e.ApplicationName,
   141  		"spy-name":            e.SpyName,
   142  		"pid":                 cmd.Process.Pid,
   143  		"detect-subprocesses": e.DetectSubprocesses,
   144  	}).Debug("starting agent session")
   145  
   146  	e.Upstream.Start()
   147  	defer e.Upstream.Stop()
   148  
   149  	if err := session.Start(); err != nil {
   150  		return fmt.Errorf("start session: %w", err)
   151  	}
   152  	defer session.Stop()
   153  
   154  	// Wait for spawned process to exit
   155  	ticker := time.NewTicker(time.Second)
   156  	defer ticker.Stop()
   157  	for {
   158  		select {
   159  		case s := <-c:
   160  			_ = process.SendSignal(cmd.Process, s)
   161  		case <-ticker.C:
   162  			if !process.Exists(cmd.Process.Pid) {
   163  				logrus.Debug("child process exited")
   164  				return cmd.Wait()
   165  			}
   166  		}
   167  	}
   168  }