feat(ttstream): implement whole stream timeout

Author: DMwangnima

pkg/remote/trans/ttstream/client_provider.go

@@ -18,6 +18,8 @@ package ttstream
 
 import (
 	"context"
+	"strconv"
+	"time"
 
 	"github.com/bytedance/gopkg/cloud/metainfo"
 	"github.com/cloudwego/gopkg/protocol/ttheader"
@@ -88,6 +90,12 @@ func (c clientProvider) NewStream(ctx context.Context, ri rpcinfo.RPCInfo) (stre
 	if strHeader == nil {
 		strHeader = map[string]string{}
 	}
+	// retrieve deadline fron context as the whole stream timeout
+	if ddl, ok := ctx.Deadline(); ok {
+		tm := time.Until(ddl)
+		intHeader[ttheader.RPCTimeout] = strconv.Itoa(int(tm.Milliseconds()))
+	}
+
 	strHeader[ttheader.HeaderIDLServiceName] = invocation.ServiceName()
 	metainfo.SaveMetaInfoToMap(ctx, strHeader)
 

pkg/remote/trans/ttstream/server_provider.go

@@ -21,13 +21,15 @@ import (
 	"fmt"
 	"net"
 	"strconv"
+	"time"
 
 	"github.com/bytedance/gopkg/cloud/metainfo"
 	"github.com/cloudwego/gopkg/protocol/thrift"
 	"github.com/cloudwego/gopkg/protocol/ttheader"
 	"github.com/cloudwego/netpoll"
 
 	"github.com/cloudwego/kitex/pkg/kerrors"
+	"github.com/cloudwego/kitex/pkg/klog"
 	"github.com/cloudwego/kitex/pkg/remote"
 	"github.com/cloudwego/kitex/pkg/remote/trans/ttstream/ktx"
 	"github.com/cloudwego/kitex/pkg/streaming"
@@ -113,6 +115,18 @@ func (s *serverProvider) OnStream(ctx context.Context, conn net.Conn) (context.C
 	// cancel ctx when OnStreamFinish
 	ctx, cancelFunc := ktx.WithCancel(ctx)
 	ctx = context.WithValue(ctx, serverStreamCancelCtxKey{}, cancelFunc)
+	// process whole stream timeout
+	var cancel context.CancelFunc
+	if tmStr, ok := st.meta[ttheader.RPCTimeout]; ok {
+		tm, err := strconv.Atoi(tmStr)
+		if err == nil {
+			ctx, cancel = context.WithTimeout(ctx, time.Duration(tm)*time.Millisecond)
+			st.cancelFunc = cancel
+		} else {
+			klog.CtxErrorf(ctx, "ttstream decode RPCTimeout failed, err: %v", err)
+		}
+	}
+
 	return ctx, ss, nil
 }
 

pkg/remote/trans/ttstream/stream.go

@@ -102,6 +102,7 @@ type stream struct {
 	recvTimeout      time.Duration
 	metaFrameHandler MetaFrameHandler
 	closeCallback    []func(error)
+	cancelFunc       context.CancelFunc // only valid in server stream
 }
 
 func (s *stream) Service() string {
@@ -126,8 +127,7 @@ func (s *stream) TransportProtocol() ktransport.Protocol {
 func (s *stream) SendMsg(ctx context.Context, msg any) (err error) {
 	if state := atomic.LoadInt32(&s.state); state == streamStateHalfCloseLocal || state == streamStateInactive {
 		if ex := s.clientStreamException.Load(); ex == nil {
-			// 这个错误设计是有问题的，不能这么玩
-			// 而且这里可以预定义，不需要每次都生成
+			// todo: predefine error here
 			return errIllegalOperation.WithCause(errors.New("stream is closed send"))
 		} else {
 			return ex.(error)
@@ -155,9 +155,11 @@ func (s *stream) RecvMsg(ctx context.Context, data any) error {
 		ctx, cancel = context.WithTimeout(ctx, s.recvTimeout)
 		defer cancel()
 	}
-	// 在这个环境检测 ctx 是否被 cancel
+	// todo: format the error returned by output
+	// like gRPC ContextError
 	payload, err := s.reader.output(ctx)
 	if err != nil {
+		s.close(err, true, clientTransport)
 		return err
 	}
 	err = DecodePayload(context.Background(), payload, data)
@@ -195,6 +197,10 @@ func (s *stream) close(exception error, rst bool, kind int32) error {
 		// stream has been closed
 		return nil
 	}
+	// todo: think about the cancel logic location
+	if s.cancelFunc != nil {
+		s.cancelFunc()
+	}
 	select {
 	case s.headerSig <- streamSigInactive:
 	default:
@@ -204,7 +210,6 @@ func (s *stream) close(exception error, rst bool, kind int32) error {
 	default:
 	}
 	s.reader.close(exception)
-	// 是否需要区分 kind？是需要判断的，因为 server stream 不需要这个东西
 	if kind == clientTransport && exception != nil {
 		s.clientStreamException.Store(exception)
 	}
@@ -310,7 +315,6 @@ func (s *stream) sendTrailer(exception error) (err error) {
 	return s.writeFrame(trailerFrameType, nil, wtrailer, payload)
 }
 
-// todo: 处理写入错误，需要感知到连接断开这个现象
 func (s *stream) sendRstFrame(exception error) (err error) {
 	klog.Debugf("stream[%d] send rst frame: err=%v", s.sid, exception)