diff --git a/cmd/connstats.go b/cmd/connstats.go new file mode 100644 index 00000000..9a86719b --- /dev/null +++ b/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() + } +} diff --git a/cmd/connstats_test.go b/cmd/connstats_test.go new file mode 100644 index 00000000..8d285ef2 --- /dev/null +++ b/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") + } +} diff --git a/cmd/root.go b/cmd/root.go index 2fa43109..81abbe81 100644 --- a/cmd/root.go +++ b/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), ), diff --git a/go.mod b/go.mod index 3a9efd46..befb0a2f 100644 --- a/go.mod +++ b/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