🔥 feat: Improve and Optimize ShutdownWithContext Func (#3162)

* 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 <yingjie.huang@fosunhn.net> Co-authored-by: Juan Calderon-Perez <835733+gaby@users.noreply.github.com>
2025-02-23 00:32:51 +08:00 · 2025-02-23 00:32:51 +08:00 · 4b62d3d592
parent 252a0221a0
commit 4b62d3d592
9 changed files with 421 additions and 251 deletions
--- 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
--- 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
--- a/docs/api/fiber.md
+++ b/docs/api/fiber.md
@ -111,11 +111,9 @@ app.Listen(":8080", fiber.ListenConfig{
 | <Reference id="enableprefork">EnablePrefork</Reference>                 | `bool`                        | When set to true, this will spawn multiple Go processes listening on the same port.                                                           | `false` |
 | <Reference id="enableprintroutes">EnablePrintRoutes</Reference>         | `bool`                        | If set to true, will print all routes with their method, path, and handler.                                                                   | `false` |
 | <Reference id="gracefulcontext">GracefulContext</Reference>             | `context.Context`             | Field to shutdown Fiber by given context gracefully.                                                                                          | `nil`   |
-| <Reference id="ShutdownTimeout">ShutdownTimeout</Reference>             | `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`   |
+| <Reference id="ShutdownTimeout">ShutdownTimeout</Reference>             | `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`   |
 | <Reference id="listeneraddrfunc">ListenerAddrFunc</Reference>           | `func(addr net.Addr)`         | Allows accessing and customizing `net.Listener`.                                                                                              | `nil`   |
 | <Reference id="listenernetwork">ListenerNetwork</Reference>             | `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`  |
-| <Reference id="onshutdownerror">OnShutdownError</Reference>             | `func(err error)`             | Allows to customize error behavior when gracefully shutting down the server by given signal.  Prints error with `log.Fatalf()`                | `nil`   |
-| <Reference id="onshutdownsuccess">OnShutdownSuccess</Reference>         | `func()`                      | Allows customizing success behavior when gracefully shutting down the server by given signal.                                                 | `nil`   |
 | <Reference id="tlsconfigfunc">TLSConfigFunc</Reference>                 | `func(tlsConfig *tls.Config)` | Allows customizing `tls.Config` as you want.                                                                                                  | `nil`   |
 | <Reference id="autocertmanager">AutoCertManager</Reference>             | `*autocert.Manager`           | Manages TLS certificates automatically using the ACME protocol. Enables integration with Let's Encrypt or other ACME-compatible providers.    | `nil`   |
 | <Reference id="tlsminversion">TLSMinVersion</Reference>                 | `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
--- 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
--- 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
--- 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)
 		}
 	}
 }
--- 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) {
--- 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)
 }
--- 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()
 }