From e309f5ebc0d7b9e5b057030643d9042acd5edaec Mon Sep 17 00:00:00 2001
From: Anubhab Ghosh <anubhabghosh.me@gmail.com>
Date: Sun, 21 Aug 2022 01:49:34 +0530
Subject: [PATCH] [Orc] Improve deintialize and shutdown logic

When deinitializing, the allocation needs to be removed from the
allocation list of its associated reservation so that remaining
allocations can be deinitialized when releasing the reservation.

Differential Revision: https://reviews.llvm.org/D132313
---
 .../ExecutorSharedMemoryMapperService.cpp          | 32 ++++++++++++----------
 1 file changed, 18 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.cpp b/llvm/lib/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.cpp
index fd487e3..23bf5a3 100644
--- a/llvm/lib/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.cpp
@@ -198,6 +198,17 @@ Error ExecutorSharedMemoryMapperService::deinitialize(
         AllErr = joinErrors(std::move(AllErr), std::move(Err));
       }
 
+      // Remove the allocation from the allocation list of its reservation
+      for (auto &Reservation : Reservations) {
+        auto AllocationIt =
+            std::find(Reservation.second.Allocations.begin(),
+                      Reservation.second.Allocations.end(), Base);
+        if (AllocationIt != Reservation.second.Allocations.end()) {
+          Reservation.second.Allocations.erase(AllocationIt);
+          break;
+        }
+      }
+
       Allocations.erase(Base);
     }
   }
@@ -264,22 +275,15 @@ Error ExecutorSharedMemoryMapperService::release(
 }
 
 Error ExecutorSharedMemoryMapperService::shutdown() {
-  std::vector<ExecutorAddr> ReservationAddrs;
-  {
-    std::lock_guard<std::mutex> Lock(Mutex);
+  if (Reservations.empty())
+    return Error::success();
 
-    if (Reservations.empty())
-      return Error::success();
-
-    ReservationAddrs.reserve(Reservations.size());
-    for (const auto &R : Reservations) {
-      ReservationAddrs.push_back(ExecutorAddr::fromPtr(R.getFirst()));
-    }
-
-    Reservations.clear();
-  }
+  std::vector<ExecutorAddr> ReservationAddrs;
+  ReservationAddrs.reserve(Reservations.size());
+  for (const auto &R : Reservations)
+    ReservationAddrs.push_back(ExecutorAddr::fromPtr(R.getFirst()));
 
-  return release(ReservationAddrs);
+  return release(std::move(ReservationAddrs));
 }
 
 void ExecutorSharedMemoryMapperService::addBootstrapSymbols(
-- 
2.7.4