[llvm] [ORC] Fix transfer to unknown ResourceTrackers (PR #114063)

[Jonas Hahnfeld via llvm-commits]

https://github.com/hahnjo created https://github.com/llvm/llvm-project/pull/114063

When transferring resources, the destination tracker key may not be in the internal map, invalidating iterators and value references. The added test creates such situation and would fail before with "Finalized allocation was not deallocated."

For good measure, fix the same pattern in `RTDyldObjectLinkingLayer` which is harder to test because it "only" results in memory managers being deleted in the wrong order.

FYI @devajithvs @vgvassilev 

>From 8acfd4e5cb30f08786172f24e98e3519d4f8ed3e Mon Sep 17 00:00:00 2001
From: Jonas Hahnfeld <jonas.hahnfeld at cern.ch>
Date: Tue, 29 Oct 2024 14:23:12 +0100
Subject: [PATCH] [ORC] Fix transfer to unknown ResourceTrackers

When transferring resources, the destination tracker key may not be
in the internal map, invalidating iterators and value references.
The added test creates such situation and would fail before with
"Finalized allocation was not deallocated."

For good measure, fix the same pattern in RTDyldObjectLinkingLayer
which is harder to test because it "only" results in memory managers
being deleted in the wrong order.
---
 .../Orc/ObjectLinkingLayer.cpp                |  9 +++---
 .../Orc/RTDyldObjectLinkingLayer.cpp          |  9 +++---
 .../Orc/ObjectLinkingLayerTest.cpp            | 30 +++++++++++++++++++
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 25ab154a01d674..86c08cbdee5f32 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -701,16 +701,15 @@ Error ObjectLinkingLayer::handleRemoveResources(JITDylib &JD, ResourceKey K) {
 void ObjectLinkingLayer::handleTransferResources(JITDylib &JD,
                                                  ResourceKey DstKey,
                                                  ResourceKey SrcKey) {
-  auto I = Allocs.find(SrcKey);
-  if (I != Allocs.end()) {
-    auto &SrcAllocs = I->second;
+  if (Allocs.contains(SrcKey)) {
+    // DstKey may not be in the DenseMap yet, so the following line may resize
+    // the container and invalidate iterators and value references.
     auto &DstAllocs = Allocs[DstKey];
+    auto &SrcAllocs = Allocs[SrcKey];
     DstAllocs.reserve(DstAllocs.size() + SrcAllocs.size());
     for (auto &Alloc : SrcAllocs)
       DstAllocs.push_back(std::move(Alloc));
 
-    // Erase SrcKey entry using value rather than iterator I: I may have been
-    // invalidated when we looked up DstKey.
     Allocs.erase(SrcKey);
   }
 
diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index bc3433d011551c..a73b2310d193aa 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -430,16 +430,15 @@ Error RTDyldObjectLinkingLayer::handleRemoveResources(JITDylib &JD,
 void RTDyldObjectLinkingLayer::handleTransferResources(JITDylib &JD,
                                                        ResourceKey DstKey,
                                                        ResourceKey SrcKey) {
-  auto I = MemMgrs.find(SrcKey);
-  if (I != MemMgrs.end()) {
-    auto &SrcMemMgrs = I->second;
+  if (MemMgrs.contains(SrcKey)) {
+    // DstKey may not be in the DenseMap yet, so the following line may resize
+    // the container and invalidate iterators and value references.
     auto &DstMemMgrs = MemMgrs[DstKey];
+    auto &SrcMemMgrs = MemMgrs[SrcKey];
     DstMemMgrs.reserve(DstMemMgrs.size() + SrcMemMgrs.size());
     for (auto &MemMgr : SrcMemMgrs)
       DstMemMgrs.push_back(std::move(MemMgr));
 
-    // Erase SrcKey entry using value rather than iterator I: I may have been
-    // invalidated when we looked up DstKey.
     MemMgrs.erase(SrcKey);
   }
 }
diff --git a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
index 63cf3a397cb30d..bc996711f7ec50 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
@@ -65,6 +65,36 @@ TEST_F(ObjectLinkingLayerTest, AddLinkGraph) {
   EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_X"), Succeeded());
 }
 
+TEST_F(ObjectLinkingLayerTest, ResourceTracker) {
+  // This test transfers allocations to previously unknown ResourceTrackers,
+  // while increasing the number of trackers in the ObjectLinkingLayer, which
+  // may invalidate some iterators internally.
+  std::vector<ResourceTrackerSP> Trackers;
+  for (unsigned I = 0; I < 64; I++) {
+    auto G = std::make_unique<LinkGraph>("foo", Triple("x86_64-apple-darwin"),
+                                         8, llvm::endianness::little,
+                                         x86_64::getEdgeKindName);
+
+    auto &Sec1 = G->createSection("__data", MemProt::Read | MemProt::Write);
+    auto &B1 = G->createContentBlock(Sec1, BlockContent,
+                                     orc::ExecutorAddr(0x1000), 8, 0);
+    llvm::SmallString<0> SymbolName;
+    SymbolName += "_X";
+    SymbolName += std::to_string(I);
+    G->addDefinedSymbol(B1, 4, SymbolName, 4, Linkage::Strong, Scope::Default,
+                        false, false);
+
+    auto RT1 = JD.createResourceTracker();
+    EXPECT_THAT_ERROR(ObjLinkingLayer.add(RT1, std::move(G)), Succeeded());
+    EXPECT_THAT_EXPECTED(ES.lookup(&JD, SymbolName), Succeeded());
+
+    auto RT2 = JD.createResourceTracker();
+    RT1->transferTo(*RT2);
+
+    Trackers.push_back(RT2);
+  }
+}
+
 TEST_F(ObjectLinkingLayerTest, ClaimLateDefinedWeakSymbols) {
   // Check that claiming weak symbols works as expected.
   //