Handle shutdown signals gracefully

cmd/order-book/main.go

@@ -1,12 +1,16 @@
 package main
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
 	"strconv"
 	"strings"
+	"syscall"
+	"time"
 
 	"github.com/ethereum/go-ethereum/ethclient"
 	"github.com/gorilla/mux"
@@ -116,7 +120,25 @@ func setup() {
 	handler.Init(userSvc.GetUserByApiKey)
 
 	server := rest.NewHTTPServer(":"+port, handler.Router)
-	server.StartServer()
-	// blocking
-	<-server.StopChannel
+
+	// Handle SIGINT and SIGTERM signals
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		server.StartServer()
+	}()
+
+	sig := <-signalChan
+	log.Printf("Received SIGTERM signal: %s. Initiating shutdown...\n", sig)
+
+	// Heroku gives 30 seconds to shutdown
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancel()
+
+	if err := server.StopServer(ctx); err != nil {
+		log.Fatalf("Failed to stop server: %v\n", err)
+	}
+
+	log.Printf("Server gracefully stopped\n")
 }

transport/rest/server.go

@@ -3,6 +3,7 @@ package rest
 import (
 	"context"
 	"fmt"
+	"log"
 	"net/http"
 	"os"
 
@@ -29,15 +30,17 @@ func NewHTTPServer(addr string, router *mux.Router) *HTTPServer {
 func (hs *HTTPServer) StartServer() {
 	go func() {
 		fmt.Printf("HTTP server started on %s\n", hs.server.Addr)
-		if err := hs.server.ListenAndServe(); err != nil {
-			fmt.Printf("HTTP server error: %v\n", err)
+		// ListenAndServe returns ErrServerClosed on graceful shutdown
+		if err := hs.server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			panic(fmt.Errorf("HTTP server error: %v", err))
 		}
 	}()
 }
 
-func (hs *HTTPServer) StopServer(ctx context.Context) {
-	fmt.Println("Shutting down the HTTP server...")
+func (hs *HTTPServer) StopServer(ctx context.Context) error {
+	log.Printf("Shutting down the HTTP server...\n")
 	if err := hs.server.Shutdown(ctx); err != nil {
-		fmt.Printf("HTTP server shutdown error: %v\n", err)
+		return fmt.Errorf("HTTP server shutdown error: %v", err)
 	}
+	return nil
 }