sampler.go

Documentation: github.com/rs/zerolog

     1  package zerolog
     2  
     3  import (
     4  	"math/rand"
     5  	"sync/atomic"
     6  	"time"
     7  )
     8  
     9  var (
    10  	// Often samples log every ~ 10 events.
    11  	Often = RandomSampler(10)
    12  	// Sometimes samples log every ~ 100 events.
    13  	Sometimes = RandomSampler(100)
    14  	// Rarely samples log every ~ 1000 events.
    15  	Rarely = RandomSampler(1000)
    16  )
    17  
    18  // Sampler defines an interface to a log sampler.
    19  type Sampler interface {
    20  	// Sample returns true if the event should be part of the sample, false if
    21  	// the event should be dropped.
    22  	Sample(lvl Level) bool
    23  }
    24  
    25  // RandomSampler use a PRNG to randomly sample an event out of N events,
    26  // regardless of their level.
    27  type RandomSampler uint32
    28  
    29  // Sample implements the Sampler interface.
    30  func (s RandomSampler) Sample(lvl Level) bool {
    31  	if s <= 0 {
    32  		return false
    33  	}
    34  	if rand.Intn(int(s)) != 0 {
    35  		return false
    36  	}
    37  	return true
    38  }
    39  
    40  // BasicSampler is a sampler that will send every Nth events, regardless of
    41  // their level.
    42  type BasicSampler struct {
    43  	N       uint32
    44  	counter uint32
    45  }
    46  
    47  // Sample implements the Sampler interface.
    48  func (s *BasicSampler) Sample(lvl Level) bool {
    49  	n := s.N
    50  	if n == 1 {
    51  		return true
    52  	}
    53  	c := atomic.AddUint32(&s.counter, 1)
    54  	return c%n == 1
    55  }
    56  
    57  // BurstSampler lets Burst events pass per Period then pass the decision to
    58  // NextSampler. If Sampler is not set, all subsequent events are rejected.
    59  type BurstSampler struct {
    60  	// Burst is the maximum number of event per period allowed before calling
    61  	// NextSampler.
    62  	Burst uint32
    63  	// Period defines the burst period. If 0, NextSampler is always called.
    64  	Period time.Duration
    65  	// NextSampler is the sampler used after the burst is reached. If nil,
    66  	// events are always rejected after the burst.
    67  	NextSampler Sampler
    68  
    69  	counter uint32
    70  	resetAt int64
    71  }
    72  
    73  // Sample implements the Sampler interface.
    74  func (s *BurstSampler) Sample(lvl Level) bool {
    75  	if s.Burst > 0 && s.Period > 0 {
    76  		if s.inc() <= s.Burst {
    77  			return true
    78  		}
    79  	}
    80  	if s.NextSampler == nil {
    81  		return false
    82  	}
    83  	return s.NextSampler.Sample(lvl)
    84  }
    85  
    86  func (s *BurstSampler) inc() uint32 {
    87  	now := time.Now().UnixNano()
    88  	resetAt := atomic.LoadInt64(&s.resetAt)
    89  	var c uint32
    90  	if now > resetAt {
    91  		c = 1
    92  		atomic.StoreUint32(&s.counter, c)
    93  		newResetAt := now + s.Period.Nanoseconds()
    94  		reset := atomic.CompareAndSwapInt64(&s.resetAt, resetAt, newResetAt)
    95  		if !reset {
    96  			// Lost the race with another goroutine trying to reset.
    97  			c = atomic.AddUint32(&s.counter, 1)
    98  		}
    99  	} else {
   100  		c = atomic.AddUint32(&s.counter, 1)
   101  	}
   102  	return c
   103  }
   104  
   105  // LevelSampler applies a different sampler for each level.
   106  type LevelSampler struct {
   107  	TraceSampler, DebugSampler, InfoSampler, WarnSampler, ErrorSampler Sampler
   108  }
   109  
   110  func (s LevelSampler) Sample(lvl Level) bool {
   111  	switch lvl {
   112  	case TraceLevel:
   113  		if s.TraceSampler != nil {
   114  			return s.TraceSampler.Sample(lvl)
   115  		}
   116  	case DebugLevel:
   117  		if s.DebugSampler != nil {
   118  			return s.DebugSampler.Sample(lvl)
   119  		}
   120  	case InfoLevel:
   121  		if s.InfoSampler != nil {
   122  			return s.InfoSampler.Sample(lvl)
   123  		}
   124  	case WarnLevel:
   125  		if s.WarnSampler != nil {
   126  			return s.WarnSampler.Sample(lvl)
   127  		}
   128  	case ErrorLevel:
   129  		if s.ErrorSampler != nil {
   130  			return s.ErrorSampler.Sample(lvl)
   131  		}
   132  	}
   133  	return true
   134  }
   135
View as plain text