rpc,server: add grpc server interceptor to instrument rpcs

Adds a new grpc.UnaryServerInterceptor that records
the duration of a grpc request. The metric is labeled
by the grpc full method name and status code. This
is disabled by default and can be enabled by setting
`server.grpc.metrics.enabled` to true. Only grpc
methods that begin with "/cockroach.server" or
"/cockroach.ts" will be recorded.

This metric uses metric.HistogramVec and is exported
to external collectors.

Epic: None
Release note: None

Author: kyle-a-wong

docs/generated/metrics/metrics.html

@@ -1594,6 +1594,7 @@
 <tr><td>APPLICATION</td><td>rpc.connection.inactive</td><td>Gauge of current connections in an inactive state and pending deletion; these are not healthy but are not tracked as unhealthy either because there is reason to believe that the connection is no longer relevant,for example if the node has since been seen under a new address</td><td>Connections</td><td>GAUGE</td><td>COUNT</td><td>AVG</td><td>NONE</td></tr>
 <tr><td>APPLICATION</td><td>rpc.connection.unhealthy</td><td>Gauge of current connections in an unhealthy state (not bidirectionally connected or heartbeating)</td><td>Connections</td><td>GAUGE</td><td>COUNT</td><td>AVG</td><td>NONE</td></tr>
 <tr><td>APPLICATION</td><td>rpc.connection.unhealthy_nanos</td><td>Gauge of nanoseconds of unhealthy connection time.<br/><br/>On the prometheus endpoint scraped with the cluster setting &#39;server.child_metrics.enabled&#39; set,<br/>the constituent parts of this metric are available on a per-peer basis and one can read off<br/>for how long a given peer has been unreachable</td><td>Nanoseconds</td><td>GAUGE</td><td>NANOSECONDS</td><td>AVG</td><td>NONE</td></tr>
+<tr><td>APPLICATION</td><td>rpc.server.request.duration.nanos</td><td>Duration of an grpc request in nanoseconds.</td><td>Duration</td><td>HISTOGRAM</td><td>NANOSECONDS</td><td>AVG</td><td>NONE</td></tr>
 <tr><td>APPLICATION</td><td>schedules.BACKUP.failed</td><td>Number of BACKUP jobs failed</td><td>Jobs</td><td>COUNTER</td><td>COUNT</td><td>AVG</td><td>NON_NEGATIVE_DERIVATIVE</td></tr>
 <tr><td>APPLICATION</td><td>schedules.BACKUP.last-completed-time</td><td>The unix timestamp of the most recently completed backup by a schedule specified as maintaining this metric</td><td>Jobs</td><td>GAUGE</td><td>TIMESTAMP_SEC</td><td>AVG</td><td>NONE</td></tr>
 <tr><td>APPLICATION</td><td>schedules.BACKUP.last-completed-time-by-virtual_cluster</td><td>The unix timestamp of the most recently completed host scheduled backup by virtual cluster specified as maintaining this metric</td><td>Jobs</td><td>GAUGE</td><td>TIMESTAMP_SEC</td><td>AVG</td><td>NONE</td></tr>

pkg/roachprod/opentelemetry/cockroachdb_metrics.go

@@ -1505,6 +1505,10 @@ var cockroachdbMetrics = map[string]string{
 	"rpc_method_transferlease_recv":                               "rpc.method.transferlease.recv",
 	"rpc_method_truncatelog_recv":                                 "rpc.method.truncatelog.recv",
 	"rpc_method_writebatch_recv":                                  "rpc.method.writebatch.recv",
+	"rpc_server_request_duration_nanos":                           "rpc.server.request.duration.nanos",
+	"rpc_server_request_duration_nanos_bucket":                    "rpc.server.request.duration.nanos.bucket",
+	"rpc_server_request_duration_nanos_count":                     "rpc.server.request.duration.nanos.count",
+	"rpc_server_request_duration_nanos_sum":                       "rpc.server.request.duration.nanos.sum",
 	"rpc_streams_mux_rangefeed_active":                            "rpc.streams.mux_rangefeed.active",
 	"rpc_streams_mux_rangefeed_recv":                              "rpc.streams.mux_rangefeed.recv",
 	"rpc_streams_rangefeed_active":                                "rpc.streams.rangefeed.active",

pkg/rpc/BUILD.bazel

@@ -74,6 +74,7 @@ go_library(
         "@com_github_golang_protobuf//proto",
         "@com_github_golang_snappy//:snappy",
         "@com_github_montanaflynn_stats//:stats",
+        "@com_github_prometheus_client_model//go",
         "@com_github_vividcortex_ewma//:ewma",
         "@io_opentelemetry_go_otel//attribute",
         "@io_storj_drpc//:drpc",

pkg/rpc/context.go

@@ -143,6 +143,12 @@ func NewServerEx(
 		})
 	})
 
