xref: /llvm-project/libc/src/__support/threads/fork_callbacks.cpp (revision 4d25761b076d19c237890ad7b7331d01bd3a84b6)
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
91