Merge pull request #102 from pyroscope-io/gospy-logging-fix

Improves logging in gospy agent

Author: petethepig

.github/workflows/tests.yml

@@ -19,6 +19,8 @@ jobs:
           key: ${{ runner.os }}-go-${{ hashFiles('**/go.sum') }}
           restore-keys: |
             ${{ runner.os }}-go-
+      - name: Ensure logrus is not used by pkg/agent/profiler
+        run: make ensure-logrus-not-used
       - name: Run Go tests
         run: go test ./...
       - name: Install Webapp dependencies

Makefile

@@ -72,6 +72,11 @@ embedded-assets: install-dev-tools $(shell echo $(EMBEDDED_ASSETS_DEPS))
 lint:
 	revive -config revive.toml -formatter stylish ./...
 
+.PHONY: ensure-logrus-not-used
+ensure-logrus-not-used:
+	@! godepgraph -nostdlib -s ./pkg/agent/profiler/ | grep ' -> "github.com/sirupsen/logrus' \
+		|| (echo "\n^ ERROR: make sure ./pkg/agent/profiler/ does not depend on logrus. We don't want users' logs to be tainted. Talk to @petethepig if have questions\n" &1>2; exit 1)
+
 .PHONY: unused
 unused:
 	staticcheck -f stylish -unused.whole-program ./...

pkg/agent/agent.go pkg/agent/cli/agent.gopkg/agent/agent.go renamed to pkg/agent/cli/agent.go

@@ -1,21 +1,21 @@
-package agent
+package cli
 
 import (
 	"os"
 
-	"github.com/sirupsen/logrus"
-
+	"github.com/pyroscope-io/pyroscope/pkg/agent"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/csock"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
 	"github.com/pyroscope-io/pyroscope/pkg/config"
 	"github.com/pyroscope-io/pyroscope/pkg/util/id"
+	"github.com/sirupsen/logrus"
 )
 
 type Agent struct {
 	cfg            *config.Config
 	cs             *csock.CSock
-	activeProfiles map[int]*ProfileSession
+	activeProfiles map[int]*agent.ProfileSession
 	id             id.ID
 	u              upstream.Upstream
 }
@@ -27,25 +27,26 @@ func New(cfg *config.Config) *Agent {
 		UpstreamAddress:        cfg.Agent.ServerAddress,
 		UpstreamRequestTimeout: cfg.Agent.UpstreamRequestTimeout,
 	})
+	r.Logger = logrus.StandardLogger()
 	return &Agent{
 		cfg:            cfg,
-		activeProfiles: make(map[int]*ProfileSession),
+		activeProfiles: make(map[int]*agent.ProfileSession),
 		u:              r,
 	}
 }
 
