ClientRequestsInFlightStatsHandler counts correctly in flight request…

…s by marking context on the OutHeader event and decrementing metric on the End event only if the context was previously marked. (#39)
piotrkowalczuk · Oct 14, 2023 · 659087a · 659087a
1 parent aaa9186
commit 659087a
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/v4/metric_client_connections.go b/v4/metric_client_connections.go
@@ -17,7 +17,7 @@ type ClientConnectionsStatsHandler struct {
 	vec *prometheus.GaugeVec
 }
 
-// NewConnectionsStatsHandler ...
+// NewClientConnectionsStatsHandler ...
 func NewClientConnectionsStatsHandler(vec *prometheus.GaugeVec) *ClientConnectionsStatsHandler {
 	return &ClientConnectionsStatsHandler{
 		baseStatsHandler: baseStatsHandler{
@@ -30,7 +30,7 @@ func NewClientConnectionsStatsHandler(vec *prometheus.GaugeVec) *ClientConnectio
 	}
 }
 
-// HandleRPC processes the RPC stats.
+// HandleConn HandleRPC processes the RPC stats.
 func (h *ClientConnectionsStatsHandler) HandleConn(ctx context.Context, stat stats.ConnStats) {
 	switch stat.(type) {
 	case *stats.ConnBegin:

diff --git a/v4/metric_client_requests_in_flight.go b/v4/metric_client_requests_in_flight.go
@@ -42,18 +42,30 @@ func NewClientRequestsInFlightStatsHandler(vec *prometheus.GaugeVec, opts ...Sta
 	return h
 }
 
+func (h *ClientRequestsInFlightStatsHandler) TagRPC(ctx context.Context, inf *stats.RPCTagInfo) context.Context {
+	ctx = h.baseStatsHandler.TagRPC(ctx, inf)
+	// LINK: https://github.com/grpc/grpc-go/issues/5823
+	ctx = context.WithValue(ctx, clientRequestInFlightKey{}, &clientRequestInFlightMark{})
+	return ctx
+}
+
 // HandleRPC processes the RPC stats.
 func (h *ClientRequestsInFlightStatsHandler) HandleRPC(ctx context.Context, stat stats.RPCStats) {
 	switch stat.(type) {
 	case *stats.OutHeader:
 		switch {
 		case stat.IsClient():
-			h.vec.WithLabelValues(h.options.handleRPCLabelFn(ctx, stat)...).Inc()
+			if mrk, ok := ctx.Value(clientRequestInFlightKey{}).(*clientRequestInFlightMark); ok {
+				mrk.started = true
+				h.vec.WithLabelValues(h.options.handleRPCLabelFn(ctx, stat)...).Inc()
+			}
 		}
 	case *stats.End:
 		switch {
 		case stat.IsClient():
-			h.vec.WithLabelValues(h.options.handleRPCLabelFn(ctx, stat)...).Dec()
+			if mrk, ok := ctx.Value(clientRequestInFlightKey{}).(*clientRequestInFlightMark); ok && mrk.started {
+				h.vec.WithLabelValues(h.options.handleRPCLabelFn(ctx, stat)...).Dec()
+			}
 		}
 	}
 }
@@ -68,3 +80,9 @@ func (h *ClientRequestsInFlightStatsHandler) labels(ctx context.Context, stat st
 		h.uas.ClientSide(ctx, stat),
 	}
 }
+
+type clientRequestInFlightKey struct{}
+
+type clientRequestInFlightMark struct {
+	started bool
+}