From 9af4b77dfa1c16655735cbae70dc4ee01f292384 Mon Sep 17 00:00:00 2001 From: Stefan Hengl Date: Thu, 21 Nov 2024 08:47:50 +0100 Subject: [PATCH 1/3] 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 | 33 +++++++++---------- internal/tenant/enforcement.go | 10 ------ .../internal/enforcement/enforcement.go | 17 ++++++++-- shards/eval.go | 10 ++++++ 4 files changed, 41 insertions(+), 29 deletions(-) diff --git a/internal/tenant/context.go b/internal/tenant/context.go index fa56f3a85..b95e3724d 100644 --- a/internal/tenant/context.go +++ b/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()) +} diff --git a/internal/tenant/enforcement.go b/internal/tenant/enforcement.go index 82aa668ce..a5cab97af 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 82a842f19..ee477cf6b 100644 --- a/internal/tenant/internal/enforcement/enforcement.go +++ b/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 diff --git a/shards/eval.go b/shards/eval.go index caa0e6440..3138609a6 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() { + 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)) From 6d00dff1c074a2a0ba5c597311a790328ed1baac Mon Sep 17 00:00:00 2001 From: Stefan Hengl Date: Thu, 21 Nov 2024 11:50:55 +0100 Subject: [PATCH 2/3] PR comments --- internal/tenant/context.go | 38 +++++++++++++++++-- .../internal/enforcement/enforcement.go | 22 ----------- 2 files changed, 35 insertions(+), 25 deletions(-) diff --git a/internal/tenant/context.go b/internal/tenant/context.go index b95e3724d..1b35b551b 100644 --- a/internal/tenant/context.go +++ b/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 + } +} diff --git a/internal/tenant/internal/enforcement/enforcement.go b/internal/tenant/internal/enforcement/enforcement.go index ee477cf6b..be2c54946 100644 --- a/internal/tenant/internal/enforcement/enforcement.go +++ b/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 - } -} From 5a6c46a05d940de30b967ba2848808e04c5b438d Mon Sep 17 00:00:00 2001 From: Stefan Hengl Date: Thu, 21 Nov 2024 12:11:05 +0100 Subject: [PATCH 3/3] IDToString -> Log --- internal/tenant/context.go | 13 +++++++------ shards/eval.go | 6 +++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/internal/tenant/context.go b/internal/tenant/context.go index 1b35b551b..afd887dab 100644 --- a/internal/tenant/context.go +++ b/internal/tenant/context.go @@ -4,13 +4,13 @@ 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" + "github.com/sourcegraph/zoekt/trace" ) var ErrMissingTenant = fmt.Errorf("missing tenant") @@ -23,9 +23,9 @@ func FromContext(ctx context.Context) (*tenanttype.Tenant, error) { return tnt, 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 { +// 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 { @@ -35,9 +35,10 @@ func IDToString(ctx context.Context) string { // skip stack for Add and this function (2). profile.Add(eventValue, 2) } - return "missing" + tr.LazyPrintf("tenant: missing") + return } - return strconv.Itoa(tnt.ID()) + tr.LazyPrintf("tenant: %d", tnt.ID()) } var pprofUniqID atomic.Int64 diff --git a/shards/eval.go b/shards/eval.go index 3138609a6..c691649d5 100644 --- a/shards/eval.go +++ b/shards/eval.go @@ -21,7 +21,7 @@ func (s *typeRepoSearcher) Search(ctx context.Context, q query.Q, opts *zoekt.Se tr.LazyLog(q, true) tr.LazyPrintf("opts: %+v", opts) if tenant.EnforceTenant() { - tr.LazyPrintf("tenant: %s", tenant.IDToString(ctx)) + tenant.Log(ctx, tr) } defer func() { if sr != nil { @@ -48,7 +48,7 @@ func (s *typeRepoSearcher) StreamSearch(ctx context.Context, q query.Q, opts *zo tr.LazyLog(q, true) tr.LazyPrintf("opts: %+v", opts) if tenant.EnforceTenant() { - tr.LazyPrintf("tenant: %s", tenant.IDToString(ctx)) + tenant.Log(ctx, tr) } var stats zoekt.Stats defer func() { @@ -76,7 +76,7 @@ func (s *typeRepoSearcher) List(ctx context.Context, q query.Q, opts *zoekt.List tr.LazyLog(q, true) tr.LazyPrintf("opts: %s", opts) if tenant.EnforceTenant() { - tr.LazyPrintf("tenant: %s", tenant.IDToString(ctx)) + tenant.Log(ctx, tr) } defer func() { if rl != nil {