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

close #52406

pkg/executor/internal/mpp/local_mpp_coordinator.go

@@ -499,9 +499,6 @@ 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)

pkg/store/copr/mpp.go

@@ -255,16 +255,23 @@ func (c *MPPClient) EstablishMPPConns(param kv.EstablishMPPConnsParam) (*tikvrpc
 	// We don't need to process any special error. When we meet errors, just let it fail.
 	rpcResp, err := c.store.GetTiKVClient().SendRequest(param.Ctx, req.Meta.GetAddress(), wrappedReq, TiFlashReadTimeoutUltraLong)
 
+	var stream *tikvrpc.MPPStreamResponse
+	if rpcResp != nil && rpcResp.Resp != nil {
+		stream = rpcResp.Resp.(*tikvrpc.MPPStreamResponse)
+	}
+
 	if err != nil {
+		if stream != nil {
+			stream.Close()
+		}
 		logutil.BgLogger().Warn("establish mpp connection meet error and cannot retry", zap.String("error", err.Error()), zap.Uint64("timestamp", taskMeta.StartTs), zap.Int64("task", taskMeta.TaskId), zap.Int64("mpp-version", taskMeta.MppVersion))
 		if config.GetGlobalConfig().DisaggregatedTiFlash && !config.GetGlobalConfig().UseAutoScaler {
 			c.store.GetRegionCache().InvalidateTiFlashComputeStores()
 		}
 		return nil, err
 	}
 
-	streamResponse := rpcResp.Resp.(*tikvrpc.MPPStreamResponse)
-	return streamResponse, nil
+	return stream, nil
 }
 
 // CheckVisibility checks if it is safe to read using given ts.