From bf0f87c8e0fe43ebe0980c8d5cc8aa45c5a9a549 Mon Sep 17 00:00:00 2001
From: Marco Pracucci <marco@pracucci.com>
Date: Mon, 4 Mar 2024 15:36:04 +0100
Subject: [PATCH] Reduced memory allocations by user.ExtractFromGRPCRequest()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

goos: darwin
goarch: amd64
pkg: github.com/grafana/dskit/user
cpu: Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
                          │ before.txt  │             after.txt              │
                          │   sec/op    │   sec/op     vs base               │
ExtractFromGRPCRequest-12   772.2n ± 2%   124.0n ± 7%  -83.94% (p=0.002 n=6)

                          │ before.txt  │             after.txt             │
                          │    B/op     │    B/op     vs base               │
ExtractFromGRPCRequest-12   560.00 ± 0%   80.00 ± 0%  -85.71% (p=0.002 n=6)

                          │ before.txt  │             after.txt             │
                          │  allocs/op  │ allocs/op   vs base               │
ExtractFromGRPCRequest-12   10.000 ± 0%   3.000 ± 0%  -70.00% (p=0.002 n=6)

Signed-off-by: Marco Pracucci <marco@pracucci.com>
---
 CHANGELOG.md      |  1 +
 user/grpc.go      |  9 ++-----
 user/grpc_test.go | 61 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 7 deletions(-)
 create mode 100644 user/grpc_test.go

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fe9d96b9b..9f5c07fce 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -194,6 +194,7 @@
 * [ENHANCEMENT] Add servicediscovery package. #469
 * [ENHANCEMENT] Expose `InstancesInZoneCount` and `ZonesCount` in `ring.ReadRing` interface. #494
 * [ENHANCEMENT] Add optimization to run `concurrency.ForEachJob()` with no parallelism when there's only 1 job. #486 #495
+* [ENHANCEMENT] Reduced memory allocations by `user.ExtractFromGRPCRequest()`. #502
 * [BUGFIX] spanlogger: Support multiple tenant IDs. #59
 * [BUGFIX] Memberlist: fixed corrupted packets when sending compound messages with more than 255 messages or messages bigger than 64KB. #85
 * [BUGFIX] Ring: `ring_member_ownership_percent` and `ring_tokens_owned` metrics are not updated on scale down. #109
diff --git a/user/grpc.go b/user/grpc.go
index 201b835ee..fcfd3d7a9 100644
--- a/user/grpc.go
+++ b/user/grpc.go
@@ -13,13 +13,8 @@ import (
 // ExtractFromGRPCRequest extracts the user ID from the request metadata and returns
 // the user ID and a context with the user ID injected.
 func ExtractFromGRPCRequest(ctx context.Context) (string, context.Context, error) {
-	md, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		return "", ctx, ErrNoOrgID
-	}
-
-	orgIDs, ok := md[lowerOrgIDHeaderName]
-	if !ok || len(orgIDs) != 1 {
+	orgIDs := metadata.ValueFromIncomingContext(ctx, lowerOrgIDHeaderName)
+	if len(orgIDs) != 1 {
 		return "", ctx, ErrNoOrgID
 	}
 
diff --git a/user/grpc_test.go b/user/grpc_test.go
new file mode 100644
index 000000000..89e685ab3
--- /dev/null
+++ b/user/grpc_test.go
@@ -0,0 +1,61 @@
+package user
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestExtractFromGRPCRequest(t *testing.T) {
+	t.Run("should return error if no org ID is set in the gRPC request context", func(t *testing.T) {
+		inputCtx := context.Background()
+
+		_, returnedCtx, err := ExtractFromGRPCRequest(inputCtx)
+		assert.Equal(t, inputCtx, returnedCtx)
+		assert.Equal(t, ErrNoOrgID, err)
+	})
+
+	t.Run("should return a context with org ID injected if org ID is set in the gRPC request context", func(t *testing.T) {
+		inputCtx := metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{
+			lowerOrgIDHeaderName: "user-1",
+		}))
+
+		// Pre-condition check: no org ID should be set in the input context.
+		_, err := ExtractOrgID(inputCtx)
+		assert.Equal(t, ErrNoOrgID, err)
+
+		actualOrgID, returnedCtx, err := ExtractFromGRPCRequest(inputCtx)
+		assert.NotEqual(t, inputCtx, returnedCtx)
+		assert.NoError(t, err)
+		assert.Equal(t, "user-1", actualOrgID)
+
+		// Org ID should be set in the returned context.
+		actualOrgID, err = ExtractOrgID(returnedCtx)
+		assert.NoError(t, err)
+		assert.Equal(t, "user-1", actualOrgID)
+	})
+}
+
+func BenchmarkExtractFromGRPCRequest(b *testing.B) {
+	// The following fixture has been generated looking at the actual metadata received by Grafana Mimir.
+	ctx := metadata.NewIncomingContext(context.Background(), metadata.New(map[string]string{
+		"authority":            "1.1.1.1",
+		"content-type":         "application/grpc",
+		"grpc-accept-encoding": "snappy,gzip",
+		"uber-trace-id":        "xxx",
+		"user-agent":           "grpc-go/1.61.1",
+		lowerOrgIDHeaderName:   "user-1",
+	}))
+
+	for n := 0; n < b.N; n++ {
+		orgID, _, err := ExtractFromGRPCRequest(ctx)
+		if orgID != "user-1" {
+			b.Fatalf("unexpected org ID: %s", orgID)
+		}
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}