Serverside metrics update attempt 2

internal/prometheusbpint/spectest/spec.go

@@ -163,6 +163,10 @@ func buildLabels(name, prefix, clientOrServer string) map[string]struct{} {
 }
 
 // thriftSpecificLabels returns the following labels:
+//
+// All serverside v0 metrics expect the following labels:
+//   - "server_name" (v0 server only)
+//
 // latency_seconds metrics expect the following labels:
 //   - "method"
 //   - "success"
@@ -177,17 +181,23 @@ func buildLabels(name, prefix, clientOrServer string) map[string]struct{} {
 // active_requests metrics expect the following labels:
 //   - "method"
 func thriftSpecificLabels(name string) []string {
+
 	labelSuffixes := []string{"method"}
+	if strings.HasSuffix(name, "v0") {
+		labelSuffixes = append(labelSuffixes, "server_name")
+	}
+
 	switch {
-	case strings.HasSuffix(name, "_latency_seconds"):
+	case strings.Contains(name, "_latency_seconds"):
 		labelSuffixes = append(labelSuffixes, "success")
-	case strings.HasSuffix(name, "_requests_total"):
+	case strings.Contains(name, "_requests_total"):
 		labelSuffixes = append(labelSuffixes, "success", "exception_type", "baseplate_status", "baseplate_status_code")
-	case strings.HasSuffix(name, "_active_requests"):
+	case strings.Contains(name, "_active_requests"):
 		// no op
 	default:
 		return nil
 	}
+
 	return labelSuffixes
 }
 

thriftbp/prometheus.go

@@ -22,6 +22,7 @@ const (
 	baseplateStatusLabel     = "thrift_baseplate_status"
 	baseplateStatusCodeLabel = "thrift_baseplate_status_code"
 	clientNameLabel          = "thrift_client_name"
+	serverNameLabel          = "thrift_server_name"
 )
 
 var (
@@ -36,6 +37,18 @@ var (
 		Buckets: prometheusbp.DefaultLatencyBuckets,
 	}, serverLatencyLabels)
 
+	serverLatencyLabelsV0 = []string{
+		serverNameLabel,
+		methodLabel,
+		successLabel,
+	}
+
+	serverLatencyDistributionV0 = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "thrift_server_latency_seconds_v0",
+		Help:    "RPC latencies",
+		Buckets: prometheusbp.DefaultLatencyBuckets,
+	}, serverLatencyLabelsV0)
+
 	serverTotalRequestLabels = []string{
 		methodLabel,
 		successLabel,
@@ -49,6 +62,20 @@ var (
 		Help: "Total RPC request count",
 	}, serverTotalRequestLabels)
 
+	serverTotalRequestLabelsV0 = []string{
+		serverNameLabel,
+		methodLabel,
+		successLabel,
+		exceptionLabel,
+		baseplateStatusLabel,
+		baseplateStatusCodeLabel,
+	}
+
+	serverTotalRequestsV0 = promauto.With(prometheusbpint.GlobalRegistry).NewCounterVec(prometheus.CounterOpts{
+		Name: "thrift_server_requests_total_v0",
+		Help: "Total RPC request count",
+	}, serverTotalRequestLabelsV0)
+
 	serverActiveRequestsLabels = []string{
 		methodLabel,
 	}
@@ -57,6 +84,16 @@ var (
 		Name: "thrift_server_active_requests",
 		Help: "The number of in-flight requests being handled by the service",
 	}, serverActiveRequestsLabels)
+
+	serverActiveRequestsLabelsV0 = []string{
+		serverNameLabel,
+		methodLabel,
+	}
+
+	serverActiveRequestsV0 = promauto.With(prometheusbpint.GlobalRegistry).NewGaugeVec(prometheus.GaugeOpts{
+		Name: "thrift_server_active_requests_v0",
+		Help: "The number of in-flight requests being handled by the service",
+	}, serverActiveRequestsLabelsV0)
 )
 
 var (
@@ -140,6 +177,12 @@ var (
 		protoLabel,
 	}
 
+	serverPayloadSizeLabelsV0 = []string{
+		serverNameLabel,
+		methodLabel,
+		protoLabel,
+	}
+
 	clientPayloadSizeLabels = []string{
 		methodLabel,
 		clientNameLabel,
@@ -158,12 +201,24 @@ var (
 		Buckets: payloadSizeBuckets,
 	}, serverPayloadSizeLabels)
 
+	serverPayloadSizeRequestBytesV0 = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "thriftbp_server_request_payload_size_bytes_v0",
+		Help:    "The (approximate) size of thrift request payloads",
+		Buckets: payloadSizeBuckets,
+	}, serverPayloadSizeLabelsV0)
+
 	serverPayloadSizeResponseBytes = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
 		Name:    "thriftbp_server_response_payload_size_bytes",
 		Help:    "The (approximate) size of thrift response payloads",
 		Buckets: payloadSizeBuckets,
 	}, serverPayloadSizeLabels)
 
