[libc] at_quick_exit function implemented (#94317)

- added at_quick_exit function 
- used helper file exit_handler which reuses code from atexit
- atexit now calls helper functions from exit_handler
- test cases and dependencies are added

---------

Co-authored-by: Aaryan Shukla <[email protected]>

Author: aaryanshukla

libc/config/linux/x86_64/entrypoints.txt

@@ -173,7 +173,6 @@ set(TARGET_LIBC_ENTRYPOINTS
     libc.src.stdlib.atoll
     libc.src.stdlib.bsearch
     libc.src.stdlib.div
-    libc.src.stdlib.quick_exit
     libc.src.stdlib.labs
     libc.src.stdlib.ldiv
     libc.src.stdlib.llabs
@@ -770,9 +769,11 @@ if(LLVM_LIBC_FULL_BUILD)
     # stdlib.h entrypoints
     libc.src.stdlib._Exit
     libc.src.stdlib.abort
+    libc.src.stdlib.at_quick_exit
     libc.src.stdlib.atexit
     libc.src.stdlib.exit
     libc.src.stdlib.getenv
+    libc.src.stdlib.quick_exit
 
     # signal.h entrypoints
     libc.src.signal.raise

libc/hdr/types/CMakeLists.txt

@@ -117,3 +117,12 @@ add_proxy_header_library(
     libc.include.llvm-libc-types.pid_t
     libc.include.sys_types
 )
+
+add_proxy_header_library(
+  atexithandler_t
+  HDRS
+    atexithandler_t.h
+  FULL_BUILD_DEPENDS
+    libc.include.llvm-libc-types.atexithandler_t
+    libc.include.stdlib
+)

libc/hdr/types/atexithandler_t.h

@@ -0,0 +1,22 @@
+//===-- Definition of macros from atexithandler_t.h -----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_HDR_ATEXITHANDLER_T_H
+#define LLVM_LIBC_HDR_ATEXITHANDLER_T_H
+
+#ifdef LIBC_FULL_BUILD
+
+#include "include/llvm-libc-types/__atexithandler_t.h"
+
+#else // overlay mode
+
+#error // type not available in overlay mode
+
+#endif // LLVM_LIBC_FULL_BUILD
+
+#endif // LLVM_LIBC_HDR_ATEXITHANDLER_T_H

libc/spec/stdc.td

@@ -1107,8 +1107,9 @@ def StdC : StandardSpec<"stdc"> {
           FunctionSpec<"free", RetValSpec<VoidType>, [ArgSpec<VoidPtr>]>,
 
           FunctionSpec<"_Exit", RetValSpec<NoReturn>, [ArgSpec<IntType>]>,
-          FunctionSpec<"exit", RetValSpec<NoReturn>, [ArgSpec<IntType>]>,
+          FunctionSpec<"at_quick_exit", RetValSpec<IntType>, [ArgSpec<AtexitHandlerT>]>,
           FunctionSpec<"atexit", RetValSpec<IntType>, [ArgSpec<AtexitHandlerT>]>,
+          FunctionSpec<"exit", RetValSpec<NoReturn>, [ArgSpec<IntType>]>,
           FunctionSpec<"quick_exit", RetValSpec<NoReturn>, [ArgSpec<IntType>]>,
       ]
   >;

libc/src/stdlib/CMakeLists.txt

@@ -50,6 +50,7 @@ add_entrypoint_object(
     quick_exit.h
   DEPENDS
     libc.src.__support.OSUtil.osutil
+    .exit_handler
 )
 
 add_entrypoint_object(
@@ -415,14 +416,14 @@ add_entrypoint_object(
     libc.src.__support.OSUtil.osutil
 )
 
-add_entrypoint_object(
-  atexit
+add_object_library(
+  exit_handler
   SRCS
-    atexit.cpp
+    exit_handler.cpp
   HDRS
-    atexit.h
+    exit_handler.h
   CXX_STANDARD
-    20 # For constinit of the atexit callback list.
+    20 # For constinit 
   DEPENDS
     libc.src.__support.CPP.mutex
     libc.src.__support.CPP.new
@@ -432,6 +433,26 @@ add_entrypoint_object(
     libc.src.__support.threads.mutex
 )
 
+add_entrypoint_object(
+  atexit
+  SRCS
+    atexit.cpp
+  HDRS
+    atexit.h
+  DEPENDS
+    .exit_handler
+)
+
+add_entrypoint_object(
+  at_quick_exit
+  SRCS
+    at_quick_exit.cpp
+  HDRS
+    at_quick_exit.h
+  DEPENDS
+    .exit_handler
+)
+
 add_entrypoint_object(
   exit
   SRCS
@@ -442,6 +463,7 @@ add_entrypoint_object(
     ._Exit
     .atexit
     libc.src.__support.OSUtil.osutil
+    .exit_handler
 )
 
 add_entrypoint_object(

libc/src/stdlib/at_quick_exit.cpp

@@ -0,0 +1,22 @@
+//===-- Implementation of at_quick_exit -----------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/stdlib/at_quick_exit.h"
+#include "hdr/types/atexithandler_t.h"
+#include "src/__support/common.h"
+#include "src/stdlib/exit_handler.h"
+
+namespace LIBC_NAMESPACE {
+
+LLVM_LIBC_FUNCTION(int, at_quick_exit, (__atexithandler_t callback)) {
+  return add_atexit_unit(
+      at_quick_exit_callbacks,
+      {&stdc_at_exit_func, reinterpret_cast<void *>(callback)});
+}
+
+} // namespace LIBC_NAMESPACE

libc/src/stdlib/at_quick_exit.h

@@ -0,0 +1,20 @@
+//===-- Implementation header for at_quick_exit -----------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC_STDLIB_AT_QUICK_EXIT_H
+#define LLVM_LIBC_SRC_STDLIB_AT_QUICK_EXIT_H
+
+#include "hdr/types/atexithandler_t.h"
+
+namespace LIBC_NAMESPACE {
+
+int at_quick_exit(__atexithandler_t);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_STDLIB_AT_QUICK_EXIT_H

libc/src/stdlib/atexit.cpp

@@ -7,95 +7,28 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/stdlib/atexit.h"
-#include "src/__support/CPP/mutex.h" // lock_guard
-#include "src/__support/blockstore.h"
+#include "hdr/types/atexithandler_t.h"
 #include "src/__support/common.h"
-#include "src/__support/fixedvector.h"
-#include "src/__support/threads/mutex.h"
+#include "src/stdlib/exit_handler.h"
 
 namespace LIBC_NAMESPACE {
 
-namespace {
-
-Mutex handler_list_mtx(/*timed=*/false, /*recursive=*/false, /*robust=*/false,
-                       /*pshared=*/false);
-
-using AtExitCallback = void(void *);
-using StdCAtExitCallback = void(void);
-
-struct AtExitUnit {
-  AtExitCallback *callback = nullptr;
-  void *payload = nullptr;
-  constexpr AtExitUnit() = default;
-  constexpr AtExitUnit(AtExitCallback *c, void *p) : callback(c), payload(p) {}
-};
-
-#if defined(LIBC_TARGET_ARCH_IS_GPU)
-// The GPU build cannot handle the potentially recursive definitions required by
-// the BlockStore class. Additionally, the liklihood that someone exceeds this
-// while executing on the GPU is extremely small.
-// FIXME: It is not generally safe to use 'atexit' on the GPU because the
-//        mutexes simply passthrough. We will need a lock free stack.
-using ExitCallbackList = FixedVector<AtExitUnit, 64>;
-#elif defined(LIBC_COPT_PUBLIC_PACKAGING)
-using ExitCallbackList = ReverseOrderBlockStore<AtExitUnit, 32>;
-#else
-// BlockStore uses dynamic memory allocation. To avoid dynamic memory
-// allocation in tests, we use a fixed size callback list when built for
-// tests.
-// If we use BlockStore, then we will have to pull in malloc etc into
-// the tests. While this is not bad, the problem we have currently is
-// that LLVM libc' allocator is SCUDO. So, we will end up pulling SCUDO's
-// deps also (some of which are not yet available in LLVM libc) into the
-// integration tests.
-using ExitCallbackList = FixedVector<AtExitUnit, CALLBACK_LIST_SIZE_FOR_TESTS>;
-#endif // LIBC_COPT_PUBLIC_PACKAGING
-
-constinit ExitCallbackList exit_callbacks;
-
-void stdc_at_exit_func(void *payload) {
-  reinterpret_cast<StdCAtExitCallback *>(payload)();
-}
-
-void call_exit_callbacks() {
-  handler_list_mtx.lock();
-  while (!exit_callbacks.empty()) {
-    AtExitUnit &unit = exit_callbacks.back();
-    exit_callbacks.pop_back();
-    handler_list_mtx.unlock();
-    unit.callback(unit.payload);
-    handler_list_mtx.lock();
-  }
-  ExitCallbackList::destroy(&exit_callbacks);
-}
-
-int add_atexit_unit(const AtExitUnit &unit) {
-  cpp::lock_guard lock(handler_list_mtx);
-  if (exit_callbacks.push_back(unit))
-    return 0;
-  return -1;
-}
-
-} // namespace
-
 extern "C" {
 
-// TODO: Handle the last dso handle argument.
 int __cxa_atexit(AtExitCallback *callback, void *payload, void *) {
-  return add_atexit_unit({callback, payload});
+  return add_atexit_unit(atexit_callbacks, {callback, payload});
 }
 
-// TODO: Handle the dso handle argument. call_exit_callbacks should only invoke
-// the callbacks from this DSO. Requires adding support for __dso_handle.
 void __cxa_finalize(void *dso) {
   if (!dso)
-    call_exit_callbacks();
+    call_exit_callbacks(atexit_callbacks);
 }
 
 } // extern "C"
 
-LLVM_LIBC_FUNCTION(int, atexit, (StdCAtExitCallback * callback)) {
+LLVM_LIBC_FUNCTION(int, atexit, (__atexithandler_t callback)) {
   return add_atexit_unit(
+      atexit_callbacks,
       {&stdc_at_exit_func, reinterpret_cast<void *>(callback)});
 }
 

libc/src/stdlib/atexit.h

@@ -9,13 +9,10 @@
 #ifndef LLVM_LIBC_SRC_STDLIB_ATEXIT_H
 #define LLVM_LIBC_SRC_STDLIB_ATEXIT_H
 
-#include <stddef.h> // For size_t
-
+#include "hdr/types/atexithandler_t.h"
 namespace LIBC_NAMESPACE {
 
-constexpr size_t CALLBACK_LIST_SIZE_FOR_TESTS = 1024;
-
-int atexit(void (*function)());
+int atexit(__atexithandler_t);
 
 } // namespace LIBC_NAMESPACE
 

libc/src/stdlib/exit_handler.cpp

@@ -0,0 +1,42 @@
+//===--- Implementation of exit_handler------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "src/stdlib/exit_handler.h"
+#include "src/__support/CPP/mutex.h" // lock_guard
+
+namespace LIBC_NAMESPACE {
+
+constinit ExitCallbackList at_quick_exit_callbacks;
+constinit ExitCallbackList atexit_callbacks;
+
+Mutex handler_list_mtx(false, false, false, false);
+
+void stdc_at_exit_func(void *payload) {
+  reinterpret_cast<StdCAtExitCallback *>(payload)();
+}
+
+void call_exit_callbacks(ExitCallbackList &callbacks) {
+  handler_list_mtx.lock();
+  while (!callbacks.empty()) {
+    AtExitUnit &unit = callbacks.back();
+    callbacks.pop_back();
+    handler_list_mtx.unlock();
+    unit.callback(unit.payload);
+    handler_list_mtx.lock();
+  }
+  ExitCallbackList::destroy(&callbacks);
+}
+
+int add_atexit_unit(ExitCallbackList &callbacks, const AtExitUnit &unit) {
+  cpp::lock_guard lock(handler_list_mtx);
+  if (callbacks.push_back(unit))
+    return 0;
+  return -1;
+}
+
+} // namespace LIBC_NAMESPACE

libc/src/stdlib/exit_handler.h

@@ -0,0 +1,53 @@
+//===-- Implementation header for exit_handler ------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC_STDLIB_EXIT_HANDLER_H
+#define LLVM_LIBC_SRC_STDLIB_EXIT_HANDLER_H
+
+#include "src/__support/CPP/mutex.h" // lock_guard
+#include "src/__support/blockstore.h"
+#include "src/__support/common.h"
+#include "src/__support/fixedvector.h"
+#include "src/__support/threads/mutex.h"
+
+namespace LIBC_NAMESPACE {
+
+using AtExitCallback = void(void *);
+using StdCAtExitCallback = void(void);
+constexpr size_t CALLBACK_LIST_SIZE_FOR_TESTS = 1024;
+
+struct AtExitUnit {
+  AtExitCallback *callback = nullptr;
+  void *payload = nullptr;
+  LIBC_INLINE constexpr AtExitUnit() = default;
+  LIBC_INLINE constexpr AtExitUnit(AtExitCallback *c, void *p)
+      : callback(c), payload(p) {}
+};
+
+#if defined(LIBC_TARGET_ARCH_IS_GPU)
+using ExitCallbackList = FixedVector<AtExitUnit, 64>;
+#elif defined(LIBC_COPT_PUBLIC_PACKAGING)
+using ExitCallbackList = ReverseOrderBlockStore<AtExitUnit, 32>;
+#else
+using ExitCallbackList = FixedVector<AtExitUnit, CALLBACK_LIST_SIZE_FOR_TESTS>;
+#endif
+
+extern ExitCallbackList atexit_callbacks;
+extern ExitCallbackList at_quick_exit_callbacks;
+
+extern Mutex handler_list_mtx;
+
+void stdc_at_exit_func(void *payload);
+
+void call_exit_callbacks(ExitCallbackList &callbacks);
+
+int add_atexit_unit(ExitCallbackList &callbacks, const AtExitUnit &unit);
+
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC_STDLIB_EXIT_HANDLER_H