tenant: set pprof label for tenant (#64338)

In the future this will allow us to attribute stack traces collected by
pprof to a tenant. This only sets it for the http middleware. I am
unsure how to achieve the same thing for grpc, since that uses
propogators.

Test Plan: captured a goroutine profile and saw some goroutines with a
tenant label

---------

Co-authored-by: Erik Seliger <[email protected]>

Author: keegancsmith

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

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,
 		),

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",

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))
+		})
 	})
 }
 

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)
+}

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))
+			})
+		}
 	})
 }