[ORC] Fix race-condition in RTDyldObjectLinkingLayer::onObjEmit.

[mkroll]

Summary

The blocking version of ExecutionSession::lookup() uses std::promise() to wait
for the result of a compilation.
In RTDyldObjectLinkingLayer::onObjEmit() another thread could have compiled
the object and notify the promise via R.notifyEmitted(), that the object was emitted, before
the memory manager of the object was registered with the resource tracker.
Thus the code could already be executed and the object removed from the
JIT, before the memory manager was registered, leading to a errors
about a defunct resource tracker.

[mkroll]

Added missing check on Expected result of lookup.
Apply ugly clang-format include reordering (this is not even alphabetic, is this just about editing distance to the containing file name?).

[lhames]

This is a really nice catch. Thank you for finding this!

I suspect the ObjectLinkingLayer has a similar issue. I'll look into that.

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

367–392

This call to notifyObjectLoaded will reference the moved-from MemMgr argument if there are any listeners attached.

The address of the MemMgr should be recorded in a local variable for use in the notifyObjectLoaded call. It would also be good if MemMgr were removed from the MemMgrs vector if the call to notifyEmitted fails.

[mkroll]

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

367–392

Ouch, how did I miss that std::move()...

But this would move the problem to the event listener. When the event listener is called, the other thread (unblocked by R.notifyEmitted()) could already have called ExecutionEngine::removeJITDylib() resulting in the memory manager to be removed in RTDyldObjectLinkingLayer::handleRemoveResources().

Hm, RtDyldObjectLinkingLayer::onObjEmit() calls three types of notification functions:

1. MaterializationResponsibility::notifyEmitted(): "Notifies the target JITDylib (and any pending queries on that JITDylib) that all symbols covered by this MaterializationResponsibility instance have been emitted."
2. JITEventListener::notifyObjectLoaded(): "Called after an object has had its sections allocated and addresses assigned to all symbols. Note: Section memory will not have been relocated yet. notifyFunctionLoaded will not be called for individual functions in the object." (notifyFunctionLoaded should probably be notifyObjectLoaded?)
3. RTDyldObjectLinkingLayer::NotifyEmitted(): "Functor for receiving finalization notifications"

Do you think, it would be possible to change the order to

1. JITEventListener::notifyObjectLoaded()
2. RTDyldObjectLinkingLayer::NotifyEmitted()
3. Add MemMgr to MemMgrs
4. MaterializationResponsibility::notifyEmitted()

?

This would then require to also undo adding the memory manager and other things by calling handleRemoveResources, if MaterializationResponsibility::notifyEmitted() failed.

The OrcJITTests seem to be OK with this...

[lhames]

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

367–392

But this would move the problem to the event listener.

Good point!

Do you think, it would be possible to change the order to

1. JITEventListener::notifyObjectLoaded()
2. RTDyldObjectLinkingLayer::NotifyEmitted()
3. Add MemMgr to MemMgrs
4. MaterializationResponsibility::notifyEmitted()

?

This would then require to also undo adding the memory manager and other things by calling handleRemoveResources, if MaterializationResponsibility::notifyEmitted() failed.

The OrcJITTests seem to be OK with this...

If the unit and regression tests pass then I'd be inclined to land that. If it breaks out-of-tree clients we may need to revert and reconsider.

What's your JIT use case and architecture? Have you looked at ObjectLinkingLayer / JITLink yet? We're going to deprecate (and eventually remove) RuntimeDyld in favor of JITLink, so it's even more important that we get this right in ObjectLinkingLayer.

[mkroll]

Finally found some time to continue looking into this...
ObjectLinkingLayer doesn't have this problem because it uses this order in ObjectLinkingLayerJITLinkContext:

1. ObjectLinkingLayer::notifyEmitted():
   • ObjectLinkingLayer::Plugin::notifyEmitted()
   • Add FinalizedAlloc to Allocs
