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

     1  package target
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"strings"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/sirupsen/logrus"
    12  
    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/remote"
    16  	"github.com/pyroscope-io/pyroscope/pkg/config"
    17  	"github.com/pyroscope-io/pyroscope/pkg/util/names"
    18  )
    19  
    20  const (
    21  	defaultBackoffPeriod = time.Second * 10
    22  )
    23  
    24  // Manager tracks targets and attaches spies to running processes.
    25  type Manager struct {
    26  	ctx    context.Context
    27  	cancel context.CancelFunc
    28  
    29  	logger *logrus.Logger
    30  	remote *remote.Remote
    31  	config *config.Agent
    32  	stop   chan struct{}
    33  	wg     sync.WaitGroup
    34  
    35  	resolve       func(config.Target) (target, bool)
    36  	backoffPeriod time.Duration
    37  }
    38  
    39  type target interface {
    40  	// attach blocks till the context cancellation or the target
    41  	// process exit, whichever occurs first.
    42  	attach(ctx context.Context)
    43  }
    44  
    45  func NewManager(l *logrus.Logger, r *remote.Remote, c *config.Agent) *Manager {
    46  	mgr := Manager{
    47  		logger:        l,
    48  		remote:        r,
    49  		config:        c,
    50  		stop:          make(chan struct{}),
    51  		backoffPeriod: defaultBackoffPeriod,
    52  	}
    53  	mgr.ctx, mgr.cancel = context.WithCancel(context.Background())
    54  	mgr.resolve = mgr.resolveTarget
    55  	return &mgr
    56  }
    57  
    58  func (mgr *Manager) canonise(t *config.Target) error {
    59  	if t.SpyName == types.GoSpy {
    60  		return fmt.Errorf("gospy can not profile other processes")
    61  	}
    62  	var found bool
    63  	for _, s := range spy.SupportedSpies {
    64  		if s == t.SpyName {
    65  			found = true
    66  			break
    67  		}
    68  	}
    69  	if !found {
    70  		return fmt.Errorf("spy %q is not supported", t.SpyName)
    71  	}
    72  	if t.SampleRate == 0 {
    73  		t.SampleRate = types.DefaultSampleRate
    74  	}
    75  	if t.ApplicationName == "" {
    76  		t.ApplicationName = string(t.SpyName) + "." + names.GetRandomName(generateSeed(t.ServiceName, t.SpyName))
    77  		logger := mgr.logger.WithField("spy-name", t.SpyName)
    78  		if t.ServiceName != "" {
    79  			logger = logger.WithField("service-name", t.ServiceName)
    80  		}
    81  		logger.Infof("we recommend specifying application name via 'application-name' parameter")
    82  		logger.Infof("for now we chose the name for you and it's %q", t.ApplicationName)
    83  	}
    84  	return nil
    85  }
    86  
    87  func (mgr *Manager) Start() {
    88  	for _, t := range mgr.config.Targets {
    89  		var tgt target
    90  		var ok bool
    91  		err := mgr.canonise(&t)
    92  		if err == nil {
    93  			tgt, ok = mgr.resolve(t)
    94  			if !ok {
    95  				err = fmt.Errorf("unknown target type")
    96  			}
    97  		}
    98  		if err != nil {
    99  			mgr.logger.
   100  				WithField("app-name", t.ApplicationName).
   101  				WithField("spy-name", t.SpyName).
   102  				WithError(err).Error("failed to setup target")
   103  			continue
   104  		}
   105  		mgr.wg.Add(1)
   106  		go mgr.runTarget(tgt)
   107  	}
   108  }
   109  
   110  func (mgr *Manager) Stop() {
   111  	mgr.cancel()
   112  	close(mgr.stop)
   113  	mgr.wg.Wait()
   114  }
   115  
   116  func (mgr *Manager) resolveTarget(t config.Target) (target, bool) {
   117  	var tgt target
   118  	switch {
   119  	case t.ServiceName != "":
   120  		tgt = newServiceTarget(mgr.logger, mgr.remote, t)
   121  	default:
   122  		return nil, false
   123  	}
   124  	return tgt, true
   125  }
   126  
   127  func (mgr *Manager) runTarget(t target) {
   128  	ticker := time.NewTicker(mgr.backoffPeriod)
   129  	defer func() {
   130  		ticker.Stop()
   131  		mgr.wg.Done()
   132  	}()
   133  	for {
   134  		select {
   135  		default:
   136  		case <-mgr.stop:
   137  			return
   138  		}
   139  		t.attach(mgr.ctx)
   140  		// Unless manager is stopped, run spy again after some backoff
   141  		// period regardless of the exit reason.
   142  		select {
   143  		case <-ticker.C:
   144  		case <-mgr.stop:
   145  			return
   146  		}
   147  	}
   148  }
   149  
   150  func generateSeed(args ...string) string {
   151  	path, err := os.Getwd()
   152  	if err != nil {
   153  		path = "<unknown>"
   154  	}
   155  	return path + "|" + strings.Join(args, "&")
   156  }