Merge #156034

156034: drpc, server: add request counter telemetry in DRPC-gateway r=cthumuluru-crdb a=Nukitt

gRPC-Gateway provides basic request counting via
https://github.com/cockroachdb/cockroach/blob/04f49359a573681a93d9724f37b6582b420eaf1b/pkg/server/grpc_gateway.go#L88

This patch adds support for request counting in drpc-gateway similar to how gRPC-Gateway provides and will help in production monitoring.

Epic: CRDB-48924
Fixes: #151399
Release note: None

Co-authored-by: Nukitt <nukit.tailor@cockroachlabs.com>

Author: craig[bot]

pkg/rpc/BUILD.bazel

@@ -41,6 +41,7 @@ go_library(
         "//pkg/security",
         "//pkg/security/certnames",
         "//pkg/security/username",
+        "//pkg/server/telemetry",
         "//pkg/settings",
         "//pkg/settings/cluster",
         "//pkg/ts/tspb",

pkg/rpc/context.go

@@ -656,6 +656,12 @@ func NewContext(ctx context.Context, opts ContextOptions) *Context {
 		rpcCtx.clientStreamInterceptorsDRPC = append(rpcCtx.clientStreamInterceptorsDRPC,
 			drpcinterceptor.StreamClientInterceptor(tracer, tagger))
 	}
+
+	// Add the DRPC gateway request counter interceptor to track telemetry for
+	// HTTP gateway requests.
+	rpcCtx.clientUnaryInterceptorsDRPC = append(rpcCtx.clientUnaryInterceptorsDRPC,
+		drpcGatewayRequestCounterInterceptor)
+
 	// Note that we do not consult rpcCtx.Knobs.StreamClientInterceptor. That knob
 	// can add another interceptor, but it can only do it dynamically, based on
 	// a connection class. Only calls going over an actual gRPC connection will

pkg/rpc/drpc.go

@@ -220,7 +220,7 @@ func NewDRPCServer(_ context.Context, rpcCtx *Context, opts ...ServerOption) (DR
 	streamInterceptors = append(streamInterceptors, stopStream)
 
 	// Recover from any uncaught panics caused by DB Console requests.
-	unaryInterceptors = append(unaryInterceptors, DRPCGatewayRequestRecoveryInterceptor)
+	unaryInterceptors = append(unaryInterceptors, drpcGatewayRequestRecoveryInterceptor)
 
 	// If the metrics interceptor is set, it should be registered second so
 	// that all other interceptors are included in the response time durations.

pkg/rpc/drpc_test.go

@@ -98,7 +98,7 @@ func TestGatewayRequestDRPCRecoveryInterceptor(t *testing.T) {
 			panic("test panic")
 		}
 
-		resp, err := DRPCGatewayRequestRecoveryInterceptor(ctx, nil, "test", handler)
+		resp, err := drpcGatewayRequestRecoveryInterceptor(ctx, nil, "test", handler)
 
 		require.Nil(t, resp)
 		require.ErrorContains(t, err, "unexpected error occurred")
@@ -118,7 +118,7 @@ func TestGatewayRequestDRPCRecoveryInterceptor(t *testing.T) {
 			}
 		}()
 
-		_, _ = DRPCGatewayRequestRecoveryInterceptor(ctx, nil, "test", handler)
+		_, _ = drpcGatewayRequestRecoveryInterceptor(ctx, nil, "test", handler)
 	})
 
 	// With gateway metadata but no panic - should pass through normally
@@ -131,7 +131,7 @@ func TestGatewayRequestDRPCRecoveryInterceptor(t *testing.T) {
 			return expectedResp, expectedErr
 		}
 
-		resp, err := DRPCGatewayRequestRecoveryInterceptor(ctx, nil, "test", handler)
+		resp, err := drpcGatewayRequestRecoveryInterceptor(ctx, nil, "test", handler)
 
 		require.Equal(t, expectedResp, resp)
 		require.ErrorIs(t, err, expectedErr)

pkg/rpc/metrics.go

@@ -7,11 +7,13 @@ package rpc
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"strings"
 
 	"github.com/VividCortex/ewma"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/util/log/logcrash"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/cockroach/pkg/util/metric/aggmetric"
@@ -23,6 +25,8 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
+	"storj.io/drpc"
+	"storj.io/drpc/drpcclient"
 	"storj.io/drpc/drpcmetadata"
 	"storj.io/drpc/drpcmux"
 )
@@ -557,10 +561,10 @@ func MarkDRPCGatewayRequest(ctx context.Context) context.Context {
 	return drpcmetadata.Add(ctx, gwRequestKey, "true")
 }
 
-// DRPCGatewayRequestRecoveryInterceptor recovers from panics in DRPC handlers
+// drpcGatewayRequestRecoveryInterceptor recovers from panics in DRPC handlers
 // that are invoked due to DB console requests. For these requests, we do not
 // want an uncaught panic to crash the node.
-func DRPCGatewayRequestRecoveryInterceptor(
+func drpcGatewayRequestRecoveryInterceptor(
 	ctx context.Context, req interface{}, rpc string, handler drpcmux.UnaryHandler,
 ) (resp interface{}, err error) {
 	if val, ok := drpcmetadata.GetValue(ctx, gwRequestKey); ok && val != "" {
@@ -574,3 +578,29 @@ func DRPCGatewayRequestRecoveryInterceptor(
 	resp, err = handler(ctx, req)
 	return resp, err
 }
+
+// drpcGatewayRequestCounterInterceptor is a client-side interceptor that
+// increments telemetry counters for DRPC requests originating from the HTTP
+// gateway. It checks for the gateway request marker and increments
+// a counter named after the RPC method.
+func drpcGatewayRequestCounterInterceptor(
+	ctx context.Context,
+	rpc string,
+	enc drpc.Encoding,
+	in, out drpc.Message,
+	cc *drpcclient.ClientConn,
+	invoker drpcclient.UnaryInvoker,
+) error {
+	// Check if this request originated from the DRPC HTTP gateway
+	if val, ok := drpcmetadata.GetValue(ctx, gwRequestKey); ok && val != "" {
+		telemetry.Inc(getDRPCGatewayEndpointCounter(rpc))
+	}
+	return invoker(ctx, rpc, enc, in, out, cc)
+}
+
+// getDRPCGatewayEndpointCounter returns a telemetry Counter corresponding to
+// the given DRPC method.
+func getDRPCGatewayEndpointCounter(method string) telemetry.Counter {
+	const counterPrefix = "http.drpc-gateway"
+	return telemetry.GetCounter(fmt.Sprintf("%s.%s", counterPrefix, method))
+}