[ORC] Make sure that queries on emitted-but-not-ready symbols fail correctly.

In r369808 the failure scheme for ORC symbols was changed to make
MaterializationResponsibility objects responsible for failing the symbols
they represented. This simplifies error logic in the case where symbols are
still covered by a MaterializationResponsibility, but left a gap in error
handling: Symbols that have been emitted but are not yet ready (due to a
dependence on some unemitted symbol) are not covered by a
MaterializationResponsibility object. Under the scheme introduced in r369808
such symbols would be moved to the error state, but queries on those symbols
were never notified. This led to deadlocks when such symbols were failed.

This commit updates error logic to immediately fail queries on any symbol that
has already been emitted if one of its dependencies fails.

llvm-svn: 369976

include/llvm/ExecutionEngine/Orc/Core.h

@@ -430,6 +430,7 @@ enum class SymbolState : uint8_t {
   NeverSearched, /// Added to the symbol table, never queried.
   Materializing, /// Queried, materialization begun.
   Resolved,      /// Assigned address, still materializing.
+  Emitted,       /// Emitted to memory, but waiting on transitive dependencies.
   Ready = 0x3f   /// Ready and safe for clients to access.
 };
 
@@ -637,7 +638,6 @@ private:
   struct MaterializingInfo {
     SymbolDependenceMap Dependants;
     SymbolDependenceMap UnemittedDependencies;
-    bool IsEmitted = false;
 
     void addQuery(std::shared_ptr<AsynchronousSymbolQuery> Q);
     void removeQuery(const AsynchronousSymbolQuery &Q);
@@ -736,13 +736,6 @@ private:
 
   SymbolNameSet getRequestedSymbols(const SymbolFlagsMap &SymbolFlags) const;
 
-  // Move a symbol to the failure state.
-  // Detaches the symbol from all dependencies, moves all dependants to the
-  // error state (but does not fail them), deletes the MaterializingInfo for
-  // the symbol (if present) and returns the set of queries that need to be
-  // notified of the failure.
-  AsynchronousSymbolQuerySet failSymbol(const SymbolStringPtr &Name);
-
   void addDependencies(const SymbolStringPtr &Name,
                        const SymbolDependenceMap &Dependants);
 
@@ -750,7 +743,9 @@ private:
 
   Error emit(const SymbolFlagsMap &Emitted);
 
-  void notifyFailed(const SymbolFlagsMap &FailedSymbols);
+  using FailedSymbolsWorklist =
+      std::vector<std::pair<JITDylib *, SymbolStringPtr>>;
+  static void notifyFailed(FailedSymbolsWorklist FailedSymbols);
 
   ExecutionSession &ES;
   std::string JITDylibName;

lib/ExecutionEngine/Orc/Core.cpp

@@ -240,6 +240,8 @@ raw_ostream &operator<<(raw_ostream &OS, const SymbolState &S) {
     return OS << "Materializing";
   case SymbolState::Resolved:
     return OS << "Resolved";
+  case SymbolState::Emitted:
+    return OS << "Emitted";
   case SymbolState::Ready:
     return OS << "Ready";
   }
@@ -423,8 +425,13 @@ void MaterializationResponsibility::failMaterialization() {
            << SymbolFlags << "\n";
   });
 
-  JD.notifyFailed(SymbolFlags);
+  JITDylib::FailedSymbolsWorklist Worklist;
+
+  for (auto &KV : SymbolFlags)
+    Worklist.push_back(std::make_pair(&JD, KV.first));
   SymbolFlags.clear();
+
+  JD.notifyFailed(std::move(Worklist));
 }
 
 void MaterializationResponsibility::replace(
@@ -841,70 +848,6 @@ JITDylib::getRequestedSymbols(const SymbolFlagsMap &SymbolFlags) const {
   });
 }
 
