Fix issue #154: Add ability to change sampling rate (#190)

Co-authored-by: alonlong <alonlong@163.com>

pkg/agent/cli/agent.go

@@ -7,6 +7,7 @@ import (
 	"github.com/pyroscope-io/pyroscope/pkg/agent"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/csock"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
+	"github.com/pyroscope-io/pyroscope/pkg/agent/types"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
 	"github.com/pyroscope-io/pyroscope/pkg/config"
@@ -66,7 +67,7 @@ func (a *Agent) controlSocketHandler(req *csock.Request) *csock.Response {
 			Upstream:         a.u,
 			AppName:          "testapp",
 			ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU, spy.ProfileAllocObjects, spy.ProfileAllocSpace, spy.ProfileInuseObjects, spy.ProfileInuseSpace},
-			SpyName:          "gospy",
+			SpyName:          types.GoSpy,
 			SampleRate:       100,
 			UploadRate:       10 * time.Second,
 			Pid:              0,

pkg/agent/gospy/gospy.go

@@ -37,7 +37,9 @@ func Start(profileType spy.ProfileType, disableGCRuns bool) (spy.Spy, error) {
 		disableGCRuns: disableGCRuns,
 	}
 	if s.profileType == spy.ProfileCPU {
-		_ = pprof.StartCPUProfile(s.buf)
+		if err := pprof.StartCPUProfile(s.buf); err != nil {
+			return nil, err
+		}
 	}
 	return s, nil
 }

pkg/agent/profiler/profiler.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/pyroscope-io/pyroscope/pkg/agent"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
+	"github.com/pyroscope-io/pyroscope/pkg/agent/types"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
 )
 
@@ -20,12 +21,11 @@ var (
 	ProfileInuseSpace   = spy.ProfileInuseSpace
 )
 
