🔥 feat: Improve and Optimize ShutdownWithContext Func

app.go

@@ -841,16 +841,18 @@ func (app *App) ShutdownWithTimeout(timeout time.Duration) error {
 //
 // 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 {
-		// TODO: check should be defered?
-		app.hooks.executeOnShutdownHooks()
-	}
-
 	app.mutex.Lock()
 	defer app.mutex.Unlock()
+
 	if app.server == nil {
 		return ErrNotRunning
 	}
+
+	// Execute shutdown hooks in a deferred function
+	if app.hooks != nil {
+		defer app.hooks.executeOnShutdownHooks()
+	}
+
 	return app.server.ShutdownWithContext(ctx)
 }
 

app_test.go

@@ -860,45 +860,61 @@
 	t.Parallel()
 
 	app := New()
+	shutdownHookCalled := false
+    app.Hooks().OnShutdown(func() error {
+        shutdownHookCalled = true
+        return nil
+    })
+
 	app.Get("/", func(ctx Ctx) error {
 		time.Sleep(5 * time.Second)
 		return ctx.SendString("body")
 	})
 
 	ln := fasthttputil.NewInmemoryListener()
 
-	go func() {
-		err := app.Listener(ln)
-		assert.NoError(t, err)
-	}()
-
-	time.Sleep(1 * time.Second)
-
-	go func() {
-		conn, err := ln.Dial()
-		assert.NoError(t, err)
-
-		_, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: google.com\r\n\r\n"))
-		assert.NoError(t, err)
-	}()
+    serverErr := make(chan error, 1)
+    go func() {
+        serverErr <- app.Listener(ln)
+    }()
+
+    time.Sleep(100 * time.Millisecond)
+
+    clientDone := make(chan struct{})
+    go func() {
+        conn, err := ln.Dial()
+        assert.NoError(t, err)
+        _, err = conn.Write([]byte("GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"))
+        assert.NoError(t, err)
+        close(clientDone)
+    }()
+
+	<-clientDone
+    time.Sleep(100 * time.Millisecond)
+
+	shutdownErr := make(chan error, 1)
+    go func() {
+        ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+        defer cancel()
+        shutdownErr <- app.ShutdownWithContext(ctx)
+    }()
 
-	time.Sleep(1 * time.Second)
+	select {
+    case <-time.After(2 * time.Second):
+        t.Fatal("shutdown did not complete in time")
+    case err := <-shutdownErr:
+        assert.Error(t, err, "Expected shutdown to return an error due to timeout")
+        assert.True(t, errors.Is(err, context.DeadlineExceeded), "Expected DeadlineExceeded error")
+    }
 
-	shutdownErr := make(chan error)
-	go func() {
-		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
-		defer cancel()
-		shutdownErr <- app.ShutdownWithContext(ctx)
-	}()
+	assert.True(t, shutdownHookCalled, "Shutdown hook was not called")
 
 	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)
-		}
-	}
+    case err := <-serverErr:
+        assert.NoError(t, err, "Server should have shut down without error")
+    default:
+        // Server is still running, which is expected as the long-running request prevented full shutdown
+    }
 }
 
 // go test -run Test_App_Mixed_Routes_WithSameLen