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()
}