diff --git a/stats.go b/stats.go index 9f167bd8..209dd52d 100644 --- a/stats.go +++ b/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 diff --git a/stats_test.go b/stats_test.go index 85482f78..7ccf0157 100644 --- a/stats_test.go +++ b/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 +}