+	// If the metrics interceptor is set, it should be registered second so
+	// that all other interceptors are included in the response time durations.
+	if o.metricsInterceptor != nil {
+		unaryInterceptor = append(unaryInterceptor, grpc.UnaryServerInterceptor(o.metricsInterceptor))
+	}
+
 	if !rpcCtx.ContextOptions.Insecure {
 		a := kvAuth{
 			sv: &rpcCtx.Settings.SV,

pkg/rpc/context_test.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"reflect"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -2343,3 +2344,37 @@ func BenchmarkGRPCPing(b *testing.B) {
 		})
 	}
 }
+
+func TestMetricsInterceptor(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	ctx := context.Background()
+
+	stopper := stop.NewStopper()
+	defer stopper.Stop(context.Background())
+	stopper.SetTracer(tracing.NewTracer())
+
+	clock := timeutil.NewManualTime(timeutil.Unix(0, 1))
+	maxOffset := time.Duration(0)
+
+	serverCtx := newTestContext(uuid.MakeV4(), clock, maxOffset, stopper)
+	serverCtx.AdvertiseAddr = "127.0.0.1:8888"
+	serverCtx.NodeID.Set(context.Background(), 1)
+
+	var interceptor RequestMetricsInterceptor = func(
+		ctx context.Context,
+		req interface{},
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		return nil, nil
+	}
+
+	_, _, serverInterceptors, err := NewServerEx(ctx, serverCtx, WithMetricsServerInterceptor(interceptor))
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, len(serverInterceptors.UnaryInterceptors), 2)
+	// make sure that the RequestMetricsInterceptor is the second registered
+	// interceptor.
+	require.Equal(t,
+		reflect.ValueOf(interceptor).Pointer(),
+		reflect.ValueOf(serverInterceptors.UnaryInterceptors[1]).Pointer())
+}

pkg/rpc/metrics.go

@@ -6,13 +6,19 @@
 package rpc
 
 import (
+	"context"
 	"strings"
 
 	"github.com/VividCortex/ewma"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/cockroach/pkg/util/metric/aggmetric"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	prometheusgo "github.com/prometheus/client_model/go"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 var (
@@ -130,6 +136,13 @@ over this connection.
 		Help:        `Counter of TCP bytes received via gRPC on connections we initiated.`,
 		Measurement: "Bytes",
 	}
