fork of https://github.com/sourcegraph/zoekt
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.
author
Stefan Hengl
committer
GitHub
date
(Nov 21, 2024, 12:20 PM +0100)
commit
c1295f90
c1295f90e53275859ad63389c4f8a7e94b45bb08
parent
c52a9cde
c52a9cdee50568ec642078b2794029735343fc35
+58
-35
+42
-10
internal/tenant/context.go
+42
-10
internal/tenant/context.go
···
4
4
"context"
5
5
"fmt"
6
6
"runtime/pprof"
7
+
"sync"
7
8
8
9
"go.uber.org/atomic"
9
10
10
11
"github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
11
12
"github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype"
13
+
"github.com/sourcegraph/zoekt/trace"
12
14
)
13
15
14
16
var ErrMissingTenant = fmt.Errorf("missing tenant")
···
16
18
func FromContext(ctx context.Context) (*tenanttype.Tenant, error) {
17
19
tnt, ok := tenanttype.GetTenant(ctx)
18
20
if !ok {
19
-
if pprofMissingTenant != nil {
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.
28
+
func Log(ctx context.Context, tr *trace.Trace) {
29
+
tnt, ok := tenanttype.GetTenant(ctx)
30
+
if !ok {
31
+
if profile := pprofMissingTenant(); profile != nil {
20
32
// We want to track every stack trace, so need a unique value for the event
21
33
eventValue := pprofUniqID.Add(1)
22
34
23
35
// skip stack for Add and this function (2).
24
-
pprofMissingTenant.Add(eventValue, 2)
36
+
profile.Add(eventValue, 2)
25
37
}
26
-
27
-
return nil, ErrMissingTenant
38
+
tr.LazyPrintf("tenant: missing")
39
+
return
28
40
}
29
-
return tnt, nil
41
+
tr.LazyPrintf("tenant: %d", tnt.ID())
30
42
}
31
43
32
44
var pprofUniqID atomic.Int64
33
-
var pprofMissingTenant = func() *pprof.Profile {
34
-
if !enforcement.ShouldLogNoTenant() {
35
-
return nil
45
+
var pprofOnce sync.Once
46
+
var pprofProfile *pprof.Profile
47
+
48
+
// pprofMissingTenant returns the pprof profile for missing tenants,
49
+
// initializing it only once.
50
+
func 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.
63
+
func shouldLogNoTenant() bool {
64
+
switch enforcement.EnforcementMode.Load() {
65
+
case "logging", "strict":
66
+
return true
67
+
default:
68
+
return false
36
69
}
37
-
return pprof.NewProfile("missing_tenant")
38
-
}()
70
+
}
-10
internal/tenant/enforcement.go
-10
internal/tenant/enforcement.go
···
1
1
package tenant
2
2
3
3
import (
4
-
"os"
5
-
6
4
"github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement"
7
5
)
8
-
9
-
func init() {
10
-
v, ok := os.LookupEnv("SRC_TENANT_ENFORCEMENT_MODE")
11
-
if !ok {
12
-
v = "disabled"
13
-
}
14
-
enforcement.EnforcementMode.Store(v)
15
-
}
16
6
17
7
func EnforceTenant() bool {
18
8
switch enforcement.EnforcementMode.Load() {
+6
-15
internal/tenant/internal/enforcement/enforcement.go
+6
-15
internal/tenant/internal/enforcement/enforcement.go
···
1
1
package enforcement
2
2
3
-
import "go.uber.org/atomic"
3
+
import (
4
+
"os"
5
+
6
+
"go.uber.org/atomic"
7
+
)
4
8
5
9
// EnforcementMode is the current tenant enforcement mode. It resides here
6
10
// instead of in the tenant package to avoid a circular dependency. See
7
11
// tenanttest.MockEnforce.
8
-
var EnforcementMode atomic.String
9
-
10
-
// ShouldLogNoTenant returns true if the tenant enforcement mode is logging or strict.
11
-
// It is used to log a warning if a request to a low-level store is made without a tenant
12
-
// so we can identify missing tenants. This will go away and only strict will be allowed
13
-
// once we are confident that all contexts carry tenants.
14
-
func ShouldLogNoTenant() bool {
15
-
switch EnforcementMode.Load() {
16
-
case "logging", "strict":
17
-
return true
18
-
default:
19
-
return false
20
-
}
21
-
}
12
+
var EnforcementMode = atomic.NewString(os.Getenv("SRC_TENANT_ENFORCEMENT_MODE"))
+10
shards/eval.go
+10
shards/eval.go
···
4
4
"context"
5
5
6
6
"github.com/sourcegraph/zoekt"
7
+
"github.com/sourcegraph/zoekt/internal/tenant"
7
8
"github.com/sourcegraph/zoekt/query"
8
9
"github.com/sourcegraph/zoekt/trace"
9
10
)
···
19
20
tr, ctx := trace.New(ctx, "typeRepoSearcher.Search", "")
20
21
tr.LazyLog(q, true)
21
22
tr.LazyPrintf("opts: %+v", opts)
23
+
if tenant.EnforceTenant() {
24
+
tenant.Log(ctx, tr)
25
+
}
22
26
defer func() {
23
27
if sr != nil {
24
28
tr.LazyPrintf("num files: %d", len(sr.Files))
···
43
47
tr, ctx := trace.New(ctx, "typeRepoSearcher.StreamSearch", "")
44
48
tr.LazyLog(q, true)
45
49
tr.LazyPrintf("opts: %+v", opts)
50
+
if tenant.EnforceTenant() {
51
+
tenant.Log(ctx, tr)
52
+
}
46
53
var stats zoekt.Stats
47
54
defer func() {
48
55
tr.LazyPrintf("stats: %+v", stats)
···
68
75
tr, ctx := trace.New(ctx, "typeRepoSearcher.List", "")
69
76
tr.LazyLog(q, true)
70
77
tr.LazyPrintf("opts: %s", opts)
78
+
if tenant.EnforceTenant() {
79
+
tenant.Log(ctx, tr)
80
+
}
71
81
defer func() {
72
82
if rl != nil {
73
83
tr.LazyPrintf("repos size: %d", len(rl.Repos))