Add GRPCServer interface to accept different grpc.Server implementations

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)
 }
 

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 {