From a97507455e9228e30d051e3fc5ad76fc6f87245b Mon Sep 17 00:00:00 2001
From: Tim Wundenberg <tim@wundenbergs.de>
Date: Sun, 29 Sep 2024 00:53:09 +0200
Subject: [PATCH] chore: gracefull shutdown

---
 main.go | 74 ++++++++++++++++++++++++++++++++++++++-------------------
 1 file changed, 49 insertions(+), 25 deletions(-)

diff --git a/main.go b/main.go
index fe9f036..1687ee3 100644
--- a/main.go
+++ b/main.go
@@ -4,10 +4,15 @@ import (
 	"me-fit/handler"
 	"me-fit/utils"
 
+	"context"
 	"database/sql"
 	"log"
 	"log/slog"
 	"net/http"
+	"os/signal"
+	"sync"
+	"syscall"
+	"time"
 
 	"github.com/joho/godotenv"
 	_ "github.com/mattn/go-sqlite3"
@@ -15,48 +20,67 @@ import (
 )
 
 func main() {
+	run(context.Background())
+}
+
+func run(ctx context.Context) {
+
+	ctx, cancel := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
+
 	slog.Info("Starting server...")
 
+	// init env
 	err := godotenv.Load()
 	if err != nil {
 		log.Fatal("Error loading .env file")
 	}
 	utils.MustInitEnv()
 
+	// init db
 	db, err := sql.Open("sqlite3", "./data.db")
 	if err != nil {
 		log.Fatal("Could not open Database data.db: ", err)
 	}
 	defer db.Close()
-
 	utils.MustRunMigrations(db, "")
 
-	startPrometheus()
-
-	var server = http.Server{
-		Addr:    ":8080",
-		Handler: handler.GetHandler(db),
-	}
-	slog.Info("Starting server on " + server.Addr)
-
-	err = server.ListenAndServe()
-	if err != nil {
-		panic(err)
-	}
-}
-
-func startPrometheus() {
-
-	var prometheusServer = http.Server{
+	// init servers
+	prometheusServer := &http.Server{
 		Addr:    ":8081",
 		Handler: promhttp.Handler(),
 	}
+	httpServer := &http.Server{
+		Addr:    ":8080",
+		Handler: handler.GetHandler(db),
+	}
+	go startServer(prometheusServer)
+	go startServer(httpServer)
 
-	go func() {
-		slog.Info("Starting prometheus server on " + prometheusServer.Addr)
-		err := prometheusServer.ListenAndServe()
-		if err != nil {
-			log.Fatal("Could not start prometheus server: ", err)
-		}
-	}()
+	// graceful shutdown
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go shutdownServer(httpServer, ctx, &wg)
+	go shutdownServer(prometheusServer, ctx, &wg)
+	wg.Wait()
+}
+
+func startServer(s *http.Server) {
+	slog.Info("Starting server on " + s.Addr)
+	if err := s.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		slog.Error("error listening and serving: " + err.Error())
+	}
+}
+
+func shutdownServer(s *http.Server, ctx context.Context, wg *sync.WaitGroup) {
+	defer wg.Done()
+	<-ctx.Done()
+	shutdownCtx := context.Background()
+	shutdownCtx, cancel := context.WithTimeout(shutdownCtx, 10*time.Second)
+	defer cancel()
+	if err := s.Shutdown(shutdownCtx); err != nil {
+		slog.Error("error shutting down http server: " + err.Error())
+	} else {
+		slog.Info("Gracefully stopped http server on " + s.Addr)
+	}
 }