+	metaRequestDuration = metric.Metadata{
+		Name:        "rpc.server.request.duration.nanos",
+		Help:        "Duration of an grpc request in nanoseconds.",
+		Measurement: "Duration",
+		Unit:        metric.Unit_NANOSECONDS,
+		MetricType:  prometheusgo.MetricType_HISTOGRAM,
+	}
 )
 
 func (m *Metrics) makeLabels(k peerKey, remoteLocality roachpb.Locality) []string {
@@ -334,3 +347,59 @@ func (m *Metrics) acquire(k peerKey, l roachpb.Locality) (peerMetrics, localityM
 	pm.ConnectionInactive.Inc(1)
 	return pm, lm
 }
+
+const (
+	RpcMethodLabel     = "methodName"
+	RpcStatusCodeLabel = "statusCode"
+)
+
+// RequestMetrics contains metrics for RPC requests.
+type RequestMetrics struct {
+	Duration *metric.HistogramVec
+}
+
+func NewRequestMetrics() *RequestMetrics {
+	return &RequestMetrics{
+		Duration: metric.NewExportedHistogramVec(
+			metaRequestDuration,
+			metric.ResponseTime30sBuckets,
+			[]string{RpcMethodLabel, RpcStatusCodeLabel}),
+	}
+}
+
+type RequestMetricsInterceptor grpc.UnaryServerInterceptor
+
+// NewRequestMetricsInterceptor creates a new gRPC server interceptor that records
+// the duration of each RPC. The metric is labeled by the method name and the
+// status code of the RPC. The interceptor will only record durations if
+// shouldRecord returns true. Otherwise, this interceptor will be a no-op.
+func NewRequestMetricsInterceptor(
+	requestMetrics *RequestMetrics, shouldRecord func(fullMethodName string) bool,
+) RequestMetricsInterceptor {
+	return func(
+		ctx context.Context,
+		req interface{},
+		info *grpc.UnaryServerInfo,
+		handler grpc.UnaryHandler,
+	) (interface{}, error) {
+		if !shouldRecord(info.FullMethod) {
+			return handler(ctx, req)
+		}
+
+		startTime := timeutil.Now()
+		resp, err := handler(ctx, req)
+		duration := timeutil.Since(startTime)
+		var code codes.Code
+		if err != nil {
+			code = status.Code(err)
+		} else {
+			code = codes.OK
+		}
+
+		requestMetrics.Duration.Observe(map[string]string{
+			RpcMethodLabel:     info.FullMethod,
+			RpcStatusCodeLabel: code.String(),
+		}, float64(duration.Nanoseconds()))
+		return resp, err
+	}
+}

pkg/rpc/metrics_test.go

@@ -6,13 +6,20 @@
 package rpc
 
 import (
+	"context"
+	"fmt"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
 	io_prometheus_client "github.com/prometheus/client_model/go"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // TestMetricsRelease verifies that peerMetrics.release() removes tracking for
@@ -82,3 +89,84 @@ func TestMetricsRelease(t *testing.T) {
 	// We added one extra peer and one extra locality, verify counts.
 	require.Equal(t, expectedCount, verifyAllFields(m, 2))
 }
+
+func TestServerRequestInstrumentInterceptor(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	requestMetrics := NewRequestMetrics()
+
+	ctx := context.Background()
+	req := struct{}{}
+
+	testcase := []struct {
+		methodName   string
+		statusCode   codes.Code
+		shouldRecord bool
+	}{
+		{"rpc/test/method", codes.OK, true},
+		{"rpc/test/method", codes.Internal, true},
+		{"rpc/test/method", codes.Aborted, true},
+		{"rpc/test/notRecorded", codes.OK, false},
+	}
+
+	for _, tc := range testcase {
+		t.Run(fmt.Sprintf("%s %s", tc.methodName, tc.statusCode), func(t *testing.T) {
+			info := &grpc.UnaryServerInfo{FullMethod: tc.methodName}
+			handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+				if tc.statusCode == codes.OK {
+					time.Sleep(time.Millisecond)
+					return struct{}{}, nil
+				}
+				return nil, status.Error(tc.statusCode, tc.statusCode.String())
+			}
+			interceptor := NewRequestMetricsInterceptor(requestMetrics, func(fullMethodName string) bool {
+				return tc.shouldRecord
+			})
+			_, err := interceptor(ctx, req, info, handler)
+			if err != nil {
+				require.Equal(t, tc.statusCode, status.Code(err))
+			}
+			var expectedCount uint64
+			if tc.shouldRecord {
+				expectedCount = 1
+			}
+			assertGrpcMetrics(t, requestMetrics.Duration.ToPrometheusMetrics(), map[string]uint64{
+				fmt.Sprintf("%s %s", tc.methodName, tc.statusCode): expectedCount,
+			})
+		})
+	}
+}
+
+func assertGrpcMetrics(
+	t *testing.T, metrics []*io_prometheus_client.Metric, expected map[string]uint64,
+) {
+	t.Helper()
+	actual := map[string]*io_prometheus_client.Histogram{}
+	for _, m := range metrics {
+		var method, statusCode string
+		for _, l := range m.Label {
+			switch *l.Name {
+			case RpcMethodLabel:
+				method = *l.Value
+			case RpcStatusCodeLabel:
+				statusCode = *l.Value
+			}
+		}
+		histogram := m.Histogram
+		require.NotNil(t, histogram, "expected histogram")
+		key := fmt.Sprintf("%s %s", method, statusCode)
+		actual[key] = histogram
+	}
+
+	for key, val := range expected {
+		histogram, ok := actual[key]
+		if val == 0 {
+			require.False(t, ok, "expected `%s` to not exist", key)
+		} else {
+			require.True(t, ok)
+			require.Greater(t, *histogram.SampleSum, float64(0), "expected `%s` to have a SampleSum > 0", key)
+			require.Equal(t, val, *histogram.SampleCount, "expected `%s` to have SampleCount of %d", key, val)
+		}
+	}
+}

pkg/rpc/settings.go

@@ -119,7 +119,8 @@ var sourceAddr = func() net.Addr {
 }()
 
 type serverOpts struct {
-	interceptor func(fullMethod string) error
+	interceptor        func(fullMethod string) error
+	metricsInterceptor RequestMetricsInterceptor
 }
 
 // ServerOption is a configuration option passed to NewServer.
@@ -142,3 +143,10 @@ func WithInterceptor(f func(fullMethod string) error) ServerOption {
 		}
 	}
 }
