[Orc] Support atexit in Orc(JITLink)

[StephenFan]

Summary

There is a bug reported at https://bugs.llvm.org/show_bug.cgi?id=48938

After looking through the glibc, I found the atexit(f) is the same as __cxa_atexit(f, NULL, NULL). In orc runtime, we identify different JITDylib by their dso_handle value, so that a NULL dso_handle is invalid. So in this patch, I added a PlatformJDDSOHandle to ELFNixRuntimeState, and functions which are registered by atexit will be registered at PlatformJD.

[StephenFan]

Add declaration.

[lhames]

This is a nice solution! LGTM.

compiler-rt/lib/orc/elfnix_platform.cpp

 //===- elfnix_platform.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
 //
 //===----------------------------------------------------------------------===//
 //
private:
   struct PerJITDylibState {
     void *Header = nullptr;
     size_t RefCount = 0;
     bool AllowReinitialization = false;
     AtExitsVector AtExits;
   };
 
 public:
-  static void initialize();
+  static void initialize(void *DSOHandle);
   static ELFNixPlatformRuntimeState &get();
   static void destroy();
 
-  ELFNixPlatformRuntimeState() = default;
+  ELFNixPlatformRuntimeState(void *DSOHandle) : PlatformJDDSOHandle(DSOHandle) {}
 
   // Delete copy and move constructors.
   ELFNixPlatformRuntimeState(const ELFNixPlatformRuntimeState &) = delete;
   ELFNixPlatformRuntimeState &
   operator=(const ELFNixPlatformRuntimeState &) = delete;
   ELFNixPlatformRuntimeState(ELFNixPlatformRuntimeState &&) = delete;
   ELFNixPlatformRuntimeState &operator=(ELFNixPlatformRuntimeState &&) = delete;
 
   Error registerObjectSections(ELFNixPerObjectSectionsToRegister POSR);
   Error deregisterObjectSections(ELFNixPerObjectSectionsToRegister POSR);
 
   const char *dlerror();
   void *dlopen(string_view Name, int Mode);
   int dlclose(void *DSOHandle);
   void *dlsym(void *DSOHandle, string_view Symbol);
 
   int registerAtExit(void (*F)(void *), void *Arg, void *DSOHandle);
   void runAtExits(void *DSOHandle);
 
   /// Returns the base address of the section containing ThreadData.
   Expected<std::pair<const char *, size_t>>
   getThreadDataSectionFor(const char *ThreadData);
 
+  void *getPlatformJDDSOHandle() { return PlatformJDDSOHandle; }
+
 private:
   PerJITDylibState *getJITDylibStateByHeaderAddr(void *DSOHandle);
   PerJITDylibState *getJITDylibStateByName(string_view Path);
   PerJITDylibState &
   getOrCreateJITDylibState(ELFNixJITDylibInitializers &MOJDIs);
 
   Error registerThreadDataSection(span<const char> ThreadDataSection);
 
   Expected<ExecutorAddr> lookupSymbolInJITDylib(void *DSOHandle,
                                                 string_view Symbol);
 
   Expected<ELFNixJITDylibInitializerSequence>
   getJITDylibInitializersByName(string_view Path);
   Expected<void *> dlopenInitialize(string_view Path, int Mode);
   Error initializeJITDylib(ELFNixJITDylibInitializers &MOJDIs);
 
   static ELFNixPlatformRuntimeState *MOPS;
 
   using InitSectionHandler =
       Error (*)(const std::vector<ExecutorAddrRange> &Sections,
                 const ELFNixJITDylibInitializers &MOJDIs);
   const std::vector<std::pair<const char *, InitSectionHandler>> InitSections =
       {{".init_array", runInitArray}};
 
+  void *PlatformJDDSOHandle;
+
   // FIXME: Move to thread-state.
   std::string DLFcnError;
 
   std::recursive_mutex JDStatesMutex;
   std::unordered_map<void *, PerJITDylibState> JDStates;
   std::unordered_map<std::string, void *> JDNameToHeader;
 
   std::mutex ThreadDataSectionsMutex;
   std::map<const char *, size_t> ThreadDataSections;
 };
 
 ELFNixPlatformRuntimeState *ELFNixPlatformRuntimeState::MOPS = nullptr;
 
