[Orc] Fix pending debug object tracking in DebugObjectManagerPlugin

[sgraenitz]

Summary

There can be multiple MaterializationResponsibilitys in-flight for a single ResourceKey. Hence, pending debug objects must be tracked by MaterializationResponsibility and not by ResourceKey.

[sgraenitz]

This is hard to reproduce with lli, because single-threaded execution with LLLazyJIT will implicitly sequence materialization. LLJIT instead materializes module dependencies recursively, which makes this issue easy to reproduce. Unfortunately, we have no LLJIT tool that I could use for a test. It's not the first time this limitation comes up for me and I wonder if it's worth adding an orc-greedy kind to lli. What do you think?

llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp

515

Can we guaranteed that affected in-flight MRs are still alive at this point? Running the code as is works, but I wonder if it's a robust assumption.

523

It's a different issue, but I realized that this might be a dangerous use of resource keys. They are only guaranteed to be valid within the lambda passed into withResourceKeyDo() right?

[lhames]

llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp

515

We can guarantee that no MR is trying to register resources with the key at the point that notifyRemovingResources is called, provided that keys are never leaked from withResourceKeyDo -- see below. :)

523

Yep. I missed this on my first read through.

You should always use MaterializationResponsibility::withResourceKeyDo to create associations between resources and keys, and never leak the key from inside withResourceKeyDo. This system guarantees that the association never ends up in a race with resource removal: Any race becomes ordered as either: ASSOCIATE then REMOVE (normal removal case), or REMOVE then ASSOCIATE (in which case withResourceKey do will return a ResourceTrackerDefunct error and not run your lambda).

[lhames]

This is hard to reproduce with lli, because single-threaded execution with LLLazyJIT will implicitly sequence materialization. LLJIT instead materializes module dependencies recursively, which makes this issue easy to reproduce. Unfortunately, we have no LLJIT tool that I could use for a test. It's not the first time this limitation comes up for me and I wonder if it's worth adding an orc-greedy kind to lli. What do you think?

I'm 100% in favor of an -orc-greedy mode for lli. The only reason it doesn't have one is that I've been short on time.

If it were built I would say that it should just be -jit-kind=orc (with -jit-kind=orc-lazy remaining as it is), and that it should be the new default. In that case all MCJIT test cases would just have to be updated to explicitly specify -jit-kind=mcjit.

[sgraenitz]

Pending resources are irrelevant in both, notifyTransferringResources() and notifyRemovingResources(). They are guaranteed to be either finalizd in notifyEmitted() or discarded in notifyFailed().

[sgraenitz]

Thanks for elaborating.

llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp

515

It turned out that I don't need to worry about pending debug objects here, which renders my question irrelevant.

523

Fixing this in D98863

[lhames]

LGTM.

llvm/include/llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h

public:
   void modifyPassConfig(MaterializationResponsibility &MR,
                         jitlink::LinkGraph &LG,
                         jitlink::PassConfiguration &PassConfig) override;
 
 private:
   ExecutionSession &ES;
 
   using OwnedDebugObject = std::unique_ptr<DebugObject>;
-  std::map<ResourceKey, OwnedDebugObject> PendingObjs;
+  std::map<MaterializationResponsibility *, OwnedDebugObject> PendingObjs;
   std::map<ResourceKey, std::vector<OwnedDebugObject>> RegisteredObjs;
 
   std::mutex PendingObjsLock;
   std::mutex RegisteredObjsLock;
 
   std::unique_ptr<DebugObjectRegistrar> Target;
 };
 
 } // namespace orc
 } // namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_DEBUGOBJECTMANAGERPLUGIN_H

llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp

