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 1/3] 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)
+		}
+	})
+}

From 5714840c0dbc29b0b0cceb0ec72586db057d1c18 Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Wed, 3 Jul 2024 23:56:43 +0200
Subject: [PATCH 2/3] lint files

---
 tenant/tenant_test.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tenant/tenant_test.go b/tenant/tenant_test.go
index f531358ff..f0b6b9e15 100644
--- a/tenant/tenant_test.go
+++ b/tenant/tenant_test.go
@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/grafana/dskit/user"
+
 	"github.com/stretchr/testify/assert"
 )
 

From c8018282abc87c0faec5f5630d1ea352909a7fec Mon Sep 17 00:00:00 2001
From: Cyril Tovena <cyril.tovena@gmail.com>
Date: Thu, 4 Jul 2024 08:38:09 +0200
Subject: [PATCH 3/3] review feedback

---
 tenant/resolver.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tenant/resolver.go b/tenant/resolver.go
index 82f7e2947..9a01d6322 100644
--- a/tenant/resolver.go
+++ b/tenant/resolver.go
@@ -27,7 +27,7 @@ func TenantID(ctx context.Context) (string, error) {
 		}
 		return orgID, nil
 	}
-	orgIDs, err := TenantIDs(ctx)
+	orgIDs, err := tenantIDsFromString(orgID)
 	if err != nil {
 		return "", err
 	}
@@ -53,6 +53,10 @@ func TenantIDs(ctx context.Context) ([]string, error) {
 		return nil, err
 	}
 
+	return tenantIDsFromString(orgID)
+}
+
+func tenantIDsFromString(orgID string) ([]string, error) {
 	orgIDs := strings.Split(orgID, tenantIDsSeparator)
 	for _, id := range orgIDs {
 		if err := ValidTenantID(id); err != nil {