-var DefaultProfileTypes = []ProfileType{ProfileCPU, ProfileAllocObjects, ProfileAllocSpace, ProfileInuseObjects, ProfileInuseSpace}
-
 type Config struct {
 	ApplicationName string // e.g backend.purchases
 	ServerAddress   string // e.g http://pyroscope.services.internal:4040
 	AuthToken       string // specify this token when using pyroscope cloud
+	SampleRate      uint32
 	Logger          agent.Logger
 	ProfileTypes    []ProfileType
 	DisableGCRuns   bool // this will disable automatic runtime.GC runs
@@ -38,42 +38,44 @@ type Profiler struct {
 // Start starts continuously profiling go code
 func Start(cfg Config) (*Profiler, error) {
 	if len(cfg.ProfileTypes) == 0 {
-		cfg.ProfileTypes = DefaultProfileTypes
+		cfg.ProfileTypes = types.DefaultProfileTypes
+	}
+	if cfg.SampleRate == 0 {
+		cfg.SampleRate = types.DefaultSampleRate
 	}
+
 	u, err := remote.New(remote.RemoteConfig{
 		AuthToken:              cfg.AuthToken,
 		UpstreamAddress:        cfg.ServerAddress,
 		UpstreamThreads:        4,
 		UpstreamRequestTimeout: 30 * time.Second,
 	})
-
-	u.Logger = cfg.Logger
-
 	if err != nil {
 		return nil, err
 	}
+	u.Logger = cfg.Logger
 
-	// TODO: add sample rate
 	c := agent.SessionConfig{
 		Upstream:         u,
 		AppName:          cfg.ApplicationName,
-		ProfilingTypes:   []ProfileType{ProfileCPU, ProfileAllocObjects, ProfileAllocSpace, ProfileInuseObjects, ProfileInuseSpace},
+		ProfilingTypes:   types.DefaultProfileTypes,
 		DisableGCRuns:    cfg.DisableGCRuns,
-		SpyName:          "gospy",
-		SampleRate:       100,
+		SpyName:          types.GoSpy,
+		SampleRate:       cfg.SampleRate,
 		UploadRate:       10 * time.Second,
 		Pid:              0,
 		WithSubprocesses: false,
 	}
 	sess := agent.NewSession(&c)
-	sess.Logger = cfg.Logger
-	sess.Start()
 
-	p := &Profiler{
-		sess: sess,
+	sess.Logger = cfg.Logger
+	if err := sess.Start(); err != nil {
+		return nil, err
 	}
 
-	return p, nil
+	return &Profiler{
+		sess: sess,
+	}, nil
 }
 
 // Stop stops continious profiling session

pkg/agent/selfprofile.go

@@ -3,20 +3,19 @@ package agent
 import (
 	"time"
 
-	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
+	"github.com/pyroscope-io/pyroscope/pkg/agent/types"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
 	"github.com/pyroscope-io/pyroscope/pkg/config"
 	"github.com/pyroscope-io/pyroscope/pkg/util/atexit"
 )
 
-func SelfProfile(_ *config.Config, u upstream.Upstream, appName string, logger Logger) error {
-	// TODO: sample rate and upload rate should come from config
+func SelfProfile(cfg *config.Config, u upstream.Upstream, appName string, logger Logger) error {
 	c := SessionConfig{
 		Upstream:         u,
 		AppName:          appName,
-		ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU, spy.ProfileAllocObjects, spy.ProfileAllocSpace, spy.ProfileInuseObjects, spy.ProfileInuseSpace},
-		SpyName:          "gospy",
-		SampleRate:       100,
+		ProfilingTypes:   types.DefaultProfileTypes,
+		SpyName:          types.GoSpy,
+		SampleRate:       uint32(cfg.Server.SampleRate),
 		UploadRate:       10 * time.Second,
 		Pid:              0,
 		WithSubprocesses: false,

pkg/agent/session.go

@@ -15,6 +15,7 @@ import (
 	_ "github.com/pyroscope-io/pyroscope/pkg/agent/phpspy"
 	_ "github.com/pyroscope-io/pyroscope/pkg/agent/pyspy"
 	_ "github.com/pyroscope-io/pyroscope/pkg/agent/rbspy"
+	"github.com/pyroscope-io/pyroscope/pkg/agent/types"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
 	"github.com/pyroscope-io/pyroscope/pkg/util/slices"
 
@@ -29,7 +30,7 @@ type ProfileSession struct {
 	upstream   upstream.Upstream
 	appName    string
 	spyName    string
-	sampleRate int
+	sampleRate uint32
 	uploadRate time.Duration
 	pids       []int
 	spies      []spy.Spy
@@ -55,7 +56,7 @@ type SessionConfig struct {
 	ProfilingTypes   []spy.ProfileType
 	DisableGCRuns    bool
 	SpyName          string
-	SampleRate       int
+	SampleRate       uint32
 	UploadRate       time.Duration
 	Pid              int
 	WithSubprocesses bool
@@ -75,7 +76,7 @@ func NewSession(c *SessionConfig) *ProfileSession {
 		withSubprocesses: c.WithSubprocesses,
 	}
 
-	if ps.spyName == "gospy" {
+	if ps.spyName == types.GoSpy {
 		ps.previousTries = make([]*transporttrie.Trie, len(ps.profileTypes))
 		ps.tries = make([]*transporttrie.Trie, len(ps.profileTypes))
 	} else {
@@ -101,11 +102,11 @@ func (ps *ProfileSession) takeSnapshots() {
 			}
 			for i, s := range ps.spies {
 				s.Snapshot(func(stack []byte, v uint64, err error) {
-					if stack != nil && len(stack) > 0 {
+					if len(stack) > 0 {
 						ps.trieMutex.Lock()
 						defer ps.trieMutex.Unlock()
 
-						if ps.spyName == "gospy" {
+						if ps.spyName == types.GoSpy {
 							ps.tries[i].Insert(stack, v, true)
 						} else {
 							ps.tries[0].Insert(stack, v, true)
@@ -130,7 +131,7 @@ func (ps *ProfileSession) takeSnapshots() {
 func (ps *ProfileSession) Start() error {
 	ps.reset()
 
-	if ps.spyName == "gospy" {
+	if ps.spyName == types.GoSpy {
 		for _, pt := range ps.profileTypes {
 			s, err := gospy.Start(pt, ps.disableGCRuns)
 			if err != nil {

pkg/agent/spy/spy.go

@@ -22,6 +22,10 @@ const (
 	ProfileAllocObjects ProfileType = "alloc_objects"
 	ProfileInuseSpace   ProfileType = "inuse_space"
 	ProfileAllocSpace   ProfileType = "alloc_space"
+
+	Go     = "gospy"
+	Python = "pyspy"
+	Ruby   = "rbspy"
 )
 
 func (t ProfileType) IsCumulative() bool {
@@ -94,7 +98,7 @@ func ResolveAutoName(s string) string {
 func SupportedExecSpies() []string {
 	supportedSpies := []string{}
 	for _, s := range SupportedSpies {
-		if s != "gospy" {
+		if s != Go {
 			supportedSpies = append(supportedSpies, s)
 		}
 	}

pkg/agent/types/types.go

@@ -0,0 +1,18 @@
+package types
+
+import "github.com/pyroscope-io/pyroscope/pkg/agent/spy"
+
+const (
+	DefaultSampleRate = 100 // 100 times per
+	GoSpy             = spy.Go
+	PySpy             = spy.Python
+	RbSpy             = spy.Ruby
+)
+
+var DefaultProfileTypes = []spy.ProfileType{
+	spy.ProfileCPU,
+	spy.ProfileAllocObjects,
+	spy.ProfileAllocSpace,
+	spy.ProfileInuseObjects,
+	spy.ProfileInuseSpace,
+}

pkg/agent/upstream/remote/remote.go

@@ -99,7 +99,7 @@ func (u *Remote) uploadProfile(j *upstream.UploadJob) {
 	q.Set("from", strconv.Itoa(int(j.StartTime.Unix())))
 	q.Set("until", strconv.Itoa(int(j.EndTime.Unix())))
 	q.Set("spyName", j.SpyName)
-	q.Set("sampleRate", strconv.Itoa(j.SampleRate))
+	q.Set("sampleRate", strconv.Itoa(int(j.SampleRate)))
 	q.Set("units", j.Units)
 	q.Set("aggregationType", j.AggregationType)
 

pkg/agent/upstream/upstream.go

@@ -11,7 +11,7 @@ type UploadJob struct {
 	StartTime       time.Time
 	EndTime         time.Time
 	SpyName         string
-	SampleRate      int
+	SampleRate      uint32
 	Units           string
 	AggregationType string
 	Trie            *transporttrie.Trie

pkg/cli/cli.go

@@ -178,6 +178,19 @@ func PopulateFlagSet(obj interface{}, flagSet *flag.FlagSet, skip ...string) *So
 				}
 			}
 			flagSet.Uint64Var(val, nameVal, defaultVal, descVal)
+		case reflect.TypeOf(uint(1)):
+			val := fieldV.Addr().Interface().(*uint)
+			var defaultVal uint
+			if defaultValStr == "" {
+				defaultVal = uint(0)
+			} else {
+				out, err := strconv.ParseUint(defaultValStr, 10, 64)
+				if err != nil {
+					logrus.Fatalf("invalid default value: %q (%s)", defaultValStr, nameVal)
+				}
+				defaultVal = uint(out)
+			}
+			flagSet.UintVar(val, nameVal, defaultVal, descVal)
 		default:
 			logrus.Fatalf("type %s is not supported", field.Type)
 		}

pkg/config/config.go

@@ -59,6 +59,7 @@ type Server struct {
 	HideApplications []string `def:"" desc:"please don't use, this will soon be deprecated"`
 
 	OutOfSpaceThreshold bytesize.ByteSize `def:"512MB" desc:"Threshold value to consider out of space in bytes"`
+	SampleRate          uint              `def:"100" desc:"sample rate for the profiler in Hz. 100 means reading 100 times per second"`
 }
 
 type Convert struct {
@@ -79,6 +80,7 @@ type DbManager struct {
 type Exec struct {
 	SpyName                string        `def:"auto" desc:"name of the profiler you want to use. Supported ones are: <supportedProfilers>"`
 	ApplicationName        string        `def:"" desc:"application name used when uploading profiling data"`
+	SampleRate             uint          `def:"100" desc:"sample rate for the profiler in Hz. 100 means reading 100 times per second"`
 	DetectSubprocesses     bool          `def:"true" desc:"makes pyroscope keep track of and profile subprocesses of the main process"`
 	LogLevel               string        `def:"info" desc:"log level: debug|info|warn|error"`
 	ServerAddress          string        `def:"http://localhost:4040" desc:"address of the pyroscope server"`

pkg/exec/cli.go

@@ -18,6 +18,7 @@ import (
 	"github.com/pyroscope-io/pyroscope/pkg/agent"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/pyspy"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
+	"github.com/pyroscope-io/pyroscope/pkg/agent/types"
 	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream/remote"
 	"github.com/pyroscope-io/pyroscope/pkg/config"
 	"github.com/pyroscope-io/pyroscope/pkg/util/atexit"
@@ -143,13 +144,17 @@ func Cli(cfg *config.Config, args []string) error {
 		"detect-subprocesses": cfg.Exec.DetectSubprocesses,
 	}).Debug("starting agent session")
 
-	// TODO: add sample rate, make it configurable
+	// if the sample rate is zero, use the default value
+	if cfg.Exec.SampleRate == 0 {
+		cfg.Exec.SampleRate = types.DefaultSampleRate
+	}
+
 	sess := agent.NewSession(&agent.SessionConfig{
 		Upstream:         u,
 		AppName:          cfg.Exec.ApplicationName,
 		ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU},
 		SpyName:          spyName,
-		SampleRate:       100,
+		SampleRate:       uint32(cfg.Exec.SampleRate),
 		UploadRate:       10 * time.Second,
 		Pid:              pid,
 		WithSubprocesses: cfg.Exec.DetectSubprocesses,
@@ -219,7 +224,7 @@ func waitForProcessToExit(pid int) {
 }
 
 func performChecks(spyName string) error {
-	if spyName == "gospy" {
+	if spyName == types.GoSpy {
 		return fmt.Errorf("gospy can not profile other processes. See our documentation on using gospy: %s", color.GreenString("https://pyroscope.io/docs/"))
 	}
 

pkg/server/ingest.go

@@ -6,6 +6,7 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/pyroscope-io/pyroscope/pkg/agent/types"
 	"github.com/pyroscope-io/pyroscope/pkg/convert"
 	"github.com/pyroscope-io/pyroscope/pkg/storage"
 	"github.com/pyroscope-io/pyroscope/pkg/storage/tree"
@@ -17,7 +18,7 @@ type ingestParams struct {
 	parserFunc      func(io.Reader) (*tree.Tree, error)
 	storageKey      *storage.Key
 	spyName         string
-	sampleRate      int
+	sampleRate      uint32
 	units           string
 	aggregationType string
 	modifiers       []string
@@ -67,10 +68,15 @@ func ingestParamsFromRequest(r *http.Request) *ingestParams {
 	}
 
 	if sr := q.Get("sampleRate"); sr != "" {
-		// TODO: error handling
-		ip.sampleRate, _ = strconv.Atoi(sr)
+		sampleRate, err := strconv.Atoi(sr)
+		if err != nil {
+			logrus.WithField("err", err).Errorf("invalid sample rate: %v", sr)
+			ip.sampleRate = types.DefaultSampleRate
+		} else {
+			ip.sampleRate = uint32(sampleRate)
+		}
 	} else {
-		ip.sampleRate = 100
+		ip.sampleRate = types.DefaultSampleRate
 	}
 
 	if sn := q.Get("spyName"); sn != "" {

pkg/storage/segment/segment.go

@@ -152,7 +152,7 @@ type Segment struct {
 	durations  []time.Duration
 
 	spyName         string
-	sampleRate      int
+	sampleRate      uint32
 	units           string
 	aggregationType string
 }
@@ -262,7 +262,7 @@ func (s *Segment) Get(st, et time.Time, cb func(depth int, samples, writes uint6
 
 // TODO: this should be refactored
 
-func (s *Segment) SetMetadata(spyName string, sampleRate int, units, aggregationType string) {
+func (s *Segment) SetMetadata(spyName string, sampleRate uint32, units, aggregationType string) {
 	s.spyName = spyName
 	s.sampleRate = sampleRate
 	s.units = units
@@ -273,7 +273,7 @@ func (s *Segment) SpyName() string {
 	return s.spyName
 }
 
-func (s *Segment) SampleRate() int {
+func (s *Segment) SampleRate() uint32 {
 	return s.sampleRate
 }
 

pkg/storage/segment/serialization.go

@@ -15,7 +15,7 @@ const currentVersion = 2
 
 func (s *Segment) populateFromMetadata(metadata map[string]interface{}) {
 	if v, ok := metadata["sampleRate"]; ok {
-		s.sampleRate = int(v.(float64))
+		s.sampleRate = uint32(v.(float64))
 	}
 	if v, ok := metadata["spyName"]; ok {
 		s.spyName = v.(string)