feat(streamv2): support ttheader streaming

client/client.go

@@ -25,6 +25,8 @@ import (
 	"strconv"
 	"sync/atomic"
 
+	"github.com/cloudwego/kitex/pkg/streamx"
+
 	"github.com/bytedance/gopkg/cloud/metainfo"
 	"github.com/cloudwego/localsession/backup"
 
@@ -70,8 +72,14 @@ type kClient struct {
 	mws     []endpoint.Middleware
 	eps     endpoint.Endpoint
 	sEps    endpoint.Endpoint
-	opt     *client.Options
-	lbf     *lbcache.BalancerFactory
+
+	// streamx
+	sxStreamMW     streamx.StreamMiddleware
+	sxStreamRecvMW streamx.StreamRecvMiddleware
+	sxStreamSendMW streamx.StreamSendMiddleware
+
+	opt *client.Options
+	lbf *lbcache.BalancerFactory
 
 	inited bool
 	closed bool
@@ -424,21 +432,30 @@ func (kc *kClient) richRemoteOption() {
 		// (newClientStreamer: call WriteMeta before remotecli.NewClient)
 		transInfoHdlr := bound.NewTransMetaHandler(kc.opt.MetaHandlers)
 		kc.opt.RemoteOpt.PrependBoundHandler(transInfoHdlr)
+
+		// add meta handlers into streaming meta handlers
+		for _, h := range kc.opt.MetaHandlers {
+			if shdlr, ok := h.(remote.StreamingMetaHandler); ok {
+				kc.opt.RemoteOpt.StreamingMetaHandlers = append(kc.opt.RemoteOpt.StreamingMetaHandlers, shdlr)
+			}
+		}
 	}
 }
 
 func (kc *kClient) buildInvokeChain() error {
+	mwchain := endpoint.Chain(kc.mws...)
+
 	innerHandlerEp, err := kc.invokeHandleEndpoint()
 	if err != nil {
 		return err
 	}
-	kc.eps = endpoint.Chain(kc.mws...)(innerHandlerEp)
+	kc.eps = mwchain(innerHandlerEp)
 
 	innerStreamingEp, err := kc.invokeStreamingEndpoint()
 	if err != nil {
 		return err
 	}
-	kc.sEps = endpoint.Chain(kc.mws...)(innerStreamingEp)
+	kc.sEps = mwchain(innerStreamingEp)
 	return nil
 }
 
@@ -720,6 +737,9 @@ func initRPCInfo(ctx context.Context, method string, opt *client.Options, svcInf
 		rpcStats.ImmutableView(),
 	)
 
+	if mi != nil {
+		ri.Invocation().(rpcinfo.InvocationSetter).SetStreamingMode(mi.StreamingMode())
+	}
 	if fromMethod := ctx.Value(consts.CtxKeyMethod); fromMethod != nil {
 		rpcinfo.AsMutableEndpointInfo(ri.From()).SetMethod(fromMethod.(string))
 	}
@@ -732,6 +752,12 @@ func initRPCInfo(ctx context.Context, method string, opt *client.Options, svcInf
 		}
 	}
 
