fix(ttstream): use consistent context of stream in ttstream.RecvMsg and fix nil message error of binary generic (cloudwego#1866)

Author: jayantxie

pkg/generic/thrift/raw.go

@@ -49,6 +49,14 @@ var _ MessageWriter = (*RawWriter)(nil)
 
 // Write returns the copy of data
 func (m *RawWriter) Write(ctx context.Context, out bufiox.Writer, msg interface{}, method string, isClient bool, requestBase *base.Base) error {
+	if msg == nil {
+		bw := thrift.NewBufferWriter(out)
+		defer bw.Recycle()
+		if err := bw.WriteFieldStop(); err != nil {
+			return err
+		}
+		return nil
+	}
 	buf, ok := msg.([]byte)
 	if !ok {
 		return fmt.Errorf("thrift binary generic msg is not []byte, method=%v", method)

pkg/generic/thrift/raw_test.go

@@ -23,6 +23,8 @@ import (
 
 	"github.com/cloudwego/gopkg/bufiox"
 	"github.com/cloudwego/gopkg/protocol/thrift"
+
+	"github.com/cloudwego/kitex/internal/test"
 )
 
 func TestRawReader_Read(t *testing.T) {
@@ -46,3 +48,24 @@ func TestRawReader_Read(t *testing.T) {
 		t.Fatalf("expect %v, got %v", nb[:off], data)
 	}
 }
+
+func TestRawWriter_Write(t *testing.T) {
+	w := NewRawWriter()
+	var buf []byte
+	out := bufiox.NewBytesWriter(&buf)
+	// nil message
+	err := w.Write(context.Background(), out, nil, "method", true, nil)
+	test.Assert(t, err == nil)
+	err = out.Flush()
+	test.Assert(t, err == nil)
+	test.Assert(t, len(buf) == 1) // field stop
+	test.Assert(t, buf[0] == byte(thrift.STOP))
+
+	// normal message
+	buf = buf[:0]
+	err = w.Write(context.Background(), out, []byte("hello world"), "method", true, nil)
+	test.Assert(t, err == nil)
+	err = out.Flush()
+	test.Assert(t, err == nil)
+	test.Assert(t, len(buf) == len("hello world"))
+}

pkg/remote/trans/ttstream/stream.go

@@ -125,16 +125,17 @@ 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.recvTimeout > 0 {
 		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, s.recvTimeout)
+		nctx, cancel = context.WithTimeout(nctx, s.recvTimeout)
 		defer cancel()
 	}
-	payload, err := s.reader.output(ctx)
+	payload, err := s.reader.output(nctx)
 	if err != nil {
 		return err
 	}
-	err = DecodePayload(ctx, payload, data)
+	err = DecodePayload(nctx, payload, data)
 	// payload will not be access after decode
 	mcache.Free(payload)
 

pkg/remote/trans/ttstream/stream_test.go

@@ -20,7 +20,9 @@ package ttstream
 
 import (
 	"context"
+	"strings"
 	"testing"
+	"time"
 
 	"github.com/cloudwego/kitex/internal/test"
 )
@@ -54,3 +56,51 @@ func TestGenericStreaming(t *testing.T) {
 	// test.Assert(t, res.A == req.A)
 	// test.Assert(t, res.B == req.B)
 }
+
+// TestStreamRecvTimeout tests that RecvMsg correctly handles timeout scenarios
+func TestStreamRecvTimeout(t *testing.T) {
+	_, ss, err := newTestStreamPipe(testServiceInfo, "Bidi")
+	test.Assert(t, err == nil, err)
+
+	// Set a very short timeout for testing
+	ss.setRecvTimeout(time.Millisecond * 10)
+
+	// Create a context that won't expire
+	ctx := context.Background()
+
+	// Try to receive message - should timeout quickly
+	res := new(testResponse)
+	err = ss.RecvMsg(ctx, res)
+	test.Assert(t, err != nil, "RecvMsg should timeout")
+	test.Assert(t, strings.Contains(err.Error(), "timeout") ||
+		strings.Contains(err.Error(), "deadline exceeded"),
+		"Error should be timeout related")
+}
+
+// TestStreamRecvWithCancellation tests that RecvMsg respects context cancellation
+func TestStreamRecvWithCancellation(t *testing.T) {
+	cs, ss, err := newTestStreamPipe(testServiceInfo, "Bidi")
+	test.Assert(t, err == nil, err)
+
+	// Create a cancellable context with short timeout
+	cancelCtx, cancel := context.WithCancel(context.Background())
+
+	cancel()
+
+	// Send message
+	req := new(testRequest)
+	req.A = 789
+	req.B = "normal_test"
+
+	err = cs.SendMsg(cancelCtx, req)
+	test.Assert(t, err == nil, err)
+
+	// Receive message
+	res := new(testResponse)
+	err = ss.RecvMsg(cancelCtx, res)
+	test.Assert(t, err == nil, err)
+
+	// Verify content
+	test.Assert(t, res.A == req.A)
+	test.Assert(t, res.B == req.B)
+}