executor: fix mppIterator memory leak when got error (#52410)

close #52406
pingcap · Apr 9, 2024 · 38f665a · 38f665a
1 parent 4a1bf18
commit 38f665a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pkg/executor/internal/mpp/local_mpp_coordinator.go b/pkg/executor/internal/mpp/local_mpp_coordinator.go
@@ -498,6 +498,9 @@ func (c *localMppCoordinator) cancelMppTasks() {
 func (c *localMppCoordinator) receiveResults(req *kv.MPPDispatchRequest, taskMeta *mpp.TaskMeta, bo *backoff.Backoffer) {
 	stream, err := c.sessionCtx.GetMPPClient().EstablishMPPConns(kv.EstablishMPPConnsParam{Ctx: bo.GetCtx(), Req: req, TaskMeta: taskMeta})
 	if err != nil {
+		if stream != nil {
+			stream.Close()
+		}
 		// if NeedTriggerFallback is true, we return timeout to trigger tikv's fallback
 		if c.needTriggerFallback {
 			c.sendError(derr.ErrTiFlashServerTimeout)