[ORC] Update Symbol Lookup / DefinitionGenerator system.

This patch moves definition generation out from the session lock, instead
running it under a per-dylib generator lock. It also makes the
DefinitionGenerator::tryToGenerate method optionally asynchronous: Generators
are handed an opaque LookupState object which can be captured to stop/restart
the lookup process.

The new scheme provides the following benefits and guarantees:

(1) Queries that do not need to attempt definition generation (because all
    requested symbols matched against existing definitions in the JITDylib)
    can proceed without being blocked by any running definition generators.

(2) Definition generators can capture the LookupState to continue their work
    asynchronously. This allows generators to run for an arbitrary amount of
    time without blocking a thread. Definition generators that do not need to
    run asynchronously can return without capturing the LookupState to eliminate
    unnecessary recursion and improve lookup performance.

(3) Definition generators still do not need to worry about concurrency or
    re-entrance: Since they are still run under a (per-dylib) lock, generators
    will never be re-entered concurrently, or given overlapping symbol sets to
    generate.

Finally, the new system distinguishes between symbols that are candidates for
generation (generation candidates) and symbols that failed to match for a query
(due to symbol visibility). This fixes a bug where an unresolved symbol could
trigger generation of a duplicate definition for an existing hidden symbol.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -36,6 +36,7 @@ class MaterializationUnit;
 class MaterializationResponsibility;
 class JITDylib;
 class ResourceTracker;
+class InProgressLookupState;
 
 enum class SymbolState : uint8_t;
 
