bug_fix/gracefulShutdown (#310)

cmd/template/framework/files/main/fiber_main.go.tmpl

@@ -14,7 +14,7 @@ import (
 	_ "github.com/joho/godotenv/autoload"
 )
 
-func gracefulShutdown(fiberServer *server.FiberServer) {
+func gracefulShutdown(fiberServer *server.FiberServer, done chan bool) {
 	// Create context that listens for the interrupt signal from the OS.
 	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
 	defer stop()
@@ -26,20 +26,27 @@ func gracefulShutdown(fiberServer *server.FiberServer) {
 
 	// The context is used to inform the server it has 5 seconds to finish
 	// the request it is currently handling
-	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 	if err := fiberServer.ShutdownWithContext(ctx); err != nil {
 		log.Printf("Server forced to shutdown with error: %v", err)
 	}
 
 	log.Println("Server exiting")
+
+	// Notify the main goroutine that the shutdown is complete
+	done <- true
 }
 
 func main() {
 
 	server := server.New()
 
 	server.RegisterFiberRoutes()
+
+	// Create a done channel to signal when the shutdown is complete
+	done := make(chan bool, 1)
+
 	go func() {
 		port, _ := strconv.Atoi(os.Getenv("PORT"))
 		err := server.Listen(fmt.Sprintf(":%d", port))
@@ -48,5 +55,10 @@ func main() {
 		}
 	}()
 
-	gracefulShutdown(server)
+	// Run graceful shutdown in a separate goroutine
+	go gracefulShutdown(server, done)
+
+	// Wait for the graceful shutdown to complete
+	<-done
+	log.Println("Graceful shutdown complete.")
 }

cmd/template/framework/files/main/main.go.tmpl

@@ -12,7 +12,7 @@ import (
 	"{{.ProjectName}}/internal/server"
 )
 
-func gracefulShutdown(apiServer *http.Server) {
+func gracefulShutdown(apiServer *http.Server, done chan bool) {
 	// Create context that listens for the interrupt signal from the OS.
 	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
 	defer stop()
@@ -24,24 +24,34 @@ func gracefulShutdown(apiServer *http.Server) {
 
 	// The context is used to inform the server it has 5 seconds to finish
 	// the request it is currently handling
-	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	defer cancel()
 	if err := apiServer.Shutdown(ctx); err != nil {
-  		log.Printf("Server forced to shutdown with error: %v", err)
+		log.Printf("Server forced to shutdown with error: %v", err)
 	}
 
 	log.Println("Server exiting")
-}
 
+	// Notify the main goroutine that the shutdown is complete
+	done <- true
+}
 
 func main() {
 
 	server := server.NewServer()
 
-	go gracefulShutdown(server)
+	// Create a done channel to signal when the shutdown is complete
+	done := make(chan bool, 1)
+
+	// Run graceful shutdown in a separate goroutine
+	go gracefulShutdown(server, done)
 
 	err := server.ListenAndServe()
 	if err != nil && err != http.ErrServerClosed {
 		panic(fmt.Sprintf("http server error: %s", err))
 	}
+
+	// Wait for the graceful shutdown to complete
+	<-done
+	log.Println("Graceful shutdown complete.")
 }