🔥 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>

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

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()
 
-	ln := fasthttputil.NewInmemoryListener()
+		// Fast request that should complete
+		app.Get("/", func(c Ctx) error {
+			return c.SendString("OK")
+		})
 
-	go func() {
-		err := app.Listener(ln)
-		assert.NoError(t, err)
-	}()
+		ln := fasthttputil.NewInmemoryListener()
+		serverStarted := make(chan bool, 1)
 
-	time.Sleep(1 * time.Second)
+		go func() {
+			serverStarted <- true
+			if err := app.Listener(ln); err != nil {
+				t.Errorf("Failed to start listener: %v", err)
+			}
+		}()
 
-	go func() {
+		<-serverStarted
+
+		// Execute normal request
 		conn, err := ln.Dial()
-		assert.NoError(t, err)
+		require.NoError(t, err)
+		_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"))
+		require.NoError(t, err)
 
-		_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n"))
-		assert.NoError(t, err)
-	}()
+		// Shutdown with sufficient timeout
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
 
-	time.Sleep(1 * time.Second)
+		err = app.ShutdownWithContext(ctx)
+		require.NoError(t, err, "Expected successful shutdown")
+	})
 
-	shutdownErr := make(chan error)
-	go func() {
-		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
-		defer cancel()
-		shutdownErr <- app.ShutdownWithContext(ctx)
-	}()
+	t.Run("shutdown with hooks", func(t *testing.T) {
+		t.Parallel()
+		app := New()
 
-	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)
+		hookOrder := make([]string, 0)
+		var hookMutex sync.Mutex
+
+		app.Hooks().OnPreShutdown(func() error {
+			hookMutex.Lock()
+			hookOrder = append(hookOrder, "pre")
+			hookMutex.Unlock()
+			return nil
+		})
+
+		app.Hooks().OnPostShutdown(func(_ error) error {
+			hookMutex.Lock()
+			hookOrder = append(hookOrder, "post")
+			hookMutex.Unlock()
+			return nil
+		})
+
+		ln := fasthttputil.NewInmemoryListener()
+		go func() {
+			if err := app.Listener(ln); err != nil {
+				t.Errorf("Failed to start listener: %v", err)
+			}
+		}()
+
+		time.Sleep(100 * time.Millisecond)
+
+		err := app.ShutdownWithContext(context.Background())
+		require.NoError(t, err)
+
+		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

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

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

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