fix: ttstream cancelFunc panic

Author: DMwangnima

pkg/remote/trans/ttstream/exception.go

@@ -34,10 +34,12 @@ var (
 
 	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)
+	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)
 )
 
 const (
@@ -103,8 +105,10 @@ func (e *Exception) Error() string {
 }
 
 func (e *Exception) withCause(cause error) *Exception {
-	e.cause = cause
-	e.bitSet |= setCause
+	if cause != nil {
+		e.cause = cause
+		e.bitSet |= setCause
+	}
 	return e
 }
 

pkg/remote/trans/ttstream/exception_test.go

@@ -39,6 +39,10 @@ func TestErrors(t *testing.T) {
 	test.Assert(t, errors.Is(appErr, errApplicationException), appErr)
 	test.Assert(t, !errors.Is(appErr, kerrors.ErrStreamingProtocol), appErr)
 	test.Assert(t, strings.Contains(appErr.Error(), causeErr.Error()))
+
+	newExWithNilErr := errIllegalFrame.newBuilder().withCause(nil)
+	test.Assert(t, !newExWithNilErr.isCauseSet(), newExWithNilErr)
+	test.Assert(t, newExWithNilErr.cause == nil, newExWithNilErr)
 }
 
 func TestCommonParentKerror(t *testing.T) {

pkg/remote/trans/ttstream/stream_server.go

@@ -50,7 +50,7 @@ func newServerStream(ctx context.Context, writer streamWriter, smeta streamFrame
 type serverStream struct {
 	*stream
 	state      int32
-	cancelFunc cancelWithReason
+	cancelFunc cancelWithReason // use serverStream.cancel to prevent atomic.Value panic caused by passing different error types
 }
 
 func (s *serverStream) SetHeader(hd streaming.Header) error {
@@ -121,7 +121,10 @@ func (s *serverStream) SendMsg(ctx context.Context, res any) error {
 // CloseSend by serverStream will be called after server handler returned
 // after CloseSend stream cannot be access again
 func (s *serverStream) CloseSend(exception error) error {
-	return s.close(exception, true)
+	if s.close(errBizHandlerReturnCancel) {
+		return nil
+	}
+	return s.sendTrailer(exception)
 }
 
 // closeRecv called only when server receiving Trailer Frame
@@ -132,22 +135,24 @@ func (s *serverStream) closeRecv(exception error) error {
 	return nil
 }
 
-func (s *serverStream) close(exception error, sendTrailer bool) error {
+// cancel ensures the ex passed in cancelFunc is of type *Exception
+func (s *serverStream) cancel(ex *Exception) {
+	s.cancelFunc(ex)
+}
+
+func (s *serverStream) close(exception *Exception) (isClosed bool) {
 	// support cascading cancel
 	// we must cancel the ctx first before changing the state
 	// otherwise Recv/Send will not get the expected exception
-	s.cancelFunc(exception)
+	s.cancel(exception)
 	if atomic.SwapInt32(&s.state, streamStateInactive) == streamStateInactive {
-		return nil
+		return true
 	}
 
 	s.reader.close(exception)
 	s.runCloseCallback(exception)
 
-	if sendTrailer {
-		return s.sendTrailer(exception)
-	}
-	return nil
+	return false
 }
 
 // === serverStream onRead callback
@@ -210,15 +215,16 @@ func (s *serverStream) onReadRstFrame(fr *Frame) (err error) {
 	}
 
 	// when receiving rst frame, we should close stream and there is no need to send rst frame
-	return s.close(rstEx, false)
+	s.close(rstEx)
+	return nil
 }
 
 // closeTest is only used in unit tests for mocking Proxy Egress/Ingress send Rst Frame to downstream and upstream
 func (s *serverStream) closeTest(exception error, cancelPath string) error {
 	// support cascading cancel
 	// we must cancel the ctx first before changing the state
 	// otherwise Recv/Send will not get the expected exception
-	s.cancelFunc(exception)
+	s.cancel(errInternalCancel.newBuilder().withCause(exception))
 	if atomic.SwapInt32(&s.state, streamStateInactive) == streamStateInactive {
 		return nil
 	}

pkg/remote/trans/ttstream/stream_server_test.go

@@ -39,22 +39,22 @@ func newTestServerStream() *serverStream {
 func Test_serverStreamStateChange(t *testing.T) {
 	t.Run("serverStream close, then RecvMsg/SendMsg returning exception", func(t *testing.T) {
 		srvSt := newTestServerStream()
-		testException := errors.New("test")
-		err := srvSt.close(testException, false)
-		test.Assert(t, err == nil, err)
+		testException := errInternalCancel.newBuilder().withCause(errors.New("test"))
+		isClosed := srvSt.close(testException)
+		test.Assert(t, !isClosed)
 		rErr := srvSt.RecvMsg(srvSt.ctx, nil)
 		test.Assert(t, rErr == testException, rErr)
 		sErr := srvSt.SendMsg(srvSt.ctx, nil)
 		test.Assert(t, sErr == testException, sErr)
 	})
 	t.Run("serverStream close twice with different exception, RecvMsg/SendMsg returning the first time exception", func(t *testing.T) {
 		srvSt := newTestServerStream()
-		testException1 := errors.New("test1")
-		err := srvSt.close(testException1, false)
-		test.Assert(t, err == nil, err)
-		testException2 := errors.New("test2")
-		err = srvSt.close(testException2, false)
-		test.Assert(t, err == nil, err)
+		testException1 := errInternalCancel.newBuilder().withCause(errors.New("test1"))
+		isClosed := srvSt.close(testException1)
+		test.Assert(t, !isClosed)
+		testException2 := errInternalCancel.newBuilder().withCause(errors.New("test2"))
+		isClosed = srvSt.close(testException2)
+		test.Assert(t, isClosed)
 
 		rErr := srvSt.RecvMsg(srvSt.ctx, nil)
 		test.Assert(t, rErr == testException1, rErr)
@@ -65,11 +65,11 @@ func Test_serverStreamStateChange(t *testing.T) {
 		srvSt := newTestServerStream()
 		var wg sync.WaitGroup
 		wg.Add(2)
-		testException := errors.New("test")
+		testException := errInternalCancel.newBuilder().withCause(errors.New("test1"))
 		go func() {
 			time.Sleep(100 * time.Millisecond)
-			err := srvSt.close(testException, false)
-			test.Assert(t, err == nil, err)
+			isClosed := srvSt.close(testException)
+			test.Assert(t, !isClosed)
 		}()
 		go func() {
 			defer wg.Done()

pkg/remote/trans/ttstream/transport_client.go

@@ -172,7 +172,7 @@ func (t *clientTransport) readFrame(reader bufiox.Reader) error {
 	var ok bool
 	s, ok = t.loadStream(fr.sid)
 	if !ok {
-		klog.Errorf("transport[%s] read a unknown stream: frame[%s]", t.Addr(), fr)
+		klog.Debugf("transport[%s] read a unknown stream: frame[%s]", t.Addr(), fr)
 		// ignore unknown stream error
 		err = nil
 	} else {

pkg/remote/trans/ttstream/transport_server.go

@@ -103,9 +103,10 @@ func (t *serverTransport) Close(exception error) (err error) {
 	}
 	klog.Debugf("server transport[%s] is closing", t.Addr())
 	// close streams first
+	ex := errConnectionClosedCancel.newBuilder().withCause(exception)
 	t.streams.Range(func(key, value any) bool {
 		s := value.(*serverStream)
-		_ = s.close(exception, false)
+		_ = s.close(ex)
 		return true
 	})
 	// then close stream and frame pipes
@@ -163,7 +164,9 @@ func (t *serverTransport) readFrame(reader bufiox.Reader) error {
 		var ok bool
 		s, ok = t.loadStream(fr.sid)
 		if !ok {
-			klog.Errorf("transport[%s] read a unknown stream: frame[%s]", t.Addr(), fr.String())
+			// 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 {

pkg/remote/trans/ttstream/transport_test.go

@@ -257,6 +257,74 @@ func TestTransportException(t *testing.T) {
 	test.Assert(t, errors.Is(err, errIllegalFrame), err)
 }
 
+func TestTransportClose(t *testing.T) {
+	cfd, sfd := netpoll.GetSysFdPairs()
+	cconn, err := netpoll.NewFDConnection(cfd)
+	test.Assert(t, err == nil, err)
+	sconn, err := netpoll.NewFDConnection(sfd)
+	test.Assert(t, err == nil, err)
+
+	intHeader := make(IntHeader)
+	intHeader[0] = "test"
+	strHeader := make(streaming.Header)
+	strHeader["key"] = "val"
+	ctrans := newClientTransport(cconn, nil)
+	defer ctrans.Close(nil)
+	ctx := context.Background()
+	cs := newClientStream(ctx, ctrans, streamFrame{sid: genStreamID(), method: "Bidi"})
+	err = ctrans.WriteStream(ctx, cs, intHeader, strHeader)
+	test.Assert(t, err == nil, err)
+	strans := newServerTransport(sconn)
+	ss, err := strans.ReadStream(context.Background())
+	test.Assert(t, err == nil, err)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	// client
+	go func() {
+		defer wg.Done()
+		for {
+			req := new(testRequest)
+			req.B = "hello"
+			sErr := cs.SendMsg(context.Background(), req)
+			if sErr != nil {
+				test.Assert(t, errors.Is(sErr, errIllegalFrame), sErr)
+				return
+			}
+
+			res := new(testResponse)
+			rErr := cs.RecvMsg(context.Background(), res)
+			if rErr != nil {
+				test.Assert(t, errors.Is(rErr, errIllegalFrame), rErr)
+				return
+			}
+			test.DeepEqual(t, req.B, res.B)
+		}
+	}()
+
+	go func() {
+		time.Sleep(100 * time.Millisecond)
+		strans.Close(nil)
+	}()
+	// server
+	for {
+		req := new(testRequest)
+		rErr := ss.RecvMsg(context.Background(), req)
+		if rErr != nil {
+			test.Assert(t, errors.Is(rErr, errConnectionClosedCancel), rErr)
+			break
+		}
+		res := new(testResponse)
+		res.B = req.B
+		sErr := ss.SendMsg(context.Background(), res)
+		if sErr != nil {
+			test.Assert(t, errors.Is(sErr, errConnectionClosedCancel), sErr)
+			break
+		}
+	}
+	wg.Wait()
+}
+
 func TestStreamID(t *testing.T) {
 	oriId := atomic.LoadInt32(&clientStreamID)
 	atomic.StoreInt32(&clientStreamID, math.MaxInt32-1)