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