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

[Stefan Gränitz via Phabricator via llvm-commits]

sgraenitz created this revision.
sgraenitz added a reviewer: lhames.
Herald added a subscriber: hiraditya.
sgraenitz requested review of this revision.
Herald added a project: LLVM.

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


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D98863

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


Index: llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
===================================================================
--- llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
+++ llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
@@ -371,16 +371,6 @@
   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 @@
   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 @@
   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();
 }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D98863.331541.patch
Type: text/x-patch
Size: 3232 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210318/da8b7449/attachment.bin>