feat: fix key in safe way (#48)

key.go

@@ -2,6 +2,7 @@ package statter
 
 import (
 	"sync"
+	"unsafe"
 )
 
 const (
@@ -19,16 +20,15 @@ type key struct {
 	b []byte
 }
 
-func newKey(name string, tags []Tag) string {
+func newKey(name string, tags []Tag) *key {
 	k := keyPool.Get().(*key)
-	defer keyPool.Put(k)
-
 	k.b = k.b[:0]
+
 	k.b = append(k.b, name...)
 
 	// Short path for no tags.
 	if len(tags) == 0 {
-		return string(k.b)
+		return k
 	}
 
 	// The tags must be sorted to create a consistent key.
@@ -41,9 +41,28 @@ func newKey(name string, tags []Tag) string {
 		k.b = append(k.b, tag[1]...)
 	}
 
+	return k
+}
+
+// String returns the key as a string.
+//
+// The returned string should only be used while
+// holding a reference to the key, nor should it be
+// stored.
+func (k *key) String() string {
+	return *(*string)(unsafe.Pointer(&k.b))
+}
+
+// SafeString returns the key as a string that
+// is safe to use after releasing the key.
+func (k *key) SafeString() string {
 	return string(k.b)
 }
 
+func putKey(k *key) {
+	keyPool.Put(k)
+}
+
 func sortTags(tags []Tag) {
 	var sorted bool
 	for !sorted {

key_internal_test.go

@@ -41,8 +41,9 @@ func TestKey(t *testing.T) {
 			t.Parallel()
 
 			k := newKey(test.keyName, test.keyTags)
+			defer putKey(k)
 
-			assert.Equal(t, test.want, k)
+			assert.Equal(t, test.want, k.String())
 		})
 	}
 }
@@ -55,7 +56,8 @@ func BenchmarkKey(b *testing.B) {
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
 			k := newKey("test", tags)
-			_ = k
+
+			putKey(k)
 		}
 	})
 }

registry.go

@@ -61,11 +61,12 @@ func (r *registry) SubStatter(parent *Statter, prefix string, tags []Tag) *Statt
 	name, tags := mergeDescriptors(parent.prefix, parent.cfg.separator, prefix, parent.tags, tags)
 
 	k := newKey(name, tags)
+	defer putKey(k)
 
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
-	if s, ok := r.statters[k]; ok {
+	if s, ok := r.statters[k.String()]; ok {
 		return s
 	}
 
@@ -79,7 +80,7 @@ func (r *registry) SubStatter(parent *Statter, prefix string, tags []Tag) *Statt
 		prefix: name,
 		tags:   tags,
 	}
-	r.statters[k] = s
+	r.statters[k.SafeString()] = s
 
 	return s
 }

statter.go

@@ -126,61 +126,69 @@ func (s *Statter) With(prefix string, tags ...Tag) *Statter {
 func (s *Statter) Counter(name string, tags ...Tag) *Counter {
 	k := newKey(name, tags)
 
-	c, ok := s.counters.Load(k)
+	c, ok := s.counters.Load(k.String())
 	if !ok {
 		n, t := s.mergeDescriptors(name, tags)
 		counter := &Counter{
 			name: n,
 			tags: t,
 		}
-		c, _ = s.counters.LoadOrStore(k, counter)
+		c, _ = s.counters.LoadOrStore(k.SafeString(), counter)
 	}
 
+	putKey(k)
+
 	return c
 }
 
 // Gauge returns a gauge for the given name and tags.
 func (s *Statter) Gauge(name string, tags ...Tag) *Gauge {
 	k := newKey(name, tags)
 
-	g, ok := s.gauges.Load(k)
+	g, ok := s.gauges.Load(k.String())
 	if !ok {
 		n, t := s.mergeDescriptors(name, tags)
 		gauge := &Gauge{
 			name: n,
 			tags: t,
 		}
-		g, _ = s.gauges.LoadOrStore(k, gauge)
+		g, _ = s.gauges.LoadOrStore(k.SafeString(), gauge)
 	}
 
+	putKey(k)
+
 	return g
 }
 
 // Histogram returns a histogram for the given name and tags.
 func (s *Statter) Histogram(name string, tags ...Tag) *Histogram {
 	k := newKey(name, tags)
 
-	h, ok := s.histograms.Load(k)
+	h, ok := s.histograms.Load(k.String())
 	if !ok {
 		n, t := s.mergeDescriptors(name, tags)
 		histogram := newHistogram(s.hr, n, t, s.pool)
-		h, _ = s.histograms.LoadOrStore(k, histogram)
+		h, _ = s.histograms.LoadOrStore(k.SafeString(), histogram)
 	}
 
+	putKey(k)
+
 	return h
 }
 
 // Timing returns a timing for the given name and tags.
 func (s *Statter) Timing(name string, tags ...Tag) *Timing {
 	k := newKey(name, tags)
 
-	t, ok := s.timings.Load(k)
+	t, ok := s.timings.Load(k.String())
 	if !ok {
 		n, newTags := s.mergeDescriptors(name, tags)
 		timing := newTiming(s.tr, n, newTags, s.pool)
-		t, _ = s.timings.LoadOrStore(k, timing)
+		t, _ = s.timings.LoadOrStore(k.SafeString(), timing)
 	}
 
+	putKey(k)
+
 	return t
 }