@@ -215,9 +216,19 @@ class SymbolLookupSet {
 
   /// Add an element to the set. The client is responsible for checking that
   /// duplicates are not added.
-  void add(SymbolStringPtr Name,
-           SymbolLookupFlags Flags = SymbolLookupFlags::RequiredSymbol) {
+  SymbolLookupSet &
+  add(SymbolStringPtr Name,
+      SymbolLookupFlags Flags = SymbolLookupFlags::RequiredSymbol) {
     Symbols.push_back(std::make_pair(std::move(Name), Flags));
+    return *this;
+  }
+
+  /// Quickly append one lookup set to another.
+  SymbolLookupSet &append(SymbolLookupSet Other) {
+    Symbols.reserve(Symbols.size() + Other.size());
+    for (auto &KV : Other)
+      Symbols.push_back(std::move(KV));
+    return *this;
   }
 
   bool empty() const { return Symbols.empty(); }
@@ -783,6 +794,7 @@ enum class SymbolState : uint8_t {
 /// makes a callback when all symbols are available.
 class AsynchronousSymbolQuery {
   friend class ExecutionSession;
+  friend class InProgressFullLookupState;
   friend class JITDylib;
   friend class JITSymbolResolverAdapter;
   friend class MaterializationResponsibility;
@@ -829,6 +841,22 @@ class AsynchronousSymbolQuery {
   SymbolState RequiredState;
 };
 
+/// Wraps state for a lookup-in-progress.
+/// DefinitionGenerators can optionally take ownership of a LookupState object
+/// to suspend a lookup-in-progress while they search for definitions.
+class LookupState {
+  friend class ExecutionSession;
+
+public:
+  /// Continue the lookup. This can be called by DefinitionGenerators
+  /// to re-start a captured query-application operation.
+  void continueLookup(Error Err);
+
+private:
+  LookupState(std::unique_ptr<InProgressLookupState> IPLS);
+  std::unique_ptr<InProgressLookupState> IPLS;
+};
+
 /// Definition generators can be attached to JITDylibs to generate new
 /// definitions for otherwise unresolved symbols during lookup.
 class DefinitionGenerator {
@@ -841,7 +869,7 @@ class DefinitionGenerator {
   /// JDLookupFlags specifies whether the search should match against
   /// hidden symbols. Finally, Symbols describes the set of unresolved
   /// symbols and their associated lookup flags.
-  virtual Error tryToGenerate(LookupKind K, JITDylib &JD,
+  virtual Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                               JITDylibLookupFlags JDLookupFlags,
                               const SymbolLookupSet &LookupSet) = 0;
 };
@@ -979,13 +1007,6 @@ class JITDylib : public ThreadSafeRefCountedBase<JITDylib> {
   /// left unmodified (no symbols are removed).
   Error remove(const SymbolNameSet &Names);
 
-  /// Search the given JITDylib for the symbols in Symbols. If found, store
-  /// the flags for each symbol in Flags. If any required symbols are not found
-  /// then an error will be returned.
-  Expected<SymbolFlagsMap> lookupFlags(LookupKind K,
-                                       JITDylibLookupFlags JDLookupFlags,
-                                       SymbolLookupSet LookupSet);
-
   /// Dump current JITDylib state to OS.
   void dump(raw_ostream &OS);
 
@@ -1104,19 +1125,19 @@ class JITDylib : public ThreadSafeRefCountedBase<JITDylib> {
   void installMaterializationUnit(std::unique_ptr<MaterializationUnit> MU,
                                   ResourceTracker &RT);
 
-  void lookupFlagsImpl(SymbolFlagsMap &Result, LookupKind K,
-                       JITDylibLookupFlags JDLookupFlags,
-                       SymbolLookupSet &Unresolved);
+  // void lookupFlagsImpl(SymbolFlagsMap &Result, LookupKind K,
+  //                      JITDylibLookupFlags JDLookupFlags,
+  //                      SymbolLookupSet &Unresolved);
 
-  Error lodgeQuery(UnmaterializedInfosList &UMIs,
-                   std::shared_ptr<AsynchronousSymbolQuery> &Q, LookupKind K,
-                   JITDylibLookupFlags JDLookupFlags,
-                   SymbolLookupSet &Unresolved);
+  // Error lodgeQuery(UnmaterializedInfosList &UMIs,
+  //                  std::shared_ptr<AsynchronousSymbolQuery> &Q, LookupKind K,
+  //                  JITDylibLookupFlags JDLookupFlags,
+  //                  SymbolLookupSet &Unresolved);
 
-  Error lodgeQueryImpl(UnmaterializedInfosList &UMIs,
-                       std::shared_ptr<AsynchronousSymbolQuery> &Q,
-                       LookupKind K, JITDylibLookupFlags JDLookupFlags,
-                       SymbolLookupSet &Unresolved);
+  // Error lodgeQueryImpl(UnmaterializedInfosList &UMIs,
+  //                      std::shared_ptr<AsynchronousSymbolQuery> &Q,
+  //                      LookupKind K, JITDylibLookupFlags JDLookupFlags,
+  //                      SymbolLookupSet &Unresolved);
 
   void detachQueryHelper(AsynchronousSymbolQuery &Q,
                          const SymbolNameSet &QuerySymbols);
@@ -1154,11 +1175,12 @@ class JITDylib : public ThreadSafeRefCountedBase<JITDylib> {
 
   ExecutionSession &ES;
   std::string JITDylibName;
+  std::mutex GeneratorsMutex;
   bool Open = true;
   SymbolTable Symbols;
   UnmaterializedInfosMap UnmaterializedInfos;
   MaterializingInfosMap MaterializingInfos;
-  std::vector<std::unique_ptr<DefinitionGenerator>> DefGenerators;
+  std::vector<std::shared_ptr<DefinitionGenerator>> DefGenerators;
   JITDylibSearchOrder LinkOrder;
   ResourceTrackerSP DefaultTracker;
 
@@ -1199,7 +1221,10 @@ class Platform {
 
 /// An ExecutionSession represents a running JIT program.
 class ExecutionSession {
+  friend class InProgressLookupFlagsState;
+  friend class InProgressFullLookupState;
   friend class JITDylib;
+  friend class LookupState;
   friend class MaterializationResponsibility;
   friend class ResourceTracker;
 
@@ -1290,7 +1315,18 @@ class ExecutionSession {
     return *this;
   }
 
-  /// Search the given JITDylib list for the given symbols.
+  /// Search the given JITDylibs to find the flags associated with each of the
+  /// given symbols.
+  void lookupFlags(LookupKind K, JITDylibSearchOrder SearchOrder,
+                   SymbolLookupSet Symbols,
+                   unique_function<void(Expected<SymbolFlagsMap>)> OnComplete);
+
+  /// Blocking version of lookupFlags.
+  Expected<SymbolFlagsMap> lookupFlags(LookupKind K,
+                                       JITDylibSearchOrder SearchOrder,
+                                       SymbolLookupSet Symbols);
+
+  /// Search the given JITDylibs for the given symbols.
   ///
   /// SearchOrder lists the JITDylibs to search. For each dylib, the associated
   /// boolean indicates whether the search should match against non-exported
@@ -1372,7 +1408,7 @@ class ExecutionSession {
     MU->materialize(std::move(MR));
   }
 
-  void runOutstandingMUs();
+  void dispatchOutstandingMUs();
 
   static std::unique_ptr<MaterializationResponsibility>
   createMaterializationResponsibility(ResourceTracker &RT,
@@ -1390,8 +1426,36 @@ class ExecutionSession {
   void transferResourceTracker(ResourceTracker &DstRT, ResourceTracker &SrcRT);
   void destroyResourceTracker(ResourceTracker &RT);
 
-
-  /// State machine functions for MaterializationResponsibility.
+  // State machine functions for query application..
+
+  /// IL_updateCandidatesFor is called to remove already-defined symbols that
+  /// match a given query from the set of candidate symbols to generate
+  /// definitions for (no need to generate a definition if one already exists).
+  Error IL_updateCandidatesFor(JITDylib &JD, JITDylibLookupFlags JDLookupFlags,
+                               SymbolLookupSet &Candidates,
+                               SymbolLookupSet *NonCandidates);
+
+  /// OL_applyQueryPhase1 is an optionally re-startable loop for triggering
+  /// definition generation. It is called when a lookup is performed, and again
+  /// each time that LookupState::continueLookup is called.
+  void OL_applyQueryPhase1(std::unique_ptr<InProgressLookupState> IPLS,
+                           Error Err);
+
+  /// OL_completeLookup is run once phase 1 successfully completes for a lookup
+  /// call. It attempts to attach the symbol to all symbol table entries and
+  /// collect all MaterializationUnits to dispatch. If this method fails then
+  /// all MaterializationUnits will be left un-materialized.
+  void OL_completeLookup(std::unique_ptr<InProgressLookupState> IPLS,
+                         std::shared_ptr<AsynchronousSymbolQuery> Q,
+                         RegisterDependenciesFunction RegisterDependencies);
+
+  /// OL_completeLookupFlags is run once phase 1 successfully completes for a
+  /// lookupFlags call.
+  void OL_completeLookupFlags(
+      std::unique_ptr<InProgressLookupState> IPLS,
+      unique_function<void(Expected<SymbolFlagsMap>)> OnComplete);
+
+  // State machine functions for MaterializationResponsibility.
   void OL_destroyMaterializationResponsibility(
       MaterializationResponsibility &MR);
   SymbolNameSet OL_getRequestedSymbols(const MaterializationResponsibility &MR);
@@ -1454,8 +1518,8 @@ Error MaterializationResponsibility::withResourceKeyDo(Func &&F) const {
 template <typename GeneratorT>
 GeneratorT &JITDylib::addGenerator(std::unique_ptr<GeneratorT> DefGenerator) {
   auto &G = *DefGenerator;
-  ES.runSessionLocked(
-      [&]() { DefGenerators.push_back(std::move(DefGenerator)); });
+  std::lock_guard<std::mutex> Lock(GeneratorsMutex);
+  DefGenerators.push_back(std::move(DefGenerator));
   return G;
 }
 
@@ -1546,7 +1610,7 @@ class ReexportsGenerator : public DefinitionGenerator {
                      JITDylibLookupFlags SourceJDLookupFlags,
                      SymbolPredicate Allow = SymbolPredicate());
 
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &LookupSet) override;
 

llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h

@@ -254,7 +254,7 @@ class DynamicLibrarySearchGenerator : public DefinitionGenerator {
     return Load(nullptr, GlobalPrefix, std::move(Allow));
   }
 
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &Symbols) override;
 
@@ -292,7 +292,7 @@ class StaticLibraryDefinitionGenerator : public DefinitionGenerator {
   static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
   Create(ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer);
 
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &Symbols) override;
 

llvm/include/llvm/ExecutionEngine/Orc/TPCDynamicLibrarySearchGenerator.h

@@ -50,7 +50,7 @@ class TPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
     return Load(TPC, nullptr);
   }
 
-  Error tryToGenerate(LookupKind K, JITDylib &JD,
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
                       JITDylibLookupFlags JDLookupFlags,
                       const SymbolLookupSet &Symbols) override;