feat: add option and mdws for localsession

cloudwego · Aug 15, 2023 · ae937ec · ae937ec
1 parent bf39157
commit ae937ec
Show file tree

Hide file tree

Showing 9 changed files with 88 additions and 7 deletions.
diff --git a/go.mod b/go.mod
@@ -4,13 +4,14 @@ go 1.16
 
 require (
 	github.com/bytedance/go-tagexpr/v2 v2.9.2
-	github.com/bytedance/gopkg v0.0.0-20220413063733-65bf48ffb3a7
+	github.com/bytedance/gopkg v0.0.0-20230728082804-614d0af6619b
 	github.com/bytedance/mockey v1.2.1
 	github.com/bytedance/sonic v1.8.1
+	github.com/cloudwego/localsession v0.0.2
 	github.com/cloudwego/netpoll v0.3.2
 	github.com/fsnotify/fsnotify v1.5.4
 	github.com/tidwall/gjson v1.13.0 // indirect
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
-	golang.org/x/sys v0.0.0-20220412211240-33da011f77ad
+	golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10
 	google.golang.org/protobuf v1.27.1
 )
diff --git a/go.sum b/go.sum
@@ -1,7 +1,8 @@
 github.com/bytedance/go-tagexpr/v2 v2.9.2 h1:QySJaAIQgOEDQBLS3x9BxOWrnhqu5sQ+f6HaZIxD39I=
 github.com/bytedance/go-tagexpr/v2 v2.9.2/go.mod h1:5qsx05dYOiUXOUgnQ7w3Oz8BYs2qtM/bJokdLb79wRM=
-github.com/bytedance/gopkg v0.0.0-20220413063733-65bf48ffb3a7 h1:PtwsQyQJGxf8iaPptPNaduEIu9BnrNms+pcRdHAxZaM=
 github.com/bytedance/gopkg v0.0.0-20220413063733-65bf48ffb3a7/go.mod h1:2ZlV9BaUH4+NXIBF0aMdKKAnHTzqH+iMU4KUjAbL23Q=
+github.com/bytedance/gopkg v0.0.0-20230728082804-614d0af6619b h1:R6PWoQtxEMpWJPHnpci+9LgFxCS7iJCfOGBvCgZeTKI=
+github.com/bytedance/gopkg v0.0.0-20230728082804-614d0af6619b/go.mod h1:FtQG3YbQG9L/91pbKSw787yBQPutC+457AvDW77fgUQ=
 github.com/bytedance/mockey v1.2.1 h1:g84ngI88hz1DR4wZTL3yOuqlEcq67MretBfQUdXwrmw=
 github.com/bytedance/mockey v1.2.1/go.mod h1:+Jm/fzWZAuhEDrPXVjDf/jLM2BlLXJkwk94zf2JZ3X4=
 github.com/bytedance/sonic v1.5.0/go.mod h1:ED5hyg4y6t3/9Ku1R6dU/4KyJ48DZ4jPhfY1O2AihPM=
@@ -10,6 +11,8 @@ github.com/bytedance/sonic v1.8.1/go.mod h1:i736AoUSYt75HyZLoJW9ERYxcy6eaN6h4BZX
 github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY=
 github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 h1:qSGYFH7+jGhDF8vLC+iwCD4WpbV1EBDSzWkJODFLams=
 github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311/go.mod h1:b583jCggY9gE99b6G5LEC39OIiVsWj+R97kbl5odCEk=
+github.com/cloudwego/localsession v0.0.2 h1:N9/IDtCPj1fCL9bCTP+DbXx3f40YjVYWcwkJG0YhQkY=
+github.com/cloudwego/localsession v0.0.2/go.mod h1:kiJxmvAcy4PLgKtEnPS5AXed3xCiXcs7Z+KBHP72Wv8=
 github.com/cloudwego/netpoll v0.3.2 h1:/998ICrNMVBo4mlul4j7qcIeY7QnEfuCCPPwck9S3X4=
 github.com/cloudwego/netpoll v0.3.2/go.mod h1:xVefXptcyheopwNDZjDPcfU6kIjZXZ4nY550k1yH9eQ=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -33,6 +36,10 @@ github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7
 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
 github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
