alternative heartbeat implementation

go.mod

@@ -5,4 +5,5 @@ go 1.20
 require (
 	github.com/sirupsen/logrus v1.9.2 // indirect
 	golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect
+	golang.org/x/time v0.3.0 // indirect
 )

go.sum

@@ -7,5 +7,7 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
+golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -8,7 +9,9 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
 	"path/filepath"
+	"syscall"
 	"time"
 
 	"github.com/sirupsen/logrus"
@@ -181,14 +184,27 @@ func main() {
 		go func() {
 			for range time.Tick(time.Duration(config.Heartbeat)) {
 				logger.Info("Server heartbeat")
-				os.Exit(0)
+				server := startServer(&config, logger)
+				stopServer(server, logger)
 			}
 		}()
 	}
 
+	// Start the initial server
+	server := startServer(&config, logger)
+
+	// Wait for termination signal
+	waitForTerminationSignal()
+
+	// Stop the server before exiting
+	stopServer(server, logger)
+}
+
+// startServer creates and starts the HTTP server
+func startServer(config *Config, logger *logrus.Logger) *http.Server {
 	// Create a new CDNHandler with the configuration
 	cdnHandler := &CDNHandler{
-		Config: config,
+		Config: *config,
 		Logger: logger,
 	}
 
@@ -201,6 +217,34 @@ func main() {
 		WriteTimeout: 10 * time.Second,
 	}
 
+	// Start the server in a separate goroutine
+	go func() {
 		logger.Infof("Starting CDN server on port %s", config.Port)
-	log.Fatal(server.ListenAndServe())
+		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			logger.Fatalf("Server error: %v", err)
+		}
+	}()
+
+	return server
+}
+
+// stopServer stops the HTTP server
+func stopServer(server *http.Server, logger *logrus.Logger) {
+	logger.Info("Stopping CDN server")
+
+	// Set a deadline for gracefully shutting down the server
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	// Shut down the server with the given context
+	if err := server.Shutdown(ctx); err != nil {
+		logger.Errorf("Server shutdown error: %v", err)
+	}
+}
+
+// waitForTerminationSignal waits for termination signals to gracefully shut down the server
+func waitForTerminationSignal() {
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, os.Interrupt, syscall.SIGTERM)
+	<-quit
 }