clt_database/storage/buffer_pool.rs
1use branches::unlikely;
2
3use super::{slot_bitmap::AtomicSlotBitmap, sqlite3_ondisk::WAL_FRAME_HEADER_SIZE};
4use crate::io::TEMP_BUFFER_CACHE;
5use crate::sync::atomic::{AtomicUsize, Ordering};
6use crate::sync::Arc;
7use crate::turso_assert;
8use crate::{Buffer, LimboError, IO};
9
10use std::cell::UnsafeCell;
11use std::ptr::NonNull;
12use std::sync::OnceLock;
13
14#[derive(Debug)]
15/// A buffer allocated from an arena from `[BufferPool]`
16pub struct ArenaBuffer {
17 /// The `Arena` the buffer came from
18 arena: Arc<Arena>,
19 /// Pointer to the start of the buffer
20 ptr: NonNull<u8>,
21 /// Identifier for the `[Arena]` the buffer came from
22 arena_id: u32,
23 /// The index of the first slot making up the buffer
24 slot_idx: u32,
25 /// The requested length of the allocation.
26 /// For pooled buffers, `len` is always `<= Arena::slot_size` and occupies exactly one slot.
27 len: usize,
28}
29
30// Unsound: write and read from different threads can be dangerous with current ArenaBuffer implementation without some additional explicit synchronization
31unsafe impl Sync for ArenaBuffer {}
32unsafe impl Send for ArenaBuffer {}
33crate::assert::assert_send_sync!(ArenaBuffer);
34
35impl ArenaBuffer {
36 fn new(arena: Arc<Arena>, ptr: NonNull<u8>, len: usize, arena_id: u32, slot_idx: u32) -> Self {
37 ArenaBuffer {
38 arena,
39 ptr,
40 arena_id,
41 slot_idx,
42 len,
43 }
44 }
45
46 #[inline(always)]
47 /// Returns the `id` of the underlying arena, only if it was registered with `io_uring`
48 pub const fn fixed_id(&self) -> Option<u32> {
49 // Arenas which are not registered will have `id`s <= UNREGISTERED_START
50 if self.arena_id < UNREGISTERED_START {
51 Some(self.arena_id)
52 } else {
53 None
54 }
55 }
56
57 /// The requested size of the allocation, the actual size of the underlying buffer is rounded up to
58 /// the arena's slot_size (and in practice is always `<= slot_size` for pooled buffers).
59 pub const fn logical_len(&self) -> usize {
60 self.len
61 }
62 pub fn as_slice(&self) -> &[u8] {
63 unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.logical_len()) }
64 }
65 pub fn as_mut_slice(&mut self) -> &mut [u8] {
66 unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.logical_len()) }
67 }
68}
69
70impl Drop for ArenaBuffer {
71 fn drop(&mut self) {
72 self.arena.free(self.slot_idx, self.logical_len());
73 }
74}
75
76impl std::ops::Deref for ArenaBuffer {
77 type Target = [u8];
78 fn deref(&self) -> &Self::Target {
79 self.as_slice()
80 }
81}
82
83impl std::ops::DerefMut for ArenaBuffer {
84 fn deref_mut(&mut self) -> &mut Self::Target {
85 self.as_mut_slice()
86 }
87}
88
89/// Static Buffer pool managing multiple memory arenas
90/// of which `[ArenaBuffer]`s are returned for requested allocations
91pub struct BufferPool {
92 inner: UnsafeCell<PoolInner>,
93}
94
95unsafe impl Sync for BufferPool {}
96unsafe impl Send for BufferPool {}
97crate::assert::assert_send_sync!(BufferPool);
98
99struct PoolInner {
100 /// An instance of the program's IO, used for registering
101 /// Arena's with io_uring.
102 io: Option<Arc<dyn IO>>,
103 /// An Arena which returns `ArenaBuffer`s of size `db_page_size`.
104 page_arena: Option<Arc<Arena>>,
105 /// An Arena which returns `ArenaBuffer`s of size `db_page_size`
106 /// plus 24 byte `WAL_FRAME_HEADER_SIZE`, preventing the fragmentation
107 /// or complex book-keeping needed to use the same arena for both sizes.
108 wal_frame_arena: Option<Arc<Arena>>,
109 /// The size of each `Arena`, in bytes.
110 arena_size: usize,
111 /// The `[Database::page_size]`, which the `page_arena` will use to
112 /// return buffers from `Self::get_page`.
113 db_page_size: OnceLock<usize>,
114}
115
116unsafe impl Sync for PoolInner {}
117unsafe impl Send for PoolInner {}
118crate::assert::assert_send_sync!(PoolInner);
119
120impl Default for BufferPool {
121 fn default() -> Self {
122 Self::new(Self::DEFAULT_ARENA_SIZE)
123 }
124}
125
126impl BufferPool {
127 /// 3MB Default size for each `Arena`. Any higher and
128 /// it will fail to register the second arena with io_uring due
129 /// to `RL_MEMLOCK` limit for un-privileged processes being 8MB total.
130 pub const DEFAULT_ARENA_SIZE: usize = 3 * 1024 * 1024;
131 /// 1MB size For testing/CI
132 pub const TEST_ARENA_SIZE: usize = 1024 * 1024;
133 /// 4KB default page_size
134 pub const DEFAULT_PAGE_SIZE: usize = 4096;
135 /// Maximum size for each Arena (64MB total)
136 const MAX_ARENA_SIZE: usize = 32 * 1024 * 1024;
137 /// 64kb Minimum arena size
138 const MIN_ARENA_SIZE: usize = 1024 * 64;
139 fn new(arena_size: usize) -> Self {
140 turso_assert!(
141 (Self::MIN_ARENA_SIZE..Self::MAX_ARENA_SIZE).contains(&arena_size),
142 "Arena size out of valid range",
143 { "arena_size": arena_size, "min": Self::MIN_ARENA_SIZE, "max": Self::MAX_ARENA_SIZE }
144 );
145 Self {
146 inner: UnsafeCell::new(PoolInner {
147 page_arena: None,
148 wal_frame_arena: None,
149 arena_size,
150 db_page_size: OnceLock::new(),
151 io: None,
152 }),
153 }
154 }
155
156 /// Request a `Buffer` of size `len`
157 #[inline]
158 pub fn allocate(&self, len: usize) -> Buffer {
159 self.inner().allocate(len)
160 }
161
162 /// Request a `Buffer` the size of the `db_page_size` the `BufferPool` was initialized with.
163 #[inline]
164 pub fn get_page(&self) -> Buffer {
165 let inner = self.inner_mut();
166 inner.get_db_page_buffer()
167 }
168
169 /// Request a `Buffer` for use with a WAL frame,
170 /// `[Database::page_size] + `WAL_FRAME_HEADER_SIZE`
171 #[inline]
172 pub fn get_wal_frame(&self) -> Buffer {
173 let inner = self.inner_mut();
174 inner.get_wal_frame_buffer()
175 }
176
177 #[inline]
178 fn inner(&self) -> &PoolInner {
179 unsafe { &*self.inner.get() }
180 }
181
182 #[inline]
183 #[allow(clippy::mut_from_ref)]
184 fn inner_mut(&self) -> &mut PoolInner {
185 unsafe { &mut *self.inner.get() }
186 }
187
188 /// Create a static `BufferPool` initialize the pool to the default page size, **without**
189 /// populating the Arenas. Arenas will not be created until `[BufferPool::finalize_page_size]`,
190 /// and the pool will temporarily return temporary buffers to prevent reallocation of the
191 /// arena if the page size is set to something other than the default value.
192 pub fn begin_init(io: &Arc<dyn IO>, arena_size: usize) -> Arc<Self> {
193 let pool = Arc::new(BufferPool::new(arena_size));
194 let inner = pool.inner_mut();
195 // Just store the IO handle, don't create arena yet
196 if inner.io.is_none() {
197 inner.io = Some(Arc::clone(io));
198 }
199 pool
200 }
201
202 /// Call when `[Database::db_state]` is initialized, providing the `page_size` to allocate
203 /// an arena for the pool. Before this call, the pool will use temporary buffers which are
204 /// cached in thread local storage.
205 pub fn finalize_with_page_size(&self, page_size: usize) -> crate::Result<()> {
206 let inner = self.inner_mut();
207 tracing::trace!("finalize page size called with size {page_size}");
208 if page_size != BufferPool::DEFAULT_PAGE_SIZE {
209 // so far we have handed out some temporary buffers, since the page size is not
210 // default, we need to clear the cache so they aren't reused for other operations.
211 TEMP_BUFFER_CACHE.with(|cache| {
212 cache.borrow_mut().reinit_cache(page_size);
213 });
214 }
215 if inner.page_arena.is_some() {
216 tracing::trace!("Buffer pool already initialized, skipping finalize");
217 return Ok(());
218 }
219
220 // Tries to atomically (guarenteed by the OnceLock) initialize the page size for the inner pool.
221 // If it succeeds, we now have to initialize the arenas.
222 // If the initialization fails, this means the arenas have already been initialized by a previous thread
223 // This avoids a potential TOCTOU race, where 2 threads could try to initalize the arena at the same time
224 // after checking the `db_page_size`
225 if inner.db_page_size.set(page_size).is_ok() {
226 inner.init_arenas()?;
227 };
228 Ok(())
229 }
230}
231
232impl PoolInner {
233 #[inline]
234 pub fn get_db_page_size(&self) -> usize {
235 *(self
236 .db_page_size
237 .get()
238 .unwrap_or(&BufferPool::DEFAULT_PAGE_SIZE))
239 }
240
241 /// Allocate a buffer of the given length from the pool, falling back to
242 /// temporary thread local buffers if the pool is not initialized or is full.
243 pub fn allocate(&self, len: usize) -> Buffer {
244 turso_assert!(len > 0, "Cannot allocate zero-length buffer");
245
246 let db_page_size = self.get_db_page_size();
247 let wal_frame_size = db_page_size + WAL_FRAME_HEADER_SIZE;
248
249 // Check if this is exactly a WAL frame size allocation
250 if len == wal_frame_size {
251 return self
252 .wal_frame_arena
253 .as_ref()
254 .and_then(|wal_arena| Arena::try_alloc(wal_arena, len))
255 .unwrap_or_else(|| Buffer::new_temporary(len));
256 }
257 // For all other sizes, use regular arena
258 self.page_arena
259 .as_ref()
260 .and_then(|arena| Arena::try_alloc(arena, len))
261 .unwrap_or_else(|| Buffer::new_temporary(len))
262 }
263
264 fn get_db_page_buffer(&mut self) -> Buffer {
265 let db_page_size = self.get_db_page_size();
266 self.page_arena
267 .as_ref()
268 .and_then(|arena| Arena::try_alloc(arena, db_page_size))
269 .unwrap_or_else(|| Buffer::new_temporary(db_page_size))
270 }
271
272 fn get_wal_frame_buffer(&mut self) -> Buffer {
273 let len = self.get_db_page_size() + WAL_FRAME_HEADER_SIZE;
274 self.wal_frame_arena
275 .as_ref()
276 .and_then(|wal_arena| Arena::try_alloc(wal_arena, len))
277 .unwrap_or_else(|| Buffer::new_temporary(len))
278 }
279
280 /// Allocate a new arena for the pool to use
281 fn init_arenas(&mut self) -> crate::Result<()> {
282 let db_page_size = self.get_db_page_size();
283 let arena_size = self.arena_size;
284
285 let io = self.io.as_ref().expect("Pool not initialized").clone();
286
287 // Create regular page arena
288 match Arena::new(db_page_size, arena_size, &io) {
289 Ok(arena) => {
290 tracing::trace!(
291 "added arena {} with size {} MB and slot size {}",
292 arena.id,
293 arena_size / (1024 * 1024),
294 db_page_size
295 );
296 self.page_arena = Some(Arc::new(arena));
297 }
298 Err(e) => {
299 tracing::error!("Failed to create arena: {:?}", e);
300 return Err(LimboError::InternalError(format!(
301 "Failed to create arena: {e}",
302 )));
303 }
304 }
305
306 // Create WAL frame arena
307 let wal_frame_size = db_page_size + WAL_FRAME_HEADER_SIZE;
308 match Arena::new(wal_frame_size, arena_size, &io) {
309 Ok(arena) => {
310 tracing::trace!(
311 "added WAL frame arena {} with size {} MB and slot size {}",
312 arena.id,
313 arena_size / (1024 * 1024),
314 wal_frame_size
315 );
316 self.wal_frame_arena = Some(Arc::new(arena));
317 }
318 Err(e) => {
319 tracing::error!("Failed to create WAL frame arena: {:?}", e);
320 return Err(LimboError::InternalError(format!(
321 "Failed to create WAL frame arena: {e}",
322 )));
323 }
324 }
325
326 Ok(())
327 }
328}
329
330/// Preallocated block of memory used by the pool to distribute `ArenaBuffer`s
331#[derive(Debug)]
332struct Arena {
333 /// Identifier to tie allocations back to the arena. If the arena is registerd
334 /// with `io_uring`, then the ID represents the index of the arena into the ring's
335 /// sparse registered buffer array created on the ring's initialization.
336 id: u32,
337 /// Base pointer to the arena returned by `mmap`
338 base: NonNull<u8>,
339 /// Total number of slots currently allocated/in use.
340 allocated_slots: AtomicUsize,
341 /// Currently free slots (lock-free atomic bitmap).
342 free_slots: AtomicSlotBitmap,
343 /// Total size of the arena in bytes
344 arena_size: usize,
345 /// Slot size the total arena is divided into.
346 slot_size: usize,
347}
348
349// SAFETY: Arena's base pointer comes from mmap and is never aliased. All mutable
350// state is behind atomics (AtomicUsize, AtomicSlotBitmap), so concurrent access is safe.
351unsafe impl Send for Arena {}
352unsafe impl Sync for Arena {}
353
354impl Drop for Arena {
355 fn drop(&mut self) {
356 unsafe { arena::dealloc(self.base.as_ptr(), self.arena_size) };
357 }
358}
359
360/// Slots 0 and 1 will be reserved for Arenas which are registered buffers
361/// with io_uring.
362const UNREGISTERED_START: u32 = 2;
363
364/// ID's for an Arena which is not registered with `io_uring`
365/// registered arena will always have id = 0..=1
366/// we purposely use std::sync::AtomicU32 instead of core::sync::AtomicU32 because
367/// this is a global static variable and can mess with shuttle tests
368static NEXT_ID: std::sync::atomic::AtomicU32 =
369 std::sync::atomic::AtomicU32::new(UNREGISTERED_START);
370
371impl Arena {
372 /// Create a new arena with the given size and page size.
373 /// NOTE: Minimum arena size is slot_size * 64
374 fn new(slot_size: usize, arena_size: usize, io: &Arc<dyn IO>) -> Result<Self, String> {
375 let min_slots = arena_size.div_ceil(slot_size);
376 let rounded_slots = (min_slots.max(64) + 63) & !63;
377 let rounded_bytes = rounded_slots * slot_size;
378 // Guard against the global cap
379 if unlikely(rounded_bytes > BufferPool::MAX_ARENA_SIZE) {
380 return Err(format!(
381 "arena size {} B exceeds hard limit of {} B",
382 rounded_bytes,
383 BufferPool::MAX_ARENA_SIZE
384 ));
385 }
386 let ptr = unsafe { arena::alloc(rounded_bytes) };
387 let base = NonNull::new(ptr).ok_or("Failed to allocate arena")?;
388 let id = io
389 .register_fixed_buffer(base, rounded_bytes)
390 .unwrap_or_else(|_| {
391 // Register with io_uring if possible, otherwise use next available ID
392 let next_id = NEXT_ID.fetch_add(1, Ordering::AcqRel);
393 tracing::trace!("Allocating arena with id {}", next_id);
394 next_id
395 });
396 let map = AtomicSlotBitmap::new(rounded_slots as u32);
397 Ok(Self {
398 id,
399 base,
400 free_slots: map,
401 allocated_slots: AtomicUsize::new(0),
402 slot_size,
403 arena_size: rounded_bytes,
404 })
405 }
406
407 /// Allocate a `Buffer` large enough for logical length `size`.
408 pub fn try_alloc(arena: &Arc<Arena>, size: usize) -> Option<Buffer> {
409 if size > arena.slot_size {
410 // The buffer pool only supports single-slot allocations. Larger requests fall back to
411 // temporary heap buffers via the caller.
412 return None;
413 }
414 let first_idx = arena.free_slots.alloc_one()?;
415 arena.allocated_slots.fetch_add(1, Ordering::AcqRel);
416 let offset = first_idx as usize * arena.slot_size;
417 let ptr = unsafe { NonNull::new_unchecked(arena.base.as_ptr().add(offset)) };
418 Some(Buffer::new_pooled(ArenaBuffer::new(
419 Arc::clone(arena),
420 ptr,
421 size,
422 arena.id,
423 first_idx,
424 )))
425 }
426
427 /// Mark all relevant slots that include `size` starting at `slot_idx` as free.
428 pub fn free(&self, slot_idx: u32, size: usize) {
429 turso_assert!(
430 size <= self.slot_size,
431 "pooled buffers must not exceed one slot"
432 );
433 turso_assert!(
434 !self.free_slots.is_free(slot_idx),
435 "must not already be marked free"
436 );
437 self.free_slots.free_one(slot_idx);
438 self.allocated_slots.fetch_sub(1, Ordering::AcqRel);
439 }
440}
441
442#[cfg(all(unix, not(miri)))]
443mod arena {
444 use libc::MAP_ANONYMOUS;
445 use libc::{mmap, munmap, MAP_PRIVATE, PROT_READ, PROT_WRITE};
446 use std::ffi::c_void;
447
448 pub unsafe fn alloc(len: usize) -> *mut u8 {
449 let ptr = mmap(
450 std::ptr::null_mut(),
451 len,
452 PROT_READ | PROT_WRITE,
453 MAP_PRIVATE | MAP_ANONYMOUS,
454 -1,
455 0,
456 );
457 if ptr == libc::MAP_FAILED {
458 panic!("mmap failed: {}", std::io::Error::last_os_error());
459 }
460 #[cfg(target_os = "linux")]
461 {
462 libc::madvise(ptr, len, libc::MADV_HUGEPAGE);
463 }
464 ptr as *mut u8
465 }
466
467 pub unsafe fn dealloc(ptr: *mut u8, len: usize) {
468 let result = munmap(ptr as *mut c_void, len);
469 if result != 0 {
470 panic!("munmap failed: {}", std::io::Error::last_os_error());
471 }
472 }
473}
474
475#[cfg(any(not(unix), miri))]
476mod arena {
477 pub unsafe fn alloc(len: usize) -> *mut u8 {
478 let layout = std::alloc::Layout::from_size_align(len, std::mem::size_of::<u8>()).unwrap();
479 unsafe { std::alloc::alloc_zeroed(layout) }
480 }
481 pub unsafe fn dealloc(ptr: *mut u8, len: usize) {
482 let layout = std::alloc::Layout::from_size_align(len, std::mem::size_of::<u8>()).unwrap();
483 unsafe { std::alloc::dealloc(ptr, layout) };
484 }
485}
486
487/// Shuttle tests for concurrent buffer pool operations.
488///
489/// These tests target the `unsafe impl Sync/Send` on:
490/// - `ArenaBuffer`: Raw pointer access across threads
491/// - `BufferPool`: UnsafeCell-based interior mutability
492/// - `PoolInner`: Shared mutable state
493///
494#[cfg(all(shuttle, clt_turso_tests))]
495mod shuttle_tests {
496 use super::*;
497 use crate::io::MemoryIO;
498 use crate::sync::*;
499 use crate::thread;
500 use rustc_hash::FxHashSet as HashSet;
501
502 fn create_test_pool() -> Arc<BufferPool> {
503 let io: Arc<dyn IO> = Arc::new(MemoryIO::new());
504 let pool = BufferPool::begin_init(&io, BufferPool::TEST_ARENA_SIZE);
505 pool.finalize_with_page_size(4096).unwrap();
506 pool
507 }
508
509 /// Test concurrent allocations from BufferPool.
510 /// Verifies that multiple threads can safely call get_page() simultaneously.
511 #[test]
512 fn shuttle_concurrent_page_allocation() {
513 shuttle::check_random(
514 || {
515 let pool = create_test_pool();
516 let mut handles = vec![];
517
518 for _ in 0..3 {
519 let pool = Arc::clone(&pool);
520 let h = thread::spawn(move || {
521 let buf = pool.get_page();
522 assert_eq!(buf.len(), 4096);
523 buf
524 });
525 handles.push(h);
526 }
527
528 let buffers: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();
529 // Verify all buffers are valid and distinct (no double allocation)
530 assert_eq!(buffers.len(), 3);
531 },
532 1000,
533 );
534 }
535
536 /// Test concurrent allocation and deallocation.
537 /// Buffers are dropped in different threads than they were allocated.
538 #[test]
539 fn shuttle_concurrent_alloc_and_drop() {
540 shuttle::check_random(
541 || {
542 let pool = create_test_pool();
543 let pool2 = Arc::clone(&pool);
544
545 // Thread 1: allocate and send buffer to be dropped elsewhere
546 let h1 = thread::spawn(move || {
547 let buf = pool.get_page();
548 buf.len() // return length, buffer dropped here
549 });
550
551 // Thread 2: allocate concurrently
552 let h2 = thread::spawn(move || {
553 let buf = pool2.get_page();
554 buf.len()
555 });
556
557 assert_eq!(h1.join().unwrap(), 4096);
558 assert_eq!(h2.join().unwrap(), 4096);
559 },
560 1000,
561 );
562 }
563
564 /// Test that ArenaBuffer can be safely sent between threads and written to.
565 /// This tests the `unsafe impl Send + Sync for ArenaBuffer`.
566 #[test]
567 fn shuttle_arena_buffer_send_and_write() {
568 shuttle::check_random(
569 || {
570 let pool = create_test_pool();
571 let buf = pool.get_page();
572
573 // Write some data
574 buf.as_mut_slice()[0] = 42;
575
576 // Send to another thread for reading
577 let h = thread::spawn(move || {
578 assert_eq!(buf.as_slice()[0], 42);
579 buf.as_slice()[0]
580 });
581
582 assert_eq!(h.join().unwrap(), 42);
583 },
584 1000,
585 );
586 }
587
588 /// Test concurrent WAL frame and page allocations.
589 /// Both arena types are exercised simultaneously.
590 #[test]
591 fn shuttle_concurrent_mixed_allocations() {
592 shuttle::check_random(
593 || {
594 let pool = create_test_pool();
595 let pool2 = Arc::clone(&pool);
596 let pool3 = Arc::clone(&pool);
597
598 let h1 = thread::spawn(move || {
599 let buf = pool.get_page();
600 assert_eq!(buf.len(), 4096);
601 });
602
603 let h2 = thread::spawn(move || {
604 let buf = pool2.get_wal_frame();
605 // WAL frame = page_size + WAL_FRAME_HEADER_SIZE (24)
606 assert_eq!(buf.len(), 4096 + WAL_FRAME_HEADER_SIZE);
607 });
608
609 let h3 = thread::spawn(move || {
610 let buf = pool3.allocate(1024);
611 assert_eq!(buf.len(), 1024);
612 });
613
614 h1.join().unwrap();
615 h2.join().unwrap();
616 h3.join().unwrap();
617 },
618 1000,
619 );
620 }
621
622 /// Stress test: many threads allocating and dropping buffers rapidly.
623 /// This helps find race conditions in the slot bitmap and arena management.
624 #[test]
625 fn shuttle_stress_concurrent_alloc_drop() {
626 shuttle::check_random(
627 || {
628 let pool = create_test_pool();
629 let mut handles = vec![];
630
631 for i in 0..4 {
632 let pool = Arc::clone(&pool);
633 let h = thread::spawn(move || {
634 // Each thread does multiple alloc/drop cycles
635 for _ in 0..2 {
636 let buf = pool.get_page();
637 // Write thread-specific data
638 buf.as_mut_slice()[0] = i as u8;
639 assert_eq!(buf.as_slice()[0], i as u8);
640 // buf dropped here, returning slot to arena
641 }
642 });
643 handles.push(h);
644 }
645
646 for h in handles {
647 h.join().unwrap();
648 }
649 },
650 1000,
651 );
652 }
653
654 /// Test that buffers allocated by one thread can be safely read by another.
655 /// Uses a channel-like pattern with Arc to share buffers.
656 #[test]
657 fn shuttle_buffer_shared_read() {
658 shuttle::check_random(
659 || {
660 let pool = create_test_pool();
661
662 // Allocate and write in main thread
663 let buf = pool.get_page();
664 for (i, byte) in buf.as_mut_slice().iter_mut().enumerate().take(100) {
665 *byte = (i % 256) as u8;
666 }
667
668 // Wrap in Arc for shared access (Buffer itself doesn't impl Clone)
669 let buf = Arc::new(buf);
670 let buf2 = Arc::clone(&buf);
671
672 // Reader thread
673 let h = thread::spawn(move || {
674 for i in 0..100 {
675 assert_eq!(buf2.as_slice()[i], (i % 256) as u8);
676 }
677 });
678
679 // Main thread also reads
680 for i in 0..100 {
681 assert_eq!(buf.as_slice()[i], (i % 256) as u8);
682 }
683
684 h.join().unwrap();
685 },
686 1000,
687 );
688 }
689
690 /// Test pool initialization race (though guarded by init_lock).
691 /// Multiple threads trying to finalize should be safe.
692 #[test]
693 fn shuttle_concurrent_finalize() {
694 let test = || {
695 let io: Arc<dyn IO> = Arc::new(MemoryIO::new());
696 let pool = BufferPool::begin_init(&io, BufferPool::TEST_ARENA_SIZE);
697 let pool2 = Arc::clone(&pool);
698 let pool3 = Arc::clone(&pool);
699
700 let h1 = thread::spawn(move || {
701 let _ = pool.finalize_with_page_size(4096).unwrap();
702 });
703
704 let h2 = thread::spawn(move || {
705 let _ = pool2.finalize_with_page_size(4096).unwrap();
706 });
707
708 // Also try to allocate while finalizing
709 let h3 = thread::spawn(move || {
710 // This may get a temporary buffer if arena isn't ready
711 let buf = pool3.allocate(4096);
712 assert_eq!(buf.len(), 4096);
713 });
714
715 h1.join().unwrap();
716 h2.join().unwrap();
717 h3.join().unwrap();
718 };
719 shuttle::check_random(test, 1000);
720 }
721
722 /// Test concurrent writes to the same ArenaBuffer at the SAME offsets.
723 /// This exercises the `unsafe impl Sync for ArenaBuffer` which is documented as potentially unsound.
724 /// Each thread writes a distinct pattern to all 4096 bytes, and we verify no torn writes occurred
725 /// (all bytes must have the same pattern value, not a mix from different threads).
726 #[test]
727 fn shuttle_concurrent_write_same_buffer_same_offset() {
728 shuttle::check_random(
729 || {
730 let pool = create_test_pool();
731 let buf = pool.get_page();
732
733 // Three distinct byte patterns that threads will race to write
734 const PATTERN_A: u8 = 0xAA;
735 const PATTERN_B: u8 = 0xBB;
736 const PATTERN_C: u8 = 0xCC;
737
738 // Use scoped threads so buf can be borrowed by multiple threads
739 thread::scope(|scope| {
740 // Thread A writes PATTERN_A to all 4096 bytes
741 scope.spawn(|| {
742 buf.as_mut_slice().fill(PATTERN_A);
743 });
744
745 // Thread B writes PATTERN_B to all 4096 bytes
746 scope.spawn(|| {
747 buf.as_mut_slice().fill(PATTERN_B);
748 });
749
750 // Thread C writes PATTERN_C to all 4096 bytes
751 scope.spawn(|| {
752 buf.as_mut_slice().fill(PATTERN_C);
753 });
754 });
755
756 // After all writes complete, verify no torn writes across the entire buffer
757 // All 4096 bytes should have the same pattern (whichever thread won)
758 let slice = buf.as_slice();
759 let first_byte = slice[0];
760
761 // First byte must be one of our patterns
762 assert!(
763 first_byte == PATTERN_A || first_byte == PATTERN_B || first_byte == PATTERN_C,
764 "Invalid pattern in buffer: 0x{:02X}",
765 first_byte
766 );
767
768 // All bytes must match the first byte (no partial/torn writes)
769 for (i, &byte) in slice.iter().enumerate() {
770 assert!(
771 byte == first_byte,
772 "Torn write at offset {}: got 0x{:02X}, expected 0xAA, 0xBB, or 0xCC",
773 i,
774 byte
775 );
776 }
777 },
778 1000,
779 );
780 }
781
782 /// Test concurrent writes to different offsets of the same buffer (non-overlapping).
783 /// This tests that writes to different parts of the buffer don't interfere.
784 #[test]
785 fn shuttle_concurrent_write_different_offsets() {
786 shuttle::check_random(
787 || {
788 let pool = create_test_pool();
789 let buf = pool.get_page();
790
791 let ptr = buf.as_ptr() as usize;
792 let len = buf.len();
793 let _buf = buf;
794
795 let h1 = thread::spawn(move || {
796 let slice = unsafe { std::slice::from_raw_parts_mut(ptr as *mut u8, len) };
797 for i in 0..100 {
798 slice[i] = 0xAA;
799 }
800 });
801
802 let h2 = thread::spawn(move || {
803 let slice = unsafe { std::slice::from_raw_parts_mut(ptr as *mut u8, len) };
804 for i in 100..200 {
805 slice[i] = 0xBB;
806 }
807 });
808
809 let h3 = thread::spawn(move || {
810 let slice = unsafe { std::slice::from_raw_parts_mut(ptr as *mut u8, len) };
811 for i in 200..300 {
812 slice[i] = 0xCC;
813 }
814 });
815
816 h1.join().unwrap();
817 h2.join().unwrap();
818 h3.join().unwrap();
819
820 // Verify each section has the correct pattern
821 let final_slice = unsafe { std::slice::from_raw_parts(ptr as *const u8, len) };
822 for i in 0..100 {
823 assert_eq!(
824 final_slice[i], 0xAA,
825 "Section 1 corrupted at offset {}: expected 0xAA, got 0x{:02X}",
826 i, final_slice[i]
827 );
828 }
829 for i in 100..200 {
830 assert_eq!(
831 final_slice[i], 0xBB,
832 "Section 2 corrupted at offset {}: expected 0xBB, got 0x{:02X}",
833 i, final_slice[i]
834 );
835 }
836 for i in 200..300 {
837 assert_eq!(
838 final_slice[i], 0xCC,
839 "Section 3 corrupted at offset {}: expected 0xCC, got 0x{:02X}",
840 i, final_slice[i]
841 );
842 }
843 },
844 1000,
845 );
846 }
847
848 /// Test allocation racing with deallocation (slot recycling).
849 /// Verifies that when a buffer is dropped and its slot is freed,
850 /// concurrent allocations correctly handle the recycled slot.
851 #[test]
852 fn shuttle_alloc_during_drop_slot_recycling() {
853 shuttle::check_random(
854 || {
855 let pool = create_test_pool();
856
857 // Pre-allocate some buffers and write identifying data
858 let mut initial_bufs: Vec<_> = (0..5).map(|_| pool.get_page()).collect();
859 let initial_ptrs: Vec<usize> =
860 initial_bufs.iter().map(|b| b.as_ptr() as usize).collect();
861
862 // Write unique patterns to initial buffers
863 for (i, buf) in initial_bufs.iter_mut().enumerate() {
864 buf.as_mut_slice()[0] = 0xDE;
865 buf.as_mut_slice()[1] = i as u8;
866 }
867
868 let pool2 = Arc::clone(&pool);
869 let pool3 = Arc::clone(&pool);
870
871 // Thread 1: drops buffers, freeing slots
872 let h1 = thread::spawn(move || {
873 drop(initial_bufs);
874 });
875
876 // Thread 2: allocates while slots are being freed
877 let h2 = thread::spawn(move || {
878 let mut bufs = Vec::new();
879 for i in 0..3 {
880 let buf = pool2.get_page();
881 assert_eq!(buf.len(), 4096, "Buffer {} has wrong length", i);
882 // Write identifying pattern
883 buf.as_mut_slice()[0] = 0xAA;
884 buf.as_mut_slice()[1] = i as u8;
885 bufs.push(buf);
886 }
887 bufs
888 });
889
890 // Thread 3: also allocates concurrently
891 let h3 = thread::spawn(move || {
892 let mut bufs = Vec::new();
893 for i in 0..3 {
894 let buf = pool3.get_page();
895 assert_eq!(buf.len(), 4096, "Buffer {} has wrong length", i);
896 // Write different identifying pattern
897 buf.as_mut_slice()[0] = 0xBB;
898 buf.as_mut_slice()[1] = i as u8;
899 bufs.push(buf);
900 }
901 bufs
902 });
903
904 h1.join().unwrap();
905 let bufs2 = h2.join().unwrap();
906 let bufs3 = h3.join().unwrap();
907
908 // Verify buffers within each thread don't overlap
909 let ptrs2: HashSet<_> = bufs2.iter().map(|b| b.as_ptr() as usize).collect();
910 assert_eq!(
911 ptrs2.len(),
912 bufs2.len(),
913 "Thread 2 got duplicate buffer pointers"
914 );
915
916 let ptrs3: HashSet<_> = bufs3.iter().map(|b| b.as_ptr() as usize).collect();
917 assert_eq!(
918 ptrs3.len(),
919 bufs3.len(),
920 "Thread 3 got duplicate buffer pointers"
921 );
922
923 // Verify no overlap between buffers from different threads
924 for ptr in &ptrs2 {
925 assert!(
926 !ptrs3.contains(ptr),
927 "Slot double-allocation: same memory 0x{:X} returned to both threads",
928 ptr
929 );
930 }
931
932 // Verify each buffer has correct identifying data (not corrupted)
933 for (i, buf) in bufs2.iter().enumerate() {
934 assert_eq!(
935 buf.as_slice()[0],
936 0xAA,
937 "Thread 2 buffer {} header corrupted",
938 i
939 );
940 assert_eq!(
941 buf.as_slice()[1],
942 i as u8,
943 "Thread 2 buffer {} index corrupted",
944 i
945 );
946 }
947 for (i, buf) in bufs3.iter().enumerate() {
948 assert_eq!(
949 buf.as_slice()[0],
950 0xBB,
951 "Thread 3 buffer {} header corrupted",
952 i
953 );
954 assert_eq!(
955 buf.as_slice()[1],
956 i as u8,
957 "Thread 3 buffer {} index corrupted",
958 i
959 );
960 }
961
962 // Verify we can still allocate after all this
963 let final_buf = pool.get_page();
964 assert_eq!(final_buf.len(), 4096, "Final allocation failed");
965
966 // Keep initial_ptrs to suppress unused warning
967 let _ = initial_ptrs;
968 },
969 1000,
970 );
971 }
972
973 /// Test arena exhaustion and recovery.
974 /// Allocates until the arena is full (falls back to temporary buffers),
975 /// then frees and verifies slots are correctly recycled.
976 #[test]
977 fn shuttle_arena_exhaustion_and_recovery() {
978 shuttle::check_random(
979 || {
980 let pool = create_test_pool();
981
982 // Allocate many buffers to exhaust the arena
983 // TEST_ARENA_SIZE = 1MB, page_size = 4KB, so ~256 slots max
984 let mut buffers: Vec<Buffer> = Vec::new();
985 let mut pooled_count = 0;
986 let mut temp_count = 0;
987
988 for i in 0..300 {
989 let buf = pool.get_page();
990 assert_eq!(buf.len(), 4096, "Buffer {} has wrong length", i);
991
992 // Write identifying data
993 buf.as_mut_slice()[0] = (i & 0xFF) as u8;
994 buf.as_mut_slice()[1] = ((i >> 8) & 0xFF) as u8;
995
996 if buf.is_pooled() {
997 pooled_count += 1;
998 } else {
999 temp_count += 1;
1000 }
1001 buffers.push(buf);
1002 }
1003
1004 assert!(temp_count > 0);
1005 // We should have some pooled and some temporary buffers
1006 assert!(pooled_count > 0, "Expected some pooled buffers, got none");
1007 // With 1MB arena and 4KB pages, we have ~256 slots
1008 // So with 300 allocations, we should have some temporary
1009 assert!(
1010 pooled_count <= 256,
1011 "Got {} pooled buffers, but arena should only have ~256 slots",
1012 pooled_count
1013 );
1014 assert!(pooled_count + temp_count >= 256);
1015
1016 // Verify all buffers have correct identifying data
1017 for (i, buf) in buffers.iter().enumerate() {
1018 assert_eq!(
1019 buf.as_slice()[0],
1020 (i & 0xFF) as u8,
1021 "Buffer {} low byte corrupted",
1022 i
1023 );
1024 assert_eq!(
1025 buf.as_slice()[1],
1026 ((i >> 8) & 0xFF) as u8,
1027 "Buffer {} high byte corrupted",
1028 i
1029 );
1030 }
1031
1032 // Drop half the buffers to free slots
1033 let dropped_count = buffers.len() - 150;
1034 buffers.truncate(150);
1035
1036 // Allocate again - should get recycled slots
1037 let pool2 = Arc::clone(&pool);
1038 let h = thread::spawn(move || {
1039 let mut new_bufs = Vec::new();
1040 for i in 0..50 {
1041 let buf = pool2.get_page();
1042 assert_eq!(buf.len(), 4096, "New buffer {} has wrong length", i);
1043 // Write new pattern
1044 buf.as_mut_slice()[0] = 0xFF;
1045 buf.as_mut_slice()[1] = i as u8;
1046 new_bufs.push(buf);
1047 }
1048 new_bufs
1049 });
1050
1051 let new_bufs = h.join().unwrap();
1052
1053 // Verify new buffers
1054 for (i, buf) in new_bufs.iter().enumerate() {
1055 assert_eq!(buf.len(), 4096, "New buffer {} length check failed", i);
1056 assert_eq!(buf.as_slice()[0], 0xFF, "New buffer {} header corrupted", i);
1057 assert_eq!(
1058 buf.as_slice()[1],
1059 i as u8,
1060 "New buffer {} index corrupted",
1061 i
1062 );
1063 }
1064
1065 // Verify remaining original buffers still have correct data
1066 for (i, buf) in buffers.iter().enumerate() {
1067 assert_eq!(
1068 buf.as_slice()[0],
1069 (i & 0xFF) as u8,
1070 "Original buffer {} corrupted after recycling",
1071 i
1072 );
1073 }
1074
1075 let _ = (temp_count, dropped_count); // suppress warnings
1076 },
1077 1000,
1078 );
1079 }
1080
1081 /// Test that allocated buffers never overlap (slot double-allocation detection).
1082 /// Multiple threads allocate concurrently and we verify all pointers are unique.
1083 #[test]
1084 fn shuttle_slot_overlap_verification() {
1085 shuttle::check_random(
1086 || {
1087 let pool = create_test_pool();
1088 let mut handles = vec![];
1089
1090 for thread_id in 0..4u8 {
1091 let pool = Arc::clone(&pool);
1092 let h = thread::spawn(move || {
1093 let mut bufs = Vec::new();
1094 for buf_id in 0..10u8 {
1095 let buf = pool.get_page();
1096 assert_eq!(buf.len(), 4096, "Buffer has wrong length");
1097
1098 // Write thread and buffer identifying data
1099 buf.as_mut_slice()[0] = thread_id;
1100 buf.as_mut_slice()[1] = buf_id;
1101 // Write a checksum pattern
1102 buf.as_mut_slice()[2] = thread_id ^ buf_id;
1103
1104 bufs.push(buf);
1105 }
1106 bufs
1107 });
1108 handles.push(h);
1109 }
1110
1111 let all_bufs: Vec<Vec<Buffer>> =
1112 handles.into_iter().map(|h| h.join().unwrap()).collect();
1113
1114 // Verify each thread got the expected number of buffers
1115 for (thread_id, thread_bufs) in all_bufs.iter().enumerate() {
1116 assert_eq!(
1117 thread_bufs.len(),
1118 10,
1119 "Thread {} got {} buffers instead of 10",
1120 thread_id,
1121 thread_bufs.len()
1122 );
1123 }
1124
1125 // Collect all pointers and verify uniqueness
1126 let mut all_ptrs: Vec<usize> = Vec::new();
1127 for thread_bufs in &all_bufs {
1128 for buf in thread_bufs {
1129 all_ptrs.push(buf.as_ptr() as usize);
1130 }
1131 }
1132
1133 let unique_ptrs: HashSet<_> = all_ptrs.iter().copied().collect();
1134 assert_eq!(
1135 all_ptrs.len(),
1136 unique_ptrs.len(),
1137 "Slot double-allocation detected: {} total buffers but only {} unique pointers",
1138 all_ptrs.len(),
1139 unique_ptrs.len()
1140 );
1141
1142 // Verify each buffer still has correct identifying data (no cross-thread corruption)
1143 for (thread_id, thread_bufs) in all_bufs.iter().enumerate() {
1144 for (buf_id, buf) in thread_bufs.iter().enumerate() {
1145 assert_eq!(
1146 buf.as_slice()[0],
1147 thread_id as u8,
1148 "Buffer [{},{}] thread_id corrupted: expected {}, got {}",
1149 thread_id,
1150 buf_id,
1151 thread_id,
1152 buf.as_slice()[0]
1153 );
1154 assert_eq!(
1155 buf.as_slice()[1],
1156 buf_id as u8,
1157 "Buffer [{},{}] buf_id corrupted: expected {}, got {}",
1158 thread_id,
1159 buf_id,
1160 buf_id,
1161 buf.as_slice()[1]
1162 );
1163 let expected_checksum = (thread_id as u8) ^ (buf_id as u8);
1164 assert_eq!(
1165 buf.as_slice()[2],
1166 expected_checksum,
1167 "Buffer [{},{}] checksum corrupted: expected {}, got {}",
1168 thread_id,
1169 buf_id,
1170 expected_checksum,
1171 buf.as_slice()[2]
1172 );
1173 }
1174 }
1175
1176 // Verify buffers don't overlap by checking memory ranges
1177 let page_size = 4096usize;
1178 for (i, ptr_i) in all_ptrs.iter().enumerate() {
1179 for (j, ptr_j) in all_ptrs.iter().enumerate() {
1180 if i != j {
1181 let range_i = *ptr_i..(*ptr_i + page_size);
1182 // Check if ptr_j falls within range_i
1183 assert!(
1184 !range_i.contains(ptr_j),
1185 "Buffer {} (0x{:X}) overlaps with buffer {} (0x{:X})",
1186 i,
1187 ptr_i,
1188 j,
1189 ptr_j
1190 );
1191 }
1192 }
1193 }
1194 },
1195 1000,
1196 );
1197 }
1198
1199 /// Test buffer content integrity under concurrent operations.
1200 /// Each thread writes a unique pattern and verifies it's not corrupted
1201 /// by other threads' operations.
1202 #[test]
1203 fn shuttle_buffer_content_integrity() {
1204 shuttle::check_random(
1205 || {
1206 let pool = create_test_pool();
1207 let mut handles = vec![];
1208
1209 for thread_id in 0u8..4 {
1210 let pool = Arc::clone(&pool);
1211 let h = thread::spawn(move || {
1212 let buf = pool.get_page();
1213 // Write thread-specific pattern
1214 let pattern = thread_id.wrapping_mul(37);
1215 for byte in buf.as_mut_slice().iter_mut() {
1216 *byte = pattern;
1217 }
1218 // Yield to allow other threads to run
1219 thread::yield_now();
1220 // Verify pattern is intact
1221 for (i, byte) in buf.as_slice().iter().enumerate() {
1222 assert_eq!(
1223 *byte, pattern,
1224 "Buffer corruption at offset {}: expected {}, got {}",
1225 i, pattern, *byte
1226 );
1227 }
1228 buf
1229 });
1230 handles.push(h);
1231 }
1232
1233 // All threads should complete without corruption
1234 for h in handles {
1235 h.join().unwrap();
1236 }
1237 },
1238 1000,
1239 );
1240 }
1241
1242 /// Test the race between ArenaBuffer::drop upgrading Weak<Arena> while
1243 /// the Arena might be getting dropped. This exercises the weak reference
1244 /// pattern used for buffer deallocation.
1245 #[test]
1246 fn shuttle_weak_reference_upgrade_during_drop() {
1247 shuttle::check_random(
1248 || {
1249 let pool = create_test_pool();
1250
1251 // Allocate multiple buffers and write identifying data
1252 let buf1 = pool.get_page();
1253 let buf2 = pool.get_page();
1254 let buf3 = pool.get_page();
1255
1256 buf1.as_mut_slice()[0] = 0x11;
1257 buf2.as_mut_slice()[0] = 0x22;
1258 buf3.as_mut_slice()[0] = 0x33;
1259
1260 let buf1_ptr = buf1.as_ptr() as usize;
1261 let buf2_ptr = buf2.as_ptr() as usize;
1262
1263 // Clone pool references
1264 let pool2 = Arc::clone(&pool);
1265 let pool3 = Arc::clone(&pool);
1266
1267 // Thread 1: drop buffer1, freeing its slot
1268 let h1 = thread::spawn(move || {
1269 // Buffer drop will try to upgrade Weak<Arena> and call free()
1270 drop(buf1);
1271 });
1272
1273 // Thread 2: drop pool reference and buffer2
1274 let h2 = thread::spawn(move || {
1275 drop(buf2);
1276 drop(pool2);
1277 });
1278
1279 // Thread 3: allocate while others are dropping
1280 let h3 = thread::spawn(move || {
1281 let new_buf = pool3.get_page();
1282 assert_eq!(new_buf.len(), 4096, "New buffer has wrong length");
1283 new_buf.as_mut_slice()[0] = 0x44;
1284 new_buf
1285 });
1286
1287 h1.join().unwrap();
1288 h2.join().unwrap();
1289 let new_buf = h3.join().unwrap();
1290
1291 // Verify buf3 is still intact
1292 assert_eq!(buf3.as_slice()[0], 0x33, "buf3 was corrupted");
1293
1294 // Verify new_buf has correct data
1295 assert_eq!(new_buf.as_slice()[0], 0x44, "new_buf was corrupted");
1296
1297 // Original pool reference keeps arena alive
1298 // Allocate more to verify arena is still functional
1299 let final_buf = pool.get_page();
1300 assert_eq!(final_buf.len(), 4096, "Final allocation failed");
1301 final_buf.as_mut_slice()[0] = 0x55;
1302 assert_eq!(final_buf.as_slice()[0], 0x55, "Final buffer write failed");
1303
1304 // The recycled slot might be one of the dropped buffers
1305 let final_ptr = final_buf.as_ptr() as usize;
1306 // This is valid - we might get a recycled slot
1307 let _ = (buf1_ptr, buf2_ptr, final_ptr);
1308 },
1309 1000,
1310 );
1311 }
1312
1313 /// Test bitmap consistency: after many concurrent operations,
1314 /// verify that allocated_slots matches actual allocations.
1315 #[test]
1316 fn shuttle_bitmap_consistency() {
1317 shuttle::check_random(
1318 || {
1319 let pool = create_test_pool();
1320 let mut handles = vec![];
1321
1322 // Many threads doing alloc/free cycles
1323 for thread_id in 0..4u8 {
1324 let pool = Arc::clone(&pool);
1325 let h = thread::spawn(move || {
1326 let mut bufs = Vec::new();
1327 // Allocate 5 buffers
1328 for i in 0..5u8 {
1329 let buf = pool.get_page();
1330 assert_eq!(
1331 buf.len(),
1332 4096,
1333 "Thread {} buf {} wrong length",
1334 thread_id,
1335 i
1336 );
1337 // Mark with identifying data
1338 buf.as_mut_slice()[0] = thread_id;
1339 buf.as_mut_slice()[1] = i;
1340 buf.as_mut_slice()[2] = 0xAA; // Initial marker
1341 bufs.push(buf);
1342 }
1343
1344 // Verify all 5 before truncation
1345 for (i, buf) in bufs.iter().enumerate() {
1346 assert_eq!(
1347 buf.as_slice()[0],
1348 thread_id,
1349 "Pre-truncate thread_id mismatch"
1350 );
1351 assert_eq!(buf.as_slice()[1], i as u8, "Pre-truncate index mismatch");
1352 }
1353
1354 // Free 3 buffers (keep first 2)
1355 bufs.truncate(2);
1356
1357 // Allocate 3 more
1358 for i in 0..3u8 {
1359 let buf = pool.get_page();
1360 assert_eq!(
1361 buf.len(),
1362 4096,
1363 "Thread {} new buf {} wrong length",
1364 thread_id,
1365 i
1366 );
1367 buf.as_mut_slice()[0] = thread_id;
1368 buf.as_mut_slice()[1] = 10 + i; // Different index range
1369 buf.as_mut_slice()[2] = 0xBB; // New marker
1370 bufs.push(buf);
1371 }
1372
1373 // Should have 5 buffers now (2 original + 3 new)
1374 assert_eq!(bufs.len(), 5, "Thread {} should have 5 buffers", thread_id);
1375 bufs
1376 });
1377 handles.push(h);
1378 }
1379
1380 let all_bufs: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();
1381
1382 // Verify each thread returned 5 buffers
1383 for (thread_id, thread_bufs) in all_bufs.iter().enumerate() {
1384 assert_eq!(
1385 thread_bufs.len(),
1386 5,
1387 "Thread {} returned {} buffers instead of 5",
1388 thread_id,
1389 thread_bufs.len()
1390 );
1391 }
1392
1393 // Count pooled vs temporary buffers
1394 let mut pooled_count = 0;
1395 let mut temp_count = 0;
1396 for thread_bufs in &all_bufs {
1397 for buf in thread_bufs {
1398 if buf.fixed_id().is_some() {
1399 pooled_count += 1;
1400 } else {
1401 temp_count += 1;
1402 }
1403 }
1404 }
1405
1406 // Total should be 20 (4 threads * 5 buffers)
1407 assert_eq!(
1408 pooled_count + temp_count,
1409 20,
1410 "Total buffer count mismatch: {} pooled + {} temp != 20",
1411 pooled_count,
1412 temp_count
1413 );
1414
1415 // Verify all buffers have valid identifying data
1416 for (thread_id, thread_bufs) in all_bufs.iter().enumerate() {
1417 for (i, buf) in thread_bufs.iter().enumerate() {
1418 assert_eq!(
1419 buf.as_slice()[0],
1420 thread_id as u8,
1421 "Buffer [{},{}] thread_id corrupted",
1422 thread_id,
1423 i
1424 );
1425 let marker = buf.as_slice()[2];
1426 assert!(
1427 marker == 0xAA || marker == 0xBB,
1428 "Buffer [{},{}] has invalid marker 0x{:02X}",
1429 thread_id,
1430 i,
1431 marker
1432 );
1433 }
1434 }
1435
1436 // Collect all pointers to verify no duplicates
1437 let all_ptrs: HashSet<_> = all_bufs
1438 .iter()
1439 .flat_map(|bufs| bufs.iter().map(|b| b.as_ptr() as usize))
1440 .collect();
1441 assert_eq!(
1442 all_ptrs.len(),
1443 20,
1444 "Found duplicate pointers: {} unique out of 20",
1445 all_ptrs.len()
1446 );
1447
1448 // Try allocating more to verify arena is consistent
1449 let mut final_bufs = Vec::new();
1450 for i in 0..10 {
1451 let buf = pool.get_page();
1452 assert_eq!(buf.len(), 4096, "Final buf {} wrong length", i);
1453 buf.as_mut_slice()[0] = 0xFF;
1454 buf.as_mut_slice()[1] = i as u8;
1455 final_bufs.push(buf);
1456 }
1457
1458 // Verify final buffers don't overlap with existing ones
1459 for buf in &final_bufs {
1460 let ptr = buf.as_ptr() as usize;
1461 assert!(
1462 !all_ptrs.contains(&ptr),
1463 "Final buffer overlaps with existing at 0x{:X}",
1464 ptr
1465 );
1466 }
1467
1468 // Keep buffers alive until end
1469 drop(all_bufs);
1470 drop(final_bufs);
1471 },
1472 1000,
1473 );
1474 }
1475
1476 /// Test concurrent access through inner_mut().
1477 /// Multiple threads calling get_page() and get_wal_frame() simultaneously
1478 /// access different PoolInner fields through the unsynchronized inner_mut().
1479 #[test]
1480 fn shuttle_concurrent_inner_mut_access() {
1481 shuttle::check_random(
1482 || {
1483 let pool = create_test_pool();
1484 let mut handles = vec![];
1485
1486 // Threads calling get_page (accesses page_arena through inner_mut)
1487 for page_thread_id in 0..2u8 {
1488 let pool = Arc::clone(&pool);
1489 let h = thread::spawn(move || {
1490 let mut bufs = Vec::new();
1491 for i in 0..5u8 {
1492 let buf = pool.get_page();
1493 assert_eq!(buf.len(), 4096, "Page buffer has wrong length");
1494 // Mark as page buffer with identifying data
1495 buf.as_mut_slice()[0] = 0xAA; // Page marker
1496 buf.as_mut_slice()[1] = page_thread_id;
1497 buf.as_mut_slice()[2] = i;
1498 bufs.push(buf);
1499 }
1500 bufs
1501 });
1502 handles.push(h);
1503 }
1504
1505 // Threads calling get_wal_frame (accesses wal_frame_arena through inner_mut)
1506 for wal_thread_id in 0..2u8 {
1507 let pool = Arc::clone(&pool);
1508 let h = thread::spawn(move || {
1509 let mut bufs = Vec::new();
1510 for i in 0..5u8 {
1511 let buf = pool.get_wal_frame();
1512 assert_eq!(
1513 buf.len(),
1514 4096 + WAL_FRAME_HEADER_SIZE,
1515 "WAL frame buffer has wrong length"
1516 );
1517 // Mark as WAL buffer with identifying data
1518 buf.as_mut_slice()[0] = 0xBB; // WAL marker
1519 buf.as_mut_slice()[1] = wal_thread_id;
1520 buf.as_mut_slice()[2] = i;
1521 bufs.push(buf);
1522 }
1523 bufs
1524 });
1525 handles.push(h);
1526 }
1527
1528 // Thread calling allocate with various sizes
1529 {
1530 let pool = Arc::clone(&pool);
1531 let h = thread::spawn(move || {
1532 let mut bufs = Vec::new();
1533 for i in 0..5u8 {
1534 let buf = pool.allocate(2048);
1535 assert_eq!(buf.len(), 2048, "Allocated buffer has wrong length");
1536 // Mark as generic allocation
1537 buf.as_mut_slice()[0] = 0xCC; // Allocate marker
1538 buf.as_mut_slice()[1] = i;
1539 bufs.push(buf);
1540 }
1541 bufs
1542 });
1543 handles.push(h);
1544 }
1545
1546 let results: Vec<Vec<Buffer>> =
1547 handles.into_iter().map(|h| h.join().unwrap()).collect();
1548
1549 // Verify we got expected number of buffers from each type
1550 // 2 page threads * 5 + 2 wal threads * 5 + 1 allocate thread * 5 = 25 total
1551 let total_bufs: usize = results.iter().map(|v| v.len()).sum();
1552 assert_eq!(
1553 total_bufs, 25,
1554 "Expected 25 total buffers, got {}",
1555 total_bufs
1556 );
1557
1558 // Verify each buffer has correct marker and data
1559 for (thread_idx, thread_bufs) in results.iter().enumerate() {
1560 for (buf_idx, buf) in thread_bufs.iter().enumerate() {
1561 let marker = buf.as_slice()[0];
1562 assert!(
1563 marker == 0xAA || marker == 0xBB || marker == 0xCC,
1564 "Buffer [{},{}] has invalid marker 0x{:02X}",
1565 thread_idx,
1566 buf_idx,
1567 marker
1568 );
1569
1570 // Verify length matches marker type
1571 match marker {
1572 0xAA => assert_eq!(buf.len(), 4096, "Page buffer wrong length"),
1573 0xBB => assert_eq!(
1574 buf.len(),
1575 4096 + WAL_FRAME_HEADER_SIZE,
1576 "WAL buffer wrong length"
1577 ),
1578 0xCC => assert_eq!(buf.len(), 2048, "Allocate buffer wrong length"),
1579 _ => unreachable!(),
1580 }
1581 }
1582 }
1583
1584 // Collect all pointers and verify no duplicates
1585 let all_ptrs: HashSet<_> = results
1586 .iter()
1587 .flat_map(|bufs| bufs.iter().map(|b| b.as_ptr() as usize))
1588 .collect();
1589 assert_eq!(
1590 all_ptrs.len(),
1591 total_bufs,
1592 "Found duplicate pointers: {} unique out of {}",
1593 all_ptrs.len(),
1594 total_bufs
1595 );
1596 },
1597 1000,
1598 );
1599 }
1600
1601 /// Stress test with higher iteration count and more threads.
1602 /// This provides better coverage for subtle race conditions.
1603 #[test]
1604 fn shuttle_high_contention_stress() {
1605 shuttle::check_random(
1606 || {
1607 let pool = create_test_pool();
1608 let mut handles = vec![];
1609
1610 for thread_id in 0..6u8 {
1611 let pool = Arc::clone(&pool);
1612 let h = thread::spawn(move || {
1613 let mut bufs = Vec::new();
1614 let mut dropped_count = 0u8;
1615
1616 for iter in 0..4u8 {
1617 // Alternate between page and WAL frame allocations
1618 let is_page = (thread_id + iter) % 2 == 0;
1619 let buf = if is_page {
1620 pool.get_page()
1621 } else {
1622 pool.get_wal_frame()
1623 };
1624
1625 // Verify length matches allocation type
1626 let expected_len = if is_page {
1627 4096
1628 } else {
1629 4096 + WAL_FRAME_HEADER_SIZE
1630 };
1631 assert_eq!(
1632 buf.len(),
1633 expected_len,
1634 "Thread {} iter {} got wrong buffer length",
1635 thread_id,
1636 iter
1637 );
1638
1639 // Write identifying data with checksum
1640 buf.as_mut_slice()[0] = thread_id;
1641 buf.as_mut_slice()[1] = iter;
1642 buf.as_mut_slice()[2] = if is_page { 0xAA } else { 0xBB };
1643 buf.as_mut_slice()[3] = thread_id ^ iter; // Checksum
1644
1645 bufs.push(buf);
1646
1647 // Occasionally drop a buffer to test recycling
1648 if bufs.len() > 2 && iter % 2 == 1 {
1649 let dropped = bufs.pop().unwrap();
1650 // Verify the dropped buffer still had valid data
1651 assert_eq!(
1652 dropped.as_slice()[0],
1653 thread_id,
1654 "Dropped buffer thread_id corrupted"
1655 );
1656 dropped_count += 1;
1657 }
1658 }
1659
1660 // Verify all remaining buffers before returning
1661 for (i, buf) in bufs.iter().enumerate() {
1662 assert_eq!(
1663 buf.as_slice()[0],
1664 thread_id,
1665 "Pre-return buffer {} thread_id corrupted",
1666 i
1667 );
1668 let checksum = buf.as_slice()[0] ^ buf.as_slice()[1];
1669 assert_eq!(
1670 buf.as_slice()[3],
1671 checksum,
1672 "Pre-return buffer {} checksum invalid",
1673 i
1674 );
1675 }
1676
1677 (bufs, dropped_count)
1678 });
1679 handles.push(h);
1680 }
1681
1682 let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();
1683
1684 // Verify results from all threads
1685 let mut total_remaining = 0;
1686 let mut total_dropped = 0u8;
1687
1688 for (thread_id, (thread_bufs, dropped)) in results.iter().enumerate() {
1689 total_remaining += thread_bufs.len();
1690 total_dropped = total_dropped.saturating_add(*dropped);
1691
1692 // Verify each buffer has correct identifying data
1693 for (buf_idx, buf) in thread_bufs.iter().enumerate() {
1694 assert_eq!(
1695 buf.as_slice()[0],
1696 thread_id as u8,
1697 "Thread {} buffer {} thread_id corrupted: expected {}, got {}",
1698 thread_id,
1699 buf_idx,
1700 thread_id,
1701 buf.as_slice()[0]
1702 );
1703
1704 let marker = buf.as_slice()[2];
1705 assert!(
1706 marker == 0xAA || marker == 0xBB,
1707 "Thread {} buffer {} has invalid marker 0x{:02X}",
1708 thread_id,
1709 buf_idx,
1710 marker
1711 );
1712
1713 // Verify checksum
1714 let expected_checksum = buf.as_slice()[0] ^ buf.as_slice()[1];
1715 assert_eq!(
1716 buf.as_slice()[3],
1717 expected_checksum,
1718 "Thread {} buffer {} checksum mismatch",
1719 thread_id,
1720 buf_idx
1721 );
1722
1723 // Verify length matches marker
1724 let expected_len = if marker == 0xAA {
1725 4096
1726 } else {
1727 4096 + WAL_FRAME_HEADER_SIZE
1728 };
1729 assert_eq!(
1730 buf.len(),
1731 expected_len,
1732 "Thread {} buffer {} length mismatch for marker",
1733 thread_id,
1734 buf_idx
1735 );
1736 }
1737 }
1738
1739 // Sanity check: we should have some buffers remaining
1740 assert!(
1741 total_remaining > 0,
1742 "No buffers remaining after stress test"
1743 );
1744
1745 // Collect all pointers and verify no duplicates among remaining buffers
1746 let all_ptrs: HashSet<_> = results
1747 .iter()
1748 .flat_map(|(bufs, _)| bufs.iter().map(|b| b.as_ptr() as usize))
1749 .collect();
1750 assert_eq!(
1751 all_ptrs.len(),
1752 total_remaining,
1753 "Found duplicate pointers: {} unique out of {} remaining",
1754 all_ptrs.len(),
1755 total_remaining
1756 );
1757
1758 // Final allocation to verify pool is still healthy
1759 let final_buf = pool.get_page();
1760 assert_eq!(final_buf.len(), 4096, "Final allocation failed");
1761 final_buf.as_mut_slice()[0] = 0xFF;
1762 assert_eq!(final_buf.as_slice()[0], 0xFF, "Final buffer write failed");
1763
1764 let _ = total_dropped; // suppress warning
1765 },
1766 1000,
1767 );
1768 }
1769}