+github.com/modern-go/gls v0.0.0-20220109145502-612d0167dce5 h1:uiS4zKYKJVj5F3ID+5iylfKPsEQmBEOucSD9Vgmn0i0=
+github.com/modern-go/gls v0.0.0-20220109145502-612d0167dce5/go.mod h1:I8AX+yW//L8Hshx6+a1m3bYkwXkpsVjA2795vP4f4oQ=
+github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
+github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
 github.com/nyaruka/phonenumbers v1.0.55 h1:bj0nTO88Y68KeUQ/n3Lo2KgK7lM1hF7L9NFuwcCl3yg=
 github.com/nyaruka/phonenumbers v1.0.55/go.mod h1:sDaTZ/KPX5f8qyV9qN+hIm+4ZBARJrupC6LuhshJq1U=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -65,13 +72,21 @@ golang.org/x/arch v0.0.0-20210923205945-b76863e36670 h1:18EFjUmQOcUvxNYSkA6jO9VA
 golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
+golang.org/x/net v0.0.0-20221014081412-f15817d10f9b h1:tvrvnPFcdzp294diPnrdZZZ8XUt2Tyj7svb7X52iDuU=
+golang.org/x/net v0.0.0-20221014081412-f15817d10f9b/go.mod h1:YDH+HFinaLZZlnHAfSS6ZXJJ9M9t4Dl22yv3iI2vPwk=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220110181412-a018aaa089fe/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220412211240-33da011f77ad h1:ntjMns5wyP/fN65tdBD4g8J5w8n015+iIIs9rtjXkY0=
 golang.org/x/sys v0.0.0-20220412211240-33da011f77ad/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10 h1:WIoqL4EROvwiPdUtaip4VcDdpZ4kha7wBWZrbVKCIZg=
+golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
+golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190328211700-ab21143f2384/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

diff --git a/pkg/app/client/option.go b/pkg/app/client/option.go
@@ -17,6 +17,7 @@
 package client
 
 import (
+	"context"
 	"crypto/tls"
 	"time"
 
@@ -174,6 +175,15 @@ func WithDialFunc(f network.DialFunc, dialers ...network.Dialer) config.ClientOp
 	}}
 }
 
