github.com/thanos-io/thanos@v0.32.5/pkg/tracing/jaeger/remote.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  // This file contains utility functions required to implement a rate limiting sampler.
     5  // Since these are part of the 'internal' package folder in the OpenTelemetry-go repo, I have added
     6  // these functions here to be used by the parent 'tracing/jaeger' package.
     7  // Ref: https://github.com/open-telemetry/opentelemetry-go-contrib/blob/main/samplers/jaegerremote/internal/utils/rate_limiter.go
     8  
     9  package jaeger
    10  
    11  import (
    12  	"math"
    13  	"sync"
    14  	"time"
    15  
    16  	tracesdk "go.opentelemetry.io/otel/sdk/trace"
    17  	oteltrace "go.opentelemetry.io/otel/trace"
    18  )
    19  
    20  type RateLimiter struct {
    21  	lock sync.Mutex
    22  
    23  	creditsPerSecond float64
    24  	balance          float64
    25  	maxBalance       float64
    26  	lastTick         time.Time
    27  
    28  	timeNow func() time.Time
    29  }
    30  
    31  type rateLimitingSampler struct {
    32  	rateLimiter        *RateLimiter
    33  	maxTracesPerSecond float64
    34  }
    35  
    36  // NewRateLimiter creates a new RateLimiter.
    37  func NewRateLimiter(creditsPerSecond, maxBalance float64) *RateLimiter {
    38  	return &RateLimiter{
    39  		creditsPerSecond: creditsPerSecond,
    40  		balance:          maxBalance,
    41  		maxBalance:       maxBalance,
    42  		lastTick:         time.Now(),
    43  		timeNow:          time.Now,
    44  	}
    45  }
    46  
    47  // CheckCredit tries to reduce the current balance by itemCost provided that the current balance
    48  // is not lest than itemCost.
    49  func (rl *RateLimiter) CheckCredit(itemCost float64) bool {
    50  	rl.lock.Lock()
    51  	defer rl.lock.Unlock()
    52  
    53  	// if we have enough credits to pay for current item, then reduce balance and allow
    54  	if rl.balance >= itemCost {
    55  		rl.balance -= itemCost
    56  		return true
    57  	}
    58  	// otherwise check if balance can be increased due to time elapsed, and try again
    59  	rl.updateBalance()
    60  	if rl.balance >= itemCost {
    61  		rl.balance -= itemCost
    62  		return true
    63  	}
    64  	return false
    65  }
    66  
    67  // updateBalance recalculates current balance based on time elapsed. Must be called while holding a lock.
    68  func (rl *RateLimiter) updateBalance() {
    69  	// calculate how much time passed since the last tick, and update current tick
    70  	currentTime := rl.timeNow()
    71  	elapsedTime := currentTime.Sub(rl.lastTick)
    72  	rl.lastTick = currentTime
    73  	// calculate how much credit have we accumulated since the last tick
    74  	rl.balance += elapsedTime.Seconds() * rl.creditsPerSecond
    75  	if rl.balance > rl.maxBalance {
    76  		rl.balance = rl.maxBalance
    77  	}
    78  }
    79  
    80  // Update changes the main parameters of the rate limiter in-place, while retaining
    81  // the current accumulated balance (pro-rated to the new maxBalance value). Using this method
    82  // instead of creating a new rate limiter helps to avoid thundering herd when sampling
    83  // strategies are updated.
    84  func (rl *RateLimiter) Update(creditsPerSecond, maxBalance float64) {
    85  	rl.lock.Lock()
    86  	defer rl.lock.Unlock()
    87  
    88  	rl.updateBalance() // get up to date balance
    89  	rl.balance = rl.balance * maxBalance / rl.maxBalance
    90  	rl.creditsPerSecond = creditsPerSecond
    91  	rl.maxBalance = maxBalance
    92  }
    93  
    94  func (r *rateLimitingSampler) Description() string {
    95  	return "rateLimitingSampler{}"
    96  }
    97  
    98  func (r *rateLimitingSampler) ShouldSample(p tracesdk.SamplingParameters) tracesdk.SamplingResult {
    99  	psc := oteltrace.SpanContextFromContext(p.ParentContext)
   100  	if r.rateLimiter.CheckCredit(1.0) {
   101  		return tracesdk.SamplingResult{
   102  			Decision:   tracesdk.RecordAndSample,
   103  			Tracestate: psc.TraceState(),
   104  		}
   105  	}
   106  	return tracesdk.SamplingResult{
   107  		Decision:   tracesdk.Drop,
   108  		Tracestate: psc.TraceState(),
   109  	}
   110  }
   111  
   112  func (r *rateLimitingSampler) init(rateLimit float64) {
   113  	if r.rateLimiter == nil {
   114  		r.rateLimiter = NewRateLimiter(rateLimit, math.Max(rateLimit, 1.0))
   115  	} else {
   116  		r.rateLimiter.Update(rateLimit, math.Max(rateLimit, 1.0))
   117  	}
   118  }
   119  
   120  func (r *rateLimitingSampler) Update(maxTracesPerSecond float64) {
   121  	if r.maxTracesPerSecond != maxTracesPerSecond {
   122  		r.init(maxTracesPerSecond)
   123  	}
   124  }