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

@@ -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.

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.

sampler.go → 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) {

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
 }