diff --git a/client/callopt/streamcall/call_options.go b/client/callopt/streamcall/call_options.go index 5de499335e..ec5c94d339 100644 --- a/client/callopt/streamcall/call_options.go +++ b/client/callopt/streamcall/call_options.go @@ -78,7 +78,7 @@ func WithRecvTimeout(d time.Duration) Option { // DisableCancelRemote: true, // }) // -// The remote peer must promptly exit the handler; otherwise, there is a risk of stream leakage! +// IMPORTANT: The remote peer must promptly exit the handler; otherwise, there is a risk of stream and Recv goroutine leakage! func WithRecvTimeoutConfig(cfg streaming.TimeoutConfig) Option { return Option{f: func(o *callopt.CallOptions, di *strings.Builder) { fmt.Fprintf(di, "WithRecvTimeoutConfig(+%v)", cfg) diff --git a/client/option_stream.go b/client/option_stream.go index 57c306a005..6569645139 100644 --- a/client/option_stream.go +++ b/client/option_stream.go @@ -68,7 +68,7 @@ func WithStreamRecvTimeout(d time.Duration) StreamOption { // DisableCancelRemote: true, // }) // -// The remote peer must promptly exit the handler; otherwise, there is a risk of stream leakage! +// IMPORTANT: The remote peer must promptly exit the handler; otherwise, there is a risk of stream and Recv goroutine leakage! func WithStreamRecvTimeoutConfig(cfg streaming.TimeoutConfig) StreamOption { return StreamOption{F: func(o *client.StreamOptions, di *utils.Slice) { di.Push(fmt.Sprintf("WithStreamRecvTimeoutConfig(%+v)", cfg)) diff --git a/client/stream.go b/client/stream.go index 2f8be11c8a..6db3df4480 100644 --- a/client/stream.go +++ b/client/stream.go @@ -21,6 +21,7 @@ import ( "fmt" "io" "runtime/debug" + "sync" "sync/atomic" "time" @@ -35,6 +36,7 @@ import ( "github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/codes" "github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata" "github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/status" + "github.com/cloudwego/kitex/pkg/remote/trans/ttstream" "github.com/cloudwego/kitex/pkg/rpcinfo" "github.com/cloudwego/kitex/pkg/serviceinfo" "github.com/cloudwego/kitex/pkg/streaming" @@ -216,6 +218,9 @@ type stream struct { streamingMode serviceinfo.StreamingMode finished uint32 isGRPC bool + + finishedErrOnce sync.Once + finishedErr atomic.Value } var ( @@ -254,6 +259,9 @@ func newStream(ctx context.Context, s streaming.ClientStream, scm *remotecli.Str // Header returns the header data sent by the server if any. func (s *stream) Header() (hd streaming.Header, err error) { + if atomic.LoadUint32(&s.finished) == 1 { + return nil, s.finishedErr.Load().(error) + } if hd, err = s.ClientStream.Header(); err != nil { s.DoFinish(err) } @@ -263,6 +271,9 @@ func (s *stream) Header() (hd streaming.Header, err error) { // RecvMsg receives a message from the server. // If an error is returned, stream.DoFinish() will be called to record the end of stream func (s *stream) RecvMsg(ctx context.Context, m interface{}) (err error) { + if atomic.LoadUint32(&s.finished) == 1 { + return s.finishedErr.Load().(error) + } if !s.recv.EqualsTo(recvEndpoint) { // If the values are not equal, it indicates the presence of custom middleware. // To prevent errors caused by middleware code that relies on rpcinfo when users @@ -293,19 +304,27 @@ func (s *stream) handleStreamRecvEvent(err error) { } func (s *stream) recvWithTimeout(ctx context.Context, m interface{}) error { - if !s.isGRPC || s.recvTmCfg.Timeout <= 0 { + if s.recvTmCfg.Timeout <= 0 { return s.recv(ctx, s.ClientStream, m) } + buildTmErr := buildTTStreamRecvTimeoutErr + buildPanicErr := buildTTStreamRecvPanicErr + if s.isGRPC { + buildTmErr = buildGRPCRecvTimeoutErr + buildPanicErr = buildGRPCRecvPanicErr + } return callWithTimeout(s.recvTmCfg, func() error { return s.recv(ctx, s.ClientStream, m) }, s.cancel, + buildTmErr, + buildPanicErr, ) } func (s *stream) cancel(err error) { - // for now, only gRPC ClientStream implements CancelableClientStream interface + // ClientStream of gRPC and ttstream both implements CancelableClientStream interface if c, ok := s.ClientStream.(internal_stream.CancelableClientStream); ok { c.CancelWithErr(err) } @@ -314,6 +333,9 @@ func (s *stream) cancel(err error) { // SendMsg sends a message to the server. // If an error is returned, stream.DoFinish() will be called to record the end of stream func (s *stream) SendMsg(ctx context.Context, m interface{}) (err error) { + if atomic.LoadUint32(&s.finished) == 1 { + return s.finishedErr.Load().(error) + } if !s.send.EqualsTo(sendEndpoint) { // same with RecvMsg ri := rpcinfo.GetRPCInfo(ctx) @@ -338,6 +360,17 @@ func (s *stream) handleStreamSendEvent(err error) { // DoFinish implements the streaming.WithDoFinish interface, and it records the end of stream // It will release the connection. func (s *stream) DoFinish(err error) { + s.finishedErrOnce.Do(func() { + // store the finished err so that subsequent Recv/Send/Header calls would fail fast + // and return the same finished err. + // When err is nil (e.g. client streaming success), use io.EOF as the sentinel + // since the stream is done and no more messages can be received. + if err != nil { + s.finishedErr.Store(err) + } else { + s.finishedErr.Store(io.EOF) + } + }) if atomic.SwapUint32(&s.finished, 1) == 1 { // already called return @@ -413,6 +446,8 @@ func (s *grpcStream) recvWithTimeout(m interface{}) error { return s.recvEndpoint(s.Stream, m) }, s.st.cancel, + buildGRPCRecvTimeoutErr, + buildGRPCRecvPanicErr, ) } @@ -429,7 +464,11 @@ func (s *grpcStream) DoFinish(err error) { s.st.DoFinish(err) } -func callWithTimeout(tmCfg streaming.TimeoutConfig, call func() error, cancel func(error)) error { +func callWithTimeout(tmCfg streaming.TimeoutConfig, + call func() error, cancel func(error), + buildTmErr func(streaming.TimeoutConfig) error, + buildPanicErr func(interface{}, []byte) error, +) error { timer := time.NewTimer(tmCfg.Timeout) defer timer.Stop() finishChan := make(chan error, 1) @@ -437,7 +476,7 @@ func callWithTimeout(tmCfg streaming.TimeoutConfig, call func() error, cancel fu var callErr error defer func() { if r := recover(); r != nil { - callErr = status.Errorf(codes.Internal, "stream Recv panic, panic=%v, stack=%s", r, debug.Stack()) + callErr = buildPanicErr(r, debug.Stack()) cancel(callErr) } finishChan <- callErr @@ -446,7 +485,8 @@ func callWithTimeout(tmCfg streaming.TimeoutConfig, call func() error, cancel fu }) select { case <-timer.C: - err := status.Errorf(codes.RecvDeadlineExceeded, recvTimeoutErrTpl, tmCfg) + err := buildTmErr(tmCfg) + // if DisableCancelRemote == true, users are responsible for ensuring that the stream is not leaked. if !tmCfg.DisableCancelRemote { // finish the stream lifecycle so that the goroutine could exit cancel(err) @@ -469,6 +509,22 @@ func isRPCError(err error) bool { return !isBizStatusError } +func buildGRPCRecvTimeoutErr(tmCfg streaming.TimeoutConfig) error { + return status.Errorf(codes.RecvDeadlineExceeded, recvTimeoutErrTpl, tmCfg) +} + +func buildGRPCRecvPanicErr(r interface{}, stack []byte) error { + return status.Errorf(codes.Internal, "stream Recv panic, panic=%v, stack=%s", r, stack) +} + +func buildTTStreamRecvTimeoutErr(tmCfg streaming.TimeoutConfig) error { + return ttstream.NewStreamRecvTimeoutException(tmCfg, true) +} + +func buildTTStreamRecvPanicErr(r interface{}, stack []byte) error { + return ttstream.NewStreamInternalException(fmt.Sprintf("stream Recv panic, panic=%v, stack=%s", r, stack), true) +} + var ( recvEndpoint cep.StreamRecvEndpoint = func(ctx context.Context, stream streaming.ClientStream, m interface{}) error { return stream.RecvMsg(ctx, m) diff --git a/client/stream_test.go b/client/stream_test.go index 67cecc4918..6657448e7f 100644 --- a/client/stream_test.go +++ b/client/stream_test.go @@ -145,9 +145,11 @@ func TestStreaming(t *testing.T) { func(ctx context.Context, stream streaming.ClientStream, message interface{}) (err error) { return stream.RecvMsg(ctx, message) }, + //nolint:staticcheck // SA1019: just for testing func(stream streaming.Stream, message interface{}) (err error) { return stream.SendMsg(message) }, + //nolint:staticcheck // SA1019: just for testing func(stream streaming.Stream, message interface{}) (err error) { return stream.RecvMsg(message) }, @@ -242,7 +244,9 @@ func (s *mockStream) CloseSend(ctx context.Context) error { } func (s *mockStream) CancelWithErr(err error) { - s.cancelWithErr(err) + if s.cancelWithErr != nil { + s.cancelWithErr(err) + } } func (s *mockStream) Close() error { @@ -280,9 +284,12 @@ func Test_newStream(t *testing.T) { }, func(ctx context.Context, stream streaming.ClientStream, message interface{}) (err error) { return recvErr }, + //nolint:staticcheck // SA1019: just for testing func(stream streaming.Stream, message interface{}) (err error) { return sendErr - }, func(stream streaming.Stream, message interface{}) (err error) { + }, + //nolint:staticcheck // SA1019: just for testing + func(stream streaming.Stream, message interface{}) (err error) { return recvErr }, ) @@ -291,7 +298,9 @@ func Test_newStream(t *testing.T) { test.Assert(t, s.kc == kc) test.Assert(t, s.streamingMode == serviceinfo.StreamingClient) test.Assert(t, s.SendMsg(context.Background(), nil) == sendErr) - test.Assert(t, s.RecvMsg(context.Background(), nil) == recvErr) + // after SendMsg fails, DoFinish is called and finished=1; + // subsequent RecvMsg returns the cached sendErr (not recvErr) due to finished guard + test.Assert(t, s.RecvMsg(context.Background(), nil) == sendErr) test.Assert(t, s.recvTmCfg.Timeout == 0, s.recvTmCfg) test.Assert(t, !s.recvTmCfg.DisableCancelRemote, s.recvTmCfg) } @@ -381,6 +390,7 @@ func Test_stream_RecvMsg(t *testing.T) { s := newStream(ctx, &mockStream{}, scm, kc, mockRPCInfo, serviceinfo.StreamingBidirectional, nil, func(ctx context.Context, stream streaming.ClientStream, message interface{}) (err error) { return nil }, nil, + //nolint:staticcheck // SA1019: just for testing func(stream streaming.Stream, message interface{}) (err error) { return nil }, @@ -416,6 +426,7 @@ func Test_stream_RecvMsg(t *testing.T) { s := newStream(ctx, st, scm, kc, ri, serviceinfo.StreamingClient, nil, func(ctx context.Context, stream streaming.ClientStream, message interface{}) (err error) { return nil }, nil, + //nolint:staticcheck // SA1019: just for testing func(stream streaming.Stream, message interface{}) (err error) { return nil }, @@ -452,6 +463,7 @@ func Test_stream_RecvMsg(t *testing.T) { s := newStream(ctx, st, scm, kc, ri, serviceinfo.StreamingBidirectional, nil, func(ctx context.Context, stream streaming.ClientStream, message interface{}) (err error) { return recvErr }, nil, + //nolint:staticcheck // SA1019: just for testing func(stream streaming.Stream, message interface{}) (err error) { return recvErr }, @@ -476,6 +488,7 @@ func Test_stream_SendMsg(t *testing.T) { s := newStream(ctx, &mockStream{}, scm, kc, rpcinfo.NewRPCInfo(nil, nil, nil, rpcinfo.NewRPCConfig(), nil), serviceinfo.StreamingBidirectional, func(ctx context.Context, stream streaming.ClientStream, message interface{}) (err error) { return nil }, nil, + //nolint:staticcheck // SA1019: just for testing func(stream streaming.Stream, message interface{}) (err error) { return nil }, @@ -515,6 +528,7 @@ func Test_stream_SendMsg(t *testing.T) { s := newStream(ctx, st, scm, kc, ri, serviceinfo.StreamingBidirectional, func(ctx context.Context, stream streaming.ClientStream, message interface{}) (err error) { return sendErr }, nil, + //nolint:staticcheck // SA1019: just for testing func(stream streaming.Stream, message interface{}) (err error) { return sendErr }, @@ -673,6 +687,93 @@ func Test_stream_DoFinish(t *testing.T) { }) } +// Test_stream_finishedGuard verifies that after DoFinish is called, subsequent +// RecvMsg/SendMsg/Header calls return the cached finishedErr immediately +// without touching the inner stream. This prevents races when a background +// goroutine from callWithTimeout is still running on the inner stream. +func Test_stream_finishedGuard(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + finishedErr := errors.New("stream finished") + + newFinishedStream := func(t *testing.T) *stream { + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("svc", "method"), rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats()) + st := &mockStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + t.Fatal("inner RecvMsg should not be called after DoFinish") + return nil + }, + send: func(ctx context.Context, msg interface{}) error { + t.Fatal("inner SendMsg should not be called after DoFinish") + return nil + }, + header: func() (streaming.Header, error) { + t.Fatal("inner Header should not be called after DoFinish") + return nil, nil + }, + } + cr := mock_remote.NewMockConnReleaser(ctrl) + cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(1) + scm := remotecli.NewStreamConnManager(cr) + kc := &kClient{opt: client.NewOptions(nil)} + s := newStream(st.ctx, st, scm, kc, ri, serviceinfo.StreamingBidirectional, + sendEndpoint, recvEndpoint, nil, nil) + s.DoFinish(finishedErr) + return s + } + + t.Run("RecvMsg() returns finishedErr", func(t *testing.T) { + s := newFinishedStream(t) + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err == finishedErr, err) + }) + t.Run("SendMsg() returns finishedErr", func(t *testing.T) { + s := newFinishedStream(t) + err := s.SendMsg(context.Background(), nil) + test.Assert(t, err == finishedErr, err) + }) + t.Run("Header() returns finishedErr", func(t *testing.T) { + s := newFinishedStream(t) + hd, err := s.Header() + test.Assert(t, err == finishedErr, err) + test.Assert(t, hd == nil) + }) + t.Run("multiple RecvMsg() calls return same error", func(t *testing.T) { + s := newFinishedStream(t) + for i := 0; i < 5; i++ { + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err == finishedErr, err) + } + }) + t.Run("DoFinish() idempotent", func(t *testing.T) { + s := newFinishedStream(t) + // second DoFinish with different error should be no-op + s.DoFinish(errors.New("other error")) + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err == finishedErr, err) + }) + t.Run("DoFinish() with nil err stores io.EOF", func(t *testing.T) { + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("svc", "method"), rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats()) + st := &mockStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + } + cr := mock_remote.NewMockConnReleaser(ctrl) + cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(1) + scm := remotecli.NewStreamConnManager(cr) + kc := &kClient{opt: client.NewOptions(nil)} + s := newStream(st.ctx, st, scm, kc, ri, serviceinfo.StreamingBidirectional, + sendEndpoint, recvEndpoint, nil, nil) + + s.DoFinish(nil) // client streaming success case + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err == io.EOF, err) + err = s.SendMsg(context.Background(), nil) + test.Assert(t, err == io.EOF, err) + }) +} + func Test_isRPCError(t *testing.T) { t.Run("nil", func(t *testing.T) { test.Assert(t, !isRPCError(nil)) @@ -830,13 +931,40 @@ func TestStreamXDoFinish(t *testing.T) { fmt.Printf("Final finishCalled status: %v\n", finishCalled) } -// TestRecvTimeout tests recv timeout scenarios +// mockTTStreamClientStream simulates ttstream's ClientStream at the outer layer. +// It implements CancelableClientStream (like real ttstream clientStream does). +type mockTTStreamClientStream struct { + ctx context.Context + recv func(ctx context.Context, msg interface{}) error + cancelWithErr func(err error) +} + +func (s *mockTTStreamClientStream) Context() context.Context { return s.ctx } + +func (s *mockTTStreamClientStream) Header() (streaming.Header, error) { return nil, nil } + +func (s *mockTTStreamClientStream) Trailer() (streaming.Trailer, error) { return nil, nil } + +func (s *mockTTStreamClientStream) CloseSend(context.Context) error { return nil } + +func (s *mockTTStreamClientStream) SendMsg(context.Context, any) error { return nil } + +func (s *mockTTStreamClientStream) RecvMsg(ctx context.Context, msg any) error { + return s.recv(ctx, msg) +} + +func (s *mockTTStreamClientStream) CancelWithErr(err error) { + if s.cancelWithErr != nil { + s.cancelWithErr(err) + } +} + +// TestRecvTimeout tests recv timeout scenarios for both gRPC and ttstream streams. func TestRecvTimeout(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() t.Run("no timeout set", func(t *testing.T) { - // no timeout, recv should not timeout even if it takes time recvCalled := false ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats()) st := &mockStream{ @@ -850,9 +978,7 @@ func TestRecvTimeout(t *testing.T) { cr := mock_remote.NewMockConnReleaser(ctrl) cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(0) scm := remotecli.NewStreamConnManager(cr) - kc := &kClient{ - opt: client.NewOptions(nil), - } + kc := &kClient{opt: client.NewOptions(nil)} s := newStream(st.ctx, st, scm, kc, ri, serviceinfo.StreamingBidirectional, sendEndpoint, recvEndpoint, nil, nil) @@ -861,212 +987,275 @@ func TestRecvTimeout(t *testing.T) { test.Assert(t, recvCalled) }) - t.Run("timeout set but finish normally", func(t *testing.T) { - // timeout set to 100ms, but recv finishes in 20ms - recvCalled := false - cfg := rpcinfo.NewRPCConfig() - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: 100 * time.Millisecond, - DisableCancelRemote: false, - }) - ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) - st := &mockStream{ - ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), - recv: func(ctx context.Context, msg interface{}) error { - recvCalled = true - time.Sleep(20 * time.Millisecond) - return nil - }, - } - cr := mock_remote.NewMockConnReleaser(ctrl) - cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(0) - scm := remotecli.NewStreamConnManager(cr) - kc := &kClient{ - opt: client.NewOptions(nil), + t.Run("grpc", func(t *testing.T) { + newGRPCStream := func(t *testing.T, ctrl *gomock.Controller, st *mockStream, ri rpcinfo.RPCInfo) *stream { + grpcInner := &mockGRPCInnerStream{ + recv: func(msg interface{}) error { + return st.recv(context.Background(), msg) + }, + ctx: st.ctx, + } + grpcSt := &mockGRPCStreamWrapper{mockStream: st, grpcStream: grpcInner} + cr := mock_remote.NewMockConnReleaser(ctrl) + cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(1) + scm := remotecli.NewStreamConnManager(cr) + kc := &kClient{opt: client.NewOptions(nil)} + return newStream(st.ctx, grpcSt, scm, kc, ri, serviceinfo.StreamingBidirectional, + sendEndpoint, recvEndpoint, + //nolint:staticcheck // SA1019: just for testing + func(stream streaming.Stream, msg interface{}) error { return nil }, + //nolint:staticcheck // SA1019: just for testing + func(stream streaming.Stream, msg interface{}) error { return stream.RecvMsg(msg) }) } - s := newStream(st.ctx, st, scm, kc, ri, serviceinfo.StreamingBidirectional, - sendEndpoint, recvEndpoint, nil, nil) - err := s.RecvMsg(context.Background(), nil) - test.Assert(t, err == nil, err) - test.Assert(t, recvCalled) - }) + t.Run("timeout set but finish normally", func(t *testing.T) { + recvCalled := false + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 100 * time.Millisecond, + }) + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) + st := &mockStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + recvCalled = true + time.Sleep(20 * time.Millisecond) + return nil + }, + } + // override: no release expected on success + cr := mock_remote.NewMockConnReleaser(ctrl) + cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(0) + scm := remotecli.NewStreamConnManager(cr) + grpcInner := &mockGRPCInnerStream{ + recv: func(msg interface{}) error { return st.recv(context.Background(), msg) }, + ctx: st.ctx, + } + grpcSt := &mockGRPCStreamWrapper{mockStream: st, grpcStream: grpcInner} + kc := &kClient{opt: client.NewOptions(nil)} + s := newStream(st.ctx, grpcSt, scm, kc, ri, serviceinfo.StreamingBidirectional, + sendEndpoint, recvEndpoint, + //nolint:staticcheck // SA1019: just for testing + func(stream streaming.Stream, msg interface{}) error { return nil }, + //nolint:staticcheck // SA1019: just for testing + func(stream streaming.Stream, msg interface{}) error { return stream.RecvMsg(msg) }) + + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err == nil, err) + test.Assert(t, recvCalled) + }) - t.Run("timeout and cancel remote", func(t *testing.T) { - // timeout set to 50ms, recv takes 200ms, should timeout and cancel remote - var recvCalled, cancelCalled atomic.Bool - cfg := rpcinfo.NewRPCConfig() - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: 50 * time.Millisecond, - DisableCancelRemote: false, + t.Run("timeout and cancel remote", func(t *testing.T) { + var recvCalled, cancelCalled atomic.Bool + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 50 * time.Millisecond, + }) + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) + st := &mockStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + recvCalled.Store(true) + time.Sleep(200 * time.Millisecond) + return nil + }, + cancelWithErr: func(err error) { + cancelCalled.Store(true) + test.Assert(t, err != nil, "cancel error should not be nil") + }, + } + s := newGRPCStream(t, ctrl, st, ri) + + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err != nil) + stat, ok := status.FromError(err) + test.Assert(t, ok, err) + test.Assert(t, stat.Code() == codes.RecvDeadlineExceeded, stat) + test.Assert(t, strings.Contains(stat.Message(), "stream Recv timeout"), stat.Message()) + test.Assert(t, cancelCalled.Load()) + test.Assert(t, recvCalled.Load()) }) - ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) - st := &mockStream{ - ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), - recv: func(ctx context.Context, msg interface{}) error { - recvCalled.Store(true) - time.Sleep(200 * time.Millisecond) - return nil - }, - cancelWithErr: func(err error) { - cancelCalled.Store(true) - test.Assert(t, err != nil, "cancel error should not be nil") - }, - } - cr := mock_remote.NewMockConnReleaser(ctrl) - cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(1) - scm := remotecli.NewStreamConnManager(cr) - kc := &kClient{ - opt: client.NewOptions(nil), - } - // create a grpc stream wrapper to enable timeout logic - grpcInner := &mockGRPCInnerStream{ - recv: func(msg interface{}) error { - // This will be called by the recv endpoint - return st.recv(context.Background(), msg) - }, - ctx: st.ctx, - } - grpcSt := &mockGRPCStreamWrapper{ - mockStream: st, - grpcStream: grpcInner, - } - s := newStream(st.ctx, grpcSt, scm, kc, ri, serviceinfo.StreamingBidirectional, - sendEndpoint, recvEndpoint, - func(stream streaming.Stream, msg interface{}) error { - return nil - }, - func(stream streaming.Stream, msg interface{}) error { - return stream.RecvMsg(msg) + t.Run("timeout but DisableCancelRemote", func(t *testing.T) { + var recvCalled, cancelCalled atomic.Bool + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 50 * time.Millisecond, + DisableCancelRemote: true, }) + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) + st := &mockStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + recvCalled.Store(true) + time.Sleep(200 * time.Millisecond) + return nil + }, + cancelWithErr: func(err error) { + cancelCalled.Store(true) + }, + } + s := newGRPCStream(t, ctrl, st, ri) + + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err != nil) + stat, ok := status.FromError(err) + test.Assert(t, ok, err) + test.Assert(t, stat.Code() == codes.RecvDeadlineExceeded, stat) + test.Assert(t, strings.Contains(stat.Message(), "stream Recv timeout"), stat.Message()) + test.Assert(t, !cancelCalled.Load()) + test.Assert(t, recvCalled.Load()) + }) - err := s.RecvMsg(context.Background(), nil) - test.Assert(t, err != nil) - stat, ok := status.FromError(err) - test.Assert(t, ok, err) - test.Assert(t, stat.Code() == codes.RecvDeadlineExceeded, stat) - test.Assert(t, strings.Contains(stat.Message(), "stream Recv timeout"), stat.Message()) - test.Assert(t, cancelCalled.Load()) - test.Assert(t, recvCalled.Load()) + t.Run("panic in recv", func(t *testing.T) { + var recvCalled, cancelCalled atomic.Bool + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 100 * time.Millisecond, + }) + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) + st := &mockStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + recvCalled.Store(true) + panic("test panic in recv") + }, + cancelWithErr: func(err error) { + cancelCalled.Store(true) + test.Assert(t, err != nil, "cancel error should not be nil") + }, + } + s := newGRPCStream(t, ctrl, st, ri) + + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err != nil) + stat, ok := status.FromError(err) + test.Assert(t, ok, err) + test.Assert(t, stat.Code() == codes.Internal, stat) + test.Assert(t, strings.Contains(stat.Message(), "stream Recv panic"), stat.Message()) + test.Assert(t, strings.Contains(stat.Message(), "test panic in recv"), stat.Message()) + test.Assert(t, cancelCalled.Load()) + test.Assert(t, recvCalled.Load()) + }) }) - t.Run("timeout but no cancel remote", func(t *testing.T) { - // timeout set to 50ms with DisableCancelRemote=true, recv takes 200ms - // should timeout but NOT cancel remote - var recvCalled, cancelCalled atomic.Bool - cfg := rpcinfo.NewRPCConfig() - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: 50 * time.Millisecond, - DisableCancelRemote: true, - }) - ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) - st := &mockStream{ - ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), - recv: func(ctx context.Context, msg interface{}) error { - recvCalled.Store(true) - time.Sleep(200 * time.Millisecond) - return nil - }, - cancelWithErr: func(err error) { - cancelCalled.Store(true) - }, - } - cr := mock_remote.NewMockConnReleaser(ctrl) - cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(1) - scm := remotecli.NewStreamConnManager(cr) - kc := &kClient{ - opt: client.NewOptions(nil), - } - // create a grpc stream wrapper to enable timeout logic - grpcInner := &mockGRPCInnerStream{ - recv: func(msg interface{}) error { - // This will be called by the recv endpoint - return st.recv(context.Background(), msg) - }, - ctx: st.ctx, - } - grpcSt := &mockGRPCStreamWrapper{ - mockStream: st, - grpcStream: grpcInner, + t.Run("ttstream", func(t *testing.T) { + newTTStream := func(ctrl *gomock.Controller, st *mockTTStreamClientStream, ri rpcinfo.RPCInfo, expectRelease int) *stream { + cr := mock_remote.NewMockConnReleaser(ctrl) + cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(expectRelease) + scm := remotecli.NewStreamConnManager(cr) + kc := &kClient{opt: client.NewOptions(nil)} + return newStream(st.ctx, st, scm, kc, ri, serviceinfo.StreamingBidirectional, + sendEndpoint, recvEndpoint, nil, nil) } - s := newStream(st.ctx, grpcSt, scm, kc, ri, serviceinfo.StreamingBidirectional, - sendEndpoint, recvEndpoint, - func(stream streaming.Stream, msg interface{}) error { - return nil - }, - func(stream streaming.Stream, msg interface{}) error { - return stream.RecvMsg(msg) + t.Run("timeout set but finish normally", func(t *testing.T) { + recvCalled := false + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 100 * time.Millisecond, }) + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) + st := &mockTTStreamClientStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + recvCalled = true + time.Sleep(20 * time.Millisecond) + return nil + }, + } + s := newTTStream(ctrl, st, ri, 0) + + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err == nil, err) + test.Assert(t, recvCalled) + }) - err := s.RecvMsg(context.Background(), nil) - stat, ok := status.FromError(err) - test.Assert(t, ok, err) - test.Assert(t, stat.Code() == codes.RecvDeadlineExceeded, stat) - test.Assert(t, strings.Contains(stat.Message(), "stream Recv timeout"), stat.Message()) - test.Assert(t, err != nil) - test.Assert(t, !cancelCalled.Load()) - test.Assert(t, recvCalled.Load()) - }) - - t.Run("panic in recv and recovered", func(t *testing.T) { - // recv panics, should be recovered and converted to error - var recvCalled, cancelCalled atomic.Bool - cfg := rpcinfo.NewRPCConfig() - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: 100 * time.Millisecond, - DisableCancelRemote: false, + t.Run("timeout fires and cancels remote", func(t *testing.T) { + var recvCalled, cancelCalled atomic.Bool + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 50 * time.Millisecond, + }) + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) + st := &mockTTStreamClientStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + recvCalled.Store(true) + time.Sleep(200 * time.Millisecond) + return nil + }, + cancelWithErr: func(err error) { + cancelCalled.Store(true) + test.Assert(t, err != nil) + }, + } + s := newTTStream(ctrl, st, ri, 1) + + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err != nil) + test.Assert(t, errors.Is(err, kerrors.ErrStreamingTimeout), err) + test.Assert(t, strings.Contains(err.Error(), "stream Recv timeout"), err.Error()) + test.Assert(t, recvCalled.Load()) + test.Assert(t, cancelCalled.Load(), "CancelWithErr should be called when DisableCancelRemote=false") }) - ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) - st := &mockStream{ - ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), - recv: func(ctx context.Context, msg interface{}) error { - recvCalled.Store(true) - panic("test panic in recv") - }, - cancelWithErr: func(err error) { - cancelCalled.Store(true) - test.Assert(t, err != nil, "cancel error should not be nil") - }, - } - cr := mock_remote.NewMockConnReleaser(ctrl) - cr.EXPECT().ReleaseConn(gomock.Any(), gomock.Any()).Times(1) - scm := remotecli.NewStreamConnManager(cr) - kc := &kClient{ - opt: client.NewOptions(nil), - } - // create a grpc stream wrapper to enable timeout logic - grpcInner := &mockGRPCInnerStream{ - recv: func(msg interface{}) error { - // This will be called by the recv endpoint - return st.recv(context.Background(), msg) - }, - ctx: st.ctx, - } - grpcSt := &mockGRPCStreamWrapper{ - mockStream: st, - grpcStream: grpcInner, - } - s := newStream(st.ctx, grpcSt, scm, kc, ri, serviceinfo.StreamingBidirectional, - sendEndpoint, recvEndpoint, - func(stream streaming.Stream, msg interface{}) error { - return nil - }, - func(stream streaming.Stream, msg interface{}) error { - return stream.RecvMsg(msg) + t.Run("timeout but DisableCancelRemote", func(t *testing.T) { + var recvCalled, cancelCalled atomic.Bool + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 50 * time.Millisecond, + DisableCancelRemote: true, }) + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) + st := &mockTTStreamClientStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + recvCalled.Store(true) + time.Sleep(200 * time.Millisecond) + return nil + }, + cancelWithErr: func(err error) { + cancelCalled.Store(true) + }, + } + s := newTTStream(ctrl, st, ri, 1) + + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err != nil) + test.Assert(t, errors.Is(err, kerrors.ErrStreamingTimeout), err) + test.Assert(t, strings.Contains(err.Error(), "stream Recv timeout"), err.Error()) + test.Assert(t, recvCalled.Load()) + test.Assert(t, !cancelCalled.Load(), "CancelWithErr should NOT be called when DisableCancelRemote=true") + }) - err := s.RecvMsg(context.Background(), nil) - test.Assert(t, err != nil) - stat, ok := status.FromError(err) - test.Assert(t, ok, err) - test.Assert(t, stat.Code() == codes.Internal, stat) - test.Assert(t, strings.Contains(stat.Message(), "stream Recv panic"), stat.Message()) - test.Assert(t, strings.Contains(stat.Message(), "test panic in recv"), stat.Message()) - test.Assert(t, cancelCalled.Load()) - test.Assert(t, recvCalled.Load()) + t.Run("panic in recv", func(t *testing.T) { + var recvCalled, cancelCalled atomic.Bool + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 100 * time.Millisecond, + }) + ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("mock_service", "mock_method"), cfg, rpcinfo.NewRPCStats()) + st := &mockTTStreamClientStream{ + ctx: rpcinfo.NewCtxWithRPCInfo(context.Background(), ri), + recv: func(ctx context.Context, msg interface{}) error { + recvCalled.Store(true) + panic("test panic in ttstream recv") + }, + cancelWithErr: func(err error) { + cancelCalled.Store(true) + test.Assert(t, err != nil) + }, + } + s := newTTStream(ctrl, st, ri, 1) + + err := s.RecvMsg(context.Background(), nil) + test.Assert(t, err != nil) + test.Assert(t, strings.Contains(err.Error(), "stream Recv panic"), err.Error()) + test.Assert(t, strings.Contains(err.Error(), "test panic in ttstream recv"), err.Error()) + test.Assert(t, recvCalled.Load()) + test.Assert(t, cancelCalled.Load(), "CancelWithErr should be called on panic") + }) }) } @@ -1076,6 +1265,7 @@ type mockGRPCStreamWrapper struct { grpcStream *mockGRPCInnerStream } +//nolint:staticcheck // SA1019: just for testing func (s *mockGRPCStreamWrapper) GetGRPCStream() streaming.Stream { return s.grpcStream } diff --git a/pkg/remote/trans/ttstream/client_handler.go b/pkg/remote/trans/ttstream/client_handler.go index 2b63c55bb7..87fd66ea4b 100644 --- a/pkg/remote/trans/ttstream/client_handler.go +++ b/pkg/remote/trans/ttstream/client_handler.go @@ -18,6 +18,7 @@ package ttstream import ( "context" + "time" "github.com/bytedance/gopkg/cloud/metainfo" "github.com/cloudwego/gopkg/protocol/ttheader" @@ -48,6 +49,7 @@ type clientTransHandler struct { } // NewStream creates a client stream +// ri.Config().StreamRecvTimeoutConfig() would not take effect, only ri.Config().StreamRecvTimeout() takes effect for compatibility. func (c clientTransHandler) NewStream(ctx context.Context, ri rpcinfo.RPCInfo) (streaming.ClientStream, error) { rconfig := ri.Config() invocation := ri.Invocation() @@ -74,6 +76,10 @@ func (c clientTransHandler) NewStream(ctx context.Context, ri rpcinfo.RPCInfo) ( } strHeader[ttheader.HeaderIDLServiceName] = invocation.ServiceName() metainfo.SaveMetaInfoToMap(ctx, strHeader) + var tm time.Duration + if ddl, ok := ctx.Deadline(); ok { + tm = time.Until(ddl) + } trans, err := c.transPool.Get(addr.Network(), addr.String()) if err != nil { @@ -84,9 +90,16 @@ func (c clientTransHandler) NewStream(ctx context.Context, ri rpcinfo.RPCInfo) ( cs := newClientStream(ctx, trans, streamFrame{sid: genStreamID(), method: method}) // stream should be configured before WriteStream or there would be a race condition for metaFrameHandler - cs.setRecvTimeoutConfig(rconfig) + // only compatible for client.WithStreamRecvTimeout | streamcall.WithRecvTimeout. + // client.WithStreamRecvTimeoutConfig | streamcall.WithRecvTimeoutConfig would only take effect in the outside client/stream layer. + // when RecvTimeout and RecvTimeoutConfig are both configured, RecvTimeout would not take effect. + if rconfig.StreamRecvTimeoutConfig().Timeout <= 0 { + cs.setRecvTimeout(rconfig.StreamRecvTimeout()) + } + cs.setMetaFrameHandler(c.metaHandler) cs.setTraceController(c.traceCtl) + cs.setStreamTimeout(tm) if err = trans.WriteStream(ctx, cs, intHeader, strHeader); err != nil { return nil, err diff --git a/pkg/remote/trans/ttstream/client_handler_test.go b/pkg/remote/trans/ttstream/client_handler_test.go new file mode 100644 index 0000000000..884d49e19e --- /dev/null +++ b/pkg/remote/trans/ttstream/client_handler_test.go @@ -0,0 +1,224 @@ +//go:build !windows + +/* + * Copyright 2026 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 ttstream + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/cloudwego/kitex/internal/test" + "github.com/cloudwego/kitex/pkg/kerrors" + "github.com/cloudwego/kitex/pkg/rpcinfo" + "github.com/cloudwego/kitex/pkg/streaming" + "github.com/cloudwego/kitex/pkg/utils" +) + +// mockTransPool implements transPool for testing clientTransHandler.NewStream +type mockTransPool struct { + getTrans *clientTransport + getErr error +} + +func (p *mockTransPool) Get(network, addr string) (*clientTransport, error) { + return p.getTrans, p.getErr +} + +func (p *mockTransPool) Put(trans *clientTransport) {} + +// mockHeaderFrameWriteHandler implements HeaderFrameWriteHandler for testing +type mockHeaderFrameWriteHandler struct { + intHeader IntHeader + strHeader streaming.Header + err error +} + +func (h *mockHeaderFrameWriteHandler) OnWriteStream(ctx context.Context) (IntHeader, StrHeader, error) { + return h.intHeader, h.strHeader, h.err +} + +func newTestRPCInfoCtx(cfg rpcinfo.RPCConfig) (context.Context, rpcinfo.RPCInfo) { + if cfg == nil { + cfg = rpcinfo.NewRPCConfig() + } + ri := rpcinfo.NewRPCInfo( + rpcinfo.NewEndpointInfo("clientSvc", "testMethod", nil, nil), + rpcinfo.NewEndpointInfo("serverSvc", "testMethod", utils.NewNetAddr("tcp", "127.0.0.1:8888"), nil), + rpcinfo.NewInvocation("serverSvc", "testMethod"), + cfg, + nil, + ) + ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri) + return ctx, ri +} + +func TestNewStream(t *testing.T) { + t.Run("no dest address", func(t *testing.T) { + handler := &clientTransHandler{ + transPool: &mockTransPool{}, + } + ri := rpcinfo.NewRPCInfo( + nil, + rpcinfo.NewEndpointInfo("serverSvc", "testMethod", nil, nil), // addr = nil + rpcinfo.NewInvocation("serverSvc", "testMethod"), + rpcinfo.NewRPCConfig(), nil, + ) + ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri) + _, err := handler.NewStream(ctx, ri) + test.Assert(t, errors.Is(err, kerrors.ErrNoDestAddress), err) + }) + + t.Run("transPool.Get fails", func(t *testing.T) { + getErr := errors.New("pool get error") + handler := &clientTransHandler{ + transPool: &mockTransPool{getErr: getErr}, + } + ctx, ri := newTestRPCInfoCtx(nil) + _, err := handler.NewStream(ctx, ri) + test.Assert(t, err == getErr, err) + }) + + t.Run("headerHandler.OnWriteStream fails", func(t *testing.T) { + hdlErr := errors.New("header handler error") + handler := &clientTransHandler{ + transPool: &mockTransPool{}, + headerHandler: &mockHeaderFrameWriteHandler{err: hdlErr}, + } + ctx, ri := newTestRPCInfoCtx(nil) + _, err := handler.NewStream(ctx, ri) + test.Assert(t, err == hdlErr, err) + }) + + t.Run("success with StreamRecvTimeout", func(t *testing.T) { + cconn, sconn := newTestConnectionPipe(t) + defer cconn.Close() + defer sconn.Close() + + ctrans := newClientTransport(cconn, nil) + handler := &clientTransHandler{ + transPool: &mockTransPool{getTrans: ctrans}, + } + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeout(100 * time.Millisecond) + ctx, ri := newTestRPCInfoCtx(cfg) + + cs, err := handler.NewStream(ctx, ri) + test.Assert(t, err == nil, err) + test.Assert(t, cs != nil) + + cliSt := cs.(*clientStream) + // StreamRecvTimeout should be set at ttstream layer + test.Assert(t, cliSt.recvTimeout == 100*time.Millisecond, cliSt.recvTimeout) + test.Assert(t, cliSt.stream.recvTimeoutCallback != nil) + }) + + t.Run("success with StreamRecvTimeoutConfig, ttstream layer should NOT set recvTimeout", func(t *testing.T) { + cconn, sconn := newTestConnectionPipe(t) + defer cconn.Close() + defer sconn.Close() + + ctrans := newClientTransport(cconn, nil) + handler := &clientTransHandler{ + transPool: &mockTransPool{getTrans: ctrans}, + } + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 200 * time.Millisecond, + DisableCancelRemote: true, + }) + ctx, ri := newTestRPCInfoCtx(cfg) + + cs, err := handler.NewStream(ctx, ri) + test.Assert(t, err == nil, err) + test.Assert(t, cs != nil) + + cliSt := cs.(*clientStream) + // StreamRecvTimeoutConfig is handled by the outer layer, ttstream should NOT set recvTimeout + test.Assert(t, cliSt.recvTimeout == 0, cliSt.recvTimeout) + test.Assert(t, cliSt.stream.recvTimeoutCallback == nil) + }) + + t.Run("both set, StreamRecvTimeoutConfig takes priority", func(t *testing.T) { + cconn, sconn := newTestConnectionPipe(t) + defer cconn.Close() + defer sconn.Close() + + ctrans := newClientTransport(cconn, nil) + handler := &clientTransHandler{ + transPool: &mockTransPool{getTrans: ctrans}, + } + cfg := rpcinfo.NewRPCConfig() + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeout(500 * time.Millisecond) + rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ + Timeout: 50 * time.Millisecond, + }) + ctx, ri := newTestRPCInfoCtx(cfg) + + cs, err := handler.NewStream(ctx, ri) + test.Assert(t, err == nil, err) + test.Assert(t, cs != nil) + + cliSt := cs.(*clientStream) + // StreamRecvTimeoutConfig takes priority, ttstream should NOT set recvTimeout + test.Assert(t, cliSt.recvTimeout == 0, cliSt.recvTimeout) + test.Assert(t, cliSt.stream.recvTimeoutCallback == nil) + }) + + t.Run("stream ctx deadline sets streamTimeout", func(t *testing.T) { + cconn, sconn := newTestConnectionPipe(t) + defer cconn.Close() + defer sconn.Close() + + ctrans := newClientTransport(cconn, nil) + handler := &clientTransHandler{ + transPool: &mockTransPool{getTrans: ctrans}, + } + ctx, ri := newTestRPCInfoCtx(nil) + ctx, cancel := context.WithTimeout(ctx, 500*time.Millisecond) + defer cancel() + + cs, err := handler.NewStream(ctx, ri) + test.Assert(t, err == nil, err) + + cliSt := cs.(*clientStream) + // streamTimeout should be set from ctx deadline + test.Assert(t, cliSt.streamTimeout > 0 && cliSt.streamTimeout <= 500*time.Millisecond, cliSt.streamTimeout) + }) + + t.Run("headerHandler provides custom headers", func(t *testing.T) { + cconn, sconn := newTestConnectionPipe(t) + defer cconn.Close() + defer sconn.Close() + + ctrans := newClientTransport(cconn, nil) + handler := &clientTransHandler{ + transPool: &mockTransPool{getTrans: ctrans}, + headerHandler: &mockHeaderFrameWriteHandler{ + intHeader: IntHeader{1: "val1"}, + strHeader: streaming.Header{"custom-key": "custom-val"}, + }, + } + ctx, ri := newTestRPCInfoCtx(nil) + + cs, err := handler.NewStream(ctx, ri) + test.Assert(t, err == nil, err) + test.Assert(t, cs != nil) + }) +} diff --git a/pkg/remote/trans/ttstream/client_trans_pool_longconn.go b/pkg/remote/trans/ttstream/client_trans_pool_longconn.go index 70db41a25e..1680f3af90 100644 --- a/pkg/remote/trans/ttstream/client_trans_pool_longconn.go +++ b/pkg/remote/trans/ttstream/client_trans_pool_longconn.go @@ -19,7 +19,7 @@ package ttstream import ( "time" - "github.com/cloudwego/kitex/pkg/remote/trans/ttstream/container" + "github.com/cloudwego/kitex/pkg/remote/trans/ttstream/internal/container" ) var DefaultLongConnConfig = LongConnConfig{ diff --git a/pkg/remote/trans/ttstream/container/pipe.go b/pkg/remote/trans/ttstream/container/pipe.go deleted file mode 100644 index daec472ae4..0000000000 --- a/pkg/remote/trans/ttstream/container/pipe.go +++ /dev/null @@ -1,152 +0,0 @@ -/* - * 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 container - -import ( - "context" - "fmt" - "io" - "sync/atomic" -) - -type pipeState = int32 - -const ( - pipeStateActive pipeState = 0 - pipeStateClosed pipeState = 1 - pipeStateCanceled pipeState = 2 -) - -var ( - ErrPipeEOF = io.EOF - ErrPipeCanceled = fmt.Errorf("pipe canceled") - stateErrors = map[pipeState]error{ - pipeStateClosed: ErrPipeEOF, - pipeStateCanceled: ErrPipeCanceled, - } -) - -type CtxDoneCallback func(ctx context.Context) - -// Pipe implement a queue that never block on Write but block on Read if there is nothing to read -type Pipe[Item any] struct { - queue *Queue[Item] - state pipeState - trigger chan struct{} - callback CtxDoneCallback -} - -func NewPipe[Item any](opts ...PipeOption) *Pipe[Item] { - options := PipeOptions{} - for _, opt := range opts { - opt(&options) - } - p := new(Pipe[Item]) - p.queue = NewQueue[Item]() - p.trigger = make(chan struct{}, 1) - p.state = pipeStateActive - p.callback = options.callback - return p -} - -// Read will block if there is nothing to read -func (p *Pipe[Item]) Read(ctx context.Context, items []Item) (n int, err error) { -READ: - // check readable items - for i := 0; i < len(items); i++ { - val, ok := p.queue.Get() - if !ok { - break - } - items[i] = val - n++ - } - if n > 0 { - return n, nil - } - - // check state - state := atomic.LoadInt32(&p.state) - if err = stateErrors[state]; err != nil { - return 0, err - } - - // no data to read, waiting writes - for { - if ctx.Done() != nil { - select { - case <-ctx.Done(): - if p.callback != nil { - p.callback(ctx) - // waiting for callback to close this pipe - // Read is not concurrent-safe; the trigger signal can only be consumed once. - <-p.trigger - goto READ - } - return 0, ctx.Err() - case <-p.trigger: - } - } else { - <-p.trigger - } - goto READ - } -} - -func (p *Pipe[Item]) Write(ctx context.Context, items ...Item) (err error) { - for _, item := range items { - p.queue.Add(item) - } - // wake up - select { - case p.trigger <- struct{}{}: - default: - } - return nil -} - -func (p *Pipe[Item]) Close() { - if !atomic.CompareAndSwapInt32(&p.state, pipeStateActive, pipeStateClosed) { - return - } - select { - case p.trigger <- struct{}{}: - default: - } -} - -func (p *Pipe[Item]) Cancel() { - if !atomic.CompareAndSwapInt32(&p.state, pipeStateActive, pipeStateCanceled) { - return - } - select { - case p.trigger <- struct{}{}: - default: - } -} - -type PipeOptions struct { - callback CtxDoneCallback -} - -type PipeOption func(*PipeOptions) - -func WithCtxDoneCallback(callback CtxDoneCallback) PipeOption { - return func(options *PipeOptions) { - options.callback = callback - } -} diff --git a/pkg/remote/trans/ttstream/container/pipe_test.go b/pkg/remote/trans/ttstream/container/pipe_test.go deleted file mode 100644 index ff4cc0fd27..0000000000 --- a/pkg/remote/trans/ttstream/container/pipe_test.go +++ /dev/null @@ -1,163 +0,0 @@ -/* - * 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 container - -import ( - "context" - "errors" - "io" - "runtime" - "sync" - "sync/atomic" - "testing" - "time" - - "github.com/cloudwego/kitex/internal/test" -) - -func TestPipeline(t *testing.T) { - ctx := context.Background() - pipe := NewPipe[int]() - var recv int - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - items := make([]int, 10) - for { - n, err := pipe.Read(ctx, items) - if err != nil { - if err != io.EOF { - t.Error(err) - } - return - } - for i := 0; i < n; i++ { - recv += items[i] - } - } - }() - round := 10000 - itemsPerRound := []int{1, 1, 1, 1, 1} - for i := 0; i < round; i++ { - err := pipe.Write(ctx, itemsPerRound...) - if err != nil { - t.Fatal(err) - } - } - t.Logf("Pipe closing") - pipe.Close() - t.Logf("Pipe closed") - wg.Wait() - if recv != len(itemsPerRound)*round { - t.Fatalf("Pipe expect %d items, got %d", len(itemsPerRound)*round, recv) - } -} - -func TestPipelineWriteCloseAndRead(t *testing.T) { - ctx := context.Background() - pipe := NewPipe[int]() - pipeSize := 10 - var pipeRead int32 - var readWG sync.WaitGroup - readWG.Add(1) - go func() { - defer readWG.Done() - readBuf := make([]int, 1) - for { - n, err := pipe.Read(ctx, readBuf) - if err != nil { - if !errors.Is(err, ErrPipeEOF) { - t.Errorf("un except err: %v", err) - } - break - } - atomic.AddInt32(&pipeRead, int32(n)) - } - }() - time.Sleep(time.Millisecond * 10) // let read goroutine start first - for i := 0; i < pipeSize; i++ { - err := pipe.Write(ctx, i) - if err != nil { - t.Error(err) - } - } - for atomic.LoadInt32(&pipeRead) != int32(pipeSize) { - runtime.Gosched() - } - pipe.Close() - readWG.Wait() -} - -func TestPipelineCancelCallback(t *testing.T) { - pipeWrapper := struct{ pipe *Pipe[int] }{} - finalNum := 100 - firstNum := 1 - pipe := NewPipe[int](WithCtxDoneCallback( - func(ctx context.Context) { - _ = pipeWrapper.pipe.Write(context.Background(), finalNum) - pipeWrapper.pipe.Close() - }, - )) - pipeWrapper.pipe = pipe - var readWG sync.WaitGroup - readWG.Add(1) - go func() { - defer readWG.Done() - var canceled bool - readBuf := make([]int, 1) - rCtx, cancel := context.WithCancel(context.Background()) - defer cancel() - for { - n, rErr := pipe.Read(rCtx, readBuf) - if canceled { - if rErr != nil { - test.Assert(t, rErr == io.EOF, rErr) - return - } - test.Assert(t, n == 1, n) - test.Assert(t, readBuf[0] == finalNum, readBuf[0]) - continue - } - test.Assert(t, n == 1, n) - test.Assert(t, readBuf[0] == firstNum, readBuf[0]) - cancel() - canceled = true - } - }() - wErr := pipe.Write(context.Background(), firstNum) - test.Assert(t, wErr == nil, wErr) - readWG.Wait() -} - -func BenchmarkPipeline(b *testing.B) { - ctx := context.Background() - pipe := NewPipe[int]() - readCache := make([]int, 8) - b.ResetTimer() - b.ReportAllocs() - for i := 0; i < b.N; i++ { - for j := 0; j < len(readCache); j++ { - go pipe.Write(ctx, 1) - } - total := 0 - for total < len(readCache) { - n, _ := pipe.Read(ctx, readCache) - total += n - } - } -} diff --git a/pkg/remote/trans/ttstream/exception.go b/pkg/remote/trans/ttstream/exception.go index 3ef94e0304..0de226aafc 100644 --- a/pkg/remote/trans/ttstream/exception.go +++ b/pkg/remote/trans/ttstream/exception.go @@ -20,33 +20,84 @@ import ( "errors" "fmt" "strings" + "time" "github.com/cloudwego/kitex/pkg/kerrors" "github.com/cloudwego/kitex/pkg/streaming" ) +const ( + errApplicationExceptionTypeId = 12001 + // protocol related + errUnexpectedHeaderTypeId = 12002 + errIllegalBizErrTypeId = 12003 + errIllegalFrameTypeId = 12004 + errIllegalOperationTypeId = 12005 + errTransportTypeId = 12006 + // cancel related + errBizCancelTypeId = 12007 + errBizCancelWithCauseTypeId = 12008 + errDownstreamCancelTypeId = 12009 + errUpstreamCancelTypeId = 12010 + errInternalCancelTypeId = 12011 + errBizHandlerReturnTypeId = 12012 + errConnectionClosedTypeId = 12013 + // timeout related + errRecvTimeoutTypeId = 12014 + errStreamTimeoutTypeId = 12015 + // other + errStreamInternalTypeId = 12016 +) + var ( - errApplicationException = newException("application exception", nil, 12001) - errUnexpectedHeader = newException("unexpected header frame", kerrors.ErrStreamingProtocol, 12002) - errIllegalBizErr = newException("illegal bizErr", kerrors.ErrStreamingProtocol, 12003) - errIllegalFrame = newException("illegal frame", kerrors.ErrStreamingProtocol, 12004) - errIllegalOperation = newException("illegal operation", kerrors.ErrStreamingProtocol, 12005) - errTransport = newException("transport is closing", kerrors.ErrStreamingProtocol, 12006) + errApplicationException = newException("application exception", nil, errApplicationExceptionTypeId) + errUnexpectedHeader = newException("unexpected header frame", kerrors.ErrStreamingProtocol, errUnexpectedHeaderTypeId) + errIllegalBizErr = newException("illegal bizErr", kerrors.ErrStreamingProtocol, errIllegalBizErrTypeId) + errIllegalFrame = newException("illegal frame", kerrors.ErrStreamingProtocol, errIllegalFrameTypeId) + errIllegalOperation = newException("illegal operation", kerrors.ErrStreamingProtocol, errIllegalOperationTypeId) + errTransport = newException("transport is closing", kerrors.ErrStreamingProtocol, errTransportTypeId) errBizCancel = newException("user code invoking stream RPC with context processed by context.WithCancel or context.WithTimeout, then invoking cancel() actively", - kerrors.ErrStreamingCanceled, 12007) - errBizCancelWithCause = newException("user code canceled with cancelCause(error)", kerrors.ErrStreamingCanceled, 12008) - errDownstreamCancel = newException("canceled by downstream", kerrors.ErrStreamingCanceled, 12009) - errUpstreamCancel = newException("canceled by upstream", kerrors.ErrStreamingCanceled, 12010) - errInternalCancel = newException("internal canceled", kerrors.ErrStreamingCanceled, 12011) - errBizHandlerReturnCancel = newException("canceled by business handler returning", kerrors.ErrStreamingCanceled, 12012) - errConnectionClosedCancel = newException("canceled by connection closed", kerrors.ErrStreamingCanceled, 12013) + kerrors.ErrStreamingCanceled, errBizCancelTypeId) + errBizCancelWithCause = newException("user code canceled with cancelCause(error)", kerrors.ErrStreamingCanceled, errBizCancelWithCauseTypeId) + errDownstreamCancel = newException("canceled by downstream", kerrors.ErrStreamingCanceled, errDownstreamCancelTypeId) + errUpstreamCancel = newException("canceled by upstream", kerrors.ErrStreamingCanceled, errUpstreamCancelTypeId) + errInternalCancel = newException("internal canceled", kerrors.ErrStreamingCanceled, errInternalCancelTypeId) + errBizHandlerReturnCancel = newException("canceled by business handler returning", kerrors.ErrStreamingCanceled, errBizHandlerReturnTypeId) + errConnectionClosedCancel = newException("canceled by connection closed", kerrors.ErrStreamingCanceled, errConnectionClosedTypeId) ) var errServerSideBizHandlerReturnCancel = errBizHandlerReturnCancel.newBuilder().withSide(serverSide) -func newStreamRecvTimeoutException(cfg streaming.TimeoutConfig) *Exception { - return newException(fmt.Sprintf("stream Recv timeout, timeout config=%+v", cfg), kerrors.ErrStreamingTimeout, 12014).withSide(clientSide) +func NewStreamRecvTimeoutException(cfg streaming.TimeoutConfig, isClient bool) *Exception { + side := clientSide + if !isClient { + side = serverSide + } + return newException(fmt.Sprintf("stream Recv timeout, timeout config=%+v", cfg), kerrors.ErrStreamingTimeout, errRecvTimeoutTypeId).withSide(side) +} + +func newStreamRecvTimeoutException(tm time.Duration) *Exception { + return newException(fmt.Sprintf("stream Recv timeout, timeout=%+v", tm), kerrors.ErrStreamingTimeout, errRecvTimeoutTypeId).withSide(clientSide) +} + +func newStreamTimeoutException(tm time.Duration) *Exception { + var msg string + if tm > 0 { + msg = fmt.Sprintf("stream timeout, timeout in ctx: %+v", tm) + } else { + msg = "stream timeout, no explicit timeout set in stream ctx" + } + return newException(msg, kerrors.ErrStreamingTimeout, errStreamTimeoutTypeId).withSide(clientSide) +} + +func NewStreamInternalException(desc string, isClient bool) *Exception { + side := clientSide + if !isClient { + side = serverSide + } + // stream internal exception does not have specified parent err, just pass nil + return newException(desc, nil, errStreamInternalTypeId).withSide(side) } const ( diff --git a/pkg/remote/trans/ttstream/exception_test.go b/pkg/remote/trans/ttstream/exception_test.go index 7ca5efc3c2..146f193f0b 100644 --- a/pkg/remote/trans/ttstream/exception_test.go +++ b/pkg/remote/trans/ttstream/exception_test.go @@ -27,7 +27,6 @@ import ( "github.com/cloudwego/kitex/internal/test" "github.com/cloudwego/kitex/pkg/kerrors" - "github.com/cloudwego/kitex/pkg/streaming" ) func TestErrors(t *testing.T) { @@ -46,7 +45,7 @@ func TestErrors(t *testing.T) { test.Assert(t, !newExWithNilErr.isCauseSet(), newExWithNilErr) test.Assert(t, newExWithNilErr.cause == nil, newExWithNilErr) - recvTmErr := newStreamRecvTimeoutException(streaming.TimeoutConfig{Timeout: 1 * time.Second}) + recvTmErr := newStreamRecvTimeoutException(1 * time.Second) test.Assert(t, errors.Is(recvTmErr, kerrors.ErrStreamingTimeout), recvTmErr) } @@ -77,7 +76,7 @@ func TestCommonParentKerror(t *testing.T) { // timeout Exception errs = []error{ - newStreamRecvTimeoutException(streaming.TimeoutConfig{Timeout: 1 * time.Second}), + newStreamRecvTimeoutException(1 * time.Second), } for _, err := range errs { test.Assert(t, errors.Is(err, kerrors.ErrStreamingTimeout), err) @@ -91,14 +90,14 @@ func TestGetTypeId(t *testing.T) { err error expectTypeId int32 }{ - {err: errApplicationException, expectTypeId: 12001}, - {err: errUnexpectedHeader, expectTypeId: 12002}, - {err: errIllegalBizErr, expectTypeId: 12003}, - {err: errIllegalFrame, expectTypeId: 12004}, - {err: errIllegalOperation, expectTypeId: 12005}, - {err: errTransport, expectTypeId: 12006}, + {err: errApplicationException, expectTypeId: errApplicationExceptionTypeId}, + {err: errUnexpectedHeader, expectTypeId: errUnexpectedHeaderTypeId}, + {err: errIllegalBizErr, expectTypeId: errIllegalBizErrTypeId}, + {err: errIllegalFrame, expectTypeId: errIllegalFrameTypeId}, + {err: errIllegalOperation, expectTypeId: errIllegalOperationTypeId}, + {err: errTransport, expectTypeId: errTransportTypeId}, {err: errApplicationException.newBuilder().withCauseAndTypeId(exception, 1000), expectTypeId: 1000}, - {err: errApplicationException.newBuilder().withCause(normalErr), expectTypeId: 12001}, + {err: errApplicationException.newBuilder().withCause(normalErr), expectTypeId: errApplicationExceptionTypeId}, } for _, testcase := range testcases { diff --git a/pkg/remote/trans/ttstream/internal/container/doc.go b/pkg/remote/trans/ttstream/internal/container/doc.go new file mode 100644 index 0000000000..1efcb40517 --- /dev/null +++ b/pkg/remote/trans/ttstream/internal/container/doc.go @@ -0,0 +1,23 @@ +/* + * Copyright 2026 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 container + +// container package contains high performance containers implementation used by ttstream. +// +// This package was previously under pkg/remote/trans/ttstream. It has been moved to pkg/remote/trans/ttstream/internal +// so that future iteration and refactoring of ttstream can proceed without unnecessary constraints. +// We sincerely apologize for any breaking changes this may introduce. diff --git a/pkg/remote/trans/ttstream/container/linklist.go b/pkg/remote/trans/ttstream/internal/container/linklist.go similarity index 100% rename from pkg/remote/trans/ttstream/container/linklist.go rename to pkg/remote/trans/ttstream/internal/container/linklist.go diff --git a/pkg/remote/trans/ttstream/container/object_pool.go b/pkg/remote/trans/ttstream/internal/container/object_pool.go similarity index 100% rename from pkg/remote/trans/ttstream/container/object_pool.go rename to pkg/remote/trans/ttstream/internal/container/object_pool.go diff --git a/pkg/remote/trans/ttstream/container/object_pool_test.go b/pkg/remote/trans/ttstream/internal/container/object_pool_test.go similarity index 100% rename from pkg/remote/trans/ttstream/container/object_pool_test.go rename to pkg/remote/trans/ttstream/internal/container/object_pool_test.go diff --git a/pkg/remote/trans/ttstream/internal/container/pipe.go b/pkg/remote/trans/ttstream/internal/container/pipe.go new file mode 100644 index 0000000000..a5aeb1730c --- /dev/null +++ b/pkg/remote/trans/ttstream/internal/container/pipe.go @@ -0,0 +1,153 @@ +/* + * 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 container + +import ( + "context" + "io" + "sync/atomic" +) + +const ( + pipeStateActive int32 = 0 + pipeStateClosed int32 = 1 +) + +type CtxDoneCallback func(ctx context.Context) error + +// Pipe implement a queue that never block on Write but block on Read if there is nothing to read +type Pipe[Item any] struct { + ctx context.Context + callback CtxDoneCallback + + queue *Queue[Item] + state int32 + trigger chan struct{} +} + +func NewPipe[Item any](ctx context.Context, callback CtxDoneCallback) *Pipe[Item] { + p := new(Pipe[Item]) + p.ctx = ctx + p.callback = callback + p.queue = NewQueue[Item]() + p.trigger = make(chan struct{}, 1) + p.state = pipeStateActive + return p +} + +// Read will block if there is nothing to read, only checks the pipe's own ctx. +func (p *Pipe[Item]) Read(items []Item) (n int, err error) { + return p.ReadCtx(context.Background(), nil, items) +} + +// ReadCtx will additionally check a per-read ctx alongside the pipe's own ctx. +func (p *Pipe[Item]) ReadCtx(perReadCtx context.Context, perReadCallback CtxDoneCallback, items []Item) (n int, err error) { + for { + n = p.read(items) + if n > 0 { + return n, nil + } + if atomic.LoadInt32(&p.state) == pipeStateClosed { + // In scenarios where the stream is closed immediately after a write operation, + // we should continue attempting to drain the remaining items to prevent data loss. + n = p.read(items) + if n > 0 { + return n, nil + } + return 0, io.EOF + } + + // no data to read, waiting writes or ctx done + readMore, wErr := p.waitCtx(perReadCtx, perReadCallback) + + if !readMore { + return 0, wErr + } + } +} + +func (p *Pipe[Item]) read(items []Item) int { + var n int + for i := 0; i < len(items); i++ { + val, ok := p.queue.Get() + if !ok { + break + } + items[i] = val + n++ + } + return n +} + +func (p *Pipe[Item]) waitCtx(ctx context.Context, callback CtxDoneCallback) (readMore bool, err error) { + select { + case <-p.ctx.Done(): + if p.callback != nil { + return false, p.callback(p.ctx) + } + return false, p.ctx.Err() + case <-ctx.Done(): + // double-check: prefer the pipe ctx + select { + case <-p.ctx.Done(): + if p.callback != nil { + return false, p.callback(p.ctx) + } + return false, p.ctx.Err() + default: + } + + if callback != nil { + return false, callback(ctx) + } + return false, ctx.Err() + case <-p.trigger: + return true, nil + } +} + +// Write enqueues items and wakes up a blocked reader. +// +// Write is NOT safe to call after Close. Once Close is called, the pipe is considered done +// and any subsequent Write may result in data loss since the reader may have already observed +// EOF. Callers must ensure all Write calls complete before calling Close. +func (p *Pipe[Item]) Write(items ...Item) { + for _, item := range items { + p.queue.Add(item) + } + // wake up + select { + case p.trigger <- struct{}{}: + default: + } +} + +// Close marks the pipe as closed and wakes up a blocked reader so it can drain remaining items +// and return io.EOF. +// +// Close is idempotent: only the first call takes effect. After Close, the reader will drain any +// items already enqueued by prior Write calls before returning io.EOF. However, items written +// after Close are not guaranteed to be delivered. +func (p *Pipe[Item]) Close() { + if !atomic.CompareAndSwapInt32(&p.state, pipeStateActive, pipeStateClosed) { + return + } + select { + case p.trigger <- struct{}{}: + default: + } +} diff --git a/pkg/remote/trans/ttstream/internal/container/pipe_test.go b/pkg/remote/trans/ttstream/internal/container/pipe_test.go new file mode 100644 index 0000000000..eeaaf7b506 --- /dev/null +++ b/pkg/remote/trans/ttstream/internal/container/pipe_test.go @@ -0,0 +1,343 @@ +/* + * 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 container + +import ( + "context" + "errors" + "io" + "runtime" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/cloudwego/kitex/internal/test" +) + +func TestPipeRead(t *testing.T) { + t.Run("concurrent write and read", func(t *testing.T) { + pipe := NewPipe[int](context.Background(), nil) + var recv int + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + items := make([]int, 10) + for { + n, err := pipe.Read(items) + if err != nil { + test.Assert(t, err == io.EOF, err) + return + } + for i := 0; i < n; i++ { + recv += items[i] + } + } + }() + round := 10000 + itemsPerRound := []int{1, 1, 1, 1, 1} + for i := 0; i < round; i++ { + pipe.Write(itemsPerRound...) + } + pipe.Close() + wg.Wait() + test.Assert(t, recv == len(itemsPerRound)*round, recv) + }) + t.Run("write then close then read drains all", func(t *testing.T) { + pipe := NewPipe[int](context.Background(), nil) + pipeSize := 10 + var pipeRead int32 + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + readBuf := make([]int, 1) + for { + n, err := pipe.Read(readBuf) + if err != nil { + test.Assert(t, errors.Is(err, io.EOF), err) + break + } + atomic.AddInt32(&pipeRead, int32(n)) + } + }() + time.Sleep(10 * time.Millisecond) + for i := 0; i < pipeSize; i++ { + pipe.Write(i) + } + for atomic.LoadInt32(&pipeRead) != int32(pipeSize) { + runtime.Gosched() + } + pipe.Close() + wg.Wait() + }) + t.Run("pipe ctx cancel without callback", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + pipe := NewPipe[int](ctx, nil) + + var readErr error + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + _, readErr = pipe.Read(make([]int, 10)) + }() + + time.Sleep(50 * time.Millisecond) + cancel() + wg.Wait() + test.Assert(t, errors.Is(readErr, context.Canceled), readErr) + }) + t.Run("pipe ctx cancel with callback", func(t *testing.T) { + pipeWrapper := struct{ pipe *Pipe[int] }{} + finalNum := 100 + firstNum := 1 + ctx, cancel := context.WithCancel(context.Background()) + pipe := NewPipe[int](ctx, func(ctx context.Context) error { + pipeWrapper.pipe.Write(finalNum) + pipeWrapper.pipe.Close() + return ctx.Err() + }) + pipeWrapper.pipe = pipe + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + var canceled bool + readBuf := make([]int, 1) + for { + n, rErr := pipe.Read(readBuf) + if canceled { + if rErr != nil { + test.Assert(t, rErr == context.Canceled, rErr) + return + } + test.Assert(t, n == 1, n) + test.Assert(t, readBuf[0] == finalNum, readBuf[0]) + continue + } + test.Assert(t, n == 1, n) + test.Assert(t, readBuf[0] == firstNum, readBuf[0]) + cancel() + canceled = true + } + }() + pipe.Write(firstNum) + wg.Wait() + }) +} + +func TestPipeReadCtx(t *testing.T) { + t.Run("read available data", func(t *testing.T) { + pipe := NewPipe[int](context.Background(), nil) + pipe.Write(1, 2, 3, 4, 5) + + items := make([]int, 10) + n, err := pipe.ReadCtx(context.Background(), nil, items) + test.Assert(t, err == nil, err) + test.Assert(t, n == 5, n) + test.Assert(t, items[0] == 1, items[0]) + test.Assert(t, items[4] == 5, items[4]) + pipe.Close() + }) + t.Run("perReadCtx cancel without callback", func(t *testing.T) { + pipe := NewPipe[int](context.Background(), nil) + readCtx, cancel := context.WithCancel(context.Background()) + + var readErr error + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + _, readErr = pipe.ReadCtx(readCtx, nil, make([]int, 10)) + }() + + time.Sleep(50 * time.Millisecond) + cancel() + wg.Wait() + test.Assert(t, errors.Is(readErr, context.Canceled), readErr) + pipe.Close() + }) + t.Run("perReadCtx cancel with callback", func(t *testing.T) { + pipe := NewPipe[int](context.Background(), nil) + readCtx, cancel := context.WithCancel(context.Background()) + var callbackCalled bool + callback := func(ctx context.Context) error { + callbackCalled = true + test.Assert(t, ctx == readCtx, ctx) + return errors.New("custom callback error") + } + + var readErr error + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + _, readErr = pipe.ReadCtx(readCtx, callback, make([]int, 10)) + }() + + time.Sleep(50 * time.Millisecond) + cancel() + wg.Wait() + test.Assert(t, callbackCalled) + test.Assert(t, readErr.Error() == "custom callback error", readErr) + pipe.Close() + }) + t.Run("pipe ctx cancel triggers pipe callback", func(t *testing.T) { + var pipeCallbackCalled int32 + pipeCtx, cancelPipe := context.WithCancel(context.Background()) + pipe := NewPipe[int](pipeCtx, func(ctx context.Context) error { + atomic.AddInt32(&pipeCallbackCalled, 1) + return errors.New("pipe canceled") + }) + + var perReadCallbackCalled int32 + perReadCallback := func(ctx context.Context) error { + atomic.AddInt32(&perReadCallbackCalled, 1) + return errors.New("should not be called") + } + + var readErr error + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // perReadCtx never cancels + _, readErr = pipe.ReadCtx(context.Background(), perReadCallback, make([]int, 10)) + }() + + time.Sleep(50 * time.Millisecond) + cancelPipe() + wg.Wait() + test.Assert(t, readErr.Error() == "pipe canceled", readErr) + test.Assert(t, atomic.LoadInt32(&pipeCallbackCalled) == 1, atomic.LoadInt32(&pipeCallbackCalled)) + test.Assert(t, atomic.LoadInt32(&perReadCallbackCalled) == 0, atomic.LoadInt32(&perReadCallbackCalled)) + }) + t.Run("perReadCtx timeout triggers perReadCallback", func(t *testing.T) { + var pipeCallbackCalled int32 + // pipe ctx never cancels + pipe := NewPipe[int](context.Background(), func(ctx context.Context) error { + atomic.AddInt32(&pipeCallbackCalled, 1) + return errors.New("should not be called") + }) + + var perReadCallbackCalled int32 + perReadCtx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + perReadCallback := func(ctx context.Context) error { + atomic.AddInt32(&perReadCallbackCalled, 1) + return errors.New("per-read timeout") + } + + _, err := pipe.ReadCtx(perReadCtx, perReadCallback, make([]int, 10)) + test.Assert(t, err.Error() == "per-read timeout", err) + test.Assert(t, atomic.LoadInt32(&perReadCallbackCalled) == 1, atomic.LoadInt32(&perReadCallbackCalled)) + test.Assert(t, atomic.LoadInt32(&pipeCallbackCalled) == 0, atomic.LoadInt32(&pipeCallbackCalled)) + pipe.Close() + }) + t.Run("write and close atomically drains data before EOF", func(t *testing.T) { + testCases := []struct { + name string + writeBuf []int + readSize int + expect int32 + }{ + { + name: "single item", + writeBuf: []int{42}, + readSize: 10, + expect: 42, + }, + { + name: "multiple items", + writeBuf: []int{1, 2, 3}, + readSize: 1, + expect: 6, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + // Verify that when Write + Close happen while reader sees state=closed, + // the reader still drains remaining items before returning EOF. + for i := 0; i < 1000; i++ { + pipe := NewPipe[int](context.Background(), nil) + var total int32 + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + items := make([]int, tc.readSize) + for { + n, err := pipe.Read(items) + for j := 0; j < n; j++ { + atomic.AddInt32(&total, int32(items[j])) + } + if err != nil { + test.Assert(t, err == io.EOF, err) + return + } + } + }() + + // Write and Close back-to-back to maximize the chance of triggering the race: + // reader's first p.read() returns 0, then sees state=closed, + // but the data written here must still be drained by the second p.read(). + pipe.Write(tc.writeBuf...) + pipe.Close() + wg.Wait() + test.Assert(t, atomic.LoadInt32(&total) == tc.expect, atomic.LoadInt32(&total)) + } + }) + } + }) + t.Run("close pipe while ReadCtx is blocked returns EOF", func(t *testing.T) { + pipe := NewPipe[int](context.Background(), nil) + + var readErr error + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + _, readErr = pipe.ReadCtx(context.Background(), nil, make([]int, 10)) + }() + + time.Sleep(50 * time.Millisecond) + pipe.Close() + wg.Wait() + test.Assert(t, readErr == io.EOF, readErr) + }) +} + +func BenchmarkPipeline(b *testing.B) { + pipe := NewPipe[int](context.Background(), nil) + readCache := make([]int, 8) + b.ResetTimer() + b.ReportAllocs() + for i := 0; i < b.N; i++ { + for j := 0; j < len(readCache); j++ { + go pipe.Write(1) + } + total := 0 + for total < len(readCache) { + n, _ := pipe.Read(readCache) + total += n + } + } +} diff --git a/pkg/remote/trans/ttstream/container/queue.go b/pkg/remote/trans/ttstream/internal/container/queue.go similarity index 100% rename from pkg/remote/trans/ttstream/container/queue.go rename to pkg/remote/trans/ttstream/internal/container/queue.go diff --git a/pkg/remote/trans/ttstream/container/queue_test.go b/pkg/remote/trans/ttstream/internal/container/queue_test.go similarity index 100% rename from pkg/remote/trans/ttstream/container/queue_test.go rename to pkg/remote/trans/ttstream/internal/container/queue_test.go diff --git a/pkg/remote/trans/ttstream/container/stack.go b/pkg/remote/trans/ttstream/internal/container/stack.go similarity index 100% rename from pkg/remote/trans/ttstream/container/stack.go rename to pkg/remote/trans/ttstream/internal/container/stack.go diff --git a/pkg/remote/trans/ttstream/container/stack_test.go b/pkg/remote/trans/ttstream/internal/container/stack_test.go similarity index 100% rename from pkg/remote/trans/ttstream/container/stack_test.go rename to pkg/remote/trans/ttstream/internal/container/stack_test.go diff --git a/pkg/remote/trans/ttstream/server_handler.go b/pkg/remote/trans/ttstream/server_handler.go index 1edf7672f0..87a51c9b75 100644 --- a/pkg/remote/trans/ttstream/server_handler.go +++ b/pkg/remote/trans/ttstream/server_handler.go @@ -144,7 +144,7 @@ func (t *svrTransHandler) OnRead(ctx context.Context, conn net.Conn) (err error) for { var st *serverStream // ReadStream will block until a stream coming or conn return error - st, err = trans.ReadStream(ctx) + st, err = trans.ReadStream() if err != nil { return } @@ -176,9 +176,9 @@ func (t *svrTransHandler) OnStream(ctx context.Context, conn net.Conn, st *serve ink := ri.Invocation().(rpcinfo.InvocationSetter) // TODO: support protobuf codec, and make `strict` true when combine service is not supported. - sinfo := t.opt.SvcSearcher.SearchService(st.Service(), st.Method(), false, serviceinfo.Thrift) + sinfo := t.opt.SvcSearcher.SearchService(st.IdlService(), st.Method(), false, serviceinfo.Thrift) if sinfo == nil { - err = remote.NewTransErrorWithMsg(remote.UnknownService, fmt.Sprintf("unknown service %s", st.Service())) + err = remote.NewTransErrorWithMsg(remote.UnknownService, fmt.Sprintf("unknown service %s", st.IdlService())) return } // TODO: pass-through grpc streaming mode. diff --git a/pkg/remote/trans/ttstream/stream.go b/pkg/remote/trans/ttstream/stream.go index 23563f5154..83543d532b 100644 --- a/pkg/remote/trans/ttstream/stream.go +++ b/pkg/remote/trans/ttstream/stream.go @@ -19,11 +19,13 @@ package ttstream import ( "context" "fmt" + "time" "github.com/bytedance/gopkg/lang/mcache" "github.com/cloudwego/gopkg/protocol/thrift" "github.com/cloudwego/gopkg/protocol/ttheader" + "github.com/cloudwego/kitex/pkg/remote/trans/ttstream/internal/container" "github.com/cloudwego/kitex/pkg/rpcinfo" "github.com/cloudwego/kitex/pkg/streaming" ktransport "github.com/cloudwego/kitex/transport" @@ -83,11 +85,12 @@ type stream struct { wheader streaming.Header // wheader == nil means it already be sent wtrailer streaming.Trailer // wtrailer == nil means it already be sent - recvTimeoutConfig streaming.TimeoutConfig - closeCallback []func(error) + recvTimeout time.Duration + recvTimeoutCallback container.CtxDoneCallback + closeCallback []func(error) } -func (s *stream) Service() string { +func (s *stream) IdlService() string { if len(s.header) == 0 { return "" } @@ -102,6 +105,13 @@ func (s *stream) TransportProtocol() ktransport.Protocol { return ktransport.TTHeaderStreaming } +func (s *stream) fromService() string { + if s.rpcInfo != nil { + return s.rpcInfo.From().ServiceName() + } + return "" +} + // SendMsg send a message to peer. // In order to avoid underlying execution errors when the context passed in by the user does not // contain information related to this RPC, the context specified when creating the stream is used @@ -125,16 +135,24 @@ func (s *stream) SendMsg(ctx context.Context, msg any) (err error) { func (s *stream) RecvMsg(ctx context.Context, data any) error { nctx := s.ctx - if s.recvTimeoutConfig.Timeout > 0 { + var payload []byte + var err error + if s.recvTimeout > 0 { + // If s.ctx and the recv context are done at the same time + // (for example, if the stream is canceled just as recv times out), + // the pipe in the reader prioritizes returning the cancellation cause from s.ctx. var cancel context.CancelFunc - nctx, cancel = context.WithTimeout(nctx, s.recvTimeoutConfig.Timeout) + nctx, cancel = context.WithTimeout(nctx, s.recvTimeout) defer cancel() + payload, err = s.reader.outputWithCtx(nctx, s.recvTimeoutCallback) + } else { + payload, err = s.reader.output() } - payload, err := s.reader.output(nctx) if err != nil { return err } err = DecodePayload(nctx, payload, data) + recvSize := len(payload) // payload will not be access after decode mcache.Free(payload) @@ -142,7 +160,7 @@ func (s *stream) RecvMsg(ctx context.Context, data any) error { ri := s.rpcInfo if ri != nil && ri.Stats() != nil { if rpcStats := rpcinfo.AsMutableRPCStats(ri.Stats()); rpcStats != nil { - rpcStats.IncrRecvSize(uint64(len(payload))) + rpcStats.IncrRecvSize(uint64(recvSize)) } } return err @@ -152,24 +170,12 @@ func (s *stream) RegisterCloseCallback(cb func(error)) { s.closeCallback = append(s.closeCallback, cb) } -func (s *stream) setRecvTimeoutConfig(cfg rpcinfo.RPCConfig) { - tmCfg := cfg.StreamRecvTimeoutConfig() - // WithStreamRecvTimeoutConfig has higher priority - if tmCfg.Timeout > 0 { - s.recvTimeoutConfig = tmCfg - return - } - // compatible for WithStreamRecvTimeout(tm), it's equivalent to - // streaming.TimeoutConfig{ - // Timeout: tm, - // DisableCancelRemote: false, - // } - tm := cfg.StreamRecvTimeout() - if tm > 0 { - s.recvTimeoutConfig = streaming.TimeoutConfig{ - Timeout: tm, - } - } +func (s *stream) setRecvTimeout(tm time.Duration) { + s.recvTimeout = tm +} + +func (s *stream) setRecvTimeoutCallback(cb container.CtxDoneCallback) { + s.recvTimeoutCallback = cb } func (s *stream) runCloseCallback(exception error) { @@ -225,6 +231,6 @@ func (s *stream) sendRst(exception error, cancelPath string) (err error) { // === Frame OnRead callback func (s *stream) onReadDataFrame(fr *Frame) (err error) { - s.reader.input(context.Background(), fr.payload) + s.reader.input(fr.payload) return nil } diff --git a/pkg/remote/trans/ttstream/stream_client.go b/pkg/remote/trans/ttstream/stream_client.go index 370788950a..5f18544555 100644 --- a/pkg/remote/trans/ttstream/stream_client.go +++ b/pkg/remote/trans/ttstream/stream_client.go @@ -22,6 +22,7 @@ import ( "fmt" "sync" "sync/atomic" + "time" "github.com/cloudwego/gopkg/protocol/thrift" "github.com/cloudwego/gopkg/protocol/ttheader" @@ -41,7 +42,7 @@ func newClientStream(ctx context.Context, writer streamWriter, smeta streamFrame headerSig: make(chan int32, 1), trailerSig: make(chan int32, 1), } - s.reader = newStreamReaderWithCtxDoneCallback(cs.ctxDoneCallback) + s.reader = newStreamReader(ctx, cs.ctxDoneCallback) return cs } @@ -69,6 +70,9 @@ type clientStream struct { trailerSig chan int32 traceCtl *rpcinfo.TraceController + // streamTimeout records the timeout value carried in the ctx when the stream was created + // it is used to return a more detailed error when the stream times out. + streamTimeout time.Duration } func (s *clientStream) Header() (streaming.Header, error) { @@ -128,17 +132,41 @@ func (s *clientStream) Context() context.Context { return s.ctx } +func (s *clientStream) CancelWithErr(err error) { + svcName := s.fromService() + cancelPath := appendCancelPath("", svcName) + s.close(err, true, cancelPath, nil) +} + // ctxDoneCallback convert ctx.Err() to ttstream related err and close client-side stream. // it is invoked in container.Pipe -func (s *clientStream) ctxDoneCallback(ctx context.Context) { - finalEx, noCancel, cancelPath := s.parseCtxErr(ctx) - +func (s *clientStream) ctxDoneCallback(ctx context.Context) error { + finalEx, noCancel, cancelPath := s.parseStreamCtxErr(ctx) s.close(finalEx, !noCancel, cancelPath, nil) + return finalEx +} + +// recvTimeoutCallback builds and returns ttstream RecvTimeout exception. +// it will close client-side stream lifecycle and rst the remote side. +// it is invoked in container.Pipe +func (s *clientStream) recvTimeoutCallback(ctx context.Context) error { + var ex *Exception + if cErr := ctx.Err(); cErr != context.DeadlineExceeded { + ex = errInternalCancel.newBuilder().withSide(clientSide).withCause(cErr) + } else { + ex = newStreamRecvTimeoutException(s.recvTimeout) + } + + svcName := s.fromService() + cancelPath := appendCancelPath("", svcName) + // send Rst Frame to finish the stream lifecycle + s.close(ex, true, cancelPath, nil) + return ex } -// parseCtxErr parses information in ctx.Err and returning ttstream Exception and cascading cancelPath -func (s *clientStream) parseCtxErr(ctx context.Context) (finalEx *Exception, noCancel bool, cancelPath string) { - svcName := s.rpcInfo.From().ServiceName() +// parseStreamCtxErr parses information in ctx.Err and returning ttstream Exception and cascading cancelPath +func (s *clientStream) parseStreamCtxErr(ctx context.Context) (finalEx *Exception, noCancel bool, cancelPath string) { + svcName := s.fromService() cErr := ctx.Err() switch cErr { // biz code invokes cancel() @@ -146,14 +174,11 @@ func (s *clientStream) parseCtxErr(ctx context.Context) (finalEx *Exception, noC finalEx = errBizCancel.newBuilder().withSide(clientSide) // the initial node sending rst, the original cancelPath is empty cancelPath = appendCancelPath("", svcName) + // stream timeout case context.DeadlineExceeded: - finalEx = newStreamRecvTimeoutException(s.recvTimeoutConfig) - if s.recvTimeoutConfig.DisableCancelRemote { - noCancel = true - } else { - // the initial node sending rst, the original cancelPath is empty - cancelPath = appendCancelPath("", svcName) - } + finalEx = newStreamTimeoutException(s.streamTimeout) + // the initial node sending rst, the original cancelPath is empty + cancelPath = appendCancelPath("", svcName) default: if tEx, ok := cErr.(*Exception); ok { // for cascading cancel case, we need to change the side from server to client @@ -239,6 +264,21 @@ func (s *clientStream) setTraceController(traceCtl *rpcinfo.TraceController) { s.traceCtl = traceCtl } +func (s *clientStream) setStreamTimeout(tm time.Duration) { + s.streamTimeout = tm +} + +func (s *clientStream) setRecvTimeout(tm time.Duration) { + s.stream.setRecvTimeout(tm) + // recv timeout takes effect, configures recvTimeoutCallback for clientStream + // serverStream now does not support recvTimeout + if s.recvTimeout > 0 { + s.setRecvTimeoutCallback(s.recvTimeoutCallback) + } +} + +// === clientStream event handler for tracing + func (s *clientStream) handleStreamStartEvent(event rpcinfo.StreamStartEvent) { if s.traceCtl == nil { return diff --git a/pkg/remote/trans/ttstream/stream_client_test.go b/pkg/remote/trans/ttstream/stream_client_test.go index 7ab350668c..fdd03d9a34 100644 --- a/pkg/remote/trans/ttstream/stream_client_test.go +++ b/pkg/remote/trans/ttstream/stream_client_test.go @@ -23,16 +23,24 @@ import ( "errors" "strings" "sync" + "sync/atomic" "testing" "time" "github.com/cloudwego/kitex/internal/test" + "github.com/cloudwego/kitex/pkg/kerrors" "github.com/cloudwego/kitex/pkg/rpcinfo" - "github.com/cloudwego/kitex/pkg/streaming" ) func newTestClientStream(ctx context.Context) *clientStream { - return newClientStream(ctx, mockStreamWriter{}, streamFrame{}) + return newClientStreamWithWriter(ctx, mockStreamWriter{}) +} + +func newClientStreamWithWriter(ctx context.Context, writer streamWriter) *clientStream { + cs := newClientStream(ctx, writer, streamFrame{method: "testMethod"}) + cs.rpcInfo = rpcinfo.NewRPCInfo( + rpcinfo.NewEndpointInfo("testService", "testMethod", nil, nil), nil, nil, nil, nil) + return cs } func Test_clientStreamStateChange(t *testing.T) { @@ -86,10 +94,11 @@ func Test_clientStreamStateChange(t *testing.T) { }) } -func Test_clientStream_parseCtxErr(t *testing.T) { +func Test_clientStream_parseStreamCtxErr(t *testing.T) { testcases := []struct { desc string ctxFunc func() (context.Context, context.CancelFunc) + waitCtxDone bool // if true, wait for ctx.Done() instead of calling cancel() expectEx *Exception expectNoCancel bool expectCancelPath string @@ -120,37 +129,16 @@ func Test_clientStream_parseCtxErr(t *testing.T) { expectCancelPath: "serviceA,serviceB", }, { - desc: "recv timeout", + desc: "stream ctx timeout", ctxFunc: func() (context.Context, context.CancelFunc) { - cfg := rpcinfo.NewRPCConfig() - tm := 100 * time.Millisecond - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: tm, - }) ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), rpcinfo.NewRPCInfo( - rpcinfo.NewEndpointInfo("serviceA", "testMethod", nil, nil), nil, nil, cfg, nil)) - return context.WithTimeout(ctx, tm) + rpcinfo.NewEndpointInfo("serviceA", "testMethod", nil, nil), nil, nil, rpcinfo.NewRPCConfig(), nil)) + return context.WithTimeout(ctx, 100*time.Millisecond) }, - expectEx: newStreamRecvTimeoutException(streaming.TimeoutConfig{Timeout: 100 * time.Millisecond}), + waitCtxDone: true, + expectEx: newStreamTimeoutException(0), expectCancelPath: "serviceA", }, - { - desc: "recv timeout with DisableCancelRemote", - ctxFunc: func() (context.Context, context.CancelFunc) { - cfg := rpcinfo.NewRPCConfig() - tm := 100 * time.Millisecond - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: tm, - DisableCancelRemote: true, - }) - ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), rpcinfo.NewRPCInfo( - rpcinfo.NewEndpointInfo("serviceA", "testMethod", nil, nil), nil, nil, cfg, nil)) - return context.WithTimeout(ctx, tm) - }, - expectEx: newStreamRecvTimeoutException(streaming.TimeoutConfig{Timeout: 100 * time.Millisecond, DisableCancelRemote: true}), - expectNoCancel: true, - expectCancelPath: "", - }, { desc: "other customized ctx Err", ctxFunc: func() (context.Context, context.CancelFunc) { @@ -170,17 +158,13 @@ func Test_clientStream_parseCtxErr(t *testing.T) { for _, tc := range testcases { t.Run(tc.desc, func(t *testing.T) { ctx, cancel := tc.ctxFunc() - cs := newClientStream(ctx, nil, streamFrame{}) - ri := rpcinfo.GetRPCInfo(ctx) - if ri.Config().StreamRecvTimeoutConfig().Timeout > 0 { - cs.setRecvTimeoutConfig(ri.Config()) - // wait for timeout + cs := newClientStream(ctx, mockStreamWriter{}, streamFrame{}) + if tc.waitCtxDone { <-ctx.Done() } else { - // cancel directly cancel() } - finalEx, noCancel, cancelPath := cs.parseCtxErr(ctx) + finalEx, noCancel, cancelPath := cs.parseStreamCtxErr(ctx) test.DeepEqual(t, finalEx, tc.expectEx) test.Assert(t, tc.expectCancelPath == cancelPath, cancelPath) test.Assert(t, tc.expectNoCancel == noCancel, noCancel) @@ -216,3 +200,198 @@ func Test_clientStream_SendMsg(t *testing.T) { err = cs.SendMsg(ctx, req) test.Assert(t, errors.Is(err, errDownstreamCancel), err) } + +// Test_CancelWithErr_concurrent tests that CancelWithErr (outer layer) racing with +// ctxDoneCallback (pipe stream ctx) is safe. This simulates the dual-timer scenario: +// outer callWithTimeout fires at the same moment as stream ctx deadline. +func Test_CancelWithErr_concurrent(t *testing.T) { + t.Run("CancelWithErr races with ctxDoneCallback", func(t *testing.T) { + // Run multiple iterations to increase chance of hitting the race window + for i := 0; i < 100; i++ { + var rstCount int32 + writer := mockStreamWriter{ + writeFrameFunc: func(f *Frame) error { + if f.typ == rstFrameType { + atomic.AddInt32(&rstCount, 1) + } + return nil + }, + } + // stream ctx with a very short deadline + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + cs := newClientStreamWithWriter(ctx, writer) + cs.setStreamTimeout(10 * time.Millisecond) + + outerErr := errBizCancel.newBuilder().withSide(clientSide) + + var wg sync.WaitGroup + wg.Add(2) + + // goroutine 1: simulate outer layer calling CancelWithErr on timeout + go func() { + defer wg.Done() + <-ctx.Done() // wait for ctx to expire, then race + cs.CancelWithErr(outerErr) + }() + + // goroutine 2: RecvMsg blocks in pipe, ctxDoneCallback fires on ctx expiry + go func() { + defer wg.Done() + _ = cs.RecvMsg(cs.ctx, nil) + }() + + wg.Wait() + cancel() + + test.Assert(t, atomic.LoadInt32(&cs.state) == streamStateInactive) + ex := cs.closeStreamException.Load() + test.Assert(t, ex != nil) + test.Assert(t, atomic.LoadInt32(&rstCount) == 1, atomic.LoadInt32(&rstCount)) + } + }) + t.Run("concurrent CancelWithErr calls", func(t *testing.T) { + for i := 0; i < 100; i++ { + writer := mockStreamWriter{ + writeFrameFunc: func(f *Frame) error { return nil }, + } + cs := newClientStreamWithWriter(context.Background(), writer) + + ex1 := errBizCancel.newBuilder().withSide(clientSide) + ex2 := errDownstreamCancel.newBuilder().withSide(clientSide) + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + cs.CancelWithErr(ex1) + }() + go func() { + defer wg.Done() + cs.CancelWithErr(ex2) + }() + wg.Wait() + + test.Assert(t, atomic.LoadInt32(&cs.state) == streamStateInactive) + stored := cs.closeStreamException.Load() + test.Assert(t, stored != nil) + storedErr := stored.(error) + test.Assert(t, errors.Is(storedErr, errBizCancel) || errors.Is(storedErr, errDownstreamCancel), storedErr) + } + }) +} + +func Test_recvTimeout(t *testing.T) { + t.Run("recv timeout closes stream and sends RST", func(t *testing.T) { + var rstFrameCount int32 + closeStreamCh := make(chan struct{}) + writer := mockStreamWriter{ + writeFrameFunc: func(f *Frame) error { + if f.typ == rstFrameType { + atomic.AddInt32(&rstFrameCount, 1) + } + return nil + }, + closeStreamFunc: func(sid int32) error { + close(closeStreamCh) + return nil + }, + } + cs := newClientStreamWithWriter(context.Background(), writer) + cs.setRecvTimeout(50 * time.Millisecond) + + rErr := cs.RecvMsg(cs.ctx, nil) + test.Assert(t, rErr != nil, rErr) + ex, ok := rErr.(*Exception) + test.Assert(t, ok, rErr) + test.Assert(t, ex.TypeId() == errRecvTimeoutTypeId, ex) + test.Assert(t, errors.Is(rErr, kerrors.ErrStreamingTimeout), rErr) + t.Log(rErr) + + select { + case <-closeStreamCh: + case <-time.After(time.Second): + t.Fatal("timed out waiting for CloseStream") + } + test.Assert(t, atomic.LoadInt32(&rstFrameCount) == 1, atomic.LoadInt32(&rstFrameCount)) + }) + t.Run("stream ctx expires before recv timeout, returns stream timeout error", func(t *testing.T) { + var rstFrameCount int32 + closeStreamCh := make(chan struct{}) + writer := mockStreamWriter{ + writeFrameFunc: func(f *Frame) error { + if f.typ == rstFrameType { + atomic.AddInt32(&rstFrameCount, 1) + } + return nil + }, + closeStreamFunc: func(sid int32) error { + close(closeStreamCh) + return nil + }, + } + // stream ctx expires in 50ms, recv timeout is 300ms + streamTimeout := 50 * time.Millisecond + ctx, cancel := context.WithTimeout(context.Background(), streamTimeout) + defer cancel() + cs := newClientStreamWithWriter(ctx, writer) + cs.setStreamTimeout(streamTimeout) + cs.setRecvTimeout(300 * time.Millisecond) + + // RecvMsg should block until stream ctx expires (50ms), NOT until recv timeout (300ms) + start := time.Now() + rErr := cs.RecvMsg(cs.ctx, nil) + elapsed := time.Since(start) + test.Assert(t, rErr != nil, rErr) + ex, ok := rErr.(*Exception) + test.Assert(t, ok, rErr) + // must be stream timeout (errStreamTimeoutTypeId), NOT recv timeout (errRecvTimeoutTypeId) + test.Assert(t, ex.TypeId() == errStreamTimeoutTypeId, ex.TypeId()) + t.Log(rErr) + + // should return around 50ms, not 300ms + test.Assert(t, elapsed < 200*time.Millisecond, elapsed) + + select { + case <-closeStreamCh: + case <-time.After(time.Second): + t.Fatal("timed out waiting for CloseStream") + } + test.Assert(t, atomic.LoadInt32(&rstFrameCount) > 0, atomic.LoadInt32(&rstFrameCount)) + }) + t.Run("recv timeout with stream ctx, recv timeout fires first", func(t *testing.T) { + var rstFrameCount int32 + closeStreamCh := make(chan struct{}) + writer := mockStreamWriter{ + writeFrameFunc: func(f *Frame) error { + if f.typ == rstFrameType { + atomic.AddInt32(&rstFrameCount, 1) + } + return nil + }, + closeStreamFunc: func(sid int32) error { + close(closeStreamCh) + return nil + }, + } + // stream ctx expires in 300ms, recv timeout is 50ms + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + cs := newClientStreamWithWriter(ctx, writer) + cs.setRecvTimeout(50 * time.Millisecond) + + rErr := cs.RecvMsg(cs.ctx, nil) + test.Assert(t, rErr != nil, rErr) + ex, ok := rErr.(*Exception) + test.Assert(t, ok, rErr) + test.Assert(t, ex.TypeId() == errRecvTimeoutTypeId, ex) + test.Assert(t, errors.Is(rErr, kerrors.ErrStreamingTimeout), rErr) + t.Log(rErr) + + select { + case <-closeStreamCh: + case <-time.After(time.Second): + t.Fatal("timed out waiting for CloseStream") + } + test.Assert(t, atomic.LoadInt32(&rstFrameCount) > 0, atomic.LoadInt32(&rstFrameCount)) + }) +} diff --git a/pkg/remote/trans/ttstream/stream_reader.go b/pkg/remote/trans/ttstream/stream_reader.go index accbcec1cf..9acbf2daba 100644 --- a/pkg/remote/trans/ttstream/stream_reader.go +++ b/pkg/remote/trans/ttstream/stream_reader.go @@ -18,11 +18,9 @@ package ttstream import ( "context" - "errors" "io" - "github.com/cloudwego/kitex/pkg/klog" - "github.com/cloudwego/kitex/pkg/remote/trans/ttstream/container" + "github.com/cloudwego/kitex/pkg/remote/trans/ttstream/internal/container" ) // streamReader is an abstraction layer for stream level IO operations @@ -37,39 +35,39 @@ type streamMsg struct { exception error } -func newStreamReader() *streamReader { +func newStreamReader(ctx context.Context, callback container.CtxDoneCallback) *streamReader { sio := new(streamReader) - sio.pipe = container.NewPipe[streamMsg]() + sio.pipe = container.NewPipe[streamMsg](ctx, callback) return sio } -func newStreamReaderWithCtxDoneCallback(callback container.CtxDoneCallback) *streamReader { - sio := new(streamReader) - sio.pipe = container.NewPipe[streamMsg](container.WithCtxDoneCallback(callback)) - return sio +func (s *streamReader) input(payload []byte) { + s.pipe.Write(streamMsg{payload: payload}) } -func (s *streamReader) input(ctx context.Context, payload []byte) { - err := s.pipe.Write(ctx, streamMsg{payload: payload}) - if err != nil { - klog.Errorf("stream pipe input failed: %v", err) +// output would return err in the following scenarios: +// - pipe finished: io.EOF +// - pipe ctx Done() triggered: callback error or ctx.Err() +// - trailer frame contains err: streamMsg.exception +func (s *streamReader) output() (payload []byte, err error) { + if s.exception != nil { + return nil, s.exception } + n, err := s.pipe.Read(s.cache[:]) + return s.handleReadResult(n, err) } -// output would return err in the following scenarios: -// - pipe finished: container.ErrPipeEOF, container.ErrPipeCanceled -// - ctx Done() triggered: ctx.Err() -// - trailer frame contains err: streamMsg.exception -func (s *streamReader) output(ctx context.Context) (payload []byte, err error) { +// outputWithCtx additionally checks a per-read ctx alongside the pipe's own ctx. +func (s *streamReader) outputWithCtx(ctx context.Context, perReadCallback container.CtxDoneCallback) (payload []byte, err error) { if s.exception != nil { return nil, s.exception } + n, err := s.pipe.ReadCtx(ctx, perReadCallback, s.cache[:]) + return s.handleReadResult(n, err) +} - n, err := s.pipe.Read(ctx, s.cache[:]) +func (s *streamReader) handleReadResult(n int, err error) ([]byte, error) { if err != nil { - if errors.Is(err, container.ErrPipeEOF) { - err = io.EOF - } s.exception = err return nil, s.exception } @@ -87,7 +85,7 @@ func (s *streamReader) output(ctx context.Context) (payload []byte, err error) { func (s *streamReader) close(exception error) { if exception != nil { - _ = s.pipe.Write(context.Background(), streamMsg{exception: exception}) + s.pipe.Write(streamMsg{exception: exception}) } s.pipe.Close() } diff --git a/pkg/remote/trans/ttstream/stream_reader_test.go b/pkg/remote/trans/ttstream/stream_reader_test.go index 2d4c651b04..701fad8ce8 100644 --- a/pkg/remote/trans/ttstream/stream_reader_test.go +++ b/pkg/remote/trans/ttstream/stream_reader_test.go @@ -31,44 +31,44 @@ func TestStreamReader(t *testing.T) { round := 10000 // basic IOs - sio := newStreamReader() + sio := newStreamReader(ctx, nil) doneCh := make(chan struct{}) go func() { for i := 0; i < round; i++ { - sio.input(ctx, msg) + sio.input(msg) } close(doneCh) }() for i := 0; i < round; i++ { - payload, err := sio.output(ctx) + payload, err := sio.output() test.Assert(t, err == nil, err) test.DeepEqual(t, msg, payload) } // wait for goroutine exited <-doneCh // exception IOs - sio.input(ctx, msg) + sio.input(msg) targetErr := errors.New("test") sio.close(targetErr) - payload, err := sio.output(ctx) + payload, err := sio.output() test.Assert(t, err == nil, err) test.DeepEqual(t, msg, payload) - payload, err = sio.output(ctx) + payload, err = sio.output() test.Assert(t, payload == nil, payload) test.Assert(t, errors.Is(err, targetErr), err) // ctx canceled IOs - ctx, cancel := context.WithCancel(ctx) - sio = newStreamReader() - sio.input(ctx, msg) + ctx2, cancel := context.WithCancel(context.Background()) + sio = newStreamReader(ctx2, nil) + sio.input(msg) go func() { time.Sleep(time.Millisecond * 10) cancel() }() - payload, err = sio.output(ctx) + payload, err = sio.output() test.Assert(t, err == nil, err) test.DeepEqual(t, msg, payload) - payload, err = sio.output(ctx) + payload, err = sio.output() test.Assert(t, payload == nil, payload) test.Assert(t, errors.Is(err, context.Canceled), err) } diff --git a/pkg/remote/trans/ttstream/stream_server.go b/pkg/remote/trans/ttstream/stream_server.go index 14d5e3b7bd..5ba942bca2 100644 --- a/pkg/remote/trans/ttstream/stream_server.go +++ b/pkg/remote/trans/ttstream/stream_server.go @@ -33,7 +33,7 @@ var _ ServerStreamMeta = (*serverStream)(nil) func newServerStream(ctx context.Context, writer streamWriter, smeta streamFrame) *serverStream { s := newBasicStream(ctx, writer, smeta) - s.reader = newStreamReader() + s.reader = newStreamReader(ctx, nil) return &serverStream{stream: s} } diff --git a/pkg/remote/trans/ttstream/stream_test.go b/pkg/remote/trans/ttstream/stream_test.go index 69c5632acc..bfa6afd262 100644 --- a/pkg/remote/trans/ttstream/stream_test.go +++ b/pkg/remote/trans/ttstream/stream_test.go @@ -25,8 +25,6 @@ import ( "time" "github.com/cloudwego/kitex/internal/test" - "github.com/cloudwego/kitex/pkg/rpcinfo" - "github.com/cloudwego/kitex/pkg/streaming" ) func TestGenericStreaming(t *testing.T) { @@ -45,18 +43,6 @@ func TestGenericStreaming(t *testing.T) { test.Assert(t, err == nil, err) test.Assert(t, res.A == req.A) test.Assert(t, res.B == req.B) - - // map generic - // client side - // reqJSON := `{"A":123, "b":"hello"}` - // err = cs.SendMsg(ctx, reqJSON) - // test.Assert(t, err == nil, err) - // server side - // res = new(testResponse) - // err = ss.RecvMsg(ctx, res) - // test.Assert(t, err == nil, err) - // test.Assert(t, res.A == req.A) - // test.Assert(t, res.B == req.B) } // TestStreamRecvTimeout tests that RecvMsg correctly handles timeout scenarios @@ -65,9 +51,7 @@ func TestStreamRecvTimeout(t *testing.T) { test.Assert(t, err == nil, err) // Set a very short timeout for testing - cfg := rpcinfo.NewRPCConfig() - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeout(time.Millisecond * 10) - ss.setRecvTimeoutConfig(cfg) + ss.setRecvTimeout(10 * time.Millisecond) // Create a context that won't expire ctx := context.Background() @@ -108,57 +92,3 @@ func TestStreamRecvWithCancellation(t *testing.T) { test.Assert(t, res.A == req.A) test.Assert(t, res.B == req.B) } - -// TestSetRecvTimeoutConfig tests setRecvTimeoutConfig method with different configuration scenarios -func TestSetRecvTimeoutConfig(t *testing.T) { - t.Run("only set StreamRecvTimeout", func(t *testing.T) { - _, ss, err := newTestStreamPipe(testServiceInfo, "Bidi") - test.Assert(t, err == nil, err) - - cfg := rpcinfo.NewRPCConfig() - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeout(100 * time.Millisecond) - - ss.setRecvTimeoutConfig(cfg) - - // verify timeout config is set from StreamRecvTimeout - test.Assert(t, ss.recvTimeoutConfig.Timeout == 100*time.Millisecond, ss.recvTimeoutConfig) - test.Assert(t, ss.recvTimeoutConfig.DisableCancelRemote == false, ss.recvTimeoutConfig) - }) - - t.Run("only set StreamRecvTimeoutConfig", func(t *testing.T) { - _, ss, err := newTestStreamPipe(testServiceInfo, "Bidi") - test.Assert(t, err == nil, err) - - cfg := rpcinfo.NewRPCConfig() - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: 200 * time.Millisecond, - DisableCancelRemote: true, - }) - - ss.setRecvTimeoutConfig(cfg) - - // verify timeout config is set from StreamRecvTimeoutConfig - test.Assert(t, ss.recvTimeoutConfig.Timeout == 200*time.Millisecond, ss.recvTimeoutConfig) - test.Assert(t, ss.recvTimeoutConfig.DisableCancelRemote == true, ss.recvTimeoutConfig) - }) - - t.Run("both set, StreamRecvTimeoutConfig has higher priority", func(t *testing.T) { - _, ss, err := newTestStreamPipe(testServiceInfo, "Bidi") - test.Assert(t, err == nil, err) - - cfg := rpcinfo.NewRPCConfig() - // set StreamRecvTimeout first with a longer duration - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeout(500 * time.Millisecond) - // then set StreamRecvTimeoutConfig with a shorter duration - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: 50 * time.Millisecond, - DisableCancelRemote: true, - }) - - ss.setRecvTimeoutConfig(cfg) - - // verify StreamRecvTimeoutConfig takes priority over StreamRecvTimeout - test.Assert(t, ss.recvTimeoutConfig.Timeout == 50*time.Millisecond, ss.recvTimeoutConfig) - test.Assert(t, ss.recvTimeoutConfig.DisableCancelRemote == true, ss.recvTimeoutConfig) - }) -} diff --git a/pkg/remote/trans/ttstream/test_utils.go b/pkg/remote/trans/ttstream/test_utils.go index 7b62e37859..a7a4f118ce 100644 --- a/pkg/remote/trans/ttstream/test_utils.go +++ b/pkg/remote/trans/ttstream/test_utils.go @@ -78,7 +78,7 @@ func newTestStreamPipe(sinfo *serviceinfo.ServiceInfo, method string) (*clientSt return nil, nil, err } strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() if err != nil { return nil, nil, err } diff --git a/pkg/remote/trans/ttstream/trace_test.go b/pkg/remote/trans/ttstream/trace_test.go index 8dabeaadf1..0011a9948f 100644 --- a/pkg/remote/trans/ttstream/trace_test.go +++ b/pkg/remote/trans/ttstream/trace_test.go @@ -137,7 +137,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -215,7 +215,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -299,7 +299,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -381,7 +381,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) // Server closes immediately with no data @@ -431,7 +431,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) // Server sends header @@ -486,7 +486,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) // Server sends one message @@ -552,7 +552,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) // Cancel immediately after stream creation @@ -609,7 +609,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) // Send one message @@ -672,7 +672,7 @@ func TestTrace_trace(t *testing.T) { cliTraceCtl.HandleStreamStartEvent(ctx, cs.rpcInfo, rpcinfo.StreamStartEvent{}) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup diff --git a/pkg/remote/trans/ttstream/transport_server.go b/pkg/remote/trans/ttstream/transport_server.go index 7c98176a10..bd93cfcaf7 100644 --- a/pkg/remote/trans/ttstream/transport_server.go +++ b/pkg/remote/trans/ttstream/transport_server.go @@ -18,8 +18,6 @@ package ttstream import ( "context" - "errors" - "io" "net" "sync" @@ -28,7 +26,7 @@ import ( "github.com/cloudwego/kitex/pkg/gofunc" "github.com/cloudwego/kitex/pkg/klog" - "github.com/cloudwego/kitex/pkg/remote/trans/ttstream/container" + "github.com/cloudwego/kitex/pkg/remote/trans/ttstream/internal/container" ) // transport is used to read/write frames and disturbed frames to different streams @@ -51,9 +49,10 @@ func newServerTransport(conn netpoll.Connection) *serverTransport { // TODO: let it configurable _ = conn.SetReadTimeout(0) t := &serverTransport{ - conn: conn, - streams: sync.Map{}, - spipe: container.NewPipe[*serverStream](), + conn: conn, + streams: sync.Map{}, + // serverTransport does not have related ctx to control lifecycle, using context.Background() is okay + spipe: container.NewPipe[*serverStream](context.Background(), nil), scache: make([]*serverStream, 0, streamCacheSize), writer: newWriterBuffer(conn.Writer()), closedTrigger: make(chan struct{}, 1), @@ -172,7 +171,7 @@ func (t *serverTransport) readFrame(reader bufiox.Reader) error { s = newServerStream(ctx, t, fr.streamFrame) s.cancelFunc = cFunc t.storeStream(s) - err = t.spipe.Write(context.Background(), s) + t.spipe.Write(s) } else { // load exist stream var ok bool @@ -181,8 +180,6 @@ func (t *serverTransport) readFrame(reader bufiox.Reader) error { // there is a race condition that server handler returns and client sends rst frame concurrently. // then serverTransport would not find the target stream when receiving the rst frame. klog.Debugf("transport[%s] read a unknown stream: frame[%s]", t.Addr(), fr.String()) - // ignore unknown stream error - err = nil } else { // process different frames switch fr.typ { @@ -245,18 +242,15 @@ func (t *serverTransport) CloseStream(sid int32) (err error) { // ReadStream wait for a new incoming stream on current connection // it's typically used by server side -func (t *serverTransport) ReadStream(ctx context.Context) (*serverStream, error) { +func (t *serverTransport) ReadStream() (*serverStream, error) { READ: if len(t.scache) > 0 { s := t.scache[len(t.scache)-1] t.scache = t.scache[:len(t.scache)-1] return s, nil } - n, err := t.spipe.Read(ctx, t.scache[0:streamCacheSize]) + n, err := t.spipe.Read(t.scache[0:streamCacheSize]) if err != nil { - if errors.Is(err, container.ErrPipeEOF) { - return nil, io.EOF - } return nil, err } if n == 0 { diff --git a/pkg/remote/trans/ttstream/transport_test.go b/pkg/remote/trans/ttstream/transport_test.go index 73f8c68ac8..8e7320f485 100644 --- a/pkg/remote/trans/ttstream/transport_test.go +++ b/pkg/remote/trans/ttstream/transport_test.go @@ -82,7 +82,7 @@ func TestTransportBasic(t *testing.T) { err = ctrans.WriteStream(ctx, cs, intHeader, strHeader) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -151,7 +151,7 @@ func TestTransportServerStreaming(t *testing.T) { err = ctrans.WriteStream(ctx, cs, intHeader, strHeader) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -216,7 +216,7 @@ func TestTransportException(t *testing.T) { err = ctrans.WriteStream(ctx, cs, make(IntHeader), make(streaming.Header)) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) res := new(testResponse) res.A = 123 @@ -245,7 +245,7 @@ func TestTransportException(t *testing.T) { cs = newClientStream(ctx, ctrans, streamFrame{sid: genStreamID(), method: "Bidi"}) err = ctrans.WriteStream(ctx, cs, make(IntHeader), make(streaming.Header)) test.Assert(t, err == nil, err) - ss, err = strans.ReadStream(context.Background()) + ss, err = strans.ReadStream() test.Assert(t, err == nil, err) test.Assert(t, ss != nil) _, err = sconn.Writer().WriteBinary([]byte("helloxxxxxxxxxxxxxxxxxxxxxx")) @@ -275,7 +275,7 @@ func TestTransportClose(t *testing.T) { err = ctrans.WriteStream(ctx, cs, intHeader, strHeader) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -353,7 +353,7 @@ func Test_clientStreamReceiveTrailer(t *testing.T) { err = ctrans.WriteStream(ctx, cs, intHeader, strHeader) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -577,7 +577,7 @@ func initTestStreams(t *testing.T, cCtx context.Context, method, cliNodeName, sr err = ctrans.WriteStream(cCtx, cs, intHeader, strHeader) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) ss.rpcInfo = rpcinfo.NewRPCInfo( rpcinfo.NewEndpointInfo(srvNodeName, method, nil, nil), nil, nil, nil, nil) @@ -1169,52 +1169,24 @@ func TestProxyCancel(t *testing.T) { } func TestRecvTimeout(t *testing.T) { - t.Run("ClientStreaming - only StreamRecvTimeout", func(t *testing.T) { - testRecvTimeoutClientStreaming(t, false, true, 50*time.Millisecond, 0) + t.Run("ClientStreaming", func(t *testing.T) { + testRecvTimeoutClientStreaming(t, 50*time.Millisecond) }) - t.Run("ClientStreaming - only StreamRecvTimeoutConfig", func(t *testing.T) { - testRecvTimeoutClientStreaming(t, true, false, 0, 50*time.Millisecond) + t.Run("ServerStreaming", func(t *testing.T) { + testRecvTimeoutServerStreaming(t, 50*time.Millisecond) }) - t.Run("ClientStreaming - both set, StreamRecvTimeoutConfig has higher priority", func(t *testing.T) { - testRecvTimeoutClientStreaming(t, true, true, 200*time.Millisecond, 50*time.Millisecond) - }) - t.Run("ServerStreaming - only StreamRecvTimeout", func(t *testing.T) { - testRecvTimeoutServerStreaming(t, false, true, 50*time.Millisecond, 0) - }) - t.Run("ServerStreaming - only StreamRecvTimeoutConfig", func(t *testing.T) { - testRecvTimeoutServerStreaming(t, true, false, 0, 50*time.Millisecond) - }) - t.Run("ServerStreaming - both set, StreamRecvTimeoutConfig has higher priority", func(t *testing.T) { - testRecvTimeoutServerStreaming(t, true, true, 200*time.Millisecond, 50*time.Millisecond) - }) - t.Run("BidiStreaming - only StreamRecvTimeout", func(t *testing.T) { - testRecvTimeoutBidiStreaming(t, false, true, 50*time.Millisecond, 0) - }) - t.Run("BidiStreaming - only StreamRecvTimeoutConfig", func(t *testing.T) { - testRecvTimeoutBidiStreaming(t, true, false, 0, 50*time.Millisecond) - }) - t.Run("BidiStreaming - both set, StreamRecvTimeoutConfig has higher priority", func(t *testing.T) { - testRecvTimeoutBidiStreaming(t, true, true, 200*time.Millisecond, 50*time.Millisecond) + t.Run("BidiStreaming", func(t *testing.T) { + testRecvTimeoutBidiStreaming(t, 50*time.Millisecond) }) } -func testRecvTimeoutClientStreaming(t *testing.T, setTimeoutConfig, setRecvTimeout bool, recvTimeout, configTimeout time.Duration) { +func testRecvTimeoutClientStreaming(t *testing.T, recvTimeout time.Duration) { cliNodeName := "ttstream client" srvNodeName := "ttstream server" method := "ClientStreaming" cliSt, srvSt := initTestStreams(t, context.Background(), method, cliNodeName, srvNodeName) - - cfg := rpcinfo.NewRPCConfig() - if setRecvTimeout { - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeout(recvTimeout) - } - if setTimeoutConfig { - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: configTimeout, - }) - } - cliSt.setRecvTimeoutConfig(cfg) + cliSt.setRecvTimeout(recvTimeout) var wg sync.WaitGroup wg.Add(1) @@ -1250,24 +1222,13 @@ func testRecvTimeoutClientStreaming(t *testing.T, setTimeoutConfig, setRecvTimeo wg.Wait() } -func testRecvTimeoutServerStreaming(t *testing.T, setTimeoutConfig, setRecvTimeout bool, recvTimeout, configTimeout time.Duration) { +func testRecvTimeoutServerStreaming(t *testing.T, recvTimeout time.Duration) { cliNodeName := "ttstream client" srvNodeName := "ttstream server" method := "ServerStreaming" cliSt, srvSt := initTestStreams(t, context.Background(), method, cliNodeName, srvNodeName) - - cfg := rpcinfo.NewRPCConfig() - if setRecvTimeout { - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeout(recvTimeout) - } - if setTimeoutConfig { - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: configTimeout, - DisableCancelRemote: false, - }) - } - cliSt.setRecvTimeoutConfig(cfg) + cliSt.setRecvTimeout(recvTimeout) var wg sync.WaitGroup wg.Add(1) @@ -1305,23 +1266,13 @@ func testRecvTimeoutServerStreaming(t *testing.T, setTimeoutConfig, setRecvTimeo wg.Wait() } -func testRecvTimeoutBidiStreaming(t *testing.T, setTimeoutConfig, setRecvTimeout bool, recvTimeout, configTimeout time.Duration) { +func testRecvTimeoutBidiStreaming(t *testing.T, recvTimeout time.Duration) { cliNodeName := "ttstream client" srvNodeName := "ttstream server" method := "BidiStreaming" cliSt, srvSt := initTestStreams(t, context.Background(), method, cliNodeName, srvNodeName) - - cfg := rpcinfo.NewRPCConfig() - if setRecvTimeout { - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeout(recvTimeout) - } - if setTimeoutConfig { - rpcinfo.AsMutableRPCConfig(cfg).SetStreamRecvTimeoutConfig(streaming.TimeoutConfig{ - Timeout: configTimeout, - }) - } - cliSt.setRecvTimeoutConfig(cfg) + cliSt.setRecvTimeout(recvTimeout) var wg sync.WaitGroup wg.Add(1) @@ -1377,7 +1328,7 @@ func testSendFailedClientStreaming(t *testing.T) { err = ctrans.WriteStream(ctx, cs, intHeader, strHeader) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -1439,7 +1390,7 @@ func testSendFailedServerStreaming(t *testing.T) { err = ctrans.WriteStream(ctx, cs, intHeader, strHeader) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup @@ -1521,7 +1472,7 @@ func testSendFailedBidiStreaming(t *testing.T) { err = ctrans.WriteStream(ctx, cs, intHeader, strHeader) test.Assert(t, err == nil, err) strans := newServerTransport(sconn) - ss, err := strans.ReadStream(context.Background()) + ss, err := strans.ReadStream() test.Assert(t, err == nil, err) var wg sync.WaitGroup diff --git a/pkg/streaming/streaming.go b/pkg/streaming/streaming.go index 9f923e95aa..576c36dd69 100644 --- a/pkg/streaming/streaming.go +++ b/pkg/streaming/streaming.go @@ -25,6 +25,16 @@ import ( ) // Stream both client and server stream +// +// On client side, once RecvMsg or SendMsg returns a non-nil error, the stream should be considered +// finished and no further calls should be made on it. +// On server side, RecvMsg returning io.EOF only indicates the client-to-server direction is done; +// the server may still call SendMsg. After RecvMsg returns io.EOF, the caller should not make +// further RecvMsg calls, but SendMsg is not affected. +// After SendMsg returns a non-nil error or RecvMsg returns any error other than io.EOF, +// the stream is aborted. +// If subsequent calls are made regardless, they will return the same error. +// // Deprecated: It's only for gRPC, use ClientStream or ServerStream instead. type Stream interface { // SetHeader sets the header metadata. It may be called multiple times. @@ -47,13 +57,21 @@ type Stream interface { Trailer() metadata.MD // Context the stream context.Context Context() context.Context - // RecvMsg recvive message from peer - // will block until an error or a message received - // not concurrent-safety + // RecvMsg receives a message from the peer. + // It blocks until a message is received or an error occurs. + // On client side, it returns io.EOF when the stream is done. + // On server side, it returns io.EOF when the client-to-server direction is done (i.e. the client + // called Close()); the server may still call SendMsg after that. + // After RecvMsg returns a non-nil error (including io.EOF), the caller should not make further + // RecvMsg calls. If subsequent calls are made regardless, they will return the same error. + // For non-EOF errors, the stream is aborted entirely. + // It is not concurrent-safe. RecvMsg(m interface{}) error - // SendMsg send message to peer - // will block until an error or enough buffer to send - // not concurrent-safety + // SendMsg sends a message to the peer. + // It blocks until the message is sent or an error occurs. + // After SendMsg returns a non-nil error, the stream is aborted and + // subsequent calls will return the same error. + // It is not concurrent-safe. SendMsg(m interface{}) error // not concurrent-safety with SendMsg io.Closer diff --git a/pkg/streaming/streamx.go b/pkg/streaming/streamx.go index 42f6ed6aff..bd12cdf24d 100644 --- a/pkg/streaming/streamx.go +++ b/pkg/streaming/streamx.go @@ -72,15 +72,24 @@ type ( Trailer map[string]string ) -// ClientStream define client stream APIs +// ClientStream defines client-side streaming APIs. +// +// Once RecvMsg or SendMsg returns a non-nil error, the stream should be considered finished +// and the caller should not make further Send/Recv calls on it. +// If subsequent calls are made regardless, they will return the same error. type ClientStream interface { - // SendMsg send message to peer - // will block until an error or enough buffer to send - // not concurrent-safety + // SendMsg sends a message to the peer. + // It blocks until the message is sent or an error occurs. + // After SendMsg returns a non-nil error, the caller should not make further SendMsg calls. + // If subsequent calls are made regardless, they will return the same error. + // It is not concurrent-safe. SendMsg(ctx context.Context, m any) error - // RecvMsg recvive message from peer - // will block until an error or a message received - // not concurrent-safety + // RecvMsg receives a message from the peer. + // It blocks until a message is received or an error occurs. + // It returns io.EOF when the stream is done. After RecvMsg returns a non-nil error (including + // io.EOF), the caller should not make further RecvMsg calls. + // If subsequent calls are made regardless, they will return the same error. + // It is not concurrent-safe. RecvMsg(ctx context.Context, m any) error // Header returns the header info received from the server if there // is any. It blocks if the header info is not ready to read. If the header info @@ -99,15 +108,28 @@ type ClientStream interface { Context() context.Context } -// ServerStream define server stream APIs +// ServerStream defines server-side streaming APIs. +// +// RecvMsg returning io.EOF only indicates the client-to-server direction is done; +// the server may still call SendMsg after that. +// Once SendMsg returns a non-nil error or RecvMsg returns any error other than io.EOF, +// the stream is aborted and the caller should not make further Send/Recv calls on it. +// If subsequent calls are made regardless, they will return the same error. type ServerStream interface { - // SendMsg send message to peer - // will block until an error or enough buffer to send - // not concurrent-safety + // SendMsg sends a message to the peer. + // It blocks until the message is sent or an error occurs. + // After SendMsg returns a non-nil error, the caller should not make further SendMsg calls. + // If subsequent calls are made regardless, they will return the same error. + // It is not concurrent-safe. SendMsg(ctx context.Context, m any) error - // RecvMsg recvive message from peer - // will block until an error or a message received - // not concurrent-safety + // RecvMsg receives a message from the peer. + // It blocks until a message is received or an error occurs. + // It returns io.EOF when the client-to-server direction is done (i.e. the client called CloseSend); + // the server may still call SendMsg after that. + // After RecvMsg returns a non-nil error (including io.EOF), the caller should not make further + // RecvMsg calls. If subsequent calls are made regardless, they will return the same error. + // For non-EOF errors, the stream is aborted entirely. + // It is not concurrent-safe. RecvMsg(ctx context.Context, m any) error // SetHeader sets the header info. It may be called multiple times. // When call multiple times, all the provided header info will be merged. diff --git a/pkg/streaming/timeout.go b/pkg/streaming/timeout.go index 03312f193a..8991b8d8bc 100644 --- a/pkg/streaming/timeout.go +++ b/pkg/streaming/timeout.go @@ -35,6 +35,8 @@ type TimeoutConfig struct { // However, in certain scenarios (e.g., resume-enabled transfers: A → B → C), A may not wish to cancel B and C upon detecting a timeout. // Instead, it expects B and C to complete one round of streaming communication and cache the results. // This allows A to resume the request from the disconnected point on the next attempt, completing the resume-from-breakpoint process. + // + // Even if DisableCancelRemote=true, no retry will occur after the Recv/Send timeout. DisableCancelRemote bool }