From e2b15245d51d95539e32204eb97ae7580be6dbcf Mon Sep 17 00:00:00 2001 From: Sailaja Kola <8sailaja@gmail.com> Date: Thu, 11 Jun 2026 22:05:50 +0530 Subject: [PATCH 1/2] feat(lifecycle): add graceful shutdown manager with lifecycle hooks --- README.md | 157 ++++++++++++++++++++++++++++ engine.go | 112 +++++++++++++++++++- lifecycle.go | 105 +++++++++++++++++++ lifecycle_test.go | 254 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 624 insertions(+), 4 deletions(-) create mode 100644 lifecycle.go create mode 100644 lifecycle_test.go diff --git a/README.md b/README.md index a9064ea..8e0fc10 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,7 @@ Unlike traditional Go frameworks that focus only on request handling, Vodka heav - [Minimal API Example](#minimal-api-example) - [Using Vodka for APIs](#using-vodka-for-apis) - [Core Concepts](#core-concepts) +- [Graceful Shutdown & Lifecycle Manager](#graceful-shutdown--lifecycle-manager) - [Middleware](#middleware) - [Validation](#validation) - [Authentication](#authentication) @@ -412,6 +413,162 @@ c.Error(400, errors.New("invalid request")) --- +# Graceful Shutdown & Lifecycle Manager + +Vodka features a production-ready application lifecycle management system that handles: +- **Startup Hooks**: Run code before the HTTP server starts. +- **Shutdown Hooks**: Run code sequentially when the application receives termination signals (`SIGINT`, `SIGTERM`). +- **Priority-Based Execution**: Control the shutdown sequence of your resources. +- **Graceful HTTP Server Shutdown**: Stop accepting new connections and finish active requests. +- **Configurable Shutdown Timeout**: Cancel remaining work if the timeout expires. +- **Error Aggregation**: Collect and report errors from all hooks. + +--- + +## Startup Hooks + +Startup hooks allow executing initialization code (e.g., establishing database connections, seeding data) before the server begins serving requests. If any startup hook returns an error, server startup is immediately aborted. + +```go +app.OnStart(func() error { + return initializeDatabase() +}) +``` + +--- + +## Shutdown Hooks & Priority + +Shutdown hooks are executed sequentially when a termination signal is received. You can register hooks with a priority value: + +- Higher priority hooks execute first. +- If priorities are equal, hooks execute in registration order. +- The default priority is `0` when calling `OnShutdown`. + +```go +// Priority 100: runs first +app.OnShutdownWithPriority(100, func(ctx context.Context) error { + return closeDatabase() +}) + +// Priority 50: runs second +app.OnShutdownWithPriority(50, func(ctx context.Context) error { + return stopWorkers() +}) + +// Default priority (0): runs last +app.OnShutdown(func(ctx context.Context) error { + return cleanupTempFiles() +}) +``` + +--- + +## Timeout Configuration + +By default, Vodka allows up to `30 seconds` for the entire shutdown sequence (including finishing active HTTP requests and running all shutdown hooks). You can customize this timeout: + +```go +app.SetShutdownTimeout(45 * time.Second) +``` + +--- + +## Production Example + +Here is a full production-ready example demonstrating database cleanup, worker queue termination, and graceful server shutdown: + +```go +package main + +import ( + "context" + "database/sql" + "fmt" + "log" + "time" + + "github.com/DevanshuTripathi/vodka" + _ "github.com/lib/pq" +) + +func main() { + app := vodka.DefaultRouter() + + var db *sql.DB + + // 1. Register Startup Hook to initialize database + app.OnStart(func() error { + var err error + db, err = sql.Open("postgres", "postgres://user:pass@localhost/db?sslmode=disable") + if err != nil { + return fmt.Errorf("failed to open database: %w", err) + } + + // Verify connection + if err := db.Ping(); err != nil { + return fmt.Errorf("failed to ping database: %w", err) + } + log.Println("Database connection established") + return nil + }) + + // 2. Start workers + workerCtx, cancelWorkers := context.WithCancel(context.Background()) + app.OnStart(func() error { + go runBackgroundWorkers(workerCtx) + log.Println("Background workers started") + return nil + }) + + // 3. Configure Shutdown Timeout + app.SetShutdownTimeout(15 * time.Second) + + // 4. Register Shutdown Hooks in priority order + + // Stop background workers first (Priority 100) + app.OnShutdownWithPriority(100, func(ctx context.Context) error { + log.Println("Stopping background workers...") + cancelWorkers() + return nil + }) + + // Close database connection (Priority 50) + app.OnShutdownWithPriority(50, func(ctx context.Context) error { + log.Println("Closing database connection...") + if db != nil { + return db.Close() + } + return nil + }) + + // Simple route + app.GET("/", func(c *vodka.Context) { + c.String(200, "Hello, Graceful Vodka!") + }) + + // 5. Start the server (handles SIGINT/SIGTERM automatically) + if err := app.Run(":8080"); err != nil { + log.Fatalf("Server stopped with error: %v", err) + } +} + +func runBackgroundWorkers(ctx context.Context) { + for { + select { + case <-ctx.Done(): + log.Println("Workers stopped") + return + default: + // Perform background work + time.Sleep(1 * time.Second) + } + } +} +``` + +--- + # Middleware Vodka middleware is simply a `vodka.HandlerFunc`. diff --git a/engine.go b/engine.go index e57f2d2..360fd72 100644 --- a/engine.go +++ b/engine.go @@ -1,15 +1,21 @@ package vodka import ( + "context" + "fmt" "html/template" "log" "net" "net/http" "os" + "os/signal" "path" "path/filepath" + "sort" "strings" "sync" + "syscall" + "time" "github.com/gorilla/websocket" "github.com/julienschmidt/httprouter" @@ -29,6 +35,7 @@ type Engine struct { templates map[string]*template.Template templatesMu sync.RWMutex *RouterGroup + lifecycle *LifecycleManager } // creates a new router @@ -36,8 +43,9 @@ func NewRouter() *Engine { router := httprouter.New() router.HandleOPTIONS = false engine := &Engine{ - router: router, - WSConfig: DefaultWSConfig(), + router: router, + WSConfig: DefaultWSConfig(), + lifecycle: NewLifecycleManager(), } engine.RouterGroup = &RouterGroup{ @@ -72,15 +80,111 @@ func (rg *RouterGroup) Use(middlewares ...HandlerFunc) { } // Runs the http server +// Runs the http server with startup/shutdown hook management and graceful shutdown func (e *Engine) Run(addr string) error { if addr == "" { addr = ":8080" } + // 1. Execute startup hooks before serving + if err := e.lifecycle.runStartupHooks(); err != nil { + return err + } + log.Printf(Green+"Pouring Vodka on %s\n"+Reset, addr) - // Using net/http - return http.ListenAndServe(addr, e) + ln, err := net.Listen("tcp", addr) + if err != nil { + return err + } + + srv := &http.Server{ + Addr: addr, + Handler: e, + } + + serverErr := make(chan error, 1) + go func() { + if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed { + serverErr <- err + } + }() + + quit := make(chan os.Signal, 1) + signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + + select { + case err := <-serverErr: + // If server failed to start/serve, perform shutdown to clean up hooks + shutdownErr := e.shutdown(srv) + if shutdownErr != nil { + return fmt.Errorf("server error: %v; shutdown errors: %v", err, shutdownErr) + } + return err + case sig := <-quit: + log.Printf(Yellow+"Received signal: %v. Initiating graceful shutdown...\n"+Reset, sig) + return e.shutdown(srv) + } +} + +// shutdown handles graceful HTTP server shutdown and executes registered shutdown hooks. +func (e *Engine) shutdown(srv *http.Server) error { + e.lifecycle.mu.Lock() + timeout := e.lifecycle.timeout + hooks := make([]lifecycleHook, len(e.lifecycle.shutdownHooks)) + copy(hooks, e.lifecycle.shutdownHooks) + e.lifecycle.mu.Unlock() + + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + var shutdownErrors []error + + // Step 2: Gracefully finish active requests + if err := srv.Shutdown(ctx); err != nil { + shutdownErrors = append(shutdownErrors, fmt.Errorf("server shutdown failed: %w", err)) + } + + // Step 3: Sort shutdown hooks by priority (descending), preserving registration order for equal priorities. + sort.SliceStable(hooks, func(i, j int) bool { + if hooks[i].priority != hooks[j].priority { + return hooks[i].priority > hooks[j].priority + } + return hooks[i].order < hooks[j].order + }) + + // Step 4: Execute registered shutdown hooks sequentially + for _, hook := range hooks { + if err := ctx.Err(); err != nil { + shutdownErrors = append(shutdownErrors, fmt.Errorf("shutdown timeout exceeded: %w", err)) + break + } + if err := hook.fn(ctx); err != nil { + shutdownErrors = append(shutdownErrors, err) + } + } + + if len(shutdownErrors) > 0 { + return &ShutdownError{Errors: shutdownErrors} + } + + return nil +} + +func (e *Engine) OnStart(fn func() error) { + e.lifecycle.RegisterStart(fn) +} + +func (e *Engine) OnShutdown(fn func(context.Context) error) { + e.lifecycle.RegisterShutdown(0, fn) +} + +func (e *Engine) OnShutdownWithPriority(priority int, fn func(context.Context) error) { + e.lifecycle.RegisterShutdown(priority, fn) +} + +func (e *Engine) SetShutdownTimeout(timeout time.Duration) { + e.lifecycle.SetTimeout(timeout) } // LoadHTMLGlob parses and caches templates from a glob pattern diff --git a/lifecycle.go b/lifecycle.go new file mode 100644 index 0000000..0f15b94 --- /dev/null +++ b/lifecycle.go @@ -0,0 +1,105 @@ +package vodka + +import ( + "context" + "fmt" + "strings" + "sync" + "time" +) + +// StartHook represents a function executed when the application starts. +type StartHook func() error + +// ShutdownHook represents a function executed when the application shuts down. +type ShutdownHook func(context.Context) error + +type lifecycleHook struct { + priority int + order int + fn ShutdownHook +} + +// LifecycleManager manages startup and shutdown hooks. +type LifecycleManager struct { + startupHooks []StartHook + shutdownHooks []lifecycleHook + timeout time.Duration + mu sync.Mutex +} + +// NewLifecycleManager creates a new LifecycleManager with a default timeout. +func NewLifecycleManager() *LifecycleManager { + return &LifecycleManager{ + startupHooks: make([]StartHook, 0), + shutdownHooks: make([]lifecycleHook, 0), + timeout: 30 * time.Second, + } +} + +// RegisterStart registers a new startup hook. +func (lm *LifecycleManager) RegisterStart(fn StartHook) { + lm.mu.Lock() + defer lm.mu.Unlock() + lm.startupHooks = append(lm.startupHooks, fn) +} + +// RegisterShutdown registers a new shutdown hook with the given priority. +func (lm *LifecycleManager) RegisterShutdown(priority int, fn ShutdownHook) { + lm.mu.Lock() + defer lm.mu.Unlock() + order := len(lm.shutdownHooks) + lm.shutdownHooks = append(lm.shutdownHooks, lifecycleHook{ + priority: priority, + order: order, + fn: fn, + }) +} + +// SetTimeout configures the maximum duration allowed for shutdown. +func (lm *LifecycleManager) SetTimeout(timeout time.Duration) { + lm.mu.Lock() + defer lm.mu.Unlock() + lm.timeout = timeout +} + +// runStartupHooks runs all startup hooks in registration order. +func (lm *LifecycleManager) runStartupHooks() error { + lm.mu.Lock() + hooks := make([]StartHook, len(lm.startupHooks)) + copy(hooks, lm.startupHooks) + lm.mu.Unlock() + + for _, hook := range hooks { + if err := hook(); err != nil { + return err + } + } + return nil +} + +// ShutdownError represents a collection of errors encountered during shutdown. +type ShutdownError struct { + Errors []error +} + +// Error formats the aggregated errors. +func (e *ShutdownError) Error() string { + if len(e.Errors) == 0 { + return "" + } + var sb strings.Builder + sb.WriteString(fmt.Sprintf("%d shutdown errors:\n", len(e.Errors))) + for i, err := range e.Errors { + sb.WriteString(fmt.Sprintf("- %v", err)) + if i < len(e.Errors)-1 { + sb.WriteString("\n") + } + } + return sb.String() +} + +// Unwrap supports Go 1.20+ multi-error unwrapping. +func (e *ShutdownError) Unwrap() []error { + return e.Errors +} diff --git a/lifecycle_test.go b/lifecycle_test.go new file mode 100644 index 0000000..3f00098 --- /dev/null +++ b/lifecycle_test.go @@ -0,0 +1,254 @@ +package vodka + +import ( + "context" + "errors" + "net/http" + "os" + "reflect" + "strings" + "sync" + "testing" + "time" +) + +func TestStartupHooks(t *testing.T) { + // 1. Single startup hook + t.Run("SingleHook", func(t *testing.T) { + app := NewRouter() + called := false + app.OnStart(func() error { + called = true + return nil + }) + + err := app.lifecycle.runStartupHooks() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !called { + t.Fatal("expected startup hook to be called") + } + }) + + // 2. Multiple startup hooks (order of execution) + t.Run("MultipleHooksOrder", func(t *testing.T) { + app := NewRouter() + var order []string + var mu sync.Mutex + + app.OnStart(func() error { + mu.Lock() + order = append(order, "first") + mu.Unlock() + return nil + }) + app.OnStart(func() error { + mu.Lock() + order = append(order, "second") + mu.Unlock() + return nil + }) + + err := app.lifecycle.runStartupHooks() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + expected := []string{"first", "second"} + if !reflect.DeepEqual(order, expected) { + t.Fatalf("expected order %v, got %v", expected, order) + } + }) + + // 3. Startup hook failure + t.Run("StartupFailure", func(t *testing.T) { + app := NewRouter() + app.OnStart(func() error { + return errors.New("startup failed") + }) + + // Verify Run returns the startup error and aborts + err := app.Run(":invalid_port_or_addr") + if err == nil || err.Error() != "startup failed" { + t.Fatalf("expected 'startup failed' error, got: %v", err) + } + }) +} + +func TestShutdownHooksPriority(t *testing.T) { + app := NewRouter() + var order []string + var mu sync.Mutex + + app.OnShutdownWithPriority(50, func(ctx context.Context) error { + mu.Lock() + order = append(order, "priority 50 (first registration)") + mu.Unlock() + return nil + }) + + app.OnShutdownWithPriority(100, func(ctx context.Context) error { + mu.Lock() + order = append(order, "priority 100") + mu.Unlock() + return nil + }) + + app.OnShutdownWithPriority(50, func(ctx context.Context) error { + mu.Lock() + order = append(order, "priority 50 (second registration)") + mu.Unlock() + return nil + }) + + app.OnShutdown(func(ctx context.Context) error { + mu.Lock() + order = append(order, "default priority (0)") + mu.Unlock() + return nil + }) + + app.OnShutdownWithPriority(-10, func(ctx context.Context) error { + mu.Lock() + order = append(order, "priority -10") + mu.Unlock() + return nil + }) + + // Create a dummy http server + srv := &http.Server{} + err := app.shutdown(srv) + if err != nil { + t.Fatalf("unexpected error during shutdown: %v", err) + } + + expected := []string{ + "priority 100", + "priority 50 (first registration)", + "priority 50 (second registration)", + "default priority (0)", + "priority -10", + } + + if !reflect.DeepEqual(order, expected) { + t.Fatalf("expected shutdown execution order %v, got %v", expected, order) + } +} + +func TestShutdownTimeout(t *testing.T) { + app := NewRouter() + app.SetShutdownTimeout(50 * time.Millisecond) + + app.OnShutdown(func(ctx context.Context) error { + select { + case <-time.After(150 * time.Millisecond): + return nil + case <-ctx.Done(): + return ctx.Err() + } + }) + + // The second hook should not run because context gets cancelled + secondCalled := false + app.OnShutdown(func(ctx context.Context) error { + secondCalled = true + return nil + }) + + srv := &http.Server{} + err := app.shutdown(srv) + if err == nil { + t.Fatal("expected error due to timeout, got nil") + } + + // The error should mention shutdown timeout exceeded/context deadline exceeded + if !strings.Contains(err.Error(), "context deadline exceeded") && !strings.Contains(err.Error(), "shutdown timeout exceeded") { + t.Fatalf("expected timeout error message, got: %v", err) + } + + if secondCalled { + t.Fatal("expected second hook not to run because context was cancelled") + } +} + +func TestErrorAggregation(t *testing.T) { + app := NewRouter() + + app.OnShutdown(func(ctx context.Context) error { + return errors.New("Database close failed") + }) + + app.OnShutdown(func(ctx context.Context) error { + return errors.New("Worker stop failed") + }) + + srv := &http.Server{} + err := app.shutdown(srv) + if err == nil { + t.Fatal("expected aggregated errors, got nil") + } + + expectedMsg := "2 shutdown errors:\n- Database close failed\n- Worker stop failed" + if err.Error() != expectedMsg { + t.Fatalf("expected formatted error:\n%q\ngot:\n%q", expectedMsg, err.Error()) + } + + // Verify Go 1.20+ Unwrap compatibility + unwrapper, ok := err.(interface{ Unwrap() []error }) + if !ok { + t.Fatal("expected error to implement Unwrap() []error") + } + + errs := unwrapper.Unwrap() + if len(errs) != 2 { + t.Fatalf("expected 2 errors, got %d: %v", len(errs), errs) + } + + if errs[0].Error() != "Database close failed" || errs[1].Error() != "Worker stop failed" { + t.Fatalf("unexpected unwrapped errors: %v", errs) + } +} + +func TestSignalHandling(t *testing.T) { + app := NewRouter() + + shutdownCalled := false + app.OnShutdown(func(ctx context.Context) error { + shutdownCalled = true + return nil + }) + + // Run the server in a goroutine + errChan := make(chan error, 1) + go func() { + errChan <- app.Run(":18080") + }() + + // Wait a moment for server to start listening + time.Sleep(100 * time.Millisecond) + + // Send SIGINT to ourselves + p, err := os.FindProcess(os.Getpid()) + if err != nil { + t.Fatalf("failed to find current process: %v", err) + } + + err = p.Signal(os.Interrupt) + if err != nil { + t.Fatalf("failed to send interrupt signal: %v", err) + } + + // Wait for Run to return + select { + case runErr := <-errChan: + if runErr != nil { + t.Fatalf("Run returned unexpected error: %v", runErr) + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for server to shutdown gracefully") + } + + if !shutdownCalled { + t.Fatal("expected shutdown hooks to be called on signal") + } +} From 420b6509114b175d36e05629acb5105783e1d978 Mon Sep 17 00:00:00 2001 From: Kola Sailaja <8sailaja@gmail.com> Date: Mon, 22 Jun 2026 10:21:20 +0530 Subject: [PATCH 2/2] feat: add graceful shutdown manager with lifecycle hooks --- engine.go | 22 ++++++++++++++++++-- lifecycle_test.go | 51 ++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/engine.go b/engine.go index 360fd72..d76bd81 100644 --- a/engine.go +++ b/engine.go @@ -36,6 +36,7 @@ type Engine struct { templatesMu sync.RWMutex *RouterGroup lifecycle *LifecycleManager + shutdownOnce sync.Once } // creates a new router @@ -86,10 +87,13 @@ func (e *Engine) Run(addr string) error { addr = ":8080" } + log.Println("Starting application...") + // 1. Execute startup hooks before serving if err := e.lifecycle.runStartupHooks(); err != nil { return err } + log.Println("Startup hooks completed") log.Printf(Green+"Pouring Vodka on %s\n"+Reset, addr) @@ -112,6 +116,7 @@ func (e *Engine) Run(addr string) error { quit := make(chan os.Signal, 1) signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) + defer signal.Stop(quit) select { case err := <-serverErr: @@ -122,13 +127,22 @@ func (e *Engine) Run(addr string) error { } return err case sig := <-quit: - log.Printf(Yellow+"Received signal: %v. Initiating graceful shutdown...\n"+Reset, sig) + log.Printf(Yellow+"Shutdown signal received: %v\n"+Reset, sig) return e.shutdown(srv) } } -// shutdown handles graceful HTTP server shutdown and executes registered shutdown hooks. +// shutdown ensures shutdown is performed only once and delegates to shutdownInternal. func (e *Engine) shutdown(srv *http.Server) error { + var err error + e.shutdownOnce.Do(func() { + err = e.shutdownInternal(srv) + }) + return err +} + +// shutdownInternal handles graceful HTTP server shutdown and executes registered shutdown hooks. +func (e *Engine) shutdownInternal(srv *http.Server) error { e.lifecycle.mu.Lock() timeout := e.lifecycle.timeout hooks := make([]lifecycleHook, len(e.lifecycle.shutdownHooks)) @@ -145,6 +159,8 @@ func (e *Engine) shutdown(srv *http.Server) error { shutdownErrors = append(shutdownErrors, fmt.Errorf("server shutdown failed: %w", err)) } + log.Println("Running shutdown hooks...") + // Step 3: Sort shutdown hooks by priority (descending), preserving registration order for equal priorities. sort.SliceStable(hooks, func(i, j int) bool { if hooks[i].priority != hooks[j].priority { @@ -165,9 +181,11 @@ func (e *Engine) shutdown(srv *http.Server) error { } if len(shutdownErrors) > 0 { + log.Println("Shutdown complete with errors") return &ShutdownError{Errors: shutdownErrors} } + log.Println("Shutdown complete") return nil } diff --git a/lifecycle_test.go b/lifecycle_test.go index 3f00098..3ae5ee2 100644 --- a/lifecycle_test.go +++ b/lifecycle_test.go @@ -221,7 +221,7 @@ func TestSignalHandling(t *testing.T) { // Run the server in a goroutine errChan := make(chan error, 1) go func() { - errChan <- app.Run(":18080") + errChan <- app.Run(":0") }() // Wait a moment for server to start listening @@ -252,3 +252,52 @@ func TestSignalHandling(t *testing.T) { t.Fatal("expected shutdown hooks to be called on signal") } } + +func TestShutdownIdempotent(t *testing.T) { + app := NewRouter() + callCount := 0 + + app.OnShutdown(func(ctx context.Context) error { + callCount++ + return nil + }) + + srv := &http.Server{} + + if err := app.shutdown(srv); err != nil { + t.Fatalf("unexpected first shutdown error: %v", err) + } + if callCount != 1 { + t.Fatalf("expected shutdown hook to run once, got %d", callCount) + } + + if err := app.shutdown(srv); err != nil { + t.Fatalf("unexpected second shutdown error: %v", err) + } + if callCount != 1 { + t.Fatalf("expected shutdown hook not to run again, got %d", callCount) + } +} + +func TestConcurrentHookRegistration(t *testing.T) { + app := NewRouter() + var wg sync.WaitGroup + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + app.OnStart(func() error { return nil }) + app.OnShutdown(func(ctx context.Context) error { return nil }) + }() + } + wg.Wait() + + if err := app.lifecycle.runStartupHooks(); err != nil { + t.Fatalf("unexpected startup hooks error: %v", err) + } + + srv := &http.Server{} + if err := app.shutdown(srv); err != nil { + t.Fatalf("unexpected shutdown error after concurrent registration: %v", err) + } +}