Add traffic stats handler to apps and proxies (#11590)

Fixes: buildbuddy-io/buildbuddy-internal#6134

Author: vanja-p

enterprise/server/cmd/cache_proxy/BUILD

@@ -30,6 +30,7 @@ go_library(
         "//enterprise/server/routing/routing_content_addressable_storage_client",
         "//enterprise/server/routing/routing_service",
         "//enterprise/server/util/proxy_util",
+        "//enterprise/server/util/trafficstats",
         "//proto:cache_service_go_proto",
         "//proto:oci_fetcher_go_proto",
         "//proto:remote_asset_go_proto",
@@ -57,6 +58,7 @@ go_library(
         "@org_golang_google_genproto_googleapis_bytestream//:bytestream",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//channelz/service",
+        "@org_golang_google_grpc//stats",
     ],
 )
 

enterprise/server/cmd/cache_proxy/cache_proxy.go

@@ -29,6 +29,7 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/routing/routing_content_addressable_storage_client"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/routing/routing_service"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/util/proxy_util"
+	"github.com/buildbuddy-io/buildbuddy/enterprise/server/util/trafficstats"
 	"github.com/buildbuddy-io/buildbuddy/server/cache_server"
 	"github.com/buildbuddy-io/buildbuddy/server/config"
 	"github.com/buildbuddy-io/buildbuddy/server/real_environment"
@@ -57,6 +58,7 @@ import (
 	http_interceptors "github.com/buildbuddy-io/buildbuddy/server/http/interceptors"
 	bspb "google.golang.org/genproto/googleapis/bytestream"
 	channelzservice "google.golang.org/grpc/channelz/service"
+	"google.golang.org/grpc/stats"
 )
 
 var (
@@ -202,6 +204,10 @@ func main() {
 }
 
 func startGRPCServers(env *real_environment.RealEnv) error {
+	trafficStatsHandler, err := trafficstats.NewServerHandler()
+	if err != nil {
+		return status.WrapError(err, "failed to create traffic stats handler")
+	}
 	// Add the API-Key, JWT, client-identity, etc... propagating interceptor.
 	grpcServerConfig := grpc_server.GRPCServerConfig{
 		ExtraChainedUnaryInterceptors: []grpc.UnaryServerInterceptor{
@@ -210,6 +216,7 @@ func startGRPCServers(env *real_environment.RealEnv) error {
 		ExtraChainedStreamInterceptors: []grpc.StreamServerInterceptor{
 			interceptors.PropagateMetadataStreamInterceptor(proxy_util.HeadersToPropagate...),
 		},
+		ExtraStatsHandlers: []stats.Handler{trafficStatsHandler},
 	}
 
 	// Start and register internal servers first because the external proxy

enterprise/server/cmd/server/BUILD

@@ -74,6 +74,7 @@ go_library(
         "//enterprise/server/usage_service",
         "//enterprise/server/util/dsingleflight",
         "//enterprise/server/util/redisutil",
+        "//enterprise/server/util/trafficstats",
         "//enterprise/server/webhooks/bitbucket",
         "//enterprise/server/webhooks/github",
         "//enterprise/server/workflow/service",
@@ -86,10 +87,12 @@ go_library(
         "//server/remote_cache/capabilities_server",
         "//server/telemetry",
         "//server/util/clickhouse",
+        "//server/util/grpc_server",
         "//server/util/healthcheck",
         "//server/util/log",
         "//server/util/tracing",
         "//server/version",
+        "@org_golang_google_grpc//stats",
     ],
 )
 

enterprise/server/cmd/server/main.go

@@ -59,6 +59,7 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/usage_service"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/util/dsingleflight"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/util/redisutil"
+	"github.com/buildbuddy-io/buildbuddy/enterprise/server/util/trafficstats"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/webhooks/bitbucket"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/webhooks/github"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/workspace"
@@ -69,12 +70,16 @@ import (
 	"github.com/buildbuddy-io/buildbuddy/server/real_environment"
 	"github.com/buildbuddy-io/buildbuddy/server/remote_cache/capabilities_server"
 	"github.com/buildbuddy-io/buildbuddy/server/telemetry"
+	"github.com/buildbuddy-io/buildbuddy/server/util/grpc_server"
+
 	"github.com/buildbuddy-io/buildbuddy/server/util/clickhouse"
 	"github.com/buildbuddy-io/buildbuddy/server/util/healthcheck"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/tracing"
 	"github.com/buildbuddy-io/buildbuddy/server/version"
 
+	"google.golang.org/grpc/stats"
+
 	enterprise_app_bundle "github.com/buildbuddy-io/buildbuddy/enterprise/app"
 	remote_execution_redis_client "github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/redis_client"
 	telserver "github.com/buildbuddy-io/buildbuddy/enterprise/server/telemetry"
@@ -342,5 +347,11 @@ func main() {
 		log.Fatalf("%v", err)
 	}
 
-	libmain.StartAndRunServices(realEnv) // Returns after graceful shutdown
+	trafficStatsHandler, err := trafficstats.NewServerHandler()
+	if err != nil {
+		log.Fatalf("Error creating traffic stats handler: %v", err)
+	}
+	libmain.StartAndRunServices(realEnv, grpc_server.GRPCServerConfig{
+		ExtraStatsHandlers: []stats.Handler{trafficStatsHandler},
+	}) // Returns after graceful shutdown
 }

server/cmd/buildbuddy/BUILD

@@ -71,6 +71,7 @@ go_library(
         "//server/libmain",
         "//server/remote_cache/capabilities_server",
         "//server/telemetry",
+        "//server/util/grpc_server",
         "//server/util/healthcheck",
         "//server/util/log",
         "//server/version",

server/cmd/buildbuddy/main.go

@@ -3,6 +3,8 @@ package main
 import (
 	"flag"
 
+	"github.com/buildbuddy-io/buildbuddy/server/util/grpc_server"
+
 	"github.com/buildbuddy-io/buildbuddy/server/config"
 	"github.com/buildbuddy-io/buildbuddy/server/janitor"
 	"github.com/buildbuddy-io/buildbuddy/server/libmain"
@@ -54,5 +56,5 @@ func main() {
 	if err := capabilities_server.Register(env); err != nil {
 		log.Fatalf("%v", err)
 	}
-	libmain.StartAndRunServices(env) // Does not return
+	libmain.StartAndRunServices(env, grpc_server.GRPCServerConfig{}) // Does not return
 }

server/libmain/libmain.go

@@ -247,8 +247,8 @@ func registerInternalServices(env *real_environment.RealEnv, grpcServer *grpc.Se
 	channelzservice.RegisterChannelzServiceToServer(grpcServer)
 }
 
-func startGRPCServers(env *real_environment.RealEnv) error {
-	b, err := grpc_server.New(env, grpc_server.GRPCPort(), false /*=ssl*/, grpc_server.GRPCServerConfig{})
+func startGRPCServers(env *real_environment.RealEnv, config grpc_server.GRPCServerConfig) error {
+	b, err := grpc_server.New(env, grpc_server.GRPCPort(), false /*=ssl*/, config)
 	if err != nil {
 		return err
 	}
@@ -260,7 +260,7 @@ func startGRPCServers(env *real_environment.RealEnv) error {
 	grpc_server.EnableGRPCOverHTTP(env, b.GetServer())
 
 	if env.GetSSLService().IsEnabled() {
-		sb, err := grpc_server.New(env, grpc_server.GRPCSPort(), true /*=ssl*/, grpc_server.GRPCServerConfig{})
+		sb, err := grpc_server.New(env, grpc_server.GRPCSPort(), true /*=ssl*/, config)
 		if err != nil {
 			return err
 		}
@@ -399,7 +399,7 @@ func RegisterLocalServersAndClients(env *real_environment.RealEnv) {
 	}
 }
 
-func StartAndRunServices(env *real_environment.RealEnv) {
+func StartAndRunServices(env *real_environment.RealEnv, grpcConfig grpc_server.GRPCServerConfig) {
 	if *maxThreads > 0 {
 		debug.SetMaxThreads(*maxThreads)
 	}
@@ -427,7 +427,7 @@ func StartAndRunServices(env *real_environment.RealEnv) {
 		log.Fatalf("%v", err)
 	}
 
-	if err := startGRPCServers(env); err != nil {
+	if err := startGRPCServers(env, grpcConfig); err != nil {
 		log.Fatalf("%v", err)
 	}