Skip to content

Commit 05963e4

Browse files
SchrodingerZhuSchrodingerZhu
committed
finish the state pool
1 parent d050054 commit 05963e4

File tree

5 files changed

+259
-15
lines changed

5 files changed

+259
-15
lines changed

‎libc/src/__support/mpmc_stack.h

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
#include "src/__support/CPP/atomic.h"
1313
#include "src/__support/CPP/new.h"
1414
#include "src/__support/CPP/optional.h"
15-
#include "src/__support/CPP/type_traits.h"
1615
#include "src/__support/aba_ptr.h"
1716

1817
namespace LIBC_NAMESPACE_DECL {

‎libc/src/stdlib/linux/CMakeLists.txt

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,3 +9,23 @@ add_entrypoint_object(
99
libc.src.signal.raise
1010
libc.src.stdlib._Exit
1111
)
12+
13+
add_header_library(
14+
vsdo_rng
15+
HDRS
16+
vsdo_rng.h
17+
DEPENDS
18+
libc.src.__support.threads.thread # For __cxa_thread_atexit_impl
19+
libc.src.__support.CPP.algorithm
20+
libc.src.__support.CPP.bit
21+
libc.src.__support.CPP.mutex
22+
libc.src.__support.CPP.optional
23+
libc.src.__support.OSUtil.linux.vdso
24+
libc.src.__support.OSUtil.osutil
25+
libc.src.__support.macros.config
26+
libc.src.__support.mpmc_stack
27+
libc.src.__support.threads.callonce
28+
libc.src.__support.threads.linux.raw_mutex
29+
libc.src.sys.auxv.getauxval
30+
libc.include.sys_syscall
31+
)

‎libc/src/stdlib/linux/vsdo_rng.h

Lines changed: 164 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
#include "src/__support/CPP/mutex.h"
1414
#include "src/__support/OSUtil/linux/vdso.h"
1515
#include "src/__support/OSUtil/syscall.h"
16-
#include "src/__support/blockstore.h"
1716
#include "src/__support/common.h"
1817
#include "src/__support/macros/config.h"
1918
#include "src/__support/mpmc_stack.h"
@@ -23,6 +22,11 @@
2322

2423
namespace LIBC_NAMESPACE_DECL {
2524
namespace vsdo_rng {
25+
extern "C" {
26+
using Destructor = void(void *);
27+
[[gnu::weak]] extern void *__dso_handle;
28+
int __cxa_thread_atexit_impl(Destructor *, void *, void *);
29+
}
2630
class GlobalState {
2731
public:
2832
struct VGetrandomOpaqueParams {
@@ -32,7 +36,6 @@ class GlobalState {
3236
unsigned int reserved[13];
3337
};
3438

35-
private:
3639
struct Config {
3740
size_t page_size;
3841
size_t pages_per_alloc;
@@ -41,12 +44,11 @@ class GlobalState {
4144
VGetrandomOpaqueParams params;
4245
};
4346

47+
private:
4448
// A lock-free stack of free opaque states.
4549
MPMCStack<void *> free_list{};
4650
// A mutex protecting the allocation of new pages.
4751
RawMutex allocation_mutex{};
48-
// A block store of allocated pages.
49-
BlockStore<void *, 16> allocations{};
5052

5153
// Shared global configuration.
5254
static CallOnceFlag config_flag;
@@ -58,15 +60,79 @@ class GlobalState {
5860

5961
// Grow available states. This function can fail if the system is out of
6062
// memory.
61-
LIBC_INLINE bool grow();
63+
// - This routine assumes that the global config is valid.
64+
// - On success, this routine returns one opaque state for direct use.
65+
LIBC_INLINE void *grow();
6266

6367
public:
6468
LIBC_INLINE constexpr GlobalState() {}
65-
LIBC_INLINE static Config &get_config();
66-
LIBC_INLINE ~GlobalState() {}
69+
LIBC_INLINE static const Config &get_config();
70+
LIBC_INLINE static const Config &get_config_unchecked() { return config; }
71+
LIBC_INLINE void *get();
72+
LIBC_INLINE void recycle(void *state);
6773
};
6874

69-
class LocalState {};
75+
LIBC_INLINE_VAR GlobalState global_state{};
76+
77+
class LocalState {
78+
bool in_flight = false;
79+
bool failed = false;
80+
void *state = nullptr;
81+
82+
public:
83+
struct Guard {
84+
LocalState *tls;
85+
LIBC_INLINE Guard(LocalState *tls) : tls(tls) {
86+
tls->in_flight = true;
87+
cpp::atomic_thread_fence(cpp::MemoryOrder::SEQ_CST);
88+
}
89+
LIBC_INLINE Guard(Guard &&other) : tls(other.tls) { other.tls = nullptr; }
90+
LIBC_INLINE ~Guard() {
91+
cpp::atomic_thread_fence(cpp::MemoryOrder::SEQ_CST);
92+
if (tls)
93+
tls->in_flight = false;
94+
}
95+
LIBC_INLINE void fill(void *buf, size_t size) const;
96+
};
97+
LIBC_INLINE constexpr LocalState() {}
98+
LIBC_INLINE cpp::optional<Guard> get() {
99+
if (in_flight)
100+
return cpp::nullopt;
101+
102+
Guard guard(this);
103+
104+
if (!failed && !state) {
105+
int register_res = __cxa_thread_atexit_impl(
106+
[](void *self) {
107+
auto *tls = static_cast<LocalState *>(self);
108+
// Reject all future attempts to get a state.
109+
void *state = tls->state;
110+
tls->in_flight = true;
111+
tls->failed = true;
112+
tls->state = nullptr;
113+
cpp::atomic_thread_fence(cpp::MemoryOrder::SEQ_CST);
114+
if (state)
115+
LIBC_NAMESPACE::vsdo_rng::global_state.recycle(state);
116+
},
117+
this, __dso_handle);
118+
if (register_res == 0)
119+
state = LIBC_NAMESPACE::vsdo_rng::global_state.get();
120+
if (!state)
121+
failed = true;
122+
}
123+
124+
if (!state)
125+
return cpp::nullopt;
126+
127+
return cpp::move(guard);
128+
}
129+
};
130+
131+
LIBC_INLINE_VAR LIBC_THREAD_LOCAL LocalState local_state{};
132+
133+
//===----------------------------------------------------------------------===//
134+
// Implementation
135+
//===----------------------------------------------------------------------===//
70136

71137
LIBC_INLINE_VAR GlobalState::Config GlobalState::config{};
72138
LIBC_INLINE_VAR CallOnceFlag GlobalState::config_flag = 0;
@@ -87,7 +153,7 @@ LIBC_INLINE size_t GlobalState::cpu_count() {
87153
return count > 0 ? count : 1;
88154
}
89155

90-
LIBC_INLINE GlobalState::Config &GlobalState::get_config() {
156+
LIBC_INLINE const GlobalState::Config &GlobalState::get_config() {
91157
callonce(&config_flag, []() {
92158
config.getrandom =
93159
LIBC_NAMESPACE::vdso::TypedSymbol<vdso::VDSOSym::GetRandom>{};
@@ -106,7 +172,7 @@ LIBC_INLINE GlobalState::Config &GlobalState::get_config() {
106172
if (!config.page_size)
107173
return;
108174

109-
size_t count = cpu_count();
175+
size_t count = cpp::max(cpu_count(), size_t{4});
110176

111177
config.states_per_page =
112178
config.page_size / config.params.size_of_opaque_states;
@@ -117,10 +183,94 @@ LIBC_INLINE GlobalState::Config &GlobalState::get_config() {
117183
return config;
118184
}
119185

120-
LIBC_INLINE bool GlobalState::grow() {
121-
// reserve a slot for the new page.
122-
if (!allocations.push_back(nullptr))
123-
return false;
186+
LIBC_INLINE void *GlobalState::grow() {
187+
cpp::lock_guard guard(allocation_mutex);
188+
189+
// It is possible that when we finally grab the lock, other threads have
190+
// successfully finished the allocation already. Hence, we first try if we
191+
// can pop anything from the free list.
192+
if (cpp::optional<void *> state = free_list.pop())
193+
return *state;
194+
195+
long mmap_res = LIBC_NAMESPACE::syscall_impl<long>(
196+
SYS_mmap, /*addr=*/nullptr,
197+
/*length=*/config.page_size * config.pages_per_alloc,
198+
/*prot=*/config.params.mmap_prot,
199+
/*flags=*/config.params.mmap_flags,
200+
/*fd=*/-1, /*offset=*/0);
201+
if (mmap_res == -1 /* MAP_FAILED */)
202+
return nullptr;
203+
204+
char *pages = reinterpret_cast<char *>(mmap_res);
205+
206+
// Initialize the page.
207+
size_t total_states = config.pages_per_alloc * config.states_per_page;
208+
size_t free_states = total_states - 1; // reserve one for direct use.
209+
__extension__ void *opaque_states[total_states];
210+
size_t index = 0;
211+
for (size_t p = 0; p < config.pages_per_alloc; ++p) {
212+
char *page = &pages[p * config.page_size];
213+
for (size_t s = 0; s < config.states_per_page; ++s) {
214+
void *state = &page[s * config.params.size_of_opaque_states];
215+
opaque_states[index++] = state;
216+
}
217+
}
218+
219+
constexpr size_t RETRY_COUNT = 64;
220+
for (size_t i = 0; i < RETRY_COUNT; ++i) {
221+
if (free_list.push_all(opaque_states, free_states))
222+
break;
223+
// Abort if we are still short in memory after all these retries.
224+
if (i + 1 == RETRY_COUNT) {
225+
LIBC_NAMESPACE::syscall_impl<long>(
226+
SYS_munmap, pages, config.page_size * config.pages_per_alloc);
227+
return nullptr;
228+
}
229+
}
230+
231+
return opaque_states[free_states];
232+
}
233+
234+
LIBC_INLINE void *GlobalState::get() {
235+
const Config &config = get_config();
236+
// If page size is not set, the global config is invalid. Early return.
237+
if (!config.page_size)
238+
return nullptr;
239+
240+
if (cpp::optional<void *> state = free_list.pop())
241+
return *state;
242+
243+
// At this stage, we know that the config is valid.
244+
return grow();
245+
}
246+
247+
LIBC_INLINE void GlobalState::recycle(void *state) {
248+
LIBC_ASSERT(state != nullptr);
249+
constexpr size_t RETRY_COUNT = 64;
250+
for (size_t i = 0; i < RETRY_COUNT; ++i)
251+
if (free_list.push(state))
252+
return;
253+
// Otherwise, we just let it leak. It won't be too bad not to reuse the state
254+
// since the OS can free the page if memory is tight.
255+
}
256+
257+
//===----------------------------------------------------------------------===//
258+
// LocalState
259+
//===----------------------------------------------------------------------===//
260+
261+
LIBC_INLINE void LocalState::Guard::fill(void *buf, size_t size) const {
262+
LIBC_ASSERT(tls->state != nullptr);
263+
char *cursor = reinterpret_cast<char *>(buf);
264+
size_t remaining = size;
265+
const auto &config = GlobalState::get_config_unchecked();
266+
while (remaining > 0) {
267+
int res = config.getrandom(cursor, remaining, /* default random flag */ 0,
268+
tls->state, config.params.size_of_opaque_states);
269+
if (res < 0)
270+
continue;
271+
remaining -= static_cast<size_t>(res);
272+
cursor += res;
273+
}
124274
}
125275

126276
} // namespace vsdo_rng

‎libc/test/integration/src/stdlib/linux/CMakeLists.txt

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
add_custom_target(stdlib-linux-integration-tests)
2+
add_dependencies(libc-integration-tests stdlib-linux-integration-tests)
3+
4+
add_integration_test(
5+
vsdo_rng_test
6+
SUITE
7+
stdlib-linux-integration-tests
8+
SRCS
9+
vsdo_rng_test.cpp
10+
DEPENDS
11+
libc.src.pthread.pthread_create
12+
libc.src.pthread.pthread_join
13+
libc.src.stdlib.linux.vsdo_rng
14+
)

‎libc/test/integration/src/stdlib/linux/vsdo_rng_test.cpp

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
//===-- Test for vsdo_rng functionality ----------------------------------===//
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 "src/pthread/pthread_create.h"
10+
#include "src/pthread/pthread_join.h"
11+
#include "src/stdlib/linux/vsdo_rng.h"
12+
#include "test/IntegrationTest/test.h"
13+
14+
using namespace LIBC_NAMESPACE;
15+
16+
void basic_test() {
17+
// Test basic functionality
18+
vsdo_rng::LocalState &local_state = vsdo_rng::local_state;
19+
20+
// Try to get a guard
21+
if (auto guard = local_state.get()) {
22+
// Fill a small buffer with random data
23+
long long buffer[32] = {0};
24+
guard->fill(buffer, sizeof(buffer));
25+
26+
// Basic sanity check - buffer should not have zero
27+
for (auto &i : buffer)
28+
if (i == 0)
29+
__builtin_trap();
30+
}
31+
// If we can't get a guard, that's okay - the vDSO might not be available
32+
// or the system might not support getrandom
33+
}
34+
35+
void multithread_test() {
36+
constexpr static size_t OUTER_REPEAT = 8;
37+
constexpr static size_t INNER_REPEAT = 32;
38+
constexpr static size_t NUM_THREADS = 16;
39+
pthread_t threads[NUM_THREADS];
40+
41+
// Repeat outer loop so that
42+
for (size_t r = 0; r < OUTER_REPEAT; ++r) {
43+
for (pthread_t &thread : threads)
44+
LIBC_NAMESPACE::pthread_create(
45+
&thread, nullptr,
46+
[](void *) -> void * {
47+
for (size_t j = 0; j < INNER_REPEAT; ++j)
48+
basic_test();
49+
return nullptr;
50+
},
51+
nullptr);
52+
for (pthread_t thread : threads)
53+
LIBC_NAMESPACE::pthread_join(thread, nullptr);
54+
}
55+
}
56+
57+
TEST_MAIN() {
58+
basic_test();
59+
multithread_test();
60+
return 0;
61+
}

0 commit comments

Comments
 (0)