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

     1  package adhoc
     2  
     3  import (
     4  	"os"
     5  	goexec "os/exec"
     6  	"os/signal"
     7  	"time"
     8  
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/pyroscope-io/pyroscope/pkg/parser"
    11  	"github.com/sirupsen/logrus"
    12  
    13  	"github.com/pyroscope-io/pyroscope/pkg/config"
    14  	"github.com/pyroscope-io/pyroscope/pkg/exec"
    15  	"github.com/pyroscope-io/pyroscope/pkg/exporter"
    16  	"github.com/pyroscope-io/pyroscope/pkg/scrape"
    17  	scrapeconfig "github.com/pyroscope-io/pyroscope/pkg/scrape/config"
    18  	"github.com/pyroscope-io/pyroscope/pkg/scrape/discovery/targetgroup"
    19  	"github.com/pyroscope-io/pyroscope/pkg/scrape/model"
    20  	"github.com/pyroscope-io/pyroscope/pkg/storage"
    21  	"github.com/pyroscope-io/pyroscope/pkg/util/process"
    22  )
    23  
    24  type pull struct {
    25  	c       chan os.Signal
    26  	cmd     *goexec.Cmd
    27  	logger  *logrus.Logger
    28  	manager *scrape.Manager
    29  	targets map[string][]*targetgroup.Group
    30  }
    31  
    32  func newPull(cfg *config.Adhoc, args []string, st *storage.Storage, logger *logrus.Logger) (runner, error) {
    33  	var c chan os.Signal
    34  	var cmd *goexec.Cmd
    35  	if len(args) > 0 {
    36  		// when arguments are specified, let's launch the target first.
    37  		c = make(chan os.Signal, 10)
    38  		signal.Notify(c)
    39  		cmd = goexec.Command(args[0], args[1:]...)
    40  		cmd.Stderr = os.Stderr
    41  		cmd.Stdout = os.Stdout
    42  		cmd.Stdin = os.Stdin
    43  	}
    44  
    45  	// build the scrape manager to retrieve data
    46  	defaultMetricsRegistry := prometheus.DefaultRegisterer
    47  	e, err := exporter.NewExporter(config.MetricsExportRules{}, defaultMetricsRegistry)
    48  	if err != nil {
    49  		return nil, err
    50  	}
    51  
    52  	p := parser.New(logger, st, e)
    53  	m := scrape.NewManager(logger, p, defaultMetricsRegistry)
    54  	scrapeCfg := &(*scrapeconfig.DefaultConfig())
    55  	scrapeCfg.JobName = "adhoc"
    56  	scrapeCfg.EnabledProfiles = []string{"cpu", "mem"}
    57  	if err := m.ApplyConfig([]*scrapeconfig.Config{scrapeCfg}); err != nil {
    58  		return nil, err
    59  	}
    60  
    61  	appName := exec.CheckApplicationName(logger, cfg.ApplicationName, "pull", args)
    62  	targets := map[string][]*targetgroup.Group{
    63  		"adhoc": {
    64  			&targetgroup.Group{
    65  				Source: "adhoc",
    66  				Labels: model.LabelSet{},
    67  				Targets: []model.LabelSet{
    68  					{
    69  						model.AddressLabel:    model.LabelValue(cfg.URL),
    70  						model.MetricNameLabel: model.LabelValue(appName),
    71  					},
    72  				},
    73  			},
    74  		},
    75  	}
    76  
    77  	return &pull{
    78  		c:       c,
    79  		cmd:     cmd,
    80  		logger:  logger,
    81  		manager: m,
    82  		targets: targets,
    83  	}, nil
    84  }
    85  
    86  func (p *pull) Run() error {
    87  	if p.cmd != nil {
    88  		if err := p.cmd.Start(); err != nil {
    89  			return err
    90  		}
    91  		defer func() {
    92  			signal.Stop(p.c)
    93  			close(p.c)
    94  		}()
    95  	}
    96  
    97  	done := make(chan error)
    98  	c := make(chan map[string][]*targetgroup.Group)
    99  	go func() {
   100  		err := p.manager.Run(c)
   101  		if err == nil {
   102  			p.manager.Stop()
   103  		}
   104  		done <- err
   105  	}()
   106  	c <- p.targets
   107  
   108  	// Wait till some exit condition happens
   109  	ticker := time.NewTicker(time.Second)
   110  	defer ticker.Stop()
   111  	// Don't check the process if there's none to check
   112  	if p.cmd == nil {
   113  		ticker.Stop()
   114  	}
   115  
   116  	for {
   117  		select {
   118  		case s := <-p.c:
   119  			if p.cmd != nil {
   120  				_ = process.SendSignal(p.cmd.Process, s)
   121  			} else {
   122  				return nil
   123  			}
   124  		case err := <-done:
   125  			return err
   126  		case <-ticker.C:
   127  			if !process.Exists(p.cmd.Process.Pid) {
   128  				p.logger.Debug("child process exited")
   129  				return p.cmd.Wait()
   130  			}
   131  		}
   132  	}
   133  }