feat: server middleware 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 @@ func WithRefuseTrafficWithoutServiceName() Option {
 // 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 @@ import (
 	"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 @@ func (s *server) invokeHandleEndpoint() endpoint.Endpoint {
 			// 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 @@ func (s *server) invokeHandleEndpoint() endpoint.Endpoint {
 	}
 }
 
+// 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