Implement graceful shutdown for the HTTP server

main.go

@@ -1,9 +1,14 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
 )
 
 func printBanner() {
@@ -43,5 +48,21 @@ func main() {
 	}
 
 	server := NewServer(addr, cacheDir)
-	log.Fatal(server.Run())
+	srv, err := server.Start()
+	if err != nil {
+		log.Fatalf("Failed to start server: %v", err)
+	}
+
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
+	sig := <-quit
+	log.Printf("Received %s, shutting down gracefully...", sig)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	if err := srv.Shutdown(ctx); err != nil {
+		log.Fatalf("Server forced to shutdown: %v", err)
+	}
+
+	log.Println("Server stopped")
 }

server.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"log"
+	"net"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -46,10 +47,11 @@ func NewServerWithCache(addr, cacheDir string) *Server {
 	return &Server{addr: addr, cacheDir: cacheDir}
 }
 
-// Run starts the HTTP server
-func (s *Server) Run() error {
+// Start starts the HTTP server and returns the *http.Server for shutdown control.
+// It begins accepting connections immediately in a background goroutine.
+func (s *Server) Start() (*http.Server, error) {
 	if err := os.MkdirAll(s.cacheDir, 0755); err != nil {
-		return fmt.Errorf("failed to create cache directory: %w", err)
+		return nil, fmt.Errorf("failed to create cache directory: %w", err)
 	}
 
 	mux := http.NewServeMux()
@@ -59,8 +61,21 @@ func (s *Server) Run() error {
 	mux.HandleFunc("GET /logo.png", s.logoHandler)
 	mux.Handle("GET /metrics", promhttp.Handler())
 
+	srv := &http.Server{Addr: s.addr, Handler: mux}
+
+	ln, err := net.Listen("tcp", s.addr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to listen on %s: %w", s.addr, err)
+	}
+
 	log.Printf("Starting server on %s (cache: %s)\n", s.addr, s.cacheDir)
-	return http.ListenAndServe(s.addr, mux)
+	go func() {
+		if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed {
+			log.Fatalf("Server error: %v", err)
+		}
+	}()
+
+	return srv, nil
 }
 
 // healthHandler handles the /health endpoint