diff --git a/app.go b/app.go index 0c48cebb96..b982249e4b 100644 --- a/app.go +++ b/app.go @@ -9,6 +9,7 @@ package fiber import ( "bufio" + "context" "encoding/json" "encoding/xml" "errors" @@ -839,12 +840,30 @@ func (app *App) HandlersCount() uint32 { } // Shutdown gracefully shuts down the server without interrupting any active connections. -// Shutdown works by first closing all open listeners and then waiting indefinitely for all connections to return to idle and then shut down. +// Shutdown works by first closing all open listeners and then waiting indefinitely for all connections to return to idle before shutting down. // // Make sure the program doesn't exit and waits instead for Shutdown to return. // // 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()) +} + +// ShutdownWithTimeout gracefully shuts down the server without interrupting any active connections. However, if the timeout is exceeded, +// ShutdownWithTimeout will forcefully close any active connections. +// ShutdownWithTimeout works by first closing all open listeners and then waiting for all connections to return to idle before shutting down. +// +// Make sure the program doesn't exit and waits instead for ShutdownWithTimeout to return. +// +// ShutdownWithTimeout does not close keepalive connections so its recommended to set ReadTimeout to something else than 0. +func (app *App) ShutdownWithTimeout(timeout time.Duration) error { + ctx, cancelFunc := context.WithTimeout(context.Background(), timeout) + defer cancelFunc() + return app.shutdownWithContext(ctx) +} + +// shutdownWithContext shuts down the server including by force if the context's deadline is exceeded. +func (app *App) shutdownWithContext(ctx context.Context) error { if app.hooks != nil { defer app.hooks.executeOnShutdownHooks() } @@ -854,7 +873,7 @@ func (app *App) Shutdown() error { if app.server == nil { return fmt.Errorf("shutdown: server is not running") } - return app.server.Shutdown() + return app.server.ShutdownWithContext(ctx) } // Server returns the underlying fasthttp server diff --git a/app_test.go b/app_test.go index 7d532f3143..343e012783 100644 --- a/app_test.go +++ b/app_test.go @@ -6,6 +6,7 @@ package fiber import ( "bytes" + "context" "crypto/tls" "errors" "fmt" @@ -23,6 +24,7 @@ import ( "github.com/gofiber/fiber/v2/utils" "github.com/valyala/fasthttp" + "github.com/valyala/fasthttp/fasthttputil" ) var testEmptyHandler = func(c *Ctx) error { @@ -707,6 +709,45 @@ func Test_App_Shutdown(t *testing.T) { }) } +func Test_App_ShutdownWithTimeout(t *testing.T) { + app := New() + app.Get("/", func(ctx *Ctx) error { + time.Sleep(5 * time.Second) + return ctx.SendString("body") + }) + ln := fasthttputil.NewInmemoryListener() + go func() { + utils.AssertEqual(t, nil, app.Listener(ln)) + }() + time.Sleep(1 * time.Second) + go func() { + conn, err := ln.Dial() + if err != nil { + t.Errorf("unexepcted error: %v", err) + } + + if _, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n")); err != nil { + t.Errorf("unexpected error: %v", err) + } + }() + time.Sleep(1 * time.Second) + + shutdownErr := make(chan error) + go func() { + shutdownErr <- app.ShutdownWithTimeout(1 * time.Second) + }() + + timer := time.NewTimer(time.Second * 5) + select { + case <-timer.C: + t.Fatal("idle connections not closed on shutdown") + case err := <-shutdownErr: + if err == nil || err != context.DeadlineExceeded { + t.Fatalf("unexpected err %v. Expecting %v", err, context.DeadlineExceeded) + } + } +} + // go test -run Test_App_Static_Index_Default func Test_App_Static_Index_Default(t *testing.T) { app := New()