From 950d09f451126a335bc6a18aa31d628d337f85c4 Mon Sep 17 00:00:00 2001
From: Saturday <minhchaurisza@gmail.com>
Date: Thu, 18 Jun 2026 21:17:18 +0800
Subject: [PATCH] =?UTF-8?q?fix(urma):=20=E4=B8=BA=E5=AF=BC=E5=85=A5?=
 =?UTF-8?q?=E6=AE=B5=E6=93=8D=E4=BD=9C=E6=B7=BB=E5=8A=A0=E4=BA=92=E6=96=A5?=
 =?UTF-8?q?=E9=94=81=E4=BF=9D=E6=8A=A4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

添加 shared_mutex 保护 imported_seg_list_ 和 remote_seg_list_ 的并发访问。
在 retrieveRemoteSeg 中修复双重检查锁定模式，并添加内存分配失败检查。
---
 .../kunpeng_transport/urma/urma_endpoint.h    |  3 ++
 .../kunpeng_transport/urma/urma_endpoint.cpp  | 44 +++++++++++++------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h
index 64e4266976..a4e8a024f3 100644
--- a/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h
+++ b/mooncake-transfer-engine/include/transport/kunpeng_transport/urma/urma_endpoint.h
@@ -16,6 +16,8 @@
 #define URMA_ENDPOINT_H
 #include <atomic>
 #include <memory>
+#include <mutex>
+#include <shared_mutex>
 #include <string>
 #include <thread>
 #include <utility>
@@ -131,6 +133,7 @@ class UrmaContext : public UbContext {
     std::vector<urma_target_seg_t*> local_tseg_list_;
     std::vector<urma_seg_t*> remote_seg_list_;
     std::vector<urma_target_seg_t*> imported_seg_list_;
+    std::shared_mutex import_tseg_mutex_;
 
     std::vector<UrmaJFR> jfr_list_;
 
diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp
index 3d5320c5b9..e14e29b1c8 100644
--- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp
+++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp
@@ -192,20 +192,23 @@ int UrmaContext::deconstruct() {
     }
     seg_region_list_.clear();
 
-    for (auto& seg : imported_seg_list_) {
-        int ret = urma_unimport_seg(seg);
-        if (ret) {
-            PLOG(ERROR) << "Failed to unimport segment";
+    {
+        std::unique_lock<std::shared_mutex> lock(import_tseg_mutex_);
+        for (auto& seg : imported_seg_list_) {
+            int ret = urma_unimport_seg(seg);
+            if (ret) {
+                PLOG(ERROR) << "Failed to unimport segment";
+            }
         }
-    }
-    imported_seg_list_.clear();
+        imported_seg_list_.clear();
 
-    for (auto& seg : remote_seg_list_) {
-        free(seg);
-    }
-    remote_seg_list_.clear();
+        for (auto& seg : remote_seg_list_) {
+            free(seg);
+        }
+        remote_seg_list_.clear();
 
-    import_tseg_map.clear();
+        import_tseg_map.clear();
+    }
 
     for (size_t i = 0; i < jfr_list_.size(); i++) {
         if (!jfr_list_[i].native) continue;
@@ -373,14 +376,25 @@ int UrmaContext::doProcessContextEvents() {
 }
 
 void* UrmaContext::retrieveRemoteSeg(const std::string& remoteSegmentStr) {
+    {
+        std::shared_lock<std::shared_mutex> lock(import_tseg_mutex_);
+        auto ret = import_tseg_map.find(remoteSegmentStr);
+        if (ret != import_tseg_map.end()) return ret->second;
+    }
+
+    std::unique_lock<std::shared_mutex> lock(import_tseg_mutex_);
     auto ret = import_tseg_map.find(remoteSegmentStr);
     if (ret != import_tseg_map.end()) return ret->second;
+
     std::vector<unsigned char> output_buffer;
     deserializeBinaryData(remoteSegmentStr, output_buffer);
-    urma_seg_t* handle;
-    handle = (urma_seg_t*)malloc(sizeof(urma_seg_t));
+    auto* handle = static_cast<urma_seg_t*>(malloc(sizeof(urma_seg_t)));
+    if (!handle) {
+        LOG(ERROR) << "Allocate remote segment handle failed";
+        return nullptr;
+    }
     memcpy(handle, output_buffer.data(), sizeof(urma_seg_t));
-    remote_seg_list_.push_back(handle);
+
     auto import_tseg =
         urma_import_seg(urma_context_, handle, &urma_token, 0, import_flag_);
     if (import_tseg == NULL) {
@@ -388,6 +402,8 @@ void* UrmaContext::retrieveRemoteSeg(const std::string& remoteSegmentStr) {
         free(handle);
         return nullptr;
     }
+
+    remote_seg_list_.push_back(handle);
     imported_seg_list_.push_back(import_tseg);
     import_tseg_map[remoteSegmentStr] = import_tseg;
     return import_tseg;