diff --git a/go.mod b/go.mod index 1e1d4427f17..2979630067e 100644 --- a/go.mod +++ b/go.mod @@ -284,6 +284,7 @@ require ( github.com/google/go-containerregistry v0.16.1 github.com/google/go-github/v48 v48.2.0 github.com/google/go-github/v55 v55.0.0 + github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.0-rc.0 github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.1.0 github.com/hashicorp/cronexpr v1.1.1 @@ -411,7 +412,6 @@ require ( github.com/gosimple/slug v1.12.0 // indirect github.com/gosimple/unidecode v1.0.1 // indirect github.com/grafana-tools/sdk v0.0.0-20220919052116-6562121319fc // indirect - github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 // indirect github.com/hashicorp/errwrap v1.1.0 // indirect github.com/hashicorp/go-multierror v1.1.1 // indirect github.com/hashicorp/go-slug v0.12.1 // indirect diff --git a/internal/grpc/defaults/defaults.go b/internal/grpc/defaults/defaults.go index 8aa8728ff6f..096481efdf9 100644 --- a/internal/grpc/defaults/defaults.go +++ b/internal/grpc/defaults/defaults.go @@ -185,6 +185,7 @@ func buildServerOptions(logger log.Logger, opts serverOptions) []grpc.ServerOpti propagator.StreamServerPropagator(tenant.TenantPropagator{}), propagator.StreamServerPropagator(actor.ActorPropagator{}), propagator.StreamServerPropagator(policy.ShouldTracePropagator{}), + tenant.StreamServerInterceptor, otelgrpc.StreamServerInterceptor(otelOpts...), //lint:ignore SA1019 the advertised replacement doesn't seem to be a drop-in replacement, use deprecated mechanism for now contextconv.StreamServerInterceptor, ), @@ -198,6 +199,7 @@ func buildServerOptions(logger log.Logger, opts serverOptions) []grpc.ServerOpti propagator.UnaryServerPropagator(tenant.TenantPropagator{}), propagator.UnaryServerPropagator(actor.ActorPropagator{}), propagator.UnaryServerPropagator(policy.ShouldTracePropagator{}), + tenant.UnaryServerInterceptor, otelgrpc.UnaryServerInterceptor(otelOpts...), //lint:ignore SA1019 the advertised replacement doesn't seem to be a drop-in replacement, use deprecated mechanism for now contextconv.UnaryServerInterceptor, ), diff --git a/internal/tenant/BUILD.bazel b/internal/tenant/BUILD.bazel index a1acea8424a..b6f267f3662 100644 --- a/internal/tenant/BUILD.bazel +++ b/internal/tenant/BUILD.bazel @@ -19,7 +19,9 @@ go_library( "//internal/testutil", "//internal/trace", "//lib/errors", + "@com_github_grpc_ecosystem_go_grpc_middleware//:go-grpc-middleware", "@com_github_sourcegraph_log//:log", + "@org_golang_google_grpc//:go_default_library", "@org_golang_google_grpc//codes", "@org_golang_google_grpc//metadata", "@org_golang_google_grpc//status", diff --git a/internal/tenant/externalhttp.go b/internal/tenant/externalhttp.go index ec7e2828208..3915791ed5a 100644 --- a/internal/tenant/externalhttp.go +++ b/internal/tenant/externalhttp.go @@ -4,6 +4,8 @@ import ( "context" "fmt" "net/http" + "runtime/pprof" + "strconv" "strings" "github.com/sourcegraph/sourcegraph/internal/errcode" @@ -37,7 +39,10 @@ func ExternalTenantFromHostnameMiddleware(tenantForHostname TenantHostnameMapper } ctx = withTenant(ctx, tenantID) - next.ServeHTTP(rw, req.WithContext(ctx)) + + pprof.Do(ctx, pprof.Labels("tenant", strconv.Itoa(tenantID)), func(ctx context.Context) { + next.ServeHTTP(rw, req.WithContext(ctx)) + }) }) } diff --git a/internal/tenant/grpc.go b/internal/tenant/grpc.go index d45082ac5ed..4f36b2cdef3 100644 --- a/internal/tenant/grpc.go +++ b/internal/tenant/grpc.go @@ -2,12 +2,16 @@ package tenant import ( "context" + "runtime/pprof" "strconv" + "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" + grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware" + "github.com/sourcegraph/sourcegraph/lib/errors" ) @@ -48,3 +52,33 @@ func (TenantPropagator) InjectContext(ctx context.Context, md metadata.MD) (cont return withTenant(ctx, uid), nil } } + +// UnaryServerInterceptor is a grpc.UnaryServerInterceptor that injects the tenant ID +// from the context into pprof labels. +func UnaryServerInterceptor(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (response any, err error) { + if tnt, err := FromContext(ctx); err == nil { + defer pprof.SetGoroutineLabels(ctx) + ctx = pprof.WithLabels(ctx, pprof.Labels("tenant", strconv.Itoa(tnt.ID()))) + pprof.SetGoroutineLabels(ctx) + } + + return handler(ctx, req) +} + +// StreamServerInterceptor is a grpc.StreamServerInterceptor that injects the tenant ID +// from the context into pprof labels. +func StreamServerInterceptor(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + if tnt, err := FromContext(ss.Context()); err == nil { + ctx := ss.Context() + defer pprof.SetGoroutineLabels(ctx) + ctx = pprof.WithLabels(ctx, pprof.Labels("tenant", strconv.Itoa(tnt.ID()))) + pprof.SetGoroutineLabels(ctx) + + ss = &grpc_middleware.WrappedServerStream{ + ServerStream: ss, + WrappedContext: ctx, + } + } + + return handler(srv, ss) +} diff --git a/internal/tenant/http.go b/internal/tenant/http.go index f664d98aa6b..8fe81ad9bd1 100644 --- a/internal/tenant/http.go +++ b/internal/tenant/http.go @@ -1,8 +1,10 @@ package tenant import ( + "context" "fmt" "net/http" + "runtime/pprof" "strconv" "github.com/sourcegraph/log" @@ -84,8 +86,10 @@ func InternalHTTPMiddleware(logger log.Logger, next http.Handler) http.Handler { // Valid tenant ctx = withTenant(ctx, tntID) - } - next.ServeHTTP(rw, req.WithContext(ctx)) + pprof.Do(ctx, pprof.Labels("tenant", strconv.Itoa(tntID)), func(ctx context.Context) { + next.ServeHTTP(rw, req.WithContext(ctx)) + }) + } }) }