diff --git a/src/observer/mysql/obmp_stmt_execute.cpp b/src/observer/mysql/obmp_stmt_execute.cpp
index c5616e3a1c..7d8be26f9d 100644
--- a/src/observer/mysql/obmp_stmt_execute.cpp
+++ b/src/observer/mysql/obmp_stmt_execute.cpp
@@ -1847,24 +1847,23 @@ int ObMPStmtExecute::process()
     }
     session.check_and_reset_retry_info(*cur_trace_id, THIS_WORKER.need_retry());
     session.set_last_trace_id(ObCurTraceId::get_trace_id());
-  }
-
-  // whether the previous error was reported, a cleanup is to be done here
-  if (NULL != sess && !async_resp_used) {
-    // async remove in ObSqlEndTransCb
-    ObPieceCache *piece_cache = static_cast<ObPieceCache*>(sess->get_piece_cache());
-    if (OB_ISNULL(piece_cache)) {
-      // do nothing
-      // piece_cache not be null in piece data protocol
-    } else {
-      for (uint64_t i = 0; OB_SUCC(ret) && i < params_num_; i++) {
-        if (OB_FAIL(piece_cache->remove_piece(
-                            piece_cache->get_piece_key(stmt_id_, i),
-                            *sess))) {
-          if (OB_HASH_NOT_EXIST == ret) {
-            ret = OB_SUCCESS;
-          } else {
-            LOG_WARN("remove piece fail", K(stmt_id_), K(i), K(ret));
+    // whether the previous error was reported, a cleanup is to be done here
+    if (!async_resp_used) {
+      // async remove in ObSqlEndTransCb
+      ObPieceCache *piece_cache = static_cast<ObPieceCache*>(session.get_piece_cache());
+      if (OB_ISNULL(piece_cache)) {
+        // do nothing
+        // piece_cache not be null in piece data protocol
+      } else {
+        for (uint64_t i = 0; OB_SUCC(ret) && i < params_num_; i++) {
+          if (OB_FAIL(piece_cache->remove_piece(
+                              piece_cache->get_piece_key(stmt_id_, i),
+                              session))) {
+            if (OB_HASH_NOT_EXIST == ret) {
+              ret = OB_SUCCESS;
+            } else {
+              LOG_WARN("remove piece fail", K(stmt_id_), K(i), K(ret));
+            }
           }
         }
       }
diff --git a/src/observer/mysql/obmp_stmt_send_piece_data.cpp b/src/observer/mysql/obmp_stmt_send_piece_data.cpp
index 65d96a2c59..524cd60bf3 100644
--- a/src/observer/mysql/obmp_stmt_send_piece_data.cpp
+++ b/src/observer/mysql/obmp_stmt_send_piece_data.cpp
@@ -488,17 +488,20 @@ int ObPieceCache::close_all(ObSQLSessionInfo &session)
 {
   int ret = OB_SUCCESS;
   if (is_inited()) {
+    common::ObSEArray<int64_t, 32> piece_keys;
     for (PieceMap::iterator iter = piece_map_.begin();  //ignore ret
         iter != piece_map_.end();
         ++iter) {
       ObPiece *piece = iter->second;
-      int64_t key = get_piece_key(piece->get_stmt_id(), piece->get_param_id());
+      piece_keys.push_back(get_piece_key(piece->get_stmt_id(), piece->get_param_id()));
+    }
+    for (int64_t i = 0; i < piece_keys.count(); i++) {
+      int64_t key = piece_keys.at(i);
       int64_t tmp_ret = remove_piece(key, session);
       // only save first error ret
       ret = ret == OB_SUCCESS ? tmp_ret : ret;
       if (OB_SUCCESS != tmp_ret) {
-        LOG_WARN("remove piece fail.", K(piece->get_stmt_id()), 
-                      K(piece->get_param_id()), K(tmp_ret));
+        LOG_WARN("remove piece fail.", K(key), K(tmp_ret));
       }
     }
   }