-void ELFNixPlatformRuntimeState::initialize() {
+void ELFNixPlatformRuntimeState::initialize(void *DSOHandle) {
   assert(!MOPS && "ELFNixPlatformRuntimeState should be null");
-  MOPS = new ELFNixPlatformRuntimeState();
+  MOPS = new ELFNixPlatformRuntimeState(DSOHandle);
 }
 
 ELFNixPlatformRuntimeState &ELFNixPlatformRuntimeState::get() {
   assert(MOPS && "ELFNixPlatformRuntimeState not initialized");
   return *MOPS;
 }
 
 void ELFNixPlatformRuntimeState::destroy() {

 } // end anonymous namespace
 
 //------------------------------------------------------------------------------
 //                             JIT entry points
 //------------------------------------------------------------------------------
 
 ORC_RT_INTERFACE __orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_platform_bootstrap(char *ArgData, size_t ArgSize) {
-  ELFNixPlatformRuntimeState::initialize();
-  return WrapperFunctionResult().release();
+  return WrapperFunction<void(uint64_t)>::handle(
+      ArgData, ArgSize, [](uint64_t &DSOHandle) {
+        ELFNixPlatformRuntimeState::initialize(reinterpret_cast<void *>(DSOHandle));
+      }).release();
 }
 
 ORC_RT_INTERFACE __orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_platform_shutdown(char *ArgData, size_t ArgSize) {
   ELFNixPlatformRuntimeState::destroy();
   return WrapperFunctionResult().release();
 }
 

 //------------------------------------------------------------------------------
 
 int __orc_rt_elfnix_cxa_atexit(void (*func)(void *), void *arg,
                                void *dso_handle) {
   return ELFNixPlatformRuntimeState::get().registerAtExit(func, arg,
                                                           dso_handle);
 }
 
+int __orc_rt_elfnix_atexit(void (*func)(void *)) {
+  auto &PlatformRTState = ELFNixPlatformRuntimeState::get();
+  return ELFNixPlatformRuntimeState::get().registerAtExit(func, NULL, PlatformRTState.getPlatformJDDSOHandle());
+}
+
 void __orc_rt_elfnix_cxa_finalize(void *dso_handle) {
   ELFNixPlatformRuntimeState::get().runAtExits(dso_handle);
 }
 
 //------------------------------------------------------------------------------
 //                        JIT'd dlfcn alternatives.
 //------------------------------------------------------------------------------
 

compiler-rt/test/orc/TestCases/Linux/x86-64/trivial-atexit.S

+// Test that the runtime correctly interposes atexit.
+//
+// RUN: %clang -c -o %t %s
+// RUN: %llvm_jitlink %t
+
+  .text
+// OnExit destructor resets the test result override to zero.
+  .section  .text._ZN6OnExitD2Ev,"axG",@progbits,_ZN6OnExitD2Ev,comdat
+  .p2align  4, 0x90
+  .type _ZN6OnExitD2Ev,@function
+_ZN6OnExitD2Ev:                         # @_ZN6OnExitD2Ev
+  .cfi_startproc
+  xorl  %edi, %edi
+  jmp llvm_jitlink_setTestResultOverride@PLT # TAILCALL
+  .cfi_endproc
+
+// main registers the atexit and sets the test result to one.
+  .globl main
+  .p2align  4, 0x90                         # -- Begin function main
+  .type main,@function
+main:              # @main
+  .cfi_startproc
+# %bb.0:
+  movq  _ZN6OnExitD2Ev@GOTPCREL(%rip), %rdi
+  callq atexit@PLT
+  movl  $1, %edi
+  callq llvm_jitlink_setTestResultOverride@PLT
+  xorl  %eax, %eax
+  retq
+.Lfunc_end1:
+  .size main, .Lfunc_end1-main
+  .cfi_endproc
+                                        # -- End function
+  .type _ZL6onExit,@object              # @_ZL6onExit
+  .local  _ZL6onExit
+  .comm _ZL6onExit,1,1

llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp

 //===------ ELFNixPlatform.cpp - Utilities for executing MachO in Orc -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
SymbolAliasMap ELFNixPlatform::standardPlatformAliases(ExecutionSession &ES) {
   addAliases(ES, Aliases, requiredCXXAliases());
   addAliases(ES, Aliases, standardRuntimeUtilityAliases());
   return Aliases;
 }
 
 ArrayRef<std::pair<const char *, const char *>>
 ELFNixPlatform::requiredCXXAliases() {
   static const std::pair<const char *, const char *> RequiredCXXAliases[] = {
-      {"__cxa_atexit", "__orc_rt_elfnix_cxa_atexit"}};
+    {"__cxa_atexit", "__orc_rt_elfnix_cxa_atexit"},
+    {"atexit", "__orc_rt_elfnix_atexit"}};
 
   return ArrayRef<std::pair<const char *, const char *>>(RequiredCXXAliases);
 }
 
 ArrayRef<std::pair<const char *, const char *>>
 ELFNixPlatform::standardRuntimeUtilityAliases() {
   static const std::pair<const char *, const char *>
       StandardRuntimeUtilityAliases[] = {
if (!RuntimeSymbolAddrs)
     return RuntimeSymbolAddrs.takeError();
 
   for (const auto &KV : AddrsToRecord) {
     auto &Name = KV.first;
     assert(RuntimeSymbolAddrs->count(Name) && "Missing runtime symbol?");
     KV.second->setValue((*RuntimeSymbolAddrs)[Name].getAddress());
   }
 
-  if (auto Err = ES.callSPSWrapper<void()>(orc_rt_elfnix_platform_bootstrap))
+  auto PJDDSOHandle = ES.lookup(
+      {{&PlatformJD, JITDylibLookupFlags::MatchAllSymbols}}, DSOHandleSymbol);
+  if (!PJDDSOHandle)
+    return PJDDSOHandle.takeError();
+
+  if (auto Err = ES.callSPSWrapper<void(uint64_t)>(orc_rt_elfnix_platform_bootstrap, PJDDSOHandle->getAddress()))
     return Err;
 
   // FIXME: Ordering is fuzzy here. We're probably best off saying
   // "behavior is undefined if code that uses the runtime is added before
   // the platform constructor returns", then move all this to the constructor.
   RuntimeBootstrapped = true;
   std::vector<ELFPerObjectSectionsToRegister> DeferredPOSRs;
   {