Skip to content

Commit 4b62d3d

Browse files
JIeJaittyingjie.huanggaby
authored
🔥 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 <[email protected]> Co-authored-by: Juan Calderon-Perez <[email protected]>
1 parent 252a022 commit 4b62d3d

File tree

9 files changed

+421
-251
lines changed

9 files changed

+421
-251
lines changed

app.go

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -894,6 +894,13 @@ func (app *App) HandlersCount() uint32 {
894894
//
895895
// Make sure the program doesn't exit and waits instead for Shutdown to return.
896896
//
897+
// Important: app.Listen() must be called in a separate goroutine, otherwise shutdown hooks will not work
898+
// as Listen() is a blocking operation. Example:
899+
//
900+
// go app.Listen(":3000")
901+
// // ...
902+
// app.Shutdown()
903+
//
897904
// Shutdown does not close keepalive connections so its recommended to set ReadTimeout to something else than 0.
898905
func (app *App) Shutdown() error {
899906
return app.ShutdownWithContext(context.Background())
@@ -918,17 +925,21 @@ func (app *App) ShutdownWithTimeout(timeout time.Duration) error {
918925
//
919926
// ShutdownWithContext does not close keepalive connections so its recommended to set ReadTimeout to something else than 0.
920927
func (app *App) ShutdownWithContext(ctx context.Context) error {
921-
if app.hooks != nil {
922-
// TODO: check should be defered?
923-
app.hooks.executeOnShutdownHooks()
924-
}
925-
926928
app.mutex.Lock()
927929
defer app.mutex.Unlock()
930+
931+
var err error
932+
928933
if app.server == nil {
929934
return ErrNotRunning
930935
}
931-
return app.server.ShutdownWithContext(ctx)
936+
937+
// Execute the Shutdown hook
938+
app.hooks.executeOnPreShutdownHooks()
939+
defer app.hooks.executeOnPostShutdownHooks(err)
940+
941+
err = app.server.ShutdownWithContext(ctx)
942+
return err
932943
}
933944

934945
// Server returns the underlying fasthttp server

app_test.go

Lines changed: 126 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"regexp"
2222
"runtime"
2323
"strings"
24+
"sync"
2425
"testing"
2526
"time"
2627

@@ -930,20 +931,29 @@ func Test_App_ShutdownWithTimeout(t *testing.T) {
930931
})
931932

932933
ln := fasthttputil.NewInmemoryListener()
934+
serverReady := make(chan struct{}) // Signal that the server is ready to start
935+
933936
go func() {
937+
serverReady <- struct{}{}
934938
err := app.Listener(ln)
935939
assert.NoError(t, err)
936940
}()
937941

938-
time.Sleep(1 * time.Second)
942+
<-serverReady // Waiting for the server to be ready
943+
944+
// Create a connection and send a request
945+
connReady := make(chan struct{})
939946
go func() {
940947
conn, err := ln.Dial()
941948
assert.NoError(t, err)
942949

943950
_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n"))
944951
assert.NoError(t, err)
952+
953+
connReady <- struct{}{} // Signal that the request has been sent
945954
}()
946-
time.Sleep(1 * time.Second)
955+
956+
<-connReady // Waiting for the request to be sent
947957

948958
shutdownErr := make(chan error)
949959
go func() {
@@ -964,46 +974,130 @@ func Test_App_ShutdownWithTimeout(t *testing.T) {
964974
func Test_App_ShutdownWithContext(t *testing.T) {
965975
t.Parallel()
966976

967-
app := New()
968-
app.Get("/", func(ctx Ctx) error {
969-
time.Sleep(5 * time.Second)
970-
return ctx.SendString("body")
971-
})
977+
t.Run("successful shutdown", func(t *testing.T) {
978+
t.Parallel()
979+
app := New()
972980

973-
ln := fasthttputil.NewInmemoryListener()
981+
// Fast request that should complete
982+
app.Get("/", func(c Ctx) error {
983+
return c.SendString("OK")
984+
})
974985

975-
go func() {
976-
err := app.Listener(ln)
977-
assert.NoError(t, err)
978-
}()
986+
ln := fasthttputil.NewInmemoryListener()
987+
serverStarted := make(chan bool, 1)
979988

980-
time.Sleep(1 * time.Second)
989+
go func() {
990+
serverStarted <- true
991+
if err := app.Listener(ln); err != nil {
992+
t.Errorf("Failed to start listener: %v", err)
993+
}
994+
}()
981995

982-
go func() {
996+
<-serverStarted
997+
998+
// Execute normal request
983999
conn, err := ln.Dial()
984-
assert.NoError(t, err)
1000+
require.NoError(t, err)
1001+
_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"))
1002+
require.NoError(t, err)
9851003

986-
_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n"))
987-
assert.NoError(t, err)
988-
}()
1004+
// Shutdown with sufficient timeout
1005+
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
1006+
defer cancel()
9891007

990-
time.Sleep(1 * time.Second)
1008+
err = app.ShutdownWithContext(ctx)
1009+
require.NoError(t, err, "Expected successful shutdown")
1010+
})
9911011

992-
shutdownErr := make(chan error)
993-
go func() {
994-
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
995-
defer cancel()
996-
shutdownErr <- app.ShutdownWithContext(ctx)
997-
}()
1012+
t.Run("shutdown with hooks", func(t *testing.T) {
1013+
t.Parallel()
1014+
app := New()
9981015

999-
select {
1000-
case <-time.After(5 * time.Second):
1001-
t.Fatal("idle connections not closed on shutdown")
1002-
case err := <-shutdownErr:
1003-
if err == nil || !errors.Is(err, context.DeadlineExceeded) {
1004-
t.Fatalf("unexpected err %v. Expecting %v", err, context.DeadlineExceeded)
1016+
hookOrder := make([]string, 0)
1017+
var hookMutex sync.Mutex
1018+
1019+
app.Hooks().OnPreShutdown(func() error {
1020+
hookMutex.Lock()
1021+
hookOrder = append(hookOrder, "pre")
1022+
hookMutex.Unlock()
1023+
return nil
1024+
})
1025+
1026+
app.Hooks().OnPostShutdown(func(_ error) error {
1027+
hookMutex.Lock()
1028+
hookOrder = append(hookOrder, "post")
1029+
hookMutex.Unlock()
1030+
return nil
1031+
})
1032+
1033+
ln := fasthttputil.NewInmemoryListener()
1034+
go func() {
1035+
if err := app.Listener(ln); err != nil {
1036+
t.Errorf("Failed to start listener: %v", err)
1037+
}
1038+
}()
1039+
1040+
time.Sleep(100 * time.Millisecond)
1041+
1042+
err := app.ShutdownWithContext(context.Background())
1043+
require.NoError(t, err)
1044+
1045+
require.Equal(t, []string{"pre", "post"}, hookOrder, "Hooks should execute in order")
1046+
})
1047+
1048+
t.Run("timeout with long running request", func(t *testing.T) {
1049+
t.Parallel()
1050+
app := New()
1051+
1052+
requestStarted := make(chan struct{})
1053+
requestProcessing := make(chan struct{})
1054+
1055+
app.Get("/", func(c Ctx) error {
1056+
close(requestStarted)
1057+
// Wait for signal to continue processing the request
1058+
<-requestProcessing
1059+
time.Sleep(2 * time.Second)
1060+
return c.SendString("OK")
1061+
})
1062+
1063+
ln := fasthttputil.NewInmemoryListener()
1064+
go func() {
1065+
if err := app.Listener(ln); err != nil {
1066+
t.Errorf("Failed to start listener: %v", err)
1067+
}
1068+
}()
1069+
1070+
// Ensure server is fully started
1071+
time.Sleep(100 * time.Millisecond)
1072+
1073+
// Start a long-running request
1074+
go func() {
1075+
conn, err := ln.Dial()
1076+
if err != nil {
1077+
t.Errorf("Failed to dial: %v", err)
1078+
return
1079+
}
1080+
if _, err := conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n")); err != nil {
1081+
t.Errorf("Failed to write: %v", err)
1082+
}
1083+
}()
1084+
1085+
// Wait for request to start
1086+
select {
1087+
case <-requestStarted:
1088+
// Request has started, signal to continue processing
1089+
close(requestProcessing)
1090+
case <-time.After(2 * time.Second):
1091+
t.Fatal("Request did not start in time")
10051092
}
1006-
}
1093+
1094+
// Attempt shutdown, should timeout
1095+
ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
1096+
defer cancel()
1097+
1098+
err := app.ShutdownWithContext(ctx)
1099+
require.ErrorIs(t, err, context.DeadlineExceeded)
1100+
})
10071101
}
10081102

10091103
// go test -run Test_App_Mixed_Routes_WithSameLen

docs/api/fiber.md

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -111,11 +111,9 @@ app.Listen(":8080", fiber.ListenConfig{
111111
| <Reference id="enableprefork">EnablePrefork</Reference> | `bool` | When set to true, this will spawn multiple Go processes listening on the same port. | `false` |
112112
| <Reference id="enableprintroutes">EnablePrintRoutes</Reference> | `bool` | If set to true, will print all routes with their method, path, and handler. | `false` |
113113
| <Reference id="gracefulcontext">GracefulContext</Reference> | `context.Context` | Field to shutdown Fiber by given context gracefully. | `nil` |
114-
| <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` |
114+
| <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` |
115115
| <Reference id="listeneraddrfunc">ListenerAddrFunc</Reference> | `func(addr net.Addr)` | Allows accessing and customizing `net.Listener`. | `nil` |
116116
| <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` |
117-
| <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` |
118-
| <Reference id="onshutdownsuccess">OnShutdownSuccess</Reference> | `func()` | Allows customizing success behavior when gracefully shutting down the server by given signal. | `nil` |
119117
| <Reference id="tlsconfigfunc">TLSConfigFunc</Reference> | `func(tlsConfig *tls.Config)` | Allows customizing `tls.Config` as you want. | `nil` |
120118
| <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` |
121119
| <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
230228

231229
ShutdownWithTimeout will forcefully close any active connections after the timeout expires.
232230

233-
ShutdownWithContext shuts down the server including by force if the context's deadline is exceeded.
231+
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.
234232

235233
```go
236234
func (app *App) Shutdown() error

docs/api/hooks.md

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@ With Fiber you can execute custom user functions at specific method execution po
1515
- [OnGroupName](#ongroupname)
1616
- [OnListen](#onlisten)
1717
- [OnFork](#onfork)
18-
- [OnShutdown](#onshutdown)
18+
- [OnPreShutdown](#onpreshutdown)
19+
- [OnPostShutdown](#onpostshutdown)
1920
- [OnMount](#onmount)
2021

2122
## Constants
@@ -28,7 +29,8 @@ type OnGroupHandler = func(Group) error
2829
type OnGroupNameHandler = OnGroupHandler
2930
type OnListenHandler = func(ListenData) error
3031
type OnForkHandler = func(int) error
31-
type OnShutdownHandler = func() error
32+
type OnPreShutdownHandler = func() error
33+
type OnPostShutdownHandler = func(error) error
3234
type OnMountHandler = func(*App) error
3335
```
3436

@@ -174,12 +176,20 @@ func main() {
174176
func (h *Hooks) OnFork(handler ...OnForkHandler)
175177
```
176178

177-
## OnShutdown
179+
## OnPreShutdown
178180

179-
`OnShutdown` is a hook to execute user functions after shutdown.
181+
`OnPreShutdown` is a hook to execute user functions before shutdown.
180182

181183
```go title="Signature"
182-
func (h *Hooks) OnShutdown(handler ...OnShutdownHandler)
184+
func (h *Hooks) OnPreShutdown(handler ...OnPreShutdownHandler)
185+
```
186+
187+
## OnPostShutdown
188+
189+
`OnPostShutdown` is a hook to execute user functions after shutdown.
190+
191+
```go title="Signature"
192+
func (h *Hooks) OnPostShutdown(handler ...OnPostShutdownHandler)
183193
```
184194

185195
## OnMount

docs/whats_new.md

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ In this guide, we'll walk you through the most important changes in Fiber `v3` a
1616
Here's a quick overview of the changes in Fiber `v3`:
1717

1818
- [🚀 App](#-app)
19+
- [🎣 Hooks](#-hooks)
20+
- [🚀 Listen](#-listen)
1921
- [🗺️ Router](#-router)
2022
- [🧠 Context](#-context)
2123
- [📎 Binding](#-binding)
@@ -158,6 +160,63 @@ app.Listen(":444", fiber.ListenConfig{
158160
})
159161
```
160162

163+
## 🎣 Hooks
164+
165+
We have made several changes to the Fiber hooks, including:
166+
167+
- Added new shutdown hooks to provide better control over the shutdown process:
168+
- `OnPreShutdown` - Executes before the server starts shutting down
169+
- `OnPostShutdown` - Executes after the server has shut down, receives any shutdown error
170+
- Deprecated `OnShutdown` in favor of the new pre/post shutdown hooks
171+
- Improved shutdown hook execution order and reliability
172+
- Added mutex protection for hook registration and execution
173+
174+
Important: When using shutdown hooks, ensure app.Listen() is called in a separate goroutine:
175+
176+
```go
177+
// Correct usage
178+
go app.Listen(":3000")
179+
// ... register shutdown hooks
180+
app.Shutdown()
181+
182+
// Incorrect usage - hooks won't work
183+
app.Listen(":3000") // This blocks
184+
app.Shutdown() // Never reached
185+
```
186+
187+
## 🚀 Listen
188+
189+
We have made several changes to the Fiber listen, including:
190+
191+
- Removed `OnShutdownError` and `OnShutdownSuccess` from `ListenerConfig` in favor of using `OnPostShutdown` hook which receives the shutdown error
192+
193+
```go
194+
app := fiber.New()
195+
196+
// Before - using ListenerConfig callbacks
197+
app.Listen(":3000", fiber.ListenerConfig{
198+
OnShutdownError: func(err error) {
199+
log.Printf("Shutdown error: %v", err)
200+
},
201+
OnShutdownSuccess: func() {
202+
log.Println("Shutdown successful")
203+
},
204+
})
205+
206+
// After - using OnPostShutdown hook
207+
app.Hooks().OnPostShutdown(func(err error) error {
208+
if err != nil {
209+
log.Printf("Shutdown error: %v", err)
210+
} else {
211+
log.Println("Shutdown successful")
212+
}
213+
return nil
214+
})
215+
go app.Listen(":3000")
216+
```
217+
218+
This change simplifies the shutdown handling by consolidating the shutdown callbacks into a single hook that receives the error status.
219+
161220
## 🗺 Router
162221

163222
We have slightly adapted our router interface

0 commit comments

Comments
 (0)