Skip to main content

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}