github.com/ManabuSeki/goa-v1@v1.4.3/middleware/sampler.go (about)

     1  package middleware
     2  
     3  import (
     4  	"math/rand"
     5  	"sync"
     6  	"sync/atomic"
     7  	"time"
     8  )
     9  
    10  type (
    11  	// Sampler is an interface for computing when a sample falls within a range.
    12  	Sampler interface {
    13  		// Sample returns true if the caller should sample now.
    14  		Sample() bool
    15  	}
    16  
    17  	adaptiveSampler struct {
    18  		sync.Mutex
    19  		lastRate        int64
    20  		maxSamplingRate int
    21  		sampleSize      uint32
    22  		start           time.Time
    23  		counter         uint32
    24  	}
    25  
    26  	fixedSampler int
    27  )
    28  
    29  const (
    30  	// adaptive upper bound has granularity in case caller becomes extremely busy.
    31  	adaptiveUpperBoundInt   = 10000
    32  	adaptiveUpperBoundFloat = float64(adaptiveUpperBoundInt)
    33  )
    34  
    35  // NewAdaptiveSampler computes the interval for sampling for tracing middleware.
    36  // it can also be used by non-web go routines to trace internal API calls.
    37  //
    38  // maxSamplingRate is the desired maximum sampling rate in requests per second.
    39  //
    40  // sampleSize sets the number of requests between two adjustments of the
    41  // sampling rate when MaxSamplingRate is set. the sample rate cannot be adjusted
    42  // until the sample size is reached at least once.
    43  func NewAdaptiveSampler(maxSamplingRate, sampleSize int) Sampler {
    44  	if maxSamplingRate <= 0 {
    45  		panic("maxSamplingRate must be greater than 0")
    46  	}
    47  	if sampleSize <= 0 {
    48  		panic("sample size must be greater than 0")
    49  	}
    50  	return &adaptiveSampler{
    51  		lastRate:        adaptiveUpperBoundInt, // samples all until initial count reaches sample size
    52  		maxSamplingRate: maxSamplingRate,
    53  		sampleSize:      uint32(sampleSize),
    54  		start:           time.Now(),
    55  	}
    56  }
    57  
    58  // NewFixedSampler sets the tracing sampling rate as a percentage value.
    59  func NewFixedSampler(samplingPercent int) Sampler {
    60  	if samplingPercent < 0 || samplingPercent > 100 {
    61  		panic("samplingPercent must be between 0 and 100")
    62  	}
    63  	return fixedSampler(samplingPercent)
    64  }
    65  
    66  // Sample implementation for adaptive rate
    67  func (s *adaptiveSampler) Sample() bool {
    68  	// adjust sampling rate whenever sample size is reached.
    69  	var currentRate int
    70  	if atomic.AddUint32(&s.counter, 1) == s.sampleSize { // exact match prevents
    71  		atomic.StoreUint32(&s.counter, 0) // race is ok
    72  		s.Lock()
    73  		{
    74  			d := time.Since(s.start).Seconds()
    75  			r := float64(s.sampleSize) / d
    76  			currentRate = int((float64(s.maxSamplingRate) * adaptiveUpperBoundFloat) / r)
    77  			if currentRate > adaptiveUpperBoundInt {
    78  				currentRate = adaptiveUpperBoundInt
    79  			} else if currentRate < 1 {
    80  				currentRate = 1
    81  			}
    82  			s.start = time.Now()
    83  		}
    84  		s.Unlock()
    85  		atomic.StoreInt64(&s.lastRate, int64(currentRate))
    86  	} else {
    87  		currentRate = int(atomic.LoadInt64(&s.lastRate))
    88  	}
    89  
    90  	// currentRate is never zero.
    91  	return currentRate == adaptiveUpperBoundInt || rand.Intn(adaptiveUpperBoundInt) < currentRate
    92  }
    93  
    94  // Sample implementation for fixed percentage
    95  func (s fixedSampler) Sample() bool {
    96  	samplingPercent := int(s)
    97  	return samplingPercent > 0 && (samplingPercent == 100 || rand.Intn(100) < samplingPercent)
    98  }