Add EMAThroughput sampler

emathroughput.go

@@ -0,0 +1,359 @@
+package dynsampler
+
+import (
+ "encoding/json"
+ "errors"
+ "fmt"
+ "math"
+ "sync"
+ "time"
+)
+
+// EMAThroughput implements Sampler and attempts to achieve a given throughput
+// rate, weighting rare traffic and frequent traffic differently so as to end up
+// with the correct value.
+//
+// Based on the EMASampleRate implementation, EMAThroughput differs in that
+// instead of trying to achieve a given sample rate, it tries to reach a given
+// throughput of events. During bursts of traffic, it will reduce sample
+// rates so as to keep the number of events per second roughly constant.
+//
+// Like the EMA sampler, it maintains an Exponential Moving Average of counts
+// seen per key, and adjusts this average at regular intervals. The weight
+// applied to more recent intervals is defined by `weight`, a number between (0,
+// 1) - larger values weight the average more toward recent observations. In
+// other words, a larger weight will cause sample rates more quickly adapt to
+// traffic patterns, while a smaller weight will result in sample rates that are
+// less sensitive to bursts or drops in traffic and thus more consistent over
+// time.
+//
+// New keys that are not found in the EMA will always have a sample
+// rate of 1. Keys that occur more frequently will be sampled on a logarithmic
+// curve. In other words, every key will be represented at least once in any
+// given window and more frequent keys will have their sample rate
+// increased proportionally to wind up with the goal throughput.
+type EMAThroughput struct {
+ // AdjustmentInterval defines how often we adjust the moving average from
+ // recent observations. Default 15s.
+ AdjustmentInterval time.Duration
+
+ // Weight is a value between (0, 1) indicating the weighting factor used to adjust
+ // the EMA. With larger values, newer data will influence the average more, and older
+ // values will be factored out more quickly. In mathematical literature concerning EMA,
+ // this is referred to as the `alpha` constant.
+ // Default is 0.5
+ Weight float64
+
+ // InitialSampleRate is the sample rate to use during startup, before we
+ // have accumulated enough data to calculate a reasonable desired sample
+ // rate. This is mainly useful in situations where unsampled throughput is
+ // high enough to cause problems.
+ // Default 10.
+ InitialSampleRate int
+
+ // GoalThroughputPerSec is the target number of events to send per second.
+ // Sample rates are generated to squash the total throughput down to match the
+ // goal throughput. Actual throughput may exceed goal throughput. default 100
+ GoalThroughputPerSec int
+
+ // MaxKeys, if greater than 0, limits the number of distinct keys tracked in EMA.
+ // Once MaxKeys is reached, new keys will not be included in the sample rate map, but
+ // existing keys will continue to be be counted.
+ MaxKeys int
+
+ // AgeOutValue indicates the threshold for removing keys from the EMA. The EMA of any key will approach 0
+ // if it is not repeatedly observed, but will never truly reach it, so we have to decide what constitutes "zero".
+ // Keys with averages below this threshold will be removed from the EMA. Default is the same as Weight, as this prevents
+ // a key with the smallest integer value (1) from being aged out immediately. This value should generally be <= Weight,
+ // unless you have very specific reasons to set it higher.
+ AgeOutValue float64
+
+ // BurstMultiple, if set, is multiplied by the sum of the running average of counts to define
+ // the burst detection threshold. If total counts observed for a given interval exceed the threshold
+ // EMA is updated immediately, rather than waiting on the AdjustmentInterval.
+ // Defaults to 2; negative value disables. With a default of 2, if your traffic suddenly doubles,
+ // burst detection will kick in.
+ BurstMultiple float64
+
+ // BurstDetectionDelay indicates the number of intervals to run after Start is called before burst detection kicks in.
+ // Defaults to 3
+ BurstDetectionDelay uint
+
+ savedSampleRates map[string]int
+ currentCounts map[string]float64
+ movingAverage map[string]float64
+ burstThreshold float64
+ currentBurstSum float64
+ intervalCount uint
+ burstSignal chan struct{}
+
+ // haveData indicates that we have gotten a sample of traffic. Before we've
+ // gotten any samples of traffic, we should use the default goal sample rate
+ // for all events instead of sampling everything at 1
+ haveData bool
+ updating bool
+ done chan struct{}
+
+ lock sync.Mutex
+
+ // used only in tests
+ testSignalMapsDone chan struct{}
+}
+
+// Ensure we implement the sampler interface
+var _ Sampler = (*EMAThroughput)(nil)
+
+func (e *EMAThroughput) Start() error {
+ // apply defaults
+ if e.AdjustmentInterval == 0 {
+ e.AdjustmentInterval = 15 * time.Second
+ }
+ if e.AdjustmentInterval < 1*time.Millisecond {
+ return fmt.Errorf("the AdjustmentInterval %v is unreasonably short for a throughput sampler", e.AdjustmentInterval)
+ }
+ if e.InitialSampleRate == 0 {
+ e.InitialSampleRate = 10
+ }
+ if e.GoalThroughputPerSec == 0 {
+ e.GoalThroughputPerSec = 100
+ }
+ if e.Weight == 0 {
+ e.Weight = 0.5
+ }
+ if e.AgeOutValue == 0 {
+ e.AgeOutValue = e.Weight
+ }
+ if e.BurstMultiple == 0 {
+ e.BurstMultiple = 2
+ }
+ if e.BurstDetectionDelay == 0 {
+ e.BurstDetectionDelay = 3
+ }
+
+ // Don't override these maps at startup in case they were loaded from a previous state
+ e.currentCounts = make(map[string]float64)
+ if e.savedSampleRates == nil {
+ e.savedSampleRates = make(map[string]int)
+ }
+ if e.movingAverage == nil {
+ e.movingAverage = make(map[string]float64)
+ }
+ e.burstSignal = make(chan struct{})
+ e.done = make(chan struct{})
+
+ go func() {
+ ticker := time.NewTicker(e.AdjustmentInterval)
+ defer ticker.Stop()
+ for {
+ select {
+ case <-e.burstSignal:
+ // reset ticker when we get a burst
+ ticker.Stop()
+ ticker = time.NewTicker(e.AdjustmentInterval)
+ e.updateMaps()
+ case <-ticker.C:
+ e.updateMaps()
+ e.intervalCount++
+ case <-e.done:
+ return
+ }
+ }
+ }()
+ return nil
+}
+
+func (e *EMAThroughput) Stop() error {
+ close(e.done)
+ return nil
+}
+
+// updateMaps calculates a new saved rate map based on the contents of the
+// counter map
+func (e *EMAThroughput) updateMaps() {
+ e.lock.Lock()
+ if e.testSignalMapsDone != nil {
+ defer func() {
+ e.testSignalMapsDone <- struct{}{}
+ }()
+ }
+ // short circuit if no traffic
+ if len(e.currentCounts) == 0 {
+ // No traffic the last interval, don't update anything. This is deliberate to avoid
+ // the average decaying when there's no traffic (comes in bursts, or there's some kind of outage).
+ e.lock.Unlock()
+ return
+ }
+ // If there is another updateMaps going, bail
+ if e.updating {
+ e.lock.Unlock()
+ return
+ }
+ e.updating = true
+ // make a local copy of the sample counters for calculation
+ tmpCounts := e.currentCounts
+ e.currentCounts = make(map[string]float64)
+ e.currentBurstSum = 0
+ e.lock.Unlock()
+
+ e.updateEMA(tmpCounts)
+
+ // Goal events to send this interval is the total count of events in the EMA
+ // divided by the desired average sample rate
+ var sumEvents float64
+ for _, count := range e.movingAverage {
+ sumEvents += math.Max(1, count)
+ }
+
+ // Store this for burst detection. This is checked in GetSampleRate
+ // so we need to grab the lock when we update it.
+ e.lock.Lock()
+ e.burstThreshold = sumEvents * e.BurstMultiple
+ e.lock.Unlock()
+
+ // Calculate the desired average sample rate per second based on the volume we've received.
+ // InitialSampleRate := float64(sumEvents) / e.AdjustmentInterval.Seconds() / float64(e.GoalThroughputPerSec)
+ // goalCount := float64(sumEvents) / InitialSampleRate
+
+ // Calculate the number of events we'd like to let through per adjustment interval
+ goalCount := float64(e.GoalThroughputPerSec) / e.AdjustmentInterval.Seconds()
+
+ // goalRatio is the goalCount divided by the sum of all the log values - it
+ // determines what percentage of the total event space belongs to each key
+ var logSum float64
+ for _, count := range e.movingAverage {
+ // We take the max of (1, count) because count * weight is < 1 for
+ // very small counts, which throws off the logSum and can cause
+ // incorrect samples rates to be computed when throughput is low
+ logSum += math.Log10(math.Max(1, count))
+ }
+ goalRatio := goalCount / logSum
+
+ newSavedSampleRates := calculateSampleRates(goalRatio, e.movingAverage)
+ e.lock.Lock()
+ defer e.lock.Unlock()
+ e.savedSampleRates = newSavedSampleRates
+ e.haveData = true
+ e.updating = false
+}
+
+// GetSampleRate takes a key and returns the appropriate sample rate for that
+// key.
+func (e *EMAThroughput) GetSampleRate(key string) int {
+ return e.GetSampleRateMulti(key, 1)
+}
+
+// GetSampleRateMulti takes a key representing count spans and returns the
+// appropriate sample rate for that key.
+func (e *EMAThroughput) GetSampleRateMulti(key string, count int) int {
+ e.lock.Lock()
+ defer e.lock.Unlock()
+
+ // Enforce MaxKeys limit on the size of the map
+ if e.MaxKeys > 0 {
+ // If a key already exists, increment it. If not, but we're under the limit, store a new key
+ if _, found := e.currentCounts[key]; found || len(e.currentCounts) < e.MaxKeys {
+ e.currentCounts[key] += float64(count)
+ e.currentBurstSum += float64(count)
+ }
+ } else {
+ e.currentCounts[key] += float64(count)
+ e.currentBurstSum += float64(count)
+ }
+
+ // Enforce the burst threshold
+ if e.burstThreshold > 0 && e.currentBurstSum >= e.burstThreshold && e.intervalCount >= e.BurstDetectionDelay {
+ // reset the burst sum to prevent additional burst updates from occurring while updateMaps is running
+ e.currentBurstSum = 0
+ // send but don't block - consuming is blocked on updateMaps, which takes the same lock we're holding
+ select {
+ case e.burstSignal <- struct{}{}:
+ default:
+ }
+ }
+
+ if !e.haveData {
+ return e.InitialSampleRate
+ }
+ if rate, found := e.savedSampleRates[key]; found {
+ return rate
+ }
+ return 1
+}
+
+func (e *EMAThroughput) updateEMA(newCounts map[string]float64) {
+ keysToUpdate := make([]string, 0, len(e.movingAverage))
+ for key := range e.movingAverage {
+ keysToUpdate = append(keysToUpdate, key)
+ }
+
+ // Update any existing keys with new values
+ for _, key := range keysToUpdate {
+ var newAvg float64
+ // Was this key seen in the last interval? Adjust by that amount
+ if val, found := newCounts[key]; found {
+ newAvg = adjustAverage(e.movingAverage[key], val, e.Weight)
+ } else {
+ // Otherwise adjust by zero
+ newAvg = adjustAverage(e.movingAverage[key], 0, e.Weight)
+ }
+
+ // Age out this value if it's too small to care about for calculating sample rates
+ // This is also necessary to keep our map from going forever.
+ if newAvg < e.AgeOutValue {
+ delete(e.movingAverage, key)
+ } else {
+ e.movingAverage[key] = newAvg
+ }
+ // We've processed this key - don't process it again when we look at new counts
+ delete(newCounts, key)
+ }
+
+ for key := range newCounts {
+ newAvg := adjustAverage(0, newCounts[key], e.Weight)
+ if newAvg >= e.AgeOutValue {
+ e.movingAverage[key] = newAvg
+ }
+ }
+}
+
+type emaThroughputState struct {
+ // These fields are exported for use by `JSON.Marshal` and `JSON.Unmarshal`
+ SavedSampleRates map[string]int `json:"saved_sample_rates"`
+ MovingAverage map[string]float64 `json:"moving_average"`
+}
+
+// SaveState returns a byte array with a JSON representation of the sampler state
+func (e *EMAThroughput) SaveState() ([]byte, error) {
+ e.lock.Lock()
+ defer e.lock.Unlock()
+
+ if e.savedSampleRates == nil {
+ return nil, errors.New("saved sample rate map is nil")
+ }
+ if e.movingAverage == nil {
+ return nil, errors.New("moving average map is nil")
+ }
+ s := &emaThroughputState{SavedSampleRates: e.savedSampleRates, MovingAverage: e.movingAverage}
+ return json.Marshal(s)
+}
+
+// LoadState accepts a byte array with a JSON representation of a previous instance's
+// state
+func (e *EMAThroughput) LoadState(state []byte) error {
+ e.lock.Lock()
+ defer e.lock.Unlock()
+
+ s := emaThroughputState{}
+ err := json.Unmarshal(state, &s)
+ if err != nil {
+ return err
+ }
+
+ // Load the previously calculated sample rates
+ e.savedSampleRates = s.SavedSampleRates
+ e.movingAverage = s.MovingAverage
+ // Allow GetSampleRate to return calculated sample rates from the loaded map
+ e.haveData = true
+
+ return nil
+}