From eb918d8daf1a8b31381f02adbf8fc31b250cc6f8 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Wed, 18 Mar 2020 19:56:55 -0700
Subject: [PATCH] [ORC] Use finer-grained and session locking in MachOPlatform
 to avoid deadlock.

In MachOPlatform, obtaining the link-order for a JITDylib requires locking the
session, but also needs to be part of a larger atomic operation that collates
initializer symbols tracked by the platform. Trying to do this under a separate
platform mutex leads to potential locking order issues, e.g.

T1 locks session then tries to lock platform to register a new init symbol
meanwhile
T2 locks platform then tries to lock session to obtain link order.

Removing the platform lock and performing all these operations under the session
lock eliminates this possibility.

At the same time we also need to collate init pointers from the
MachOPlatform::InitScraperPlugin, and we don't need or want to lock the session
for that. The new InitSeqMutex has been added to guard these init pointers, and
the session mutex is never obtained while the InitSeqMutex is held.
---
 .../llvm/ExecutionEngine/Orc/MachOPlatform.h     |  5 ++++-
 llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp   | 16 +++++++---------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
index f869ebdfbe4e..15fe079eccaf 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/MachOPlatform.h
@@ -143,12 +143,15 @@ private:
                         MachOJITDylibInitializers::SectionExtent ObjCSelRefs,
                         MachOJITDylibInitializers::SectionExtent ObjCClassList);
 
-  std::mutex PlatformMutex;
   ExecutionSession &ES;
   ObjectLinkingLayer &ObjLinkingLayer;
   std::unique_ptr<MemoryBuffer> StandardSymbolsObject;
 
   DenseMap<JITDylib *, SymbolLookupSet> RegisteredInitSymbols;
+
+  // InitSeqs gets its own mutex to avoid locking the whole session when
+  // aggregating data from the jitlink.
+  std::mutex InitSeqsMutex;
   DenseMap<JITDylib *, MachOJITDylibInitializers> InitSeqs;
 };
 
diff --git a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
index 9a836677ef15..13a187c66387 100644
--- a/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp
@@ -163,7 +163,6 @@ Error MachOPlatform::notifyAdding(JITDylib &JD, const MaterializationUnit &MU) {
   if (!InitSym)
     return Error::success();
 
-  std::lock_guard<std::mutex> Lock(PlatformMutex);
   RegisteredInitSymbols[&JD].add(InitSym);
   LLVM_DEBUG({
     dbgs() << "MachOPlatform: Registered init symbol " << *InitSym << " for MU "
@@ -187,11 +186,10 @@ MachOPlatform::getInitializerSequence(JITDylib &JD) {
   std::vector<JITDylib *> DFSLinkOrder;
 
   while (true) {
-    // Lock the platform while we search for any initializer symbols to
-    // look up.
+
     DenseMap<JITDylib *, SymbolLookupSet> NewInitSymbols;
-    {
-      std::lock_guard<std::mutex> Lock(PlatformMutex);
+
+    ES.runSessionLocked([&]() {
       DFSLinkOrder = getDFSLinkOrder(JD);
 
       for (auto *InitJD : DFSLinkOrder) {
@@ -201,7 +199,7 @@ MachOPlatform::getInitializerSequence(JITDylib &JD) {
           RegisteredInitSymbols.erase(RISItr);
         }
       }
-    }
+    });
 
     if (NewInitSymbols.empty())
       break;
@@ -228,7 +226,7 @@ MachOPlatform::getInitializerSequence(JITDylib &JD) {
   // Lock again to collect the initializers.
   InitializerSequence FullInitSeq;
   {
-    std::lock_guard<std::mutex> Lock(PlatformMutex);
+    std::lock_guard<std::mutex> Lock(InitSeqsMutex);
     for (auto *InitJD : reverse(DFSLinkOrder)) {
       LLVM_DEBUG({
         dbgs() << "MachOPlatform: Appending inits for \"" << InitJD->getName()
@@ -251,7 +249,7 @@ MachOPlatform::getDeinitializerSequence(JITDylib &JD) {
 
   DeinitializerSequence FullDeinitSeq;
   {
-    std::lock_guard<std::mutex> Lock(PlatformMutex);
+    std::lock_guard<std::mutex> Lock(InitSeqsMutex);
     for (auto *DeinitJD : DFSLinkOrder) {
       FullDeinitSeq.emplace_back(DeinitJD, MachOJITDylibDeinitializers());
     }
@@ -285,7 +283,7 @@ void MachOPlatform::registerInitInfo(
     MachOJITDylibInitializers::SectionExtent ModInits,
     MachOJITDylibInitializers::SectionExtent ObjCSelRefs,
     MachOJITDylibInitializers::SectionExtent ObjCClassList) {
-  std::lock_guard<std::mutex> Lock(PlatformMutex);
+  std::lock_guard<std::mutex> Lock(InitSeqsMutex);
 
   auto &InitSeq = InitSeqs[&JD];