perf: optimize server-side handling of timeout control

TarsCloud · Jan 8, 2024 · a7f73e0 · a7f73e0
1 parent 9319267
commit a7f73e0
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tars/transport/tarsserver.go b/tars/transport/tarsserver.go
@@ -112,12 +112,12 @@ func (ts *TarsServer) invoke(ctx context.Context, pkg []byte) []byte {
 	if cfg.HandleTimeout == 0 {
 		rsp = ts.protocol.Invoke(ctx, pkg)
 	} else {
-		invokeDone, cancelFunc := context.WithTimeout(context.Background(), cfg.HandleTimeout)
+		invokeCtx, cancelFunc := context.WithTimeout(ctx, cfg.HandleTimeout)
 		go func() {
-			rsp = ts.protocol.Invoke(ctx, pkg)
+			rsp = ts.protocol.Invoke(invokeCtx, pkg)
 			cancelFunc()
 		}()
-		<-invokeDone.Done()
+		<-invokeCtx.Done()
 		if len(rsp) == 0 { // The rsp must be none-empty
 			rsp = ts.protocol.InvokeTimeout(pkg)
 		}