From 562d15db863248fe4d249b76560843c3a29511d3 Mon Sep 17 00:00:00 2001
From: "Juan C. Yamacho H" <juancarlos.yamacho@blacklane.com>
Date: Sun, 9 Apr 2023 15:08:03 +0200
Subject: [PATCH] :rocket: Feature: Public ShutdownWithContext (#2407)

* feat: public shutdown with context

* docs: add server shutdown option

* chore: revert spacing changes

* test: app shutdown with context
---
 app.go          | 12 ++++++++----
 app_test.go     | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 docs/api/app.md |  3 +++
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/app.go b/app.go
index 7857da34..ce23ee7a 100644
--- a/app.go
+++ b/app.go
@@ -847,7 +847,7 @@ func (app *App) HandlersCount() uint32 {
 //
 // 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())
+	return app.ShutdownWithContext(context.Background())
 }
 
 // ShutdownWithTimeout gracefully shuts down the server without interrupting any active connections. However, if the timeout is exceeded,
@@ -860,11 +860,15 @@ func (app *App) Shutdown() error {
 func (app *App) ShutdownWithTimeout(timeout time.Duration) error {
 	ctx, cancelFunc := context.WithTimeout(context.Background(), timeout)
 	defer cancelFunc()
-	return app.shutdownWithContext(ctx)
+	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 {
+// ShutdownWithContext shuts down the server including by force if the context's deadline is exceeded.
+//
+// Make sure the program doesn't exit and waits instead for ShutdownWithTimeout to return.
+//
+// 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 {
 		defer app.hooks.executeOnShutdownHooks()
 	}
diff --git a/app_test.go b/app_test.go
index aeb946b8..3ef57948 100644
--- a/app_test.go
+++ b/app_test.go
@@ -788,6 +788,53 @@ 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")
+	})
+
+	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() {
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		defer cancel()
+		shutdownErr <- app.ShutdownWithContext(ctx)
+	}()
+
+	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)
+		}
+	}
+}
+
 // go test -run Test_App_Static_Index_Default
 func Test_App_Static_Index_Default(t *testing.T) {
 	t.Parallel()
diff --git a/docs/api/app.md b/docs/api/app.md
index 21e631a4..5f788ca3 100644
--- a/docs/api/app.md
+++ b/docs/api/app.md
@@ -236,9 +236,12 @@ 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.
+
 ```go
 func (app *App) Shutdown() error
 func (app *App) ShutdownWithTimeout(timeout time.Duration) error
+func (app *App) ShutdownWithContext(ctx context.Context) error
 ```
 
 ## HandlersCount