🔥 Feature: Using executeOnPreShutdownHooks and executeOnPostShutdownH…

…ooks Instead of OnShutdownSuccess and OnShutdownError

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())
@@ -921,16 +928,18 @@ func (app *App) ShutdownWithContext(ctx context.Context) error {
 	app.mutex.Lock()
 	defer app.mutex.Unlock()
 
+	var err error
+
 	if app.server == nil {
 		return ErrNotRunning
 	}
 
 	// Execute the Shutdown hook
-	if app.hooks != nil {
-		app.hooks.executeOnPreShutdownHooks()
-	}
+	app.hooks.executeOnPreShutdownHooks()
+	defer app.hooks.executeOnPostShutdownHooks(err)
 
-	return app.server.ShutdownWithContext(ctx)
+	err = app.server.ShutdownWithContext(ctx)
+	return err
 }
 
 // Server returns the underlying fasthttp server

app_test.go

@@ -21,7 +21,7 @@ import (
 	"regexp"
 	"runtime"
 	"strings"
-	"sync/atomic"
+	"sync"
 	"testing"
 	"time"
 
@@ -880,62 +880,103 @@ func Test_App_ShutdownWithTimeout(t *testing.T) {
 func Test_App_ShutdownWithContext(t *testing.T) {
 	t.Parallel()
 
-	app := New()
-	var shutdownHookCalled atomic.Int32
-
-	app.Hooks().OnPreShutdown(func() error {
-		shutdownHookCalled.Store(1)
-		return nil
-	})
+	t.Run("successful shutdown", func(t *testing.T) {
+		t.Parallel()
+		app := New()
 
-	app.Get("/", func(ctx Ctx) error {
-		time.Sleep(5 * time.Second)
-		return ctx.SendString("body")
-	})
+		// Fast request that should complete
+		app.Get("/", func(c Ctx) error {
+			return c.SendString("OK")
+		})
 
-	ln := fasthttputil.NewInmemoryListener()
+		ln := fasthttputil.NewInmemoryListener()
+		serverStarted := make(chan bool, 1)
 
-	serverErr := make(chan error, 1)
-	go func() {
-		serverErr <- app.Listener(ln)
-	}()
+		go func() {
+			serverStarted <- true
+			_ = app.Listener(ln)
+		}()
 
-	time.Sleep(100 * time.Millisecond)
+		<-serverStarted
 
-	clientDone := make(chan struct{})
-	go func() {
+		// 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"))
-		assert.NoError(t, err)
-		close(clientDone)
-	}()
-
-	<-clientDone
-	// Sleep to ensure the server has started processing the request
-	time.Sleep(100 * time.Millisecond)
+		require.NoError(t, err)
 
-	shutdownErr := make(chan error, 1)
-	go func() {
-		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		// Shutdown with sufficient timeout
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
-		shutdownErr <- app.ShutdownWithContext(ctx)
-	}()
 
-	select {
-	case <-time.After(2 * time.Second):
-		t.Fatal("shutdown did not complete in time")
-	case err := <-shutdownErr:
-		require.Error(t, err, "Expected shutdown to return an error due to timeout")
-		require.ErrorIs(t, err, context.DeadlineExceeded, "Expected DeadlineExceeded error")
-	}
+		err = app.ShutdownWithContext(ctx)
+		require.NoError(t, err, "Expected successful shutdown")
+	})
+
+	t.Run("shutdown with hooks", func(t *testing.T) {
+		t.Parallel()
+		app := New()
+
+		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(err error) error {
+			hookMutex.Lock()
+			hookOrder = append(hookOrder, "post")
+			hookMutex.Unlock()
+			return nil
+		})
+
+		ln := fasthttputil.NewInmemoryListener()
+		go func() {
+			_ = app.Listener(ln)
+		}()
+
+		time.Sleep(100 * time.Millisecond)
+
+		err := app.ShutdownWithContext(context.Background())
+		require.NoError(t, err)
 
-	assert.Equal(t, int32(1), shutdownHookCalled.Load(), "Shutdown hook was not called")
+		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()
+
+		app.Get("/", func(c Ctx) error {
+			time.Sleep(2 * time.Second)
+			return c.SendString("OK")
+		})
+
+		ln := fasthttputil.NewInmemoryListener()
+		go func() {
+			_ = app.Listener(ln)
+		}()
 
-	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
+		time.Sleep(100 * time.Millisecond)
+
+		// Start long request
+		go func() {
+			conn, _ := ln.Dial()
+			_, _ = conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"))
+		}()
+
+		time.Sleep(100 * time.Millisecond) // Wait for request to start
+
+		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

listen.go

@@ -134,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
 	}