From 79a5c450ff3b8db8ecabf0ddae1370f638cd5826 Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Wed, 3 Jul 2024 23:54:47 +0200
Subject: [PATCH] chore: Improve performance of TenantID

---
 tenant/resolver.go    | 11 +++++++++++
 tenant/tenant_test.go | 22 ++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/tenant/resolver.go b/tenant/resolver.go
index 35e95b1c8..82f7e2947 100644
--- a/tenant/resolver.go
+++ b/tenant/resolver.go
@@ -16,6 +16,17 @@ import (
 //
 //nolint:revive
 func TenantID(ctx context.Context) (string, error) {
+	//lint:ignore faillint wrapper around upstream method
+	orgID, err := user.ExtractOrgID(ctx)
+	if err != nil {
+		return "", err
+	}
+	if !strings.Contains(orgID, tenantIDsSeparator) {
+		if err := ValidTenantID(orgID); err != nil {
+			return "", err
+		}
+		return orgID, nil
+	}
 	orgIDs, err := TenantIDs(ctx)
 	if err != nil {
 		return "", err
diff --git a/tenant/tenant_test.go b/tenant/tenant_test.go
index bc3a60b54..f531358ff 100644
--- a/tenant/tenant_test.go
+++ b/tenant/tenant_test.go
@@ -1,9 +1,11 @@
 package tenant
 
 import (
+	"context"
 	"strings"
 	"testing"
 
+	"github.com/grafana/dskit/user"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -48,3 +50,23 @@ func TestValidTenantIDs(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkTenantID(b *testing.B) {
+	singleCtx := context.Background()
+	singleCtx = user.InjectOrgID(singleCtx, "tenant-a")
+	multiCtx := context.Background()
+	multiCtx = user.InjectOrgID(multiCtx, "tenant-a|tenant-b|tenant-c")
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	b.Run("single", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			_, _ = TenantID(singleCtx)
+		}
+	})
+	b.Run("multi", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			_, _ = TenantID(multiCtx)
+		}
+	})
+}