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  }