Refactor metrics interceptor and fix tests (#413)

Author: ash2k

interceptors/client.go

@@ -7,6 +7,7 @@ package interceptors
 
 import (
 	"context"
+	"io"
 	"time"
 
 	"google.golang.org/grpc"
@@ -74,7 +75,10 @@ func (s *monitoredClientStream) RecvMsg(m interface{}) error {
 	if err == nil {
 		return nil
 	}
-
-	s.reporter.PostCall(err, time.Since(s.startTime))
+	var postErr error
+	if err != io.EOF {
+		postErr = err
+	}
+	s.reporter.PostCall(postErr, time.Since(s.startTime))
 	return err
 }

interceptors/client_test.go

@@ -50,7 +50,7 @@ func (m *mockReportable) Equal(t *testing.T, expected []*mockReport) {
 				require.NoError(t, err)
 				continue
 			}
-			require.Equal(t, expected[i].postCalls[k].Error(), err.Error(), "%v %v", i, k)
+			require.EqualError(t, err, expected[i].postCalls[k].Error(), "%v %v", i, k)
 		}
 		require.Len(t, expected[i].postMsgSends, len(e.postMsgSends), "%v", i)
 		for k, err := range e.postMsgSends {
@@ -276,7 +276,7 @@ func (s *ClientInterceptorTestSuite) TestListReporting() {
 		typ:             ServerStream,
 		svcName:         testpb.TestServiceFullName,
 		methodName:      "PingList",
-		postCalls:       []error{io.EOF},
+		postCalls:       []error{nil},
 		postMsgReceives: append(make([]error, testpb.ListResponseCount), io.EOF),
 		postMsgSends:    []error{nil},
 	}})
@@ -329,7 +329,9 @@ func (s *ClientInterceptorTestSuite) TestBiStreamingReporting() {
 			if err == io.EOF {
 				break
 			}
-			require.NoError(s.T(), err, "reading pingStream shouldn't fail")
+			if !s.Assert().NoError(err, "reading pingStream shouldn't fail") {
+				break
+			}
 			count++
 		}
 	}()