DebugObjectManagerPlugin::DebugObjectManagerPlugin(
     ExecutionSession &ES, std::unique_ptr<DebugObjectRegistrar> Target)
     : ES(ES), Target(std::move(Target)) {}
 
 DebugObjectManagerPlugin::~DebugObjectManagerPlugin() = default;
 
 void DebugObjectManagerPlugin::notifyMaterializing(
     MaterializationResponsibility &MR, LinkGraph &G, JITLinkContext &Ctx,
     MemoryBufferRef ObjBuffer) {
-  assert(PendingObjs.count(getResourceKey(MR)) == 0 &&
+  std::lock_guard<std::mutex> Lock(PendingObjsLock);
+  assert(PendingObjs.count(&MR) == 0 &&
          "Cannot have more than one pending debug object per "
          "MaterializationResponsibility");
 
-  std::lock_guard<std::mutex> Lock(PendingObjsLock);
   if (auto DebugObj = createDebugObjectFromBuffer(ES, G, Ctx, ObjBuffer)) {
     // Not all link artifacts allow debugging.
-    if (*DebugObj != nullptr) {
-      ResourceKey Key = getResourceKey(MR);
-      PendingObjs[Key] = std::move(*DebugObj);
-    }
+    if (*DebugObj != nullptr)
+      PendingObjs[&MR] = std::move(*DebugObj);
   } else {
     ES.reportError(DebugObj.takeError());
   }
 }
 
 void DebugObjectManagerPlugin::modifyPassConfig(
     MaterializationResponsibility &MR, LinkGraph &G,
     PassConfiguration &PassConfig) {
   // Not all link artifacts have associated debug objects.
   std::lock_guard<std::mutex> Lock(PendingObjsLock);
-  auto It = PendingObjs.find(getResourceKey(MR));
+  auto It = PendingObjs.find(&MR);
   if (It == PendingObjs.end())
     return;
 
   DebugObject &DebugObj = *It->second;
   if (DebugObj.has(Requirement::ReportFinalSectionLoadAddresses)) {
     PassConfig.PostAllocationPasses.push_back(
         [&DebugObj](LinkGraph &Graph) -> Error {
           for (const Section &GraphSection : Graph.sections())
             DebugObj.reportSectionTargetMemoryRange(GraphSection.getName(),
                                                     SectionRange(GraphSection));
           return Error::success();
         });
   }
 }
 
 Error DebugObjectManagerPlugin::notifyEmitted(
     MaterializationResponsibility &MR) {
-  ResourceKey Key = getResourceKey(MR);
-
   std::lock_guard<std::mutex> Lock(PendingObjsLock);
-  auto It = PendingObjs.find(Key);
+  auto It = PendingObjs.find(&MR);
   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
   // info.
   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, Key, UnownedDebugObj,
+        [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);
           std::lock_guard<std::mutex> Lock(RegisteredObjsLock);
           RegisteredObjs[Key].push_back(std::move(ReownedDebugObj));
         });
 
   return FinalizeErr.get();
 }
 
 Error DebugObjectManagerPlugin::notifyFailed(
     MaterializationResponsibility &MR) {
   std::lock_guard<std::mutex> Lock(PendingObjsLock);
-  PendingObjs.erase(getResourceKey(MR));
+  PendingObjs.erase(&MR);
   return Error::success();
 }
 
 void DebugObjectManagerPlugin::notifyTransferringResources(ResourceKey DstKey,
                                                            ResourceKey SrcKey) {
-  {
-    std::lock_guard<std::mutex> Lock(RegisteredObjsLock);
-    auto SrcIt = RegisteredObjs.find(SrcKey);
-    if (SrcIt != RegisteredObjs.end()) {
-      // Resources from distinct MaterializationResponsibilitys can get merged
-      // after emission, so we can have multiple debug objects per resource key.
-      for (std::unique_ptr<DebugObject> &DebugObj : SrcIt->second)
-        RegisteredObjs[DstKey].push_back(std::move(DebugObj));
-      RegisteredObjs.erase(SrcIt);
-    }
-  }
-  {
-    std::lock_guard<std::mutex> Lock(PendingObjsLock);
-    auto SrcIt = PendingObjs.find(SrcKey);
-    if (SrcIt != PendingObjs.end()) {
-      assert(PendingObjs.count(DstKey) == 0 &&
-             "Cannot have more than one pending debug object per "
-             "MaterializationResponsibility");
-      PendingObjs[DstKey] = std::move(SrcIt->second);
-      PendingObjs.erase(SrcIt);
-    }
-  }
-}
+  // Debug objects are stored by ResourceKey only after registration.
+  // Thus, pending objects don't need to be updated here.
 
-Error DebugObjectManagerPlugin::notifyRemovingResources(ResourceKey K) {
-  {
-    std::lock_guard<std::mutex> Lock(RegisteredObjsLock);
-    RegisteredObjs.erase(K);
-    // TODO: Implement unregister notifications.
-  }
-  std::lock_guard<std::mutex> Lock(PendingObjsLock);
-  PendingObjs.erase(K);
+Error DebugObjectManagerPlugin::notifyRemovingResources(ResourceKey Key) {
+  // Removing the resource for a pending object fails materialization, so they
+  // get cleaned up in the notifyFailed() handler.
+  RegisteredObjs.erase(Key);
 
+  // TODO: Implement unregister notifications.
   return Error::success();
 }
 
 } // namespace orc
 } // namespace llvm