tenant: pprof for missing tenant in FromContext (#64405)

Instead of adding instrumentation for each place we expect a tenant, we
instead assume that we always expect a tenant if FromContext is called.
This should allow us to not miss any places.

Note: this does remove the logging we had in diskcache, but it seems
more reasonable to just use pprof for something that can be noisy?

Test Plan: just gonna rely on CI since this defaults to off.

internal/database/dbconn/BUILD.bazel

@@ -15,7 +15,6 @@ go_library(
         "metrics.go",
         "open.go",
         "postgres_version.go",
-        "tenant_pprof.go",
     ],
     importpath = "github.com/sourcegraph/sourcegraph/internal/database/dbconn",
     visibility = ["//:__subpackages__"],
@@ -23,7 +22,6 @@ go_library(
         "//internal/database/dbconn/rds",
         "//internal/env",
         "//internal/lazyregexp",
-        "//internal/tenant",
         "//lib/errors",
         "//lib/output",
         "@com_github_jackc_pgx_v4//:pgx",

internal/database/dbconn/open.go

@@ -220,13 +220,11 @@ func (n *extendedConn) CheckNamedValue(namedValue *driver.NamedValue) error {
 }
 
 func (n *extendedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
-	pprofCheckTenantlessQuery(ctx)
 	ctx, query = instrumentQuery(ctx, query, len(args))
 	return n.execerContext.ExecContext(ctx, query, args)
 }
 
 func (n *extendedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
-	pprofCheckTenantlessQuery(ctx)
 	ctx, query = instrumentQuery(ctx, query, len(args))
 	return n.queryerContext.QueryContext(ctx, query, args)
 }

internal/database/dbconn/tenant_pprof.go

@@ -1,33 +0,0 @@
-package dbconn
-
-import (
-	"context"
-	"runtime/pprof"
-	"sync/atomic"
-
-	"github.com/sourcegraph/sourcegraph/internal/tenant"
-)
-
-var pprofUniqID atomic.Int64
-var pprofTenantlessQueries = func() *pprof.Profile {
-	if !tenant.ShouldLogNoTenant() {
-		return nil
-	}
-	return pprof.NewProfile("tenantless_queries")
-}()
-
-func pprofCheckTenantlessQuery(ctx context.Context) {
-	if pprofTenantlessQueries == nil {
-		return
-	}
-
-	if _, err := tenant.FromContext(ctx); err == nil {
-		return
-	}
-
-	// 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).
-	pprofTenantlessQueries.Add(eventValue, 2)
-}

internal/diskcache/cache.go

@@ -10,7 +10,6 @@ import (
 	"log" //nolint:logging // TODO move all logging to sourcegraph/log
 	"os"
 	"path/filepath"
-	"runtime"
 	"sort"
 	"strconv"
 	"strings"
@@ -205,11 +204,6 @@ func (s *store) OpenWithPath(ctx context.Context, key []string, fetcher FetcherW
 
 // path returns the path for key.
 func (s *store) path(ctx context.Context, key []string) (string, error) {
-	if tenant.ShouldLogNoTenant() {
-		if _, err := tenant.FromContext(ctx); err != nil {
-			log.Printf("diskcache: %s:\n%s\n", err, captureStackTrace())
-		}
-	}
 	if !tenant.EnforceTenant() {
 		return s.pathNoTenant(key), nil
 	}
@@ -434,15 +428,3 @@ func fsync(path string) error {
 	}
 	return err
 }
-
-func captureStackTrace() string {
-	// Allocate a large enough buffer to capture the stack trace
-	buf := make([]byte, 1024)
-	for {
-		n := runtime.Stack(buf, false)
-		if n < len(buf) {
-			return string(buf[:n])
-		}
-		buf = make([]byte, len(buf)*2)
-	}
-}

internal/tenant/context.go

@@ -2,6 +2,7 @@ package tenant
 
 import (
 	"context"
+	"runtime/pprof"
 
 	"go.uber.org/atomic"
 
@@ -20,6 +21,14 @@ var ErrNoTenantInContext = errors.New("no tenant in context")
 func FromContext(ctx context.Context) (*Tenant, error) {
 	tnt, ok := ctx.Value(tenantKey).(*Tenant)
 	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, ErrNoTenantInContext
 	}
 	return tnt, nil
@@ -73,3 +82,11 @@ func NewTestContext() context.Context {
 	}
 	return withTenant(context.Background(), int(tenantCounter.Inc()))
 }
+
+var pprofUniqID atomic.Int64
+var pprofMissingTenant = func() *pprof.Profile {
+	if !shouldLogNoTenant() {
+		return nil
+	}
+	return pprof.NewProfile("missing_tenant")
+}()

internal/tenant/enforcement.go

@@ -9,11 +9,11 @@ import (
 
 var enforcementMode = env.Get("SRC_TENANT_ENFORCEMENT_MODE", "disabled", "INTERNAL: enforcement mode for tenant isolation. Valid values: disabled, logging, strict")
 
-// ShouldLogNoTenant returns true if the tenant enforcement mode is logging or strict.
+// 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 {
+func shouldLogNoTenant() bool {
 	switch enforcementMode {
 	case "logging", "strict":
 		return true