use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::Duration;
use fsqlite_types::sync_primitives::{Instant, SystemTime};
use tracing::{debug, info, warn};
pub const DEFAULT_TABLE_CAPACITY: u32 = 1 << 20;
const MAX_LOAD_FACTOR: f64 = 0.70;
const DRAINING_NONE: u32 = 0xFFFF_FFFF;
const DEFAULT_LEASE_SECS: u64 = 5;
#[cfg(any(target_os = "linux", test))]
const PID_BIRTH_PROCFS_TAG: u64 = 1_u64 << 63;
const OCCUPANCY_STRIPE_COUNT: usize = 16;
const REBUILD_DRAIN_FULL_SCAN_INTERVAL: Duration = Duration::from_millis(100);
const REBUILD_DRAIN_HANDOFF_BASE_SPINS: u32 = 64;
const REBUILD_DRAIN_HANDOFF_MAX_SPINS: u32 = 2_048;
const REBUILD_DRAIN_HANDOFF_YIELD_EVERY: u32 = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct RebuildDrainWait {
spin_loops: u32,
yielded: bool,
}
#[derive(Debug, Clone, Copy, Default)]
struct RebuildDrainHandoff {
next_attempt: u32,
}
impl RebuildDrainHandoff {
fn next_wait(&mut self, started: Instant, timeout: Duration) -> Option<RebuildDrainWait> {
if started.elapsed() >= timeout {
return None;
}
let attempt = self.next_attempt.saturating_add(1);
self.next_attempt = attempt;
Some(RebuildDrainWait {
spin_loops: rebuild_drain_spin_loops(attempt),
yielded: rebuild_drain_should_yield(attempt),
})
}
}
const fn rebuild_drain_spin_loops(attempt: u32) -> u32 {
let growth = attempt.saturating_sub(1);
let shift = if growth > 5 { 5 } else { growth };
let spins = REBUILD_DRAIN_HANDOFF_BASE_SPINS << shift;
if spins > REBUILD_DRAIN_HANDOFF_MAX_SPINS {
REBUILD_DRAIN_HANDOFF_MAX_SPINS
} else {
spins
}
}
const fn rebuild_drain_should_yield(attempt: u32) -> bool {
attempt >= REBUILD_DRAIN_HANDOFF_YIELD_EVERY
&& attempt.is_multiple_of(REBUILD_DRAIN_HANDOFF_YIELD_EVERY)
}
fn perform_rebuild_drain_handoff(wait: RebuildDrainWait) {
for _ in 0..wait.spin_loops {
std::hint::spin_loop();
}
if wait.yielded {
std::thread::yield_now();
}
}
#[cfg(target_os = "linux")]
fn read_proc_start_time_ticks(pid: u32) -> Option<u64> {
let stat_path = std::path::Path::new("/proc")
.join(pid.to_string())
.join("stat");
let stat = std::fs::read_to_string(stat_path).ok()?;
let comm_end = stat.rfind(')')?;
let tail = stat.get(comm_end + 1..)?.trim_start();
tail.split_whitespace().nth(19)?.parse::<u64>().ok()
}
fn current_process_birth_token(now_fallback: u64) -> u64 {
#[cfg(target_os = "linux")]
{
if !std::path::Path::new("/proc").exists() {
return now_fallback;
}
if let Some(start_ticks) = read_proc_start_time_ticks(std::process::id()) {
return PID_BIRTH_PROCFS_TAG | (start_ticks & !PID_BIRTH_PROCFS_TAG);
}
now_fallback
}
#[cfg(any(target_os = "macos", windows))]
{
fsqlite_vfs::process::current_process_birth_token().unwrap_or(now_fallback)
}
#[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
{
now_fallback
}
}
fn process_alive_os(pid: u32, pid_birth: u64) -> bool {
#[cfg(target_os = "linux")]
{
if pid == 0 {
return false;
}
if !std::path::Path::new("/proc").exists() {
return true;
}
let proc_dir = std::path::Path::new("/proc").join(pid.to_string());
if !proc_dir.exists() {
return false;
}
if pid_birth & PID_BIRTH_PROCFS_TAG == 0 {
return true;
}
let expected_ticks = pid_birth & !PID_BIRTH_PROCFS_TAG;
read_proc_start_time_ticks(pid).is_some_and(|start_ticks| start_ticks == expected_ticks)
}
#[cfg(any(target_os = "macos", windows))]
{
!matches!(
fsqlite_vfs::process::process_alive(pid, pid_birth),
fsqlite_vfs::process::ProcessLiveness::Dead
)
}
#[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
{
let _ = (pid, pid_birth);
true
}
}
pub struct PageLockEntry {
page_number: AtomicU32,
owner_txn: AtomicU64,
}
impl PageLockEntry {
fn new() -> Self {
Self {
page_number: AtomicU32::new(0),
owner_txn: AtomicU64::new(0),
}
}
}
#[repr(align(64))]
struct CacheAlignedAtomicU32(AtomicU32);
impl CacheAlignedAtomicU32 {
const fn new(value: u32) -> Self {
Self(AtomicU32::new(value))
}
fn load(&self, ordering: Ordering) -> u32 {
self.0.load(ordering)
}
fn fetch_add(&self, value: u32, ordering: Ordering) {
self.0.fetch_add(value, ordering);
}
fn store(&self, value: u32, ordering: Ordering) {
self.0.store(value, ordering);
}
}
struct StripedOccupancyCounter {
stripes: [CacheAlignedAtomicU32; OCCUPANCY_STRIPE_COUNT],
}
impl StripedOccupancyCounter {
fn new() -> Self {
Self {
stripes: std::array::from_fn(|_| CacheAlignedAtomicU32::new(0)),
}
}
#[inline]
fn load(&self) -> u32 {
self.stripes.iter().fold(0_u32, |sum, stripe| {
sum.saturating_add(stripe.load(Ordering::Relaxed))
})
}
#[inline]
fn increment(&self, page_number: u32) {
let stripe = Self::stripe_index(page_number);
self.stripes[stripe].fetch_add(1, Ordering::Relaxed);
}
fn clear(&self) {
for stripe in &self.stripes {
stripe.store(0, Ordering::Release);
}
}
#[inline]
fn stripe_index(page_number: u32) -> usize {
let mixed = page_number.wrapping_mul(2_654_435_769);
mixed as usize & (OCCUPANCY_STRIPE_COUNT - 1)
}
}
struct LockTableInstance {
entries: Vec<PageLockEntry>,
occupied_count: StripedOccupancyCounter,
}
impl LockTableInstance {
fn new(capacity: u32) -> Self {
let entries: Vec<PageLockEntry> = (0..capacity).map(|_| PageLockEntry::new()).collect();
Self {
entries,
occupied_count: StripedOccupancyCounter::new(),
}
}
#[inline]
fn occupied_count(&self) -> u32 {
self.occupied_count.load()
}
#[inline]
fn increment_occupied(&self, page_number: u32) {
self.occupied_count.increment(page_number);
}
#[allow(dead_code)]
fn occupied_count_full_scan(&self) -> u32 {
let mut count = 0_u32;
for entry in &self.entries {
if entry.page_number.load(Ordering::Relaxed) != 0 {
count += 1;
}
}
count
}
fn locked_count(&self) -> u32 {
let mut count = 0_u32;
for entry in &self.entries {
if entry.owner_txn.load(Ordering::Relaxed) != 0 {
count += 1;
}
}
count
}
fn is_quiescent(&self) -> bool {
self.entries
.iter()
.all(|e| e.owner_txn.load(Ordering::Acquire) == 0)
}
fn clear_all(&self) {
for entry in &self.entries {
entry.page_number.store(0, Ordering::Release);
entry.owner_txn.store(0, Ordering::Release);
}
self.occupied_count.clear();
}
fn release_all_for_txn(&self, txn_id: u64) -> u32 {
let mut released = 0_u32;
for entry in &self.entries {
if entry
.owner_txn
.compare_exchange(txn_id, 0, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
{
released += 1;
}
}
released
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AcquireResult {
Acquired,
AlreadyHeld,
Busy { holder: u64 },
CapacityExhausted,
}
impl AcquireResult {
#[must_use]
pub fn is_ok(&self) -> bool {
matches!(self, Self::Acquired | Self::AlreadyHeld)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RebuildLeaseError {
LeaseHeld { pid: u32 },
NoDrainInProgress,
NotQuiescent { remaining: u32 },
TargetNotEmpty,
}
impl std::fmt::Display for RebuildLeaseError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::LeaseHeld { pid } => write!(f, "rebuild lease held by PID {pid}"),
Self::NoDrainInProgress => f.write_str("no drain in progress"),
Self::NotQuiescent { remaining } => {
write!(f, "draining table not quiescent: {remaining} locks remain")
}
Self::TargetNotEmpty => {
f.write_str("target table not empty from prior incomplete rebuild")
}
}
}
}
impl std::error::Error for RebuildLeaseError {}
pub struct SharedPageLockTable {
capacity: u32,
mask: u32,
active_table: AtomicU32,
draining_table: AtomicU32,
rebuild_pid: AtomicU32,
rebuild_pid_birth: AtomicU64,
rebuild_lease_expiry: AtomicU64,
rebuild_epoch: AtomicU32,
tables: [LockTableInstance; 2],
}
impl SharedPageLockTable {
#[must_use]
pub fn new(capacity: u32) -> Self {
assert!(
capacity > 0 && capacity.is_power_of_two(),
"capacity must be a power of two"
);
info!(capacity, "SharedPageLockTable: created");
Self {
capacity,
mask: capacity - 1,
active_table: AtomicU32::new(0),
draining_table: AtomicU32::new(DRAINING_NONE),
rebuild_pid: AtomicU32::new(0),
rebuild_pid_birth: AtomicU64::new(0),
rebuild_lease_expiry: AtomicU64::new(0),
rebuild_epoch: AtomicU32::new(0),
tables: [
LockTableInstance::new(capacity),
LockTableInstance::new(capacity),
],
}
}
#[must_use]
pub fn with_default_capacity() -> Self {
Self::new(DEFAULT_TABLE_CAPACITY)
}
#[inline]
fn hash_index(&self, page_number: u32) -> u32 {
let h = page_number.wrapping_mul(2_654_435_769);
let shift = 32 - self.capacity.trailing_zeros();
h >> shift
}
pub fn try_acquire(&self, page_number: u32, txn_id: u64) -> AcquireResult {
debug_assert!(page_number != 0, "page_number 0 is the empty sentinel");
debug_assert!(txn_id != 0, "txn_id 0 is the unlocked sentinel");
let start_epoch = self.rebuild_epoch.load(Ordering::Acquire);
let active_idx = self.active_table.load(Ordering::Acquire);
let draining_idx = self.draining_table.load(Ordering::Acquire);
if draining_idx != DRAINING_NONE {
let draining = &self.tables[draining_idx as usize];
match self.probe_for_existing(draining, page_number) {
ProbeResult::FoundOwnedBy(owner) if owner == txn_id => {
debug!(
page_number,
requester_txn_id = txn_id,
holder_txn_id = owner,
table = "draining",
lock_intent = "exclusive",
"page lock already held by requester in draining table"
);
return AcquireResult::AlreadyHeld;
}
ProbeResult::FoundOwnedBy(holder) => {
debug!(
page_number,
requester_txn_id = txn_id,
holder_txn_id = holder,
table = "draining",
lock_intent = "exclusive",
"page lock acquisition blocked by draining-table holder"
);
return AcquireResult::Busy { holder };
}
ProbeResult::FoundUnlocked | ProbeResult::NotFound => {
}
}
}
let mut active = &self.tables[active_idx as usize];
let mut occupied = active.occupied_count();
let mut at_capacity = f64::from(occupied) / f64::from(self.capacity) > MAX_LOAD_FACTOR;
if at_capacity {
self.try_start_best_effort_rebuild();
let refreshed_active_idx = self.active_table.load(Ordering::Acquire);
active = &self.tables[refreshed_active_idx as usize];
occupied = active.occupied_count();
at_capacity = f64::from(occupied) / f64::from(self.capacity) > MAX_LOAD_FACTOR;
}
debug!(
page_number,
requester_txn_id = txn_id,
lock_intent = "exclusive",
rebuild_epoch = start_epoch,
active_table = active_idx,
draining_table = ?(draining_idx != DRAINING_NONE).then_some(draining_idx),
occupied,
capacity = self.capacity,
"page lock acquisition requested"
);
let mut idx = self.hash_index(page_number);
let mut probes = 0_u32;
let first_empty;
loop {
if probes >= self.capacity {
warn!(
page_number,
requester_txn_id = txn_id,
lock_intent = "exclusive",
capacity = self.capacity,
"SharedPageLockTable: full table probe wrap"
);
return AcquireResult::CapacityExhausted;
}
let entry = &active.entries[idx as usize];
let current_page = entry.page_number.load(Ordering::Acquire);
if current_page == page_number {
return match entry.owner_txn.compare_exchange(
0,
txn_id,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
if self.rebuild_epoch.load(Ordering::Acquire) != start_epoch {
let _ = entry.owner_txn.compare_exchange(
txn_id,
0,
Ordering::AcqRel,
Ordering::Relaxed,
);
return self.try_acquire(page_number, txn_id);
}
debug!(
page_number,
requester_txn_id = txn_id,
table = "active",
slot_idx = idx,
reused_slot_key = true,
lock_intent = "exclusive",
"page lock acquired"
);
AcquireResult::Acquired
}
Err(holder) if holder == txn_id => {
debug!(
page_number,
requester_txn_id = txn_id,
holder_txn_id = holder,
table = "active",
slot_idx = idx,
lock_intent = "exclusive",
"page lock already held by requester"
);
AcquireResult::AlreadyHeld
}
Err(holder) => {
debug!(
page_number,
requester_txn_id = txn_id,
holder_txn_id = holder,
table = "active",
slot_idx = idx,
lock_intent = "exclusive",
"page lock acquisition conflicted with active holder"
);
AcquireResult::Busy { holder }
}
};
}
if current_page == 0 {
first_empty = Some(idx);
break;
}
idx = (idx + 1) & self.mask;
probes += 1;
}
if at_capacity {
warn!(
page_number,
requester_txn_id = txn_id,
lock_intent = "exclusive",
occupied,
capacity = self.capacity,
"SharedPageLockTable: capacity exhausted (load factor > 0.70)"
);
return AcquireResult::CapacityExhausted;
}
let idx = first_empty.unwrap();
let entry = &active.entries[idx as usize];
if entry
.page_number
.compare_exchange(0, page_number, Ordering::AcqRel, Ordering::Acquire)
.is_err()
{
return self.try_acquire(page_number, txn_id);
}
active.increment_occupied(page_number);
match entry
.owner_txn
.compare_exchange(0, txn_id, Ordering::AcqRel, Ordering::Acquire)
{
Ok(_) => {
if self.rebuild_epoch.load(Ordering::Acquire) != start_epoch {
let _ = entry.owner_txn.compare_exchange(
txn_id,
0,
Ordering::AcqRel,
Ordering::Relaxed,
);
return self.try_acquire(page_number, txn_id);
}
debug!(
page_number,
requester_txn_id = txn_id,
table = "active",
slot_idx = idx,
reused_slot_key = false,
lock_intent = "exclusive",
"page lock acquired"
);
AcquireResult::Acquired
}
Err(_) => {
let holder = entry.owner_txn.load(Ordering::Acquire);
debug!(
page_number,
requester_txn_id = txn_id,
holder_txn_id = holder,
table = "active",
slot_idx = idx,
lock_intent = "exclusive",
"page lock acquisition conflicted after slot claim race"
);
AcquireResult::Busy { holder }
}
}
}
fn try_start_best_effort_rebuild(&self) {
let draining_idx = self.draining_table.load(Ordering::Acquire);
if draining_idx != DRAINING_NONE {
let draining = &self.tables[draining_idx as usize];
if !draining.is_quiescent() {
return;
}
let now_secs = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map_or(0, |duration| duration.as_secs());
let pid = std::process::id();
let pid_birth = current_process_birth_token(now_secs);
if self
.acquire_rebuild_lease(pid, pid_birth, now_secs)
.is_err()
{
return;
}
if let Err(error) = self.finalize_rebuild(pid) {
warn!(
error = %error,
"SharedPageLockTable: failed to finalize quiescent draining rebuild"
);
self.rebuild_pid.store(0, Ordering::Release);
self.rebuild_pid_birth.store(0, Ordering::Release);
self.rebuild_lease_expiry.store(0, Ordering::Release);
}
return;
}
let now_secs = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map_or(0, |duration| duration.as_secs());
let pid = std::process::id();
let pid_birth = current_process_birth_token(now_secs);
let _ = self.full_rebuild(
pid,
pid_birth,
now_secs,
|_txn_id| true,
Duration::from_millis(1),
);
}
fn probe_for_existing(&self, table: &LockTableInstance, page_number: u32) -> ProbeResult {
let mut idx = self.hash_index(page_number);
let mut probes = 0_u32;
loop {
if probes >= self.capacity {
return ProbeResult::NotFound;
}
let entry = &table.entries[idx as usize];
let current_page = entry.page_number.load(Ordering::Acquire);
if current_page == page_number {
let owner = entry.owner_txn.load(Ordering::Acquire);
if owner == 0 {
return ProbeResult::FoundUnlocked;
}
return ProbeResult::FoundOwnedBy(owner);
}
if current_page == 0 {
return ProbeResult::NotFound;
}
idx = (idx + 1) & self.mask;
probes += 1;
}
}
pub fn release(&self, page_number: u32, txn_id: u64) -> bool {
let active_idx = self.active_table.load(Ordering::Acquire);
let draining_idx = self.draining_table.load(Ordering::Acquire);
if self.release_in_table(&self.tables[active_idx as usize], page_number, txn_id) {
return true;
}
if draining_idx != DRAINING_NONE {
return self.release_in_table(&self.tables[draining_idx as usize], page_number, txn_id);
}
false
}
fn release_in_table(&self, table: &LockTableInstance, page_number: u32, txn_id: u64) -> bool {
let mut idx = self.hash_index(page_number);
let mut probes = 0_u32;
loop {
if probes >= self.capacity {
return false;
}
let entry = &table.entries[idx as usize];
let current_page = entry.page_number.load(Ordering::Acquire);
if current_page == page_number {
return entry
.owner_txn
.compare_exchange(txn_id, 0, Ordering::Release, Ordering::Relaxed)
.is_ok();
}
if current_page == 0 {
return false;
}
idx = (idx + 1) & self.mask;
probes += 1;
}
}
pub fn release_set(&self, pages: &[u32], txn_id: u64) {
for &page in pages {
self.release(page, txn_id);
}
}
pub fn release_all_for_txn(&self, txn_id: u64) -> u32 {
let mut total = 0_u32;
for table in &self.tables {
total += table.release_all_for_txn(txn_id);
}
total
}
#[must_use]
pub fn holder(&self, page_number: u32) -> Option<u64> {
let active_idx = self.active_table.load(Ordering::Acquire);
let draining_idx = self.draining_table.load(Ordering::Acquire);
if let ProbeResult::FoundOwnedBy(owner) =
self.probe_for_existing(&self.tables[active_idx as usize], page_number)
{
return Some(owner);
}
if draining_idx != DRAINING_NONE
&& let ProbeResult::FoundOwnedBy(owner) =
self.probe_for_existing(&self.tables[draining_idx as usize], page_number)
{
return Some(owner);
}
None
}
pub fn acquire_rebuild_lease(
&self,
pid: u32,
pid_birth: u64,
now_secs: u64,
) -> Result<(), RebuildLeaseError> {
let lease_expires_at = now_secs + DEFAULT_LEASE_SECS;
match self
.rebuild_pid
.compare_exchange(0, pid, Ordering::AcqRel, Ordering::Acquire)
{
Ok(_) => {
self.rebuild_pid_birth.store(pid_birth, Ordering::Release);
self.rebuild_lease_expiry
.store(lease_expires_at, Ordering::Release);
info!(
pid,
pid_birth,
request_now_secs = now_secs,
lease_expires_at,
epoch = self.rebuild_epoch.load(Ordering::Relaxed),
"rebuild lease acquired"
);
Ok(())
}
Err(current_pid) => {
let expiry = self.rebuild_lease_expiry.load(Ordering::Acquire);
let current_pid_birth = self.rebuild_pid_birth.load(Ordering::Acquire);
let lease_expired = expiry <= now_secs;
let process_dead = current_pid != 0
&& current_pid_birth != 0
&& !process_alive_os(current_pid, current_pid_birth);
if lease_expired || process_dead {
match self.rebuild_pid.compare_exchange(
current_pid,
pid,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
self.rebuild_pid_birth.store(pid_birth, Ordering::Release);
self.rebuild_lease_expiry
.store(lease_expires_at, Ordering::Release);
warn!(
old_pid = current_pid,
old_pid_birth = current_pid_birth,
old_lease_expiry = expiry,
new_pid = pid,
new_pid_birth = pid_birth,
request_now_secs = now_secs,
lease_expires_at,
lease_expired,
process_dead,
"rebuild lease stolen from stale holder"
);
Ok(())
}
Err(_) => Err(RebuildLeaseError::LeaseHeld { pid: current_pid }),
}
} else {
Err(RebuildLeaseError::LeaseHeld { pid: current_pid })
}
}
}
}
pub fn renew_rebuild_lease(&self, pid: u32, now_secs: u64) -> bool {
if self.rebuild_pid.load(Ordering::Acquire) == pid {
let lease_expires_at = now_secs + DEFAULT_LEASE_SECS;
self.rebuild_lease_expiry
.store(lease_expires_at, Ordering::Release);
debug!(
pid,
request_now_secs = now_secs,
lease_expires_at,
"rebuild lease renewed"
);
true
} else {
false
}
}
pub fn rotate(&self) -> Result<(), RebuildLeaseError> {
let draining = self.draining_table.load(Ordering::Acquire);
if draining != DRAINING_NONE {
return Err(RebuildLeaseError::NoDrainInProgress);
}
let active = self.active_table.load(Ordering::Acquire);
let new_active = 1 - active;
if self.tables[new_active as usize].occupied_count() > 0 {
return Err(RebuildLeaseError::TargetNotEmpty);
}
self.draining_table.store(active, Ordering::Release);
self.active_table.store(new_active, Ordering::Release);
debug!(
old_active = active,
new_active, "SharedPageLockTable: tables rotated"
);
Ok(())
}
#[must_use]
pub fn drain_progress(&self) -> Option<DrainStatus> {
let draining_idx = self.draining_table.load(Ordering::Acquire);
if draining_idx == DRAINING_NONE {
return None;
}
let draining = &self.tables[draining_idx as usize];
let remaining = draining.locked_count();
let quiescent = remaining == 0;
Some(DrainStatus {
remaining,
quiescent,
})
}
pub fn drain_orphaned(&self, is_active_txn: impl Fn(u64) -> bool) -> u32 {
let draining_idx = self.draining_table.load(Ordering::Acquire);
if draining_idx == DRAINING_NONE {
return 0;
}
let draining = &self.tables[draining_idx as usize];
let mut cleaned = 0_u32;
for entry in &draining.entries {
let owner = entry.owner_txn.load(Ordering::Acquire);
if owner != 0 && !is_active_txn(owner) {
if entry
.owner_txn
.compare_exchange(owner, 0, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
{
cleaned += 1;
}
}
}
if cleaned > 0 {
debug!(
cleaned,
draining_table = draining_idx,
"SharedPageLockTable: orphaned locks cleaned during drain"
);
}
cleaned
}
pub fn finalize_rebuild(&self, pid: u32) -> Result<u32, RebuildLeaseError> {
let draining_idx = self.draining_table.load(Ordering::Acquire);
if draining_idx == DRAINING_NONE {
return Err(RebuildLeaseError::NoDrainInProgress);
}
let draining = &self.tables[draining_idx as usize];
if !draining.is_quiescent() {
let remaining = draining.locked_count();
return Err(RebuildLeaseError::NotQuiescent { remaining });
}
self.draining_table.store(DRAINING_NONE, Ordering::Release);
let new_epoch = self.rebuild_epoch.fetch_add(1, Ordering::AcqRel) + 1;
let cleared = draining.occupied_count();
draining.clear_all();
self.rebuild_pid.store(0, Ordering::Release);
self.rebuild_pid_birth.store(0, Ordering::Release);
self.rebuild_lease_expiry.store(0, Ordering::Release);
info!(
epoch = new_epoch,
cleared,
pid,
draining_table = draining_idx,
"SharedPageLockTable: rebuild finalized"
);
Ok(cleared)
}
#[allow(clippy::too_many_lines)]
pub fn full_rebuild(
&self,
pid: u32,
pid_birth: u64,
now_secs: u64,
is_active_txn: impl Fn(u64) -> bool,
timeout: Duration,
) -> Result<RebuildResult, RebuildLeaseError> {
let started = Instant::now();
let mut next_full_scan_at = Duration::ZERO;
let mut handoff = RebuildDrainHandoff::default();
self.acquire_rebuild_lease(pid, pid_birth, now_secs)?;
if let Err(e) = self.rotate() {
self.rebuild_pid.store(0, Ordering::Release);
return Err(e);
}
let mut orphaned_cleaned = 0_u32;
loop {
let elapsed = started.elapsed();
if elapsed >= next_full_scan_at {
orphaned_cleaned += self.drain_orphaned(&is_active_txn);
next_full_scan_at = elapsed.saturating_add(REBUILD_DRAIN_FULL_SCAN_INTERVAL);
if let Some(status) = self.drain_progress() {
if status.quiescent {
break;
}
debug!(remaining = status.remaining, "drain in progress");
}
} else {
let draining_idx = self.draining_table.load(Ordering::Acquire);
if draining_idx != DRAINING_NONE {
let draining = &self.tables[draining_idx as usize];
if draining.is_quiescent() {
break;
}
}
}
if elapsed >= timeout {
let draining_idx = self.draining_table.load(Ordering::Acquire);
let quiescent = draining_idx != DRAINING_NONE
&& self.tables[draining_idx as usize].is_quiescent();
if !quiescent {
self.rebuild_pid.store(0, Ordering::Release);
return Ok(RebuildResult {
cleared: 0,
orphaned_cleaned,
elapsed,
epoch: self.rebuild_epoch.load(Ordering::Acquire),
timed_out: true,
});
}
break;
}
let Some(wait) = handoff.next_wait(started, timeout) else {
let draining_idx = self.draining_table.load(Ordering::Acquire);
let quiescent = draining_idx != DRAINING_NONE
&& self.tables[draining_idx as usize].is_quiescent();
if !quiescent {
self.rebuild_pid.store(0, Ordering::Release);
return Ok(RebuildResult {
cleared: 0,
orphaned_cleaned,
elapsed: started.elapsed(),
epoch: self.rebuild_epoch.load(Ordering::Acquire),
timed_out: true,
});
}
break;
};
perform_rebuild_drain_handoff(wait);
}
let cleared = self.finalize_rebuild(pid)?;
Ok(RebuildResult {
cleared,
orphaned_cleaned,
elapsed: started.elapsed(),
epoch: self.rebuild_epoch.load(Ordering::Acquire),
timed_out: false,
})
}
#[must_use]
pub fn active_load_factor(&self) -> f64 {
let active_idx = self.active_table.load(Ordering::Acquire);
let occupied = self.tables[active_idx as usize].occupied_count();
f64::from(occupied) / f64::from(self.capacity)
}
#[must_use]
pub fn needs_rebuild(&self) -> bool {
self.active_load_factor() > MAX_LOAD_FACTOR
}
#[must_use]
pub fn is_rebuild_in_progress(&self) -> bool {
self.draining_table.load(Ordering::Acquire) != DRAINING_NONE
}
#[must_use]
pub fn rebuild_epoch(&self) -> u32 {
self.rebuild_epoch.load(Ordering::Acquire)
}
#[must_use]
pub fn capacity(&self) -> u32 {
self.capacity
}
#[must_use]
pub fn active_occupied(&self) -> u32 {
let active_idx = self.active_table.load(Ordering::Acquire);
self.tables[active_idx as usize].occupied_count()
}
#[must_use]
pub fn total_locked(&self) -> u32 {
self.tables[0].locked_count() + self.tables[1].locked_count()
}
}
enum ProbeResult {
FoundOwnedBy(u64),
FoundUnlocked,
NotFound,
}
#[derive(Debug, Clone, Copy)]
pub struct DrainStatus {
pub remaining: u32,
pub quiescent: bool,
}
#[derive(Debug, Clone)]
pub struct RebuildResult {
pub cleared: u32,
pub orphaned_cleaned: u32,
pub elapsed: Duration,
pub epoch: u32,
pub timed_out: bool,
}
#[cfg(test)]
mod tests {
use super::*;
use std::io;
use std::sync::Arc;
use std::sync::Mutex;
const TEST_CAP: u32 = 64;
#[derive(Clone)]
struct BufMakeWriter(Arc<Mutex<Vec<u8>>>);
impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for BufMakeWriter {
type Writer = BufWriter;
fn make_writer(&'a self) -> Self::Writer {
BufWriter(Arc::clone(&self.0))
}
}
struct BufWriter(Arc<Mutex<Vec<u8>>>);
impl io::Write for BufWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let mut guard = self.0.lock().expect("shared_lock_table log buffer lock");
guard.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
fn with_tracing_capture<F, R>(
_capture_guard: &crate::test_support::TracingCaptureGuard,
f: F,
) -> (R, String)
where
F: FnOnce() -> R,
{
let buf = Arc::new(Mutex::new(Vec::new()));
let subscriber = tracing_subscriber::fmt()
.with_ansi(false)
.with_max_level(tracing::Level::DEBUG)
.with_writer(BufMakeWriter(Arc::clone(&buf)))
.finish();
let result = tracing::subscriber::with_default(subscriber, f);
let bytes = buf
.lock()
.expect("shared_lock_table log buffer lock")
.clone();
(result, String::from_utf8_lossy(&bytes).to_string())
}
#[derive(Debug)]
struct RebuildDrainScheduleSummary {
total_spin_loops: u64,
max_spin_loops: u32,
yield_count: u32,
p95_spin_loops: u32,
full_scan_interval: Duration,
}
fn summarize_rebuild_drain_schedule(attempts: u32) -> RebuildDrainScheduleSummary {
let waits = (1..=attempts)
.map(|attempt| RebuildDrainWait {
spin_loops: rebuild_drain_spin_loops(attempt),
yielded: rebuild_drain_should_yield(attempt),
})
.collect::<Vec<_>>();
let total_spin_loops = waits.iter().map(|wait| u64::from(wait.spin_loops)).sum();
let max_spin_loops = waits.iter().map(|wait| wait.spin_loops).max().unwrap_or(0);
let yield_count = u32::try_from(waits.iter().filter(|wait| wait.yielded).count())
.expect("test wait count fits u32");
let mut spin_samples = waits.iter().map(|wait| wait.spin_loops).collect::<Vec<_>>();
spin_samples.sort_unstable();
RebuildDrainScheduleSummary {
total_spin_loops,
max_spin_loops,
yield_count,
p95_spin_loops: percentile_rebuild_drain_spin_loops(&spin_samples, 95),
full_scan_interval: REBUILD_DRAIN_FULL_SCAN_INTERVAL,
}
}
fn percentile_rebuild_drain_spin_loops(sorted_spin_loops: &[u32], percentile: u8) -> u32 {
assert!(
!sorted_spin_loops.is_empty(),
"percentile requires at least one spin sample"
);
let index = usize::from(percentile)
.saturating_mul(sorted_spin_loops.len())
.saturating_sub(1)
/ 100;
sorted_spin_loops[index.min(sorted_spin_loops.len().saturating_sub(1))]
}
#[test]
fn test_rebuild_drain_handoff_spin_loops_grow_then_cap() {
assert_eq!(
rebuild_drain_spin_loops(1),
REBUILD_DRAIN_HANDOFF_BASE_SPINS
);
assert_eq!(
rebuild_drain_spin_loops(2),
REBUILD_DRAIN_HANDOFF_BASE_SPINS * 2
);
assert_eq!(
rebuild_drain_spin_loops(3),
REBUILD_DRAIN_HANDOFF_BASE_SPINS * 4
);
assert_eq!(rebuild_drain_spin_loops(6), REBUILD_DRAIN_HANDOFF_MAX_SPINS);
assert_eq!(
rebuild_drain_spin_loops(32),
REBUILD_DRAIN_HANDOFF_MAX_SPINS
);
}
#[test]
fn test_rebuild_drain_handoff_yield_cadence_is_bounded() {
for attempt in 1..REBUILD_DRAIN_HANDOFF_YIELD_EVERY {
assert!(
!rebuild_drain_should_yield(attempt),
"attempt {attempt} should stay on CPU before the first bounded yield"
);
}
assert!(rebuild_drain_should_yield(
REBUILD_DRAIN_HANDOFF_YIELD_EVERY
));
assert!(!rebuild_drain_should_yield(
REBUILD_DRAIN_HANDOFF_YIELD_EVERY + 1
));
assert!(rebuild_drain_should_yield(
REBUILD_DRAIN_HANDOFF_YIELD_EVERY * 2
));
}
#[test]
fn test_rebuild_drain_handoff_respects_deadline() {
let mut handoff = RebuildDrainHandoff::default();
let started = Instant::now();
assert!(
handoff.next_wait(started, Duration::ZERO).is_none(),
"expired deadline must stop rebuild drain handoff"
);
let fresh_wait = handoff
.next_wait(Instant::now(), Duration::from_secs(1))
.expect("fresh deadline should allow a bounded handoff wait");
assert_eq!(fresh_wait.spin_loops, REBUILD_DRAIN_HANDOFF_BASE_SPINS);
assert!(!fresh_wait.yielded);
}
#[test]
fn test_rebuild_drain_schedule_summary_bounds_wake_amplification() {
let summary = summarize_rebuild_drain_schedule(8);
assert_eq!(summary.total_spin_loops, 8_128);
assert_eq!(summary.max_spin_loops, REBUILD_DRAIN_HANDOFF_MAX_SPINS);
assert_eq!(summary.yield_count, 2);
assert_eq!(summary.p95_spin_loops, REBUILD_DRAIN_HANDOFF_MAX_SPINS);
assert_eq!(summary.full_scan_interval, REBUILD_DRAIN_FULL_SCAN_INTERVAL);
}
#[test]
fn test_rebuild_rotate_swaps_active_table() {
let table = SharedPageLockTable::new(TEST_CAP);
assert_eq!(table.active_table.load(Ordering::Relaxed), 0);
assert_eq!(table.draining_table.load(Ordering::Relaxed), DRAINING_NONE);
assert!(table.try_acquire(1, 100).is_ok());
assert!(table.try_acquire(2, 200).is_ok());
table.acquire_rebuild_lease(1234, 0, 1000).unwrap();
table.rotate().unwrap();
assert_eq!(table.active_table.load(Ordering::Relaxed), 1);
assert_eq!(table.draining_table.load(Ordering::Relaxed), 0);
assert!(table.try_acquire(3, 300).is_ok());
assert_eq!(table.holder(1), Some(100));
assert_eq!(table.holder(2), Some(200));
assert_eq!(table.holder(3), Some(300));
}
#[test]
fn test_rebuild_drain_reaches_quiescence() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.try_acquire(10, 1).is_ok());
assert!(table.try_acquire(20, 2).is_ok());
table.acquire_rebuild_lease(1, 0, 1000).unwrap();
table.rotate().unwrap();
let status = table.drain_progress().unwrap();
assert_eq!(status.remaining, 2);
assert!(!status.quiescent);
assert!(table.release(10, 1));
assert!(table.release(20, 2));
let status = table.drain_progress().unwrap();
assert_eq!(status.remaining, 0);
assert!(status.quiescent);
}
#[test]
fn test_try_acquire_conflict_logs_lock_intent_and_holder() {
let table = SharedPageLockTable::new(TEST_CAP);
assert_eq!(table.try_acquire(7, 11), AcquireResult::Acquired);
let capture_guard = crate::test_support::tracing_capture_guard();
let (result, logs) = with_tracing_capture(&capture_guard, || table.try_acquire(7, 22));
assert_eq!(result, AcquireResult::Busy { holder: 11 });
assert!(logs.contains("lock_intent=\"exclusive\""));
assert!(logs.contains("requester_txn_id=22"));
assert!(logs.contains("holder_txn_id=11"));
assert!(logs.contains("page_number=7"));
}
#[test]
fn test_rebuild_lease_logs_pid_birth_and_expiry() {
let table = SharedPageLockTable::new(TEST_CAP);
let capture_guard = crate::test_support::tracing_capture_guard();
let ((), logs) = with_tracing_capture(&capture_guard, || {
table
.acquire_rebuild_lease(1234, 5678, 100)
.expect("lease acquisition must succeed");
});
assert!(logs.contains("pid=1234"));
assert!(logs.contains("pid_birth=5678"));
assert!(logs.contains("request_now_secs=100"));
assert!(logs.contains("lease_expires_at=105"));
}
#[test]
fn test_rebuild_no_abort_guarantee() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.try_acquire(10, 1).is_ok());
assert!(table.try_acquire(20, 1).is_ok());
table.acquire_rebuild_lease(1, 0, 1000).unwrap();
table.rotate().unwrap();
assert!(table.try_acquire(30, 1).is_ok());
assert!(table.release(10, 1));
assert!(table.release(20, 1));
assert!(table.try_acquire(10, 2).is_ok());
assert_eq!(table.holder(10), Some(2));
}
#[test]
fn test_rebuild_lease_prevents_concurrent_rebuilds() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.acquire_rebuild_lease(1001, 0, 1000).is_ok());
let err = table.acquire_rebuild_lease(1002, 0, 1000).unwrap_err();
assert_eq!(err, RebuildLeaseError::LeaseHeld { pid: 1001 });
}
#[test]
fn test_rebuild_stale_lease_stolen() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.acquire_rebuild_lease(1001, 0, 1000).is_ok());
let err = table.acquire_rebuild_lease(1002, 0, 1003).unwrap_err();
assert_eq!(err, RebuildLeaseError::LeaseHeld { pid: 1001 });
assert!(table.acquire_rebuild_lease(1002, 0, 1006).is_ok());
assert_eq!(table.rebuild_pid.load(Ordering::Relaxed), 1002);
}
#[test]
fn test_rebuild_stale_lease_stolen_at_expiry_boundary() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.acquire_rebuild_lease(1001, 0, 1000).is_ok());
let err = table.acquire_rebuild_lease(1002, 0, 1004).unwrap_err();
assert_eq!(err, RebuildLeaseError::LeaseHeld { pid: 1001 });
assert!(table.acquire_rebuild_lease(1002, 0, 1005).is_ok());
assert_eq!(table.rebuild_pid.load(Ordering::Relaxed), 1002);
}
#[cfg(unix)]
#[test]
fn test_process_alive_os_rejects_tagged_birth_mismatch() {
let pid = std::process::id();
let now_secs = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map_or(0, |duration| duration.as_secs());
let birth = current_process_birth_token(now_secs);
if birth & PID_BIRTH_PROCFS_TAG == 0 {
return;
}
assert!(process_alive_os(pid, birth));
let mismatched = PID_BIRTH_PROCFS_TAG | ((birth & !PID_BIRTH_PROCFS_TAG).wrapping_add(1));
assert!(!process_alive_os(pid, mismatched));
}
#[cfg(unix)]
#[test]
fn test_rebuild_lease_stolen_when_holder_process_dead_before_expiry() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.acquire_rebuild_lease(u32::MAX, 42, 1000).is_ok());
assert_eq!(table.rebuild_lease_expiry.load(Ordering::Relaxed), 1005);
assert!(table.acquire_rebuild_lease(1002, 7, 1001).is_ok());
assert_eq!(table.rebuild_pid.load(Ordering::Relaxed), 1002);
}
#[test]
fn test_rebuild_cancellation_safety() {
let table = SharedPageLockTable::new(TEST_CAP);
for i in 1..=5_u32 {
assert!(table.try_acquire(i, u64::from(i)).is_ok());
}
table.acquire_rebuild_lease(1, 0, 1000).unwrap();
table.rotate().unwrap();
for i in 1..=5_u32 {
assert!(table.release(i, u64::from(i)));
}
assert!(table.drain_progress().unwrap().quiescent);
let cleared = table.finalize_rebuild(1).unwrap();
assert_eq!(cleared, 5);
assert_eq!(table.rebuild_epoch(), 1);
assert_eq!(table.draining_table.load(Ordering::Relaxed), DRAINING_NONE);
assert_eq!(table.rebuild_pid.load(Ordering::Relaxed), 0);
assert_eq!(table.tables[0].occupied_count(), 0);
}
#[test]
fn test_striped_occupied_counter_matches_full_scan_across_rebuild() {
let table = SharedPageLockTable::new(TEST_CAP);
for page in 1..=20_u32 {
assert!(table.try_acquire(page, u64::from(page)).is_ok());
}
let active = table.active_table.load(Ordering::Relaxed);
let active_table = &table.tables[active as usize];
assert_eq!(
active_table.occupied_count(),
active_table.occupied_count_full_scan(),
"striped occupied counter must match exact active-table slot count"
);
table.acquire_rebuild_lease(1, 0, 1000).unwrap();
table.rotate().unwrap();
for page in 1..=20_u32 {
assert!(table.release(page, u64::from(page)));
}
assert!(table.drain_progress().unwrap().quiescent);
let cleared = table.finalize_rebuild(1).unwrap();
assert_eq!(cleared, 20);
assert_eq!(table.tables[0].occupied_count(), 0);
assert_eq!(table.tables[0].occupied_count_full_scan(), 0);
}
#[test]
fn test_rebuild_resource_exhaustion_busy() {
let table = SharedPageLockTable::new(16);
let process_id = std::process::id();
let now_secs = SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.map_or(0, |duration| duration.as_secs());
let pid_birth = current_process_birth_token(now_secs);
assert!(
table
.acquire_rebuild_lease(process_id, pid_birth, now_secs)
.is_ok()
);
for i in 1..=12_u32 {
let result = table.try_acquire(i, u64::from(i));
assert!(
result.is_ok(),
"should be able to acquire page {i}, got {result:?}"
);
}
let result = table.try_acquire(100, 100);
assert_eq!(
result,
AcquireResult::CapacityExhausted,
"new key insertion beyond 70% load factor must fail"
);
assert!(table.release(1, 1));
let result = table.try_acquire(1, 50);
assert!(
result.is_ok(),
"re-acquiring existing key slot must succeed even at capacity"
);
}
#[test]
fn test_rebuild_try_acquire_consults_draining_first() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.try_acquire(42, 1).is_ok());
table.acquire_rebuild_lease(1, 0, 1000).unwrap();
table.rotate().unwrap();
let result = table.try_acquire(42, 2);
assert_eq!(
result,
AcquireResult::Busy { holder: 1 },
"locked page must return SQLITE_BUSY immediately"
);
let result = table.try_acquire(42, 1);
assert_eq!(
result,
AcquireResult::AlreadyHeld,
"idempotent re-acquire in draining table"
);
assert!(table.release(42, 1));
let result = table.try_acquire(42, 2);
assert!(
result.is_ok(),
"should succeed after draining table release"
);
}
#[test]
fn test_e2e_lock_table_rolling_rebuild_under_load() {
let table = Arc::new(SharedPageLockTable::new(256));
let mut txns: Vec<(u64, Vec<u32>)> = Vec::new();
for txn_id in 1..=10_u64 {
let pages: Vec<u32> = (1..=5)
.map(|i| u32::try_from(txn_id).unwrap() * 10 + i)
.collect();
for &page in &pages {
assert!(table.try_acquire(page, txn_id).is_ok());
}
txns.push((txn_id, pages));
}
assert_eq!(table.total_locked(), 50);
let table2 = Arc::clone(&table);
let reader = std::thread::spawn(move || {
for _ in 0..100 {
for txn_id in 1..=10_u64 {
let page = u32::try_from(txn_id).unwrap() * 10 + 1;
let holder = table2.holder(page);
assert!(
holder.is_none() || holder == Some(txn_id),
"unexpected holder for page {page}: {holder:?}"
);
}
std::thread::yield_now();
}
});
table.acquire_rebuild_lease(999, 0, 1000).unwrap();
table.rotate().unwrap();
for (txn_id, pages) in &txns[0..5] {
for &page in pages {
table.release(page, *txn_id);
}
}
for txn_id in 11..=15_u64 {
let page = u32::try_from(txn_id).unwrap() * 10 + 1;
assert!(table.try_acquire(page, txn_id).is_ok());
}
for (txn_id, pages) in &txns[5..10] {
for &page in pages {
table.release(page, *txn_id);
}
}
let status = table.drain_progress().unwrap();
assert!(
status.quiescent,
"draining table must be quiescent after all releases"
);
let cleared = table.finalize_rebuild(999).unwrap();
assert!(
cleared > 0,
"should have cleared entries from drained table"
);
assert_eq!(table.rebuild_epoch(), 1);
reader.join().unwrap();
assert_eq!(table.total_locked(), 5);
for txn_id in 11..=15_u64 {
let page = u32::try_from(txn_id).unwrap() * 10 + 1;
assert_eq!(table.holder(page), Some(txn_id));
}
}
#[test]
fn test_try_acquire_free_page_succeeds() {
let table = SharedPageLockTable::new(TEST_CAP);
let result = table.try_acquire(42, 1);
assert_eq!(result, AcquireResult::Acquired);
assert_eq!(table.holder(42), Some(1));
assert_eq!(table.total_locked(), 1);
}
#[test]
fn test_try_acquire_locked_page_returns_busy() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.try_acquire(42, 1).is_ok());
let result = table.try_acquire(42, 2);
assert_eq!(
result,
AcquireResult::Busy { holder: 1 },
"locked page must return SQLITE_BUSY immediately"
);
assert_eq!(table.holder(42), Some(1));
}
#[test]
fn test_try_acquire_idempotent_same_txn() {
let table = SharedPageLockTable::new(TEST_CAP);
assert_eq!(table.try_acquire(10, 1), AcquireResult::Acquired);
assert_eq!(table.try_acquire(10, 1), AcquireResult::AlreadyHeld);
assert_eq!(table.total_locked(), 1);
assert_eq!(table.holder(10), Some(1));
}
#[test]
fn test_release_key_stable_no_page_number_deletion() {
let table = SharedPageLockTable::new(TEST_CAP);
assert!(table.try_acquire(42, 1).is_ok());
assert!(table.release(42, 1));
assert_eq!(table.holder(42), None);
assert_eq!(
table.active_occupied(),
1,
"page_number must NOT be deleted during release (key-stable invariant)"
);
let result = table.try_acquire(42, 2);
assert_eq!(result, AcquireResult::Acquired);
assert_eq!(table.holder(42), Some(2));
assert_eq!(table.active_occupied(), 1);
}
#[test]
fn test_linear_probing_collision_handling() {
let table = SharedPageLockTable::new(TEST_CAP);
let page_a = 1_u32;
let bucket_a = table.hash_index(page_a);
let mut page_b = 2_u32;
while table.hash_index(page_b) != bucket_a {
page_b += 1;
}
assert_eq!(table.try_acquire(page_a, 100), AcquireResult::Acquired);
assert_eq!(table.try_acquire(page_b, 200), AcquireResult::Acquired);
assert_eq!(table.holder(page_a), Some(100));
assert_eq!(table.holder(page_b), Some(200));
assert_eq!(table.active_occupied(), 2);
assert!(table.release(page_a, 100));
assert_eq!(table.holder(page_a), None);
assert_eq!(table.holder(page_b), Some(200));
}
#[test]
fn test_load_factor_70_percent_guard() {
let table = SharedPageLockTable::new(16);
let process_id = std::process::id();
assert!(table.acquire_rebuild_lease(process_id, 0, 0).is_ok());
for i in 1..=12_u32 {
let result = table.try_acquire(i, u64::from(i));
assert!(result.is_ok(), "should acquire page {i}, got {result:?}");
}
let old_active = table
.active_table
.load(std::sync::atomic::Ordering::Relaxed);
let result = table.try_acquire(200, 200);
let new_active = table
.active_table
.load(std::sync::atomic::Ordering::Relaxed);
assert_eq!(
result,
AcquireResult::Acquired,
"new key beyond 70% load factor triggers rebuild and succeeds"
);
assert_ne!(
old_active, new_active,
"table must have rotated due to load factor"
);
assert!(table.release(1, 1));
let result = table.try_acquire(1, 50);
assert!(
result.is_ok(),
"re-acquiring existing key slot must succeed even at capacity"
);
}
#[test]
fn test_released_locks_do_not_permanently_exhaust_capacity() {
let table = SharedPageLockTable::new(256);
for page in 1..=180_u32 {
assert_eq!(table.try_acquire(page, 1), AcquireResult::Acquired);
assert!(table.release(page, 1));
}
let result = table.try_acquire(9_999, 2);
assert!(
result.is_ok(),
"released historical keys should not cause permanent capacity exhaustion (got {result:?})"
);
assert_eq!(table.holder(9_999), Some(2));
}
#[test]
fn test_best_effort_rebuild_finalizes_after_timeout_once_quiescent() {
let table = SharedPageLockTable::new(TEST_CAP);
assert_eq!(table.try_acquire(11, 77), AcquireResult::Acquired);
let timed_out = table
.full_rebuild(1001, 0, 1_000, |_txn_id| true, Duration::ZERO)
.expect("full rebuild should return timeout result");
assert!(
timed_out.timed_out,
"rebuild should time out with active lock"
);
assert!(
table.is_rebuild_in_progress(),
"draining state should remain after timed-out rebuild"
);
assert!(table.release(11, 77), "lock release should succeed");
table.try_start_best_effort_rebuild();
assert!(
!table.is_rebuild_in_progress(),
"quiescent drain must be finalized"
);
assert_eq!(table.rebuild_epoch(), 1, "epoch must increment on rebuild");
}
#[test]
fn test_crash_cleanup_release_all_for_txn() {
let table = SharedPageLockTable::new(TEST_CAP);
let crashed_txn: u64 = 999;
assert!(table.try_acquire(1, crashed_txn).is_ok());
assert!(table.try_acquire(2, crashed_txn).is_ok());
assert!(table.try_acquire(3, crashed_txn).is_ok());
assert!(table.try_acquire(10, 500).is_ok());
assert_eq!(table.total_locked(), 4);
table.acquire_rebuild_lease(1, 0, 1000).unwrap();
table.rotate().unwrap();
assert!(table.try_acquire(4, crashed_txn).is_ok());
assert_eq!(table.total_locked(), 5);
let released = table.release_all_for_txn(crashed_txn);
assert_eq!(
released, 4,
"must release all 4 locks held by crashed txn across both tables"
);
assert_eq!(table.holder(1), None);
assert_eq!(table.holder(2), None);
assert_eq!(table.holder(3), None);
assert_eq!(table.holder(4), None);
assert_eq!(table.holder(10), Some(500));
}
#[test]
fn test_rolling_rebuild_rotate_drain_clear() {
let table = SharedPageLockTable::new(TEST_CAP);
for i in 1..=8_u32 {
assert!(table.try_acquire(i, u64::from(i)).is_ok());
}
assert_eq!(table.total_locked(), 8);
assert_eq!(table.rebuild_epoch(), 0);
table.acquire_rebuild_lease(42, 0, 1000).unwrap();
table.rotate().unwrap();
assert!(table.try_acquire(100, 100).is_ok());
let status = table.drain_progress().unwrap();
assert_eq!(status.remaining, 8);
assert!(!status.quiescent);
for i in 1..=8_u32 {
assert!(table.release(i, u64::from(i)));
}
let status = table.drain_progress().unwrap();
assert_eq!(status.remaining, 0);
assert!(status.quiescent, "drain must reach lock-quiescence");
let cleared = table.finalize_rebuild(42).unwrap();
assert_eq!(
cleared, 8,
"must clear all 8 occupied slots from drained table"
);
assert_eq!(table.rebuild_epoch(), 1, "epoch must increment on rebuild");
assert!(!table.is_rebuild_in_progress());
assert_eq!(table.holder(100), Some(100));
}
#[test]
fn test_e2e_shared_page_lock_table_cross_thread_contention() {
use std::sync::Barrier;
use std::sync::atomic::AtomicBool;
let table = Arc::new(SharedPageLockTable::new(256));
let barrier = Arc::new(Barrier::new(2));
let done = Arc::new(AtomicBool::new(false));
let table1 = Arc::clone(&table);
let barrier1 = Arc::clone(&barrier);
let done1 = Arc::clone(&done);
let proc1 = std::thread::spawn(move || {
let mut acquired_count = 0_u32;
for page in 1..=20_u32 {
let result = table1.try_acquire(page, 1);
assert!(result.is_ok(), "proc1 should acquire page {page}");
acquired_count += 1;
}
barrier1.wait();
std::thread::yield_now();
for page in 1..=20_u32 {
table1.release(page, 1);
}
done1.store(true, Ordering::Release);
acquired_count
});
let table2 = Arc::clone(&table);
let barrier2 = Arc::clone(&barrier);
let done2 = Arc::clone(&done);
let proc2 = std::thread::spawn(move || {
let mut busy_count = 0_u32;
let mut acquired_count = 0_u32;
barrier2.wait();
for page in 1..=20_u32 {
let result = table2.try_acquire(page, 2);
match result {
AcquireResult::Busy { holder } => {
assert_eq!(
holder, 1,
"if busy, holder must be proc1's txn (page {page})"
);
busy_count += 1;
}
AcquireResult::Acquired => {
acquired_count += 1;
}
other => {
assert!(
matches!(&other, AcquireResult::Busy { .. } | AcquireResult::Acquired),
"unexpected result for page {page}: {other:?}"
);
}
}
}
while !done2.load(Ordering::Acquire) {
std::thread::yield_now();
}
for page in 1..=20_u32 {
if table2.holder(page) != Some(2) {
let result = table2.try_acquire(page, 2);
assert!(
result.is_ok(),
"liveness: page {page} should be acquirable after proc1 released"
);
}
}
(busy_count, acquired_count)
});
let p1_acquired = proc1.join().unwrap();
let (p2_busy, p2_early_acquired) = proc2.join().unwrap();
assert_eq!(p1_acquired, 20, "proc1 acquired all 20 pages");
assert!(
p2_busy + p2_early_acquired == 20,
"proc2 saw exactly 20 results: {p2_busy} busy + {p2_early_acquired} acquired"
);
for page in 1..=20_u32 {
assert_eq!(
table.holder(page),
Some(2),
"after proc1 release, proc2 should hold page {page}"
);
}
let rebuild_result = table.full_rebuild(
100,
0,
2000,
|_| false, Duration::from_secs(5),
);
assert!(rebuild_result.is_ok(), "rebuild must complete");
let result = rebuild_result.unwrap();
assert!(!result.timed_out, "rebuild must not time out");
assert_eq!(result.epoch, 1, "rebuild epoch should be 1");
}
}
#[cfg(all(test, target_os = "linux"))]
mod crash_cleanup_process_liveness_tests {
use super::*;
use std::collections::HashMap;
use std::process::{Child, Command, Stdio};
use std::sync::{Mutex, MutexGuard, PoisonError};
use std::time::{Duration as StdDuration, Instant as StdInstant};
use nix::sys::signal::{Signal, kill};
use nix::unistd::Pid;
const TEST_CAP: u32 = 64;
const CRASH_ITERS: u32 = 30;
const REUSE_SPAWN_CAP: u32 = 200;
static PROC_TEST_LOCK: Mutex<()> = Mutex::new(());
fn proc_test_guard() -> MutexGuard<'static, ()> {
PROC_TEST_LOCK
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
fn child_birth_token(pid: u32) -> Option<u64> {
let ticks = read_proc_start_time_ticks(pid)?;
Some(PID_BIRTH_PROCFS_TAG | (ticks & !PID_BIRTH_PROCFS_TAG))
}
fn spawn_blocker() -> Option<Child> {
Command::new("sleep")
.arg("86400")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.ok()
}
fn spawn_blocker_with_birth() -> Option<(Child, u32, u64)> {
let child = spawn_blocker()?;
let pid = child.id();
match child_birth_token(pid) {
Some(birth) => Some((child, pid, birth)),
None => {
terminate(child);
None
}
}
}
fn terminate(mut child: Child) {
let _ = child.kill(); let _ = child.wait(); }
fn send_signal(pid: u32, sig: Signal) -> bool {
i32::try_from(pid).is_ok_and(|raw| kill(Pid::from_raw(raw), Some(sig)).is_ok())
}
fn wait_until_dead(pid: u32, birth: u64, timeout: StdDuration) -> bool {
let start = StdInstant::now();
loop {
if !process_alive_os(pid, birth) {
return true;
}
if start.elapsed() >= timeout {
return false;
}
std::thread::sleep(StdDuration::from_millis(5));
}
}
fn stale_birth(birth: u64) -> u64 {
PID_BIRTH_PROCFS_TAG | ((birth & !PID_BIRTH_PROCFS_TAG).wrapping_add(1))
}
#[test]
fn test_process_alive_live_child_then_dead_after_reap() {
let _guard = proc_test_guard();
let Some((child, pid, birth)) = spawn_blocker_with_birth() else {
eprintln!("skip: could not spawn `sleep` child in this environment");
return;
};
assert!(
process_alive_os(pid, birth),
"freshly spawned child pid={pid} must read as alive"
);
assert!(
!process_alive_os(pid, stale_birth(birth)),
"live child with a mismatched birth token must read as dead (reuse-safe)"
);
terminate(child);
assert!(
wait_until_dead(pid, birth, StdDuration::from_secs(1)),
"killed+reaped child pid={pid} must read as dead within 1s of reap"
);
}
#[test]
fn test_pid_reuse_birth_disambiguation_is_reuse_safe() {
let _guard = proc_test_guard();
let Some((child, pid, birth)) = spawn_blocker_with_birth() else {
eprintln!("skip: could not spawn `sleep` child in this environment");
return;
};
assert!(
process_alive_os(pid, birth),
"current birth for a live PID must read Alive"
);
assert!(
!process_alive_os(pid, stale_birth(birth)),
"stale birth for a reused PID must read Dead (no false-positive alive)"
);
terminate(child);
let mut retired: HashMap<u32, u64> = HashMap::new();
let mut observed_reuse = false;
for _ in 0..REUSE_SPAWN_CAP {
let Some((child, pid, birth)) = spawn_blocker_with_birth() else {
continue;
};
if let Some(&old_birth) = retired.get(&pid)
&& old_birth != birth
{
assert!(
!process_alive_os(pid, old_birth),
"reused pid={pid} with the OLD birth must read Dead"
);
assert!(
process_alive_os(pid, birth),
"reused pid={pid} with the NEW birth must read Alive"
);
observed_reuse = true;
terminate(child);
break;
}
retired.insert(pid, birth);
terminate(child);
}
if !observed_reuse {
eprintln!(
"note: OS did not reuse a PID within {REUSE_SPAWN_CAP} spawns; \
deterministic birth-token reuse-safety (Part A) still verified"
);
}
}
#[test]
fn test_crash_recovery_dead_holder_cas_cleared() {
let _guard = proc_test_guard();
let table = SharedPageLockTable::new(TEST_CAP);
const PAGE: u32 = 7;
let mut ran = 0_u32;
for iter in 0..CRASH_ITERS {
let Some((child, pid, birth)) = spawn_blocker_with_birth() else {
eprintln!("skip: could not spawn `sleep` child (iter {iter})");
break;
};
let holder_txn = u64::from(iter) + 1_000;
assert_eq!(
table.try_acquire(PAGE, holder_txn),
AcquireResult::Acquired,
"holder txn={holder_txn} should acquire page {PAGE}"
);
assert_eq!(table.holder(PAGE), Some(holder_txn));
terminate(child);
assert!(
wait_until_dead(pid, birth, StdDuration::from_secs(1)),
"crashed holder pid={pid} must be observably dead before cleanup"
);
assert!(
!process_alive_os(pid, birth),
"cleanup must see the holder as dead"
);
let released = table.release_all_for_txn(holder_txn);
assert_eq!(released, 1, "exactly the crashed holder's lock is cleared");
assert_eq!(
table.holder(PAGE),
None,
"page {PAGE} must be free after crash cleanup"
);
let next_txn = holder_txn + 500;
assert_eq!(
table.try_acquire(PAGE, next_txn),
AcquireResult::Acquired,
"post-cleanup acquire by txn={next_txn} must succeed"
);
assert_eq!(table.holder(PAGE), Some(next_txn));
assert!(
table.release(PAGE, next_txn),
"reset page for next iteration"
);
ran += 1;
}
assert!(ran > 0, "crash-recovery loop must run at least once");
}
#[test]
fn test_lease_expiry_reclaims_stopped_holder_no_double_ownership() {
let _guard = proc_test_guard();
let table = SharedPageLockTable::new(TEST_CAP);
let Some((child, holder_pid, holder_birth)) = spawn_blocker_with_birth() else {
eprintln!("skip: could not spawn `sleep` child in this environment");
return;
};
let now0 = 1_000_u64;
assert!(
table
.acquire_rebuild_lease(holder_pid, holder_birth, now0)
.is_ok(),
"holder should take the rebuild lease"
);
let expiry = now0 + DEFAULT_LEASE_SECS;
assert_eq!(table.rebuild_lease_expiry.load(Ordering::Relaxed), expiry);
assert!(
send_signal(holder_pid, Signal::SIGSTOP),
"SIGSTOP should succeed on the live holder"
);
assert!(
process_alive_os(holder_pid, holder_birth),
"a SIGSTOPped holder is still alive to the liveness probe — \
only the lease/hard-timeout can reclaim it"
);
let worker_pid = std::process::id();
let worker_birth = current_process_birth_token(now0);
let err = table
.acquire_rebuild_lease(worker_pid, worker_birth, expiry - 1)
.expect_err("lease must not be stealable before expiry from a live holder");
assert_eq!(err, RebuildLeaseError::LeaseHeld { pid: holder_pid });
let hard_deadline = now0 + 5 * DEFAULT_LEASE_SECS;
assert!(
table
.acquire_rebuild_lease(worker_pid, worker_birth, hard_deadline)
.is_ok(),
"lease must be reclaimable after the hard timeout"
);
assert_eq!(
table.rebuild_pid.load(Ordering::Relaxed),
worker_pid,
"worker now owns the rebuild lease"
);
assert!(
send_signal(holder_pid, Signal::SIGCONT),
"SIGCONT the holder"
);
assert!(
!table.renew_rebuild_lease(holder_pid, hard_deadline),
"resumed holder must fail to renew a lease it no longer owns"
);
let reacquire = table.acquire_rebuild_lease(holder_pid, holder_birth, hard_deadline);
assert_eq!(
reacquire,
Err(RebuildLeaseError::LeaseHeld { pid: worker_pid }),
"resumed holder must not double-own the lease held by the worker"
);
terminate(child);
}
}