diff --git a/sampler_test.go b/sampler_test.go index 48b4b4a7..c3914f50 100644 --- a/sampler_test.go +++ b/sampler_test.go @@ -466,7 +466,11 @@ func TestUpdateSampler(t *testing.T) { assert.NotEqual(t, initSampler, sampler.sampler, "Sampler should have been updated") assert.Equal(t, test.defaultProbability, s.defaultSampler.SamplingRate()) + // First call is always sampled sampled, tags := sampler.IsSampled(TraceID{Low: testMaxID + 10}, testOperationName) + assert.True(t, sampled) + + sampled, tags = sampler.IsSampled(TraceID{Low: testMaxID + 10}, testOperationName) assert.False(t, sampled) sampled, tags = sampler.IsSampled(TraceID{Low: testMaxID - 10}, testOperationName) assert.True(t, sampled) diff --git a/utils/rate_limiter.go b/utils/rate_limiter.go index c14e2ba9..58927586 100644 --- a/utils/rate_limiter.go +++ b/utils/rate_limiter.go @@ -56,7 +56,7 @@ type rateLimiter struct { func NewRateLimiter(creditsPerSecond, maxBalance float64) RateLimiter { return &rateLimiter{ creditsPerSecond: creditsPerSecond, - balance: creditsPerSecond, + balance: maxBalance, maxBalance: maxBalance, lastTick: time.Now(), timeNow: time.Now} diff --git a/utils/rate_limiter_test.go b/utils/rate_limiter_test.go index 307beca6..02fab159 100644 --- a/utils/rate_limiter_test.go +++ b/utils/rate_limiter_test.go @@ -69,7 +69,8 @@ func TestMaxBalance(t *testing.T) { limiter.(*rateLimiter).timeNow = func() time.Time { return ts } - assert.False(t, limiter.CheckCredit(1.0)) + // on initialization, should have enough credits for 1 message + assert.True(t, limiter.CheckCredit(1.0)) // move time 20s forward, enough to accumulate credits for 2 messages, but it should still be capped at 1 limiter.(*rateLimiter).timeNow = func() time.Time {