Add StoreOption/WithSinkWrap to expose the sink NewDefaultStore builds

stats.go

@@ -217,19 +217,49 @@ func NewStore(sink Sink, _ bool) Store {
 	return &statStore{sink: sink}
 }
 
+// A StoreOption configures a Store returned by NewDefaultStore.
+type StoreOption func(*storeOptions)
+
+type storeOptions struct {
+	wrapSink func(FlushableSink) FlushableSink
+}
+
+// WithSinkWrap returns a StoreOption that passes the sink NewDefaultStore
+// would otherwise have used unwrapped (chosen from Settings exactly as
+// NewDefaultStore always does: a TCP statsd sink, or a logging/null sink
+// depending on UseStatsd and LoggingSinkDisabled) through wrap before
+// constructing the Store.
+//
+// This lets a caller decorate that sink -- for example to filter which
+// stats actually get flushed -- without having to reimplement
+// NewDefaultStore's sink-selection logic themselves just to get at the
+// underlying sink, which NewDefaultStore doesn't otherwise expose.
+func WithSinkWrap(wrap func(FlushableSink) FlushableSink) StoreOption {
+	return func(o *storeOptions) {
+		o.wrapSink = wrap
+	}
+}
+
 // NewDefaultStore returns a Store with a TCP statsd sink, and a running flush timer.
-func NewDefaultStore() Store {
+func NewDefaultStore(opts ...StoreOption) Store {
+	so := storeOptions{wrapSink: func(sink FlushableSink) FlushableSink { return sink }}
+	for _, opt := range opts {
+		opt(&so)
+	}
+
 	var newStore Store
 	settings := GetSettings()
 	if !settings.UseStatsd {
+		var inner FlushableSink
 		if settings.LoggingSinkDisabled {
-			newStore = NewStore(NewNullSink(), false)
+			inner = NewNullSink()
 		} else {
-			newStore = NewStore(NewLoggingSink(), false)
+			inner = NewLoggingSink()
 		}
+		newStore = NewStore(so.wrapSink(inner), false)
 		go newStore.Start(time.NewTicker(10 * time.Second))
 	} else {
-		newStore = NewStore(NewTCPStatsdSink(), false)
+		newStore = NewStore(so.wrapSink(NewTCPStatsdSink()), false)
 		go newStore.Start(time.NewTicker(time.Duration(settings.FlushIntervalS) * time.Second))
 	}
 	return newStore

stats_test.go

@@ -521,3 +521,76 @@ func BenchmarkStoreNewPerInstanceCounter(b *testing.B) {
 		}
 	})
 }
+
+func TestWithSinkWrap_CallsWrapWithTheSinkNewDefaultStoreWouldHaveUsed(t *testing.T) {
+	cases := []struct {
+		name                string
+		useStatsd           string
+		loggingSinkDisabled string
+		assertType          func(t *testing.T, got FlushableSink)
+	}{
+		{"statsd", "true", "false", func(t *testing.T, got FlushableSink) {
+			if _, ok := got.(*netSink); !ok {
+				t.Errorf("wrap called with %T, want *netSink", got)
+			}
+		}},
+		{"logging", "false", "false", func(t *testing.T, got FlushableSink) {
+			if _, ok := got.(*loggingSink); !ok {
+				t.Errorf("wrap called with %T, want *loggingSink", got)
+			}
+		}},
+		{"null", "false", "true", func(t *testing.T, got FlushableSink) {
+			if _, ok := got.(nullSink); !ok {
+				t.Errorf("wrap called with %T, want nullSink", got)
+			}
+		}},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Setenv("USE_STATSD", tc.useStatsd)
+			t.Setenv("GOSTATS_LOGGING_SINK_DISABLED", tc.loggingSinkDisabled)
+
+			var wrapped FlushableSink
+			NewDefaultStore(WithSinkWrap(func(sink FlushableSink) FlushableSink {
+				wrapped = sink
+				return sink
+			}))
+
+			if wrapped == nil {
+				t.Fatal("wrap was never called")
+			}
+			tc.assertType(t, wrapped)
+		})
+	}
+}
+
+// TestWithSinkWrap_InstallsWrappedSink proves the Store actually flushes to
+// whatever wrap returns, not to the sink NewDefaultStore would otherwise
+// have used unwrapped -- the whole point of the option.
+func TestWithSinkWrap_InstallsWrappedSink(t *testing.T) {
+	t.Setenv("USE_STATSD", "false")
+	t.Setenv("GOSTATS_LOGGING_SINK_DISABLED", "true")
+
+	replacement := mock.NewSink()
+	store := NewDefaultStore(WithSinkWrap(func(FlushableSink) FlushableSink { return replacement }))
+
+	store.NewCounter("test_counter").Inc()
+	store.Flush()
+
+	if v, ok := replacement.LoadCounter("test_counter"); !ok || v != 1 {
+		t.Errorf("test_counter: got %v, %v want 1, true", v, ok)
+	}
+}
+
+// TestNewDefaultStore_NoOptions proves the zero-args call form still works
+// after NewDefaultStore became variadic.
+func TestNewDefaultStore_NoOptions(t *testing.T) {
+	t.Setenv("USE_STATSD", "false")
+	t.Setenv("GOSTATS_LOGGING_SINK_DISABLED", "true")
+
+	store := NewDefaultStore()
+
+	store.NewCounter("test_counter").Inc()
+	store.Flush() // must not panic
+}