diff --git a/server.go b/server.go index 322f990..7de56f8 100644 --- a/server.go +++ b/server.go @@ -30,11 +30,16 @@ func init() { prom.MustRegister(DefaultServerMetrics.serverStreamMsgSent) } +// GRPCServer lets register any grpc server implementation, e.g. *grpc.Server and *xds.GRPCServer. +type GRPCServer interface { + GetServiceInfo() map[string]grpc.ServiceInfo +} + // Register takes a gRPC server and pre-initializes all counters to 0. This // allows for easier monitoring in Prometheus (no missing metrics), and should // be called *after* all services have been registered with the server. This // function acts on the DefaultServerMetrics variable. -func Register(server *grpc.Server) { +func Register(server GRPCServer) { DefaultServerMetrics.InitializeMetrics(server) } diff --git a/server_metrics.go b/server_metrics.go index d28a46e..b66e576 100644 --- a/server_metrics.go +++ b/server_metrics.go @@ -2,6 +2,7 @@ package grpc_prometheus import ( "context" + "github.com/grpc-ecosystem/go-grpc-prometheus/packages/grpcstatus" prom "github.com/prometheus/client_golang/prometheus" @@ -129,7 +130,7 @@ func (m *ServerMetrics) StreamServerInterceptor() func(srv interface{}, ss grpc. // InitializeMetrics initializes all metrics, with their appropriate null // value, for all gRPC methods registered on a gRPC server. This is useful, to // ensure that all metrics exist when collecting and querying. -func (m *ServerMetrics) InitializeMetrics(server *grpc.Server) { +func (m *ServerMetrics) InitializeMetrics(server GRPCServer) { serviceInfo := server.GetServiceInfo() for serviceName, info := range serviceInfo { for _, mInfo := range info.Methods {