From 7f9314b55bd7efb29ea4ea6ae76b2110831783e4 Mon Sep 17 00:00:00 2001
From: Anurag Mittal <anurag.mittal@scality.com>
Date: Thu, 19 Dec 2024 02:21:08 +0100
Subject: [PATCH] Integrate Prometheus metrics server in COSI driver

- In main.go, allow graceful shutdown of the metrics server.
- Added `metricsAddress` flag to configure the Prometheus metrics endpoint.
- Integrated `metrics.StartMetricsServer` to expose metrics at the
  configured address.
- Ensured graceful shutdown of the metrics server during service
  termination.
- Updated the `run` function to include metrics server lifecycle
  management.
- Maintains flexibility for metrics configuration using constants from
  the `pkg/constants` package.

Issue: COSI-65
---
 cmd/scality-cosi-driver/cmd.go  | 22 +++++++++++++++++++---
 cmd/scality-cosi-driver/main.go |  1 +
 pkg/constants/constants.go      |  1 +
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/cmd/scality-cosi-driver/cmd.go b/cmd/scality-cosi-driver/cmd.go
index 59ea62c2..3528ebca 100644
--- a/cmd/scality-cosi-driver/cmd.go
+++ b/cmd/scality-cosi-driver/cmd.go
@@ -20,11 +20,14 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"time"
 
 	"github.com/scality/cosi-driver/pkg/driver"
 	"k8s.io/klog/v2"
 
+	c "github.com/scality/cosi-driver/pkg/constants"
 	"github.com/scality/cosi-driver/pkg/grpcfactory"
+	"github.com/scality/cosi-driver/pkg/metrics"
 )
 
 const (
@@ -33,8 +36,9 @@ const (
 )
 
 var (
-	driverAddress = flag.String("driver-address", "unix:///var/lib/cosi/cosi.sock", "driver address for the socket")
-	driverPrefix  = flag.String("driver-prefix", "", "prefix for COSI driver, e.g. <prefix>.scality.com")
+	driverAddress  = flag.String("driver-address", "unix:///var/lib/cosi/cosi.sock", "driver address for the socket")
+	driverPrefix   = flag.String("driver-prefix", "", "prefix for COSI driver, e.g. <prefix>.scality.com")
+	metricsAddress = flag.String("metrics-address", c.MetricsAddress, "The address to expose Prometheus metrics.")
 )
 
 func init() {
@@ -53,6 +57,11 @@ func init() {
 }
 
 func run(ctx context.Context) error {
+	metricsServer, err := metrics.StartMetricsServer(*metricsAddress)
+	if err != nil {
+		return fmt.Errorf("failed to start metrics server: %w", err)
+	}
+
 	driverName := *driverPrefix + "." + provisionerName
 
 	identityServer, bucketProvisioner, err := driver.CreateDriver(ctx, driverName)
@@ -65,5 +74,12 @@ func run(ctx context.Context) error {
 		return fmt.Errorf("failed to start the provisioner server: %w", err)
 	}
 
-	return server.Run(ctx)
+	err = server.Run(ctx)
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer shutdownCancel()
+	if shutdownErr := metricsServer.Shutdown(shutdownCtx); shutdownErr != nil {
+		klog.ErrorS(shutdownErr, "Failed to gracefully shutdown metrics server")
+	}
+
+	return err
 }
diff --git a/cmd/scality-cosi-driver/main.go b/cmd/scality-cosi-driver/main.go
index edc930fb..28835b9a 100644
--- a/cmd/scality-cosi-driver/main.go
+++ b/cmd/scality-cosi-driver/main.go
@@ -50,5 +50,6 @@ func main() {
 	// Call the run function (defined in cmd.go)
 	if err := run(ctx); err != nil {
 		klog.ErrorS(err, "Scality COSI driver encountered an error, shutting down")
+		os.Exit(1)
 	}
 }
diff --git a/pkg/constants/constants.go b/pkg/constants/constants.go
index b4250895..3b3c8a89 100644
--- a/pkg/constants/constants.go
+++ b/pkg/constants/constants.go
@@ -14,4 +14,5 @@ const (
 // Service initialization constants
 const (
 	MetricsPath = "/metrics"
+	MetricsAddress = ":8080"
 )