[Orc] Support atexit in Orc(JITLink)

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.

Reviewed By: lhames

Differential Revision: https://reviews.llvm.org/D111413

4 files changed, 68 insertions, 8 deletions

diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index fae93b4..0352f6c 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -84,11 +84,12 @@ private:
   };
 
 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;
@@ -112,6 +113,8 @@ public:
   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);
@@ -136,6 +139,8 @@ private:
   const std::vector<std::pair<const char *, InitSectionHandler>> InitSections =
       {{".init_array", runInitArray}};
 
+  void *PlatformJDDSOHandle;
+
   // FIXME: Move to thread-state.
   std::string DLFcnError;
 
@@ -149,9 +154,9 @@ private:
 
 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() {
@@ -434,8 +439,13 @@ void destroyELFNixTLVMgr(void *ELFNixTLVMgr) {
 
 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
@@ -511,6 +521,12 @@ int __orc_rt_elfnix_cxa_atexit(void (*func)(void *), void *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);
 }
diff --git a/compiler-rt/lib/orc/elfnix_platform.h b/compiler-rt/lib/orc/elfnix_platform.h
index 5afeac5..12b9591 100644
--- a/compiler-rt/lib/orc/elfnix_platform.h
+++ b/compiler-rt/lib/orc/elfnix_platform.h
@@ -19,6 +19,7 @@
 // Atexit functions.
 ORC_RT_INTERFACE int __orc_rt_elfnix_cxa_atexit(void (*func)(void *), void *arg,
                                                 void *dso_handle);
+ORC_RT_INTERFACE int __orc_rt_elfnix_atexit(void (*func)(void *));
 ORC_RT_INTERFACE void __orc_rt_elfnix_cxa_finalize(void *dso_handle);
 
 // dlfcn functions.
diff --git a/compiler-rt/test/orc/TestCases/Linux/x86-64/trivial-atexit.S b/compiler-rt/test/orc/TestCases/Linux/x86-64/trivial-atexit.S
new file mode 100644
index 0000000..16852c6
--- /dev/null
+++ b/compiler-rt/test/orc/TestCases/Linux/x86-64/trivial-atexit.S
@@ -0,0 +1,36 @@
+// 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
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index 94499de..79c602d 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -194,7 +194,8 @@ SymbolAliasMap ELFNixPlatform::standardPlatformAliases(ExecutionSession &ES) {
 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);
 }
@@ -474,7 +475,13 @@ Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
     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