From 1e34aa359ff6cfe4ba2a44775d261311ea024197 Mon Sep 17 00:00:00 2001
From: Ethan <ZSHYOU@GMAIL.COM>
Date: Thu, 7 Apr 2022 07:11:27 +0800
Subject: [PATCH] feature: use peroperationsamplerupdater first (#2137)

* feature: use peroperationsamplerupdater first; ratelimiter balance init with min(maxbalance, creditspersecond)

* feautre: polish ratelimiter test and add TestRemotelyControlledSampler_multiStrategyResponse

* feature: update changelog regarding PR #2137

* feature: recover ratelimiter first pass strategy

* feature: fix tests of withUpdaters

Co-authored-by: Chester Cheung <cheung.zhy.csu@gmail.com>
---
 CHANGELOG.md                                  |  3 ++
 .../jaegerremote/sampler_remote_options.go    | 10 ++---
 samplers/jaegerremote/sampler_remote_test.go  | 41 ++++++++++++++++++-
 3 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9e39fca8fa4..1d0f274394d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,9 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased]
 
+### Fixed
+- Fixed jaegerremote sampler not behaving properly with per operation strategy set. (#2137)
+
 ## [1.6.0/0.31.0] - 2022-03-28
 
 ### Added
diff --git a/samplers/jaegerremote/sampler_remote_options.go b/samplers/jaegerremote/sampler_remote_options.go
index 354a258a58c..cc80e3aad17 100644
--- a/samplers/jaegerremote/sampler_remote_options.go
+++ b/samplers/jaegerremote/sampler_remote_options.go
@@ -52,12 +52,10 @@ func newConfig(options ...Option) config {
 	for _, option := range options {
 		option.apply(&c)
 	}
-	c.updaters = append(c.updaters,
-		&perOperationSamplerUpdater{
-			MaxOperations:            c.posParams.MaxOperations,
-			OperationNameLateBinding: c.posParams.OperationNameLateBinding,
-		})
-
+	c.updaters = append([]samplerUpdater{&perOperationSamplerUpdater{
+		MaxOperations:            c.posParams.MaxOperations,
+		OperationNameLateBinding: c.posParams.OperationNameLateBinding,
+	}}, c.updaters...)
 	return c
 }
 
diff --git a/samplers/jaegerremote/sampler_remote_test.go b/samplers/jaegerremote/sampler_remote_test.go
index 4fe4e8bee7c..84a5d66d58f 100644
--- a/samplers/jaegerremote/sampler_remote_test.go
+++ b/samplers/jaegerremote/sampler_remote_test.go
@@ -125,7 +125,7 @@ func TestRemoteSamplerOptions(t *testing.T) {
 	assert.Equal(t, 42*time.Second, sampler.samplingRefreshInterval)
 	assert.Same(t, fetcher, sampler.samplingFetcher)
 	assert.Same(t, parser, sampler.samplingParser)
-	assert.Same(t, updaters[0], sampler.updaters[0])
+	assert.EqualValues(t, sampler.updaters[0], &perOperationSamplerUpdater{MaxOperations: 42, OperationNameLateBinding: true})
 }
 
 func TestRemoteSamplerOptionsDefaults(t *testing.T) {
@@ -281,6 +281,45 @@ func TestRemotelyControlledSampler_updateSampler(t *testing.T) {
 	}
 }
 
+func TestRemotelyControlledSampler_multiStrategyResponse(t *testing.T) {
+	agent, sampler := initAgent(t)
+	defer agent.Close()
+	initSampler, ok := sampler.sampler.(*probabilisticSampler)
+	assert.True(t, ok)
+
+	defaultSampingRate := 1.0
+	testUnusedOpName := "unused_op"
+	testUnusedOpSamplingRate := 0.0
+
+	res := &jaeger_api_v2.SamplingStrategyResponse{
+		StrategyType:          jaeger_api_v2.SamplingStrategyType_PROBABILISTIC,
+		ProbabilisticSampling: &jaeger_api_v2.ProbabilisticSamplingStrategy{SamplingRate: defaultSampingRate},
+		OperationSampling: &jaeger_api_v2.PerOperationSamplingStrategies{
+			DefaultSamplingProbability:       defaultSampingRate,
+			DefaultLowerBoundTracesPerSecond: 0.001,
+			PerOperationStrategies: []*jaeger_api_v2.OperationSamplingStrategy{
+				{
+					Operation: testUnusedOpName,
+					ProbabilisticSampling: &jaeger_api_v2.ProbabilisticSamplingStrategy{
+						SamplingRate: testUnusedOpSamplingRate,
+					}},
+			},
+		},
+	}
+
+	agent.AddSamplingStrategy("client app", res)
+	sampler.UpdateSampler()
+	s, ok := sampler.sampler.(*perOperationSampler)
+	assert.True(t, ok)
+	assert.NotEqual(t, initSampler, sampler.sampler, "Sampler should have been updated")
+	assert.Equal(t, defaultSampingRate, s.defaultSampler.SamplingRate())
+
+	result := sampler.ShouldSample(makeSamplingParameters(testMaxID-10, testUnusedOpName))
+	assert.Equal(t, trace.RecordAndSample, result.Decision) // first call always pass
+	result = sampler.ShouldSample(makeSamplingParameters(testMaxID, testUnusedOpName))
+	assert.Equal(t, trace.Drop, result.Decision)
+}
+
 func TestSamplerQueryError(t *testing.T) {
 	agent, sampler := initAgent(t)
 	defer agent.Close()