core/sampler: Support decision hook

benchmarks/zap_test.go

@@ -116,7 +116,7 @@ func newZapLogger(lvl zapcore.Level) *zap.Logger {
 }
 
 func newSampledLogger(lvl zapcore.Level) *zap.Logger {
-	return zap.New(zapcore.NewSampler(
+	return zap.New(zapcore.NewSamplerWithOptions(
 		newZapLogger(zap.DebugLevel).Core(),
 		100*time.Millisecond,
 		10, // first

config.go

@@ -32,10 +32,15 @@ import (
 // global CPU and I/O load that logging puts on your process while attempting
 // to preserve a representative subset of your logs.
 //
-// Values configured here are per-second. See zapcore.NewSampler for details.
+// Hook is called whenever a Sampler makes a decision. Currently, whenever a
+// log is dropped.
+//
+// Values configured here are per-second. See zapcore.NewSamplerWithOptions for
+// details.
 type SamplingConfig struct {
 	Initial    int `json:"initial" yaml:"initial"`
 	Thereafter int `json:"thereafter" yaml:"thereafter"`
+	Hook       func(zapcore.Entry, zapcore.SamplingDecision) error
 }
 
 // Config offers a declarative way to construct a logger. It doesn't do
@@ -207,10 +212,20 @@ func (cfg Config) buildOptions(errSink zapcore.WriteSyncer) []Option {
 	if !cfg.DisableStacktrace {
 		opts = append(opts, AddStacktrace(stackLevel))
 	}
+	if cfg.Sampling != nil && cfg.Sampling.Hook == nil {
+		// Assign a default nop sampling hook.
+		cfg.Sampling.Hook = zapcore.NopSamplingHook
+	}
 
 	if cfg.Sampling != nil {
 		opts = append(opts, WrapCore(func(core zapcore.Core) zapcore.Core {
-			return zapcore.NewSampler(core, time.Second, int(cfg.Sampling.Initial), int(cfg.Sampling.Thereafter))
+			return zapcore.NewSamplerWithOptions(
+				core,
+				time.Second,
+				cfg.Sampling.Initial,
+				cfg.Sampling.Thereafter,
+				zapcore.SamplerHook(cfg.Sampling.Hook),
+			)
 		}))
 	}
 

config_test.go

@@ -27,6 +27,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 	"go.uber.org/zap/zapcore"
 )
 
@@ -144,3 +145,65 @@ func TestConfigWithMissingAttributes(t *testing.T) {
 		})
 	}
 }
+
+func makeSamplerCountingHook() (func(zapcore.Entry, zapcore.SamplingDecision) error, *atomic.Int64) {
+	count := &atomic.Int64{}
+	h := func(zapcore.Entry, zapcore.SamplingDecision) error {
+		count.Inc()
+		return nil
+	}
+	return h, count
+}
+
+func TestConfigWithSamplingHook(t *testing.T) {
+	shook, scount := makeSamplerCountingHook()
+	cfg := Config{
+		Level:       NewAtomicLevelAt(InfoLevel),
+		Development: false,
+		Sampling: &SamplingConfig{
+			Initial:    100,
+			Thereafter: 100,
+			Hook:       shook,
+		},
+		Encoding:         "json",
+		EncoderConfig:    NewProductionEncoderConfig(),
+		OutputPaths:      []string{"stderr"},
+		ErrorOutputPaths: []string{"stderr"},
+	}
+	expectN := 2 + 100 + 1 // 2 from initial logs, 100 initial sampled logs, 1 from off-by-one in sampler
+	expectRe := `{"level":"info","caller":"zap/config_test.go:\d+","msg":"info","k":"v","z":"zz"}` + "\n" +
+		`{"level":"warn","caller":"zap/config_test.go:\d+","msg":"warn","k":"v","z":"zz"}` + "\n"
+	expectDropped := 99 // 200 - 100 initial - 1 thereafter
+
+	temp, err := ioutil.TempFile("", "zap-prod-config-test")
+	require.NoError(t, err, "Failed to create temp file.")
+	defer func() {
+		err := os.Remove(temp.Name())
+		if err != nil {
+			return
+		}
+	}()
+
+	cfg.OutputPaths = []string{temp.Name()}
+	cfg.EncoderConfig.TimeKey = "" // no timestamps in tests
+	cfg.InitialFields = map[string]interface{}{"z": "zz", "k": "v"}
+
+	hook, count := makeCountingHook()
+	logger, err := cfg.Build(Hooks(hook))
+	require.NoError(t, err, "Unexpected error constructing logger.")
+
+	logger.Debug("debug")
+	logger.Info("info")
+	logger.Warn("warn")
+
+	byteContents, err := ioutil.ReadAll(temp)
+	require.NoError(t, err, "Couldn't read log contents from temp file.")
+	logs := string(byteContents)
+	assert.Regexp(t, expectRe, logs, "Unexpected log output.")
+
+	for i := 0; i < 200; i++ {
+		logger.Info("sampling")
+	}
+	assert.Equal(t, int64(expectN), count.Load(), "Hook called an unexpected number of times.")
+	assert.Equal(t, int64(expectDropped), scount.Load())
+}

zapcore/sampler.go

@@ -81,17 +81,87 @@ func (c *counter) IncCheckReset(t time.Time, tick time.Duration) uint64 {
 	return 1
 }
 
+// SamplingDecision represents a decision made by sampler.
+type SamplingDecision uint8
+
+const (
+	// Dropped means that a log was dropped.
+	Dropped SamplingDecision = iota
+)
+
+// optionFunc wraps a func so it satisfies the SamplerOption interface.
+type optionFunc func(*sampler)
+
+func (f optionFunc) apply(s *sampler) {
+	f(s)
+}
+
+// SamplerOption configures a Sampler option.
+type SamplerOption interface {
+	apply(*sampler)
+}
+
+// NopSamplingHook is the default hook used by sampler.
+func NopSamplingHook(_ Entry, _ SamplingDecision) error {
+	return nil
+}
+
+// SamplerHook registers a which will be called when Sampler makes a decision.
+// Currently a hook is called when a log is dropped and zapcore.Dropped decision
+// is emitted.
+//
+// This hook is useful for side effects, for example emitting number of dropped
+// logs. Note, there is no access to Fields in this hook. In the future, this
+// hook can be expanded to emit whether this is first entry that was dropped,
+// first after a period, etc.
+func SamplerHook(hook func(entry Entry, dec SamplingDecision) error) SamplerOption {
+	return optionFunc(func(s *sampler) {
+		s.hook = hook
+	})
+}
+
+// NewSamplerWithOptions creates a Core that samples incoming entries, which
+// caps the CPU and I/O load of logging while attempting to preserve a
+// representative subset of your logs.
+//
+// Zap samples by logging the first N entries with a given level and message
+// 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.
+//
+// Sampler also accepts an optional hook that can be used to count number of
+// dropped logs.
+//
+// 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 {
+	s := &sampler{
+		Core:       core,
+		tick:       tick,
+		counts:     newCounters(),
+		first:      uint64(first),
+		thereafter: uint64(thereafter),
+		hook:       NopSamplingHook,
+	}
+	for _, opt := range opts {
+		opt.apply(s)
+	}
+
+	return s
+}
+
 type sampler struct {
 	Core
 
 	counts            *counters
 	tick              time.Duration
 	first, thereafter uint64
+	hook              func(Entry, SamplingDecision) error
 }
 
-// NewSampler creates a Core that samples incoming entries, which caps the CPU
-// and I/O load of logging while attempting to preserve a representative subset
-// of your logs.
+// NewSampler creates a Core that samples incoming entries, which
+// caps the CPU and I/O load of logging while attempting to preserve a
+// representative subset of your logs.
 //
 // Zap samples by logging the first N entries with a given level and message
 // each tick. If more Entries with the same level and message are seen during
@@ -100,13 +170,15 @@ type sampler struct {
 // 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.
+// Deprecated: use NewSamplerWithOptions.
 func NewSampler(core Core, tick time.Duration, first, thereafter int) Core {
 	return &sampler{
 		Core:       core,
 		tick:       tick,
 		counts:     newCounters(),
 		first:      uint64(first),
 		thereafter: uint64(thereafter),
+		hook:       NopSamplingHook,
 	}
 }
 
@@ -117,6 +189,7 @@ func (s *sampler) With(fields []Field) Core {
 		counts:     s.counts,
 		first:      s.first,
 		thereafter: s.thereafter,
+		hook:       s.hook,
 	}
 }
 
@@ -128,7 +201,9 @@ func (s *sampler) Check(ent Entry, ce *CheckedEntry) *CheckedEntry {
 	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 {
+		_ = s.hook(ent, Dropped)
 		return ce
 	}
+
 	return s.Core.Check(ent, ce)
 }

zapcore/sampler_bench_test.go

@@ -203,7 +203,7 @@ var counterTestCases = [][]string{
 func BenchmarkSampler_Check(b *testing.B) {
 	for _, keys := range counterTestCases {
 		b.Run(fmt.Sprintf("%v keys", len(keys)), func(b *testing.B) {
-			fac := NewSampler(
+			fac := NewSamplerWithOptions(
 				NewCore(
 					NewJSONEncoder(testEncoderConfig()),
 					&ztest.Discarder{},

zapcore/sampler_test.go

@@ -37,6 +37,7 @@ import (
 
 func fakeSampler(lvl LevelEnabler, tick time.Duration, first, thereafter int) (Core, *observer.ObservedLogs) {
 	core, logs := observer.New(lvl)
+	// Keep using deprecated constructor for cc.
 	core = NewSampler(core, tick, first, thereafter)
 	return core, logs
 }
@@ -162,7 +163,7 @@ func TestSamplerConcurrent(t *testing.T) {
 
 	tick := ztest.Timeout(10 * time.Millisecond)
 	cc := &countingCore{}
-	sampler := NewSampler(cc, tick, logsPerTick, 100000)
+	sampler := NewSamplerWithOptions(cc, tick, logsPerTick, 100000)
 
 	var (
 		done atomic.Bool