github.com/pinpoint-apm/pinpoint-go-agent@v1.4.1-0.20240110120318-a50c2eb18c8c/sampler.go (about) 1 package pinpoint 2 3 import ( 4 "golang.org/x/time/rate" 5 "sync/atomic" 6 "time" 7 ) 8 9 const ( 10 samplingMaxPercentRate = 100 * 100 11 ) 12 13 type sampler interface { 14 isSampled() bool 15 } 16 17 type rateSampler struct { 18 rate uint64 19 counter uint64 20 } 21 22 func newRateSampler(rate int) *rateSampler { 23 if rate < 0 { 24 rate = 0 25 } 26 return &rateSampler{ 27 rate: uint64(rate), 28 counter: 0, 29 } 30 } 31 32 func (s *rateSampler) isSampled() bool { 33 if s.rate == 0 { 34 return false 35 } 36 samplingCount := atomic.AddUint64(&s.counter, 1) 37 isSampled := samplingCount % s.rate 38 return isSampled == 0 39 } 40 41 type percentSampler struct { 42 rate uint64 43 counter uint64 44 } 45 46 func newPercentSampler(percent float64) *percentSampler { 47 if percent < 0 { 48 percent = 0 49 } else if percent < 0.01 { 50 percent = 0.01 51 } else if percent > 100 { 52 percent = 100 53 } 54 55 return &percentSampler{ 56 rate: uint64(percent * 100), 57 counter: 0, 58 } 59 } 60 61 func (s *percentSampler) isSampled() bool { 62 if s.rate == 0 { 63 return false 64 } 65 samplingCount := atomic.AddUint64(&s.counter, s.rate) 66 r := samplingCount % samplingMaxPercentRate 67 return r < s.rate 68 } 69 70 type traceSampler interface { 71 isNewSampled() bool 72 isContinueSampled() bool 73 } 74 75 type basicTraceSampler struct { 76 baseSampler sampler 77 } 78 79 func newBasicTraceSampler(base sampler) *basicTraceSampler { 80 return &basicTraceSampler{ 81 baseSampler: base, 82 } 83 } 84 85 func (s *basicTraceSampler) isNewSampled() bool { 86 sampled := s.baseSampler.isSampled() 87 if sampled { 88 incrSampleNew() 89 } else { 90 incrUnSampleNew() 91 } 92 return sampled 93 } 94 95 func (s *basicTraceSampler) isContinueSampled() bool { 96 incrSampleCont() 97 return true 98 } 99 100 type throughputLimitTraceSampler struct { 101 baseSampler sampler 102 newSampleLimiter *rate.Limiter 103 continueSampleLimiter *rate.Limiter 104 } 105 106 func newThroughputLimitTraceSampler(base sampler, newTps int, continueTps int) *throughputLimitTraceSampler { 107 var ( 108 newLimiter *rate.Limiter 109 contLimiter *rate.Limiter 110 ) 111 112 if newTps > 0 { 113 newLimiter = rate.NewLimiter(per(newTps, time.Second), 1) 114 } 115 if continueTps > 0 { 116 contLimiter = rate.NewLimiter(per(continueTps, time.Second), 1) 117 } 118 return &throughputLimitTraceSampler{ 119 baseSampler: base, 120 newSampleLimiter: newLimiter, 121 continueSampleLimiter: contLimiter, 122 } 123 } 124 125 func per(throughput int, d time.Duration) rate.Limit { 126 return rate.Every(d / time.Duration(throughput)) 127 } 128 129 func (s *throughputLimitTraceSampler) isNewSampled() bool { 130 sampled := s.baseSampler.isSampled() 131 if sampled { 132 if s.newSampleLimiter != nil { 133 sampled = s.newSampleLimiter.Allow() 134 if sampled { 135 incrSampleNew() 136 } else { 137 incrSkipNew() 138 } 139 } else { 140 incrSampleNew() 141 } 142 } else { 143 incrUnSampleNew() 144 } 145 146 return sampled 147 } 148 149 func (s *throughputLimitTraceSampler) isContinueSampled() bool { 150 sampled := true 151 if s.continueSampleLimiter != nil { 152 sampled = s.continueSampleLimiter.Allow() 153 if sampled { 154 incrSampleCont() 155 } else { 156 incrSkipCont() 157 } 158 } else { 159 incrSampleCont() 160 } 161 162 return sampled 163 }