Graceful Shutdown

package main

import (
    "context"
    "fmt"
    "log"
    "net/http"
    "os"
    "os/signal"
    "syscall"
    "time"
)

type Server struct {
    httpServer *http.Server
    workers    []Worker
}

type Worker interface {
    Start(ctx context.Context) error
    Shutdown(ctx context.Context) error
}

// BackgroundWorker processes jobs in the background
type BackgroundWorker struct {
    name string
}

func (w *BackgroundWorker) Start(ctx context.Context) error {
    log.Printf("Starting worker: %s", w.name)
    for {
        select {
        case <-ctx.Done():
            log.Printf("Worker %s received shutdown signal", w.name)
            return nil
        case <-time.After(2 * time.Second):
            log.Printf("Worker %s processing...", w.name)
        }
    }
}

func (w *BackgroundWorker) Shutdown(ctx context.Context) error {
    log.Printf("Worker %s shutting down gracefully", w.name)
    // Perform cleanup: flush buffers, save state, etc.
    time.Sleep(500 * time.Millisecond) // Simulate cleanup
    return nil
}

func NewServer(addr string, workers ...Worker) *Server {
    mux := http.NewServeMux()
    mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        // Simulate work
        time.Sleep(3 * time.Second)
        fmt.Fprintf(w, "Request completed at %s\n", time.Now().Format(time.RFC3339))
    })

    return &Server{
        httpServer: &http.Server{
            Addr:    addr,
            Handler: mux,
        },
        workers: workers,
    }
}

func (s *Server) Run() error {
    // Create context that will be cancelled on shutdown signal
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    // Channel to listen for interrupt signals
    quit := make(chan os.Signal, 1)
    signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

    // Start background workers
    for _, worker := range s.workers {
        w := worker // Capture loop variable
        go func() {
            if err := w.Start(ctx); err != nil {
                log.Printf("Worker error: %v", err)
            }
        }()
    }

    // Start HTTP server in a goroutine
    serverErrors := make(chan error, 1)
    go func() {
        log.Printf("Server listening on %s", s.httpServer.Addr)
        serverErrors <- s.httpServer.ListenAndServe()
    }()

    // Block until we receive a signal or server error
    select {
    case err := <-serverErrors:
        return fmt.Errorf("server error: %w", err)
    case sig := <-quit:
        log.Printf("Received signal: %v. Starting graceful shutdown...", sig)

        // Cancel context to stop workers
        cancel()

        // Give workers time to finish
        shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
        defer shutdownCancel()

        // Shutdown HTTP server
        log.Println("Shutting down HTTP server...")
        if err := s.httpServer.Shutdown(shutdownCtx); err != nil {
            log.Printf("HTTP server shutdown error: %v", err)
            return err
        }

        // Shutdown workers
        for _, worker := range s.workers {
            if err := worker.Shutdown(shutdownCtx); err != nil {
                log.Printf("Worker shutdown error: %v", err)
            }
        }

        log.Println("Graceful shutdown complete")
        return nil
    }
}

package main

import (
    "log"
)

func main() {
    // Create workers
    worker1 := &BackgroundWorker{name: "metrics-flusher"}
    worker2 := &BackgroundWorker{name: "event-processor"}

    // Create server with workers
    server := NewServer(":8080", worker1, worker2)

    // Run server (blocks until shutdown)
    if err := server.Run(); err != nil && err != http.ErrServerClosed {
        log.Fatalf("Server failed: %v", err)
    }
}

// Test it:
// 1. Run the server: go run main.go
// 2. Make a long request: curl http://localhost:8080
// 3. Press Ctrl+C while request is running
// 4. Server waits for request to complete before shutting down