use alloc::rc::Rc;
use core::cell::Cell;
#[cfg(debug_assertions)]
use alloc::rc::Weak;
#[cfg(debug_assertions)]
use alloc::vec::Vec;
#[cfg(debug_assertions)]
use core::cell::RefCell;
use crate::scope::OwnerFrame;
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct OwnershipStats {
pub live_owners: usize,
pub cleanup_registrations: usize,
pub owned_children: usize,
pub owners_created: u64,
pub owners_disposed: u64,
pub cleanups_registered: u64,
pub cleanups_run: u64,
}
impl OwnershipStats {
#[must_use = "this answers a question, it does not assert it; wrap it in `assert!`"]
pub fn is_empty(&self) -> bool {
self.live_owners == 0 && self.cleanup_registrations == 0 && self.owned_children == 0
}
}
#[must_use = "this only reads the gauges; in statement position it checks nothing"]
pub fn ownership_stats() -> OwnershipStats {
OWNERSHIP.try_with(Counters::snapshot).unwrap_or_default()
}
#[derive(Default)]
struct Counters {
live_owners: Cell<usize>,
cleanup_registrations: Cell<usize>,
owned_children: Cell<usize>,
owners_created: Cell<u64>,
owners_disposed: Cell<u64>,
cleanups_registered: Cell<u64>,
cleanups_run: Cell<u64>,
#[cfg(debug_assertions)]
registry: RefCell<Vec<Weak<OwnerFrame>>>,
}
impl Counters {
fn snapshot(&self) -> OwnershipStats {
OwnershipStats {
live_owners: self.live_owners.get(),
cleanup_registrations: self.cleanup_registrations.get(),
owned_children: self.owned_children.get(),
owners_created: self.owners_created.get(),
owners_disposed: self.owners_disposed.get(),
cleanups_registered: self.cleanups_registered.get(),
cleanups_run: self.cleanups_run.get(),
}
}
}
thread_local! {
static OWNERSHIP: Counters = Counters::default();
}
fn with_counters(f: impl FnOnce(&Counters)) {
let _ = OWNERSHIP.try_with(f);
}
fn add(cell: &Cell<usize>, n: usize) {
cell.set(cell.get().saturating_add(n));
}
fn sub(cell: &Cell<usize>, n: usize) {
cell.set(cell.get().saturating_sub(n));
}
fn tick(cell: &Cell<u64>, n: u64) {
cell.set(cell.get().saturating_add(n));
}
pub(crate) struct OwnerTally;
impl OwnerTally {
pub(crate) fn new() -> Self {
with_counters(|counters| {
add(&counters.live_owners, 1);
tick(&counters.owners_created, 1);
});
Self
}
}
impl Drop for OwnerTally {
fn drop(&mut self) {
with_counters(|counters| sub(&counters.live_owners, 1));
}
}
pub(crate) fn register(frame: &Rc<OwnerFrame>) {
#[cfg(debug_assertions)]
with_counters(|counters| {
let mut registry = counters.registry.borrow_mut();
if registry.len() >= 8 && registry.len() >= counters.live_owners.get().saturating_mul(2) {
registry.retain(|weak| weak.strong_count() > 0);
}
registry.push(Rc::downgrade(frame));
});
#[cfg(not(debug_assertions))]
let _ = frame;
}
pub(crate) fn cleanup_registered() {
with_counters(|counters| {
add(&counters.cleanup_registrations, 1);
tick(&counters.cleanups_registered, 1);
});
}
pub(crate) fn cleanups_taken(count: usize) {
with_counters(|counters| {
sub(&counters.cleanup_registrations, count);
tick(&counters.cleanups_run, count as u64);
});
}
pub(crate) fn cleanup_run_immediately() {
with_counters(|counters| {
tick(&counters.cleanups_registered, 1);
tick(&counters.cleanups_run, 1);
});
}
pub(crate) fn child_adopted() {
with_counters(|counters| add(&counters.owned_children, 1));
}
pub(crate) fn children_taken(count: usize) {
with_counters(|counters| sub(&counters.owned_children, count));
}
pub(crate) fn owner_disposed() {
with_counters(|counters| tick(&counters.owners_disposed, 1));
}
#[cfg(all(test, debug_assertions))]
pub(crate) fn registry_len() -> usize {
OWNERSHIP.with(|counters| counters.registry.borrow().len())
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum OwnershipGauge {
LiveOwners,
CleanupRegistrations,
OwnedChildren,
}
#[non_exhaustive]
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct OwnershipDrift {
pub gauge: OwnershipGauge,
pub reported: usize,
pub actual: usize,
}
#[non_exhaustive]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum OwnershipAudit {
Unavailable,
Consistent,
Drifted(Vec<OwnershipDrift>),
}
#[must_use = "the audit reports, it does not assert: `audit_ownership();` checks nothing. Use \
`debug_assert_ownership_consistent()` for the assertion form"]
pub fn audit_ownership() -> OwnershipAudit {
#[cfg(not(debug_assertions))]
return OwnershipAudit::Unavailable;
#[cfg(debug_assertions)]
OWNERSHIP
.try_with(|counters| {
let mut live = Vec::new();
counters.registry.borrow_mut().retain(|weak| {
let Some(frame) = weak.upgrade() else {
return false;
};
live.push(frame);
true
});
let mut drift = Vec::new();
let mut check = |gauge: OwnershipGauge, reported: usize, actual: usize| {
if reported != actual {
drift.push(OwnershipDrift {
gauge,
reported,
actual,
});
}
};
check(
OwnershipGauge::LiveOwners,
counters.live_owners.get(),
live.len(),
);
check(
OwnershipGauge::CleanupRegistrations,
counters.cleanup_registrations.get(),
live.iter().map(|frame| frame.pending_cleanups()).sum(),
);
check(
OwnershipGauge::OwnedChildren,
counters.owned_children.get(),
live.iter().map(|frame| frame.owned_children()).sum(),
);
if drift.is_empty() {
OwnershipAudit::Consistent
} else {
OwnershipAudit::Drifted(drift)
}
})
.unwrap_or(OwnershipAudit::Unavailable)
}
#[track_caller]
pub fn debug_assert_ownership_consistent() {
if let OwnershipAudit::Drifted(drift) = audit_ownership() {
panic!("ownership counters drifted from the live owner tree: {drift:?}");
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Reactor, on_cleanup, scope};
#[cfg(debug_assertions)]
#[test]
fn the_audit_detects_a_gauge_that_has_drifted() {
std::thread::spawn(|| {
let (_handle, ()) = scope(|| {});
assert_eq!(audit_ownership(), OwnershipAudit::Consistent);
let before = ownership_stats().owned_children;
child_adopted();
match audit_ownership() {
OwnershipAudit::Drifted(drift) => assert_eq!(
drift,
vec![OwnershipDrift {
gauge: OwnershipGauge::OwnedChildren,
reported: before + 1,
actual: before,
}]
),
other => panic!("the audit missed a drift it was built to catch: {other:?}"),
}
children_taken(1); assert_eq!(audit_ownership(), OwnershipAudit::Consistent);
})
.join()
.expect("test thread panicked");
}
#[cfg(debug_assertions)]
#[test]
fn the_audit_registry_does_not_grow_with_owners_long_dead() {
std::thread::spawn(|| {
for _ in 0..2_000 {
let (handle, ()) = scope(|| on_cleanup(|| {}));
handle.dispose();
}
let live = ownership_stats().live_owners;
assert!(
registry_len() < 64,
"registry retained {} entries for {live} live owners",
registry_len()
);
})
.join()
.expect("test thread panicked");
}
#[test]
fn the_audit_reports_unavailable_rather_than_consistent_without_a_registry() {
std::thread::spawn(|| {
let (_handle, ()) = scope(|| {});
let audit = audit_ownership();
if cfg!(debug_assertions) {
assert_eq!(audit, OwnershipAudit::Consistent);
} else {
assert_eq!(
audit,
OwnershipAudit::Unavailable,
"without a registry the honest answer is `cannot say`"
);
}
debug_assert_ownership_consistent();
})
.join()
.expect("test thread panicked");
}
#[test]
fn accounting_survives_a_thread_local_destroyed_before_the_frames_it_counts() {
thread_local! {
static HOLDER: std::cell::RefCell<Option<crate::ScopeHandle>> =
const { std::cell::RefCell::new(None) };
}
std::thread::spawn(|| {
HOLDER.with(|holder| drop(holder.borrow())); let reactor = Reactor::new();
let _guard = reactor.enter();
let (handle, ()) = scope(|| on_cleanup(|| {}));
HOLDER.with(|holder| *holder.borrow_mut() = Some(handle));
})
.join()
.expect("the thread must unwind cleanly rather than abort");
}
}