From 7b6fef333f4c6480d9f3e24ebb2923e8a439f3f4 Mon Sep 17 00:00:00 2001
From: Kent Quirk <kentquirk@honeycomb.io>
Date: Fri, 9 Jun 2023 15:20:04 -0400
Subject: [PATCH] Update dynamic samplers to use GetSampleRateMulti

---
 sample/dynamic.go             |  4 +++-
 sample/dynamic_ema.go         |  4 +++-
 sample/ema_throughput.go      |  4 +++-
 sample/totalthroughput.go     |  4 +++-
 sample/windowed_throughput.go | 10 +++++++++-
 5 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/sample/dynamic.go b/sample/dynamic.go
index 2f3c3232fa..5021940ee5 100644
--- a/sample/dynamic.go
+++ b/sample/dynamic.go
@@ -65,7 +65,8 @@ func (d *DynamicSampler) Start() error {
 
 func (d *DynamicSampler) GetSampleRate(trace *types.Trace) (rate uint, keep bool, reason string, key string) {
 	key = d.key.build(trace)
-	rate = uint(d.dynsampler.GetSampleRate(key))
+	count := int(trace.DescendantCount())
+	rate = uint(d.dynsampler.GetSampleRateMulti(key, count))
 	if rate < 1 { // protect against dynsampler being broken even though it shouldn't be
 		rate = 1
 	}
@@ -75,6 +76,7 @@ func (d *DynamicSampler) GetSampleRate(trace *types.Trace) (rate uint, keep bool
 		"sample_rate": rate,
 		"sample_keep": shouldKeep,
 		"trace_id":    trace.TraceID,
+		"span_count":  count,
 	}).Logf("got sample rate and decision")
 	if shouldKeep {
 		d.Metrics.Increment(d.prefix + "num_kept")
diff --git a/sample/dynamic_ema.go b/sample/dynamic_ema.go
index d698fc63e0..9935902c97 100644
--- a/sample/dynamic_ema.go
+++ b/sample/dynamic_ema.go
@@ -73,7 +73,8 @@ func (d *EMADynamicSampler) Start() error {
 
 func (d *EMADynamicSampler) GetSampleRate(trace *types.Trace) (rate uint, keep bool, reason string, key string) {
 	key = d.key.build(trace)
-	rate = uint(d.dynsampler.GetSampleRate(key))
+	count := int(trace.DescendantCount())
+	rate = uint(d.dynsampler.GetSampleRateMulti(key, count))
 	if rate < 1 { // protect against dynsampler being broken even though it shouldn't be
 		rate = 1
 	}
@@ -83,6 +84,7 @@ func (d *EMADynamicSampler) GetSampleRate(trace *types.Trace) (rate uint, keep b
 		"sample_rate": rate,
 		"sample_keep": shouldKeep,
 		"trace_id":    trace.TraceID,
+		"span_count":  count,
 	}).Logf("got sample rate and decision")
 	if shouldKeep {
 		d.Metrics.Increment(d.prefix + "num_kept")
diff --git a/sample/ema_throughput.go b/sample/ema_throughput.go
index ad6e8f67b1..e047aaa728 100644
--- a/sample/ema_throughput.go
+++ b/sample/ema_throughput.go
@@ -77,7 +77,8 @@ func (d *EMAThroughputSampler) Start() error {
 
 func (d *EMAThroughputSampler) GetSampleRate(trace *types.Trace) (rate uint, keep bool, reason string, key string) {
 	key = d.key.build(trace)
-	rate = uint(d.dynsampler.GetSampleRate(key))
+	count := int(trace.DescendantCount())
+	rate = uint(d.dynsampler.GetSampleRateMulti(key, count))
 	if rate < 1 { // protect against dynsampler being broken even though it shouldn't be
 		rate = 1
 	}
@@ -87,6 +88,7 @@ func (d *EMAThroughputSampler) GetSampleRate(trace *types.Trace) (rate uint, kee
 		"sample_rate": rate,
 		"sample_keep": shouldKeep,
 		"trace_id":    trace.TraceID,
+		"span_count":  count,
 	}).Logf("got sample rate and decision")
 	if shouldKeep {
 		d.Metrics.Increment(d.prefix + "num_kept")
diff --git a/sample/totalthroughput.go b/sample/totalthroughput.go
index 7834954c9e..713c146054 100644
--- a/sample/totalthroughput.go
+++ b/sample/totalthroughput.go
@@ -69,7 +69,8 @@ func (d *TotalThroughputSampler) Start() error {
 
 func (d *TotalThroughputSampler) GetSampleRate(trace *types.Trace) (rate uint, keep bool, reason string, key string) {
 	key = d.key.build(trace)
-	rate = uint(d.dynsampler.GetSampleRate(key))
+	count := int(trace.DescendantCount())
+	rate = uint(d.dynsampler.GetSampleRateMulti(key, count))
 	if rate < 1 { // protect against dynsampler being broken even though it shouldn't be
 		rate = 1
 	}
@@ -79,6 +80,7 @@ func (d *TotalThroughputSampler) GetSampleRate(trace *types.Trace) (rate uint, k
 		"sample_rate": rate,
 		"sample_keep": shouldKeep,
 		"trace_id":    trace.TraceID,
+		"span_count":  count,
 	}).Logf("got sample rate and decision")
 	if shouldKeep {
 		d.Metrics.Increment(d.prefix + "num_kept")
diff --git a/sample/windowed_throughput.go b/sample/windowed_throughput.go
index c971091f44..d3e6c1aab3 100644
--- a/sample/windowed_throughput.go
+++ b/sample/windowed_throughput.go
@@ -65,11 +65,19 @@ func (d *WindowedThroughputSampler) Start() error {
 
 func (d *WindowedThroughputSampler) GetSampleRate(trace *types.Trace) (rate uint, keep bool, reason string, key string) {
 	key = d.key.build(trace)
-	rate = uint(d.dynsampler.GetSampleRate(key))
+	count := int(trace.DescendantCount())
+	rate = uint(d.dynsampler.GetSampleRateMulti(key, count))
 	if rate < 1 { // protect against dynsampler being broken even though it shouldn't be
 		rate = 1
 	}
 	shouldKeep := rand.Intn(int(rate)) == 0
+	d.Logger.Debug().WithFields(map[string]interface{}{
+		"sample_key":  key,
+		"sample_rate": rate,
+		"sample_keep": shouldKeep,
+		"trace_id":    trace.TraceID,
+		"span_count":  count,
+	}).Logf("got sample rate and decision")
 	if shouldKeep {
 		d.Metrics.Increment(d.prefix + "num_kept")
 	} else {