From 4b62d3d592248077c38f53d6117828f278e8eb3f Mon Sep 17 00:00:00 2001 From: JIeJaitt <498938874@qq.com> Date: Sun, 23 Feb 2025 00:32:51 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A5=20feat:=20Improve=20and=20Optimize?= =?UTF-8?q?=20ShutdownWithContext=20Func=20(#3162)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: Optimize ShutdownWithContext method in app.go - Reorder mutex lock acquisition to the start of the function - Early return if server is not running - Use defer for executing shutdown hooks - Simplify nil check for hooks - Remove TODO comment This commit improves the readability, robustness, and execution order of the shutdown process. It ensures consistent state throughout the shutdown and guarantees hook execution even in error cases. * feat: Enhance ShutdownWithContext test for improved reliability - Add shutdown hook verification - Implement better synchronization with channels - Improve error handling and assertions - Adjust timeouts for more consistent results - Add server state check after shutdown attempt - Include comments explaining expected behavior This commit improves the comprehensiveness and reliability of the ShutdownWithContext test, ensuring proper verification of shutdown hooks, timeout behavior, and server state during long-running requests. * πŸ“š Doc: update the docs to explain shutdown & hook execution order * 🩹 Fix: Possible Data Race on shutdownHookCalled Variable * 🩹 Fix: Remove the default Case * 🩹 Fix: Import sync/atomic * 🩹 Fix: golangci-lint problem * 🎨 Style: add block in api.md * 🩹 Fix: go mod tidy * feat: Optimize ShutdownWithContext method in app.go - Reorder mutex lock acquisition to the start of the function - Early return if server is not running - Use defer for executing shutdown hooks - Simplify nil check for hooks - Remove TODO comment This commit improves the readability, robustness, and execution order of the shutdown process. It ensures consistent state throughout the shutdown and guarantees hook execution even in error cases. * feat: Enhance ShutdownWithContext test for improved reliability - Add shutdown hook verification - Implement better synchronization with channels - Improve error handling and assertions - Adjust timeouts for more consistent results - Add server state check after shutdown attempt - Include comments explaining expected behavior This commit improves the comprehensiveness and reliability of the ShutdownWithContext test, ensuring proper verification of shutdown hooks, timeout behavior, and server state during long-running requests. * πŸ“š Doc: update the docs to explain shutdown & hook execution order * 🩹 Fix: Possible Data Race on shutdownHookCalled Variable * 🩹 Fix: Remove the default Case * 🩹 Fix: Import sync/atomic * 🩹 Fix: golangci-lint problem * 🎨 Style: add block in api.md * 🩹 Fix: go mod tidy * ♻️ Refactor: replaced OnShutdown by OnPreShutdown and OnPostShutdown * ♻️ Refactor: streamline post-shutdown hook execution in graceful shutdown process * 🚨 Test: add test for gracefulShutdown * πŸ”₯ Feature: Using executeOnPreShutdownHooks and executeOnPostShutdownHooks Instead of OnShutdownSuccess and OnShutdownError * 🩹 Fix: deal Listener err * 🩹 Fix: go lint error * 🩹 Fix: reduced memory alignment * 🩹 Fix: reduced memory alignment * 🩹 Fix: context should be created inside the concatenation. * πŸ“š Doc: update what_new.md and hooks.md * ♻️ Refactor: use blocking channel instead of time.Sleep * 🩹 Fix: Improve synchronization in error propagation test. * 🩹 Fix: Replace sleep with proper synchronization. * 🩹 Fix: Server but not shut down properly * 🩹 Fix: Using channels to synchronize and pass results * 🩹 Fix: timeout with long running request * πŸ“š Doc: remove OnShutdownError and OnShutdownSuccess from fiber.md * Update hooks.md * 🚨 Test: Add graceful shutdown timeout error test case * πŸ“ Doc: Restructure hooks documentation for OnPreShutdown and OnPostShutdown * πŸ“ Doc: Remove extra whitespace in hooks documentation --------- Co-authored-by: yingjie.huang Co-authored-by: Juan Calderon-Perez <835733+gaby@users.noreply.github.com> --- app.go | 23 +++-- app_test.go | 158 +++++++++++++++++++++++++------- docs/api/fiber.md | 6 +- docs/api/hooks.md | 20 +++-- docs/whats_new.md | 59 ++++++++++++ hooks.go | 80 ++++++++++------- hooks_test.go | 75 +++++++++++++++- listen.go | 28 +----- listen_test.go | 223 ++++++++++++++++------------------------------ 9 files changed, 421 insertions(+), 251 deletions(-) diff --git a/app.go b/app.go index c867cc4e..4bbd9efb 100644 --- a/app.go +++ b/app.go @@ -894,6 +894,13 @@ func (app *App) HandlersCount() uint32 { // // Make sure the program doesn't exit and waits instead for Shutdown to return. // +// Important: app.Listen() must be called in a separate goroutine, otherwise shutdown hooks will not work +// as Listen() is a blocking operation. Example: +// +// go app.Listen(":3000") +// // ... +// app.Shutdown() +// // Shutdown does not close keepalive connections so its recommended to set ReadTimeout to something else than 0. func (app *App) Shutdown() error { return app.ShutdownWithContext(context.Background()) @@ -918,17 +925,21 @@ func (app *App) ShutdownWithTimeout(timeout time.Duration) error { // // ShutdownWithContext does not close keepalive connections so its recommended to set ReadTimeout to something else than 0. func (app *App) ShutdownWithContext(ctx context.Context) error { - if app.hooks != nil { - // TODO: check should be defered? - app.hooks.executeOnShutdownHooks() - } - app.mutex.Lock() defer app.mutex.Unlock() + + var err error + if app.server == nil { return ErrNotRunning } - return app.server.ShutdownWithContext(ctx) + + // Execute the Shutdown hook + app.hooks.executeOnPreShutdownHooks() + defer app.hooks.executeOnPostShutdownHooks(err) + + err = app.server.ShutdownWithContext(ctx) + return err } // Server returns the underlying fasthttp server diff --git a/app_test.go b/app_test.go index d64f9a68..18f6cef9 100644 --- a/app_test.go +++ b/app_test.go @@ -21,6 +21,7 @@ import ( "regexp" "runtime" "strings" + "sync" "testing" "time" @@ -930,20 +931,29 @@ func Test_App_ShutdownWithTimeout(t *testing.T) { }) ln := fasthttputil.NewInmemoryListener() + serverReady := make(chan struct{}) // Signal that the server is ready to start + go func() { + serverReady <- struct{}{} err := app.Listener(ln) assert.NoError(t, err) }() - time.Sleep(1 * time.Second) + <-serverReady // Waiting for the server to be ready + + // Create a connection and send a request + connReady := make(chan struct{}) go func() { conn, err := ln.Dial() assert.NoError(t, err) _, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n")) assert.NoError(t, err) + + connReady <- struct{}{} // Signal that the request has been sent }() - time.Sleep(1 * time.Second) + + <-connReady // Waiting for the request to be sent shutdownErr := make(chan error) go func() { @@ -964,46 +974,130 @@ func Test_App_ShutdownWithTimeout(t *testing.T) { func Test_App_ShutdownWithContext(t *testing.T) { t.Parallel() - app := New() - app.Get("/", func(ctx Ctx) error { - time.Sleep(5 * time.Second) - return ctx.SendString("body") + t.Run("successful shutdown", func(t *testing.T) { + t.Parallel() + app := New() + + // Fast request that should complete + app.Get("/", func(c Ctx) error { + return c.SendString("OK") + }) + + ln := fasthttputil.NewInmemoryListener() + serverStarted := make(chan bool, 1) + + go func() { + serverStarted <- true + if err := app.Listener(ln); err != nil { + t.Errorf("Failed to start listener: %v", err) + } + }() + + <-serverStarted + + // Execute normal request + conn, err := ln.Dial() + require.NoError(t, err) + _, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")) + require.NoError(t, err) + + // Shutdown with sufficient timeout + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + err = app.ShutdownWithContext(ctx) + require.NoError(t, err, "Expected successful shutdown") }) - ln := fasthttputil.NewInmemoryListener() + t.Run("shutdown with hooks", func(t *testing.T) { + t.Parallel() + app := New() - go func() { - err := app.Listener(ln) - assert.NoError(t, err) - }() + hookOrder := make([]string, 0) + var hookMutex sync.Mutex - time.Sleep(1 * time.Second) + app.Hooks().OnPreShutdown(func() error { + hookMutex.Lock() + hookOrder = append(hookOrder, "pre") + hookMutex.Unlock() + return nil + }) - go func() { - conn, err := ln.Dial() - assert.NoError(t, err) + app.Hooks().OnPostShutdown(func(_ error) error { + hookMutex.Lock() + hookOrder = append(hookOrder, "post") + hookMutex.Unlock() + return nil + }) - _, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n")) - assert.NoError(t, err) - }() + ln := fasthttputil.NewInmemoryListener() + go func() { + if err := app.Listener(ln); err != nil { + t.Errorf("Failed to start listener: %v", err) + } + }() - time.Sleep(1 * time.Second) + time.Sleep(100 * time.Millisecond) - shutdownErr := make(chan error) - go func() { - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) - defer cancel() - shutdownErr <- app.ShutdownWithContext(ctx) - }() + err := app.ShutdownWithContext(context.Background()) + require.NoError(t, err) - select { - case <-time.After(5 * time.Second): - t.Fatal("idle connections not closed on shutdown") - case err := <-shutdownErr: - if err == nil || !errors.Is(err, context.DeadlineExceeded) { - t.Fatalf("unexpected err %v. Expecting %v", err, context.DeadlineExceeded) + require.Equal(t, []string{"pre", "post"}, hookOrder, "Hooks should execute in order") + }) + + t.Run("timeout with long running request", func(t *testing.T) { + t.Parallel() + app := New() + + requestStarted := make(chan struct{}) + requestProcessing := make(chan struct{}) + + app.Get("/", func(c Ctx) error { + close(requestStarted) + // Wait for signal to continue processing the request + <-requestProcessing + time.Sleep(2 * time.Second) + return c.SendString("OK") + }) + + ln := fasthttputil.NewInmemoryListener() + go func() { + if err := app.Listener(ln); err != nil { + t.Errorf("Failed to start listener: %v", err) + } + }() + + // Ensure server is fully started + time.Sleep(100 * time.Millisecond) + + // Start a long-running request + go func() { + conn, err := ln.Dial() + if err != nil { + t.Errorf("Failed to dial: %v", err) + return + } + if _, err := conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")); err != nil { + t.Errorf("Failed to write: %v", err) + } + }() + + // Wait for request to start + select { + case <-requestStarted: + // Request has started, signal to continue processing + close(requestProcessing) + case <-time.After(2 * time.Second): + t.Fatal("Request did not start in time") } - } + + // Attempt shutdown, should timeout + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + + err := app.ShutdownWithContext(ctx) + require.ErrorIs(t, err, context.DeadlineExceeded) + }) } // go test -run Test_App_Mixed_Routes_WithSameLen diff --git a/docs/api/fiber.md b/docs/api/fiber.md index 70320984..a79b2ba0 100644 --- a/docs/api/fiber.md +++ b/docs/api/fiber.md @@ -111,11 +111,9 @@ app.Listen(":8080", fiber.ListenConfig{ | EnablePrefork | `bool` | When set to true, this will spawn multiple Go processes listening on the same port. | `false` | | EnablePrintRoutes | `bool` | If set to true, will print all routes with their method, path, and handler. | `false` | | GracefulContext | `context.Context` | Field to shutdown Fiber by given context gracefully. | `nil` | -| ShutdownTimeout | `time.Duration` | Specifies the maximum duration to wait for the server to gracefully shutdown. When the timeout is reached, the graceful shutdown process is interrupted and forcibly terminated, and the `context.DeadlineExceeded` error is passed to the `OnShutdownError` callback. Set to 0 to disable the timeout and wait indefinitely. | `10 * time.Second` | +| ShutdownTimeout | `time.Duration` | Specifies the maximum duration to wait for the server to gracefully shutdown. When the timeout is reached, the graceful shutdown process is interrupted and forcibly terminated, and the `context.DeadlineExceeded` error is passed to the `OnPostShutdown` callback. Set to 0 to disable the timeout and wait indefinitely. | `10 * time.Second` | | ListenerAddrFunc | `func(addr net.Addr)` | Allows accessing and customizing `net.Listener`. | `nil` | | ListenerNetwork | `string` | Known networks are "tcp", "tcp4" (IPv4-only), "tcp6" (IPv6-only). WARNING: When prefork is set to true, only "tcp4" and "tcp6" can be chosen. | `tcp4` | -| OnShutdownError | `func(err error)` | Allows to customize error behavior when gracefully shutting down the server by given signal. Prints error with `log.Fatalf()` | `nil` | -| OnShutdownSuccess | `func()` | Allows customizing success behavior when gracefully shutting down the server by given signal. | `nil` | | TLSConfigFunc | `func(tlsConfig *tls.Config)` | Allows customizing `tls.Config` as you want. | `nil` | | AutoCertManager | `*autocert.Manager` | Manages TLS certificates automatically using the ACME protocol. Enables integration with Let's Encrypt or other ACME-compatible providers. | `nil` | | TLSMinVersion | `uint16` | Allows customizing the TLS minimum version. | `tls.VersionTLS12` | @@ -230,7 +228,7 @@ Shutdown gracefully shuts down the server without interrupting any active connec ShutdownWithTimeout will forcefully close any active connections after the timeout expires. -ShutdownWithContext shuts down the server including by force if the context's deadline is exceeded. +ShutdownWithContext shuts down the server including by force if the context's deadline is exceeded. Shutdown hooks will still be executed, even if an error occurs during the shutdown process, as they are deferred to ensure cleanup happens regardless of errors. ```go func (app *App) Shutdown() error diff --git a/docs/api/hooks.md b/docs/api/hooks.md index 48528666..a6e6c1ac 100644 --- a/docs/api/hooks.md +++ b/docs/api/hooks.md @@ -15,7 +15,8 @@ With Fiber you can execute custom user functions at specific method execution po - [OnGroupName](#ongroupname) - [OnListen](#onlisten) - [OnFork](#onfork) -- [OnShutdown](#onshutdown) +- [OnPreShutdown](#onpreshutdown) +- [OnPostShutdown](#onpostshutdown) - [OnMount](#onmount) ## Constants @@ -28,7 +29,8 @@ type OnGroupHandler = func(Group) error type OnGroupNameHandler = OnGroupHandler type OnListenHandler = func(ListenData) error type OnForkHandler = func(int) error -type OnShutdownHandler = func() error +type OnPreShutdownHandler = func() error +type OnPostShutdownHandler = func(error) error type OnMountHandler = func(*App) error ``` @@ -174,12 +176,20 @@ func main() { func (h *Hooks) OnFork(handler ...OnForkHandler) ``` -## OnShutdown +## OnPreShutdown -`OnShutdown` is a hook to execute user functions after shutdown. +`OnPreShutdown` is a hook to execute user functions before shutdown. ```go title="Signature" -func (h *Hooks) OnShutdown(handler ...OnShutdownHandler) +func (h *Hooks) OnPreShutdown(handler ...OnPreShutdownHandler) +``` + +## OnPostShutdown + +`OnPostShutdown` is a hook to execute user functions after shutdown. + +```go title="Signature" +func (h *Hooks) OnPostShutdown(handler ...OnPostShutdownHandler) ``` ## OnMount diff --git a/docs/whats_new.md b/docs/whats_new.md index 1958632a..7f0b6322 100644 --- a/docs/whats_new.md +++ b/docs/whats_new.md @@ -16,6 +16,8 @@ In this guide, we'll walk you through the most important changes in Fiber `v3` a Here's a quick overview of the changes in Fiber `v3`: - [πŸš€ App](#-app) +- [🎣 Hooks](#-hooks) +- [πŸš€ Listen](#-listen) - [πŸ—ΊοΈ Router](#-router) - [🧠 Context](#-context) - [πŸ“Ž Binding](#-binding) @@ -158,6 +160,63 @@ app.Listen(":444", fiber.ListenConfig{ }) ``` +## 🎣 Hooks + +We have made several changes to the Fiber hooks, including: + +- Added new shutdown hooks to provide better control over the shutdown process: + - `OnPreShutdown` - Executes before the server starts shutting down + - `OnPostShutdown` - Executes after the server has shut down, receives any shutdown error +- Deprecated `OnShutdown` in favor of the new pre/post shutdown hooks +- Improved shutdown hook execution order and reliability +- Added mutex protection for hook registration and execution + +Important: When using shutdown hooks, ensure app.Listen() is called in a separate goroutine: + +```go +// Correct usage +go app.Listen(":3000") +// ... register shutdown hooks +app.Shutdown() + +// Incorrect usage - hooks won't work +app.Listen(":3000") // This blocks +app.Shutdown() // Never reached +``` + +## πŸš€ Listen + +We have made several changes to the Fiber listen, including: + +- Removed `OnShutdownError` and `OnShutdownSuccess` from `ListenerConfig` in favor of using `OnPostShutdown` hook which receives the shutdown error + +```go +app := fiber.New() + +// Before - using ListenerConfig callbacks +app.Listen(":3000", fiber.ListenerConfig{ + OnShutdownError: func(err error) { + log.Printf("Shutdown error: %v", err) + }, + OnShutdownSuccess: func() { + log.Println("Shutdown successful") + }, +}) + +// After - using OnPostShutdown hook +app.Hooks().OnPostShutdown(func(err error) error { + if err != nil { + log.Printf("Shutdown error: %v", err) + } else { + log.Println("Shutdown successful") + } + return nil +}) +go app.Listen(":3000") +``` + +This change simplifies the shutdown handling by consolidating the shutdown callbacks into a single hook that receives the error status. + ## πŸ—Ί Router We have slightly adapted our router interface diff --git a/hooks.go b/hooks.go index 3da5c671..314717d0 100644 --- a/hooks.go +++ b/hooks.go @@ -6,14 +6,15 @@ import ( // OnRouteHandler Handlers define a function to create hooks for Fiber. type ( - OnRouteHandler = func(Route) error - OnNameHandler = OnRouteHandler - OnGroupHandler = func(Group) error - OnGroupNameHandler = OnGroupHandler - OnListenHandler = func(ListenData) error - OnShutdownHandler = func() error - OnForkHandler = func(int) error - OnMountHandler = func(*App) error + OnRouteHandler = func(Route) error + OnNameHandler = OnRouteHandler + OnGroupHandler = func(Group) error + OnGroupNameHandler = OnGroupHandler + OnListenHandler = func(ListenData) error + OnPreShutdownHandler = func() error + OnPostShutdownHandler = func(error) error + OnForkHandler = func(int) error + OnMountHandler = func(*App) error ) // Hooks is a struct to use it with App. @@ -22,14 +23,15 @@ type Hooks struct { app *App // Hooks - onRoute []OnRouteHandler - onName []OnNameHandler - onGroup []OnGroupHandler - onGroupName []OnGroupNameHandler - onListen []OnListenHandler - onShutdown []OnShutdownHandler - onFork []OnForkHandler - onMount []OnMountHandler + onRoute []OnRouteHandler + onName []OnNameHandler + onGroup []OnGroupHandler + onGroupName []OnGroupNameHandler + onListen []OnListenHandler + onPreShutdown []OnPreShutdownHandler + onPostShutdown []OnPostShutdownHandler + onFork []OnForkHandler + onMount []OnMountHandler } // ListenData is a struct to use it with OnListenHandler @@ -41,15 +43,16 @@ type ListenData struct { func newHooks(app *App) *Hooks { return &Hooks{ - app: app, - onRoute: make([]OnRouteHandler, 0), - onGroup: make([]OnGroupHandler, 0), - onGroupName: make([]OnGroupNameHandler, 0), - onName: make([]OnNameHandler, 0), - onListen: make([]OnListenHandler, 0), - onShutdown: make([]OnShutdownHandler, 0), - onFork: make([]OnForkHandler, 0), - onMount: make([]OnMountHandler, 0), + app: app, + onRoute: make([]OnRouteHandler, 0), + onGroup: make([]OnGroupHandler, 0), + onGroupName: make([]OnGroupNameHandler, 0), + onName: make([]OnNameHandler, 0), + onListen: make([]OnListenHandler, 0), + onPreShutdown: make([]OnPreShutdownHandler, 0), + onPostShutdown: make([]OnPostShutdownHandler, 0), + onFork: make([]OnForkHandler, 0), + onMount: make([]OnMountHandler, 0), } } @@ -96,10 +99,17 @@ func (h *Hooks) OnListen(handler ...OnListenHandler) { h.app.mutex.Unlock() } -// OnShutdown is a hook to execute user functions after Shutdown. -func (h *Hooks) OnShutdown(handler ...OnShutdownHandler) { +// OnPreShutdown is a hook to execute user functions before Shutdown. +func (h *Hooks) OnPreShutdown(handler ...OnPreShutdownHandler) { h.app.mutex.Lock() - h.onShutdown = append(h.onShutdown, handler...) + h.onPreShutdown = append(h.onPreShutdown, handler...) + h.app.mutex.Unlock() +} + +// OnPostShutdown is a hook to execute user functions after Shutdown. +func (h *Hooks) OnPostShutdown(handler ...OnPostShutdownHandler) { + h.app.mutex.Lock() + h.onPostShutdown = append(h.onPostShutdown, handler...) h.app.mutex.Unlock() } @@ -191,10 +201,18 @@ func (h *Hooks) executeOnListenHooks(listenData ListenData) error { return nil } -func (h *Hooks) executeOnShutdownHooks() { - for _, v := range h.onShutdown { +func (h *Hooks) executeOnPreShutdownHooks() { + for _, v := range h.onPreShutdown { if err := v(); err != nil { - log.Errorf("failed to call shutdown hook: %v", err) + log.Errorf("failed to call pre shutdown hook: %v", err) + } + } +} + +func (h *Hooks) executeOnPostShutdownHooks(err error) { + for _, v := range h.onPostShutdown { + if err := v(err); err != nil { + log.Errorf("failed to call post shutdown hook: %v", err) } } } diff --git a/hooks_test.go b/hooks_test.go index f96f5707..b95e0733 100644 --- a/hooks_test.go +++ b/hooks_test.go @@ -181,22 +181,89 @@ func Test_Hook_OnGroupName_Error(t *testing.T) { grp.Get("/test", testSimpleHandler) } -func Test_Hook_OnShutdown(t *testing.T) { +func Test_Hook_OnPrehutdown(t *testing.T) { t.Parallel() app := New() buf := bytebufferpool.Get() defer bytebufferpool.Put(buf) - app.Hooks().OnShutdown(func() error { - _, err := buf.WriteString("shutdowning") + app.Hooks().OnPreShutdown(func() error { + _, err := buf.WriteString("pre-shutdowning") require.NoError(t, err) return nil }) require.NoError(t, app.Shutdown()) - require.Equal(t, "shutdowning", buf.String()) + require.Equal(t, "pre-shutdowning", buf.String()) +} + +func Test_Hook_OnPostShutdown(t *testing.T) { + t.Run("should execute post shutdown hook with error", func(t *testing.T) { + app := New() + expectedErr := errors.New("test shutdown error") + + hookCalled := make(chan error, 1) + defer close(hookCalled) + + app.Hooks().OnPostShutdown(func(err error) error { + hookCalled <- err + return nil + }) + + go func() { + if err := app.Listen(":0"); err != nil { + return + } + }() + + time.Sleep(100 * time.Millisecond) + + app.hooks.executeOnPostShutdownHooks(expectedErr) + + select { + case err := <-hookCalled: + require.Equal(t, expectedErr, err) + case <-time.After(time.Second): + t.Fatal("hook execution timeout") + } + + require.NoError(t, app.Shutdown()) + }) + + t.Run("should execute multiple hooks in order", func(t *testing.T) { + app := New() + + execution := make([]int, 0) + + app.Hooks().OnPostShutdown(func(_ error) error { + execution = append(execution, 1) + return nil + }) + + app.Hooks().OnPostShutdown(func(_ error) error { + execution = append(execution, 2) + return nil + }) + + app.hooks.executeOnPostShutdownHooks(nil) + + require.Len(t, execution, 2, "expected 2 hooks to execute") + require.Equal(t, []int{1, 2}, execution, "hooks executed in wrong order") + }) + + t.Run("should handle hook error", func(_ *testing.T) { + app := New() + hookErr := errors.New("hook error") + + app.Hooks().OnPostShutdown(func(_ error) error { + return hookErr + }) + + // Should not panic + app.hooks.executeOnPostShutdownHooks(nil) + }) } func Test_Hook_OnListen(t *testing.T) { diff --git a/listen.go b/listen.go index 793d36d2..f33c9daf 100644 --- a/listen.go +++ b/listen.go @@ -60,17 +60,6 @@ type ListenConfig struct { // Default: nil BeforeServeFunc func(app *App) error `json:"before_serve_func"` - // OnShutdownError allows to customize error behavior when to graceful shutdown server by given signal. - // - // Print error with log.Fatalf() by default. - // Default: nil - OnShutdownError func(err error) - - // OnShutdownSuccess allows to customize success behavior when to graceful shutdown server by given signal. - // - // Default: nil - OnShutdownSuccess func() - // AutoCertManager manages TLS certificates automatically using the ACME protocol, // Enables integration with Let's Encrypt or other ACME-compatible providers. // @@ -102,7 +91,7 @@ type ListenConfig struct { CertClientFile string `json:"cert_client_file"` // When the graceful shutdown begins, use this field to set the timeout - // duration. If the timeout is reached, OnShutdownError will be called. + // duration. If the timeout is reached, OnPostShutdown will be called with the error. // Set to 0 to disable the timeout and wait indefinitely. // // Default: 10 * time.Second @@ -136,9 +125,6 @@ func listenConfigDefault(config ...ListenConfig) ListenConfig { return ListenConfig{ TLSMinVersion: tls.VersionTLS12, ListenerNetwork: NetworkTCP4, - OnShutdownError: func(err error) { - log.Fatalf("shutdown: %v", err) //nolint:revive // It's an option - }, ShutdownTimeout: 10 * time.Second, } } @@ -148,12 +134,6 @@ func listenConfigDefault(config ...ListenConfig) ListenConfig { cfg.ListenerNetwork = NetworkTCP4 } - if cfg.OnShutdownError == nil { - cfg.OnShutdownError = func(err error) { - log.Fatalf("shutdown: %v", err) //nolint:revive // It's an option - } - } - if cfg.TLSMinVersion == 0 { cfg.TLSMinVersion = tls.VersionTLS12 } @@ -517,11 +497,9 @@ func (app *App) gracefulShutdown(ctx context.Context, cfg ListenConfig) { } if err != nil { - cfg.OnShutdownError(err) + app.hooks.executeOnPostShutdownHooks(err) return } - if success := cfg.OnShutdownSuccess; success != nil { - success() - } + app.hooks.executeOnPostShutdownHooks(nil) } diff --git a/listen_test.go b/listen_test.go index 032f7d32..1a5bd77f 100644 --- a/listen_test.go +++ b/listen_test.go @@ -15,6 +15,7 @@ import ( "testing" "time" + "github.com/gofiber/utils/v2" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/valyala/fasthttp" @@ -37,98 +38,42 @@ func Test_Listen(t *testing.T) { // go test -run Test_Listen_Graceful_Shutdown func Test_Listen_Graceful_Shutdown(t *testing.T) { - var mu sync.Mutex - var shutdown bool - - app := New() - - app.Get("/", func(c Ctx) error { - return c.SendString(c.Hostname()) + t.Run("Basic Graceful Shutdown", func(t *testing.T) { + testGracefulShutdown(t, 0) }) - ln := fasthttputil.NewInmemoryListener() - errs := make(chan error) + t.Run("Shutdown With Timeout", func(t *testing.T) { + testGracefulShutdown(t, 500*time.Millisecond) + }) - go func() { - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - - errs <- app.Listener(ln, ListenConfig{ - DisableStartupMessage: true, - GracefulContext: ctx, - OnShutdownSuccess: func() { - mu.Lock() - shutdown = true - mu.Unlock() - }, - }) - }() - - // Server readiness check - for i := 0; i < 10; i++ { - conn, err := ln.Dial() - if err == nil { - conn.Close() //nolint:errcheck // ignore error - break - } - // Wait a bit before retrying - time.Sleep(100 * time.Millisecond) - if i == 9 { - t.Fatalf("Server did not become ready in time: %v", err) - } - } - - testCases := []struct { - ExpectedErr error - ExpectedBody string - Time time.Duration - ExpectedStatusCode int - }{ - {Time: 500 * time.Millisecond, ExpectedBody: "example.com", ExpectedStatusCode: StatusOK, ExpectedErr: nil}, - {Time: 3 * time.Second, ExpectedBody: "", ExpectedStatusCode: StatusOK, ExpectedErr: fasthttputil.ErrInmemoryListenerClosed}, - } - - for _, tc := range testCases { - time.Sleep(tc.Time) - - req := fasthttp.AcquireRequest() - req.SetRequestURI("http://example.com") - - client := fasthttp.HostClient{} - client.Dial = func(_ string) (net.Conn, error) { return ln.Dial() } - - resp := fasthttp.AcquireResponse() - err := client.Do(req, resp) - - require.Equal(t, tc.ExpectedErr, err) - require.Equal(t, tc.ExpectedStatusCode, resp.StatusCode()) - require.Equal(t, tc.ExpectedBody, string(resp.Body())) - - fasthttp.ReleaseRequest(req) - fasthttp.ReleaseResponse(resp) - } - - mu.Lock() - err := <-errs - require.True(t, shutdown) - require.NoError(t, err) - mu.Unlock() + t.Run("Shutdown With Timeout Error", func(t *testing.T) { + testGracefulShutdown(t, 1*time.Nanosecond) + }) } -// go test -run Test_Listen_Graceful_Shutdown_Timeout -func Test_Listen_Graceful_Shutdown_Timeout(t *testing.T) { +func testGracefulShutdown(t *testing.T, shutdownTimeout time.Duration) { + t.Helper() + var mu sync.Mutex - var shutdownSuccess bool - var shutdownTimeoutError error + var shutdown bool + var receivedErr error app := New() - app.Get("/", func(c Ctx) error { + time.Sleep(10 * time.Millisecond) return c.SendString(c.Hostname()) }) ln := fasthttputil.NewInmemoryListener() - errs := make(chan error) + errs := make(chan error, 1) + + app.hooks.OnPostShutdown(func(err error) error { + mu.Lock() + defer mu.Unlock() + shutdown = true + receivedErr = err + return nil + }) go func() { ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) @@ -137,93 +82,83 @@ func Test_Listen_Graceful_Shutdown_Timeout(t *testing.T) { errs <- app.Listener(ln, ListenConfig{ DisableStartupMessage: true, GracefulContext: ctx, - ShutdownTimeout: 500 * time.Millisecond, - OnShutdownSuccess: func() { - mu.Lock() - shutdownSuccess = true - mu.Unlock() - }, - OnShutdownError: func(err error) { - mu.Lock() - shutdownTimeoutError = err - mu.Unlock() - }, + ShutdownTimeout: shutdownTimeout, }) }() - // Server readiness check - for i := 0; i < 10; i++ { + require.Eventually(t, func() bool { conn, err := ln.Dial() - // To test a graceful shutdown timeout, do not close the connection. if err == nil { - _ = conn - break - } - // Wait a bit before retrying - time.Sleep(100 * time.Millisecond) - if i == 9 { - t.Fatalf("Server did not become ready in time: %v", err) + if err := conn.Close(); err != nil { + t.Logf("error closing connection: %v", err) + } + return true } + return false + }, time.Second, 100*time.Millisecond, "Server failed to become ready") + + client := fasthttp.HostClient{ + Dial: func(_ string) (net.Conn, error) { return ln.Dial() }, } - testCases := []struct { - ExpectedErr error - ExpectedShutdownError error - ExpectedBody string - Time time.Duration - ExpectedStatusCode int - ExpectedShutdownSuccess bool - }{ + type testCase struct { + expectedErr error + expectedBody string + name string + waitTime time.Duration + expectedStatusCode int + closeConnection bool + } + + testCases := []testCase{ { - Time: 100 * time.Millisecond, - ExpectedBody: "example.com", - ExpectedStatusCode: StatusOK, - ExpectedErr: nil, - ExpectedShutdownError: nil, - ExpectedShutdownSuccess: false, + name: "Server running normally", + waitTime: 500 * time.Millisecond, + expectedBody: "example.com", + expectedStatusCode: StatusOK, + expectedErr: nil, + closeConnection: true, }, { - Time: 3 * time.Second, - ExpectedBody: "", - ExpectedStatusCode: StatusOK, - ExpectedErr: fasthttputil.ErrInmemoryListenerClosed, - ExpectedShutdownError: context.DeadlineExceeded, - ExpectedShutdownSuccess: false, + name: "Server shutdown complete", + waitTime: 3 * time.Second, + expectedBody: "", + expectedStatusCode: StatusOK, + expectedErr: fasthttputil.ErrInmemoryListenerClosed, + closeConnection: true, }, } for _, tc := range testCases { - time.Sleep(tc.Time) + t.Run(tc.name, func(t *testing.T) { + time.Sleep(tc.waitTime) - req := fasthttp.AcquireRequest() - req.SetRequestURI("http://example.com") + req := fasthttp.AcquireRequest() + defer fasthttp.ReleaseRequest(req) + req.SetRequestURI("http://example.com") - client := fasthttp.HostClient{} - client.Dial = func(_ string) (net.Conn, error) { return ln.Dial() } + resp := fasthttp.AcquireResponse() + defer fasthttp.ReleaseResponse(resp) - resp := fasthttp.AcquireResponse() - err := client.Do(req, resp) + err := client.Do(req, resp) - if err == nil { - require.NoError(t, err) - require.Equal(t, tc.ExpectedStatusCode, resp.StatusCode()) - require.Equal(t, tc.ExpectedBody, string(resp.Body())) - } else { - require.ErrorIs(t, err, tc.ExpectedErr) - } - - mu.Lock() - require.Equal(t, tc.ExpectedShutdownSuccess, shutdownSuccess) - require.Equal(t, tc.ExpectedShutdownError, shutdownTimeoutError) - mu.Unlock() - - fasthttp.ReleaseRequest(req) - fasthttp.ReleaseResponse(resp) + if tc.expectedErr == nil { + require.NoError(t, err) + require.Equal(t, tc.expectedStatusCode, resp.StatusCode()) + require.Equal(t, tc.expectedBody, utils.UnsafeString(resp.Body())) + } else { + require.ErrorIs(t, err, tc.expectedErr) + } + }) } mu.Lock() - err := <-errs - require.NoError(t, err) + require.True(t, shutdown) + if shutdownTimeout == 1*time.Nanosecond { + require.Error(t, receivedErr) + require.ErrorIs(t, receivedErr, context.DeadlineExceeded) + } + require.NoError(t, <-errs) mu.Unlock() }