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

     1  package exec
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"os/signal"
     7  	"syscall"
     8  	"time"
     9  
    10  	"github.com/sirupsen/logrus"
    11  
    12  	"github.com/pyroscope-io/pyroscope/pkg/agent"
    13  	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
    14  	"github.com/pyroscope-io/pyroscope/pkg/agent/types"
    15  	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
    16  	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
    17  	"github.com/pyroscope-io/pyroscope/pkg/config"
    18  	"github.com/pyroscope-io/pyroscope/pkg/util/process"
    19  )
    20  
    21  type Connect struct {
    22  	Logger             *logrus.Logger
    23  	Upstream           upstream.Upstream
    24  	SpyName            string
    25  	ApplicationName    string
    26  	SampleRate         uint32
    27  	DetectSubprocesses bool
    28  	Tags               map[string]string
    29  	Pid                int
    30  	PHPSpyArgs         string
    31  }
    32  
    33  func NewConnect(cfg *config.Connect) (*Connect, error) {
    34  	spyName := cfg.SpyName
    35  	if cfg.Pid == -1 {
    36  		if spyName != "" && spyName != "ebpfspy" {
    37  			return nil, fmt.Errorf("pid -1 can only be used with ebpfspy")
    38  		}
    39  		spyName = "ebpfspy"
    40  	}
    41  	if err := PerformChecks(spyName); err != nil {
    42  		return nil, err
    43  	}
    44  
    45  	logger := NewLogger(cfg.LogLevel, cfg.NoLogging)
    46  
    47  	rc := remote.RemoteConfig{
    48  		AuthToken:              cfg.AuthToken,
    49  		TenantID:               cfg.TenantID,
    50  		BasicAuthUser:          cfg.BasicAuthUser,
    51  		BasicAuthPassword:      cfg.BasicAuthPassword,
    52  		HTTPHeaders:            cfg.Headers,
    53  		UpstreamThreads:        cfg.UpstreamThreads,
    54  		UpstreamAddress:        cfg.ServerAddress,
    55  		UpstreamRequestTimeout: cfg.UpstreamRequestTimeout,
    56  	}
    57  	up, err := remote.New(rc, logger)
    58  	if err != nil {
    59  		return nil, fmt.Errorf("new remote upstream: %v", err)
    60  	}
    61  
    62  	// if the sample rate is zero, use the default value
    63  	sampleRate := uint32(types.DefaultSampleRate)
    64  	if cfg.SampleRate != 0 {
    65  		sampleRate = uint32(cfg.SampleRate)
    66  	}
    67  
    68  	return &Connect{
    69  		Logger:             logger,
    70  		Upstream:           up,
    71  		SpyName:            spyName,
    72  		ApplicationName:    CheckApplicationName(logger, cfg.ApplicationName, spyName, []string{}),
    73  		SampleRate:         sampleRate,
    74  		DetectSubprocesses: cfg.DetectSubprocesses,
    75  		Tags:               cfg.Tags,
    76  		Pid:                cfg.Pid,
    77  		PHPSpyArgs:         cfg.PHPSpyArgs,
    78  	}, nil
    79  }
    80  
    81  func (c *Connect) Run() error {
    82  	c.Logger.Debug("starting command")
    83  
    84  	// The channel buffer capacity should be sufficient to be keep up with
    85  	// the expected signal rate (in case of Exec all the signals to be relayed
    86  	// to the child process)
    87  	ch := make(chan os.Signal, 10)
    88  	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
    89  	defer func() {
    90  		signal.Stop(ch)
    91  		close(ch)
    92  	}()
    93  
    94  	sc := agent.SessionConfig{
    95  		Upstream:         c.Upstream,
    96  		AppName:          c.ApplicationName,
    97  		Tags:             c.Tags,
    98  		ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU},
    99  		SpyName:          c.SpyName,
   100  		SampleRate:       c.SampleRate,
   101  		UploadRate:       10 * time.Second,
   102  		Pid:              c.Pid,
   103  		WithSubprocesses: c.DetectSubprocesses,
   104  		Logger:           c.Logger,
   105  		PHPSpyArgs:       c.PHPSpyArgs,
   106  	}
   107  	session, err := agent.NewSession(sc)
   108  	if err != nil {
   109  		return fmt.Errorf("new session: %w", err)
   110  	}
   111  
   112  	logrus.WithFields(logrus.Fields{
   113  		"app-name":            c.ApplicationName,
   114  		"spy-name":            c.SpyName,
   115  		"pid":                 c.Pid,
   116  		"detect-subprocesses": c.DetectSubprocesses,
   117  	}).Debug("starting agent session")
   118  
   119  	c.Upstream.Start()
   120  	defer c.Upstream.Stop()
   121  
   122  	if err = session.Start(); err != nil {
   123  		return fmt.Errorf("start session: %w", err)
   124  	}
   125  	defer session.Stop()
   126  
   127  	// wait for process to exit
   128  	// pid == -1 means we're profiling whole system
   129  	if c.Pid == -1 {
   130  		<-ch
   131  		return nil
   132  	}
   133  	ticker := time.NewTicker(time.Second)
   134  	defer ticker.Stop()
   135  	for {
   136  		select {
   137  		case <-ch:
   138  			return nil
   139  		case <-ticker.C:
   140  			if !process.Exists(c.Pid) {
   141  				c.Logger.Debugf("child process exited")
   142  				return nil
   143  			}
   144  		}
   145  	}
   146  }