From 9b0a3c775363a81eb5b3b7ee6c896d27087339f1 Mon Sep 17 00:00:00 2001
From: Stanislav Khalash <stanislav.khalash@sap.com>
Date: Wed, 23 Oct 2024 14:05:46 +0200
Subject: [PATCH] fix: Implement graceful shutdown

---
 Dockerfile |  2 +-
 main.go    | 64 ++++++++++++++++++++++++++++++++++++++++++++----------
 2 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index ab273f0..d21dd42 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM europe-docker.pkg.dev/kyma-project/prod/external/library/golang:1.23.1-alpine3.20 as build
+FROM europe-docker.pkg.dev/kyma-project/prod/external/library/golang:1.23.1-alpine3.20 AS build
 
 
 WORKDIR /src/
diff --git a/main.go b/main.go
index 95caae3..1706b53 100644
--- a/main.go
+++ b/main.go
@@ -1,19 +1,32 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"flag"
+	"fmt"
 	"log/slog"
 	"net/http"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+
 	"github.com/kyma-project/directory-size-exporter/internal/exporter"
+)
 
-	"github.com/prometheus/client_golang/prometheus/promhttp"
+const (
+	defaultLogFormat  = "json"
+	defaultLogLevel   = "info"
+	readHeaderTimeout = 1 * time.Second
+	shutdownTimeout   = 10 * time.Second
 )
 
 var (
+	logger = createLogger(defaultLogFormat, defaultLogLevel)
+
 	storagePath string
 	metricName  string
 	logFormat   string
@@ -23,8 +36,15 @@ var (
 )
 
 func main() {
-	flag.StringVar(&logFormat, "log-format", "json", "Log format (json or text)")
-	flag.StringVar(&logLevel, "log-level", "info", "Log level (debug, info, warn, error)")
+	if err := run(); err != nil {
+		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+		os.Exit(1)
+	}
+}
+
+func run() error {
+	flag.StringVar(&logFormat, "log-format", defaultLogFormat, "Log format (json or text)")
+	flag.StringVar(&logLevel, "log-level", defaultLogLevel, "Log level (debug, info, warn, error)")
 
 	flag.StringVar(&storagePath, "storage-path", "", "Path to the observed data folder")
 	flag.StringVar(&metricName, "metric-name", "", "Metric name used for exporting the folder size")
@@ -33,13 +53,11 @@ func main() {
 
 	flag.Parse()
 	if err := validateFlags(); err != nil {
-		panic(err)
+		return fmt.Errorf("invalid flags: %s", err)
 	}
 
-	logger := createLogger()
-
+	logger = createLogger(logFormat, logLevel)
 	exp := exporter.NewExporter(storagePath, metricName, logger)
-	logger.Info("Exporter is initialized")
 
 	exp.RecordMetrics(interval)
 	logger.Info("Started recording metrics")
@@ -47,13 +65,35 @@ func main() {
 	http.Handle("/metrics", promhttp.Handler())
 	server := &http.Server{
 		Addr:              ":" + port,
-		ReadHeaderTimeout: 1 * time.Second,
+		ReadHeaderTimeout: readHeaderTimeout,
 	}
 
-	if err := server.ListenAndServe(); err != nil {
-		panic(err)
+	go func() {
+		logger.Info("Listening on port '" + port + "'")
+
+		// When Shutdown is called, ListenAndServe will return http.ErrServerClosed, do not log it as an error
+		if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
+			logger.Error("HTTP server error: %v", slog.Any("err", err))
+		}
+
+		logger.Info("Stopped serving new connections.")
+	}()
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+	<-sigChan
+
+	shutdownCtx, shutdownRelease := context.WithTimeout(context.Background(), shutdownTimeout)
+	defer shutdownRelease()
+
+	if err := server.Shutdown(shutdownCtx); err != nil {
+		logger.Error("HTTP shutdown error: %v", slog.Any("err", err))
+		return err
 	}
-	logger.Info("Listening on port '" + port + "'")
+
+	logger.Info("Graceful shutdown complete.")
+
+	return nil
 }
 
 func validateFlags() error {
@@ -72,7 +112,7 @@ func validateFlags() error {
 	return nil
 }
 
-func createLogger() *slog.Logger {
+func createLogger(logFormat, logLevel string) *slog.Logger {
 	level := slog.LevelInfo
 	switch logLevel {
 	case "debug":