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

     1  package cli
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"os"
     7  	"path/filepath"
     8  
     9  	"github.com/kardianos/service"
    10  	"github.com/sirupsen/logrus"
    11  	"gopkg.in/yaml.v2"
    12  
    13  	"github.com/pyroscope-io/pyroscope/pkg/agent/target"
    14  	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
    15  	"github.com/pyroscope-io/pyroscope/pkg/config"
    16  )
    17  
    18  type agentService struct {
    19  	remote *remote.Remote
    20  	tgtMgr *target.Manager
    21  }
    22  
    23  func newAgentService(logger *logrus.Logger, cfg *config.Agent) (*agentService, error) {
    24  	rc := remote.RemoteConfig{
    25  		AuthToken:              cfg.AuthToken,
    26  		BasicAuthUser:          cfg.BasicAuthUser,
    27  		BasicAuthPassword:      cfg.BasicAuthPassword,
    28  		TenantID:               cfg.TenantID,
    29  		HTTPHeaders:            cfg.Headers,
    30  		UpstreamThreads:        cfg.UpstreamThreads,
    31  		UpstreamAddress:        cfg.ServerAddress,
    32  		UpstreamRequestTimeout: cfg.UpstreamRequestTimeout,
    33  	}
    34  	upstream, err := remote.New(rc, logger)
    35  	if err != nil {
    36  		return nil, fmt.Errorf("upstream configuration: %w", err)
    37  	}
    38  	s := agentService{
    39  		tgtMgr: target.NewManager(logger, upstream, cfg),
    40  		remote: upstream,
    41  	}
    42  	return &s, nil
    43  }
    44  
    45  func (svc *agentService) Start(_ service.Service) error {
    46  	svc.remote.Start()
    47  	svc.tgtMgr.Start()
    48  	return nil
    49  }
    50  
    51  func (svc *agentService) Stop(_ service.Service) error {
    52  	svc.tgtMgr.Stop()
    53  	svc.remote.Stop()
    54  	return nil
    55  }
    56  
    57  // loadAgentConfig is a hack for viper parser, which can't merge maps:
    58  // https://github.com/spf13/viper#accessing-nested-keys.
    59  // TODO(kolesnikovae): find a way to get rid of the function.
    60  func loadAgentConfig(c *config.Agent) error {
    61  	b, err := os.ReadFile(c.Config)
    62  	switch {
    63  	case err == nil:
    64  	case os.IsNotExist(err):
    65  		return nil
    66  	default:
    67  		return err
    68  	}
    69  	var a config.Agent
    70  	if err = yaml.Unmarshal(b, &a); err != nil {
    71  		return err
    72  	}
    73  	// Override tags from config file with flags.
    74  	c.Tags = mergeTags(a.Tags, c.Tags)
    75  	for _, t := range a.Targets {
    76  		t.Tags = mergeTags(t.Tags, c.Tags)
    77  		c.Targets = append(c.Targets, t)
    78  	}
    79  	return nil
    80  }
    81  
    82  // mergeTags creates a new map with tags from a and b.
    83  // Values from b take precedence. Returned map is never nil.
    84  func mergeTags(a, b map[string]string) map[string]string {
    85  	t := make(map[string]string, len(a))
    86  	for k, v := range a {
    87  		t[k] = v
    88  	}
    89  	for k, v := range b {
    90  		t[k] = v
    91  	}
    92  	return t
    93  }
    94  
    95  func createLogger(cfg *config.Agent) (*logrus.Logger, error) {
    96  	if cfg.NoLogging {
    97  		logrus.SetOutput(io.Discard)
    98  		return logrus.StandardLogger(), nil
    99  	}
   100  	l, err := logrus.ParseLevel(cfg.LogLevel)
   101  	if err != nil {
   102  		return nil, fmt.Errorf("parsing log level: %w", err)
   103  	}
   104  	logrus.SetLevel(l)
   105  	if service.Interactive() || cfg.LogFilePath == "" {
   106  		return logrus.StandardLogger(), nil
   107  	}
   108  	f, err := ensureLogFile(cfg.LogFilePath)
   109  	if err != nil {
   110  		return nil, fmt.Errorf("log file: %w", err)
   111  	}
   112  	logrus.SetOutput(f)
   113  	return logrus.StandardLogger(), nil
   114  }
   115  
   116  func ensureLogFile(p string) (*os.File, error) {
   117  	if err := os.MkdirAll(filepath.Dir(p), 0770); err != nil {
   118  		return nil, err
   119  	}
   120  	return os.OpenFile(p, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0666)
   121  }