-
Notifications
You must be signed in to change notification settings - Fork 1
/
sampler.go
90 lines (77 loc) · 2.61 KB
/
sampler.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
package tracing
import (
"context"
"sync/atomic"
"go.opentelemetry.io/otel/sdk/trace"
// "go.opentelemetry.io/otel/trace"
)
func WithSkipSpan(ctx context.Context) context.Context {
return context.WithValue(ctx, TracePaused, "true")
}
// CustomSampler struct holds different samplers.
type CustomSampler struct {
defaultSampler trace.Sampler
samplers map[string]trace.Sampler
}
// NewCustomSampler creates a new instance of CustomSampler.
// `defaultSampler` is used when no specific sampler is found for a span name.
// `samplers` is a map where keys are span names and values are the specific samplers for those span names.
func NewCustomSampler(defaultSampler trace.Sampler, samplers map[string]trace.Sampler) *CustomSampler {
return &CustomSampler{
defaultSampler: defaultSampler,
samplers: samplers,
}
}
// ShouldSample implements the Sampler interface.
// It delegates the decision to a sampler based on the span name.
func (cs *CustomSampler) ShouldSample(params trace.SamplingParameters) trace.SamplingResult {
if skip := params.ParentContext.Value(TracePaused); skip != nil {
return trace.SamplingResult{
Decision: trace.Drop,
Attributes: nil,
}
}
if sampler, ok := cs.samplers[params.Name]; ok {
return sampler.ShouldSample(params)
}
return cs.defaultSampler.ShouldSample(params)
}
// Description returns a description of the sampler.
func (cs *CustomSampler) Description() string {
return "CustomSampler"
}
// CounterSampler samples spans based on a counter, ensuring a percentage of spans are sampled.
type CounterSampler struct {
percentage float64
counter atomic.Int64
rate int64
}
// NewCounterSampler creates a new CounterSampler.
// `percentage` should be a value between 0 and 100, representing the percentage of spans to sample.
func NewCounterSampler(percentage float64) *CounterSampler {
return &CounterSampler{
percentage: percentage,
counter: atomic.Int64{},
rate: int64(100.0 / percentage),
}
}
// ShouldSample decides whether a span should be sampled based on a counter.
func (cs *CounterSampler) ShouldSample(params trace.SamplingParameters) trace.SamplingResult {
// Atomically increment the counter
count := cs.counter.Add(1)
// Sample the span if the count modulo the rate equals 1 (to sample the first span in each cycle)
if count%cs.rate == 1 {
return trace.SamplingResult{
Decision: trace.RecordAndSample,
Attributes: nil,
}
}
return trace.SamplingResult{
Decision: trace.Drop,
Attributes: nil,
}
}
// Description returns the description of the sampler.
func (cs *CounterSampler) Description() string {
return "CounterSampler"
}