[orc-rt] Add NativeDylibManager. (#194792)

NativeDylibManager is an orc_rt::Service that supports loading,
unloading, and lookup of symbols via the system dynamic loader's native
APIs.

The current implementation only supports the POSIX dlfcn.h APIs (dlopen,
dlclose, dlsym), but it should be straightforward to extend to Windows.

orc-rt/include/CMakeLists.txt

@@ -16,6 +16,7 @@ set(ORC_RT_HEADERS
     orc-rt/LockedAccess.h
     orc-rt/Math.h
     orc-rt/MemoryFlags.h
+    orc-rt/NativeDylibManager.h
     orc-rt/QueueingTaskDispatcher.h
     orc-rt/RTTI.h
     orc-rt/ScopeExit.h
@@ -25,6 +26,7 @@ set(ORC_RT_HEADERS
     orc-rt/SimplePackedSerialization.h
     orc-rt/SPSAllocAction.h
     orc-rt/sps-ci/AllSPSCI.h
+    orc-rt/sps-ci/NativeDylibManager.h
     orc-rt/sps-ci/SimpleNativeMemoryMapSPSCI.h
     orc-rt/SPSMemoryFlags.h
     orc-rt/SPSWrapperFunction.h

orc-rt/include/orc-rt/NativeDylibManager.h

@@ -0,0 +1,94 @@
+//===--- NativeDylibManager.h - Manage dylibs via native APIs ---*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Manage dynamic libraries via the native OS APIs in the executor.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_NATIVEDYLIBMANAGER_H
+#define ORC_RT_NATIVEDYLIBMANAGER_H
+
+#include "orc-rt/BootstrapInfo.h"
+#include "orc-rt/Service.h"
+#include "orc-rt/sps-ci/NativeDylibManagerSPSCI.h"
+
+#include <mutex>
+#include <unordered_map>
+
+namespace orc_rt {
+
+class Session;
+
+/// Dylib loading / unloading / symbol lookup service.
+///
+/// Any dynamic libraries loaded through this service that are not manually
+/// unloaded will be automatically unloaded at shutdown time in LIFO order.
+class NativeDylibManager : public Service {
+public:
+  /// Create a NativeDylibManager, adding associated symbols to the given
+  /// SimpleSymbolTable (typically the BootstrapInfo table).
+  static Expected<std::unique_ptr<NativeDylibManager>>
+  Create(Session &S, SimpleSymbolTable &ST,
+         const char *InstanceName = "orc_rt_ci_NativeDylibManager_Instance",
+         SimpleSymbolTable::MutatorFn AddInterface =
+             sps_ci::addNativeDylibManager);
+
+  /// Convenience constructor that adds default symbols to the given
+  /// BootstrapInfo's symbols map.
+  static Expected<std::unique_ptr<NativeDylibManager>>
+  Create(Session &S, BootstrapInfo &BI) {
+    return Create(S, BI.symbols());
+  }
+
+  /// NativeDylibManager is not copyable / moveable.
+  NativeDylibManager(const NativeDylibManager &) = delete;
+  NativeDylibManager &operator=(const NativeDylibManager &) = delete;
+  NativeDylibManager(NativeDylibManager &&) = delete;
+  NativeDylibManager &operator=(NativeDylibManager &&) = delete;
+
+  /// Load the given library.
+  ///
+  /// Returns an Expected handle.
+  using OnLoadCompleteFn = move_only_function<void(Expected<void *>)>;
+  void load(OnLoadCompleteFn &&OnComplete, std::string Path);
+
+  /// Unload the given library handle.
+  ///
+  /// Returns an error on failure.
+  using OnUnloadCompleteFn = move_only_function<void(Error)>;
+  void unload(OnUnloadCompleteFn &&OnComplete, void *Handle);
+
+  /// Lookup addresses of the given symbols.
+  ///
+  /// Returns a sequence of addresses.
+  using OnLookupCompleteFn =
+      move_only_function<void(Expected<std::vector<void *>>)>;
+  void lookup(OnLookupCompleteFn &&OnLookupComplete, void *Handle,
+              std::vector<std::string> Names);
+
+  void onDetach(Service::OnCompleteFn OnComplete,
+                bool ShutdownRequested) override;
+  void onShutdown(Service::OnCompleteFn OnComplete) override;
+
+private:
+  NativeDylibManager(Session &S) : S(S) {}
+
+  Session &S;
+
+  struct LoadInfo {
+    size_t Ordinal = 0;
+    size_t RefCount = 0;
+  };
+
+  std::mutex M;
+  std::unordered_map<void *, LoadInfo> LoadInfos;
+};
+
+} // namespace orc_rt
+
+#endif // ORC_RT_NATIVEDYLIBMANAGER_H

orc-rt/include/orc-rt/sps-ci/NativeDylibManagerSPSCI.h

@@ -0,0 +1,25 @@
+//===-- NativeDylibManagerSPSCI.h -- NativeDylibManager SPS CI --*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// SPS Controller Interface registration for NativeDylibManager.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_SPS_CI_NATIVEDYLIBMANAGERSPSCI_H
+#define ORC_RT_SPS_CI_NATIVEDYLIBMANAGERSPSCI_H
+
+#include "orc-rt/SimpleSymbolTable.h"
+
+namespace orc_rt::sps_ci {
+
+/// Add the NativeDylibManager SPS interface to the controller interface.
+Error addNativeDylibManager(SimpleSymbolTable &ST);
+
+} // namespace orc_rt::sps_ci
+
+#endif // ORC_RT_SPS_CI_NATIVEDYLIBMANAGERSPSCI_H

orc-rt/lib/executor/CMakeLists.txt

@@ -4,6 +4,7 @@ set(files
   SimpleSymbolTable.cpp
   Error.cpp
   ExecutorProcessInfo.cpp
+  NativeDylibManager.cpp
   QueueingTaskDispatcher.cpp
   RTTI.cpp
   Service.cpp
@@ -12,6 +13,7 @@ set(files
   TaskDispatcher.cpp
   ThreadPoolTaskDispatcher.cpp
   sps-ci/AllSPSCI.cpp
+  sps-ci/NativeDylibManagerSPSCI.cpp
   sps-ci/SimpleNativeMemoryMapSPSCI.cpp
   )
 

orc-rt/lib/executor/NativeDylibManager.cpp

@@ -0,0 +1,151 @@
+//===- NativeDylibManager.cpp ---------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// NativeDylibManager and related APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/NativeDylibManager.h"
+#include "orc-rt/Session.h"
+
+#include <sstream>
+
+#if defined(__APPLE__) || defined(__linux__)
+#include "Unix/NativeDylibAPIs.inc"
+#else
+#error "Target OS dylib APIs unsupported"
+#endif
+
+namespace orc_rt {
+
+Expected<std::unique_ptr<NativeDylibManager>>
+NativeDylibManager::Create(Session &S, SimpleSymbolTable &ST,
+                           const char *InstanceName,
+                           SimpleSymbolTable::MutatorFn AddInterface) {
+
+  std::unique_ptr<NativeDylibManager> Instance(new NativeDylibManager(S));
+
+  SimpleSymbolTable NDMST;
+  if (auto Err = AddInterface(NDMST))
+    return Err;
+  std::pair<const char *, const void *> InstanceSym[] = {
+      {InstanceName, static_cast<const void *>(Instance.get())}};
+  if (auto Err = NDMST.addUnique(InstanceSym))
+    return std::move(Err);
+
+  if (auto Err = ST.addUnique(NDMST))
+    return std::move(Err);
+
+  return std::move(Instance);
+}
+
+void NativeDylibManager::load(OnLoadCompleteFn &&OnComplete, std::string Path) {
+
+  if (auto H = hostOSLoadLibrary(Path)) {
+    {
+      std::scoped_lock<std::mutex> Lock(M);
+      auto &LI = LoadInfos[*H];
+      if (LI.Ordinal == 0) // new entry.
+        LI.Ordinal = LoadInfos.size();
+      ++LI.RefCount;
+    }
+    OnComplete(std::move(H));
+  } else
+    OnComplete(H.takeError());
+}
+
+void NativeDylibManager::unload(OnUnloadCompleteFn &&OnComplete, void *Handle) {
+  std::unique_lock<std::mutex> Lock(M);
+
+  auto LIItr = LoadInfos.find(Handle);
+  if (LIItr == LoadInfos.end()) {
+    Lock.unlock();
+    std::ostringstream ErrMsg;
+    ErrMsg << "error: attempt to unload unrecognized handle " << Handle;
+    OnComplete(make_error<StringError>(ErrMsg.str()));
+    return;
+  }
+
+  auto &LI = LIItr->second;
+
+  if (LI.RefCount == 0) {
+    Lock.unlock();
+    std::ostringstream ErrMsg;
+    ErrMsg << "error: cannot close handle " << Handle
+           << ", refcount is already zero";
+    OnComplete(make_error<StringError>(ErrMsg.str()));
+    return;
+  }
+
+  --LI.RefCount;
+
+  Lock.unlock();
+  OnComplete(hostOSUnloadLibrary(Handle));
+}
+
+void NativeDylibManager::lookup(OnLookupCompleteFn &&OnLookupComplete,
+                                void *Handle, std::vector<std::string> Names) {
+  {
+    std::unique_lock<std::mutex> Lock(M);
+    auto LIItr = LoadInfos.find(Handle);
+    if (LIItr == LoadInfos.end()) {
+      Lock.unlock();
+      std::ostringstream ErrMsg;
+      ErrMsg << "error: cannot perform lookup on unrecognized handle "
+             << Handle;
+      OnLookupComplete(make_error<StringError>(ErrMsg.str()));
+      return;
+    }
+
+    if (LIItr->second.RefCount == 0) {
+      Lock.unlock();
+      std::ostringstream ErrMsg;
+      ErrMsg << "error: cannot perform lookup on closed handle " << Handle;
+      OnLookupComplete(make_error<StringError>(ErrMsg.str()));
+      return;
+    }
+  }
+
+  OnLookupComplete(hostOSLibraryLookup(Handle, Names));
+}
+
+void NativeDylibManager::onDetach(Service::OnCompleteFn OnComplete,
+                                  bool ShutdownRequested) {
+  // Detach is a noop for now. If/when we add bloom-filter support this will be
+  // a good time to update filters.
+  OnComplete();
+}
+
+void NativeDylibManager::onShutdown(Service::OnCompleteFn OnComplete) {
+
+  // Unload in reverse load order (LIFO).
+  std::vector<void *> ToUnload;
+  ToUnload.reserve(LoadInfos.size());
+
+  for (auto &[Handle, Info] : LoadInfos)
+    ToUnload.push_back(Handle);
+
+  std::sort(ToUnload.begin(), ToUnload.end(), [this](void *LHS, void *RHS) {
+    assert(LoadInfos.count(LHS));
+    assert(LoadInfos.count(RHS));
+    return LoadInfos[LHS].Ordinal < LoadInfos[RHS].Ordinal;
+  });
+
+  while (!ToUnload.empty()) {
+    void *H = ToUnload.back();
+    ToUnload.pop_back();
+    size_t UnloadCount = LoadInfos[H].RefCount;
+    for (size_t I = 0; I != UnloadCount; ++I)
+      if (auto Err = hostOSUnloadLibrary(H))
+        S.reportError(std::move(Err));
+  }
+
+  OnComplete();
+}
+
+} // namespace orc_rt

orc-rt/lib/executor/Unix/NativeDylibAPIs.inc

@@ -0,0 +1,53 @@
+//===- NativeDylibAPIs.inc --------------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Generic wrappers for POSIX dlfcn.h APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/Error.h"
+
+#include <dlfcn.h>
+
+namespace {
+
+orc_rt::Expected<void *> hostOSLoadLibrary(const std::string &Path) {
+  void *H = dlopen(Path.c_str(), RTLD_LAZY | RTLD_LOCAL);
+  if (H == nullptr) {
+    std::ostringstream ErrMsg;
+    ErrMsg << "error loading ";
+    if (!Path.empty())
+      ErrMsg << "\"" << Path << "\"";
+    else
+      ErrMsg << "process symbols";
+    ErrMsg << ": " << dlerror();
+    return orc_rt::make_error<orc_rt::StringError>(ErrMsg.str());
+  }
+
+  return H;
+}
+
+orc_rt::Error hostOSUnloadLibrary(void *Handle) {
+  if (dlclose(Handle) != 0)
+    return orc_rt::make_error<orc_rt::StringError>(
+        (std::ostringstream()
+         << "error unloading " << Handle << ": " << dlerror())
+            .str());
+  return orc_rt::Error::success();
+}
+
+std::vector<void *> hostOSLibraryLookup(void *Handle,
+                                        const std::vector<std::string> &Names) {
+  std::vector<void *> Result;
+  Result.reserve(Names.size());
+  for (const auto &Name : Names)
+    Result.push_back(dlsym(Handle, Name.c_str()));
+  return Result;
+}
+
+} // anonymous namespace

orc-rt/lib/executor/sps-ci/NativeDylibManagerSPSCI.cpp

@@ -0,0 +1,46 @@
+//===- NativeDylibManagerSPSCI.cpp ----------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// SPS Controller Interface implementation for NativeDylibManager.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/sps-ci/NativeDylibManagerSPSCI.h"
+#include "orc-rt/NativeDylibManager.h"
+#include "orc-rt/SPSWrapperFunction.h"
+
+namespace orc_rt::sps_ci {
+
+ORC_RT_SPS_WRAPPER(
+    orc_rt_sps_ci_NativeDylibManager_load_sps_wrapper,
+    SPSExpected<SPSExecutorAddr>(SPSExecutorAddr, SPSString),
+    WrapperFunction::handleWithAsyncMethod(&NativeDylibManager::load))
+
+ORC_RT_SPS_WRAPPER(
+    orc_rt_sps_ci_NativeDylibManager_unload_sps_wrapper,
+    SPSError(SPSExecutorAddr, SPSExecutorAddr),
+    WrapperFunction::handleWithAsyncMethod(&NativeDylibManager::unload))
+
+ORC_RT_SPS_WRAPPER(
+    orc_rt_sps_ci_NativeDylibManager_lookup_sps_wrapper,
+    SPSExpected<SPSSequence<SPSExecutorAddr>>(SPSExecutorAddr, SPSExecutorAddr,
+                                              SPSSequence<SPSString>),
+    WrapperFunction::handleWithAsyncMethod(&NativeDylibManager::lookup))
+
+static std::pair<const char *, const void *>
+    orc_rt_sps_ci_NativeDylibManager_sps_interface[] = {
+        ORC_RT_SYMTAB_PAIR(orc_rt_sps_ci_NativeDylibManager_load_sps_wrapper),
+        ORC_RT_SYMTAB_PAIR(orc_rt_sps_ci_NativeDylibManager_unload_sps_wrapper),
+        ORC_RT_SYMTAB_PAIR(
+            orc_rt_sps_ci_NativeDylibManager_lookup_sps_wrapper)};
+
+Error addNativeDylibManager(SimpleSymbolTable &ST) {
+  return ST.addUnique(orc_rt_sps_ci_NativeDylibManager_sps_interface);
+}
+
+} // namespace orc_rt::sps_ci

orc-rt/unittests/CMakeLists.txt

@@ -28,6 +28,8 @@ add_orc_rt_unittest(CoreTests
   LockedAccessTest.cpp
   MathTest.cpp
   MemoryFlagsTest.cpp
+  NativeDylibManagerTest.cpp
+  NativeDylibManagerSPSCITest.cpp
   QueueingTaskDispatcherTest.cpp
   RTTITest.cpp
   ScopeExitTest.cpp
@@ -52,3 +54,12 @@ add_orc_rt_unittest(CoreTests
   )
 target_compile_options(CoreTests PRIVATE ${ORC_RT_COMPILE_FLAGS})
 target_link_libraries(CoreTests PRIVATE orc-rt-executor)
+
+# Build a shared library for NativeDylibManager tests.
+add_library(NativeDylibManagerTestLib SHARED
+  Inputs/NativeDylibManagerTestLib.cpp)
+set_target_properties(NativeDylibManagerTestLib PROPERTIES
+  PREFIX "")
+target_compile_definitions(CoreTests PRIVATE
+  "NDM_TEST_LIB_PATH=\"$<TARGET_FILE:NativeDylibManagerTestLib>\"")
+add_dependencies(CoreTests NativeDylibManagerTestLib)

orc-rt/unittests/Inputs/NativeDylibManagerTestLib.cpp

@@ -0,0 +1,10 @@
+// A minimal shared library for NativeDylibManager tests.
+
+#if defined(_WIN32)
+#define TEST_EXPORT __declspec(dllexport)
+#else
+#define TEST_EXPORT __attribute__((visibility("default")))
+#endif
+
+extern "C" TEST_EXPORT int NativeDylibManagerTestFunc() { return 42; }
+extern "C" TEST_EXPORT int NativeDylibManagerTestFunc2() { return 7; }

orc-rt/unittests/NativeDylibManagerSPSCITest.cpp

@@ -0,0 +1,183 @@
+//===- NativeDylibManagerSPSCITest.cpp ------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Tests for NativeDylibManager's SPS Controller Interface.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/sps-ci/NativeDylibManagerSPSCI.h"
+#include "orc-rt/NativeDylibManager.h"
+#include "orc-rt/SPSWrapperFunction.h"
+#include "orc-rt/Session.h"
+
+#include "CommonTestUtils.h"
+#include "DirectCaller.h"
+#include "gtest/gtest.h"
+
+using namespace orc_rt;
+
+#ifndef NDM_TEST_LIB_PATH
+#error                                                                         \
+    "NDM_TEST_LIB_PATH must be defined to the path of the test shared library"
+#endif
+
+class NativeDylibManagerSPSCITest : public ::testing::Test {
+protected:
+  void SetUp() override {
+    S = std::make_unique<Session>(mockExecutorProcessInfo(),
+                                  std::make_unique<NoDispatcher>(), noErrors);
+    NDM = cantFail(NativeDylibManager::Create(*S, CI));
+  }
+
+  void TearDown() override {
+    if (NDM) {
+      std::future<void> F;
+      NDM->onShutdown(waitFor(F));
+      F.get();
+    }
+  }
+
+  DirectCaller caller(const char *Name) {
+    return DirectCaller(nullptr, reinterpret_cast<orc_rt_WrapperFunction>(
+                                     const_cast<void *>(CI.at(Name))));
+  }
+
+  template <typename OnCompleteFn>
+  void spsLoad(OnCompleteFn &&OnComplete, std::string Path) {
+    using SPSSig = SPSExpected<SPSExecutorAddr>(SPSExecutorAddr, SPSString);
+    SPSWrapperFunction<SPSSig>::call(
+        caller("orc_rt_sps_ci_NativeDylibManager_load_sps_wrapper"),
+        std::forward<OnCompleteFn>(OnComplete), NDM.get(), std::move(Path));
+  }
+
+  template <typename OnCompleteFn>
+  void spsUnload(OnCompleteFn &&OnComplete, void *Handle) {
+    using SPSSig = SPSError(SPSExecutorAddr, SPSExecutorAddr);
+    SPSWrapperFunction<SPSSig>::call(
+        caller("orc_rt_sps_ci_NativeDylibManager_unload_sps_wrapper"),
+        std::forward<OnCompleteFn>(OnComplete), NDM.get(), Handle);
+  }
+
+  template <typename OnCompleteFn>
+  void spsLookup(OnCompleteFn &&OnComplete, void *Handle,
+                 std::vector<std::string> Names) {
+    using SPSSig = SPSExpected<SPSSequence<SPSExecutorAddr>>(
+        SPSExecutorAddr, SPSExecutorAddr, SPSSequence<SPSString>);
+    SPSWrapperFunction<SPSSig>::call(
+        caller("orc_rt_sps_ci_NativeDylibManager_lookup_sps_wrapper"),
+        std::forward<OnCompleteFn>(OnComplete), NDM.get(), Handle,
+        std::move(Names));
+  }
+
+  SimpleSymbolTable CI;
+  std::unique_ptr<Session> S;
+  std::unique_ptr<NativeDylibManager> NDM;
+};
+
+TEST_F(NativeDylibManagerSPSCITest, Registration) {
+  EXPECT_TRUE(CI.count("orc_rt_sps_ci_NativeDylibManager_load_sps_wrapper"));
+  EXPECT_TRUE(CI.count("orc_rt_sps_ci_NativeDylibManager_unload_sps_wrapper"));
+  EXPECT_TRUE(CI.count("orc_rt_sps_ci_NativeDylibManager_lookup_sps_wrapper"));
+}
+
+TEST_F(NativeDylibManagerSPSCITest, LoadAndUnload) {
+  std::future<Expected<Expected<void *>>> LoadResult;
+  spsLoad(waitFor(LoadResult), NDM_TEST_LIB_PATH);
+  void *Handle = cantFail(cantFail(LoadResult.get()));
+  EXPECT_NE(Handle, nullptr);
+
+  std::future<Expected<Error>> UnloadResult;
+  spsUnload(waitFor(UnloadResult), Handle);
+  cantFail(cantFail(UnloadResult.get()));
+}
+
+TEST_F(NativeDylibManagerSPSCITest, LoadNonExistent) {
+  std::future<Expected<Expected<void *>>> LoadResult;
+  spsLoad(waitFor(LoadResult), "/no/such/library.dylib");
+  auto Handle = cantFail(LoadResult.get());
+  EXPECT_FALSE(!!Handle);
+  consumeError(Handle.takeError());
+}
+
+TEST_F(NativeDylibManagerSPSCITest, UnloadUnrecognizedHandle) {
+  // Use Session object address as bogus handle.
+  void *BadHandle = reinterpret_cast<void *>(&S);
+  std::future<Expected<Error>> UnloadResult;
+  spsUnload(waitFor(UnloadResult), BadHandle);
+  auto Handle = cantFail(UnloadResult.get());
+  EXPECT_TRUE(!!Handle);
+  consumeError(std::move(Handle));
+}
+
+TEST_F(NativeDylibManagerSPSCITest, LookupSingleSymbol) {
+  std::future<Expected<Expected<void *>>> LoadResult;
+  spsLoad(waitFor(LoadResult), NDM_TEST_LIB_PATH);
+  void *Handle = cantFail(cantFail(LoadResult.get()));
+
+  std::future<Expected<Expected<std::vector<void *>>>> LookupResult;
+  spsLookup(waitFor(LookupResult), Handle, {"NativeDylibManagerTestFunc"});
+  auto Addrs = cantFail(cantFail(LookupResult.get()));
+  ASSERT_EQ(Addrs.size(), 1U);
+  EXPECT_NE(Addrs[0], nullptr);
+
+  auto *Func = reinterpret_cast<int (*)()>(Addrs[0]);
+  EXPECT_EQ(Func(), 42);
+
+  std::future<Expected<Error>> UnloadResult;
+  spsUnload(waitFor(UnloadResult), Handle);
+  cantFail(cantFail(UnloadResult.get()));
+}
+
+TEST_F(NativeDylibManagerSPSCITest, LookupMultipleSymbols) {
+  std::future<Expected<Expected<void *>>> LoadResult;
+  spsLoad(waitFor(LoadResult), NDM_TEST_LIB_PATH);
+  void *Handle = cantFail(cantFail(LoadResult.get()));
+
+  std::future<Expected<Expected<std::vector<void *>>>> LookupResult;
+  spsLookup(waitFor(LookupResult), Handle,
+            {"NativeDylibManagerTestFunc", "NativeDylibManagerTestFunc2"});
+  auto Addrs = cantFail(cantFail(LookupResult.get()));
+  ASSERT_EQ(Addrs.size(), 2U);
+  EXPECT_NE(Addrs[0], nullptr);
+  EXPECT_NE(Addrs[1], nullptr);
+
+  auto *Func1 = reinterpret_cast<int (*)()>(Addrs[0]);
+  auto *Func2 = reinterpret_cast<int (*)()>(Addrs[1]);
+  EXPECT_EQ(Func1(), 42);
+  EXPECT_EQ(Func2(), 7);
+
+  std::future<Expected<Error>> UnloadResult;
+  spsUnload(waitFor(UnloadResult), Handle);
+  cantFail(cantFail(UnloadResult.get()));
+}
+
+TEST_F(NativeDylibManagerSPSCITest, LookupNonExistentSymbol) {
+  std::future<Expected<Expected<void *>>> LoadResult;
+  spsLoad(waitFor(LoadResult), NDM_TEST_LIB_PATH);
+  void *Handle = cantFail(cantFail(LoadResult.get()));
+
+  std::future<Expected<Expected<std::vector<void *>>>> LookupResult;
+  spsLookup(waitFor(LookupResult), Handle, {"no_such_symbol"});
+  auto Addrs = cantFail(cantFail(LookupResult.get()));
+  ASSERT_EQ(Addrs.size(), 1U);
+  EXPECT_EQ(Addrs[0], nullptr);
+
+  std::future<Expected<Error>> UnloadResult;
+  spsUnload(waitFor(UnloadResult), Handle);
+  cantFail(cantFail(UnloadResult.get()));
+}
+
+TEST_F(NativeDylibManagerSPSCITest, LookupOnUnrecognizedHandle) {
+  // Use Session object address as bogus handle.
+  void *BadHandle = reinterpret_cast<void *>(&S);
+  std::future<Expected<Expected<std::vector<void *>>>> LookupResult;
+  spsLookup(waitFor(LookupResult), BadHandle, {"NativeDylibManagerTestFunc"});
+  auto Addrs = cantFail(LookupResult.get());
+  EXPECT_FALSE(!!Addrs);
+  consumeError(Addrs.takeError());
+}

orc-rt/unittests/NativeDylibManagerTest.cpp

@@ -0,0 +1,199 @@
+//===- NativeDylibManagerTest.cpp -----------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Test NativeDylibManager APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/NativeDylibManager.h"
+#include "orc-rt/Session.h"
+
+#include "gtest/gtest.h"
+
+#include "CommonTestUtils.h"
+
+#include <optional>
+
+using namespace orc_rt;
+
+#ifndef NDM_TEST_LIB_PATH
+#error                                                                         \
+    "NDM_TEST_LIB_PATH must be defined to the path of the test shared library"
+#endif
+
+// Helper: synchronously run load and return result.
+static Expected<void *> syncLoad(NativeDylibManager &NDM, std::string Path) {
+  std::optional<Expected<void *>> Result;
+  NDM.load([&](Expected<void *> R) { Result = std::move(R); }, std::move(Path));
+  return std::move(*Result);
+}
+
+// Helper: synchronously run unload and return result.
+static Error syncUnload(NativeDylibManager &NDM, void *Handle) {
+  std::optional<Error> Result;
+  NDM.unload([&](Error R) { Result = std::move(R); }, Handle);
+  return std::move(*Result);
+}
+
+// Helper: synchronously run lookup and return results.
+static Expected<std::vector<void *>>
+syncLookup(NativeDylibManager &NDM, void *Handle,
+           std::vector<std::string> Names) {
+  std::optional<Expected<std::vector<void *>>> Result;
+  NDM.lookup([&](Expected<std::vector<void *>> R) { Result = std::move(R); },
+             Handle, std::move(Names));
+  return std::move(*Result);
+}
+
+TEST(NativeDylibManagerTest, Create) {
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = NativeDylibManager::Create(S, ST);
+  ASSERT_TRUE(!!NDM) << toString(NDM.takeError());
+}
+
+TEST(NativeDylibManagerTest, LoadAndUnload) {
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = cantFail(NativeDylibManager::Create(S, ST));
+
+  // Load the test library.
+  auto LoadResult = syncLoad(*NDM, NDM_TEST_LIB_PATH);
+  ASSERT_TRUE(!!LoadResult) << toString(LoadResult.takeError());
+  void *Handle = *LoadResult;
+  EXPECT_NE(Handle, nullptr);
+
+  // Unload it.
+  auto UnloadResult = syncUnload(*NDM, Handle);
+  EXPECT_FALSE(!!UnloadResult) << toString(std::move(UnloadResult));
+}
+
+TEST(NativeDylibManagerTest, LoadNonExistent) {
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = cantFail(NativeDylibManager::Create(S, ST));
+
+  auto LoadResult = syncLoad(*NDM, "/no/such/library.dylib");
+  EXPECT_FALSE(!!LoadResult);
+  consumeError(LoadResult.takeError());
+}
+
+TEST(NativeDylibManagerTest, UnloadUnrecognizedHandle) {
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = cantFail(NativeDylibManager::Create(S, ST));
+
+  void *Bogus = reinterpret_cast<void *>(0xDEADBEEF);
+  auto UnloadResult = syncUnload(*NDM, Bogus);
+  EXPECT_TRUE(!!UnloadResult);
+  consumeError(std::move(UnloadResult));
+}
+
+TEST(NativeDylibManagerTest, LoadSameLibraryTwice) {
+  // Loading the same library twice should succeed both times and return
+  // the same handle (since dlopen refcounts).
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = cantFail(NativeDylibManager::Create(S, ST));
+
+  auto R1 = syncLoad(*NDM, NDM_TEST_LIB_PATH);
+  ASSERT_TRUE(!!R1) << toString(R1.takeError());
+  void *H1 = *R1;
+
+  auto R2 = syncLoad(*NDM, NDM_TEST_LIB_PATH);
+  ASSERT_TRUE(!!R2) << toString(R2.takeError());
+  void *H2 = *R2;
+
+  EXPECT_EQ(H1, H2);
+
+  // Unload both references.
+  auto UR1 = syncUnload(*NDM, H1);
+  EXPECT_FALSE(!!UR1) << toString(std::move(UR1));
+  auto UR2 = syncUnload(*NDM, H2);
+  EXPECT_FALSE(!!UR2) << toString(std::move(UR2));
+}
+
+TEST(NativeDylibManagerTest, LookupSingleSymbol) {
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = cantFail(NativeDylibManager::Create(S, ST));
+
+  void *Handle = cantFail(syncLoad(*NDM, NDM_TEST_LIB_PATH));
+
+  auto Result = syncLookup(*NDM, Handle, {"NativeDylibManagerTestFunc"});
+  ASSERT_TRUE(!!Result) << toString(Result.takeError());
+  ASSERT_EQ(Result->size(), 1U);
+  EXPECT_NE((*Result)[0], nullptr);
+
+  // Verify the symbol points to the right function.
+  auto *Func = reinterpret_cast<int (*)()>((*Result)[0]);
+  EXPECT_EQ(Func(), 42);
+
+  auto UR = syncUnload(*NDM, Handle);
+  EXPECT_FALSE(!!UR) << toString(std::move(UR));
+}
+
+TEST(NativeDylibManagerTest, LookupMultipleSymbols) {
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = cantFail(NativeDylibManager::Create(S, ST));
+
+  void *Handle = cantFail(syncLoad(*NDM, NDM_TEST_LIB_PATH));
+
+  auto Result =
+      syncLookup(*NDM, Handle,
+                 {"NativeDylibManagerTestFunc", "NativeDylibManagerTestFunc2"});
+  ASSERT_TRUE(!!Result) << toString(Result.takeError());
+  ASSERT_EQ(Result->size(), 2U);
+  EXPECT_NE((*Result)[0], nullptr);
+  EXPECT_NE((*Result)[1], nullptr);
+
+  auto *Func1 = reinterpret_cast<int (*)()>((*Result)[0]);
+  auto *Func2 = reinterpret_cast<int (*)()>((*Result)[1]);
+  EXPECT_EQ(Func1(), 42);
+  EXPECT_EQ(Func2(), 7);
+
+  auto UR = syncUnload(*NDM, Handle);
+  EXPECT_FALSE(!!UR) << toString(std::move(UR));
+}
+
+TEST(NativeDylibManagerTest, LookupNonExistentSymbol) {
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = cantFail(NativeDylibManager::Create(S, ST));
+
+  void *Handle = cantFail(syncLoad(*NDM, NDM_TEST_LIB_PATH));
+
+  auto Result = syncLookup(*NDM, Handle, {"no_such_symbol"});
+  ASSERT_TRUE(!!Result) << toString(Result.takeError());
+  ASSERT_EQ(Result->size(), 1U);
+  EXPECT_EQ((*Result)[0], nullptr);
+
+  auto UR = syncUnload(*NDM, Handle);
+  EXPECT_FALSE(!!UR) << toString(std::move(UR));
+}
+
+TEST(NativeDylibManagerTest, LookupOnUnrecognizedHandle) {
+  Session S(mockExecutorProcessInfo(), std::make_unique<NoDispatcher>(),
+            noErrors);
+  SimpleSymbolTable ST;
+  auto NDM = cantFail(NativeDylibManager::Create(S, ST));
+
+  void *Bogus = reinterpret_cast<void *>(0xDEADBEEF);
+  auto Result = syncLookup(*NDM, Bogus, {"NativeDylibManagerTestFunc"});
+  EXPECT_FALSE(!!Result);
+  consumeError(Result.takeError());
+}