[libc] Add implementation of pthread_atfork.

[michaelrj]

LGTM with nit

libc/src/__support/fork_callbacks.cpp

30

nit: "can"

[sivachandra]

Fix typo in comment.

libc/config/linux/api.td

let Enumerations = [
     "thrd_busy",
     "thrd_error",
     "thrd_nomem",
   ];
 }
 
 def PThreadAPI : PublicAPI<"pthread.h"> {
   let Types = [
+      "__atfork_callback_t",
       "__pthread_once_func_t",
       "__pthread_start_t",
       "__pthread_tss_dtor_t",
       "pthread_attr_t",
       "pthread_mutex_t",
       "pthread_mutexattr_t",
       "pthread_t",
       "pthread_key_t",

libc/config/linux/x86_64/entrypoints.txt

list(APPEND TARGET_LIBC_ENTRYPOINTS
 
     # dirent.h entrypoints
     libc.src.dirent.closedir
     libc.src.dirent.dirfd
     libc.src.dirent.opendir
     libc.src.dirent.readdir
 
     # pthread.h entrypoints
+    libc.src.pthread.pthread_atfork
     libc.src.pthread.pthread_attr_destroy
     libc.src.pthread.pthread_attr_init
     libc.src.pthread.pthread_attr_getdetachstate
     libc.src.pthread.pthread_attr_getguardsize
     libc.src.pthread.pthread_attr_getstack
     libc.src.pthread.pthread_attr_getstacksize
     libc.src.pthread.pthread_attr_setdetachstate
     libc.src.pthread.pthread_attr_setguardsize

libc/include/CMakeLists.txt

 )
 
 add_gen_header(
   pthread
   DEF_FILE pthread.h.def
   GEN_HDR pthread.h
   DEPENDS
     .llvm_libc_common_h
+    .llvm-libc-types.__atfork_callback_t
     .llvm-libc-types.__pthread_start_t
     .llvm-libc-types.__pthread_tss_dtor_t
     .llvm-libc-types.pthread_attr_t
     .llvm-libc-types.pthread_key_t
     .llvm-libc-types.pthread_mutex_t
     .llvm-libc-types.pthread_mutexattr_t
     .llvm-libc-types.pthread_once_t
     .llvm-libc-types.__pthread_once_func_t

libc/include/llvm-libc-types/CMakeLists.txt

 add_header(off64_t HDR off64_t.h)
 add_header(size_t HDR size_t.h)
+add_header(__atfork_callback_t HDR __atfork_callback_t.h)
 add_header(__bsearchcompare_t HDR __bsearchcompare_t.h)
 add_header(__call_once_func_t HDR __call_once_func_t.h)
 add_header(__exec_argv_t HDR __exec_argv_t.h)
 add_header(__exec_envp_t HDR __exec_envp_t.h)
 add_header(__futex_word HDR __futex_word.h)
 add_header(__mutex_type HDR __mutex_type.h DEPENDS .__futex_word)
 add_header(__pthread_once_func_t HDR __pthread_once_func_t.h)
 add_header(__pthread_start_t HDR __pthread_start_t.h)

libc/include/llvm-libc-types/__atfork_callback_t.h

+//===-- Definition of type __atfork_callback_t ----------------------------===//
+//
+// 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_TYPES_ATFORK_CALLBACK_T_H__
+#define __LLVM_LIBC_TYPES_ATFORK_CALLBACK_T_H__
+
+typedef void (*__atfork_callback_t)(void);
+
+#endif // __LLVM_LIBC_TYPES_ATFORK_CALLBACK_T_H__

libc/spec/posix.td

 def ConstStructDirentPtrPtr : ConstType<StructDirentPtrPtr>;
 
 def StructTimeSpec : NamedType<"struct timespec">;
 def StructTimeSpecPtr : PtrType<StructTimeSpec>;
 
 def ExecArgvT : NamedType<"__exec_argv_t">;
 def ExecEnvpT : NamedType<"__exec_envp_t">;
 
+def AtForkCallbackT : NamedType<"__atfork_callback_t">;
+
 def POSIX : StandardSpec<"POSIX"> {
   PtrType CharPtr = PtrType<CharType>;
   RestrictedPtrType RestrictedCharPtr = RestrictedPtrType<CharType>;
   RestrictedPtrType CharRestrictedDoublePtr = RestrictedPtrType<CharPtr>;
   ConstType ConstCharPtr = ConstType<CharPtr>;
   ConstType ConstRestrictedCharPtr = ConstType<RestrictedCharPtr>;
 
   NamedType ModeTType = NamedType<"mode_t">;
[
         >,
     ]
   >;
 
   HeaderSpec PThread = HeaderSpec<
     "pthread.h",
     [], // Macros
     [
+        AtForkCallbackT,
         PThreadAttrTType,
         PThreadKeyT,
         PThreadMutexAttrTType,
         PThreadMutexTType,
         PThreadOnceCallback,
         PThreadOnceT,
         PThreadStartT,
         PThreadTSSDtorT,
         PThreadTType,
     ], // Types
     [], // Enumerations
     [
       FunctionSpec<
+          "pthread_atfork",
+          RetValSpec<IntType>,
+          [ArgSpec<AtForkCallbackT>, ArgSpec<AtForkCallbackT>, ArgSpec<AtForkCallbackT>]
+      >,
+      FunctionSpec<
           "pthread_attr_init",
           RetValSpec<IntType>,
           [ArgSpec<PThreadAttrTPtr>]
       >,
       FunctionSpec<
           "pthread_attr_destroy",
           RetValSpec<IntType>,
           [ArgSpec<PThreadAttrTPtr>]

libc/src/__support/CMakeLists.txt

DEPENDS
     .uint128
     libc.include.errno
     libc.src.__support.CPP.limits
     libc.src.__support.FPUtil.fp_bits
     libc.src.__support.builtin_wrappers
     libc.src.errno.errno
 )
 
+add_object_library(
+  fork_callbacks
+  SRCS
+    fork_callbacks.cpp
+  HDRS
+    fork_callbacks.h
+  DEPENDS
+    .threads.mutex
+)
+
 add_header_library(
   integer_operations
   HDRS
     integer_operations.h
 )
 
 add_header_library(
   arg_list

libc/src/__support/fork_callbacks.h

+//===-- At-fork callback helpers  -------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+namespace __llvm_libc {
+
+using ForkCallback = void(void);
+
+bool register_atfork_callbacks(ForkCallback *prepare_cd,
+                               ForkCallback *parent_cb, ForkCallback *child_cb);
+void invoke_prepare_callbacks();
+void invoke_parent_callbacks();
+void invoke_child_callbacks();
+
+} // namespace __llvm_libc

libc/src/__support/fork_callbacks.cpp

+//===-- Implementation of at-fork callback helpers  -----------------------===//
+//
+// 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 "fork_callbacks.h"
+
+#include "src/__support/threads/mutex.h"
+
+#include <stddef.h> // For size_t
+
+namespace __llvm_libc {
+
+namespace {
+
+struct ForkCallbackTriple {
+  ForkCallback *prepare = nullptr;
+  ForkCallback *parent = nullptr;
+  ForkCallback *child = nullptr;
+  constexpr ForkCallbackTriple() = default;
+};
+
+class AtForkCallbackManager {
+  static constexpr size_t CALLBACK_SIZE = 32;
+  // TODO: Replace this with block store when integration tests
+  // can use allocators.
+  ForkCallbackTriple list[CALLBACK_SIZE];
+  Mutex mtx;
+  size_t next_index;
+
+public:
+  constexpr AtForkCallbackManager() : mtx(false, false, false), next_index(0) {}
+
+  bool register_triple(const ForkCallbackTriple &triple) {
+    MutexLock lock(&mtx);
+    if (next_index >= CALLBACK_SIZE)
+      return false;
+    list[next_index] = triple;
+    ++next_index;
+    return true;
+  }
+
+  void invoke_prepare() {
+    MutexLock lock(&mtx);
+    for (size_t i = 0; i < next_index; ++i) {
+      auto prepare = list[i].prepare;
+      if (prepare)
+        prepare();
+    }
+  }
+
+  void invoke_parent() {
+    MutexLock lock(&mtx);
+    for (size_t i = 0; i < next_index; ++i) {
+      auto parent = list[i].parent;
+      if (parent)
+        parent();
+    }
+  }
+
+  void invoke_child() {
+    MutexLock lock(&mtx);
+    for (size_t i = 0; i < next_index; ++i) {
+      auto child = list[i].child;
+      if (child)
+        child();
+    }
+  }
+};
+
+AtForkCallbackManager cb_manager;
+
+} // Anonymous namespace
+
+bool register_atfork_callbacks(ForkCallback *prepare_cb,
+                               ForkCallback *parent_cb,
+                               ForkCallback *child_cb) {
+  return cb_manager.register_triple({prepare_cb, parent_cb, child_cb});
+}
+
+void invoke_child_callbacks() { cb_manager.invoke_child(); }
+
+void invoke_prepare_callbacks() { cb_manager.invoke_prepare(); }
+
+void invoke_parent_callbacks() { cb_manager.invoke_parent(); }
+
+} // namespace __llvm_libc

libc/src/pthread/CMakeLists.txt

add_entrypoint_object(
   SRCS
     pthread_once.cpp
   HDRS
     pthread_once.h
   DEPENDS
     libc.include.pthread
     libc.src.__support.threads.callonce
 )
+
+add_entrypoint_object(
+  pthread_atfork
+  SRCS
+    pthread_atfork.cpp
+  HDRS
+    pthread_atfork.h
+  DEPENDS
+    libc.include.errno
+    libc.include.pthread
+    libc.src.__support.fork_callbacks
+)

libc/src/pthread/pthread_atfork.h

+//===-- Implementation header for pthread_atfork function -------*- 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_THREADS_PTHREAD_ATFORK_H
+#define LLVM_LIBC_SRC_THREADS_PTHREAD_ATFORK_H
+
+#include <pthread.h>
+
+namespace __llvm_libc {
+
+int pthread_atfork(__atfork_callback_t prepare, __atfork_callback_t parent,
+                   __atfork_callback_t child);
+
+} // namespace __llvm_libc
+
+#endif // LLVM_LIBC_SRC_THREADS_PTHREAD_ATFORK_H

libc/src/pthread/pthread_atfork.cpp

+//===-- Linux implementation of the pthread_atfork function ---------------===//
+//
+// 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 "pthread_atfork.h"
+
+#include "src/__support/common.h"
+#include "src/__support/fork_callbacks.h"
+
+#include <errno.h>
+#include <pthread.h> // For pthread_* type definitions.
+
+namespace __llvm_libc {
+
+LLVM_LIBC_FUNCTION(int, pthread_atfork,
+                   (__atfork_callback_t prepare, __atfork_callback_t parent,
+                    __atfork_callback_t child)) {
+  return register_atfork_callbacks(prepare, parent, child) ? 0 : ENOMEM;
+}
+
+} // namespace __llvm_libc

libc/src/unistd/linux/CMakeLists.txt

add_entrypoint_object(
   SRCS
     fork.cpp
   HDRS
     ../fork.h
   DEPENDS
     libc.include.errno
     libc.include.unistd
     libc.include.sys_syscall
+    libc.src.__support.fork_callbacks
     libc.src.__support.OSUtil.osutil
     libc.src.__support.threads.thread
     libc.src.errno.errno
 )
 
 add_entrypoint_object(
   execv
   SRCS

libc/src/unistd/linux/fork.cpp

 //===-- Linux implementation of fork --------------------------------------===//
 //
 // 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/unistd/fork.h"
 
 #include "src/__support/OSUtil/syscall.h" // For internal syscall function.
 #include "src/__support/common.h"
+#include "src/__support/fork_callbacks.h"
 #include "src/__support/threads/thread.h" // For thread self object
 
 #include <errno.h>
 #include <sys/syscall.h> // For syscall numbers.
 
 namespace __llvm_libc {
 
 // The implementation of fork here is very minimal. We will add more
 // functionality and standard compliance in future.
 
 LLVM_LIBC_FUNCTION(pid_t, fork, (void)) {
+  invoke_prepare_callbacks();
 #ifdef SYS_fork
   pid_t ret = __llvm_libc::syscall_impl(SYS_fork);
 #elif defined(SYS_clone)
   pid_t ret = __llvm_libc::syscall_impl(SYS_clone, SIGCHLD, 0);
 #else
 #error "SYS_fork or SYS_clone not available."
 #endif
   if (ret == 0) {
     // Return value is 0 in the child process.
     // The child is created with a single thread whose self object will be a
     // copy of parent process' thread which called fork. So, we have to fix up
     // the child process' self object with the new process' tid.
     self.attrib->tid = __llvm_libc::syscall_impl(SYS_gettid);
+    invoke_child_callbacks();
     return 0;
   }
 
   if (ret < 0) {
     // Error case, a child process was not created.
     errno = -ret;
     return -1;
   }
 
+  invoke_parent_callbacks();
   return ret;
 }
 
 } // namespace __llvm_libc

libc/test/integration/src/unistd/CMakeLists.txt

 add_custom_target(unistd-integration-tests)
 add_dependencies(libc-integration-tests unistd-integration-tests)
 
 add_integration_test(
   fork_test
   SUITE
     unistd-integration-tests
   SRCS
     fork_test.cpp
   LOADER
     libc.loader.linux.crt1
   DEPENDS
     libc.include.errno
     libc.include.signal
     libc.include.sys_wait
     libc.include.unistd
+    libc.src.pthread.pthread_atfork
     libc.src.signal.raise
     libc.src.sys.wait.wait
     libc.src.sys.wait.wait4
     libc.src.sys.wait.waitpid
     libc.src.unistd.fork
 )
 
 add_executable(

libc/test/integration/src/unistd/fork_test.cpp

 //===-- Unittests for fork ------------------------------------------------===//
 //
 // 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/pthread/pthread_atfork.h"
 #include "src/signal/raise.h"
 #include "src/sys/wait/wait.h"
 #include "src/sys/wait/wait4.h"
 #include "src/sys/wait/waitpid.h"
 #include "src/unistd/fork.h"
 
 #include "utils/IntegrationTest/test.h"
 
void fork_and_waitpid_signal_exit() {
   int status;
   pid_t cpid = __llvm_libc::waitpid(pid, &status, 0);
   ASSERT_TRUE(cpid > 0);
   ASSERT_EQ(cpid, pid);
   ASSERT_FALSE(WIFEXITED(status));
   ASSERT_TRUE(WTERMSIG(status) == SIGUSR1);
 }
 
+static int prepare = 0;
+static int parent = 0;
+static int child = 0;
+static constexpr int DONE = 0x600D;
+
+static void prepare_cb() { prepare = DONE; }
+
+static void parent_cb() { parent = DONE; }
+
+static void child_cb() { child = DONE; }
+
+void fork_with_atfork_callbacks() {
+  ASSERT_EQ(__llvm_libc::pthread_atfork(&prepare_cb, &parent_cb, &child_cb), 0);
+  pid_t pid = __llvm_libc::fork();
+  if (pid == 0) {
+    // Raise a signal from the child if unexpected at-fork
+    // behavior is observed.
+    if (child != DONE || prepare != DONE || parent == DONE)
+      __llvm_libc::raise(SIGUSR1);
+    return;
+  }
+
+  ASSERT_TRUE(pid > 0);
+  int status;
+  pid_t cpid = __llvm_libc::waitpid(pid, &status, 0);
+  ASSERT_TRUE(cpid > 0);
+  ASSERT_EQ(cpid, pid);
+  ASSERT_TRUE(WIFEXITED(status));
+  ASSERT_EQ(prepare, DONE);
+  ASSERT_EQ(parent, DONE);
+  ASSERT_NE(child, DONE);
+}
+
 TEST_MAIN(int argc, char **argv, char **envp) {
   fork_and_wait_normal_exit();
   fork_and_wait4_normal_exit();
   fork_and_waitpid_normal_exit();
   fork_and_wait_signal_exit();
   fork_and_wait4_signal_exit();
   fork_and_waitpid_signal_exit();
+  fork_with_atfork_callbacks();
   return 0;
 }