sampler.go

package rz

import (
	"math/rand"
	"sync/atomic"
	"time"
)

const (
	// SampleOften samples log every ~ 10 events.
	SampleOften = SamplerRandom(10)
	// SampleSometimes samples log every ~ 100 events.
	SampleSometimes = SamplerRandom(100)
	// SampleRarely samples log every ~ 1000 events.
	SampleRarely = SamplerRandom(1000)
)

// LogSampler defines an interface to a log sampler.
type LogSampler interface {
	// Sample returns true if the event should be part of the sample, false if
	// the event should be dropped.
	Sample(lvl LogLevel) bool
}

// SamplerRandom use a PRNG to randomly sample an event out of N events,
// regardless of their level.
type SamplerRandom uint32

// Sample implements the Sampler interface.
func (s SamplerRandom) Sample(lvl LogLevel) bool {
	if s <= 0 {
		return false
	}
	if rand.Intn(int(s)) != 0 {
		return false
	}
	return true
}

// SamplerBasic is a sampler that will send every Nth events, regardless of
// there level.
type SamplerBasic struct {
	N       uint32
	counter uint32
}

// Sample implements the Sampler interface.
func (s *SamplerBasic) Sample(lvl LogLevel) bool {
	c := atomic.AddUint32(&s.counter, 1)
	return c%s.N == s.N-1
}

// SamplerBurst lets Burst events pass per Period then pass the decision to
// NextSampler. If Sampler is not set, all subsequent events are rejected.
type SamplerBurst struct {
	// Burst is the maximum number of event per period allowed before calling
	// NextSampler.
	Burst uint32
	// Period defines the burst period. If 0, NextSampler is always called.
	Period time.Duration
	// NextSampler is the sampler used after the burst is reached. If nil,
	// events are always rejected after the burst.
	NextSampler LogSampler

	counter uint32
	resetAt int64
}

// Sample implements the Sampler interface.
func (s *SamplerBurst) Sample(lvl LogLevel) bool {
	if s.Burst > 0 && s.Period > 0 {
		if s.inc() <= s.Burst {
			return true
		}
	}
	if s.NextSampler == nil {
		return false
	}
	return s.NextSampler.Sample(lvl)
}

func (s *SamplerBurst) inc() uint32 {
	now := time.Now().UnixNano()
	resetAt := atomic.LoadInt64(&s.resetAt)
	var c uint32
	if now > resetAt {
		c = 1
		atomic.StoreUint32(&s.counter, c)
		newResetAt := now + s.Period.Nanoseconds()
		reset := atomic.CompareAndSwapInt64(&s.resetAt, resetAt, newResetAt)
		if !reset {
			// Lost the race with another goroutine trying to reset.
			c = atomic.AddUint32(&s.counter, 1)
		}
	} else {
		c = atomic.AddUint32(&s.counter, 1)
	}
	return c
}

// SamplerLevel applies a different sampler for each level.
type SamplerLevel struct {
	DebugSampler LogSampler
	InfoSampler  LogSampler
	WarnSampler  LogSampler
	ErrorSampler LogSampler
}

// Sample implements the Sampler interface.
func (s SamplerLevel) Sample(lvl LogLevel) bool {
	switch lvl {
	case DebugLevel:
		if s.DebugSampler != nil {
			return s.DebugSampler.Sample(lvl)
		}
	case InfoLevel:
		if s.InfoSampler != nil {
			return s.InfoSampler.Sample(lvl)
		}
	case WarnLevel:
		if s.WarnSampler != nil {
			return s.WarnSampler.Sample(lvl)
		}
	case ErrorLevel:
		if s.ErrorSampler != nil {
			return s.ErrorSampler.Sample(lvl)
		}
	}
	return true
}