fork of https://github.com/sourcegraph/zoekt
1package tenant
2
3import (
4 "context"
5 "fmt"
6 "runtime/pprof"
7 "sync"
8
9 "go.uber.org/atomic"
10
11 "github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
12 "github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype"
13 "github.com/sourcegraph/zoekt/trace"
14)
15
16var ErrMissingTenant = fmt.Errorf("missing tenant")
17
18func FromContext(ctx context.Context) (*tenanttype.Tenant, error) {
19 tnt, ok := tenanttype.GetTenant(ctx)
20 if !ok {
21 return nil, ErrMissingTenant
22 }
23 return tnt, nil
24}
25
26// Log logs the tenant ID to the trace. If tenant logging is enabled, it also
27// logs a stack trace to a pprof profile.
28func Log(ctx context.Context, tr *trace.Trace) {
29 tnt, ok := tenanttype.GetTenant(ctx)
30 if !ok {
31 if profile := pprofMissingTenant(); profile != nil {
32 // We want to track every stack trace, so need a unique value for the event
33 eventValue := pprofUniqID.Add(1)
34
35 // skip stack for Add and this function (2).
36 profile.Add(eventValue, 2)
37 }
38 tr.LazyPrintf("tenant: missing")
39 return
40 }
41 tr.LazyPrintf("tenant: %d", tnt.ID())
42}
43
44var pprofUniqID atomic.Int64
45var pprofOnce sync.Once
46var pprofProfile *pprof.Profile
47
48// pprofMissingTenant returns the pprof profile for missing tenants,
49// initializing it only once.
50func pprofMissingTenant() *pprof.Profile {
51 pprofOnce.Do(func() {
52 if shouldLogNoTenant() {
53 pprofProfile = pprof.NewProfile("missing_tenant")
54 }
55 })
56 return pprofProfile
57}
58
59// shouldLogNoTenant returns true if the tenant enforcement mode is logging or strict.
60// It is used to log a warning if a request to a low-level store is made without a tenant
61// so we can identify missing tenants. This will go away and only strict will be allowed
62// once we are confident that all contexts carry tenants.
63func shouldLogNoTenant() bool {
64 switch enforcement.EnforcementMode.Load() {
65 case "logging", "strict":
66 return true
67 default:
68 return false
69 }
70}