+	serverPayloadSizeResponseBytesV0 = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "thriftbp_server_response_payload_size_bytes_v0",
+		Help:    "The (approximate) size of thrift response payloads",
+		Buckets: payloadSizeBuckets,
+	}, serverPayloadSizeLabelsV0)
+
 	clientPayloadSizeRequestBytes = promauto.With(prometheusbpint.GlobalRegistry).NewHistogramVec(prometheus.HistogramOpts{
 		Name:    "thriftbp_client_request_payload_size_bytes",
 		Help:    "The size of thrift request payloads",

thriftbp/prometheus_test.go

@@ -53,6 +53,7 @@ func TestPrometheusServerMiddleware(t *testing.T) {
 	}
 
 	const method = "testmethod"
+	const serverName = "testserver"
 
 	for _, tt := range testCases {
 		t.Run(tt.name, func(t *testing.T) {
@@ -95,7 +96,72 @@ func TestPrometheusServerMiddleware(t *testing.T) {
 					return tt.wantOK, tt.wantErr
 				},
 			}
-			gotOK, gotErr := PrometheusServerMiddleware(method, next).Process(context.Background(), 1, nil, nil)
+			gotOK, gotErr := PrometheusServerMiddlewareWithArgs(PrometheusServerMiddlewareArgs{ServerName: serverName})(method, next).Process(context.Background(), 1, nil, nil)
+
+			if gotOK != tt.wantOK {
+				t.Errorf("wanted %v, got %v", tt.wantOK, gotOK)
+			}
+			if gotErr != tt.wantErr {
+				t.Errorf("wanted %v, got %v", tt.wantErr, gotErr)
+			}
+		})
+	}
+}
+
+func TestReportPayloadSizeMetrics(t *testing.T) {
+	testCases := []struct {
+		name    string
+		wantErr thrift.TException
+		wantOK  bool
+	}{
+		{
+			name:    "success",
+			wantErr: nil,
+			wantOK:  true,
+		},
+		{
+			name:    "error",
+			wantErr: thrift.NewTApplicationException(thrift.UNKNOWN_METHOD, "unknown err msg"),
+			wantOK:  false,
+		},
+	}
+
+	const method = "testmethod"
+	const serverName = "testserver"
+	const proto = "header-binary"
+
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			serverPayloadSizeRequestBytes.Reset()
+			serverPayloadSizeResponseBytes.Reset()
+
+			payloadLabels := prometheus.Labels{
+				methodLabel: method,
+				protoLabel:  proto,
+			}
+
+			defer promtest.NewPrometheusMetricTest(t, "request payload", serverPayloadSizeRequestBytes, payloadLabels).CheckSampleCountDelta(1)
+			defer promtest.NewPrometheusMetricTest(t, "response payload", serverPayloadSizeResponseBytes, payloadLabels).CheckSampleCountDelta(1)
+
+			next := thrift.WrappedTProcessorFunction{
+				Wrapped: func(ctx context.Context, seqId int32, in, out thrift.TProtocol) (bool, thrift.TException) {
+					return tt.wantOK, tt.wantErr
+				},
+			}
+
+			// Create THeader transport and protocol to trigger payload size tracking
+			cfg := &thrift.TConfiguration{}
+			transport := thrift.NewTMemoryBuffer()
+			headerTransport := thrift.NewTHeaderTransportConf(transport, cfg)
+			headerProtocol := thrift.NewTHeaderProtocolConf(headerTransport, cfg)
+
+			middleware := ReportPayloadSizeMetricsWithArgs(ReportPayloadSizeMetricsArgs{
+				ServerName: serverName,
+				SampleRate: 0,
+			})
+
+			wrapped := middleware(method, next)
+			gotOK, gotErr := wrapped.Process(context.Background(), 1, headerProtocol, headerProtocol)
 
 			if gotOK != tt.wantOK {
 				t.Errorf("wanted %v, got %v", tt.wantOK, gotOK)