fix: rpcinfo pool disable also for the server

seabyte7 · Dec 12, 2023 · 00b0718 · 00b0718
1 parent 4db8525
commit 00b0718
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 8 deletions.
diff --git a/pkg/remote/trans/nphttp2/server_handler.go b/pkg/remote/trans/nphttp2/server_handler.go
@@ -125,9 +125,11 @@ func (t *svrTransHandler) OnRead(ctx context.Context, conn net.Conn) error {
 			ri := svrTrans.pool.Get().(rpcinfo.RPCInfo)
 			rCtx := rpcinfo.NewCtxWithRPCInfo(s.Context(), ri)
 			defer func() {
-				// reset rpcinfo
+				// reset rpcinfo for performance (PR #584)
 				ri = t.opt.InitOrResetRPCInfoFunc(ri, conn.RemoteAddr())
-				svrTrans.pool.Put(ri)
+				if rpcinfo.PoolEnabled() {
+					svrTrans.pool.Put(ri)
+				}
 			}()
 
 			// set grpc transport flag before execute metahandler

diff --git a/pkg/rpcinfo/rpcinfo.go b/pkg/rpcinfo/rpcinfo.go
@@ -51,7 +51,8 @@ func EnablePool(enable bool) {
 	}
 }
 
-func poolEnabled() bool {
+// PoolEnabled returns true if rpcInfoPool is enabled.
+func PoolEnabled() bool {
 	return atomic.LoadInt32(&enablePool) == 1
 }
 
@@ -88,7 +89,7 @@ func (r *rpcInfo) zero() {
 
 // Recycle reuses the rpcInfo.
 func (r *rpcInfo) Recycle() {
-	if !poolEnabled() {
+	if !PoolEnabled() {
 		return
 	}
 	if v, ok := r.from.(internal.Reusable); ok {

diff --git a/pkg/rpcinfo/rpcinfo_test.go b/pkg/rpcinfo/rpcinfo_test.go
@@ -25,12 +25,12 @@ import (
 func TestEnablePool(t *testing.T) {
 	t.Run("disable", func(t *testing.T) {
 		EnablePool(false)
-		test.Assert(t, !poolEnabled())
+		test.Assert(t, !PoolEnabled())
 	})
 
 	t.Run("disable-enable", func(t *testing.T) {
 		EnablePool(false)
 		EnablePool(true)
-		test.Assert(t, poolEnabled())
+		test.Assert(t, PoolEnabled())
 	})
 }
diff --git a/server/server.go b/server/server.go
@@ -138,8 +138,8 @@ func newErrorHandleMW(errHandle func(context.Context, error) error) endpoint.Mid
 
 func (s *server) initOrResetRPCInfoFunc() func(rpcinfo.RPCInfo, net.Addr) rpcinfo.RPCInfo {
 	return func(ri rpcinfo.RPCInfo, rAddr net.Addr) rpcinfo.RPCInfo {
-		// Reset rpcinfo if it exists in ctx.
-		if ri != nil {
+		// Reset existing rpcinfo to improve performance for long connections (PR #584).
+		if ri != nil && rpcinfo.PoolEnabled() {
 			fi := rpcinfo.AsMutableEndpointInfo(ri.From())
 			fi.Reset()
 			fi.SetAddress(rAddr)
@@ -155,6 +155,8 @@ func (s *server) initOrResetRPCInfoFunc() func(rpcinfo.RPCInfo, net.Addr) rpcinf
 			}
 			return ri
 		}
+
+		// allocate a new rpcinfo if it's the connection's first request or rpcInfoPool is disabled
 		rpcStats := rpcinfo.AsMutableRPCStats(rpcinfo.NewRPCStats())
 		if s.opt.StatsLevel != nil {
 			rpcStats.SetLevel(*s.opt.StatsLevel)

diff --git a/server/server_test.go b/server/server_test.go
@@ -194,7 +194,10 @@ func TestInitOrResetRPCInfo(t *testing.T) {
 	test.Assert(t, ri.Stats().Level() == stats.LevelDetailed)
 
 	// test reset
+	pOld := reflect.ValueOf(ri).Pointer()
 	ri = rpcInfoInitFunc(ri, conn.RemoteAddr())
+	pNew := reflect.ValueOf(ri).Pointer()
+	test.Assert(t, pOld == pNew, pOld, pNew)
 	test.Assert(t, ri.From().Address().String() == remoteAddr.String())
 	test.Assert(t, ri.Config().ReadWriteTimeout() == rwTimeout)
 
@@ -226,6 +229,17 @@ func TestInitOrResetRPCInfo(t *testing.T) {
 	test.Assert(t, panicked == nil)
 	test.Assert(t, ri.Stats().Error() == nil)
 	test.Assert(t, ri.Stats().Level() == 0)
+
+	// test reset after rpcInfoPool is disabled
+	t.Run("reset with pool disabled", func(t *testing.T) {
+		backupState := rpcinfo.PoolEnabled()
+		defer rpcinfo.EnablePool(backupState)
+		rpcinfo.EnablePool(false)
+
+		riNew := rpcInfoInitFunc(ri, conn.RemoteAddr())
+		pOld, pNew := reflect.ValueOf(ri).Pointer(), reflect.ValueOf(riNew).Pointer()
+		test.Assert(t, pOld != pNew, pOld, pNew)
+	})
 }
 
 func TestServiceRegisterFailed(t *testing.T) {