From 0fdc92c5552ee578c7b6d27505cd8b00c5d2468a Mon Sep 17 00:00:00 2001
From: Cerek Hillen <cerekh@gmail.com>
Date: Wed, 21 Jun 2023 10:13:51 -0700
Subject: [PATCH] let people configure cacheSize if they want

---
 stats.go | 30 +++++++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/stats.go b/stats.go
index 0128072..333b393 100644
--- a/stats.go
+++ b/stats.go
@@ -59,6 +59,9 @@ type Store interface {
 	Scope
 }
 
+// StoreOption configures a Store when constructed with NewStoreWithOptions.
+type StoreOption func(*statStore)
+
 // A Scope namespaces Statistics.
 //
 //	store := stats.NewDefaultStore()
@@ -216,16 +219,32 @@ type StatGenerator interface {
 // NewStore returns an Empty store that flushes to Sink passed as an argument.
 // Note: the export argument is unused.
 func NewStore(sink Sink, _ bool) Store {
+	return NewStoreWithOptions(sink)
+}
+
+// NewStoreWithOptions returns an Empty store that flushes to Sink passed as an argument.
+func NewStoreWithOptions(sink Sink, options ...StoreOption) Store {
 	s := &statStore{sink: sink}
+	for _, option := range options {
+		option(s)
+	}
 
 	// lru.NewWithEvict can only return a non-nil error when cacheSize < 0.
-	cacheSize := 65536
-	s.counters, _ = lru.NewWithEvict(cacheSize, s.flushCounter)
-	s.timers, _ = lru.New[string, *timer](cacheSize)
+	s.counters, _ = lru.NewWithEvict(s.cacheSize, s.flushCounter)
+	s.timers, _ = lru.New[string, *timer](s.cacheSize)
 
 	return s
 }
 
+// WithCacheSize sets the maximum number of counters and timers gostats will cache.
+// By default this value is 65,536 and it should be kept higher
+// than the expected maximum number of counters or timers.
+func WithCacheSize(cacheSize int) StoreOption {
+	return func(store *statStore) {
+		store.cacheSize = cacheSize
+	}
+}
+
 // NewDefaultStore returns a Store with a TCP statsd sink, and a running flush timer.
 func NewDefaultStore() Store {
 	var newStore Store
@@ -345,8 +364,9 @@ func (ts *timespan) CompleteWithDuration(value time.Duration) {
 }
 
 type statStore struct {
-	counters *lru.Cache[string, *counter]
-	timers   *lru.Cache[string, *timer]
+	cacheSize int
+	counters  *lru.Cache[string, *counter]
+	timers    *lru.Cache[string, *timer]
 	// Gauges must not be expunged because they are client-side stateful.
 	// We use a sync.Map instead of a cache to ensure they are kept indefinitely.
 	gauges sync.Map