tenant: add tenant to trace

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

Note: Before, 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
zoekt_missing_tenant shows up as pprof profile.

internal/tenant/context.go

@@ -3,9 +3,7 @@ package tenant
 import (
 	"context"
 	"fmt"
-	"runtime/pprof"
-
-	"go.uber.org/atomic"
+	"strconv"
 
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
 	"github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype"
@@ -16,23 +14,24 @@ var ErrMissingTenant = fmt.Errorf("missing tenant")
 func FromContext(ctx context.Context) (*tenanttype.Tenant, error) {
 	tnt, ok := tenanttype.GetTenant(ctx)
 	if !ok {
-		if pprofMissingTenant != 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)
-		}
-
 		return nil, ErrMissingTenant
 	}
 	return tnt, nil
 }
 
-var pprofUniqID atomic.Int64
-var pprofMissingTenant = func() *pprof.Profile {
-	if !enforcement.ShouldLogNoTenant() {
-		return nil
+// IDToString is a helper function that returns a printable string of the tenant
+// ID in the context. This is useful for logging.
+func IDToString(ctx context.Context) string {
+	tnt, ok := tenanttype.GetTenant(ctx)
+	if !ok {
+		if enforcement.PPROFMissingTenant != nil {
+			// We want to track every stack trace, so need a unique value for the event
+			eventValue := enforcement.PPROFUniqID.Add(1)
+
+			// skip stack for Add and this function (2).
+			enforcement.PPROFMissingTenant.Add(eventValue, 2)
+		}
+		return "missing"
 	}
-	return pprof.NewProfile("missing_tenant")
-}()
+	return strconv.Itoa(tnt.ID())
+}

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,11 +1,24 @@
 package enforcement
 
-import "go.uber.org/atomic"
+import (
+	"os"
+	"runtime/pprof"
+
+	"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
+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

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() {
+		tr.LazyPrintf("tenant: %s", tenant.IDToString(ctx))
+	}
 	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() {
+		tr.LazyPrintf("tenant: %s", tenant.IDToString(ctx))
+	}
 	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() {
+		tr.LazyPrintf("tenant: %s", tenant.IDToString(ctx))
+	}
 	defer func() {
 		if rl != nil {
 			tr.LazyPrintf("repos size: %d", len(rl.Repos))