2. MaterializationResponsiblity::notifyEmitted() (note: The FinalizedAlloc is not explicitly removed from Allocs, if this fails. I didn't check, whether this happens indirectly)

We use the JIT to compile and execute shader code on the CPU (we also have backends for GPU). The problem occurred, when a mesh with 3 mio triangles was displaced on ARM with 8 CPUs and the same (we fixed that) displacement shader code was compiled 300 times in parallel.

Our JIT is based on the BuildingAJIT tutorial, not emitting lazily, using the ConcurrentIRCompiler, searching for symbols in specific modules to avoid problems with duplicate names and using an own copy of DefaultMMapper to avoid a use-after-free problem in the SectionMemoryManager, when things happen in the wrong order during application shutdown. All modules are linked with one "runtime" BareJITDylib, which only consists of absolute symbols pointing to hand-picked runtime functions. We're currently still using LLVM 12, but good to know, that I will have to rewrite the JIT integration again as with every LLVM upgrade ^_^

[mkroll]

Update order as discussed to first EventListeners, then NotifyEmitted, then adding the MemMgr and then MaterializationResponsibility::notifyEmitted().
This is basically the same order as in ObjectLinkingLayer now.

ORCJitTests and LLVM regression tests are successful for me.

[mkroll]

Ping!
I won't have access to my computer for two months starting on Wednesday, so if you still want anything to be changed before you can approve this, there's not much time for me left ^_^

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

 
 using namespace llvm;
 using namespace llvm::orc;
 
 class JITDylibSearchOrderResolver : public JITSymbolResolver {
 public:
   JITDylibSearchOrderResolver(MaterializationResponsibility &MR) : MR(MR) {}
 
-  void lookup(const LookupSet &Symbols, OnResolvedFunction OnResolved) override {
+  void lookup(const LookupSet &Symbols,
+              OnResolvedFunction OnResolved) override {
     auto &ES = MR.getTargetJITDylib().getExecutionSession();
     SymbolLookupSet InternedSymbols;
 
     // Intern the requested symbols: lookup takes interned strings.
     for (auto &S : Symbols)
       InternedSymbols.add(ES.intern(S));
 
     // Build an OnResolve callback to unwrap the interned strings and pass them
void RTDyldObjectLinkingLayer::onObjEmit(
     std::unique_ptr<RuntimeDyld::MemoryManager> MemMgr,
     std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo, Error Err) {
   if (Err) {
     getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();
     return;
   }
 
-  if (auto Err = R.notifyEmitted()) {
-    getExecutionSession().reportError(std::move(Err));
-    R.failMaterialization();
-    return;
-  }
-
   std::unique_ptr<object::ObjectFile> Obj;
   std::unique_ptr<MemoryBuffer> ObjBuffer;
   std::tie(Obj, ObjBuffer) = O.takeBinary();
 
   // Run EventListener notifyLoaded callbacks.
   {
     std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
     for (auto *L : EventListeners)
       L->notifyObjectLoaded(pointerToJITTargetAddress(MemMgr.get()), *Obj,
                             *LoadedObjInfo);
   }
 
   if (NotifyEmitted)
     NotifyEmitted(R, std::move(ObjBuffer));
 
+  // Add to MemMgrs before potentially notifying other threads, that the object
+  // was emitted. Otherwise these other threads could already execute the code
+  // and remove the object, leading to defunct resource trackers.
   if (auto Err = R.withResourceKeyDo(
           [&](ResourceKey K) { MemMgrs[K].push_back(std::move(MemMgr)); })) {
     getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();
+    return;
+  }
+
+  if (auto Err = R.notifyEmitted()) {
+    // Undo adding the memory manager and calling the event listeners.
+    consumeError(R.withResourceKeyDo([&](ResourceKey K) {
+      consumeError(handleRemoveResources(R.getTargetJITDylib(), K));
+    }));
+
+    getExecutionSession().reportError(std::move(Err));
+    R.failMaterialization();
   }
 }
 
 Error RTDyldObjectLinkingLayer::handleRemoveResources(JITDylib &JD,
                                                       ResourceKey K) {
 
   std::vector<MemoryManagerUP> MemMgrsToRemove;
 

llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp

 //===--- RTDyldObjectLinkingLayerTest.cpp - RTDyld linking layer tests ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "OrcTestCommon.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
-#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 #include <string>
 
 using namespace llvm;
 using namespace llvm::orc;
 
 namespace {
 
ES.lookup(
       SymbolState::Resolved,
       [](Expected<SymbolMap> R) { cantFail(std::move(R)); },
       NoDependenciesToRegister);
 
   if (auto Err = ES.endSession())
     ES.reportError(std::move(Err));
 }
 
+TEST(RTDyldObjectLinkingLayerTest, RemoveJITDylibMaterialzedInOtherThread) {
+#if LLVM_ENABLE_THREADS
+  OrcNativeTarget::initialize();
+
+  JITTargetMachineBuilder JTMB =
+      cantFail(JITTargetMachineBuilder::detectHost());
+
+  // Try to build a target machine.
+  Expected<std::unique_ptr<TargetMachine>> TM = JTMB.createTargetMachine();
+  if (TM.takeError())
+    GTEST_SKIP();
+
+  // Create a module with a void() function foo.
+  ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
+  ThreadSafeModule MFoo;
+  {
+    ModuleBuilder MB(*TSCtx.getContext(), (*TM)->getTargetTriple().str(),
+                     "FooMod");
+    MB.getModule()->setDataLayout((*TM)->createDataLayout());
+
+    Function *Func = MB.createFunctionDecl(
+        FunctionType::get(Type::getVoidTy(*TSCtx.getContext()), {}, false),
+        "foo");
+    Func->setLinkage(GlobalValue::ExternalLinkage);
+    BasicBlock *FooEntry =
+        BasicBlock::Create(*TSCtx.getContext(), "entry", Func);
+    IRBuilder<> B(FooEntry);
+    B.CreateRetVoid();
+
+    MFoo = ThreadSafeModule(MB.takeModule(), TSCtx);
+  }
+
+  // Create a simple stack.
+  ExecutionSession ES{std::make_unique<UnsupportedExecutorProcessControl>()};
+  RTDyldObjectLinkingLayer ObjLayer(
+      ES, []() { return std::make_unique<SectionMemoryManager>(); });
+  IRCompileLayer CompileLayer(ES, ObjLayer,
+                              std::make_unique<ConcurrentIRCompiler>(JTMB));
+
+  // By default, the Exported flag is not set for symbols.
+  // Ensure, the flags are overridden, if necessary
+  if (JTMB.getTargetTriple().isOSBinFormatCOFF())
+    ObjLayer.setOverrideObjectFlagsWithResponsibilityFlags(true);
+
+  auto &JDFoo = ES.createBareJITDylib("JDFoo");
+  cantFail(CompileLayer.add(JDFoo, std::move(MFoo)));
+
+  // What we want to achieve:
+  //  - Thread1 calls lookup(Foo) and gets stuck in waiting for the promise
+  //  - Thread2 is created to dispatch the materialization task for FooMod
+  //  - Thread2 materializes Foo and sends notifyEmitted in
+  //    RTDyldObjectLinkingLayer::onObjEmit()
+  //  - Thread1 reads the promised result and removes the module JD
+  //  - Thread2 continues execution in onObjEmit().
+  //    In the original version, onObjEmit() would now try to add the
+  //    MemoryManager for the already removed module which results in an error
+  //    due to a defunct resource tracker
+  //    ("JIT session error: Resource tracker <some address> became defunct").
+
+  // Always run tasks on an own thread.
+  std::mutex WorkThreadsMutex;
+  std::vector<std::thread> WorkThreads;
+  ES.setDispatchTask([&](std::unique_ptr<Task> T) {
+    std::promise<void> WaitP;
+    std::lock_guard<std::mutex> Lock(WorkThreadsMutex);
+    WorkThreads.push_back(
+        std::thread([T = std::move(T), WaitF = WaitP.get_future()]() mutable {
+          WaitF.get();
+          T->run();
+        }));
+    WaitP.set_value();
+  });
+
+  // No errors should happen in this session.
+  ES.setErrorReporter([&](Error Err) {
+    EXPECT_THAT_ERROR(std::move(Err), Succeeded()) << "Test failed";
+  });
+
+  // Install NotifyEmitted handler to delay execution in
+  // RTDyldObjectLinkingLayer::onObjEmit().
+  ObjLayer.setNotifyEmitted(
+      [](MaterializationResponsibility &R, std::unique_ptr<MemoryBuffer>) {
+        std::this_thread::sleep_for(std::chrono::milliseconds(1000));
+      });
+
+  // Do a blocking lookup on foo.
+  auto Foo = ES.intern("foo");
+  cantFail(ES.lookup(makeJITDylibSearchOrder(&JDFoo), SymbolLookupSet({Foo})));
+
+  // Remove JDFoo while emitting thread is stuck.
+  cantFail(ES.removeJITDylib(JDFoo));
+
+  for (auto &WT : WorkThreads)
+    WT.join();
+
+  if (auto Err = ES.endSession())
+    ES.reportError(std::move(Err));
+#endif
+}
+
 } // end anonymous namespace