From 9ee0a0ae73c109fb862d38e17e34ddf2cca792b8 Mon Sep 17 00:00:00 2001
From: "Hui, Chunyang" <sanqian.hcy@antfin.com>
Date: Fri, 22 Dec 2023 09:46:00 +0000
Subject: [PATCH] Refine replace_mem_chunks to lock once

---
 src/libos/src/vm/process_vm.rs | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/libos/src/vm/process_vm.rs b/src/libos/src/vm/process_vm.rs
index 06d4125a..e637e42a 100644
--- a/src/libos/src/vm/process_vm.rs
+++ b/src/libos/src/vm/process_vm.rs
@@ -355,8 +355,9 @@ impl ProcessVM {
     }
 
     pub fn replace_mem_chunk(&self, old_chunk: &ChunkRef, new_chunk: ChunkRef) {
-        self.remove_mem_chunk(old_chunk);
-        self.add_mem_chunk(new_chunk)
+        let mut mem_chunks = self.mem_chunks.write().unwrap();
+        mem_chunks.remove(old_chunk);
+        mem_chunks.insert(new_chunk);
     }
 
     // Try merging all connecting single VMAs of the process.