-func (a *Agent) Start() {
+func (a *Agent) Start() error {
 	sockPath := a.cfg.Agent.UNIXSocketPath
 	cs, err := csock.NewUnixCSock(sockPath, a.controlSocketHandler)
 	if err != nil {
-		logrus.Fatal(err)
+		return err
 	}
 	a.cs = cs
 	defer os.Remove(sockPath)
 
-	go SelfProfile(a.cfg, a.u, "pyroscope.agent.cpu{}")
-	logrus.WithField("addr", cs.CanonicalAddr()).Info("Starting control socket")
+	go agent.SelfProfile(a.cfg, a.u, "pyroscope.agent.cpu{}", logrus.StandardLogger())
 	cs.Start()
+	return nil
 }
 
 func (a *Agent) Stop() {
@@ -59,7 +60,8 @@ func (a *Agent) controlSocketHandler(req *csock.Request) *csock.Response {
 		// TODO: pass withSubprocesses from somewhere
 		// TODO: pass appName from somewhere
 		// TODO: add sample rate
-		s := NewSession(a.u, "testapp.cpu", "gospy", 100, 0, false)
+		s := agent.NewSession(a.u, "testapp.cpu", "gospy", 100, 0, false)
+		s.Logger = logrus.StandardLogger()
 		a.activeProfiles[profileID] = s
 		s.Start()
 		return &csock.Response{ProfileID: profileID}

pkg/agent/logger.go

@@ -0,0 +1,10 @@
+package agent
+
+// Logger is an interface that library users can use
+// It is based on logrus, but much smaller — That's because we don't want library users to have to implement
+//   all of the logrus's methods
+type Logger interface {
+	Infof(format string, args ...interface{})
+	Debugf(format string, args ...interface{})
+	Errorf(format string, args ...interface{})
+}

pkg/agent/profiler/profiler.go

@@ -13,6 +13,7 @@ type Config struct {
 	ApplicationName string // e.g backend.purchases
 	ServerAddress   string // e.g http://pyroscope.services.internal:4040
 	AuthToken       string
+	Logger          agent.Logger
 }
 
 type Profiler struct {
@@ -27,11 +28,16 @@ func Start(cfg Config) (*Profiler, error) {
 		UpstreamThreads:        4,
 		UpstreamRequestTimeout: 30 * time.Second,
 	})
+
+	u.Logger = cfg.Logger
+
 	if err != nil {
 		return nil, err
 	}
+
 	// TODO: add sample rate
 	sess := agent.NewSession(u, cfg.ApplicationName, "gospy", 100, 0, false)
+	sess.Logger = cfg.Logger
 	sess.Start()
 
 	p := &Profiler{

pkg/agent/selfprofile.go

@@ -4,18 +4,19 @@ import (
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
 	"github.com/pyroscope-io/pyroscope/pkg/config"
 	"github.com/pyroscope-io/pyroscope/pkg/util/atexit"
-	"github.com/sirupsen/logrus"
 )
 
-func SelfProfile(_ *config.Config, u upstream.Upstream, appName string) {
+func SelfProfile(_ *config.Config, u upstream.Upstream, appName string, logger Logger) error {
 	// TODO: add sample rate
 	s := NewSession(u, appName, "gospy", 100, 0, false)
 	err := s.Start()
 
+	s.Logger = logger
+
 	if err != nil {
-		logrus.Errorf("failed to start profiling session: %s", err)
-		return
+		return err
 	}
 
 	atexit.Register(s.Stop)
+	return nil
 }

pkg/agent/session.go

@@ -13,7 +13,6 @@ import (
 	_ "github.com/pyroscope-io/pyroscope/pkg/agent/pyspy"
 	_ "github.com/pyroscope-io/pyroscope/pkg/agent/rbspy"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
-	"github.com/sirupsen/logrus"
 
 	// revive:enable:blank-imports
 
@@ -36,6 +35,8 @@ type ProfileSession struct {
 
 	startTime time.Time
 	stopTime  time.Time
+
+	Logger Logger
 }
 
 func NewSession(upstream upstream.Upstream, appName string, spyName string, sampleRate int, pid int, withSubprocesses bool) *ProfileSession {
@@ -155,15 +156,13 @@ func (ps *ProfileSession) addSubprocesses() {
 			ps.pids = append(ps.pids, newPid)
 			newSpy, err := spy.SpyFromName(ps.spyName, newPid)
 			if err != nil {
-				logrus.WithFields(logrus.Fields{
-					"spy-name": ps.spyName,
-					"pid":      newPid,
-				}).Error("failed to initialize a spy")
+				if ps.Logger != nil {
+					ps.Logger.Errorf("failed to initialize a spy %d [%s]", newPid, ps.spyName)
+				}
 			} else {
-				logrus.WithFields(logrus.Fields{
-					"spy-name": ps.spyName,
-					"pid":      newPid,
-				}).Debug("started spy for subprocess")
+				if ps.Logger != nil {
+					ps.Logger.Debugf("started spy for subprocess %d [%s]", newPid, ps.spyName)
+				}
 				ps.spies = append(ps.spies, newSpy)
 			}
 		}

pkg/agent/upstream/remote/remote.go

@@ -12,8 +12,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/sirupsen/logrus"
-
+	"github.com/pyroscope-io/pyroscope/pkg/agent"
 	"github.com/pyroscope-io/pyroscope/pkg/structs/transporttrie"
 )
 
@@ -34,6 +33,8 @@ type Remote struct {
 	todo   chan *uploadJob
 	done   chan *sync.WaitGroup
 	client *http.Client
+
+	Logger agent.Logger
 }
 
 type RemoteConfig struct {
@@ -97,7 +98,9 @@ func (u *Remote) Upload(name string, startTime time.Time, endTime time.Time, spy
 	select {
 	case u.todo <- job:
 	default:
-		logrus.Error("Remote upload queue is full, dropping a profile")
+		if u.Logger != nil {
+			u.Logger.Errorf("Remote upload queue is full, dropping a profile")
+		}
 	}
 }
 
@@ -115,11 +118,15 @@ func (u *Remote) uploadProfile(j *uploadJob) {
 	urlObj.Path = path.Join(urlObj.Path, "/ingest")
 	urlObj.RawQuery = q.Encode()
 	buf := j.t.Bytes()
-	logrus.Info("uploading at ", urlObj.String())
+	if u.Logger != nil {
+		u.Logger.Infof("uploading at %s", urlObj.String())
+	}
 
 	req, err := http.NewRequest("POST", urlObj.String(), bytes.NewReader(buf))
 	if err != nil {
-		logrus.Error("Error happened when uploading a profile:", err)
+		if u.Logger != nil {
+			u.Logger.Errorf("Error happened when uploading a profile: %v", err)
+		}
 		return
 	}
 	req.Header.Set("Content-Type", "binary/octet-stream+trie")
@@ -129,15 +136,19 @@ func (u *Remote) uploadProfile(j *uploadJob) {
 	resp, err := u.client.Do(req)
 
 	if err != nil {
-		logrus.Error("Error happened when uploading a profile:", err)
+		if u.Logger != nil {
+			u.Logger.Errorf("Error happened when uploading a profile: %v", err)
+		}
 		return
 	}
 
 	if resp != nil {
 		defer resp.Body.Close()
 		_, err := ioutil.ReadAll(resp.Body)
 		if err != nil {
-			logrus.Error("Error happened while reading server response:", err)
+			if u.Logger != nil {
+				u.Logger.Errorf("Error happened while reading server response: %v", err)
+			}
 		}
 	}
 }

pkg/cli/cli.go

@@ -177,7 +177,6 @@ func resolvePath(path string) string {
 
 func Start(cfg *config.Config) error {
 	var (
-		agentFlagSet     = flag.NewFlagSet("pyroscope agent", flag.ExitOnError)
 		serverFlagSet    = flag.NewFlagSet("pyroscope server", flag.ExitOnError)
 		convertFlagSet   = flag.NewFlagSet("pyroscope convert", flag.ExitOnError)
 		execFlagSet      = flag.NewFlagSet("pyroscope exec", flag.ExitOnError)
@@ -187,7 +186,6 @@ func Start(cfg *config.Config) error {
 		rootFlagSet      = flag.NewFlagSet("pyroscope", flag.ExitOnError)
 	)
 
-	agentSortedFlags := PopulateFlagSet(&cfg.Agent, agentFlagSet)
 	serverSortedFlags := PopulateFlagSet(&cfg.Server, serverFlagSet)
 	convertSortedFlags := PopulateFlagSet(&cfg.Convert, convertFlagSet)
 	execSortedFlags := PopulateFlagSet(&cfg.Exec, execFlagSet, "pid")
@@ -203,15 +201,6 @@ func Start(cfg *config.Config) error {
 		ff.WithConfigFileFlag("config"),
 	}
 
-	agentCmd := &ffcli.Command{
-		UsageFunc:  agentSortedFlags.printUsage,
-		Options:    options,
-		Name:       "agent",
-		ShortUsage: "pyroscope agent [flags]",
-		ShortHelp:  "starts pyroscope agent. Run this one on the machines you want to profile",
-		FlagSet:    agentFlagSet,
-	}
-
 	serverCmd := &ffcli.Command{
 		UsageFunc:  serverSortedFlags.printUsage,
 		Options:    options,
@@ -272,7 +261,6 @@ func Start(cfg *config.Config) error {
 		ShortUsage: "pyroscope [flags] <subcommand>",
 		FlagSet:    rootFlagSet,
 		Subcommands: []*ffcli.Command{
-			agentCmd,
 			convertCmd,
 			serverCmd,
 			execCmd,
@@ -282,15 +270,6 @@ func Start(cfg *config.Config) error {
 		},
 	}
 
-	agentCmd.Exec = func(_ context.Context, args []string) error {
-		if l, err := logrus.ParseLevel(cfg.Agent.LogLevel); err == nil {
-			logrus.SetLevel(l)
-		}
-		a := agent.New(cfg)
-		atexit.Register(a.Stop)
-		a.Start()
-		return nil
-	}
 	serverCmd.Exec = func(_ context.Context, args []string) error {
 		if l, err := logrus.ParseLevel(cfg.Server.LogLevel); err == nil {
 			logrus.SetLevel(l)
@@ -374,7 +353,7 @@ func startServer(cfg *config.Config) {
 		panic(err)
 	}
 	u := direct.New(cfg, s)
-	go agent.SelfProfile(cfg, u, "pyroscope.server.cpu{}")
+	go agent.SelfProfile(cfg, u, "pyroscope.server.cpu{}", logrus.StandardLogger())
 	go printRAMUsage()
 	go printDiskUsage(cfg)
 	c := server.New(cfg, s)

pkg/dbmanager/cli.go

@@ -9,6 +9,7 @@ import (
 	"github.com/pyroscope-io/pyroscope/pkg/config"
 	"github.com/pyroscope-io/pyroscope/pkg/storage"
 	"github.com/pyroscope-io/pyroscope/pkg/util/atexit"
+	"github.com/sirupsen/logrus"
 
 	"github.com/cheggaaa/pb/v3"
 )
@@ -69,7 +70,7 @@ func copyData(cfg *config.Config) error {
 
 	if cfg.DbManager.EnableProfiling {
 		u := direct.New(cfg, s)
-		go agent.SelfProfile(cfg, u, "pyroscope.dbmanager.cpu{}")
+		go agent.SelfProfile(cfg, u, "pyroscope.dbmanager.cpu{}", logrus.StandardLogger())
 	}
 
 	st := srcSt