From 65e213ff21f85c908775c7de9a82e448861ba9d0 Mon Sep 17 00:00:00 2001
From: Abhinav Gupta <abg@uber.com>
Date: Wed, 8 Dec 2021 12:21:00 -0800
Subject: [PATCH] sampler: Support thereafter of zero

The Zap sampling logger accepts three configuration parameters:
interval, first, and thereafter.

After we see `first` log with the same message in `interval`,
the sampler kicks in.
Following that, we let through every `thereafter`-th log.

So for example,

    NewSamplerWithOptions(core, time.Second, 100, 50)

This will allow 100 logs with the same message in a second,
and following that, every 50th message with that message.

In #1032, the user wanted the sampler to reject *all* duplicate messages
once the limit was reached, but our sampler panics
if `thereafter = 0`.

This change adds support for setting `thereafter` to 0,
dropping all messages in that interval once the limit is reached.
It also adds further explanation to the documentation to address the
misuse in #1032.

Resolves #1032
---
 zapcore/sampler.go      | 15 +++++++++++++--
 zapcore/sampler_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/zapcore/sampler.go b/zapcore/sampler.go
index 31ed96e12..8c116049d 100644
--- a/zapcore/sampler.go
+++ b/zapcore/sampler.go
@@ -133,10 +133,21 @@ func SamplerHook(hook func(entry Entry, dec SamplingDecision)) SamplerOption {
 // each tick. If more Entries with the same level and message are seen during
 // the same interval, every Mth message is logged and the rest are dropped.
 //
+// For example,
+//
+//   core = NewSamplerWithOptions(core, time.Second, 10, 5)
+//
+// This will log the first 10 log entries with the same level and message
+// in a one second interval as-is. Following that, it will allow through
+// every 5th log entry with the same level and message in that interval.
+//
+// If thereafter is zero, the Core will drop all log entries after the first N
+// in that interval.
+//
 // Sampler can be configured to report sampling decisions with the SamplerHook
 // option.
 //
-// Keep in mind that zap's sampling implementation is optimized for speed over
+// Keep in mind that Zap's sampling implementation is optimized for speed over
 // absolute precision; under load, each tick may be slightly over- or
 // under-sampled.
 func NewSamplerWithOptions(core Core, tick time.Duration, first, thereafter int, opts ...SamplerOption) Core {
@@ -200,7 +211,7 @@ func (s *sampler) Check(ent Entry, ce *CheckedEntry) *CheckedEntry {
 	if ent.Level >= _minLevel && ent.Level <= _maxLevel {
 		counter := s.counts.get(ent.Level, ent.Message)
 		n := counter.IncCheckReset(ent.Time, s.tick)
-		if n > s.first && (n-s.first)%s.thereafter != 0 {
+		if n > s.first && (s.thereafter == 0 || (n-s.first)%s.thereafter != 0) {
 			s.hook(ent, LogDropped)
 			return ce
 		}
diff --git a/zapcore/sampler_test.go b/zapcore/sampler_test.go
index fbcdd46cc..f19a6fc9e 100644
--- a/zapcore/sampler_test.go
+++ b/zapcore/sampler_test.go
@@ -275,3 +275,42 @@ func TestSamplerUnknownLevels(t *testing.T) {
 		})
 	}
 }
+
+func TestSamplerWithZeroThereafter(t *testing.T) {
+	var counter countingCore
+
+	// Logs two messages per second.
+	sampler := NewSamplerWithOptions(&counter, time.Second, 2, 0)
+
+	now := time.Now()
+
+	for i := 0; i < 1000; i++ {
+		ent := Entry{
+			Level:   InfoLevel,
+			Message: "msg",
+			Time:    now,
+		}
+		if ce := sampler.Check(ent, nil); ce != nil {
+			ce.Write()
+		}
+	}
+
+	assert.Equal(t, 2, int(counter.logs.Load()),
+		"Unexpected number of logs")
+
+	now = now.Add(time.Second)
+
+	for i := 0; i < 1000; i++ {
+		ent := Entry{
+			Level:   InfoLevel,
+			Message: "msg",
+			Time:    now,
+		}
+		if ce := sampler.Check(ent, nil); ce != nil {
+			ce.Write()
+		}
+	}
+
+	assert.Equal(t, 4, int(counter.logs.Load()),
+		"Unexpected number of logs")
+}