+	// streamx config
+	sopt := opt.StreamX
+	if sopt.RecvTimeout > 0 {
+		cfg.SetStreamRecvTimeout(sopt.RecvTimeout)
+	}
+
 	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)
 
 	if callOpts != nil && callOpts.CompressorName != "" {

client/client_streamx.go

@@ -0,0 +1,85 @@
+/*
+ * Copyright 2024 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package client
+
+import (
+	"context"
+
+	"github.com/cloudwego/kitex/client/streamxclient/streamxcallopt"
+	"github.com/cloudwego/kitex/pkg/rpcinfo"
+	"github.com/cloudwego/kitex/pkg/streamx"
+)
+
+// NewStream create stream for streamx mode
+func (kc *kClient) NewStream(ctx context.Context, method string, streamArgs streamx.StreamArgs, callOptions ...streamxcallopt.CallOption) (context.Context, streamx.ClientStream, error) {
+	if !kc.inited {
+		panic("client not initialized")
+	}
+	if kc.closed {
+		panic("client is already closed")
+	}
+	if ctx == nil {
+		panic("ctx is nil")
+	}
+	var ri rpcinfo.RPCInfo
+	ctx, ri, _ = kc.initRPCInfo(ctx, method, 0, nil)
+
+	err := rpcinfo.AsMutableRPCConfig(ri.Config()).SetInteractionMode(rpcinfo.Streaming)
+	if err != nil {
+		return nil, nil, err
+	}
+	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)
+
+	// tracing
+	ctx = kc.opt.TracerCtl.DoStart(ctx, ri)
+	ctx, copts := streamxcallopt.NewCtxWithCallOptions(ctx)
+	callOptions = append(callOptions, streamxcallopt.WithStreamCloseCallback(func() {
+		kc.opt.TracerCtl.DoFinish(ctx, ri, err)
+	}))
+	copts.Apply(callOptions)
+
+	if msargs := streamx.AsMutableStreamArgs(streamArgs); msargs != nil {
+		msargs.SetStreamMiddleware(kc.sxStreamMW)
+
+		eventHandler := kc.opt.TracerCtl.GetStreamEventHandler()
+		if eventHandler == nil {
+			msargs.SetStreamRecvMiddleware(kc.sxStreamRecvMW)
+			msargs.SetStreamSendMiddleware(kc.sxStreamSendMW)
+		} else {
+			traceRecvMW := streamx.NewStreamRecvStatMiddleware(ctx, eventHandler)
+			traceSendMW := streamx.NewStreamSendStatMiddleware(ctx, eventHandler)
+			if kc.sxStreamRecvMW == nil {
+				msargs.SetStreamRecvMiddleware(traceRecvMW)
+			} else {
+				msargs.SetStreamRecvMiddleware(streamx.StreamRecvMiddlewareChain(traceRecvMW, kc.sxStreamRecvMW))
+			}
+			if kc.sxStreamSendMW == nil {
+				msargs.SetStreamSendMiddleware(traceSendMW)
+			} else {
+				msargs.SetStreamSendMiddleware(streamx.StreamSendMiddlewareChain(traceSendMW, kc.sxStreamSendMW))
+			}
+		}
+	}
+	// with streamx mode, req is nil and resp is streamArgs
+	// it's an ugly trick but if we don't want to refactor too much,
+	// this is the only way to compatible with current endpoint API.
+	err = kc.sEps(ctx, nil, streamArgs)
+	if err != nil {
+		return nil, nil, err
+	}
+	return ctx, streamArgs.Stream().(streamx.ClientStream), nil
+}

client/stream.go

@@ -24,9 +24,9 @@ import (
 	"github.com/cloudwego/kitex/pkg/kerrors"
 	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata"
 	"github.com/cloudwego/kitex/pkg/serviceinfo"
+	"github.com/cloudwego/kitex/pkg/streamx"
 
 	"github.com/cloudwego/kitex/pkg/endpoint"
-	"github.com/cloudwego/kitex/pkg/remote"
 	"github.com/cloudwego/kitex/pkg/remote/remotecli"
 	"github.com/cloudwego/kitex/pkg/rpcinfo"
 	"github.com/cloudwego/kitex/pkg/streaming"
@@ -75,24 +75,41 @@ func (kc *kClient) invokeStreamingEndpoint() (endpoint.Endpoint, error) {
 	if err != nil {
 		return nil, err
 	}
-	for _, h := range kc.opt.MetaHandlers {
-		if shdlr, ok := h.(remote.StreamingMetaHandler); ok {
-			kc.opt.RemoteOpt.StreamingMetaHandlers = append(kc.opt.RemoteOpt.StreamingMetaHandlers, shdlr)
-		}
-	}
 
+	// old version streaming mw
 	recvEndpoint := kc.opt.Streaming.BuildRecvInvokeChain(kc.invokeRecvEndpoint())
 	sendEndpoint := kc.opt.Streaming.BuildSendInvokeChain(kc.invokeSendEndpoint())
 
+	// streamx version streaming mw
+	kc.sxStreamMW = streamx.StreamMiddlewareChain(kc.opt.StreamX.StreamMWs...)
+	if len(kc.opt.StreamX.StreamRecvMWs) > 0 {
+		kc.sxStreamRecvMW = streamx.StreamRecvMiddlewareChain(kc.opt.StreamX.StreamRecvMWs...)
+	}
+	if len(kc.opt.StreamX.StreamSendMWs) > 0 {
+		kc.sxStreamSendMW = streamx.StreamSendMiddlewareChain(kc.opt.StreamX.StreamSendMWs...)
+	}
+
 	return func(ctx context.Context, req, resp interface{}) (err error) {
 		// req and resp as &streaming.Stream
 		ri := rpcinfo.GetRPCInfo(ctx)
 		st, scm, err := remotecli.NewStream(ctx, ri, handler, kc.opt.RemoteOpt)
 		if err != nil {
 			return
 		}
-		clientStream := newStream(st, scm, kc, ri, kc.getStreamingMode(ri), sendEndpoint, recvEndpoint)
-		resp.(*streaming.Result).Stream = clientStream
+
+		// streamx API
+		if cs, ok := st.(streamx.Stream); ok {
+			streamArgs := resp.(streamx.StreamArgs)
+			// the middlewares can get real Stream after set
+			streamx.AsMutableStreamArgs(streamArgs).SetStream(cs)
+			return nil
+		}
+
+		// old version streaming
+		if cs, ok := st.(streaming.Stream); ok {
+			clientStream := newStream(cs, scm, kc, ri, kc.getStreamingMode(ri), sendEndpoint, recvEndpoint)
+			resp.(*streaming.Result).Stream = clientStream
+		}
 		return
 	}, nil
 }

client/stream_test.go

@@ -104,7 +104,7 @@ func TestStreaming(t *testing.T) {
 	cliInfo.ConnPool = connpool
 	s, cr, _ := remotecli.NewStream(ctx, mockRPCInfo, new(mocks.MockCliTransHandler), cliInfo)
 	stream := newStream(
-		s, cr, kc, mockRPCInfo, serviceinfo.StreamingBidirectional,
+		s.(streaming.Stream), cr, kc, mockRPCInfo, serviceinfo.StreamingBidirectional,
 		func(stream streaming.Stream, message interface{}) (err error) {
 			return stream.SendMsg(message)
 		},

client/streamxclient/client_gen.go

@@ -0,0 +1,102 @@
+/*
+ * Copyright 2024 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package streamxclient
+
+import (
+	"context"
+	"errors"
+
+	"github.com/cloudwego/kitex/client"
+	"github.com/cloudwego/kitex/client/streamxclient/streamxcallopt"
+	"github.com/cloudwego/kitex/internal/streamx/streamxclient"
+	"github.com/cloudwego/kitex/pkg/serviceinfo"
+	"github.com/cloudwego/kitex/pkg/streamx"
+)
+
+// InvokeStream create a new client stream and wrapped related middlewares
+func InvokeStream[Req, Res any](
+	ctx context.Context, cli client.Client, smode serviceinfo.StreamingMode, method string,
+	req *Req, res *Res, callOptions ...streamxcallopt.CallOption,
+) (context.Context, *streamx.GenericClientStream[Req, Res], error) {
+	reqArgs, resArgs := streamx.NewStreamReqArgs(nil), streamx.NewStreamResArgs(nil)
+	streamArgs := streamx.NewStreamArgs(nil)
+	// important notes: please don't set a typed nil value into interface arg like NewStreamReqArgs({typ: *Res, ptr: nil})
+	// otherwise, reqArgs.Req() may not ==nil forever
+	if req != nil {
+		reqArgs.SetReq(req)
+	}
+	if res != nil {
+		resArgs.SetRes(res)
+	}
+
+	scli, ok := cli.(streamxclient.StreamXClient)
+	if !ok {
+		return nil, nil, errors.New("current client is not support streamx interface")
+	}
+	// NewStream should register client middlewares into stream Args
+	ctx, cs, err := scli.NewStream(ctx, method, streamArgs, callOptions...)
+	if err != nil {
+		return nil, nil, err
+	}
+	stream := streamx.NewGenericClientStream[Req, Res](cs)
+	var streamMW streamx.StreamMiddleware
+	if streamMWs, ok := streamArgs.(streamx.StreamMiddlewaresArgs); ok {
+		var recvMW streamx.StreamRecvMiddleware
+		var sendMW streamx.StreamSendMiddleware
+		streamMW, recvMW, sendMW = streamMWs.Middlewares()
+		stream.SetStreamRecvMiddleware(recvMW)
+		stream.SetStreamSendMiddleware(sendMW)
+	}
+
+	streamInvoke := func(ctx context.Context, streamArgs streamx.StreamArgs, reqArgs streamx.StreamReqArgs, resArgs streamx.StreamResArgs) (err error) {
+		// assemble streaming args depend on each stream mode
+		switch smode {
+		case serviceinfo.StreamingUnary:
+			// client should call CloseSend and server should call SendAndClose
+			if err = stream.SendMsg(ctx, req); err != nil {
+				return err
+			}
+			if err = stream.RecvMsg(ctx, res); err != nil {
+				return err
+			}
+			resArgs.SetRes(res)
+			if err = stream.CloseSend(ctx); err != nil {
+				return err
+			}
+		case serviceinfo.StreamingClient:
+		case serviceinfo.StreamingServer:
+			if err = stream.SendMsg(ctx, req); err != nil {
+				return err
+			}
+			if err = stream.CloseSend(ctx); err != nil {
+				return err
+			}
+		case serviceinfo.StreamingBidirectional:
+		default:
+		}
+		return nil
+	}
+	if streamMW != nil {
+		err = streamMW(streamInvoke)(ctx, streamArgs, reqArgs, resArgs)
+	} else {
+		err = streamInvoke(ctx, streamArgs, reqArgs, resArgs)
+	}
+	if err != nil {
+		return nil, nil, err
+	}
+	return ctx, stream, nil
+}

client/streamxclient/client_option.go

@@ -0,0 +1,53 @@
+/*
+ * Copyright 2024 CloudWeGo Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package streamxclient
+
+import (
+	"time"
+
+	internal_client "github.com/cloudwego/kitex/internal/client"
+	"github.com/cloudwego/kitex/pkg/streamx"
+	"github.com/cloudwego/kitex/pkg/utils"
+)
+
+// Deprecated: Note that it maybe refactor in the next version
+func WithProvider(pvd streamx.ClientProvider) internal_client.Option {
+	return internal_client.Option{F: func(o *internal_client.Options, di *utils.Slice) {
+		o.RemoteOpt.Provider = pvd
+	}}
+}
+
+// WithStreamRecvTimeout add recv timeout for stream.Recv function
+func WithStreamRecvTimeout(timeout time.Duration) internal_client.Option {
+	return internal_client.Option{F: func(o *internal_client.Options, di *utils.Slice) {
+		o.StreamX.RecvTimeout = timeout
+	}}
+}
+
+// WithStreamRecvMiddleware add recv middleware
+func WithStreamRecvMiddleware(smw streamx.StreamRecvMiddleware) internal_client.Option {
+	return internal_client.Option{F: func(o *internal_client.Options, di *utils.Slice) {
+		o.StreamX.StreamRecvMWs = append(o.StreamX.StreamRecvMWs, smw)
+	}}
+}
+
+// WithStreamSendMiddleware add send middleware
+func WithStreamSendMiddleware(smw streamx.StreamSendMiddleware) internal_client.Option {
+	return internal_client.Option{F: func(o *internal_client.Options, di *utils.Slice) {
+		o.StreamX.StreamSendMWs = append(o.StreamX.StreamSendMWs, smw)
+	}}
+}