1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
#pragma once
#include <snmalloc/aal/aal.h>
#include <snmalloc/stl/atomic.h>
#include <stddef.h>
#ifdef SNMALLOC_PTHREAD_ATFORK_WORKS
# include <pthread.h>
#endif
namespace snmalloc
{
#ifdef SNMALLOC_PTHREAD_FORK_PROTECTION
// This is a simple implementation of a class that can be
// used to prevent a process from forking. Holding a lock
// in the allocator while forking can lead to deadlocks.
// This causes the fork to wait out any other threads inside
// the allocators locks.
//
// The use is
// ```
// {
// PreventFork pf;
// // Code that should not be running during a fork.
// }
// ```
class PreventFork
{
// Global atomic counter of the number of threads currently preventing the
// system from forking. The bottom bit is used to signal that a thread is
// wanting to fork.
static inline stl::Atomic<size_t> threads_preventing_fork{0};
// The depth of the current thread's prevention of forking.
// This is used to enable reentrant prevention of forking.
static inline thread_local size_t depth_of_prevention{0};
// There could be multiple copies of the atfork handler installed.
// Only perform work for the first prefork and final postfork.
static inline thread_local size_t depth_of_handlers{0};
// This function ensures that the fork handler has been installed at least
// once. It might be installed more than once, this is safe. As subsequent
// calls would be ignored.
static void ensure_init()
{
# ifdef SNMALLOC_PTHREAD_ATFORK_WORKS
static stl::Atomic<bool> initialised{false};
if (initialised.load(stl::memory_order_acquire))
return;
pthread_atfork(prefork, postfork_parent, postfork_child);
initialised.store(true, stl::memory_order_release);
# endif
};
public:
PreventFork()
{
if (depth_of_prevention++ == 0)
{
// Ensure that the system is initialised before we start.
// Don't do this on nested Prevent calls.
ensure_init();
while (true)
{
auto prev = threads_preventing_fork.fetch_add(2);
if (prev % 2 == 0)
break;
threads_preventing_fork.fetch_sub(2);
while ((threads_preventing_fork.load() % 2) == 1)
{
Aal::pause();
}
};
}
}
~PreventFork()
{
if (--depth_of_prevention == 0)
{
threads_preventing_fork -= 2;
}
}
// The function that notifies new threads not to enter PreventFork regions
// It waits until all threads are no longer in a PreventFork region before
// returning.
static void prefork()
{
if (depth_of_handlers++ != 0)
return;
if (depth_of_prevention != 0)
error("Fork attempted while in PreventFork region.");
while (true)
{
auto current = threads_preventing_fork.load();
if (
(current % 2 == 0) &&
(threads_preventing_fork.compare_exchange_weak(current, current + 1)))
{
break;
}
Aal::pause();
};
while (threads_preventing_fork.load() != 1)
{
Aal::pause();
}
// Finally set the flag that allows this thread to enter PreventFork
// regions This is safe as the only other calls here are to other prefork
// handlers.
depth_of_prevention++;
}
// Unsets the flag that allows threads to enter PreventFork regions
// and for another thread to request a fork.
static void postfork_child()
{
// Count out the number of handlers that have been called, and
// only perform on the last.
if (--depth_of_handlers != 0)
return;
// This thread is no longer preventing a fork, so decrement the counter.
depth_of_prevention--;
// Allow other threads to allocate
// There could have been threads spinning in the prefork handler having
// optimistically increasing thread_preventing_fork by 2, but now the
// threads do not exist due to the fork. So restart the counter in the
// child.
threads_preventing_fork = 0;
}
// Unsets the flag that allows threads to enter PreventFork regions
// and for another thread to request a fork.
static void postfork_parent()
{
// Count out the number of handlers that have been called, and
// only perform on the last.
if (--depth_of_handlers != 0)
return;
// This thread is no longer preventing a fork, so decrement the counter.
depth_of_prevention--;
// Allow other threads to allocate
// Just remove the bit, and let the potential other threads in prefork
// remove their counts.
threads_preventing_fork--;
}
};
#else
// The fork protection can cost a lot and it is generally not required.
// This is a dummy implementation of the PreventFork class that does nothing.
class PreventFork
{
public:
PreventFork() {}
~PreventFork() {}
};
#endif
} // namespace snmalloc