chore: Improve performance of TenantID

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

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)
+		}
+	})
+}