From 4f4d2c2f8e3c43711fe922f807fa4223becd524d Mon Sep 17 00:00:00 2001
From: John Guo <john@johng.cn>
Date: Sat, 6 Jan 2024 13:03:49 +0800
Subject: [PATCH] add `MiddlewareNeverDoneCtx` for package `ghttp` (#3250)

---
 net/ghttp/ghttp_middleware_never_done_ctx.go  | 13 +++
 .../ghttp_z_unit_feature_request_ctx_test.go  | 90 +++++++++++++++++++
 os/gctx/gctx_never_done.go                    |  6 +-
 3 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 net/ghttp/ghttp_middleware_never_done_ctx.go

diff --git a/net/ghttp/ghttp_middleware_never_done_ctx.go b/net/ghttp/ghttp_middleware_never_done_ctx.go
new file mode 100644
index 00000000000..9f6cee44713
--- /dev/null
+++ b/net/ghttp/ghttp_middleware_never_done_ctx.go
@@ -0,0 +1,13 @@
+// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
+//
+// This Source Code Form is subject to the terms of the MIT License.
+// If a copy of the MIT was not distributed with this file,
+// You can obtain one at https://github.com/gogf/gf.
+
+package ghttp
+
+// MiddlewareNeverDoneCtx sets the context never done for current process.
+func MiddlewareNeverDoneCtx(r *Request) {
+	r.SetCtx(r.GetNeverDoneCtx())
+	r.Middleware.Next()
+}
diff --git a/net/ghttp/ghttp_z_unit_feature_request_ctx_test.go b/net/ghttp/ghttp_z_unit_feature_request_ctx_test.go
index 823d1d2fc7e..642e50901e8 100644
--- a/net/ghttp/ghttp_z_unit_feature_request_ctx_test.go
+++ b/net/ghttp/ghttp_z_unit_feature_request_ctx_test.go
@@ -13,6 +13,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gogf/gf/v2/container/garray"
 	"github.com/gogf/gf/v2/encoding/gbase64"
 	"github.com/gogf/gf/v2/frame/g"
 	"github.com/gogf/gf/v2/net/ghttp"
@@ -345,3 +346,92 @@ func Test_Request_Form(t *testing.T) {
 		}), "john")
 	})
 }
+
+func Test_Request_NeverDoneCtx_Done(t *testing.T) {
+	var array = garray.New(true)
+	s := g.Server(guid.S())
+	s.BindHandler("/done", func(r *ghttp.Request) {
+		var (
+			ctx    = r.Context()
+			ticker = time.NewTimer(time.Millisecond * 1500)
+		)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				array.Append(1)
+				return
+			case <-ticker.C:
+				array.Append(1)
+				return
+			}
+		}
+	})
+	s.SetDumpRouterMap(false)
+	s.Start()
+	defer s.Shutdown()
+
+	time.Sleep(100 * time.Millisecond)
+
+	c := g.Client()
+	c.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", s.GetListenedPort()))
+	gtest.C(t, func(t *gtest.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+		go func() {
+			result := c.GetContent(ctx, "/done")
+			fmt.Println(result)
+		}()
+		time.Sleep(time.Millisecond * 100)
+
+		t.Assert(array.Len(), 0)
+		cancel()
+
+		time.Sleep(time.Millisecond * 500)
+		t.Assert(array.Len(), 1)
+	})
+}
+
+func Test_Request_NeverDoneCtx_NeverDone(t *testing.T) {
+	var array = garray.New(true)
+	s := g.Server(guid.S())
+	s.Use(ghttp.MiddlewareNeverDoneCtx)
+	s.BindHandler("/never-done", func(r *ghttp.Request) {
+		var (
+			ctx    = r.Context()
+			ticker = time.NewTimer(time.Millisecond * 1500)
+		)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				array.Append(1)
+				return
+			case <-ticker.C:
+				array.Append(1)
+				return
+			}
+		}
+	})
+	s.SetDumpRouterMap(false)
+	s.Start()
+	defer s.Shutdown()
+
+	time.Sleep(100 * time.Millisecond)
+
+	c := g.Client()
+	c.SetPrefix(fmt.Sprintf("http://127.0.0.1:%d", s.GetListenedPort()))
+	gtest.C(t, func(t *gtest.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+		go func() {
+			result := c.GetContent(ctx, "/never-done")
+			fmt.Println(result)
+		}()
+		time.Sleep(time.Millisecond * 100)
+
+		t.Assert(array.Len(), 0)
+		cancel()
+
+		time.Sleep(time.Millisecond * 1500)
+		t.Assert(array.Len(), 1)
+	})
+}
diff --git a/os/gctx/gctx_never_done.go b/os/gctx/gctx_never_done.go
index 8d08e53a1bb..2e28c7809dc 100644
--- a/os/gctx/gctx_never_done.go
+++ b/os/gctx/gctx_never_done.go
@@ -32,7 +32,11 @@ func (c *neverDoneCtx) Err() error {
 }
 
 // NeverDone wraps and returns a new context object that will be never done,
-// which forbids the context manually done, to make the context can be propagated to asynchronous goroutines.
+// which forbids the context manually done, to make the context can be propagated
+// to asynchronous goroutines.
+//
+// Note that, it does not affect the closing (canceling) of the parent context,
+// as it is a wrapper for its parent, which only affects the next context handling.
 func NeverDone(ctx context.Context) context.Context {
 	return &neverDoneCtx{ctx}
 }