feat: add ShutdownWithTimeout function (#2228)

* add ShutdownWithTimeout function

* no message

* fix func documentation

* test: add Test_App_ShutdownWithTimeout

Co-authored-by: rocketlaunchr-cto <rocketlaunchr.cloud@gmail.com>
Co-authored-by: kinggo <lilong.21@bytedance.com>

app.go

@@ -9,6 +9,7 @@ package fiber
 
 import (
 	"bufio"
+	"context"
 	"encoding/json"
 	"encoding/xml"
 	"errors"
@@ -838,12 +839,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()
 	}
@@ -853,7 +872,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

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 {
@@ -717,6 +719,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()