Add more advanced sampling modes

README.md

@@ -154,12 +154,31 @@ log.Logger = log.With().Str("foo", "bar").Logger()
 ### Log Sampling
 
 ```go
-sampled := log.Sample(10)
+sampled := log.Sample(&zerolog.BasicSampler{N: 10})
 sampled.Info().Msg("will be logged every 10 messages")
 
-// Output: {"time":1494567715,"sample":10,"message":"will be logged every 10 messages"}
+// Output: {"time":1494567715,"level":"info","message":"will be logged every 10 messages"}
 ```
 
+More advanced sampling:
+
+```go
+// Will let 5 debug messages per period of 1 second.
+// Over 5 debug message, 1 every 100 debug messages are logged.
+// Other levels are not sampled.
+sampled := log.Sample(zerolog.LevelSampler{
+    DebugSampler: &zerolog.BurstSampler{
+        Burst: 5,
+        Period: 1*time.Second,
+        NextSampler: &zerolog.BasicSampler{N: 100},
+    },
+})
+sampled.Debug().Msg("hello world")
+
+// Output: {"time":1494567715,"level":"debug","message":"hello world"}
+```
+
+
 ### Pass a sub-logger by context
 
 ```go

log.go

@@ -72,7 +72,6 @@ import (
 	"io/ioutil"
 	"os"
 	"strconv"
-	"sync/atomic"
 
 	"github.com/rs/zerolog/internal/json"
 )
@@ -115,15 +114,6 @@ func (l Level) String() string {
 	return ""
 }
 
-const (
-	// Often samples log every 10 events.
-	Often = 10
-	// Sometimes samples log every 100 events.
-	Sometimes = 100
-	// Rarely samples log every 1000 events.
-	Rarely = 1000
-)
-
 var disabledEvent = newEvent(levelWriterAdapter{ioutil.Discard}, 0, false)
 
 // A Logger represents an active logging object that generates lines
@@ -134,8 +124,7 @@ var disabledEvent = newEvent(levelWriterAdapter{ioutil.Discard}, 0, false)
 type Logger struct {
 	w       LevelWriter
 	level   Level
-	sample  uint32
-	counter *uint32
+	sampler Sampler
 	context []byte
 }
 
@@ -166,7 +155,7 @@ func Nop() Logger {
 func (l Logger) Output(w io.Writer) Logger {
 	l2 := New(w)
 	l2.level = l.level
-	l2.sample = l.sample
+	l2.sampler = l.sampler
 	l2.context = make([]byte, len(l.context))
 	copy(l2.context, l.context)
 	return l2
@@ -190,27 +179,20 @@ func (l Logger) Level(lvl Level) Logger {
 	return Logger{
 		w:       l.w,
 		level:   lvl,
-		sample:  l.sample,
-		counter: l.counter,
+		sampler: l.sampler,
 		context: l.context,
 	}
 }
 
-// Sample returns a logger that only let one message out of every to pass thru.
-func (l Logger) Sample(every int) Logger {
-	if every == 0 {
-		// Create a child with no sampling.
-		return Logger{
-			w:       l.w,
-			level:   l.level,
-			context: l.context,
-		}
+// Sample returns a logger with the s sampler.
+func (l Logger) Sample(s Sampler) Logger {
+	if l.sampler == s {
+		return l
 	}
 	return Logger{
 		w:       l.w,
 		level:   l.level,
-		sample:  uint32(every),
-		counter: new(uint32),
+		sampler: s,
 		context: l.context,
 	}
 }
@@ -314,7 +296,7 @@ func (l Logger) newEvent(level Level, addLevelField bool, done func(string)) *Ev
 	if addLevelField {
 		lvl = level
 	}
-	e := newEvent(l.w, lvl, enabled)
+	e := newEvent(l.w, lvl, true)
 	e.done = done
 	if l.context != nil && len(l.context) > 0 && l.context[0] > 0 {
 		// first byte of context is ts flag
@@ -323,9 +305,6 @@ func (l Logger) newEvent(level Level, addLevelField bool, done func(string)) *Ev
 	if addLevelField {
 		e.Str(LevelFieldName, level.String())
 	}
-	if l.sample > 0 && SampleFieldName != "" {
-		e.Uint32(SampleFieldName, l.sample)
-	}
 	if l.context != nil && len(l.context) > 1 {
 		if len(e.buf) > 1 {
 			e.buf = append(e.buf, ',')
@@ -340,9 +319,8 @@ func (l Logger) should(lvl Level) bool {
 	if lvl < l.level || lvl < globalLevel() {
 		return false
 	}
-	if l.sample > 0 && l.counter != nil && !samplingDisabled() {
-		c := atomic.AddUint32(l.counter, 1)
-		return c%l.sample == 0
+	if l.sampler != nil && !samplingDisabled() {
+		return l.sampler.Sample(lvl)
 	}
 	return true
 }

log/log.go

@@ -27,9 +27,9 @@ func Level(level zerolog.Level) zerolog.Logger {
 	return Logger.Level(level)
 }
 
-// Sample returns a logger that only let one message out of every to pass thru.
-func Sample(every int) zerolog.Logger {
-	return Logger.Sample(every)
+// Sample returns a logger with the s sampler.
+func Sample(s zerolog.Sampler) zerolog.Logger {
+	return Logger.Sample(s)
 }
 
 // Debug starts a new message with debug level.

log_example_test.go

@@ -38,15 +38,15 @@ func ExampleLogger_Level() {
 }
 
 func ExampleLogger_Sample() {
-	log := zerolog.New(os.Stdout).Sample(2)
+	log := zerolog.New(os.Stdout).Sample(&zerolog.BasicSampler{N: 2})
 
 	log.Info().Msg("message 1")
 	log.Info().Msg("message 2")
 	log.Info().Msg("message 3")
 	log.Info().Msg("message 4")
 
-	// Output: {"level":"info","sample":2,"message":"message 2"}
-	// {"level":"info","sample":2,"message":"message 4"}
+	// Output: {"level":"info","message":"message 2"}
+	// {"level":"info","message":"message 4"}
 }
 
 func ExampleLogger_Debug() {

log_test.go

@@ -310,12 +310,12 @@ func TestLevel(t *testing.T) {
 
 func TestSampling(t *testing.T) {
 	out := &bytes.Buffer{}
-	log := New(out).Sample(2)
+	log := New(out).Sample(&BasicSampler{N: 2})
 	log.Log().Int("i", 1).Msg("")
 	log.Log().Int("i", 2).Msg("")
 	log.Log().Int("i", 3).Msg("")
 	log.Log().Int("i", 4).Msg("")
-	if got, want := out.String(), "{\"sample\":2,\"i\":2}\n{\"sample\":2,\"i\":4}\n"; got != want {
+	if got, want := out.String(), "{\"i\":2}\n{\"i\":4}\n"; got != want {
 		t.Errorf("invalid log output:\ngot:  %v\nwant: %v", got, want)
 	}
 }

sampler.go

@@ -0,0 +1,128 @@
+package zerolog
+
+import (
+	"math/rand"
+	"sync/atomic"
+	"time"
+)
+
+var (
+	// Often samples log every ~ 10 events.
+	Often = RandomSampler(10)
+	// Sometimes samples log every ~ 100 events.
+	Sometimes = RandomSampler(100)
+	// Rarely samples log every ~ 1000 events.
+	Rarely = RandomSampler(1000)
+)
+
+// Sampler defines an interface to a log sampler.
+type Sampler interface {
+	// Sample returns true if the event should be part of the sample, false if
+	// the event should be dropped.
+	Sample(lvl Level) bool
+}
+
+// RandomSampler use a PRNG to randomly sample an event out of N events,
+// regardless of their level.
+type RandomSampler uint32
+
+// Sample implements the Sampler interface.
+func (s RandomSampler) Sample(lvl Level) bool {
+	if s <= 0 {
+		return false
+	}
+	if s > 0 {
+		if rand.Intn(int(s)) != 0 {
+			return false
+		}
+	}
+	return true
+}
+
+// BasicSampler is a sampler that will send every Nth events, regardless of
+// there level.
+type BasicSampler struct {
+	N       uint32
+	counter uint32
+}
+
+// Sample implements the Sampler interface.
+func (s *BasicSampler) Sample(lvl Level) bool {
+	c := atomic.AddUint32(&s.counter, 1)
+	return c%s.N == 0
+}
+
+// BurstSampler lets Burst events pass per Period then pass the decision to
+// NextSampler. If Sampler is not set, all subsequent events are rejected.
+type BurstSampler 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 Sampler
+
+	counter uint32
+	resetAt int64
+}
+
+// Sample implements the Sampler interface.
+func (s *BurstSampler) Sample(lvl Level) bool {
+	if s.Burst > 9 && s.Period > 0 {
+		if s.inc() <= s.Burst {
+			return true
+		}
+	}
+	if s.NextSampler == nil {
+		return false
+	}
+	return s.NextSampler.Sample(lvl)
+}
+
+func (s *BurstSampler) 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
+}
+
+// LevelSampler applies a different sampler for each level.
+type LevelSampler struct {
+	DebugSampler, InfoSampler, WarnSampler, ErrorSampler Sampler
+}
+
+func (s LevelSampler) Sample(lvl Level) 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
+}

sampler_test.go

@@ -0,0 +1,75 @@
+package zerolog
+
+import (
+	"testing"
+	"time"
+)
+
+var samplers = []struct {
+	name    string
+	sampler func() Sampler
+	total   int
+	wantMin int
+	wantMax int
+}{
+	{
+		"BasicSampler",
+		func() Sampler {
+			return &BasicSampler{N: 5}
+		},
+		100, 20, 20,
+	},
+	{
+		"RandomSampler",
+		func() Sampler {
+			return RandomSampler(5)
+		},
+		100, 10, 30,
+	},
+	{
+		"BurstSampler",
+		func() Sampler {
+			return &BurstSampler{Burst: 20, Period: time.Second}
+		},
+		100, 20, 20,
+	},
+	{
+		"BurstSamplerNext",
+		func() Sampler {
+			return &BurstSampler{Burst: 20, Period: time.Second, NextSampler: &BasicSampler{N: 5}}
+		},
+		120, 40, 40,
+	},
+}
+
+func TestSamplers(t *testing.T) {
+	for i := range samplers {
+		s := samplers[i]
+		t.Run(s.name, func(t *testing.T) {
+			sampler := s.sampler()
+			got := 0
+			for t := s.total; t > 0; t-- {
+				if sampler.Sample(0) {
+					got++
+				}
+			}
+			if got < s.wantMin || got > s.wantMax {
+				t.Errorf("%s.Sample(0) == true %d on %d, want [%d, %d]", s.name, got, s.total, s.wantMin, s.wantMax)
+			}
+		})
+	}
+}
+
+func BenchmarkSamplers(b *testing.B) {
+	for i := range samplers {
+		s := samplers[i]
+		b.Run(s.name, func(b *testing.B) {
+			sampler := s.sampler()
+			b.RunParallel(func(pb *testing.PB) {
+				for pb.Next() {
+					sampler.Sample(0)
+				}
+			})
+		})
+	}
+}