Repository: cloudwego/hertz
Pull request #1101 (2024-04-25 12:01:57)

feat: add method to exile requestContext
#### What type of PR is this?
feat

#### Check the PR title.
<!--
The description of the title will be attached in Release Notes, 
so please describe it from user-oriented, what this PR does / why we need it.
Please check your PR title with the below requirements:
-->
- [x] This PR title match the format: \<type\>(optional scope): \<description\>
- [x] The description of this PR title is user-oriented and clear enough for others to understand.
- [ ] Attach the PR updating the user documentation if the current PR requires user awareness at the usage level. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)


#### (Optional) Translate the PR title into Chinese.
增加一个方法能够放逐掉当前的 RequestContext（请求结束不入池）

#### (Optional) More detailed description for this PR(en: English/zh: Chinese).
<!--
Provide more detailed info for review(e.g., it's recommended to provide perf data if this is a perf type PR).
-->
en:
zh(optional): Experimental features: Use with caution, it may have a slight impact on performance.
实验性功能：谨慎使用，可能会对性能有少许影响。

#### (Optional) Which issue(s) this PR fixes:
<!--
Automatically closes linked issue when PR is merged.
Eg: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
-->

#### (Optional) The PR that updates user documentation:
<!--
If the current PR requires user awareness at the usage level, please submit a PR to update user docs. [User docs repo](https://github.com/cloudwego/cloudwego.github.io)
-->


Patch:
diff --git a/cmd/hz/app/app.go b/cmd/hz/app/app.go
index bc52efe..116df93 100644
--- a/cmd/hz/app/app.go
+++ b/cmd/hz/app/app.go
@@ -179,6 +179,7 @@ func Init() *cli.App {
 	noRecurseFlag := cli.BoolFlag{Name: "no_recurse", Usage: "Generate master model only.", Destination: &globalArgs.NoRecurse}
 	forceNewFlag := cli.BoolFlag{Name: "force", Aliases: []string{"f"}, Usage: "Force new a project, which will overwrite the generated files", Destination: &globalArgs.ForceNew}
 	enableExtendsFlag := cli.BoolFlag{Name: "enable_extends", Usage: "Parse 'extends' for thrift IDL", Destination: &globalArgs.EnableExtends}
+	sortRouterFlag := cli.BoolFlag{Name: "sort_router", Usage: "Sort router register code, to avoid code difference", Destination: &globalArgs.SortRouter}
 
 	jsonEnumStrFlag := cli.BoolFlag{Name: "json_enumstr", Usage: "Use string instead of num for json enums when idl is thrift.", Destination: &globalArgs.JSONEnumStr}
 	queryEnumIntFlag := cli.BoolFlag{Name: "query_enumint", Usage: "Use num instead of string for query enum parameter.", Destination: &globalArgs.QueryEnumAsInt}
@@ -227,6 +228,7 @@ func Init() *cli.App {
 				&noRecurseFlag,
 				&forceNewFlag,
 				&enableExtendsFlag,
+				&sortRouterFlag,
 
 				&jsonEnumStrFlag,
 				&unsetOmitemptyFlag,
@@ -261,6 +263,7 @@ func Init() *cli.App {
 				&optPkgFlag,
 				&noRecurseFlag,
 				&enableExtendsFlag,
+				&sortRouterFlag,
 
 				&jsonEnumStrFlag,
 				&unsetOmitemptyFlag,
diff --git a/cmd/hz/config/argument.go b/cmd/hz/config/argument.go
index 2e208d2..4e6d75e 100644
--- a/cmd/hz/config/argument.go
+++ b/cmd/hz/config/argument.go
@@ -72,6 +72,7 @@ type Argument struct {
 	ForceNew             bool
 	SnakeStyleMiddleware bool
 	EnableExtends        bool
+	SortRouter           bool
 
 	CustomizeLayout     string
 	CustomizeLayoutData string
diff --git a/cmd/hz/generator/handler.go b/cmd/hz/generator/handler.go
index 609e1d6..eeab0cf 100644
--- a/cmd/hz/generator/handler.go
+++ b/cmd/hz/generator/handler.go
@@ -161,7 +161,7 @@ func (pkgGen *HttpPackageGenerator) processHandler(handler *Handler, root *Route
 			}
 			handler.Imports[mm.PackageName] = mm
 		}
-		err := root.Update(m, handler.PackageName, singleHandlerPackage)
+		err := root.Update(m, handler.PackageName, singleHandlerPackage, pkgGen.SortRouter)
 		if err != nil {
 			return err
 		}
diff --git a/cmd/hz/generator/package.go b/cmd/hz/generator/package.go
index 22e3113..db4d976 100644
--- a/cmd/hz/generator/package.go
+++ b/cmd/hz/generator/package.go
@@ -63,12 +63,13 @@ type HttpPackageGenerator struct {
 	IdlClientDir   string // client dir for "client" command
 	ForceClientDir string // client dir without namespace for "client" command
 	BaseDomain     string // request domain for "client" command
-	QueryEnumAsInt bool // client code use number for query parameter
+	QueryEnumAsInt bool   // client code use number for query parameter
 	ServiceGenDir  string
 
 	NeedModel            bool
 	HandlerByMethod      bool // generate handler files with method dimension
 	SnakeStyleMiddleware bool // use snake name style for middleware
+	SortRouter           bool
 
 	loadedBackend   Backend
 	curModel        *model.Model
diff --git a/cmd/hz/generator/router.go b/cmd/hz/generator/router.go
index 10f4316..1300087 100644
--- a/cmd/hz/generator/router.go
+++ b/cmd/hz/generator/router.go
@@ -20,10 +20,13 @@ import (
 	"bytes"
 	"fmt"
 	"io/ioutil"
+	"math"
 	"path/filepath"
 	"regexp"
 	"sort"
+	"strconv"
 	"strings"
+	"unicode"
 
 	"github.com/cloudwego/hertz/cmd/hz/util"
 )
@@ -73,7 +76,7 @@ func (routerNode *RouterNode) Sort() {
 	sort.Sort(routerNode.Children)
 }
 
-func (routerNode *RouterNode) Update(method *HttpMethod, handlerType, handlerPkg string) error {
+func (routerNode *RouterNode) Update(method *HttpMethod, handlerType, handlerPkg string, sortRouter bool) error {
 	if method.Path == "" {
 		return fmt.Errorf("empty path for method '%s'", method.Name)
 	}
@@ -81,12 +84,12 @@ func (routerNode *RouterNode) Update(method *HttpMethod, handlerType, handlerPkg
 	if paths[0] == "" {
 		paths = paths[1:]
 	}
-	parent, last := routerNode.FindNearest(paths)
+	parent, last := routerNode.FindNearest(paths, method.HTTPMethod, sortRouter)
 	if last == len(paths) {
 		return fmt.Errorf("path '%s' has been registered", method.Path)
 	}
 	name := util.ToVarName(paths[:last])
-	parent.Insert(name, method, handlerType, paths[last:], handlerPkg)
+	parent.Insert(name, method, handlerType, paths[last:], handlerPkg, sortRouter)
 	parent.Sort()
 	return nil
 }
@@ -192,7 +195,7 @@ func (routerNode *RouterNode) DFS(i int, hook func(layer int, node *RouterNode)
 
 var handlerPkgMap map[string]string
 
-func (routerNode *RouterNode) Insert(name string, method *HttpMethod, handlerType string, paths []string, handlerPkg string) {
+func (routerNode *RouterNode) Insert(name string, method *HttpMethod, handlerType string, paths []string, handlerPkg string, sortRouter bool) {
 	cur := routerNode
 	for i, p := range paths {
 		c := &RouterNode{
@@ -229,6 +232,9 @@ func (routerNode *RouterNode) Insert(name string, method *HttpMethod, handlerTyp
 			cur.Children = make([]*RouterNode, 0, 1)
 		}
 		cur.Children = append(cur.Children, c)
+		if sortRouter {
+			sort.Sort(cur.Children)
+		}
 		cur = c
 	}
 }
@@ -240,14 +246,21 @@ func getHttpMethod(method string) string {
 	return strings.ToUpper(method)
 }
 
-func (routerNode *RouterNode) FindNearest(paths []string) (*RouterNode, int) {
+func (routerNode *RouterNode) FindNearest(paths []string, method string, sortRouter bool) (*RouterNode, int) {
 	ns := len(paths)
 	cur := routerNode
 	i := 0
 	path := paths[i]
 	for j := 0; j < len(cur.Children); j++ {
 		c := cur.Children[j]
+		tmpMethod := "" // group do not have http method
+		if i == ns {    // only i==ns, the path is http method node
+			tmpMethod = method
+		}
 		if ("/" + path) == c.Path {
+			if sortRouter && strings.EqualFold(c.HttpMethod, tmpMethod) {
+				continue
+			}
 			i++
 			if i == ns {
 				return cur, i - 1
@@ -270,17 +283,35 @@ func (c childrenRouterInfo) Len() int {
 // Less reports whether the element with
 // index i should sort before the element with index j.
 func (c childrenRouterInfo) Less(i, j int) bool {
-	ci := c[i].Path
-	if len(c[i].Children) != 0 {
-		ci = ci[1:]
+	if c[i].HttpMethod == "" && c[j].HttpMethod != "" {
+		return false
+	}
+	if c[i].HttpMethod != "" && c[j].HttpMethod == "" {
+		return true
 	}
-	cj := c[j].Path
-	if len(c[j].Children) != 0 {
-		cj = cj[1:]
+	// remove non-litter char
+	// eg. /a -> a
+	//     /:a -> a
+	ci := removeNonLetterPrefix(c[i].Path)
+	cj := removeNonLetterPrefix(c[j].Path)
+
+	// if ci == cj, use HTTP mothod for sort, preventing sorting inconsistencies
+	if ci == cj {
+		return c[i].HttpMethod < c[j].HttpMethod
 	}
+
 	return ci < cj
 }
 
+func removeNonLetterPrefix(str string) string {
+	for i, char := range str {
+		if unicode.IsLetter(char) || unicode.IsDigit(char) {
+			return str[i:]
+		}
+	}
+	return str
+}
+
 // Swap swaps the elements with indexes i and j.
 func (c childrenRouterInfo) Swap(i, j int) {
 	c[i], c[j] = c[j], c[i]
@@ -341,6 +372,29 @@ func (pkgGen *HttpPackageGenerator) updateRegister(pkg, rDir, pkgName string) er
 	return nil
 }
 
+func appendMw(mws []string, mw string) ([]string, string) {
+	for i := 0; true; i++ {
+		if i == math.MaxInt {
+			break
+		}
+		if !stringsIncludes(mws, mw) {
+			mws = append(mws, mw)
+			break
+		}
+		mw += strconv.Itoa(i)
+	}
+	return mws, mw
+}
+
+func stringsIncludes(strs []string, str string) bool {
+	for _, s := range strs {
+		if s == str {
+			return true
+		}
+	}
+	return false
+}
+
 func (pkgGen *HttpPackageGenerator) genRouter(pkg *HttpPackage, root *RouterNode, handlerPackage, routerDir, routerPackage string) error {
 	err := root.DyeGroupName(pkgGen.SnakeStyleMiddleware)
 	if err != nil {
@@ -367,6 +421,31 @@ func (pkgGen *HttpPackageGenerator) genRouter(pkg *HttpPackage, root *RouterNode
 		router.HandlerPackages = handlerMap
 	}
 
+	if pkgGen.SnakeStyleMiddleware { // unique middleware name for SnakeStyleMiddleware
+		mws := []string{}
+		hook := func(layer int, node *RouterNode) error {
+			if len(node.Children) == 0 {
+				return nil
+			}
+			groupMwName := node.GroupMiddleware
+			handlerMwName := node.HandlerMiddleware
+			if len(groupMwName) != 0 {
+				mws, groupMwName = appendMw(mws, groupMwName)
+			}
+			if len(handlerMwName) != 0 {
+				mws, handlerMwName = appendMw(mws, handlerMwName)
+			}
+			if groupMwName != node.GroupMiddleware {
+				node.GroupMiddleware = groupMwName
+			}
+			if handlerMwName != node.HandlerMiddleware {
+				node.HandlerMiddleware = handlerMwName
+			}
+			return nil
+		}
+		root.DFS(0, hook)
+	}
+
 	// store router info
 	pkg.RouterInfo = &router
 
diff --git a/cmd/hz/protobuf/plugin.go b/cmd/hz/protobuf/plugin.go
index 95d2473..d6f775c 100644
--- a/cmd/hz/protobuf/plugin.go
+++ b/cmd/hz/protobuf/plugin.go
@@ -621,6 +621,7 @@ func (plugin *Plugin) genHttpPackage(ast *descriptorpb.FileDescriptorProto, deps
 		BaseDomain:           args.BaseDomain,
 		QueryEnumAsInt:       args.QueryEnumAsInt,
 		SnakeStyleMiddleware: args.SnakeStyleMiddleware,
+		SortRouter:           args.SortRouter,
 	}
 
 	if args.ModelBackend != "" {
diff --git a/cmd/hz/thrift/plugin.go b/cmd/hz/thrift/plugin.go
index 956a0a6..91fab02 100644
--- a/cmd/hz/thrift/plugin.go
+++ b/cmd/hz/thrift/plugin.go
@@ -152,6 +152,7 @@ func (plugin *Plugin) Run() int {
 		BaseDomain:           args.BaseDomain,
 		QueryEnumAsInt:       args.QueryEnumAsInt,
 		SnakeStyleMiddleware: args.SnakeStyleMiddleware,
+		SortRouter:           args.SortRouter,
 	}
 	if args.ModelBackend != "" {
 		sg.Backend = meta.Backend(args.ModelBackend)
diff --git a/pkg/app/context.go b/pkg/app/context.go
index 9117e3b..369ece7 100644
--- a/pkg/app/context.go
+++ b/pkg/app/context.go
@@ -236,6 +236,17 @@ type RequestContext struct {
 
 	binder    binding.Binder
 	validator binding.StructValidator
+	exiled    bool
+}
+
+// Exile marks this RequestContext as not to be recycled.
+// Experimental features: Use with caution, it may have a slight impact on performance.
+func (ctx *RequestContext) Exile() {
+	ctx.exiled = true
+}
+
+func (ctx *RequestContext) IsExiled() bool {
+	return ctx.exiled
 }
 
 // Flush is the shortcut for ctx.Response.GetHijackWriter().Flush().
@@ -1227,6 +1238,10 @@ func (ctx *RequestContext) Cookie(key string) []byte {
 //	4. ctx.SetCookie("user", "", 10, "/", "localhost",protocol.CookieSameSiteLaxMode, false, false)
 //	add response header --->  Set-Cookie: user=; max-age=10; domain=localhost; path=/; SameSite=Lax;
 func (ctx *RequestContext) SetCookie(name, value string, maxAge int, path, domain string, sameSite protocol.CookieSameSite, secure, httpOnly bool) {
+	ctx.setCookie(name, value, maxAge, path, domain, sameSite, secure, httpOnly, false)
+}
+
+func (ctx *RequestContext) setCookie(name, value string, maxAge int, path, domain string, sameSite protocol.CookieSameSite, secure, httpOnly, partitioned bool) {
 	if path == "" {
 		path = "/"
 	}
@@ -1240,9 +1255,20 @@ func (ctx *RequestContext) SetCookie(name, value string, maxAge int, path, domai
 	cookie.SetSecure(secure)
 	cookie.SetHTTPOnly(httpOnly)
 	cookie.SetSameSite(sameSite)
+	cookie.SetPartitioned(partitioned)
 	ctx.Response.Header.SetCookie(cookie)
 }
 
+// SetPartitionedCookie adds a partitioned cookie to the Response's headers.
+// Use protocol.CookieSameSiteNoneMode for cross-site cookies to work.
+//
+// Usage: ctx.SetPartitionedCookie("user", "name", 10, "/", "localhost", protocol.CookieSameSiteNoneMode, true, true)
+//
+// This adds the response header: Set-Cookie: user=name; Max-Age=10; Domain=localhost; Path=/; HttpOnly; Secure; SameSite=None; Partitioned
+func (ctx *RequestContext) SetPartitionedCookie(name, value string, maxAge int, path, domain string, sameSite protocol.CookieSameSite, secure, httpOnly bool) {
+	ctx.setCookie(name, value, maxAge, path, domain, sameSite, secure, httpOnly, true)
+}
+
 // UserAgent returns the value of the request user_agent.
 func (ctx *RequestContext) UserAgent() []byte {
 	return ctx.Request.Header.UserAgent()
diff --git a/pkg/protocol/http1/server.go b/pkg/protocol/http1/server.go
index 535d1d5..3ea6596 100644
--- a/pkg/protocol/http1/server.go
+++ b/pkg/protocol/http1/server.go
@@ -133,6 +133,11 @@ func (s Server) Serve(c context.Context, conn network.Conn) (err error) {
 			zr.Release() //nolint:errcheck
 			zr = nil
 		}
+
+		if ctx.IsExiled() {
+			return
+		}
+
 		ctx.Reset()
 		s.Core.GetCtxPool().Put(ctx)
 	}()


Test patch:
diff --git a/cmd/hz/thrift/plugin_test.go b/cmd/hz/thrift/plugin_test.go
index 5e471b6..1cd4d8b 100644
--- a/cmd/hz/thrift/plugin_test.go
+++ b/cmd/hz/thrift/plugin_test.go
@@ -79,12 +79,21 @@ func TestRun(t *testing.T) {
 		HandlerDir: handlerDir,
 		RouterDir:  routerDir,
 		ModelDir:   modelDir,
+		UseDir:     args.Use,
 		ClientDir:  clientDir,
 		TemplateGenerator: generator.TemplateGenerator{
 			OutputDir: args.OutDir,
+			Excludes:  args.Excludes,
 		},
-		ProjPackage: pkg,
-		Options:     options,
+		ProjPackage:          pkg,
+		Options:              options,
+		HandlerByMethod:      args.HandlerByMethod,
+		CmdType:              args.CmdType,
+		ForceClientDir:       args.ForceClientDir,
+		BaseDomain:           args.BaseDomain,
+		QueryEnumAsInt:       args.QueryEnumAsInt,
+		SnakeStyleMiddleware: args.SnakeStyleMiddleware,
+		SortRouter:           args.SortRouter,
 	}
 	if args.ModelBackend != "" {
 		sg.Backend = meta.Backend(args.ModelBackend)
diff --git a/pkg/app/context_test.go b/pkg/app/context_test.go
index 7cbe71c..1317a06 100644
--- a/pkg/app/context_test.go
+++ b/pkg/app/context_test.go
@@ -1615,8 +1615,14 @@ func TestSetBinder(t *testing.T) {
 
 func TestRequestContext_SetCookie(t *testing.T) {
 	c := NewContext(0)
-	c.SetCookie("user", "hertz", 1, "/", "localhost", protocol.CookieSameSiteLaxMode, true, true)
-	assert.DeepEqual(t, "user=hertz; max-age=1; domain=localhost; path=/; HttpOnly; secure; SameSite=Lax", c.Response.Header.Get("Set-Cookie"))
+	c.SetCookie("user", "hertz", 1, "/", "localhost", protocol.CookieSameSiteNoneMode, true, true)
+	assert.DeepEqual(t, "user=hertz; max-age=1; domain=localhost; path=/; HttpOnly; secure; SameSite=None", c.Response.Header.Get("Set-Cookie"))
+}
+
+func TestRequestContext_SetPartitionedCookie(t *testing.T) {
+	c := NewContext(0)
+	c.SetPartitionedCookie("user", "hertz", 1, "/", "localhost", protocol.CookieSameSiteNoneMode, true, true)
+	assert.DeepEqual(t, "user=hertz; max-age=1; domain=localhost; path=/; HttpOnly; secure; SameSite=None; Partitioned", c.Response.Header.Get("Set-Cookie"))
 }
 
 func TestRequestContext_SetCookiePathEmpty(t *testing.T) {