From 1d756429465347ed98b84b0cd868a9125a74111e Mon Sep 17 00:00:00 2001
From: Yutong Zhao <zhaoyt0208@gmail.com>
Date: Tue, 3 Feb 2026 19:59:55 +0800
Subject: [PATCH] Fix the double-free corruption under high-concurrency
 multi-threaded scenarios.

---
 .../astra-sim/system/collective/NcclTreeFlowModel.cc     | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/astra-sim-alibabacloud/astra-sim/system/collective/NcclTreeFlowModel.cc b/astra-sim-alibabacloud/astra-sim/system/collective/NcclTreeFlowModel.cc
index 4c84a873..ad5092f7 100644
--- a/astra-sim-alibabacloud/astra-sim/system/collective/NcclTreeFlowModel.cc
+++ b/astra-sim-alibabacloud/astra-sim/system/collective/NcclTreeFlowModel.cc
@@ -433,9 +433,14 @@ bool NcclTreeFlowModel::iteratable(int channel_id) {
       break;
     }
   }
+  bool should_exit = false;
+  if (all_channel_finished && all_packets_freed) {
+    if (!judge_exit_flag.exchange(true)) {
+      should_exit = true;
+    }
+  }
   cs.ExitSection();
-  if (all_channel_finished == true &&
-      all_packets_freed == true) {
+  if (should_exit) {
     exit();
     return false;
   }