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