fix: gonet transServer processes gRPC transHandler

ppzqh · Jan 31, 2025 · 49b8cdf · 49b8cdf
1 parent 047444c
commit 49b8cdf
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 0 deletions.
diff --git a/pkg/remote/trans/common.go b/pkg/remote/trans/common.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync/atomic"
 	"time"
 
 	"github.com/cloudwego/kitex/pkg/remote"
@@ -29,6 +30,29 @@ import (
 
 var readMoreTimeout = 5 * time.Millisecond
 
+type (
+	CtxKeyOnRead   struct{}
+	CtxValueOnRead struct {
+		onlyOnce int32
+	}
+)
+
+func (v *CtxValueOnRead) SetOnlyOnce(b bool) {
+	if b {
+		atomic.StoreInt32(&v.onlyOnce, 1)
+	} else {
+		atomic.StoreInt32(&v.onlyOnce, 0)
+	}
+}
+
+func (v *CtxValueOnRead) GetOnlyOnce() bool {
+	if atomic.LoadInt32(&v.onlyOnce) == 1 {
+		return true
+	} else {
+		return false
+	}
+}
+
 // Extension is the interface that trans extensions need to implement, it will make the extension of trans more easily.
 // Normally if we want to extend transport layer we need to implement the trans interfaces which are defined in trans_handler.go.
 // In fact most code logic is similar in same mode, so the Extension interface is the the differentiated part that need to

diff --git a/pkg/remote/trans/gonet/trans_server.go b/pkg/remote/trans/gonet/trans_server.go
@@ -96,9 +96,18 @@ func (ts *transServer) BootstrapServer(ln net.Listener) (err error) {
 				klog.CtxErrorf(ctx, "KITEX: OnActive error=%s", err)
 				return
 			}
+			ctxValueOnRead := &trans.CtxValueOnRead{}
+			ctx = context.WithValue(ctx, trans.CtxKeyOnRead{}, ctxValueOnRead)
+			onReadOnlyOnceCheck := false
 			for {
 				ts.refreshDeadline(rpcinfo.GetRPCInfo(ctx), bc)
 				err := ts.transHdlr.OnRead(ctx, bc)
+				if !onReadOnlyOnceCheck {
+					onReadOnlyOnceCheck = true
+					if ctxValueOnRead.GetOnlyOnce() {
+						break
+					}
+				}
 				if err != nil {
 					ts.onError(ctx, err, bc)
 					_ = bc.Close()

diff --git a/pkg/remote/trans/nphttp2/server_handler.go b/pkg/remote/trans/nphttp2/server_handler.go
@@ -29,6 +29,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/cloudwego/kitex/pkg/remote/trans"
+
 	"github.com/cloudwego/netpoll"
 
 	"github.com/cloudwego/kitex/pkg/endpoint"
@@ -128,6 +130,10 @@ func (t *svrTransHandler) Read(ctx context.Context, conn net.Conn, msg remote.Me
 
 // 只 return write err
 func (t *svrTransHandler) OnRead(ctx context.Context, conn net.Conn) error {
+	if v := ctx.Value(trans.CtxKeyOnRead{}); v != nil {
+		value := v.(*trans.CtxValueOnRead)
+		value.SetOnlyOnce(true)
+	}
 	svrTrans := ctx.Value(ctxKeySvrTransport).(*SvrTrans)
 	tr := svrTrans.tr
 	tr.HandleStreams(func(s *grpcTransport.Stream) {