PR comments

internal/tenant/context.go

@@ -3,7 +3,11 @@ package tenant
 import (
 	"context"
 	"fmt"
+	"runtime/pprof"
 	"strconv"
+	"sync"
+
+	"go.uber.org/atomic"
 
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype"
@@ -24,14 +28,42 @@ func FromContext(ctx context.Context) (*tenanttype.Tenant, error) {
 func IDToString(ctx context.Context) string {
 	tnt, ok := tenanttype.GetTenant(ctx)
 	if !ok {
-		if enforcement.PPROFMissingTenant != nil {
+		if profile := pprofMissingTenant(); profile != nil {
 			// We want to track every stack trace, so need a unique value for the event
-			eventValue := enforcement.PPROFUniqID.Add(1)
+			eventValue := pprofUniqID.Add(1)
 
 			// skip stack for Add and this function (2).
-			enforcement.PPROFMissingTenant.Add(eventValue, 2)
+			profile.Add(eventValue, 2)
 		}
 		return "missing"
 	}
 	return strconv.Itoa(tnt.ID())
 }
+
+var pprofUniqID atomic.Int64
+var pprofOnce sync.Once
+var pprofProfile *pprof.Profile
+
+// pprofMissingTenant returns the pprof profile for missing tenants,
+// initializing it only once.
+func pprofMissingTenant() *pprof.Profile {
+	pprofOnce.Do(func() {
+		if shouldLogNoTenant() {
+			pprofProfile = pprof.NewProfile("missing_tenant")
+		}
+	})
+	return pprofProfile
+}
+
+// shouldLogNoTenant returns true if the tenant enforcement mode is logging or strict.
+// It is used to log a warning if a request to a low-level store is made without a tenant
+// so we can identify missing tenants. This will go away and only strict will be allowed
+// once we are confident that all contexts carry tenants.
+func shouldLogNoTenant() bool {
+	switch enforcement.EnforcementMode.Load() {
+	case "logging", "strict":
+		return true
+	default:
+		return false
+	}
+}

internal/tenant/internal/enforcement/enforcement.go

@@ -2,7 +2,6 @@ package enforcement
 
 import (
 	"os"
-	"runtime/pprof"
 
 	"go.uber.org/atomic"
 )
@@ -11,24 +10,3 @@ import (
 // instead of in the tenant package to avoid a circular dependency. See
 // tenanttest.MockEnforce.
 var EnforcementMode = atomic.NewString(os.Getenv("SRC_TENANT_ENFORCEMENT_MODE"))
-
-var PPROFUniqID atomic.Int64
-var PPROFMissingTenant = func() *pprof.Profile {
-	if !ShouldLogNoTenant() {
-		return nil
-	}
-	return pprof.NewProfile("zoekt_missing_tenant")
-}()
-
-// ShouldLogNoTenant returns true if the tenant enforcement mode is logging or strict.
-// It is used to log a warning if a request to a low-level store is made without a tenant
-// so we can identify missing tenants. This will go away and only strict will be allowed
-// once we are confident that all contexts carry tenants.
-func ShouldLogNoTenant() bool {
-	switch EnforcementMode.Load() {
-	case "logging", "strict":
-		return true
-	default:
-		return false
-	}
-}