From 27e114b290f5e43b9fcf90ac09475a6d891b5f3d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8E=8B=E6=B5=A9=E4=BC=9F?=
 <50612703+wanghaowei0107@users.noreply.github.com>
Date: Thu, 25 Jan 2024 13:40:12 +0800
Subject: [PATCH] Add graceful shutdown (#867)

---
 main.go | 51 ++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/main.go b/main.go
index 6ab8424c..390cc54a 100644
--- a/main.go
+++ b/main.go
@@ -1,12 +1,16 @@
 package main
 
 import (
+	"context"
+	"errors"
 	"flag"
 	"net/http"
 	"os"
+	"os/signal"
 	"runtime"
 	"strconv"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -208,20 +212,41 @@ func main() {
 
 	log.Infof("Providing metrics at %s%s", *listenAddress, *metricPath)
 	log.Debugf("Configured redis addr: %#v", *redisAddr)
-	if *tlsServerCertFile != "" && *tlsServerKeyFile != "" {
-		log.Debugf("Bind as TLS using cert %s and key %s", *tlsServerCertFile, *tlsServerKeyFile)
+	server := &http.Server{
+		Addr:    *listenAddress,
+		Handler: exp,
+	}
+	go func() {
+		if *tlsServerCertFile != "" && *tlsServerKeyFile != "" {
+			log.Debugf("Bind as TLS using cert %s and key %s", *tlsServerCertFile, *tlsServerKeyFile)
 
-		tlsConfig, err := exp.CreateServerTLSConfig(*tlsServerCertFile, *tlsServerKeyFile, *tlsServerCaCertFile, *tlsServerMinVersion)
-		if err != nil {
-			log.Fatal(err)
+			tlsConfig, err := exp.CreateServerTLSConfig(*tlsServerCertFile, *tlsServerKeyFile, *tlsServerCaCertFile, *tlsServerMinVersion)
+			if err != nil {
+				log.Fatal(err)
+			}
+			server.TLSConfig = tlsConfig
+			if err := server.ListenAndServeTLS("", ""); err != nil && !errors.Is(err, http.ErrServerClosed) {
+				log.Fatalf("TLS Server error: %v", err)
+			}
+		} else {
+			if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+				log.Fatalf("Server error: %v", err)
+			}
 		}
-
-		server := &http.Server{
-			Addr:      *listenAddress,
-			TLSConfig: tlsConfig,
-			Handler:   exp}
-		log.Fatal(server.ListenAndServeTLS("", ""))
-	} else {
-		log.Fatal(http.ListenAndServe(*listenAddress, exp))
+	}()
+
+	// graceful shutdown
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
+	_quit := <-quit
+	log.Infof("Received %s signal, exiting", _quit.String())
+	// Create a context with a timeout
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	// Shutdown the HTTP server gracefully
+	if err := server.Shutdown(ctx); err != nil {
+		log.Fatalf("Server shutdown failed: %v", err)
 	}
+	log.Infof("Server shut down gracefully")
 }