@@ -340,12 +342,12 @@ func (s *ClientInterceptorTestSuite) TestBiStreamingReporting() {
 	require.NoError(s.T(), ss.CloseSend())
 	wg.Wait()
 
-	require.EqualValues(s.T(), count, 100, "Number of received msg on the wire must match")
+	require.EqualValues(s.T(), 100, count, "Number of received msg on the wire must match")
 	s.mock.Equal(s.T(), []*mockReport{{
 		typ:             BidiStream,
 		svcName:         testpb.TestServiceFullName,
 		methodName:      "PingStream",
-		postCalls:       []error{io.EOF},
+		postCalls:       []error{nil},
 		postMsgReceives: append(make([]error, 100), io.EOF),
 		postMsgSends:    make([]error, 100),
 	}})

providers/openmetrics/client_interceptor.go

@@ -1,35 +1,21 @@
 package metrics
 
 import (
-	openmetrics "github.com/prometheus/client_golang/prometheus"
 	"google.golang.org/grpc"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors"
 )
 
-// RegisterClientMetrics returns a custom ClientMetrics object registered
-// with the user's registry, and registers some common metrics associated
-// with every instance.
-func RegisterClientMetrics(registry openmetrics.Registerer) *ClientMetrics {
-	customClientMetrics := NewClientMetrics(registry)
-	customClientMetrics.MustRegister(customClientMetrics.clientStartedCounter)
-	customClientMetrics.MustRegister(customClientMetrics.clientHandledCounter)
-	customClientMetrics.MustRegister(customClientMetrics.clientStreamMsgReceived)
-	customClientMetrics.MustRegister(customClientMetrics.clientStreamMsgSent)
-
-	return customClientMetrics
-}
-
 // UnaryClientInterceptor is a gRPC client-side interceptor that provides Prometheus monitoring for Unary RPCs.
-func UnaryClientInterceptor(clientRegister openmetrics.Registerer) grpc.UnaryClientInterceptor {
+func UnaryClientInterceptor(clientMetrics *ClientMetrics) grpc.UnaryClientInterceptor {
 	return interceptors.UnaryClientInterceptor(&reportable{
-		registry: clientRegister,
+		clientMetrics: clientMetrics,
 	})
 }
 
 // StreamClientInterceptor is a gRPC client-side interceptor that provides Prometheus monitoring for Streaming RPCs.
-func StreamClientInterceptor(clientRegister openmetrics.Registerer) grpc.StreamClientInterceptor {
+func StreamClientInterceptor(clientMetrics *ClientMetrics) grpc.StreamClientInterceptor {
 	return interceptors.StreamClientInterceptor(&reportable{
-		registry: clientRegister,
+		clientMetrics: clientMetrics,
 	})
 }

providers/openmetrics/client_metrics.go

@@ -1,99 +1,103 @@
 package metrics
 
 import (
-	"github.com/prometheus/client_golang/prometheus"
 	openmetrics "github.com/prometheus/client_golang/prometheus"
 )
 
 // ClientMetrics represents a collection of metrics to be registered on a
 // Prometheus metrics registry for a gRPC client.
 type ClientMetrics struct {
-	clientRegister openmetrics.Registerer
-
 	clientStartedCounter    *openmetrics.CounterVec
 	clientHandledCounter    *openmetrics.CounterVec
 	clientStreamMsgReceived *openmetrics.CounterVec
 	clientStreamMsgSent     *openmetrics.CounterVec
 
-	clientHandledHistogramEnabled bool
-	clientHandledHistogramOpts    openmetrics.HistogramOpts
-	clientHandledHistogram        *openmetrics.HistogramVec
-
-	clientStreamRecvHistogramEnabled bool
-	clientStreamRecvHistogramOpts    openmetrics.HistogramOpts
-	clientStreamRecvHistogram        *openmetrics.HistogramVec
-
-	clientStreamSendHistogramEnabled bool
-	clientStreamSendHistogramOpts    openmetrics.HistogramOpts
-	clientStreamSendHistogram        *openmetrics.HistogramVec
+	// clientHandledHistogram can be nil
+	clientHandledHistogram *openmetrics.HistogramVec
+	// clientStreamRecvHistogram can be nil
+	clientStreamRecvHistogram *openmetrics.HistogramVec
+	// clientStreamSendHistogram can be nil
+	clientStreamSendHistogram *openmetrics.HistogramVec
 }
 
-// NewClientMetrics returns a ClientMetrics object. Use a new instance of
-// ClientMetrics when not using the default Prometheus metrics registry, for
-// example when wanting to control which metrics are added to a registry as
-// opposed to automatically adding metrics via init functions.
-func NewClientMetrics(clientRegistry prometheus.Registerer, counterOpts ...CounterOption) *ClientMetrics {
-	opts := counterOptions(counterOpts)
+// NewClientMetrics returns a new ClientMetrics object.
+func NewClientMetrics(opts ...ClientMetricsOption) *ClientMetrics {
+	var config clientMetricsConfig
+	config.apply(opts)
 	return &ClientMetrics{
-		clientRegister: clientRegistry,
 		clientStartedCounter: openmetrics.NewCounterVec(
-			opts.apply(openmetrics.CounterOpts{
+			config.counterOpts.apply(openmetrics.CounterOpts{
 				Name: "grpc_client_started_total",
 				Help: "Total number of RPCs started on the client.",
 			}), []string{"grpc_type", "grpc_service", "grpc_method"}),
 
 		clientHandledCounter: openmetrics.NewCounterVec(
-			opts.apply(openmetrics.CounterOpts{
+			config.counterOpts.apply(openmetrics.CounterOpts{
 				Name: "grpc_client_handled_total",
 				Help: "Total number of RPCs completed by the client, regardless of success or failure.",
 			}), []string{"grpc_type", "grpc_service", "grpc_method", "grpc_code"}),
 
 		clientStreamMsgReceived: openmetrics.NewCounterVec(
-			opts.apply(openmetrics.CounterOpts{
+			config.counterOpts.apply(openmetrics.CounterOpts{
 				Name: "grpc_client_msg_received_total",
 				Help: "Total number of RPC stream messages received by the client.",
 			}), []string{"grpc_type", "grpc_service", "grpc_method"}),
 
 		clientStreamMsgSent: openmetrics.NewCounterVec(
-			opts.apply(openmetrics.CounterOpts{
+			config.counterOpts.apply(openmetrics.CounterOpts{
 				Name: "grpc_client_msg_sent_total",
 				Help: "Total number of gRPC stream messages sent by the client.",
 			}), []string{"grpc_type", "grpc_service", "grpc_method"}),
 
-		clientHandledHistogramEnabled: false,
-		clientHandledHistogramOpts: openmetrics.HistogramOpts{
-			Name:    "grpc_client_handling_seconds",
-			Help:    "Histogram of response latency (seconds) of the gRPC until it is finished by the application.",
-			Buckets: openmetrics.DefBuckets,
-		},
-		clientHandledHistogram:           nil,
-		clientStreamRecvHistogramEnabled: false,
-		clientStreamRecvHistogramOpts: openmetrics.HistogramOpts{
-			Name:    "grpc_client_msg_recv_handling_seconds",
-			Help:    "Histogram of response latency (seconds) of the gRPC single message receive.",
-			Buckets: openmetrics.DefBuckets,
-		},
-		clientStreamRecvHistogram:        nil,
-		clientStreamSendHistogramEnabled: false,
-		clientStreamSendHistogramOpts: openmetrics.HistogramOpts{
-			Name:    "grpc_client_msg_send_handling_seconds",
-			Help:    "Histogram of response latency (seconds) of the gRPC single message send.",
-			Buckets: openmetrics.DefBuckets,
-		},
-		clientStreamSendHistogram: nil,
+		clientHandledHistogram:    config.clientHandledHistogram,
+		clientStreamRecvHistogram: config.clientStreamRecvHistogram,
+		clientStreamSendHistogram: config.clientStreamSendHistogram,
 	}
 }
 
-// Register registers the provided Collector with the custom register.
+// NewRegisteredClientMetrics returns a custom ClientMetrics object registered
+// with the user's registry, and registers some common metrics associated
+// with every instance.
+func NewRegisteredClientMetrics(registry openmetrics.Registerer, opts ...ClientMetricsOption) *ClientMetrics {
+	customClientMetrics := NewClientMetrics(opts...)
+	customClientMetrics.MustRegister(registry)
+	return customClientMetrics
+}
+
+// Register registers the metrics with the registry.
 // returns error much like DefaultRegisterer of Prometheus.
-func (m *ClientMetrics) Register(c openmetrics.Collector) error {
-	return m.clientRegister.Register(c)
+func (m *ClientMetrics) Register(registry openmetrics.Registerer) error {
+	for _, collector := range m.toRegister() {
+		if err := registry.Register(collector); err != nil {
+			return err
+		}
+	}
+	return nil
 }
 
-// MustRegister registers the provided Collectors with the custom Registerer
+// MustRegister registers the metrics with the registry
 // and panics if any error occurs much like DefaultRegisterer of Prometheus.
-func (m *ClientMetrics) MustRegister(c openmetrics.Collector) {
-	m.clientRegister.MustRegister(c)
+func (m *ClientMetrics) MustRegister(registry openmetrics.Registerer) {
+	registry.MustRegister(m.toRegister()...)
+}
+
+func (m *ClientMetrics) toRegister() []openmetrics.Collector {
+	res := []openmetrics.Collector{
+		m.clientStartedCounter,
+		m.clientHandledCounter,
+		m.clientStreamMsgReceived,
+		m.clientStreamMsgSent,
+	}
+	if m.clientHandledHistogram != nil {
+		res = append(res, m.clientHandledHistogram)
+	}
+	if m.clientStreamRecvHistogram != nil {
+		res = append(res, m.clientStreamRecvHistogram)
+	}
+	if m.clientStreamSendHistogram != nil {
+		res = append(res, m.clientStreamSendHistogram)
+	}
+	return res
 }
 
 // Describe sends the super-set of all possible descriptors of metrics
@@ -104,13 +108,13 @@ func (m *ClientMetrics) Describe(ch chan<- *openmetrics.Desc) {
 	m.clientHandledCounter.Describe(ch)
 	m.clientStreamMsgReceived.Describe(ch)
 	m.clientStreamMsgSent.Describe(ch)
-	if m.clientHandledHistogramEnabled {
+	if m.clientHandledHistogram != nil {
 		m.clientHandledHistogram.Describe(ch)
 	}
-	if m.clientStreamRecvHistogramEnabled {
+	if m.clientStreamRecvHistogram != nil {
 		m.clientStreamRecvHistogram.Describe(ch)
 	}
-	if m.clientStreamSendHistogramEnabled {
+	if m.clientStreamSendHistogram != nil {
 		m.clientStreamSendHistogram.Describe(ch)
 	}
 }
@@ -123,65 +127,13 @@ func (m *ClientMetrics) Collect(ch chan<- openmetrics.Metric) {
 	m.clientHandledCounter.Collect(ch)
 	m.clientStreamMsgReceived.Collect(ch)
 	m.clientStreamMsgSent.Collect(ch)
-	if m.clientHandledHistogramEnabled {
+	if m.clientHandledHistogram != nil {
 		m.clientHandledHistogram.Collect(ch)
 	}
-	if m.clientStreamRecvHistogramEnabled {
+	if m.clientStreamRecvHistogram != nil {
 		m.clientStreamRecvHistogram.Collect(ch)
 	}
-	if m.clientStreamSendHistogramEnabled {
+	if m.clientStreamSendHistogram != nil {
 		m.clientStreamSendHistogram.Collect(ch)
 	}
 }
-
-// EnableClientHandlingTimeHistogram turns on recording of handling time of RPCs.
-// Histogram metrics can be very expensive for Prometheus to retain and query.
-func (m *ClientMetrics) EnableClientHandlingTimeHistogram(opts ...HistogramOption) error {
-	for _, o := range opts {
-		o(&m.clientHandledHistogramOpts)
-	}
-	if !m.clientHandledHistogramEnabled {
-		m.clientHandledHistogram = openmetrics.NewHistogramVec(
-			m.clientHandledHistogramOpts,
-			[]string{"grpc_type", "grpc_service", "grpc_method"},
-		)
-	}
-	m.clientHandledHistogramEnabled = true
-	return m.clientRegister.Register(m.clientHandledHistogram)
-}
-
-// EnableClientStreamReceiveTimeHistogram turns on recording of single message receive time of streaming RPCs.
-// Histogram metrics can be very expensive for Prometheus to retain and query.
-func (m *ClientMetrics) EnableClientStreamReceiveTimeHistogram(opts ...HistogramOption) error {
-	for _, o := range opts {
-		o(&m.clientStreamRecvHistogramOpts)
-	}
-
-	if !m.clientStreamRecvHistogramEnabled {
-		m.clientStreamRecvHistogram = openmetrics.NewHistogramVec(
-			m.clientStreamRecvHistogramOpts,
-			[]string{"grpc_type", "grpc_service", "grpc_method"},
-		)
-	}
-
-	m.clientStreamRecvHistogramEnabled = true
-	return m.clientRegister.Register(m.clientStreamRecvHistogram)
-}
-
-// EnableClientStreamSendTimeHistogram turns on recording of single message send time of streaming RPCs.
-// Histogram metrics can be very expensive for Prometheus to retain and query.
-func (m *ClientMetrics) EnableClientStreamSendTimeHistogram(opts ...HistogramOption) error {
-	for _, o := range opts {
-		o(&m.clientStreamSendHistogramOpts)
-	}
-
-	if !m.clientStreamSendHistogramEnabled {
-		m.clientStreamSendHistogram = openmetrics.NewHistogramVec(
-			m.clientStreamSendHistogramOpts,
-			[]string{"grpc_type", "grpc_service", "grpc_method"},
-		)
-	}
-
-	m.clientStreamSendHistogramEnabled = true
-	return m.clientRegister.Register(m.clientStreamSendHistogram)
-}