tenant: add tenant to trace (#862)

This adds the tenant ID to the trace. I also move the pprof logging from `FromContext` to the higher level `typeRepoSearcher`. The number of events was too high, because we logged missing tenants per document.

I also fixed a bug where pprof logging didn't work at all, because we read the tenant enforcemnt ENV after we set the pprof profile, so the profile was always nil.

Test plan:
Checked locally that tenants show up in the traces and "missing_tenant" shows up as pprof profile.

internal/tenant/context.go

@@ -4,35 +4,67 @@ import (
 	"context"
 	"fmt"
 	"runtime/pprof"
+	"sync"
 
 	"go.uber.org/atomic"
 
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype"
+	"github.com/sourcegraph/zoekt/trace"
 )
 
 var ErrMissingTenant = fmt.Errorf("missing tenant")
 
 func FromContext(ctx context.Context) (*tenanttype.Tenant, error) {
 	tnt, ok := tenanttype.GetTenant(ctx)
 	if !ok {
-		if pprofMissingTenant != nil {
+		return nil, ErrMissingTenant
+	}
+	return tnt, nil
+}
+
+// Log logs the tenant ID to the trace. If tenant logging is enabled, it also
+// logs a stack trace to a pprof profile.
+func Log(ctx context.Context, tr *trace.Trace) {
+	tnt, ok := tenanttype.GetTenant(ctx)
+	if !ok {
+		if profile := pprofMissingTenant(); profile != nil {
 			// We want to track every stack trace, so need a unique value for the event
 			eventValue := pprofUniqID.Add(1)
 
 			// skip stack for Add and this function (2).
-			pprofMissingTenant.Add(eventValue, 2)
+			profile.Add(eventValue, 2)
 		}
-
-		return nil, ErrMissingTenant
+		tr.LazyPrintf("tenant: missing")
+		return
 	}
-	return tnt, nil
+	tr.LazyPrintf("tenant: %d", tnt.ID())
 }
 
 var pprofUniqID atomic.Int64
-var pprofMissingTenant = func() *pprof.Profile {
-	if !enforcement.ShouldLogNoTenant() {
-		return nil
+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
 	}
-	return pprof.NewProfile("missing_tenant")
-}()
+}

internal/tenant/enforcement.go

@@ -1,19 +1,9 @@
 package tenant
 
 import (
-	"os"
-
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
 )
 
-func init() {
-	v, ok := os.LookupEnv("SRC_TENANT_ENFORCEMENT_MODE")
-	if !ok {
-		v = "disabled"
-	}
-	enforcement.EnforcementMode.Store(v)
-}
-
 func EnforceTenant() bool {
 	switch enforcement.EnforcementMode.Load() {
 	case "strict":

internal/tenant/internal/enforcement/enforcement.go

@@ -1,21 +1,12 @@
 package enforcement
 
-import "go.uber.org/atomic"
+import (
+	"os"
+
+	"go.uber.org/atomic"
+)
 
 // EnforcementMode is the current tenant enforcement mode. It resides here
 // instead of in the tenant package to avoid a circular dependency. See
 // tenanttest.MockEnforce.
-var EnforcementMode atomic.String
-
-// 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
-	}
-}
+var EnforcementMode = atomic.NewString(os.Getenv("SRC_TENANT_ENFORCEMENT_MODE"))

shards/eval.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/sourcegraph/zoekt"
+	"github.com/sourcegraph/zoekt/internal/tenant"
 	"github.com/sourcegraph/zoekt/query"
 	"github.com/sourcegraph/zoekt/trace"
 )
@@ -19,6 +20,9 @@ func (s *typeRepoSearcher) Search(ctx context.Context, q query.Q, opts *zoekt.Se
 	tr, ctx := trace.New(ctx, "typeRepoSearcher.Search", "")
 	tr.LazyLog(q, true)
 	tr.LazyPrintf("opts: %+v", opts)
+	if tenant.EnforceTenant() {
+		tenant.Log(ctx, tr)
+	}
 	defer func() {
 		if sr != nil {
 			tr.LazyPrintf("num files: %d", len(sr.Files))
@@ -43,6 +47,9 @@ func (s *typeRepoSearcher) StreamSearch(ctx context.Context, q query.Q, opts *zo
 	tr, ctx := trace.New(ctx, "typeRepoSearcher.StreamSearch", "")
 	tr.LazyLog(q, true)
 	tr.LazyPrintf("opts: %+v", opts)
+	if tenant.EnforceTenant() {
+		tenant.Log(ctx, tr)
+	}
 	var stats zoekt.Stats
 	defer func() {
 		tr.LazyPrintf("stats: %+v", stats)
@@ -68,6 +75,9 @@ func (s *typeRepoSearcher) List(ctx context.Context, q query.Q, opts *zoekt.List
 	tr, ctx := trace.New(ctx, "typeRepoSearcher.List", "")
 	tr.LazyLog(q, true)
 	tr.LazyPrintf("opts: %s", opts)
+	if tenant.EnforceTenant() {
+		tenant.Log(ctx, tr)
+	}
 	defer func() {
 		if rl != nil {
 			tr.LazyPrintf("repos size: %d", len(rl.Repos))