[WIP] feat: server middleware context pass through

internal/server/option.go

@@ -100,6 +100,8 @@ type Options struct {
 
  RefuseTrafficWithoutServiceName bool
  EnableContextTimeout bool
+
+ EnableStreamingContextPassThrough bool
 }
 
 type Limit struct {

server/option.go

@@ -386,6 +386,15 @@
 // so there's no need to use this option.
 func WithEnableContextTimeout(enable bool) Option {
  return Option{F: func(o *internal_server.Options, di *utils.Slice) {
+ di.Push(fmt.Sprintf("WithEnableContextTimeout({%+v})", enable))
  o.EnableContextTimeout = enable
  }}
 }
+
+// WithEnableStreamingContextPassThrough enables passing through context modified by server middleware to handler.
+func WithEnableStreamingContextPassThrough() Option {
+ return Option{F: func(o *internal_server.Options, di *utils.Slice) {
+ di.Push(fmt.Sprintf("EnableStreamingContextPassThrough()"))
+ o.EnableStreamingContextPassThrough = true
+ }}
+}

server/server.go

@@ -45,6 +45,7 @@
  "github.com/cloudwego/kitex/pkg/rpcinfo"
  "github.com/cloudwego/kitex/pkg/serviceinfo"
  "github.com/cloudwego/kitex/pkg/stats"
+ "github.com/cloudwego/kitex/pkg/streaming"
 )
 
 // Server is an abstraction of an RPC server. It accepts connections and dispatches them to the service
@@ -337,7 +338,16 @@
  // clear session
  backup.ClearCtx()
  }()
- implHandlerFunc := svcInfo.MethodInfo(methodName).Handler()
+ methodInfo := svcInfo.MethodInfo(methodName)
+ if methodInfo.IsStreaming() && s.opt.EnableStreamingContextPassThrough {
+ if sArg, ok := args.(*streaming.Args); ok {
+ sArg.Stream = &contextStream{
+ Stream: sArg.Stream,
+ ctx: ctx,
+ }
+ }
+ }
+ implHandlerFunc := methodInfo.Handler()
  rpcinfo.Record(ctx, ri, stats.ServerHandleStart, nil)
  // set session
  backup.BackupCtx(ctx)
@@ -355,6 +365,17 @@
  }
 }
 
+// contextStream is used to return the ctx node passed in by server middlewares which might contain new key/values
+type contextStream struct {
+ streaming.Stream
+ ctx context.Context
+}
+
+// Context implements streaming.Stream
+func (s *contextStream) Context() context.Context {
+ return s.ctx
+}
+
 func (s *server) initBasicRemoteOption() {
  remoteOpt := s.opt.RemoteOpt
  remoteOpt.TargetSvcInfo = s.targetSvcInfo