[llvm] cbcc1c9 - [Orc] Make usage of ResourceKeys thread-safe in DebugObjectManagerPlugin

[Stefan Gränitz via llvm-commits]

Author: Stefan Gränitz
Date: 2021-03-22T17:47:33+01:00
New Revision: cbcc1c9f87080779bf138259a9177e8f2fe674c7

URL: https://github.com/llvm/llvm-project/commit/cbcc1c9f87080779bf138259a9177e8f2fe674c7
DIFF: https://github.com/llvm/llvm-project/commit/cbcc1c9f87080779bf138259a9177e8f2fe674c7.diff

LOG: [Orc] Make usage of ResourceKeys thread-safe in DebugObjectManagerPlugin

Don't leak ResourceKeys from MaterializationResponsibility::withResourceKeyDo() in notifyEmitted().
Also make some improvements in the overall implementation.

Differential Revision: https://reviews.llvm.org/D98863

Added: 
    

Modified: 
    llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
index 957de3586677..36efc744bf30 100644
--- a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
@@ -371,16 +371,6 @@ DebugObjectSection *ELFDebugObject::getSection(StringRef Name) {
   return It == Sections.end() ? nullptr : It->second.get();
 }
 
-static ResourceKey getResourceKey(MaterializationResponsibility &MR) {
-  ResourceKey Key;
-  if (auto Err = MR.withResourceKeyDo([&](ResourceKey K) { Key = K; })) {
-    MR.getExecutionSession().reportError(std::move(Err));
-    return ResourceKey{};
-  }
-  assert(Key && "Invalid key");
-  return Key;
-}
-
 /// Creates a debug object based on the input object file from
 /// ObjectLinkingLayerJITLinkContext.
 ///
@@ -449,9 +439,6 @@ Error DebugObjectManagerPlugin::notifyEmitted(
   if (It == PendingObjs.end())
     return Error::success();
 
-  DebugObject *UnownedDebugObj = It->second.release();
-  PendingObjs.erase(It);
-
   // During finalization the debug object is registered with the target.
   // Materialization must wait for this process to finish. Otherwise we might
   // start running code before the debugger processed the corresponding debug
@@ -459,30 +446,27 @@ Error DebugObjectManagerPlugin::notifyEmitted(
   std::promise<MSVCPError> FinalizePromise;
   std::future<MSVCPError> FinalizeErr = FinalizePromise.get_future();
 
-  // FIXME: We released ownership of the DebugObject, so we can easily capture
-  // the raw pointer in the continuation function, which re-owns it immediately.
-  if (UnownedDebugObj)
-    UnownedDebugObj->finalizeAsync(
-        [this, UnownedDebugObj, &MR,
-         &FinalizePromise](Expected<sys::MemoryBlock> TargetMem) {
-          std::unique_ptr<DebugObject> ReownedDebugObj(UnownedDebugObj);
-          if (!TargetMem) {
-            FinalizePromise.set_value(TargetMem.takeError());
-            return;
-          }
-          if (Error Err = Target->registerDebugObject(*TargetMem)) {
-            FinalizePromise.set_value(std::move(Err));
-            return;
-          }
-
-          // Registration successful, notifyEmitted() can return now and
-          // materialization can finish.
-          FinalizePromise.set_value(Error::success());
-
-          ResourceKey Key = getResourceKey(MR);
+  It->second->finalizeAsync(
+      [this, &FinalizePromise, &MR](Expected<sys::MemoryBlock> TargetMem) {
+        // Any failure here will fail materialization.
+        if (!TargetMem) {
+          FinalizePromise.set_value(TargetMem.takeError());
+          return;
+        }
+        if (Error Err = Target->registerDebugObject(*TargetMem)) {
+          FinalizePromise.set_value(std::move(Err));
+          return;
+        }
+
+        // Once our tracking info is updated, notifyEmitted() can return and
+        // finish materialization.
+        FinalizePromise.set_value(MR.withResourceKeyDo([&](ResourceKey K) {
+          assert(PendingObjs.count(&MR) && "We still hold PendingObjsLock");
           std::lock_guard<std::mutex> Lock(RegisteredObjsLock);
-          RegisteredObjs[Key].push_back(std::move(ReownedDebugObj));
-        });
+          RegisteredObjs[K].push_back(std::move(PendingObjs[&MR]));
+          PendingObjs.erase(&MR);
+        }));
+      });
 
   return FinalizeErr.get();
 }