Add Prometheus gauge for current gRPC client connections

cmd/connstats.go

@@ -0,0 +1,39 @@
+// Copyright (c) 2019-present The Zcash developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or https://www.opensource.org/licenses/mit-license.php .
+package cmd
+
+import (
+	"context"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"google.golang.org/grpc/stats"
+)
+
+var grpcServerConnectionsCurrent = promauto.NewGauge(prometheus.GaugeOpts{
+	Name: "grpc_server_connections_current",
+	Help: "Number of currently active gRPC client connections.",
+})
+
+// connStatsHandler implements stats.Handler to track gRPC connection lifecycle.
+type connStatsHandler struct{}
+
+func (h *connStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+	return ctx
+}
+
+func (h *connStatsHandler) HandleRPC(ctx context.Context, s stats.RPCStats) {}
+
+func (h *connStatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
+	return ctx
+}
+
+func (h *connStatsHandler) HandleConn(ctx context.Context, s stats.ConnStats) {
+	switch s.(type) {
+	case *stats.ConnBegin:
+		grpcServerConnectionsCurrent.Inc()
+	case *stats.ConnEnd:
+		grpcServerConnectionsCurrent.Dec()
+	}
+}

cmd/connstats_test.go

@@ -0,0 +1,62 @@
+// Copyright (c) 2019-present The Zcash developers
+// Distributed under the MIT software license, see the accompanying
+// file COPYING or https://www.opensource.org/licenses/mit-license.php .
+package cmd
+
+import (
+	"context"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
+	"google.golang.org/grpc/stats"
+)
+
+func gaugeValue(g prometheus.Gauge) float64 {
+	var m dto.Metric
+	g.Write(&m)
+	return m.GetGauge().GetValue()
+}
+
+func TestHandleConnBeginIncrementsGauge(t *testing.T) {
+	grpcServerConnectionsCurrent.Set(0)
+	h := &connStatsHandler{}
+	ctx := context.Background()
+
+	h.HandleConn(ctx, &stats.ConnBegin{})
+	if v := gaugeValue(grpcServerConnectionsCurrent); v != 1 {
+		t.Fatalf("expected gauge 1 after ConnBegin, got %v", v)
+	}
+
+	h.HandleConn(ctx, &stats.ConnEnd{})
+	if v := gaugeValue(grpcServerConnectionsCurrent); v != 0 {
+		t.Fatalf("expected gauge 0 after ConnEnd, got %v", v)
+	}
+}
+
+func TestMultipleConnections(t *testing.T) {
+	grpcServerConnectionsCurrent.Set(0)
+	h := &connStatsHandler{}
+	ctx := context.Background()
+
+	h.HandleConn(ctx, &stats.ConnBegin{})
+	h.HandleConn(ctx, &stats.ConnBegin{})
+	h.HandleConn(ctx, &stats.ConnBegin{})
+	if v := gaugeValue(grpcServerConnectionsCurrent); v != 3 {
+		t.Fatalf("expected gauge 3 after 3 ConnBegin, got %v", v)
+	}
+
+	h.HandleConn(ctx, &stats.ConnEnd{})
+	if v := gaugeValue(grpcServerConnectionsCurrent); v != 2 {
+		t.Fatalf("expected gauge 2 after 1 ConnEnd, got %v", v)
+	}
+}
+
+func TestTagRPCPassthrough(t *testing.T) {
+	h := &connStatsHandler{}
+	ctx := context.Background()
+	got := h.TagRPC(ctx, &stats.RPCTagInfo{})
+	if got != ctx {
+		t.Fatal("TagRPC should return the same context")
+	}
+}

cmd/root.go

@@ -140,6 +140,7 @@ func startServer(opts *common.Options) error {
 		common.Log.Warningln("Starting insecure no-TLS (plaintext) server")
 		fmt.Println("Starting insecure server")
 		server = grpc.NewServer(
+			grpc.StatsHandler(&connStatsHandler{}),
 			grpc.StreamInterceptor(
 				grpc_middleware.ChainStreamServer(
 					grpc_prometheus.StreamServerInterceptor),
@@ -168,6 +169,7 @@ func startServer(opts *common.Options) error {
 		}
 		server = grpc.NewServer(
 			grpc.Creds(transportCreds),
+			grpc.StatsHandler(&connStatsHandler{}),
 			grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
 				grpc_prometheus.StreamServerInterceptor),
 			),

go.mod

@@ -11,9 +11,11 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 	github.com/prometheus/client_golang v1.18.0
+	github.com/prometheus/client_model v0.5.0
 	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/cobra v1.8.0
 	github.com/spf13/viper v1.18.2
+	golang.org/x/exp v0.0.0-20230905200255-921286631fa9
 	google.golang.org/grpc v1.61.0
 	google.golang.org/protobuf v1.33.0
 	gopkg.in/ini.v1 v1.67.0
@@ -37,7 +39,6 @@ require (
 	github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/pelletier/go-toml/v2 v2.1.0 // indirect
-	github.com/prometheus/client_model v0.5.0 // indirect
 	github.com/prometheus/common v0.45.0 // indirect
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/sagikazarmark/locafero v0.4.0 // indirect
@@ -50,7 +51,6 @@ require (
 	go.uber.org/atomic v1.9.0 // indirect
 	go.uber.org/multierr v1.9.0 // indirect
 	golang.org/x/crypto v0.45.0 // indirect
-	golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect
 	golang.org/x/net v0.47.0 // indirect
 	golang.org/x/sys v0.38.0 // indirect
 	golang.org/x/text v0.31.0 // indirect