+
+// WithMetricsServerInterceptor adds a RequestMetricsInterceptor to the grpc server.
+func WithMetricsServerInterceptor(interceptor RequestMetricsInterceptor) ServerOption {
+	return func(opts *serverOpts) {
+		opts.metricsInterceptor = interceptor
+	}
+}

pkg/server/BUILD.bazel

@@ -451,6 +451,7 @@ go_test(
         "get_local_files_test.go",
         "graphite_test.go",
         "grpc_gateway_test.go",
+        "grpc_server_test.go",
         "helpers_test.go",
         "http_metrics_test.go",
         "index_usage_stats_test.go",

pkg/server/grpc_server.go

@@ -7,10 +7,14 @@ package server
 
 import (
 	"context"
+	"strings"
 	"sync/atomic"
 
 	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/server/srverrors"
+	"github.com/cockroachdb/cockroach/pkg/settings"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -27,13 +31,21 @@ type grpcServer struct {
 	mode                   serveMode
 }
 
-func newGRPCServer(ctx context.Context, rpcCtx *rpc.Context) (*grpcServer, error) {
+func newGRPCServer(
+	ctx context.Context, rpcCtx *rpc.Context, metricsRegistry *metric.Registry,
+) (*grpcServer, error) {
 	s := &grpcServer{}
 	s.mode.set(modeInitializing)
+	requestMetrics := rpc.NewRequestMetrics()
+	metricsRegistry.AddMetricStruct(requestMetrics)
 	srv, dsrv, interceptorInfo, err := rpc.NewServerEx(
 		ctx, rpcCtx, rpc.WithInterceptor(func(path string) error {
 			return s.intercept(path)
-		}))
+		}), rpc.WithMetricsServerInterceptor(
+			rpc.NewRequestMetricsInterceptor(requestMetrics, func(method string) bool {
+				return shouldRecordRequestDuration(rpcCtx.Settings, method)
+			},
+			)))
 	if err != nil {
 		return nil, err
 	}
@@ -116,3 +128,25 @@ func NewWaitingForInitError(methodName string) error {
 	// NB: this error string is sadly matched in grpcutil.IsWaitingForInit().
 	return grpcstatus.Errorf(codes.Unavailable, "node waiting for init; %s not available", methodName)
 }
+
+const (
+	serverPrefix = "/cockroach.server"
+	tsdbPrefix   = "/cockroach.ts"
+)
+
+// serverGRPCRequestMetricsEnabled is a cluster setting that enables the
+// collection of gRPC request duration metrics. This uses export only
+// metrics so the metrics are only exported to external sources such as
+// /_status/vars and DataDog.
+var serverGRPCRequestMetricsEnabled = settings.RegisterBoolSetting(
+	settings.ApplicationLevel,
+	"server.grpc.request_metrics.enabled",
+	"enables the collection of grpc metrics",
+	false,
+)
+
+func shouldRecordRequestDuration(settings *cluster.Settings, method string) bool {
+	return serverGRPCRequestMetricsEnabled.Get(&settings.SV) &&
+		(strings.HasPrefix(method, serverPrefix) ||
+			strings.HasPrefix(method, tsdbPrefix))
+}