libc/src/__support/threads/fork_callbacks.cpp

   1 //===-- Implementation of at-fork callback helpers  -----------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "fork_callbacks.h"
  10
  11 #include "src/__support/threads/mutex.h"
  12
  13 #include <stddef.h> // For size_t
  14
  15 namespace __llvm_libc {
  16
  17 namespace {
  18
  19 struct ForkCallbackTriple {
  20   ForkCallback *prepare = nullptr;
  21   ForkCallback *parent = nullptr;
  22   ForkCallback *child = nullptr;
  23   constexpr ForkCallbackTriple() = default;
  24 };
  25
  26 class AtForkCallbackManager {
  27   static constexpr size_t CALLBACK_SIZE = 32;
  28   // TODO: Replace this with block store when integration tests
  29   // can use allocators.
  30   ForkCallbackTriple list[CALLBACK_SIZE];
  31   Mutex mtx;
  32   size_t next_index;
  33
  34 public:
  35   constexpr AtForkCallbackManager() : mtx(false, false, false), next_index(0) {}
  36
  37   bool register_triple(const ForkCallbackTriple &triple) {
  38     MutexLock lock(&mtx);
  39     if (next_index >= CALLBACK_SIZE)
  40       return false;
  41     list[next_index] = triple;
  42     ++next_index;
  43     return true;
  44   }
  45
  46   void invoke_prepare() {
  47     MutexLock lock(&mtx);
  48     for (size_t i = 0; i < next_index; ++i) {
  49       auto prepare = list[i].prepare;
  50       if (prepare)
  51         prepare();
  52     }
  53   }
  54
  55   void invoke_parent() {
  56     MutexLock lock(&mtx);
  57     for (size_t i = 0; i < next_index; ++i) {
  58       auto parent = list[i].parent;
  59       if (parent)
  60         parent();
  61     }
  62   }
  63
  64   void invoke_child() {
  65     MutexLock lock(&mtx);
  66     for (size_t i = 0; i < next_index; ++i) {
  67       auto child = list[i].child;
  68       if (child)
  69         child();
  70     }
  71   }
  72 };
  73
  74 AtForkCallbackManager cb_manager;
  75
  76 } // Anonymous namespace
  77
  78 bool register_atfork_callbacks(ForkCallback *prepare_cb,
  79                                ForkCallback *parent_cb,
  80                                ForkCallback *child_cb) {
  81   return cb_manager.register_triple({prepare_cb, parent_cb, child_cb});
  82 }
  83
  84 void invoke_child_callbacks() { cb_manager.invoke_child(); }
  85
  86 void invoke_prepare_callbacks() { cb_manager.invoke_prepare(); }
  87
  88 void invoke_parent_callbacks() { cb_manager.invoke_parent(); }
  89
  90 } // namespace __llvm_libc