From c570fd3e969b18f6f82030f6f844629aa624abb0 Mon Sep 17 00:00:00 2001
From: alanprot <alanprot@gmail.com>
Date: Fri, 26 Apr 2024 11:36:23 -0700
Subject: [PATCH] Optimize Write Signature calculation

---
 pkg/cortexpb/extensions.go      | 17 +++++------------
 pkg/cortexpb/extensions_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/pkg/cortexpb/extensions.go b/pkg/cortexpb/extensions.go
index f900b775918..716fafcc79d 100644
--- a/pkg/cortexpb/extensions.go
+++ b/pkg/cortexpb/extensions.go
@@ -21,9 +21,8 @@ var signerPool = sync.Pool{
 }
 
 type signer struct {
-	h         *xxhash.Digest
-	b         []byte
-	optimized bool
+	h *xxhash.Digest
+	b []byte
 }
 
 func newSigner() *signer {
@@ -38,18 +37,15 @@ func newSigner() *signer {
 func (s *signer) Reset() {
 	s.h.Reset()
 	s.b = s.b[:0]
-	s.optimized = true
 }
 
 func (s *signer) WriteString(val string) {
 	switch {
-	case !s.optimized:
-		_, _ = s.h.WriteString(val)
 	case len(s.b)+len(val) > cap(s.b):
 		// If labels val does not fit in the []byte we fall back to not allocate the whole entry.
 		_, _ = s.h.Write(s.b)
-		_, _ = s.h.WriteString(val)
-		s.optimized = false
+		s.b = s.b[:0]
+		s.b = append(s.b, val...)
 	default:
 		// Use xxhash.Sum64(b) for fast path as it's faster.
 		s.b = append(s.b, val...)
@@ -57,10 +53,7 @@ func (s *signer) WriteString(val string) {
 }
 
 func (s *signer) Sum64() uint64 {
-	if s.optimized {
-		return xxhash.Sum64(s.b)
-	}
-
+	_, _ = s.h.Write(s.b)
 	return s.h.Sum64()
 }
 
diff --git a/pkg/cortexpb/extensions_test.go b/pkg/cortexpb/extensions_test.go
index 0f8b6385264..94a5f76d48d 100644
--- a/pkg/cortexpb/extensions_test.go
+++ b/pkg/cortexpb/extensions_test.go
@@ -10,6 +10,32 @@ import (
 	"github.com/weaveworks/common/user"
 )
 
+func BenchmarkSignRequest(b *testing.B) {
+	ctx := context.Background()
+	ctx = user.InjectOrgID(ctx, "user-1")
+
+	tests := []struct {
+		size int
+	}{
+		{size: 10},
+		{size: 100},
+		{size: 1000},
+		{size: 10000},
+	}
+
+	for _, tc := range tests {
+		b.Run(fmt.Sprintf("WriteRequestSize: %v", tc.size), func(b *testing.B) {
+			wr := createWriteRequest(tc.size, true, "family1", "help1", "unit")
+			b.ResetTimer()
+			b.ReportAllocs()
+			for i := 0; i < b.N; i++ {
+				_, err := wr.Sign(ctx)
+				require.NoError(b, err)
+			}
+		})
+	}
+}
+
 func TestWriteRequest_Sign(t *testing.T) {
 	ctx := context.Background()
 	ctx = user.InjectOrgID(ctx, "user-1")