From c1295f90e53275859ad63389c4f8a7e94b45bb08 Mon Sep 17 00:00:00 2001 From: Stefan Hengl Date: Thu, 21 Nov 2024 12:20:56 +0100 Subject: [PATCH] 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 | 52 +++++++++++++++---- internal/tenant/enforcement.go | 10 ---- .../internal/enforcement/enforcement.go | 21 +++----- shards/eval.go | 10 ++++ 4 files changed, 58 insertions(+), 35 deletions(-) diff --git a/internal/tenant/context.go b/internal/tenant/context.go index fa56f3a8..afd887da 100644 --- a/internal/tenant/context.go +++ b/internal/tenant/context.go @@ -4,11 +4,13 @@ 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") @@ -16,23 +18,53 @@ 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") -}() +} diff --git a/internal/tenant/enforcement.go b/internal/tenant/enforcement.go index 82aa668c..a5cab97a 100644 --- a/internal/tenant/enforcement.go +++ b/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": diff --git a/internal/tenant/internal/enforcement/enforcement.go b/internal/tenant/internal/enforcement/enforcement.go index 82a842f1..be2c5494 100644 --- a/internal/tenant/internal/enforcement/enforcement.go +++ b/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")) diff --git a/shards/eval.go b/shards/eval.go index caa0e644..c691649d 100644 --- a/shards/eval.go +++ b/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))