From 5b5fd159dfdf501a710f911c572eface83ea0858 Mon Sep 17 00:00:00 2001
From: Achille <achille.roussel@gmail.com>
Date: Fri, 12 May 2023 18:52:34 -0700
Subject: [PATCH] sampler: do not share sampling state across function calls
 (#73)

This PR is another follow-up from #71, which contains two fixes:

The first fix consists in not applying the sampling scale to time
measures of the CPU profiler. This is needed because the time deltas
between call and return of functions are not affected by the sampling
rate, it's an absolute value.

The second fix changes the sampling mechanism to track state per
function instead of sharing it across all functions, which was causing
too much skew. We now record `1 / sample rate` calls to each function
instead of `1 / sample rate` function calls, which gives a much more
accurate profiles.

---------

Signed-off-by: Achille Roussel <achille.roussel@gmail.com>
---
 cpu.go                            |  9 ++++++++-
 mem.go                            |  5 ++++-
 sampler.go => sample.go           | 28 +++++++++++++---------------
 sampler_test.go => sample_test.go |  0
 wzprof.go                         |  6 +++---
 5 files changed, 28 insertions(+), 20 deletions(-)
 rename sampler.go => sample.go (89%)
 rename sampler_test.go => sample_test.go (100%)

diff --git a/cpu.go b/cpu.go
index 195cecd..1839368 100644
--- a/cpu.go
+++ b/cpu.go
@@ -104,7 +104,14 @@ func (p *CPUProfiler) StopProfile(sampleRate float64, symbols Symbolizer) *profi
 		}
 	}
 
-	return buildProfile(sampleRate, symbols, samples, start, duration, p.SampleType())
+	ratios := []float64{
+		1 / sampleRate,
+		// Time values are not influenced by the sampling rate so we don't have
+		// to scale them out.
+		1,
+	}
+
+	return buildProfile(symbols, samples, start, duration, p.SampleType(), ratios)
 }
 
 // Name returns "profile" to match the name of the CPU profiler in pprof.
diff --git a/mem.go b/mem.go
index 135a57b..c06cf04 100644
--- a/mem.go
+++ b/mem.go
@@ -67,7 +67,10 @@ func NewMemoryProfiler(options ...MemoryProfilerOption) *MemoryProfiler {
 // NewProfile takes a snapshot of the current memory allocation state and builds
 // a profile representing the state of the program memory.
 func (p *MemoryProfiler) NewProfile(sampleRate float64, symbols Symbolizer) *profile.Profile {
-	return buildProfile(sampleRate, symbols, p.snapshot(), p.start, time.Since(p.start), p.SampleType())
+	ratio := 1 / sampleRate
+	return buildProfile(symbols, p.snapshot(), p.start, time.Since(p.start), p.SampleType(),
+		[]float64{ratio, ratio, ratio, ratio},
+	)
 }
 
 // Name returns "allocs" to match the name of the memory profiler in pprof.
diff --git a/sampler.go b/sample.go
similarity index 89%
rename from sampler.go
rename to sample.go
index 24af81d..7994aac 100644
--- a/sampler.go
+++ b/sample.go
@@ -23,18 +23,19 @@ func Sample(sampleRate float64, factory experimental.FunctionListenerFactory) ex
 	if sampleRate >= 1 {
 		return factory
 	}
-	sampler := new(sampler)
-	sampler.cycle = uint64(math.Ceil(1 / sampleRate))
-	sampler.count = sampler.cycle
+	cycle := uint32(math.Ceil(1 / sampleRate))
 	return experimental.FunctionListenerFactoryFunc(func(def api.FunctionDefinition) experimental.FunctionListener {
 		lstn := factory.NewFunctionListener(def)
 		if lstn == nil {
 			return nil
 		}
-		return &sampledFunctionListener{
-			sampler: sampler,
-			lstn:    lstn,
+		sampled := &sampledFunctionListener{
+			cycle: cycle,
+			count: cycle,
+			lstn:  lstn,
 		}
+		sampled.stack.bits = sampled.bits[:]
+		return sampled
 	})
 }
 
@@ -44,15 +45,12 @@ func (emptyFunctionListenerFactory) NewFunctionListener(api.FunctionDefinition)
 	return nil
 }
 
-type sampler struct {
-	count uint64
-	cycle uint64
-	stack bitstack
-}
-
 type sampledFunctionListener struct {
-	*sampler
-	lstn experimental.FunctionListener
+	count uint32
+	cycle uint32
+	bits  [1]uint64
+	stack bitstack
+	lstn  experimental.FunctionListener
 }
 
 func (s *sampledFunctionListener) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, stack experimental.StackIterator) {
@@ -80,8 +78,8 @@ func (s *sampledFunctionListener) Abort(ctx context.Context, mod api.Module, def
 }
 
 type bitstack struct {
-	bits []uint64
 	size uint
+	bits []uint64
 }
 
 func (s *bitstack) push(bit uint) {
diff --git a/sampler_test.go b/sample_test.go
similarity index 100%
rename from sampler_test.go
rename to sample_test.go
diff --git a/wzprof.go b/wzprof.go
index 49a0a9a..87bc40a 100644
--- a/wzprof.go
+++ b/wzprof.go
@@ -271,7 +271,7 @@ type sampleType interface {
 	sampleValue() []int64
 }
 
-func buildProfile[T sampleType](sampleRate float64, symbols Symbolizer, samples map[uint64]T, start time.Time, duration time.Duration, sampleType []*profile.ValueType) *profile.Profile {
+func buildProfile[T sampleType](symbols Symbolizer, samples map[uint64]T, start time.Time, duration time.Duration, sampleType []*profile.ValueType, ratios []float64) *profile.Profile {
 	prof := &profile.Profile{
 		SampleType:    sampleType,
 		Sample:        make([]*profile.Sample, 0, len(samples)),
@@ -325,8 +325,8 @@ func buildProfile[T sampleType](sampleRate float64, symbols Symbolizer, samples
 		prof.Function[fn.ID-1] = fn
 	}
 
-	if sampleRate < 1 {
-		prof.Scale(1 / sampleRate)
+	if err := prof.ScaleN(ratios[:len(sampleType)]); err != nil {
+		panic(err)
 	}
 	return prof
 }