From 1c76e56a0dc4644d1cc43f27f1d9b5ff3804a7fe Mon Sep 17 00:00:00 2001
From: Wenbo Han <hwbrzzl@gmail.com>
Date: Sat, 7 Sep 2024 16:42:42 +0800
Subject: [PATCH] feat: Optmize Shutdown (#99)

---
 go.mod        | 2 +-
 go.sum        | 4 ++--
 route.go      | 9 +++++++--
 route_test.go | 5 ++---
 4 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/go.mod b/go.mod
index 612f783..4176ed0 100644
--- a/go.mod
+++ b/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/gofiber/fiber/v2 v2.52.5
 	github.com/gofiber/template/html/v2 v2.1.2
 	github.com/gookit/validate v1.5.2
-	github.com/goravel/framework v1.14.5-0.20240904064435-05476fa00c9c
+	github.com/goravel/framework v1.14.5
 	github.com/savioxavier/termlink v1.3.0
 	github.com/spf13/cast v1.6.0
 	github.com/stretchr/testify v1.9.0
diff --git a/go.sum b/go.sum
index a699821..5153aa8 100644
--- a/go.sum
+++ b/go.sum
@@ -386,8 +386,8 @@ github.com/gookit/validate v1.5.2 h1:i5I2OQ7WYHFRPRATGu9QarR9snnNHydvwSuHXaRWAV0
 github.com/gookit/validate v1.5.2/go.mod h1:yuPy2WwDlwGRa06fFJ5XIO8QEwhRnTC2LmxmBa5SE14=
 github.com/goravel/file-rotatelogs/v2 v2.4.2 h1:g68AzbePXcm0V2CpUMc9j4qVzcDn7+7aoWSjZ51C0m4=
 github.com/goravel/file-rotatelogs/v2 v2.4.2/go.mod h1:23VuSW8cBS4ax5cmbV+5AaiLpq25b8UJ96IhbAkdo8I=
-github.com/goravel/framework v1.14.5-0.20240904064435-05476fa00c9c h1:RiMpkbasXxmNmE/pLCy6lqmu3uOadyhICiqMtgroA7I=
-github.com/goravel/framework v1.14.5-0.20240904064435-05476fa00c9c/go.mod h1:rScDXGQZdoVfyxemNPmijlz/2a+lWNOa4jTuak5GGVg=
+github.com/goravel/framework v1.14.5 h1:FItqxRGkBK0h/TIknF24TuMZCtBRaSr3DnQLEzhfvXc=
+github.com/goravel/framework v1.14.5/go.mod h1:rScDXGQZdoVfyxemNPmijlz/2a+lWNOa4jTuak5GGVg=
 github.com/gorilla/securecookie v1.1.1/go.mod h1:ra0sb63/xPlUeL+yeDciTfxMRAA+MP+HVt/4epWDjd4=
 github.com/gorilla/sessions v1.2.1/go.mod h1:dk2InVEVJ0sfLlnXv9EAgkf6ecYs/i80K/zI+bUmuGM=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 h1:UH//fgunKIs4JdUbpDl1VZCDaL56wXCB/5+wF6uHfaI=
diff --git a/route.go b/route.go
index c5197a4..2c27243 100644
--- a/route.go
+++ b/route.go
@@ -187,8 +187,13 @@ func (r *Route) RunTLSWithCert(host, certFile, keyFile string) error {
 
 // Shutdown gracefully shuts down the server
 // Shutdown 优雅退出HTTP Server
-func (r *Route) Shutdown(ctx context.Context) error {
-	return r.instance.ShutdownWithContext(ctx)
+func (r *Route) Shutdown(ctx ...context.Context) error {
+	c := context.Background()
+	if len(ctx) > 0 {
+		c = ctx[0]
+	}
+
+	return r.instance.ShutdownWithContext(c)
 }
 
 // ServeHTTP serve http request (Not support)
diff --git a/route_test.go b/route_test.go
index 4eb76e8..aa8879c 100644
--- a/route_test.go
+++ b/route_test.go
@@ -1,7 +1,6 @@
 package fiber
 
 import (
-	"context"
 	"crypto/tls"
 	"errors"
 	"fmt"
@@ -429,7 +428,7 @@ func TestShutdown(t *testing.T) {
 
 				assertHttpNormal(t, addr, true)
 
-				assert.Nil(t, route.Shutdown(context.Background()))
+				assert.Nil(t, route.Shutdown())
 
 				assertHttpNormal(t, addr, false)
 				return nil
@@ -454,7 +453,7 @@ func TestShutdown(t *testing.T) {
 					}()
 				}
 				time.Sleep(100 * time.Millisecond)
-				assert.Nil(t, route.Shutdown(context.Background()))
+				assert.Nil(t, route.Shutdown())
 				assertHttpNormal(t, addr, false)
 				wg.Wait()
 				assert.Equal(t, count.Load(), int64(3))