Fix grpc prometheus metrics (#4172)

* fix grpc prom metrics * fix test --------- Co-authored-by: JamesMurkin <[email protected]>
armadaproject · Jan 29, 2025 · a5afc8a · a5afc8a
1 parent d43a9d9
commit a5afc8a
Show file tree

Hide file tree

Showing 9 changed files with 36 additions and 30 deletions.
diff --git a/go.mod b/go.mod
@@ -21,7 +21,6 @@ require (
 	github.com/golang/protobuf v1.5.4
 	github.com/google/go-cmp v0.6.0 // indirect
 	github.com/google/uuid v1.6.0
-	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
 	github.com/grpc-ecosystem/grpc-gateway v1.16.0
 	github.com/hashicorp/go-memdb v1.3.4
 	github.com/hashicorp/go-multierror v1.1.1

diff --git a/go.sum b/go.sum
@@ -300,8 +300,6 @@ github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.1 h1:qnpS
 github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.1/go.mod h1:lXGCsh6c22WGtjr+qGHj1otzZpV/1kwTMAqkwZsnWRU=
 github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.2.0 h1:kQ0NI7W1B3HwiN5gAYtY+XFItDPbLBwYRxAqbFTyDes=
 github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.2.0/go.mod h1:zrT2dxOAjNFPRGjTUe2Xmb4q4YdUwVvQFV6xiCSf+z0=
-github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 h1:Ovs26xHkKqVztRpIrF/92BcuyuQ/YW4NSIpoGtfXNho=
-github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0 h1:gmcG1KaJ57LophUzW0Hy8NmPhnMZb4M0+kPpLofRdBo=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
 github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c h1:6rhixN/i8ZofjG1Y75iExal34USq5p+wiN1tpie8IrU=

diff --git a/internal/binoculars/server.go b/internal/binoculars/server.go
@@ -4,8 +4,6 @@ import (
 	"os"
 	"sync"
 
-	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
-
 	"github.com/armadaproject/armada/internal/binoculars/configuration"
 	"github.com/armadaproject/armada/internal/binoculars/server"
 	"github.com/armadaproject/armada/internal/binoculars/service"
@@ -48,7 +46,6 @@ func StartUp(config *configuration.BinocularsConfig) (func(), *sync.WaitGroup) {
 	cordonService := service.NewKubernetesCordonService(config.Cordon, permissionsChecker, kubernetesClientProvider)
 	binocularsServer := server.NewBinocularsServer(logService, cordonService)
 	binoculars.RegisterBinocularsServer(grpcServer, binocularsServer)
-	grpc_prometheus.Register(grpcServer)
 
 	grpcCommon.Listen(config.GrpcPort, grpcServer, wg)
 

diff --git a/internal/common/grpc/grpc.go b/internal/common/grpc/grpc.go
@@ -13,8 +13,8 @@ import (
 	grpc_auth "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/auth"
 	grpc_logging "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
 	grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/recovery"
-
 	"github.com/prometheus/client_golang/prometheus"
+
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
@@ -76,8 +76,7 @@ func setupPromMetrics() *grpc_prometheus.ServerMetrics {
 			grpc_prometheus.WithHistogramBuckets([]float64{0.001, 0.01, 0.1, 0.3, 0.6, 1, 3, 6, 9, 20, 30, 60, 90, 120}),
 		),
 	)
-	reg := prometheus.NewRegistry()
-	reg.MustRegister(srvMetrics)
+	prometheus.MustRegister(srvMetrics)
 	return srvMetrics
 }
 

diff --git a/internal/executor/application.go b/internal/executor/application.go
@@ -9,7 +9,7 @@ import (
 	"time"
 
 	"github.com/go-playground/validator/v10"
-	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
 	"github.com/prometheus/client_golang/prometheus"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
@@ -284,12 +284,15 @@ func setupExecutorApiComponents(
 }
 
 func createConnectionToApi(connectionDetails client.ApiConnectionDetails, maxMessageSizeBytes int, grpcConfig keepalive.ClientParameters) (*grpc.ClientConn, error) {
-	grpc_prometheus.EnableClientHandlingTimeHistogram()
+	clientMetrics := grpc_prometheus.NewClientMetrics(
+		grpc_prometheus.WithClientHandlingTimeHistogram(),
+	)
+	prometheus.MustRegister(clientMetrics)
 	return client.CreateApiConnectionWithCallOptions(
 		&connectionDetails,
 		[]grpc.CallOption{grpc.MaxCallRecvMsgSize(maxMessageSizeBytes)},
-		grpc.WithChainUnaryInterceptor(grpc_prometheus.UnaryClientInterceptor),
-		grpc.WithChainStreamInterceptor(grpc_prometheus.StreamClientInterceptor),
+		grpc.WithChainUnaryInterceptor(clientMetrics.UnaryClientInterceptor()),
+		grpc.WithChainStreamInterceptor(clientMetrics.StreamClientInterceptor()),
 		grpc.WithKeepaliveParams(grpcConfig),
 	)
 }

diff --git a/internal/scheduler/leader/leader_client.go b/internal/scheduler/leader/leader_client.go
@@ -5,7 +5,7 @@ import (
 	"strings"
 	"sync"
 
-	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
 	"github.com/pkg/errors"
 	"google.golang.org/grpc"
 
@@ -24,14 +24,16 @@ type LeaderConnectionProvider struct {
 	leaderConfig     configuration.LeaderConfig
 	connectionLock   sync.Mutex
 	connectionByName map[string]*grpc.ClientConn
+	metrics          *grpc_prometheus.ClientMetrics
 }
 
-func NewLeaderConnectionProvider(leaderController LeaderController, leaderConfig configuration.LeaderConfig) *LeaderConnectionProvider {
+func NewLeaderConnectionProvider(leaderController LeaderController, leaderConfig configuration.LeaderConfig, metrics *grpc_prometheus.ClientMetrics) *LeaderConnectionProvider {
 	return &LeaderConnectionProvider{
 		leaderController: leaderController,
 		leaderConfig:     leaderConfig,
 		connectionLock:   sync.Mutex{},
 		connectionByName: map[string]*grpc.ClientConn{},
+		metrics:          metrics,
 	}
 }
 
@@ -60,7 +62,7 @@ func (l *LeaderConnectionProvider) getClientByName(currentLeaderName string) (*g
 	leaderConnectionDetails := l.leaderConfig.LeaderConnection
 	leaderConnectionDetails.ArmadaUrl = strings.ReplaceAll(leaderConnectionDetails.ArmadaUrl, leaseHolderNameToken, currentLeaderName)
 
-	apiConnection, err := createApiConnection(leaderConnectionDetails)
+	apiConnection, err := createApiConnection(leaderConnectionDetails, l.metrics)
 	if err != nil {
 		return nil, errors.Wrapf(err, "error creating connection to leader")
 	}
@@ -69,12 +71,11 @@ func (l *LeaderConnectionProvider) getClientByName(currentLeaderName string) (*g
 	return apiConnection, nil
 }
 
-func createApiConnection(connectionDetails client.ApiConnectionDetails) (*grpc.ClientConn, error) {
-	grpc_prometheus.EnableClientHandlingTimeHistogram()
+func createApiConnection(connectionDetails client.ApiConnectionDetails, clientMetrics *grpc_prometheus.ClientMetrics) (*grpc.ClientConn, error) {
 	return client.CreateApiConnectionWithCallOptions(
 		&connectionDetails,
 		[]grpc.CallOption{},
-		grpc.WithChainUnaryInterceptor(grpc_prometheus.UnaryClientInterceptor),
-		grpc.WithChainStreamInterceptor(grpc_prometheus.StreamClientInterceptor),
+		grpc.WithChainUnaryInterceptor(clientMetrics.UnaryClientInterceptor()),
+		grpc.WithChainStreamInterceptor(clientMetrics.StreamClientInterceptor()),
 	)
 }
diff --git a/internal/scheduler/leader/leader_client_test.go b/internal/scheduler/leader/leader_client_test.go
@@ -3,6 +3,7 @@ package leader
 import (
 	"testing"
 
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/armadaproject/armada/internal/common/armadacontext"
@@ -28,7 +29,7 @@ var templatedLeader = configuration.LeaderConfig{
 
 func TestGetCurrentLeaderClientConnection(t *testing.T) {
 	leaderController := &FakeLeaderController{}
-	clientProvider := NewLeaderConnectionProvider(leaderController, defaultLeaderConfig)
+	clientProvider := NewLeaderConnectionProvider(leaderController, defaultLeaderConfig, &grpc_prometheus.ClientMetrics{})
 	leaderController.LeaderName = "new-leader"
 
 	isCurrentProcessLeader, result, err := clientProvider.GetCurrentLeaderClientConnection()
@@ -40,7 +41,7 @@ func TestGetCurrentLeaderClientConnection(t *testing.T) {
 
 func TestGetCurrentLeaderClientConnection_WithTemplatedConnection(t *testing.T) {
 	leaderController := &FakeLeaderController{}
-	clientProvider := NewLeaderConnectionProvider(leaderController, templatedLeader)
+	clientProvider := NewLeaderConnectionProvider(leaderController, templatedLeader, &grpc_prometheus.ClientMetrics{})
 
 	leaderController.LeaderName = "new-leader"
 	isCurrentProcessLeader, result, err := clientProvider.GetCurrentLeaderClientConnection()
@@ -59,7 +60,7 @@ func TestGetCurrentLeaderClientConnection_WithTemplatedConnection(t *testing.T)
 
 func TestGetCurrentLeaderClientConnection_NoLeader(t *testing.T) {
 	leaderController := &FakeLeaderController{}
-	clientProvider := NewLeaderConnectionProvider(leaderController, defaultLeaderConfig)
+	clientProvider := NewLeaderConnectionProvider(leaderController, defaultLeaderConfig, &grpc_prometheus.ClientMetrics{})
 
 	isCurrentProcessLeader, result, err := clientProvider.GetCurrentLeaderClientConnection()
 	assert.Nil(t, result)
@@ -69,7 +70,7 @@ func TestGetCurrentLeaderClientConnection_NoLeader(t *testing.T) {
 
 func TestGetCurrentLeaderClientConnection_OnCurrentProcessIsLeader(t *testing.T) {
 	leaderController := &FakeLeaderController{}
-	clientProvider := NewLeaderConnectionProvider(leaderController, defaultLeaderConfig)
+	clientProvider := NewLeaderConnectionProvider(leaderController, defaultLeaderConfig, &grpc_prometheus.ClientMetrics{})
 	leaderController.IsCurrentlyLeader = true
 
 	isCurrentProcessLeader, result, err := clientProvider.GetCurrentLeaderClientConnection()

diff --git a/internal/scheduler/schedulerapp.go b/internal/scheduler/schedulerapp.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/apache/pulsar-client-go/pulsar"
 	"github.com/google/uuid"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
 	grpc_logging "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
@@ -224,7 +225,12 @@ func Run(config schedulerconfig.Configuration) error {
 	schedulingContextRepository := reports.NewSchedulingContextRepository()
 	reportServer := reports.NewServer(schedulingContextRepository)
 
-	leaderClientConnectionProvider := leader.NewLeaderConnectionProvider(leaderController, config.Leader)
+	clientMetrics := grpc_prometheus.NewClientMetrics(
+		grpc_prometheus.WithClientHandlingTimeHistogram(),
+	)
+	prometheus.MustRegister(clientMetrics)
+
+	leaderClientConnectionProvider := leader.NewLeaderConnectionProvider(leaderController, config.Leader, clientMetrics)
 	schedulingSchedulerReportingServer := reports.NewLeaderProxyingSchedulingReportsServer(reportServer, leaderClientConnectionProvider)
 	schedulerobjects.RegisterSchedulerReportingServer(grpcServer, schedulingSchedulerReportingServer)
 

diff --git a/internal/server/server.go b/internal/server/server.go
@@ -7,7 +7,7 @@ import (
 
 	"github.com/apache/pulsar-client-go/pulsar"
 	"github.com/google/uuid"
-	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/redis/go-redis/extra/redisprometheus/v9"
@@ -194,7 +194,6 @@ func Serve(ctx *armadacontext.Context, config *configuration.ArmadaConfig, healt
 	api.RegisterExecutorServer(grpcServer, executorServer)
 
 	schedulerobjects.RegisterSchedulerReportingServer(grpcServer, schedulingReportsServer)
-	grpc_prometheus.Register(grpcServer)
 
 	// Cancel the errgroup if grpcServer.Serve returns an error.
 	log.Infof("Armada gRPC server listening on %d", config.GrpcPort)
@@ -235,11 +234,14 @@ func validateSubmissionConfig(config configuration.SubmissionConfig) error {
 }
 
 func createApiConnection(connectionDetails client.ApiConnectionDetails) (*grpc.ClientConn, error) {
-	grpc_prometheus.EnableClientHandlingTimeHistogram()
+	clientMetrics := grpc_prometheus.NewClientMetrics(
+		grpc_prometheus.WithClientHandlingTimeHistogram(),
+	)
+	prometheus.MustRegister(clientMetrics)
 	return client.CreateApiConnectionWithCallOptions(
 		&connectionDetails,
 		[]grpc.CallOption{},
-		grpc.WithChainUnaryInterceptor(grpc_prometheus.UnaryClientInterceptor),
-		grpc.WithChainStreamInterceptor(grpc_prometheus.StreamClientInterceptor),
+		grpc.WithChainUnaryInterceptor(clientMetrics.UnaryClientInterceptor()),
+		grpc.WithChainStreamInterceptor(clientMetrics.StreamClientInterceptor()),
 	)
 }