+// WithContextBackup enables local-session to retrieve context which is backed up by server,
+// in case of user don't correctly pass context into next HTTP/RPC call.
+//   - backupHandler pass a handler to check and handler user-defined key-values according to current context, returning backup==false means no need further operations.
+func WithContextBackup(backupHandler func(prev, cur context.Context) (ctx context.Context, backup bool)) config.ClientOption {
+	return config.ClientOption{F: func(o *config.ClientOptions) {
+		o.CtxBackupHandler = backupHandler
+	}}
+}
+
 // customDialer set customDialerFunc and params to set dailFunc
 type customDialer struct {
 	network.Dialer

diff --git a/pkg/app/middlewares/client/localsession/localsession.go b/pkg/app/middlewares/client/localsession/localsession.go
@@ -0,0 +1,18 @@
+package localsession
+
+import (
+	"context"
+	"github.com/cloudwego/hertz/pkg/app/client"
+	"github.com/cloudwego/hertz/pkg/protocol"
+	"github.com/cloudwego/localsession/backup"
+)
+
+// CtxRecover will merge backup context if no metainfo found in ctx
+func CtxRecover(hdl backup.BackupHandler) client.Middleware {
+	return func(next client.Endpoint) client.Endpoint {
+		return func(ctx context.Context, req *protocol.Request, resp *protocol.Response) (err error) {
+			ctx = backup.RecoverCtxOnDemands(ctx, hdl)
+			return next(ctx, req, resp)
+		}
+	}
+}
diff --git a/pkg/app/middlewares/server/localsession/localsession.go b/pkg/app/middlewares/server/localsession/localsession.go
@@ -0,0 +1,17 @@
+package localsession
+
+import (
+	"context"
+	"github.com/cloudwego/hertz/pkg/app"
+	"github.com/cloudwego/localsession/backup"
+)
+
+func BackupSession(opts backup.Options) app.HandlerFunc {
+	backup.Init(opts)
+
+	return func(c context.Context, ctx *app.RequestContext) {
+		backup.BackupCtx(c)
+		defer backup.ClearCtx()
+		ctx.Next(c)
+	}
+}
diff --git a/pkg/app/server/hertz.go b/pkg/app/server/hertz.go
@@ -49,7 +49,6 @@ func New(opts ...config.Option) *Hertz {
 func Default(opts ...config.Option) *Hertz {
 	h := New(opts...)
 	h.Use(recovery.Recovery())
-
 	return h
 }
 

diff --git a/pkg/app/server/option.go b/pkg/app/server/option.go
@@ -19,6 +19,7 @@ package server
 import (
 	"context"
 	"crypto/tls"
+	"github.com/cloudwego/localsession/backup"
 	"net"
 	"strings"
 	"time"
@@ -339,10 +340,22 @@ func WithOnAccept(fn func(conn net.Conn) context.Context) config.Option {
 	}}
 }
 
-// WithOnConnect sets the onConnect function. It can received data from connection in netpoll.
+// WithOnConnect sets the onConnect function. It can receive data from connection in netpoll.
 // In go net, it will be called after converting tls connection.
 func WithOnConnect(fn func(ctx context.Context, conn network.Conn) context.Context) config.Option {
 	return config.Option{F: func(o *config.Options) {
 		o.OnConnect = fn
 	}}
 }
+
+// WithContextBackup enables local-session to back up hertz server's context,
+// in case of user don't correctly pass context into next HTTP/RPC call.
+//   - enable means enable context backup
+//   - async means backup session will propagate to children goroutines, otherwise it only works on the same goroutine handler
+func WithContextBackup(enable, async bool) config.Option {
+	return config.Option{F: func(o *config.Options) {
+		o.BackupOpt = backup.DefaultOptions()
+		o.BackupOpt.Enable = enable
+		o.BackupOpt.EnableImplicitlyTransmitAsync = async
+	}}
+}
diff --git a/pkg/common/config/client_option.go b/pkg/common/config/client_option.go
@@ -18,6 +18,7 @@ package config
 
 import (
 	"crypto/tls"
+	"github.com/cloudwego/localsession/backup"
 	"time"
 
 	"github.com/cloudwego/hertz/pkg/app/client/retry"
@@ -117,7 +118,7 @@ type ClientOptions struct {
 	// Disabled path normalization may be useful for proxying incoming requests
 	// to servers that are expecting paths to be forwarded as-is.
 	//
-	// By default path values are normalized, i.e.
+	// By default, path values are normalized, i.e.
 	// extra slashes are removed, special characters are encoded.
 	DisablePathNormalizing bool
 
@@ -128,6 +129,9 @@ type ClientOptions struct {
 
 	// StateObserve execution interval
 	ObservationInterval time.Duration
+
+	// local session for context backup
+	CtxBackupHandler backup.BackupHandler
 }
 
 func NewClientOptions(opts []ClientOption) *ClientOptions {

diff --git a/pkg/common/config/option.go b/pkg/common/config/option.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/cloudwego/hertz/pkg/app/server/registry"
 	"github.com/cloudwego/hertz/pkg/network"
+	"github.com/cloudwego/localsession/backup"
 )
 
 // Option is the only struct that can be used to set Options.
@@ -97,6 +98,9 @@ type Options struct {
 	// The HTML template will reload according to files' changing event
 	// otherwise it will reload after AutoReloadInterval.
 	AutoReloadInterval time.Duration
+
+	// local session for context backup
+	BackupOpt backup.Options
 }
 
 func (o *Options) Apply(opts []Option) {