🔥 feat: Improve and Optimize ShutdownWithContext Func

app.go

@@ -900,16 +900,19 @@ 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()
+
 	if app.server == nil {
 		return ErrNotRunning
 	}
+
+	// Execute the Shutdown hook
+	if app.hooks != nil {
+		app.hooks.executeOnPreShutdownHooks()
+	}
+	defer app.hooks.executeOnPostShutdownHooks(nil)
+
 	return app.server.ShutdownWithContext(ctx)
 }
 

app_test.go

@@ -21,6 +21,7 @@ import (
 	"regexp"
 	"runtime"
 	"strings"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -861,45 +862,61 @@ func Test_App_ShutdownWithContext(t *testing.T) {
 	t.Parallel()
 
 	app := New()
+	var shutdownHookCalled atomic.Int32
+	// TODO: add test
+	// app.Hooks().OnShutdown(func() error {
+	// 	shutdownHookCalled.Store(1)
+	// 	return nil
+	// })
+
 	app.Get("/", func(ctx Ctx) error {
 		time.Sleep(5 * time.Second)
 		return ctx.SendString("body")
 	})
 
 	ln := fasthttputil.NewInmemoryListener()
 
+	serverErr := make(chan error, 1)
 	go func() {
-		err := app.Listener(ln)
-		assert.NoError(t, err)
+		serverErr <- app.Listener(ln)
 	}()
 
-	time.Sleep(1 * time.Second)
+	time.Sleep(100 * time.Millisecond)
 
+	clientDone := 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"))
+		_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"))
 		assert.NoError(t, err)
+		close(clientDone)
 	}()
 
-	time.Sleep(1 * time.Second)
+	<-clientDone
+	// Sleep to ensure the server has started processing the request
+	time.Sleep(100 * time.Millisecond)
 
-	shutdownErr := make(chan error)
+	shutdownErr := make(chan error, 1)
 	go func() {
 		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
 		defer cancel()
 		shutdownErr <- app.ShutdownWithContext(ctx)
 	}()
 
 	select {
-	case <-time.After(5 * time.Second):
-		t.Fatal("idle connections not closed on shutdown")
+	case <-time.After(2 * time.Second):
+		t.Fatal("shutdown did not complete in time")
 	case err := <-shutdownErr:
-		if err == nil || !errors.Is(err, context.DeadlineExceeded) {
-			t.Fatalf("unexpected err %v. Expecting %v", err, context.DeadlineExceeded)
-		}
+		require.Error(t, err, "Expected shutdown to return an error due to timeout")
+		require.ErrorIs(t, err, context.DeadlineExceeded, "Expected DeadlineExceeded error")
 	}
+
+	assert.Equal(t, int32(1), shutdownHookCalled.Load(), "Shutdown hook was not called")
+
+	err := <-serverErr
+	assert.NoError(t, err, "Server should have shut down without error")
+	// default:
+	// Server is still running, which is expected as the long-running request prevented full shutdown
 }
 
 // go test -run Test_App_Mixed_Routes_WithSameLen

docs/api/fiber.md

@@ -228,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

hooks.go

@@ -11,9 +11,11 @@ type (
 	OnGroupHandler     = func(Group) error
 	OnGroupNameHandler = OnGroupHandler
 	OnListenHandler    = func(ListenData) error
-	OnShutdownHandler  = func() error
-	OnForkHandler      = func(int) error
-	OnMountHandler     = func(*App) error
+	// OnShutdownHandler     = func() 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.
@@ -27,9 +29,11 @@ type Hooks struct {
 	onGroup     []OnGroupHandler
 	onGroupName []OnGroupNameHandler
 	onListen    []OnListenHandler
-	onShutdown  []OnShutdownHandler
-	onFork      []OnForkHandler
-	onMount     []OnMountHandler
+	// onShutdown     []OnShutdownHandler
+	onPreShutdown  []OnPreShutdownHandler
+	onPostShutdown []OnPostShutdownHandler
+	onFork         []OnForkHandler
+	onMount        []OnMountHandler
 }
 
 // ListenData is a struct to use it with OnListenHandler
@@ -47,9 +51,11 @@ func newHooks(app *App) *Hooks {
 		onGroupName: make([]OnGroupNameHandler, 0),
 		onName:      make([]OnNameHandler, 0),
 		onListen:    make([]OnListenHandler, 0),
-		onShutdown:  make([]OnShutdownHandler, 0),
-		onFork:      make([]OnForkHandler, 0),
-		onMount:     make([]OnMountHandler, 0),
+		// onShutdown:     make([]OnShutdownHandler, 0),
+		onPreShutdown:  make([]OnPreShutdownHandler, 0),
+		onPostShutdown: make([]OnPostShutdownHandler, 0),
+		onFork:         make([]OnForkHandler, 0),
+		onMount:        make([]OnMountHandler, 0),
 	}
 }
 
@@ -96,10 +102,25 @@ func (h *Hooks) OnListen(handler ...OnListenHandler) {
 	h.app.mutex.Unlock()
 }
 
+// TODO：To be deleted, replaced by OnPreShutdown and OnPostShutdown
 // OnShutdown is a hook to execute user functions after Shutdown.
-func (h *Hooks) OnShutdown(handler ...OnShutdownHandler) {
+// func (h *Hooks) OnShutdown(handler ...OnShutdownHandler) {
+// 	h.app.mutex.Lock()
+// 	h.onShutdown = append(h.onShutdown, handler...)
+// 	h.app.mutex.Unlock()
+// }
+
+// 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 +212,26 @@ func (h *Hooks) executeOnListenHooks(listenData ListenData) error {
 	return nil
 }
 
-func (h *Hooks) executeOnShutdownHooks() {
-	for _, v := range h.onShutdown {
+// func (h *Hooks) executeOnShutdownHooks() {
+// 	for _, v := range h.onShutdown {
+// 		if err := v(); err != nil {
+// 			log.Errorf("failed to call shutdown hook: %v", err)
+// 		}
+// 	}
+// }
+
+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 pre shutdown hook: %v", err)
 		}
 	}
 }

hooks_test.go

@@ -188,12 +188,13 @@ func Test_Hook_OnShutdown(t *testing.T) {
 	buf := bytebufferpool.Get()
 	defer bytebufferpool.Put(buf)
 
-	app.Hooks().OnShutdown(func() error {
-		_, err := buf.WriteString("shutdowning")
-		require.NoError(t, err)
+	// TODO: add test
+	// app.Hooks().OnShutdown(func() error {
+	// 	_, err := buf.WriteString("shutdowning")
+	// 	require.NoError(t, err)
 
-		return nil
-	})
+	// 	return nil
+	// })
 
 	require.NoError(t, app.Shutdown())
 	require.Equal(t, "shutdowning", buf.String())