-JITDylib::AsynchronousSymbolQuerySet
-JITDylib::failSymbol(const SymbolStringPtr &Name) {
-  assert(Symbols.count(Name) && "Name not in symbol table");
-  assert(Symbols[Name].getFlags().hasError() &&
-         "Failing symbol not in the error state");
-
-  auto MII = MaterializingInfos.find(Name);
-  if (MII == MaterializingInfos.end())
-    return AsynchronousSymbolQuerySet();
-
-  auto &MI = MII->second;
-
-  // Visit all dependants.
-  for (auto &KV : MI.Dependants) {
-    auto &DependantJD = *KV.first;
-    for (auto &DependantName : KV.second) {
-      assert(DependantJD.Symbols.count(DependantName) &&
-             "No symbol with DependantName in DependantJD");
-      auto &DependantSymTabEntry = DependantJD.Symbols[DependantName];
-      DependantSymTabEntry.setFlags(DependantSymTabEntry.getFlags() |
-                                    JITSymbolFlags::HasError);
-
-      assert(DependantJD.MaterializingInfos.count(DependantName) &&
-             "Dependant symbol does not have MaterializingInfo?");
-      auto &DependantMI = DependantJD.MaterializingInfos[DependantName];
-
-      assert(DependantMI.UnemittedDependencies.count(this) &&
-             "No unemitted dependency recorded for this JD?");
-      auto UnemittedDepsI = DependantMI.UnemittedDependencies.find(this);
-      assert(UnemittedDepsI != DependantMI.UnemittedDependencies.end() &&
-             "No unemitted dependency on this JD");
-      assert(UnemittedDepsI->second.count(Name) &&
-             "No unemitted dependency on symbol Name in this JD");
-
-      UnemittedDepsI->second.erase(Name);
-      if (UnemittedDepsI->second.empty())
-        DependantMI.UnemittedDependencies.erase(UnemittedDepsI);
-    }
-  }
-
-  // Visit all unemitted dependencies and disconnect from them.
-  for (auto &KV : MI.UnemittedDependencies) {
-    auto &DependencyJD = *KV.first;
-    for (auto &DependencyName : KV.second) {
-      assert(DependencyJD.MaterializingInfos.count(DependencyName) &&
-             "Dependency does not have MaterializingInfo");
-      auto &DependencyMI = DependencyJD.MaterializingInfos[DependencyName];
-      auto DependantsI = DependencyMI.Dependants.find(this);
-      assert(DependantsI != DependencyMI.Dependants.end() &&
-             "No dependnts entry recorded for this JD");
-      assert(DependantsI->second.count(Name) &&
-             "No dependants entry recorded for Name");
-      DependantsI->second.erase(Name);
-      if (DependantsI->second.empty())
-        DependencyMI.Dependants.erase(DependantsI);
-    }
-  }
-
-  AsynchronousSymbolQuerySet QueriesToFail;
-  for (auto &Q : MI.takeAllPendingQueries())
-    QueriesToFail.insert(std::move(Q));
-  return QueriesToFail;
-}
-
 void JITDylib::addDependencies(const SymbolStringPtr &Name,
                                const SymbolDependenceMap &Dependencies) {
   assert(Symbols.count(Name) && "Name not in symbol table");
@@ -916,7 +859,8 @@ void JITDylib::addDependencies(const SymbolStringPtr &Name,
     return;
 
   auto &MI = MaterializingInfos[Name];
-  assert(!MI.IsEmitted && "Can not add dependencies to an emitted symbol");
+  assert(Symbols[Name].getState() != SymbolState::Emitted &&
+         "Can not add dependencies to an emitted symbol");
 
   bool DependsOnSymbolInErrorState = false;
 
@@ -930,18 +874,21 @@ void JITDylib::addDependencies(const SymbolStringPtr &Name,
     for (auto &OtherSymbol : KV.second) {
 
       // Check the sym entry for the dependency.
-      auto SymI = OtherJITDylib.Symbols.find(OtherSymbol);
+      auto OtherSymI = OtherJITDylib.Symbols.find(OtherSymbol);
 
 #ifndef NDEBUG
-      // Assert that this symbol exists and has not been emitted already.
+      // Assert that this symbol exists and has not reached the ready state
-      assert(SymI != OtherJITDylib.Symbols.end() &&
+      // already.
-             (SymI->second.getState() != SymbolState::Ready &&
+      assert(OtherSymI != OtherJITDylib.Symbols.end() &&
-              "Dependency on emitted symbol"));
+             (OtherSymI->second.getState() != SymbolState::Ready &&
+              "Dependency on emitted/ready symbol"));
 #endif
 
+      auto &OtherSymEntry = OtherSymI->second;
+
       // If the dependency is in an error state then note this and continue,
       // we will move this symbol to the error state below.
-      if (SymI->second.getFlags().hasError()) {
+      if (OtherSymEntry.getFlags().hasError()) {
         DependsOnSymbolInErrorState = true;
         continue;
       }
@@ -952,7 +899,7 @@ void JITDylib::addDependencies(const SymbolStringPtr &Name,
              "No MaterializingInfo for dependency");
       auto &OtherMI = OtherJITDylib.MaterializingInfos[OtherSymbol];
 
-      if (OtherMI.IsEmitted)
+      if (OtherSymEntry.getState() == SymbolState::Emitted)
         transferEmittedNodeDependencies(MI, Name, OtherMI);
       else if (&OtherJITDylib != this || OtherSymbol != Name) {
         OtherMI.Dependants[this].insert(Name);
@@ -1087,6 +1034,12 @@ Error JITDylib::emit(const SymbolFlagsMap &Emitted) {
       Worklist.pop_back();
 
       auto &Name = SymI->first;
+      auto &SymEntry = SymI->second;
+
+      // Move symbol to the emitted state.
+      assert(SymEntry.getState() == SymbolState::Resolved &&
+             "Emitting from state other than Resolved");
+      SymEntry.setState(SymbolState::Emitted);
 
       auto MII = MaterializingInfos.find(Name);
       assert(MII != MaterializingInfos.end() &&
@@ -1121,20 +1074,22 @@ Error JITDylib::emit(const SymbolFlagsMap &Emitted) {
           DependantJD.transferEmittedNodeDependencies(DependantMI,
                                                       DependantName, MI);
 
+          auto DependantSymI = DependantJD.Symbols.find(DependantName);
+          assert(DependantSymI != DependantJD.Symbols.end() &&
+                 "Dependant has no entry in the Symbols table");
+          auto &DependantSymEntry = DependantSymI->second;
+
           // If the dependant is emitted and this node was the last of its
           // unemitted dependencies then the dependant node is now ready, so
           // notify any pending queries on the dependant node.
-          if (DependantMI.IsEmitted &&
+          if (DependantSymEntry.getState() == SymbolState::Emitted &&
               DependantMI.UnemittedDependencies.empty()) {
             assert(DependantMI.Dependants.empty() &&
                    "Dependants should be empty by now");
 
             // Since this dependant is now ready, we erase its MaterializingInfo
             // and update its materializing state.
-            auto DependantSymI = DependantJD.Symbols.find(DependantName);
+            DependantSymEntry.setState(SymbolState::Ready);
-            assert(DependantSymI != DependantJD.Symbols.end() &&
-                   "Dependant has no entry in the Symbols table");
-            DependantSymI->second.setState(SymbolState::Ready);
 
             for (auto &Q : DependantMI.takeQueriesMeeting(SymbolState::Ready)) {
               Q->notifySymbolMetRequiredState(
@@ -1148,9 +1103,8 @@ Error JITDylib::emit(const SymbolFlagsMap &Emitted) {
           }
         }
       }
-      MI.Dependants.clear();
-      MI.IsEmitted = true;
 
+      MI.Dependants.clear();
       if (MI.UnemittedDependencies.empty()) {
         SymI->second.setState(SymbolState::Ready);
         for (auto &Q : MI.takeQueriesMeeting(SymbolState::Ready)) {
@@ -1183,17 +1137,27 @@ Error JITDylib::emit(const SymbolFlagsMap &Emitted) {
   return Error::success();
 }
 
-void JITDylib::notifyFailed(const SymbolFlagsMap &FailedSymbols) {
+void JITDylib::notifyFailed(FailedSymbolsWorklist Worklist) {
   AsynchronousSymbolQuerySet FailedQueries;
+  auto FailedSymbolsMap = std::make_shared<SymbolDependenceMap>();
+
+  // Failing no symbols is a no-op.
+  if (Worklist.empty())
+    return;
+
+  auto &ES = Worklist.front().first->getExecutionSession();
 
   ES.runSessionLocked([&]() {
-    std::vector<const SymbolStringPtr *> MaterializerNamesToFail;
+    while (!Worklist.empty()) {
+      assert(Worklist.back().first && "Failed JITDylib can not be null");
+      auto &JD = *Worklist.back().first;
+      auto Name = std::move(Worklist.back().second);
+      Worklist.pop_back();
 
-    for (auto &KV : FailedSymbols) {
+      (*FailedSymbolsMap)[&JD].insert(Name);
-      auto &Name = KV.first;
 
-      assert(Symbols.count(Name) && "No symbol table entry for Name");
+      assert(JD.Symbols.count(Name) && "No symbol table entry for Name");
-      auto &Sym = Symbols[Name];
+      auto &Sym = JD.Symbols[Name];
 
       // Move the symbol into the error state.
       // Note that this may be redundant: The symbol might already have been
@@ -1203,13 +1167,12 @@ void JITDylib::notifyFailed(const SymbolFlagsMap &FailedSymbols) {
       // FIXME: Come up with a sane mapping of state to
       // presence-of-MaterializingInfo so that we can assert presence / absence
       // here, rather than testing it.
-      auto MII = MaterializingInfos.find(Name);
+      auto MII = JD.MaterializingInfos.find(Name);
 
-      if (MII == MaterializingInfos.end())
+      if (MII == JD.MaterializingInfos.end())
         continue;
 
       auto &MI = MII->second;
-      MaterializerNamesToFail.push_back(&KV.first);
 
       // Move all dependants to the error state and disconnect from them.
       for (auto &KV : MI.Dependants) {
@@ -1225,7 +1188,7 @@ void JITDylib::notifyFailed(const SymbolFlagsMap &FailedSymbols) {
                  "No MaterializingInfo for dependant");
           auto &DependantMI = DependantJD.MaterializingInfos[DependantName];
 
-          auto UnemittedDepI = DependantMI.UnemittedDependencies.find(this);
+          auto UnemittedDepI = DependantMI.UnemittedDependencies.find(&JD);
           assert(UnemittedDepI != DependantMI.UnemittedDependencies.end() &&
                  "No UnemittedDependencies entry for this JITDylib");
           assert(UnemittedDepI->second.count(Name) &&
@@ -1233,8 +1196,18 @@ void JITDylib::notifyFailed(const SymbolFlagsMap &FailedSymbols) {
           UnemittedDepI->second.erase(Name);
           if (UnemittedDepI->second.empty())
             DependantMI.UnemittedDependencies.erase(UnemittedDepI);
+
+          // If this symbol is already in the emitted state then we need to
+          // take responsibility for failing its queries, so add it to the
+          // worklist.
+          if (DependantSym.getState() == SymbolState::Emitted) {
+            assert(DependantMI.Dependants.empty() &&
+                   "Emitted symbol should not have dependants");
+            Worklist.push_back(std::make_pair(&DependantJD, DependantName));
+          }
         }
       }
+      MI.Dependants.clear();
 
       // Disconnect from all unemitted depenencies.
       for (auto &KV : MI.UnemittedDependencies) {
@@ -1244,35 +1217,35 @@ void JITDylib::notifyFailed(const SymbolFlagsMap &FailedSymbols) {
               UnemittedDepJD.MaterializingInfos.find(UnemittedDepName);
           assert(UnemittedDepMII != UnemittedDepJD.MaterializingInfos.end() &&
                  "Missing MII for unemitted dependency");
-          assert(UnemittedDepMII->second.Dependants.count(this) &&
+          assert(UnemittedDepMII->second.Dependants.count(&JD) &&
                  "JD not listed as a dependant of unemitted dependency");
-          assert(UnemittedDepMII->second.Dependants[this].count(Name) &&
+          assert(UnemittedDepMII->second.Dependants[&JD].count(Name) &&
                  "Name is not listed as a dependant of unemitted dependency");
-          UnemittedDepMII->second.Dependants[this].erase(Name);
+          UnemittedDepMII->second.Dependants[&JD].erase(Name);
-          if (UnemittedDepMII->second.Dependants[this].empty())
+          if (UnemittedDepMII->second.Dependants[&JD].empty())
-            UnemittedDepMII->second.Dependants.erase(this);
+            UnemittedDepMII->second.Dependants.erase(&JD);
         }
       }
+      MI.UnemittedDependencies.clear();
 
       // Collect queries to be failed for this MII.
-      for (auto &Q : MII->second.pendingQueries())
+      for (auto &Q : MII->second.pendingQueries()) {
+        // Add the query to the list to be failed and detach it.
         FailedQueries.insert(Q);
-    }
+        Q->detach();
+      }
 
-    // Detach failed queries.
+      assert(MI.Dependants.empty() &&
-    for (auto &Q : FailedQueries)
+             "Can not delete MaterializingInfo with dependants still attached");
-      Q->detach();
+      assert(MI.UnemittedDependencies.empty() &&
-
+             "Can not delete MaterializingInfo with unemitted dependencies "
-    // Remove the MaterializingInfos.
+             "still attached");
-    while (!MaterializerNamesToFail.empty()) {
+      assert(!MI.hasQueriesPending() &&
-      MaterializingInfos.erase(*MaterializerNamesToFail.back());
+             "Can not delete MaterializingInfo with queries pending");
-      MaterializerNamesToFail.pop_back();
+      JD.MaterializingInfos.erase(MII);
     }
   });
 
-  auto FailedSymbolsMap = std::make_shared<SymbolDependenceMap>();
-  for (auto &KV : FailedSymbols)
-    (*FailedSymbolsMap)[this].insert(KV.first);
   for (auto &Q : FailedQueries)
     Q->handleFailed(make_error<FailedToMaterialize>(FailedSymbolsMap));
 }
@@ -1701,8 +1674,6 @@ void JITDylib::dump(raw_ostream &OS) {
       OS << "  MaterializingInfos entries:\n";
     for (auto &KV : MaterializingInfos) {
       OS << "    \"" << *KV.first << "\":\n"
-         << "      IsEmitted = " << (KV.second.IsEmitted ? "true" : "false")
-         << "\n"
          << "      " << KV.second.pendingQueries().size()
          << " pending queries: { ";
       for (const auto &Q : KV.second.pendingQueries())

unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

@@ -718,6 +718,60 @@ TEST_F(CoreAPIsStandardTest, AddDependencyOnFailedSymbol) {
       << "Lookup on failed symbol should fail";
 }
 
+TEST_F(CoreAPIsStandardTest, FailAfterMaterialization) {
+  Optional<MaterializationResponsibility> FooR;
+  Optional<MaterializationResponsibility> BarR;
+
+  // Create a MaterializationUnit for each symbol that moves the
+  // MaterializationResponsibility into one of the locals above.
+  auto FooMU = std::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
+      [&](MaterializationResponsibility R) { FooR.emplace(std::move(R)); });
+
+  auto BarMU = std::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Bar, BarSym.getFlags()}}),
+      [&](MaterializationResponsibility R) { BarR.emplace(std::move(R)); });
+
+  // Define the symbols.
+  cantFail(JD.define(FooMU));
+  cantFail(JD.define(BarMU));
+
+  bool OnFooReadyRun = false;
+  auto OnFooReady = [&](Expected<SymbolMap> Result) {
+    EXPECT_THAT_EXPECTED(std::move(Result), Failed());
+    OnFooReadyRun = true;
+  };
+
+  ES.lookup(JITDylibSearchList({{&JD, false}}), {Foo}, SymbolState::Ready,
+            std::move(OnFooReady), NoDependenciesToRegister);
+
+  bool OnBarReadyRun = false;
+  auto OnBarReady = [&](Expected<SymbolMap> Result) {
+    EXPECT_THAT_EXPECTED(std::move(Result), Failed());
+    OnBarReadyRun = true;
+  };
+
+  ES.lookup(JITDylibSearchList({{&JD, false}}), {Bar}, SymbolState::Ready,
+            std::move(OnBarReady), NoDependenciesToRegister);
+
+  // Add a dependency by Foo on Bar and vice-versa.
+  FooR->addDependenciesForAll({{&JD, SymbolNameSet({Bar})}});
+  BarR->addDependenciesForAll({{&JD, SymbolNameSet({Foo})}});
+
+  // Materialize Foo.
+  EXPECT_THAT_ERROR(FooR->notifyResolved({{Foo, FooSym}}), Succeeded())
+      << "Expected resolution for \"Foo\" to succeed.";
+  EXPECT_THAT_ERROR(FooR->notifyEmitted(), Succeeded())
+      << "Expected emission for \"Foo\" to succeed.";
+
+  // Fail bar.
+  BarR->failMaterialization();
+
+  // Verify that both queries failed.
+  EXPECT_TRUE(OnFooReadyRun) << "Query for Foo did not run";
+  EXPECT_TRUE(OnBarReadyRun) << "Query for Bar did not run";
+}
+
 TEST_F(CoreAPIsStandardTest, FailMaterializerWithUnqueriedSymbols) {
   // Make sure that symbols with no queries aganist them still
   // fail correctly.