diff options
Diffstat (limited to 'crates/tor-memquota/src')
| -rw-r--r-- | crates/tor-memquota/src/config.rs | 103 | ||||
| -rw-r--r-- | crates/tor-memquota/src/drop_bomb.rs | 515 | ||||
| -rw-r--r-- | crates/tor-memquota/src/drop_reentrancy.rs | 143 | ||||
| -rw-r--r-- | crates/tor-memquota/src/error.rs | 144 | ||||
| -rw-r--r-- | crates/tor-memquota/src/internal_prelude.rs | 63 | ||||
| -rw-r--r-- | crates/tor-memquota/src/lib.rs | 206 | ||||
| -rw-r--r-- | crates/tor-memquota/src/mtracker.rs | 1047 | ||||
| -rw-r--r-- | crates/tor-memquota/src/mtracker/bookkeeping.rs | 308 | ||||
| -rw-r--r-- | crates/tor-memquota/src/mtracker/reclaim.rs | 396 | ||||
| -rw-r--r-- | crates/tor-memquota/src/mtracker/reclaim/deferred_drop.rs | 81 | ||||
| -rw-r--r-- | crates/tor-memquota/src/mtracker/test.rs | 678 | ||||
| -rw-r--r-- | crates/tor-memquota/src/mtracker/total_qty_notifier.rs | 70 | ||||
| -rw-r--r-- | crates/tor-memquota/src/refcount.rs | 315 | ||||
| -rw-r--r-- | crates/tor-memquota/src/utils.rs | 69 |
14 files changed, 4138 insertions, 0 deletions
diff --git a/crates/tor-memquota/src/config.rs b/crates/tor-memquota/src/config.rs new file mode 100644 index 000000000..ce8fd364c --- /dev/null +++ b/crates/tor-memquota/src/config.rs @@ -0,0 +1,103 @@ +//! Configuration (private module) + +use crate::internal_prelude::*; + +/// We want to support at least this many participants with a cache each +/// +/// This is not a recommended value; it's probably too lax +const MIN_MAX_PARTICIPANTS: usize = 10; + +/// Minimum hysteresis +/// +/// This is not a recommended value; it's probably far too lax for sensible performance! +const MAX_LOW_WATER_RATIO: f32 = 0.98; + +define_derive_deftly! { + /// Define setters on the builder for every field of type `Qty` + /// + /// The field type must be spelled precisely that way: + /// we use `approx_equal(...)`. + QtySetters = + + impl $< $ttype Builder > { + $( + ${when approx_equal($ftype, Qty)} + + ${fattrs doc} + /// + /// (Setter method.) + pub fn $fname(&mut self, value: usize) -> &mut Self { + self.$fname = Some(Qty(value)); + self + } + ) + } +} + +/// Configuration for a memory data tracker +/// +/// This is where the quota is specified. +#[derive(Debug, Clone, Builder, Eq, PartialEq, Deftly)] +#[derive_deftly(QtySetters)] +#[builder(build_fn(private, name = "build_unvalidated", error = "ConfigBuildError"))] +#[builder(derive(Serialize, Deserialize, Debug, Deftly, Eq, PartialEq))] +#[builder_struct_attr(derive_deftly(tor_config::Flattenable))] +pub struct Config { + /// Maximum memory usage tolerated before reclamation starts + /// + /// Note that this is not a hard limit. + /// See Approximate in [the overview](crate). + /// + /// + #[builder(setter(custom))] + pub(crate) max: Qty, + + /// Reclamation will stop when memory use is reduced to below this value + /// + /// Default is 75% of the maximum. + #[builder(setter(custom))] + pub(crate) low_water: Qty, +} + +impl Config { + /// Start building a [`Config`] + /// + /// Returns a fresh default [`ConfigBuilder`]. + pub fn builder() -> ConfigBuilder { + ConfigBuilder::default() + } +} + +impl ConfigBuilder { + /// Builds a new `Config` from a builder + /// + /// Returns an error unless at least `max` has been specified, + /// or if the fields values are invalid or inconsistent. + pub fn build(&self) -> Result<Config, ConfigBuildError> { + let mut builder = self.clone(); + if let (Some(max), None) = (builder.max, builder.low_water) { + builder.low_water = Some(Qty((*max as f32 * 0.75) as _)); + } + let config = builder.build_unvalidated()?; + + let min_low_water = crate::mtracker::MAX_CACHE.as_usize() * MIN_MAX_PARTICIPANTS; + if *config.low_water < min_low_water { + return Err(ConfigBuildError::Invalid { + field: "low_water".into(), + problem: format!("must be at least {min_low_water}"), + }); + } + + let ratio: f32 = *config.low_water as f32 / *config.max as f32; + if ratio > MAX_LOW_WATER_RATIO { + return Err(ConfigBuildError::Inconsistent { + fields: vec!["low_water".into(), "max".into()], + problem: format!( + "low_water / max = {ratio}; must be <= {MAX_LOW_WATER_RATIO}, ideally considerably lower" + ), + }); + } + + Ok(config) + } +} diff --git a/crates/tor-memquota/src/drop_bomb.rs b/crates/tor-memquota/src/drop_bomb.rs new file mode 100644 index 000000000..6a7b55fa6 --- /dev/null +++ b/crates/tor-memquota/src/drop_bomb.rs @@ -0,0 +1,515 @@ +//! Drop bombs, for assurance of postconditions when types are dropped +//! +//! Provides two drop bomb types: [`DropBomb`] and [`DropBombCondition`]. +//! +//! These help assure that our algorithms are correct, +//! by detecting when types that contain the bomb are dropped inappropriately. +//! +//! # No-op outside `#[cfg(test)]` +//! +//! When used outside test code, these types are unit ZSTs, +//! and are completely inert. +//! They won't cause panics or detect bugs, in production. +//! +//! # Panics (in tests), and simulation +//! +//! These types work by panicking in drop, when a bug is detected. +//! This will then cause a test failure. +//! Such panics are described as "explodes (panics)" in the documentation. +//! +//! There are also simulated drop bombs, whose explosions do not actually panic. +//! Instead, they record that a panic would have occurred, +//! and print a message to stderr. +//! The constructors provide a handle to allow the caller to enquire about explosions. +//! This allows for testing a containing type's drop bomb logic. +//! +//! Certain misuses result in actual panics, even with simulated bombs. +//! This is described as "panics (actually)". +//! +//! # Choosing a bomb +//! +//! [`DropBomb`] is for assuring the runtime context or appropriate timing of drops +//! (and could be used for implementing general conditions). +//! +//! [`DropBombCondition`] is for assuring the properties of a value that is being dropped. + +use crate::internal_prelude::*; + +#[cfg(test)] +use std::sync::atomic::{AtomicBool, Ordering}; + +//---------- macros used in this module, and supporting trait ---------- + +define_derive_deftly! { + /// Helper for common impls on bombs + /// + /// * Provides `fn new_armed` + /// * Provides `fn new_simulated` + /// * Implements `Drop`, using `TestableDrop::drop_impl` + BombImpls = + + impl $ttype { + /// Create a new drop bomb, which must be properly disposed of + pub(crate) const fn new_armed() -> Self { + let status = Status::ARMED_IN_TESTS; + $ttype { status } + } + } + + #[cfg(test)] + impl $ttype { + /// Create a simulated drop bomb + pub(crate) fn new_simulated() -> (Self, SimulationHandle) { + let handle = SimulationHandle::new(); + let status = S::ArmedSimulated(handle.clone()); + ($ttype { status }, handle) + } + + /// Turn an existing armed drop bomb into a simulated one + /// + /// This is useful for writing test cases, without having to make a `new_simulated` + /// constructor for whatever type contains the drop bomb. + /// Instead, construct it normally, and then reach in and call this on the bomb. + /// + /// # Panics + /// + /// `self` must be armed. Otherwise, (actually) panics. + pub(crate) fn make_simulated(&mut self) -> SimulationHandle { + let handle = SimulationHandle::new(); + let new_status = S::ArmedSimulated(handle.clone()); + let old_status = mem::replace(&mut self.status, new_status); + assert!(matches!(old_status, S::Armed)); + handle + } + + /// Implemnetation of `Drop::drop`, split out for testability. + /// + /// Calls `drop_status`, and replaces `self.status` with `S::Disarmed`, + /// so that `self` can be actually dropped (if we didn't panic). + fn drop_impl(&mut self) { + // Do the replacement first, so that if drop_status unwinds, we don't panic in panic. + let status = mem::replace(&mut self.status, S::Disarmed); + <$ttype as DropStatus>::drop_status(status); + } + } + + + #[cfg(test)] + impl Drop for $ttype { + fn drop(&mut self) { + // We don't check for unwinding. + // We shouldn't drop a nonzero one of these even if we're panicking. + // If we do, it'll be a double panic => abort. + self.drop_impl(); + } + } +} + +/// Core of `Drop`, that can be called separately, for testing +/// +/// To use: implement this, and derive deftly +/// [`BombImpls`](derive_deftly_template_BombImpls). +trait DropStatus { + /// Handles dropping of a `Self` with this `status` field value + fn drop_status(status: Status); +} + +//---------- public types ---------- + +/// Drop bomb: for assuring that drops happen only when expected +/// +/// Obtained from [`DropBomb::new_armed()`]. +/// +/// # Explosions +/// +/// Explodes (panicking) if dropped, +/// unless [`.disarm()`](DropBomb::disarm) is called first. +#[derive(Deftly, Debug)] +#[derive_deftly(BombImpls)] +pub(crate) struct DropBomb { + /// What state are we in + status: Status, +} + +/// Drop condition: for ensuring that a condition is true, on drop +/// +/// Obtained from [`DropBombCondition::new_armed()`]. +/// +/// Instead of dropping this, you must call +/// `drop_bomb_disarm_assert!` +/// (or its internal function `disarm_assert()`. +// rustdoc can't manage to make a link to this crate-private macro or cfg-test item. +/// +/// It will often be necessary to add `#[allow(dead_code)]` +/// on the `DropBombCondition` field of a containing type, +/// since outside tests, the `Drop` impl will usually be configured out, +/// and that's the only place this field is actually read. +/// +/// # Panics +/// +/// Panics (actually) if it is simply dropped. +#[derive(Deftly, Debug)] +#[derive_deftly(BombImpls)] +pub(crate) struct DropBombCondition { + /// What state are we in + #[allow(dead_code)] // not read outside tests + status: Status, +} + +/// Handle onto a simulated [`DropBomb`] or [`DropCondition`] +/// +/// Can be used to tell whether the bomb "exploded" +/// (ie, whether `drop` would have panicked, if this had been a non-simulated bomb). +#[cfg(test)] +#[derive(Debug)] +pub(crate) struct SimulationHandle { + exploded: Arc<AtomicBool>, +} + +/// Unit token indicating that a simulated drop bomb did explode, and would have panicked +#[cfg(test)] +#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)] +pub(crate) struct SimulationExploded; + +//---------- internal types ---------- + +/// State of some kind of drop bomb +/// +/// This type is inert; the caller is responsible for exploding or panicking. +#[derive(Debug)] +enum Status { + /// This bomb is disarmed and will not panic. + /// + /// This is always the case outside `#[cfg(test)]` + Disarmed, + + /// This bomb is armed. It will (or may) panic on drop. + #[cfg(test)] + Armed, + + /// This bomb is armed, but we're running in simulation. + #[cfg(test)] + ArmedSimulated(SimulationHandle), +} + +use Status as S; + +//---------- DropBomb impls ---------- + +impl DropBomb { + /// Disarm this bomb. + /// + /// It will no longer explode (panic) when dropped. + pub(crate) fn disarm(&mut self) { + self.status = S::Disarmed; + } +} + +#[cfg(test)] +impl DropStatus for DropBomb { + fn drop_status(status: Status) { + match status { + S::Disarmed => {} + S::Armed => panic!("DropBomb dropped without a previous call to .disarm()"), + S::ArmedSimulated(handle) => handle.set_exploded(), + } + } +} + +//---------- DropCondition impls ---------- + +/// Check the condition, and disarm the bomb +/// +/// If `CONDITION` is true, disarms the bomb; otherwise, explodes (panics). +/// +/// # Syntax +/// +/// ``` +/// drop_bomb_disarm_assert!(BOMB, CONDITION); +/// drop_bomb_disarm_assert!(BOMB, CONDITION, "FORMAT", FORMAT_ARGS..); +/// ``` +/// +/// where +/// +/// * `BOMB: &mut DropCondition` (or something that derefs to that). +/// * `CONDITION: bool` +/// +/// # Example +/// +/// ``` +/// # struct S { drop_bomb: DropCondition }; +/// # impl S { fn f(&mut self) { +/// drop_bomb_disarm_assert!(self.drop_bomb, self.raw, Qty(0)); +/// # } } +/// ``` +/// +/// # Explodes +/// +/// Explodes unless the condition is satisfied. +// +// This macro has this long name because we can't do scoping of macro-rules macros. +#[cfg(test)] // Should not be used outside tests, since the drop impls should be conditional +macro_rules! drop_bomb_disarm_assert { + { $bomb:expr, $condition:expr $(,)? } => { + $bomb.disarm_assert( + || $condition, + format_args!(concat!("condition = ", stringify!($condition))), + ) + }; + { $bomb:expr, $condition:expr, $fmt:literal $($rest:tt)* } => { + $bomb.disarm_assert( + || $condition, + format_args!(concat!("condition = ", stringify!($condition), ": ", $fmt), + $($rest)*), + ) + }; +} + +impl DropBombCondition { + /// Check a condition, and disarm the bomb + /// + /// If `call()` returns true, disarms the bomb; otherwise, explodes (panics). + /// + /// # Explodes + /// + /// Explodes unless the condition is satisfied. + #[inline] + #[cfg(test)] // Should not be used outside tests, since the drop impls should be conditional + pub(crate) fn disarm_assert(&mut self, call: impl FnOnce() -> bool, msg: fmt::Arguments) { + match mem::replace(&mut self.status, S::Disarmed) { + S::Disarmed => { + // outside cfg(test), this is the usual path. + // placate the compiler: we ignore all our arguments + let _ = call; + let _ = msg; + + #[cfg(test)] + panic!("disarm_assert called more than once!"); + } + #[cfg(test)] + S::Armed => { + if !call() { + panic!("drop condition violated: dropped, but condition is false: {msg}"); + } + } + #[cfg(test)] + #[allow(clippy::print_stderr)] + S::ArmedSimulated(handle) => { + if !call() { + eprintln!("drop condition violated in simulation: {msg}"); + handle.set_exploded(); + } + } + } + } +} + +/// Ideally, if you use this, your struct's other default values meet your drop condition! +impl Default for DropBombCondition { + fn default() -> DropBombCondition { + Self::new_armed() + } +} + +#[cfg(test)] +impl DropStatus for DropBombCondition { + fn drop_status(status: Status) { + assert!(matches!(status, S::Disarmed)); + } +} + +//---------- SimulationHandle impls ---------- + +#[cfg(test)] +impl SimulationHandle { + /// Determine whether a drop bomb would have been triggered + /// + /// If the corresponding [`DropBomb]` or [`DropCondition`] + /// would have panicked (if we weren't simulating), + /// returns `Err`. + /// + /// # Panics + /// + /// The corresponding `DropBomb` or `DropCondition` must have been dropped. + /// Otherwise, calling `outcome` will (actually) panic. + pub(crate) fn outcome(mut self) -> Result<(), SimulationExploded> { + let panicked = Arc::into_inner(mem::take(&mut self.exploded)) + .expect("bomb has not yet been dropped") + .into_inner(); + if panicked { + Err(SimulationExploded) + } else { + Ok(()) + } + } + + /// Require that this bomb did *not* explode + /// + /// # Panics + /// + /// Panics if corresponding `DropBomb` hasn't yet been dropped, + /// or if it exploded when it was dropped. + pub(crate) fn expect_ok(self) { + let () = self.outcome().expect("bomb unexpectedly exploded"); + } + + /// Require that this bomb *did* explode + /// + /// # Panics + /// + /// Panics if corresponding `DropBomb` hasn't yet been dropped, + /// or if it did *not* explode when it was dropped. + pub(crate) fn expect_exploded(self) { + let SimulationExploded = self + .outcome() + .expect_err("bomb unexpectedly didn't explode"); + } + + /// Return a new handle with no explosion recorded + fn new() -> Self { + SimulationHandle { + exploded: Default::default(), + } + } + + /// Return a clone of this handle + // + // Deliberately not a public Clone impl + fn clone(&self) -> Self { + SimulationHandle { + exploded: self.exploded.clone(), + } + } + + /// Mark this simulated bomb as having exploded + fn set_exploded(&self) { + self.exploded.store(true, Ordering::Release); + } +} + +//---------- internal impls ---------- + +impl Status { + /// Armed, in tests + #[cfg(test)] + const ARMED_IN_TESTS: Status = S::Armed; + + /// "Armed", outside tests, is in fact not armed + #[cfg(not(test))] + const ARMED_IN_TESTS: Status = S::Disarmed; +} + +#[cfg(test)] +mod test { + // @@ begin test lint list maintained by maint/add_warning @@ + #![allow(clippy::bool_assert_comparison)] + #![allow(clippy::clone_on_copy)] + #![allow(clippy::dbg_macro)] + #![allow(clippy::mixed_attributes_style)] + #![allow(clippy::print_stderr)] + #![allow(clippy::print_stdout)] + #![allow(clippy::single_char_pattern)] + #![allow(clippy::unwrap_used)] + #![allow(clippy::unchecked_duration_subtraction)] + #![allow(clippy::useless_vec)] + #![allow(clippy::needless_pass_by_value)] + //! <!-- @@ end test lint list maintained by maint/add_warning @@ --> + #![allow(clippy::let_and_return)] // TODO this lint is annoying and we should disable it + + use super::*; + use std::any::Any; + use std::panic::catch_unwind; + + #[test] + fn bomb_disarmed() { + let mut b = DropBomb::new_armed(); + b.disarm(); + drop(b); + } + + #[test] + fn bomb_panic() { + let mut b = DropBomb::new_armed(); + let _: Box<dyn Any> = catch_unwind(AssertUnwindSafe(|| b.drop_impl())).unwrap_err(); + } + + #[test] + fn bomb_sim_disarmed() { + let (mut b, h) = DropBomb::new_simulated(); + b.disarm(); + drop(b); + h.expect_ok(); + } + + #[test] + fn bomb_sim_explosion() { + let (b, h) = DropBomb::new_simulated(); + drop(b); + h.expect_exploded(); + } + + #[test] + fn bomb_make_sim_explosion() { + let mut b = DropBomb::new_armed(); + let h = b.make_simulated(); + drop(b); + h.expect_exploded(); + } + + struct HasBomb { + on_drop: Result<(), ()>, + bomb: DropBombCondition, + } + + impl Drop for HasBomb { + fn drop(&mut self) { + drop_bomb_disarm_assert!(self.bomb, self.on_drop.is_ok()); + } + } + + #[test] + fn cond_ok() { + let hb = HasBomb { + on_drop: Ok(()), + bomb: DropBombCondition::new_armed(), + }; + drop(hb); + } + + #[test] + fn cond_sim_explosion() { + let (bomb, h) = DropBombCondition::new_simulated(); + let hb = HasBomb { + on_drop: Err(()), + bomb, + }; + drop(hb); + h.expect_exploded(); + } + + #[test] + fn cond_explosion_panic() { + // make an actual panic + let mut bomb = DropBombCondition::new_armed(); + let _: Box<dyn Any> = catch_unwind(AssertUnwindSafe(|| { + bomb.disarm_assert(|| false, format_args!("testing")); + })) + .unwrap_err(); + } + + #[test] + fn cond_forgot_drop_impl() { + // pretend that we put a DropBombCondition on this, + // but we forgot to impl Drop and call drop_bomb_disarm_assert + struct ForgotDropImpl { + bomb: DropBombCondition, + } + let fdi = ForgotDropImpl { + bomb: DropBombCondition::new_armed(), + }; + // pretend that fdi is being dropped + let mut bomb = fdi.bomb; // move out + + let _: Box<dyn Any> = catch_unwind(AssertUnwindSafe(|| bomb.drop_impl())).unwrap_err(); + } +} diff --git a/crates/tor-memquota/src/drop_reentrancy.rs b/crates/tor-memquota/src/drop_reentrancy.rs new file mode 100644 index 000000000..3e157cf88 --- /dev/null +++ b/crates/tor-memquota/src/drop_reentrancy.rs @@ -0,0 +1,143 @@ +//! Newtype which helps assure lack of drop entrance hazards +//! +//! Provides a drop bomb which will help tests detect latent bugs. +//! +//! We want this because there are places where we handle an Arc containing +//! a user-provided trait object, but where we want to prevent invoking +//! the user's Drop impl since that may lead to reentrancy. +//! +//! See the section on "Reentrancy" in the docs for `mtracker::State`. +//! +//! Outside tests, the types in this module are equivalent to `std::sync`'s. +//! So, we never panic in a drop in production. +//! Dropping in the wrong place might lead to a deadlock +//! (due to mutex reentrancy) +//! but this is far from certain: +//! probably, while we're running, the caller has another live reference, +//! so the drop of the underlying type won't happen now anyway. +//! +//! In any case, drop bombs mustn't be used in production. +//! Not only can they escalate the severity of problems, +//! where the program might blunder on, +//! but also +//! because Rust upstream are seriously considering +//! [turning them into aborts](https://github.com/rust-lang/rfcs/pull/3288)! +// +// There are no separate tests for this module. Drop bombs are hard to test for. +// However, in an ad-hoc test, the bomb has been shown to be able to explode, +// if a `ProtectedArc` is dropped. + +use crate::internal_prelude::*; + +/// A `Weak<P>`, but upgradeable only to a `ProtectedArc`, not a raw `Arc`. +#[derive(Debug)] +pub(crate) struct ProtectedWeak<P: ?Sized>(Weak<P>); + +/// An `Arc`, but containing a type which should only be dropped in certain places +/// +/// In non `#[cfg(test)]` builds, this is just `Arc<P>`. +/// +/// When testing, it has a drop bomb. You must call `.promise_dropping_is_ok`. +/// It will panic if it's simply dropped. +#[derive(Debug, Deref, DerefMut)] +pub(crate) struct ProtectedArc<P: ?Sized> { + /// The actual explosive (might be armed or disarmed) + bomb: DropBomb, + + /// The underlying `Arc` + #[deref(forward)] + #[deref_mut(forward)] + arc: Arc<P>, +} + +impl<P: ?Sized> ProtectedWeak<P> { + /// Make a new `ProtectedWeak` + pub(crate) fn new(p: Weak<P>) -> Self { + ProtectedWeak(p) + } + + /// Upgrade a `ProtectedWeak` to a `ProtectedArc`, if it's not been garbage collected + pub(crate) fn upgrade(&self) -> Option<ProtectedArc<P>> { + Some(ProtectedArc::new(self.0.upgrade()?)) + } + + /// Convert back into an unprotected `Weak`. + /// + /// # CORRECTNESS + /// + /// You must arrange that the drop reentrancy requirements aren't violated + /// by `Arc`s made from the returned `Weak`. + pub(crate) fn unprotect(self) -> Weak<P> { + self.0 + } +} + +impl<P: ?Sized> ProtectedArc<P> { + /// Make a new `ProtectedArc` from a raw `Arc` + /// + /// # CORRECTNESS + /// + /// Presumably the `Arc` came from an uncontrolled external source, such as user code. + pub(crate) fn new(arc: Arc<P>) -> Self { + let bomb = DropBomb::new_armed(); + ProtectedArc { arc, bomb } + } + + /// Obtain a `ProtectedWeak` from a `&ProtectedArc` + // + // If this were a more general-purpose library, we'd avoid this and other methods on `self`. + pub(crate) fn downgrade(&self) -> ProtectedWeak<P> { + ProtectedWeak(Arc::downgrade(&self.arc)) + } + + /// Convert back into an unprotected `Arc` + /// + /// # CORRECTNESS + /// + /// If the return value is dropped, the location must be suitable for that. + /// Or, maybe the returned value is going to calling code in the external user, + /// (which, therefore, wouldn't pose a reentrancy hazard). + pub(crate) fn promise_dropping_is_ok(mut self) -> Arc<P> { + self.bomb.disarm(); + self.arc + } +} + +#[cfg(test)] +mod test { + // @@ begin test lint list maintained by maint/add_warning @@ + #![allow(clippy::bool_assert_comparison)] + #![allow(clippy::clone_on_copy)] + #![allow(clippy::dbg_macro)] + #![allow(clippy::mixed_attributes_style)] + #![allow(clippy::print_stderr)] + #![allow(clippy::print_stdout)] + #![allow(clippy::single_char_pattern)] + #![allow(clippy::unwrap_used)] + #![allow(clippy::unchecked_duration_subtraction)] + #![allow(clippy::useless_vec)] + #![allow(clippy::needless_pass_by_value)] + //! <!-- @@ end test lint list maintained by maint/add_warning @@ --> + #![allow(clippy::let_and_return)] // TODO this lint is annoying and we should disable it + + use super::*; + + struct Payload; + + #[test] + fn fine() { + let arc = Arc::new(Payload); + let prot = ProtectedArc::new(arc); + let arc = prot.promise_dropping_is_ok(); + drop(arc); + } + + #[test] + fn bad() { + let arc = Arc::new(Payload); + let mut prot = ProtectedArc::new(arc); + let h = prot.bomb.make_simulated(); + drop(prot); + h.expect_exploded(); + } +} diff --git a/crates/tor-memquota/src/error.rs b/crates/tor-memquota/src/error.rs new file mode 100644 index 000000000..4cca5bb10 --- /dev/null +++ b/crates/tor-memquota/src/error.rs @@ -0,0 +1,144 @@ +//! Errors arising from memory tracking + +use crate::internal_prelude::*; + +/// An error occurring when tracking memory usage +#[derive(Debug, Clone, Error)] +#[non_exhaustive] +pub enum Error { + /// The memory quota tracker has been torn down + #[error("attempted to use shut down memory tracker")] + TrackerShutdown, + + /// The Account has been torn down + /// + /// This can happen if the account or participant has Collapsed due to reclamation + #[error("memory pressure (attempted to use closed memory tracking account)")] + AccountClosed, + + /// The Participant has been torn down + /// + /// This can happen if the account or participant has Collapsed due to reclamation + #[error("memory pressure (attempt to allocate by torn-down memory tracking participant)")] + ParticipantShutdown, + + /// Previous bug, memory quota tracker is corrupted + #[error("memory tracker is corrupted due to previous bug")] + TrackerCorrupted, + + /// Bug + #[error("internal error")] + Bug(#[from] Bug), +} + +/// Memory pressure means this data structure (or other facility) was torn down +/// +/// Error type suitable for use by data structures and facilities +/// which participate in memory tracking. +/// +/// Convertible from a [`tor_memtrack::Error`](enum@Error), +/// or constructible via `Default` or [`new`](MemoryReclaimedError::new). +#[derive(Debug, Clone, Error, Default)] +#[non_exhaustive] +#[error("{0}")] +pub struct MemoryReclaimedError(ReclaimedErrorInner); + +/// Content of a [`MemoryReclaimedError`] +// Separate struct so we don't expose the variants +#[derive(Debug, Clone, Error, Default)] +enum ReclaimedErrorInner { + /// Collapsed, from `ReclaimedError::new` + #[error("data structure discarded due to memory pressure")] + #[default] + Collapsed, + + /// Othere error from tracker + #[error("{0}")] + TrackerError(#[from] Error), +} + +/// An error occurring when setting up a memory quota tracker +#[derive(Debug, Clone, Error)] +#[non_exhaustive] +pub enum StartupError { + /// Task spawn failed + #[error("couldn't spawn reclamation task")] + Spawn(#[source] Arc<SpawnError>), +} + +impl From<SpawnError> for StartupError { + fn from(e: SpawnError) -> StartupError { + StartupError::Spawn(Arc::new(e)) + } +} + +/// Tracker corrupted +/// +/// Separate type so we don't expose `PoisonError -> crate::Error` conversion +#[derive(Debug, Clone, Error)] +#[error("poisoned (corrupted)")] +pub(crate) struct TrackerCorrupted; + +impl<T> From<PoisonError<T>> for TrackerCorrupted { + fn from(_: PoisonError<T>) -> TrackerCorrupted { + TrackerCorrupted + } +} + +impl From<TrackerCorrupted> for Error { + fn from(_: TrackerCorrupted) -> Error { + Error::TrackerCorrupted + } +} + +/// Error returned when reclaim task crashes +/// +/// Does not escape the crate; is used for logging. +#[derive(Debug, Clone, Error)] +pub(crate) enum ReclaimCrashed { + /// Previous bug, memory quota tracker is corrupted + #[error("memory tracker corrupted due to previous bug")] + TrackerCorrupted(#[from] TrackerCorrupted), + + /// Bug + #[error("internal error")] + Bug(#[from] Bug), +} + +impl MemoryReclaimedError { + /// Create a new `MemoryReclaimedError` (with no additional information) + pub fn new() -> Self { + MemoryReclaimedError::default() + } +} + +impl From<Error> for MemoryReclaimedError { + fn from(e: Error) -> MemoryReclaimedError { + MemoryReclaimedError(e.into()) + } +} + +impl HasKind for MemoryReclaimedError { + fn kind(&self) -> ErrorKind { + use ErrorKind as EK; + use ReclaimedErrorInner as REI; + match &self.0 { + REI::Collapsed => EK::LocalResourceExhausted, + REI::TrackerError(e) => e.kind(), + } + } +} + +impl HasKind for Error { + fn kind(&self) -> ErrorKind { + use Error as E; + use ErrorKind as EK; + match self { + E::TrackerShutdown => EK::ArtiShuttingDown, + E::AccountClosed => EK::LocalResourceExhausted, + E::ParticipantShutdown => EK::LocalResourceExhausted, + E::TrackerCorrupted => EK::Internal, + E::Bug(e) => e.kind(), + } + } +} diff --git a/crates/tor-memquota/src/internal_prelude.rs b/crates/tor-memquota/src/internal_prelude.rs new file mode 100644 index 000000000..c727496f4 --- /dev/null +++ b/crates/tor-memquota/src/internal_prelude.rs @@ -0,0 +1,63 @@ +//! Internal prelude +//! +//! This file contains most of the imports we wish to use, throughout this crate. +//! +//! Every module does `use crate::internal_prelude::*;` +//! +//! Exceptions: +//! +//! * Names that are private to a module and its submodules +//! are imported to the sub-modules via `use super::*`. +//! (Thus, the sub-module inherits the prelude from its parent.) +//! +//! * Broad names from specific contexts, that are unsuitable for wide imports. +//! For example, individual cell and message names from `tor-cell`, +//! and the types from `tor_proto::stream` other than the high-level `DataStream`. + +pub(crate) use std::{ + cmp::{Ordering, Reverse}, + collections::{BinaryHeap, HashSet}, + fmt::{self, Debug, Display}, + future::Future, + marker::PhantomData, + mem, + ops::{Deref, DerefMut}, + panic::{catch_unwind, AssertUnwindSafe}, + pin::Pin, + sync::{Arc, Mutex, MutexGuard, PoisonError, Weak}, +}; + +pub(crate) use futures::{ + channel::mpsc, + task::{Spawn, SpawnError, SpawnExt as _}, + FutureExt as _, StreamExt as _, +}; + +pub(crate) use { + derive_builder::Builder, + derive_deftly::{define_derive_deftly, Deftly}, + derive_more::{Deref, DerefMut, From, Into}, + educe::Educe, + serde::{Deserialize, Serialize}, + slotmap::SlotMap, + static_assertions::assert_not_impl_any, + thiserror::Error, + tracing::{error, info}, + void::{ResultVoidExt as _, Void}, +}; + +pub(crate) use { + tor_config::ConfigBuildError, + tor_error::{error_report, internal, into_internal, Bug, ErrorKind, HasKind}, + tor_log_ratelim::log_ratelim, + tor_rtcompat::CoarseInstant, +}; + +pub(crate) use crate::{ + config::Config, + drop_bomb::{DropBomb, DropBombCondition}, + drop_reentrancy, + error::{Error, ReclaimCrashed, StartupError, TrackerCorrupted}, + refcount, + utils::{DefaultExtTake, Qty}, +}; diff --git a/crates/tor-memquota/src/lib.rs b/crates/tor-memquota/src/lib.rs new file mode 100644 index 000000000..f6d5ad21b --- /dev/null +++ b/crates/tor-memquota/src/lib.rs @@ -0,0 +1,206 @@ +#![cfg_attr(docsrs, feature(doc_auto_cfg, doc_cfg))] +#![doc = include_str!("../README.md")] + +//! ## Intended behavour +//! +//! In normal operation we try to track as little state as possible, cheaply. +//! We do track total memory use in nominal bytes +//! (but a little approximately). +//! +//! When we exceed the quota, we engage a more expensive algorithm: +//! we build a heap to select oldest victims, and +//! we use the heap to keep reducing memory +//! until we go below a low-water mark (hysteresis). +//! +//! ## Key concepts +//! +//! * **Tracker**: +//! Instance of the memory quota system +//! Each tracker has a notion of how much memory its participants +//! are allowed to use, in aggregate. +//! Tracks memory usage by all the Accounts and Participants. +//! Different Trackers are completely independent. +//! +//! * **Account**: +//! all memory used within the same Account is treated equally, +//! and reclamation also happens on an account-by-account basis. +//! (Each Account is with one Tracker.) +//! +//! * **Participant**: +//! one data structure that uses memory. +//! Each Participant is linked to *one* Account. An account has *one or more* Participants. +//! (An Account can exist with zero Participants, but can't then claim memory.) +//! A Participant provides a `dyn IsParticipant` to the memory system; +//! in turn, the memory system provides the Participant with a `Participation` - +//! a handle for tracking memory alloc/free. +//! +//! * **Child Account**/**Parent Account**: +//! An Account may have a Parent. +//! When a tracker requests memory reclamation from a Parent, +//! it will also request it of all that Parent's Children (but not vice versa). +//! +//! * **Data age**: +//! Each Participant must be able to say what the oldest data is, that it is storing. +//! The reclamation policy is to try to free the oldest data. +//! +//! * **Reclamation**: +//! When a Tracker decides that too much memory is being used, +//! it will select a victim Account based on the data age. +//! It will then ask *every Participant* in that Account, +//! and every Participant in every Child of that Account, +//! to reclaim memory. +//! A Participant responds by freeing at least some memory, +//! according to the reclamation request, and tells the Tracker when it has done so. +//! +//! * **Reclamation strategy**: +//! To avoid too-frequent reclamation, once reclamation has started, +//! it will continue until a low-water mark is reached, significantly lower than the quota. +//! I.e. the system has a hysteresis. +// TODO we haven't implemented the queue wrapper yet +// ! The only currently implemented higher-level Participant is +// ! a queue which responds to a reclamation request +// ! by completely destroying itself and freeing all its data. +//! +//! * **Approximate** (both in time and space): +//! The memory quota system is not completely precise. +//! Participants need not report their use precisely, +//! but the errors should be reasonably small, and bounded. +//! Likewise, the enforcement is not precise: +//! reclamation may start slightly too early, or too late; +//! but the memory use will be bounded below by O(number of participants) +//! and above by O(1) (plus errors from the participants). +//! Reclamation is not immediate, and is dependent on task scheduling; +//! during memory pressure the quota may be exceeded; +//! new allocations are not prevented while attempts at reclamation are ongoing. +//! +// TODO we haven't implemented the queue wrapper yet +// ! * **Queues**: +// ! We provide a higher-level API that wraps an mpsc queue and turns it into a Participant. +// ! +//! ## Ownership and Arc keeping-alive +//! +//! * Somewhere, someone must keep an `Account` to keep the account open. +//! Ie, the principal object corresponding to the accountholder should contain an `Account`. +//! +//! * `Arc<MemoryTracker>` holds `Weak<dyn IsParticipant>`. +//! If the tracker finds the `IsParticipant` has vanished, +//! it assumes this means that the Participant is being destroyed and +//! it can treat all of the memory it claimed as freed. +//! +//! * Each participant holds a `Participation`. +//! A `Participation` may be invalidated by collapse of the underlying Account, +//! which may be triggered in any number of ways. +//! +//! * A `Participation` does *not* keep its `Account` alive. +//! Ie, it has only a weak reference to the Account. +//! +//! * A Participant's implementor of `IsParticipant` may hold a `Participation`. +//! If the `impl IsParticipant` is also the principal accountholder object, +//! it must hold an `Account` too. +//! +//! * Child/parent accounts do not imply any keeping-alive relationship. +//! It's just that a reclamation request to a parent (if it still exists) +//! will also be made to its children. +//! +//! +//! ```text +//! accountholder =======================================>* Participant +//! (impl IsParticipant) +//! || +//! || ^ || +//! || | || +//! || global Weak<dyn>| || +//! || || | || +//! \/* \/ | || +//! | || +//! Account *===========> MemoryTracker ------------------' || +//! || +//! ^ || +//! | \/ +//! | +//! `-------------------------------------------------* Participation +//! +//! +//! +//! accountholder which is also directly the Participant ==============\ +//! (impl IsParticipant) || +//! ^ || +//! || | || +//! || | || +//! || global |Weak<dyn> || +//! || || | || +//! \/ \/ | || +//! || +//! Account *===========> MemoryTracker || +//! || +//! ^ || +//! | \/ +//! | +//! `-------------------------------------------------* Participation +//! +//! ``` + +// @@ begin lint list maintained by maint/add_warning @@ +#![cfg_attr(not(ci_arti_stable), allow(renamed_and_removed_lints))] +#![cfg_attr(not(ci_arti_nightly), allow(unknown_lints))] +#![warn(missing_docs)] +#![warn(noop_method_call)] +#![warn(unreachable_pub)] +#![warn(clippy::all)] +#![deny(clippy::await_holding_lock)] +#![deny(clippy::cargo_common_metadata)] +#![deny(clippy::cast_lossless)] +#![deny(clippy::checked_conversions)] +#![warn(clippy::cognitive_complexity)] +#![deny(clippy::debug_assert_with_mut_call)] +#![deny(clippy::exhaustive_enums)] +#![deny(clippy::exhaustive_structs)] +#![deny(clippy::expl_impl_clone_on_copy)] +#![deny(clippy::fallible_impl_from)] +#![deny(clippy::implicit_clone)] +#![deny(clippy::large_stack_arrays)] +#![warn(clippy::manual_ok_or)] +#![deny(clippy::missing_docs_in_private_items)] +#![warn(clippy::needless_borrow)] +#![warn(clippy::needless_pass_by_value)] +#![warn(clippy::option_option)] +#![deny(clippy::print_stderr)] +#![deny(clippy::print_stdout)] +#![warn(clippy::rc_buffer)] +#![deny(clippy::ref_option_ref)] +#![warn(clippy::semicolon_if_nothing_returned)] +#![warn(clippy::trait_duplication_in_bounds)] +#![deny(clippy::unchecked_duration_subtraction)] +#![deny(clippy::unnecessary_wraps)] +#![warn(clippy::unseparated_literal_suffix)] +#![deny(clippy::unwrap_used)] +#![allow(clippy::let_unit_value)] // This can reasonably be done for explicitness +#![allow(clippy::uninlined_format_args)] +#![allow(clippy::significant_drop_in_scrutinee)] // arti/-/merge_requests/588/#note_2812945 +#![allow(clippy::result_large_err)] // temporary workaround for arti#587 +#![allow(clippy::needless_raw_string_hashes)] // complained-about code is fine, often best +//! <!-- @@ end lint list maintained by maint/add_warning @@ --> +#![allow(clippy::blocks_in_conditions)] // TODO #1176 + +// Internal supporting modules +#[macro_use] +mod drop_bomb; +mod drop_reentrancy; +mod internal_prelude; +#[macro_use] +mod refcount; +mod utils; + +// Modules with public items +mod config; +mod error; +pub mod mtracker; + +//---------- re-exports at the crate root ---------- + +pub use config::{Config, ConfigBuilder}; +pub use error::{Error, MemoryReclaimedError, StartupError}; +pub use mtracker::MemoryQuotaTracker; + +/// `Result` whose `Err` is [`tor_memtrack::Error`](Error) +pub type Result<T> = std::result::Result<T, Error>; diff --git a/crates/tor-memquota/src/mtracker.rs b/crates/tor-memquota/src/mtracker.rs new file mode 100644 index 000000000..33333a930 --- /dev/null +++ b/crates/tor-memquota/src/mtracker.rs @@ -0,0 +1,1047 @@ +//! Memory quota tracker, core and low-level API +//! +//! # Example +//! +//! ``` +//! use std::{collections::VecDeque, sync::{Arc, Mutex}}; +//! use tor_rtcompat::{CoarseInstant, CoarseTimeProvider, PreferredRuntime}; +//! use tor_memquota::{mtracker, MemoryQuotaTracker, MemoryReclaimedError}; +//! use void::{ResultVoidExt, Void}; +//! +//! #[derive(Debug)] +//! struct TrackingQueue(Mutex<Result<Inner, MemoryReclaimedError>>); +//! #[derive(Debug)] +//! struct Inner { +//! partn: mtracker::Participation, +//! data: VecDeque<(Box<[u8]>, CoarseInstant)>, +//! } +//! +//! impl TrackingQueue { +//! fn push(&self, now: CoarseInstant, bytes: Box<[u8]>) -> Result<(), MemoryReclaimedError> { +//! let mut inner = self.0.lock().unwrap(); +//! let inner = inner.as_mut().map_err(|e| e.clone())?; +//! inner.partn.claim(bytes.len())?; +//! inner.data.push_back((bytes, now)); +//! Ok(()) +//! } +//! } +//! +//! impl mtracker::IsParticipant for TrackingQueue { +//! fn get_oldest(&self) -> Option<CoarseInstant> { +//! let inner = self.0.lock().unwrap(); +//! Some(inner.as_ref().ok()?.data.front()?.1) +//! } +//! fn reclaim(self: Arc<Self>) -> mtracker::ReclaimFuture { +//! let mut inner = self.0.lock().unwrap(); +//! *inner = Err(MemoryReclaimedError::new()); +//! Box::pin(async { mtracker::Reclaimed::Collapsing }) +//! } +//! } +//! +//! let runtime = PreferredRuntime::create().unwrap(); +//! let config = tor_memquota::Config::builder().max(1024*1024*1024).build().unwrap(); +//! let trk = MemoryQuotaTracker::new(&runtime, config).unwrap(); +//! let account = trk.new_account(None).unwrap(); +//! +//! let queue: Arc<TrackingQueue> = account.register_participant_with( +//! runtime.now_coarse(), +//! |partn| { +//! Ok::<_, Void>(Arc::new(TrackingQueue(Mutex::new(Ok(Inner { +//! partn, +//! data: VecDeque::new(), +//! }))))) +//! }, +//! ).unwrap().void_unwrap(); +//! +//! queue.push(runtime.now_coarse(), Box::new([0; 24])).unwrap(); +//! ``` +// +// For key internal documentation about the data structure, see the doc comment for +// `struct State` (down in the middle of the file). + +use crate::internal_prelude::*; + +mod bookkeeping; +mod reclaim; +mod total_qty_notifier; + +#[cfg(test)] +mod test; + +use bookkeeping::{BookkeepableQty, ClaimedQty, ParticipQty, TotalQty}; +use total_qty_notifier::TotalQtyNotifier; + +/// Maximum amount we'll "cache" locally in a [`Participation`] +/// +/// ie maximum value of `Participation.cache`. +// +// TODO is this a good amount? should it be configurable? +pub(crate) const MAX_CACHE: Qty = Qty(16384); + +/// Target cache size when we seem to be claiming +const TARGET_CACHE_CLAIMING: Qty = Qty(MAX_CACHE.as_usize() * 3 / 4); +/// Target cache size when we seem to be releasing +#[allow(clippy::identity_op)] // consistency +const TARGET_CACHE_RELEASING: Qty = Qty(MAX_CACHE.as_usize() * 1 / 4); + +//---------- public data types ---------- + +/// Memory data tracker +/// +/// Instance of the memory quota system. +/// +/// Usually found as `Arc<MemoryQuotaTracker>`. +#[derive(Debug)] +pub struct MemoryQuotaTracker { + /// The actual tracker state etc. + state: Mutex<State>, +} + +/// Handle onto an Account +/// +/// An `Account` is a handle. All clones refer to the same underlying conceptual Account. +/// +/// `Account`s are created using [`MemoryQuotaTracker::new_account`]. +#[derive(Educe)] +#[educe(Debug)] +pub struct Account { + /// The account ID + aid: refcount::Ref<AId>, + + /// The underlying tracker + #[educe(Debug(ignore))] + tracker: Arc<MemoryQuotaTracker>, +} + +/// Weak handle onto an Account +/// +/// Like [`Account`], but doesn't keep the account alive. +/// Must be upgraded before use. +// +// Doesn't count for ARecord.account_clones +// +// We can't lift out Arc, so that the caller sees `Arc<Account>`, +// because an Account is Arc<MemoryQuotaTracker> plus AId, +// not Arc of something account-specific. +#[derive(Clone, Educe)] +#[educe(Debug)] +pub struct WeakAccount { + /// The account ID + aid: AId, + + /// The underlying tracker + #[educe(Debug(ignore))] + tracker: Weak<MemoryQuotaTracker>, +} + +/// Handle onto a participant's participation in a tracker +/// +/// `Participation` is a handle. All clones are for use by the same conceptual Participant. +/// It doesn't keep the underlying Account alive. +/// +/// `Participation`s are created by registering new participants, +/// for example using [`Account::register_participant`]. +/// +/// Variables of this type are often named `partn`. +#[derive(Debug)] +pub struct Participation { + /// Participant id + pid: refcount::Ref<PId>, + + /// Account id + aid: AId, + + /// The underlying tracker + tracker: Weak<MemoryQuotaTracker>, + + /// Quota we have preemptively claimed for use by this Account + /// + /// Has been added to `PRecord.used`, + /// but not yet returned by `Participation::claim`. + /// + /// This cache field arranges that most of the time we don't have to hammer a + /// single cache line. + /// + /// The value here is bounded by a configured limit. + /// + /// Invariants on memory accounting: + /// + /// * `Participation.cache < configured limit` + /// * `PRecord.used = Participation.cache + Σ Participation::claim - Σ P'n::release` + /// except if `PRecord` has been deleted + /// (ie when we aren't tracking any more and think the Participant is `Collapsing`). + /// * `Σ PRecord.used = State.total_used` + /// + /// Enforcement of these invariants is partially assured by + /// types in [`bookkeeping`]. + cache: ClaimedQty, +} + +/// Participants provide an impl of the hooks in this trait +/// +/// Trait implemented by client of the memtrack API. +/// +/// # Panic handling, "unwind safety" +/// +/// If these methods panic, the memory tracker will tear down its records of the +/// participant, preventing future allocations. +/// +/// But, it's not guaranteed that these methods on `IsParticipant` won't be called again, +/// even if they have already panicked on a previous occasion. +/// Thus the implementations might see "broken invariants" +/// as discussed in the docs for `std::panic::UnwindSafe`. +/// +/// Nevertheless we don't make `RefUnwindSafe` a supertrait of `IsParticipant`. +/// That would force the caller to mark *all* their methods unwind-safe, +/// which is unreasonable (and probably undesirable). +/// +/// Variables which are `IsParticipant` are often named `particip`. +pub trait IsParticipant: Debug + Send + Sync + 'static { + /// Return the age of the oldest data held by this Participant + /// + /// `None` means this Participant holds no data. + /// + /// # Performance and reentrancy + /// + /// This function runs with the `MemoryQuotaTracker`'s internal global lock held. + /// Therefore: + /// + /// * It must be fast. + /// * it *must not* call back into methods from [`tracker`](crate::mtracker). + /// * It *must not* even `Clone` or `Drop` a [`MemoryQuotaTracker`], + /// [`Account`], or [`Participation`]. + fn get_oldest(&self) -> Option<CoarseInstant>; + + /// Start memory reclamation + /// + /// The Participant should start to free all of its memory, + /// and then return `Reclaimed::Collapsing`. + // + // In the future: + // + // Should free *at least* all memory at least as old as discard_... + // + // v1 of the actual implementation might not have `discard_everything_as_old_as` + // and `but_can_stop_discarding_...`, + // and might therefore only support Reclaimed::Collapsing + fn reclaim( + self: Arc<Self>, + // Future: + // discard_everything_as_old_as_this: RoughTime, + // but_can_stop_discarding_after_freeing_this_much: Qty, + ) -> ReclaimFuture; +} + +/// Future returned by the [`IsParticipant::reclaim`] reclamation request +pub type ReclaimFuture = Pin<Box<dyn Future<Output = Reclaimed> + Send + Sync>>; + +/// Outcome of [`IsParticipant::reclaim`] +#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)] +#[non_exhaustive] +pub enum Reclaimed { + /// Participant is responding to reclamation by collapsing completely. + /// + /// All memory will be freed and `release`'d soon (if it hasn't been already). + /// `MemoryQuotaTracker` should forget the Participant and all memory it used, right away. + /// + /// Currently this is the only supported behaviour. + Collapsing, + // Future: + // /// Participant has now reclaimed some memory as instructed + // /// + // /// If this is not sufficient, tracker must call reclaim() again. + // /// (We may not want to implement Partial right away but the API + // /// ought to support it so let's think about it now, even if we don't implement it.) + // Partial, +} + +//---------- principal data structure ---------- + +slotmap::new_key_type! { + /// Identifies an Account + /// + /// After an account is torn down, the `AId` becomes invalid + /// and attempts to use it will give an error. + /// + /// The same `AId` won't be reused for a later Account. + struct AId; + + /// Identifies a Participant within an Account + /// + /// Ie, PId is scoped within in the context of an account. + /// + /// As with `AId`, a `PId` is invalid after the + /// participation is torn down, and is not reused. + struct PId; +} + +/// Memory tracker inner, including mutable state +/// +/// # Module internal documentation +/// +/// ## Data structure +/// +/// * [`MemoryQuotaTracker`] contains mutex-protected `State`. +/// * The `State` contains a [`SlotMap`] of account records [`ARecord`]. +/// * Each `ARecord` contains a `SlotMap` of participant records [`PRecord`]. +/// +/// The handles [`Account`], [`WeakAccount`], and [`Participation`], +/// each contain a reference (`Arc`/`Weak`) to the `MemoryQuotaTracker`, +/// and the necessary slotmap keys. +/// +/// The `ARecord` and `PRecord` each contain a reference count, +/// which is used to clean up when all the handles are gone. +/// +/// The slotmap keys which count for the reference count (ie, strong references) +/// are stored as [`refcount::Ref`], +/// which helps assure correct reference counting. +/// (Bare ids [`AId`] and [`PId`] are weak references.) +/// +/// ## Data structure lookup +/// +/// Given a reference to the tracker, and some ids, the macro `find_in_tracker!` +/// is used to obtain mutable references to the `ARecord` and (if applicable) `PRecord`. +/// +/// ## Bookkeeping +/// +/// We use separate types for quantities of memory in various "states", +/// rather than working with raw quantities. +/// +/// The types, and the legitimate transactions, are in `bookkeeping`. +/// +/// ## Reentrancy (esp. `Drop` and `Clone`) +/// +/// When the handle structs are dropped or cloned, they must manipulate the refcount(s). +/// So they must take the lock. +/// Therefore, an `Account` and `Participation` may not be dropped with the lock held! +/// +/// Internally, this is actually fairly straightforward: +/// we take handles by reference, and constructors only make them at the last moment on return, +/// so our internal code here, in this module, doesn't have owned handles. +/// +/// We also need to worry about reentrantly reentering the tracker code, from user code. +/// The user supplies a `dyn IsParticipant`. +/// The principal methods are from [`IsParticipant`], +/// for which we handle reentrancy in the docs. +/// But we also implicitly invoke its `Drop` impl, which might in turn drop stuff of ours, +/// such as [`Account`]s and [`Participation`]s, whose `Drop` impls need to take our lock. +/// To make sure this isn't done reentrantly, we have a special newtype around it, +/// and defer some of our drops during reclaim. +/// That's in `drop_reentrancy` and `tracker::reclaim::deferred_drop`. +/// +/// The `Debug` impl isn't of concern, since we don't call it ourselves. +/// And we don't rely on it being `Clone`, since it's in an `Arc`. +/// +/// ## Drop bombs +/// +/// With `#[cfg(test)]`, several of our types have "drop bombs": +/// they cause a panic if dropped inappropriately. +/// This is intended to detect bad code paths during testing. +#[derive(Debug, Deref, DerefMut)] +struct State { + /// Global parts of state + /// + /// Broken out to allow passing both + /// `&mut Global` and `&mut ARecord`/`&mut PRecord` + /// to some function(s). + #[deref] + #[deref_mut] + global: Global, + + /// Accounts + accounts: SlotMap<AId, ARecord>, +} + +/// Global parts of `State` +#[derive(Debug)] +struct Global { + /// Total memory used + /// + /// Wrapper type for ensuring we wake up the reclaimation task + total_used: TotalQtyNotifier, + + /// Configuration + config: Config, +} + +/// Account record, within `State.accounts` +#[derive(Debug)] +#[must_use = "don't just drop, call auto_release"] +struct ARecord { + /// Number of clones of `Account`; to know when to tear down the account + refcount: refcount::Count<AId>, + + /// Child accounts + children: Vec<AId>, + + /// Participants linked to this Account + ps: SlotMap<PId, PRecord>, +} + +/// Participant record, within `ARecord.ps` +#[derive(Debug)] +#[must_use = "don't just drop, call auto_release"] +struct PRecord { + /// Number of clones of `Participation`; to know when to tear down the participant + refcount: refcount::Count<PId>, + + /// Memory usage of this participant + /// + /// Not 100% accurate, can lag, and be (boundedly) an overestimate + used: ParticipQty, + + /// The hooks provided by the Participant + particip: drop_reentrancy::ProtectedWeak<dyn IsParticipant>, +} + +//#################### IMPLEMENTATION #################### + +/// Given a `&Weak<MemoryQuotaTracker>`, find an account and maybe participant +/// +/// ### Usage templates +/// +/// ```rust,ignore +/// find_in_tracker! { +/// weak_tracker => + tracker, state; +/// aid => arecord; +/// [ pid => precord; ] +/// [ ?Error | ?None ] +/// }; +/// +/// find_in_tracker! { +/// strong_tracker => state; +/// .. // as above +/// }; +/// ``` +/// +/// ### Input expressions (value arguments to the macro0 +/// +/// * `weak_tracker: &Weak<MemoryQuotaTracker>` (or equivalent) +/// * `strong_tracker: &MemoryQuotaTracker` (or equivalent) +/// * `aid: AId` +/// * `pid: PId` +/// +/// ### Generated bindings (identifier arguments to the macro) +/// +/// * `tracker: Arc<MemoryQuotaTracker>` +/// * `state: &mut State` (borrowed from a `MutexGuard<State>` borrowed from `tracker`) +/// * `arecord: &mut ARecord` (mut borrowed from `state.accounts`) +/// * `precord: &mut PRecord` (mut borrowed from `arecord.ps`) +/// +/// There is no access to the `MutexGuard` itself. +/// For control of the mutex release point, place `find_in_tracker!` in an enclosing block. +/// +/// ### Error handling +/// +/// If the tracker, account, or participant, can't be found, +/// the macro returns early from the enclosing scope (using `?`). +/// +/// If `Error` is specified, applies `?` to `Err(Error::...)`. +/// If `None` is specified, just returns `None` (by applying `?` to None`). +// +// This has to be a macro because it makes a self-referential set of bindings. +// Input syntax is a bit janky because macro_rules is so bad. +// For an internal macro with ~9 call sites it's not worth making a big parsing contraption. +macro_rules! find_in_tracker { { + // This `+` is needed because otherwise it's LL1-ambiguous and macro_rules can't cope + $tracker_input:expr => $( + $tracker:ident, )? $state:ident; + $aid:expr => $arecord:ident; + $( $pid:expr => $precord:ident; )? + // Either `Error` or None, to be passed to `find_in_tracker_eh!($eh ...: ...)` + // (We need this to be an un-repeated un-optional binding, because + // it is used within some other $( ... )?, and macro_rules gets confused.) + ? $eh:tt +} => { + let tracker = &$tracker_input; + $( + let $tracker: Arc<MemoryQuotaTracker> = find_in_tracker_eh!( + $eh TrackerShutdown: + tracker.upgrade() + ); + let tracker = &$tracker; + )? + let mut state: MutexGuard<State> = find_in_tracker_eh!( + $eh TrackerCorrupted: + tracker.state.lock().ok() + ); + let $state: &mut State = &mut *state; + let aid: AId = $aid; + let $arecord: &mut ARecord = find_in_tracker_eh!( + $eh AccountClosed: + $state.accounts.get_mut(aid) + ); + $( + let pid: PId = $pid; + let $precord: &mut PRecord = find_in_tracker_eh!( + $eh ParticipantShutdown: + $arecord.ps.get_mut(pid) + ); + )? +} } +/// Error handling helper for `find_in_tracker` +macro_rules! find_in_tracker_eh { + { None $variant:ident: $result:expr } => { $result? }; + { Error $variant:ident: $result:expr } => { $result.ok_or(Error::$variant)? }; +} + +//========== impls on public types, including public methods and trait impls ========== + +//---------- MemoryQuotaTracker ---------- + +impl MemoryQuotaTracker { + /// Set up a new `MemoryDataTracker` + pub fn new<R: Spawn>(runtime: &R, config: Config) -> Result<Arc<Self>, StartupError> { + let (reclaim_tx, reclaim_rx) = mpsc::channel(0 /* plus num_senders, ie 1 */); + let total_used = TotalQtyNotifier::new_zero(reclaim_tx); + + let global = Global { total_used, config }; + let accounts = SlotMap::default(); + let state = Mutex::new(State { global, accounts }); + let tracker = Arc::new(MemoryQuotaTracker { state }); + + // We don't provide a separate `launch_background_tasks`, because this task doesn't + // wake up periodically, or, indeed, do anything until the tracker is used. + + let for_task = Arc::downgrade(&tracker); + runtime.spawn(reclaim::task(for_task, reclaim_rx))?; + + Ok(tracker) + } + + /// Make a new `Account` + /// + /// To actually record memory usage, a Participant must be added. + // + // Right now, parent can't be changed after construction of an Account, + // so circular accounts are impossible. + // But, we might choose to support that in the future. + // Circular parent relationships might need just a little care + // in the reclamation loop (to avoid infinitely looping), + // but aren't inherently unsupportable. + #[allow(clippy::redundant_closure_call)] // We have IEFEs for good reaons + pub fn new_account(self: &Arc<Self>, parent: Option<&Account>) -> crate::Result<Account> { + let mut state = self.lock()?; + + let parent_aid_good = parent + .map(|parent| { + // Find and check the requested parent's Accountid + + let parent_aid = *parent.aid; + let parent_arecord = state + .accounts + .get_mut(parent_aid) + .ok_or(Error::AccountClosed)?; + + // Can we insert the new child without reallocating? + if !parent_arecord.children.spare_capacity_mut().is_empty() { + return Ok(parent_aid); + } + + // No. Well, let's do some garbage collection. + // (Otherwise .children might grow without bound as accounts come and go) + // + // We would like to scan the accounts array while mutating this account. + // Instead, steal the children array temporarily and put the filtered one back. + // Must be infallible! + // + // The next line can't be in the closure (confuses borrowck) + let mut parent_children = mem::take(&mut parent_arecord.children); + (|| { + parent_children.retain(|child_aid| state.accounts.contains_key(*child_aid)); + + // Put the filtered list back, so sanity is restored. + state + .accounts + .get_mut(parent_aid) + .expect("parent vanished!") + .children = parent_children; + })(); + + Ok::<_, Error>(parent_aid) + }) + .transpose()?; + + // We have resolved the parent AId and prepared to add the new account to its list of + // children. We still hold the lock, so nothing can have changed. + + // commitment - infallible IEFE assures that so we don't do half of it + Ok((|| { + let aid = refcount::slotmap_insert(&mut state.accounts, |refcount| ARecord { + refcount, + children: vec![], + ps: SlotMap::default(), + }); + + if let Some(parent_aid_good) = parent_aid_good { + state + .accounts + .get_mut(parent_aid_good) + .expect("parent vanished!") + .children + .push(*aid); + } + + let tracker = self.clone(); + Account { aid, tracker } // don't make this fallible, see above. + })()) + } + + /// Obtain the lock on the state + fn lock(&self) -> Result<MutexGuard<State>, TrackerCorrupted> { + Ok(self.state.lock()?) + } +} + +//---------- Account ---------- + +impl Account { + /// Register a new Participant + /// + /// Returns the [`Participation`], which can be used to record memory allocations. + /// + /// Often, your implementation of [`IsParticipant`] wants to contain the [`Participation`]. + /// If so, use [`register_participant_with`](Account::register_participant_with) instead. + pub fn register_participant( + &self, + particip: Weak<dyn IsParticipant>, + ) -> Result<Participation, Error> { + let aid = *self.aid; + find_in_tracker! { + self.tracker => state; + aid => arecord; + ?Error + } + + let (pid, cache) = refcount::slotmap_try_insert(&mut arecord.ps, |refcount| { + let mut precord = PRecord { + refcount, + used: ParticipQty::ZERO, + particip: drop_reentrancy::ProtectedWeak::new(particip), + }; + let cache = + state + .global + .total_used + .claim(&mut precord, MAX_CACHE, &state.global.config)?; + Ok::<_, Error>((precord, cache)) + })?; + + let tracker = Arc::downgrade(&self.tracker); + Ok(Participation { + tracker, + pid, + aid, + cache, + }) + } + + /// Set the callbacks for a Participant (identified by its weak ids) + fn set_participant_callbacks( + &self, + aid: AId, + pid: PId, + particip: drop_reentrancy::ProtectedWeak<dyn IsParticipant>, + ) -> Result<(), Error> { + find_in_tracker! { + self.tracker => state; + aid => arecord; + pid => precord; + ?Error + } + precord.particip = particip; + Ok(()) + } + + /// Register a new Participant using a constructor + /// + /// Passes `constructor` a [`Participation`] for the nascent Participant. + /// Returns the `P: IsParticipant` provided by the constructor. + /// + /// For use when your `impl `[`IsParticipant`] wants to own the `Participation`. + /// + /// # Re-entrancy guarantees + /// + /// The `Participation` *may* be used by `constructor` for claiming memory use, + /// even during construction. + /// `constructor` may also clone the `Participation`, etc. + /// + /// Reclamation callbacks (via the `P as IsParticipant` impl) cannot occur + /// until `constructor` returns. + /// + /// # Error handling + /// + /// Failures can occur before `constructor` is called, + /// or be detected afterwards. + /// If a failure is detected after `constructor` returns, + /// the `Arc<P>` from `constructor` will be dropped + /// (resulting in `P` being dropped, unless `constructor` kept another clone of it). + /// + /// `constructor` may also fail (throwing a different error type, `E`), + /// in which case `register_participant_with` returns `Ok(Err(E))`. + /// + /// On successful setup of the Participant, returns `Ok(Ok(Arc<P>))`. + pub fn register_participant_with<P: IsParticipant, E>( + &self, + now: CoarseInstant, + constructor: impl FnOnce(Participation) -> Result<Arc<P>, E>, + ) -> Result<Result<Arc<P>, E>, Error> { + use std::sync::atomic::{AtomicBool, Ordering}; + + /// Temporary participant, which stands in during constructon + #[derive(Debug)] + struct TemporaryParticipant { + /// The age, which is right now. We hope this is all fast! + now: CoarseInstant, + /// Did someone call reclaim() ? + collapsing: AtomicBool, + } + + impl IsParticipant for TemporaryParticipant { + fn get_oldest(&self) -> Option<CoarseInstant> { + Some(self.now) + } + fn reclaim(self: Arc<Self>) -> ReclaimFuture { + self.collapsing.store(true, Ordering::Release); + Box::pin(async { Reclaimed::Collapsing }) + } + } + + let temp_particip = Arc::new(TemporaryParticipant { + now, + collapsing: false.into(), + }); + + let partn = self.register_participant(Arc::downgrade(&temp_particip) as _)?; + let aid = partn.aid; + let pid_weak = *partn.pid; + + // We don't hold the state lock here. register_participant took it and released it. + // This is important, because the constructor might call claim! + // (And, also, we don't want the constructor panicking to poison the whole tracker.) + // But it means there can be quite a lot of concurrent excitement, + // including, theoretically, a possible reclaim. + let particip = match constructor(partn) { + Ok(y) => y, + Err(e) => return Ok(Err(e)), + }; + let particip = drop_reentrancy::ProtectedArc::new(particip); + + // IEFE prevents use from accidentally dropping `particip` until we mean to + let r = (|| { + let weak = { + let weak = particip.downgrade(); + + // Trait cast, from Weak<P> to Weak<dyn IsParticipant>. + // We can only do this for a primitive, so we must unprotect + // the Weak, converr it, and protect it again. + drop_reentrancy::ProtectedWeak::new(weak.unprotect() as _) + }; + self.set_participant_callbacks(aid, pid_weak, weak)?; + + if temp_particip.collapsing.load(Ordering::Acquire) { + return Err(Error::ParticipantShutdown); + } + Ok(()) + })(); + + let particip = particip.promise_dropping_is_ok(); + r?; + Ok(Ok(particip)) + } + + /// Obtains a handle for the `MemoryQuotaTracker` + pub fn tracker(&self) -> Arc<MemoryQuotaTracker> { + self.tracker.clone() + } + + /// Downgrade to a weak handle for the same Account + pub fn downgrade(&self) -> WeakAccount { + WeakAccount { + aid: *self.aid, + tracker: Arc::downgrade(&self.tracker), + } + } +} + +impl Clone for Account { + fn clone(&self) -> Account { + let tracker = self.tracker.clone(); + let aid = (|| { + let aid = *self.aid; + find_in_tracker! { + tracker => state; + aid => arecord; + ?None + } + let aid = refcount::Ref::new(aid, &mut arecord.refcount).ok()?; + // commitment point + Some(aid) + })() + .unwrap_or_else(|| { + // Either the account has been closed, or our refcount overflowed. + // Return a busted `Account`, which always fails when we try to use it. + // + // If the problem was a refcount overflow, we're technically violating the + // documented behaviour, since the returned `Account` isn't equivalent + // to the original. We could instead choose to tear down the Account; + // that would be legal; but it's a lot of code to marginally change the + // behaviour for a very unlikely situation. + refcount::Ref::null() + }); + Account { aid, tracker } + } +} + +impl Drop for Account { + fn drop(&mut self) { + (|| { + find_in_tracker! { + self.tracker => state; + *self.aid => arecord; + ?None + } + if let Some(refcount::Garbage(mut removed)) = + slotmap_dec_ref!(&mut state.accounts, self.aid.take(), &mut arecord.refcount) + { + // This account is gone. Automatically release everything. + removed.auto_release(state); + } + Some(()) + })() + .unwrap_or_else(|| { + // Account has been torn down. Dispose of the strong ref. + // (This has no effect except in cfg(test), when it defuses the drop bombs) + self.aid.take().dispose_container_destroyed(); + }); + } +} + +//---------- WeakAccount ---------- + +impl WeakAccount { + /// Upgrade to an `Account`, if the account still exists + pub fn upgrade(&self) -> crate::Result<Account> { + let aid = self.aid; + // (we must use a block, and can't use find_in_tracker's upgrade, because borrowck) + let tracker = self.tracker.upgrade().ok_or(Error::TrackerShutdown)?; + let aid = { + find_in_tracker! { + tracker => state; + aid => arecord; + ?Error + } + refcount::Ref::new(aid, &mut arecord.refcount)? + // commitment point + }; + Ok(Account { aid, tracker }) + } + + /// Obtains a handle onto the `MemoryQuotaTracker` + /// + /// The returned handle is itself weak, and needs to be upgraded before use. + pub fn tracker(&self) -> Weak<MemoryQuotaTracker> { + self.tracker.clone() + } +} + +//---------- Participation ---------- + +impl Participation { + /// Record that some memory has been (or will be) allocated + pub fn claim(&mut self, want: usize) -> crate::Result<()> { + self.claim_qty(Qty(want)) + } + + /// Record that some memory has been (or will be) allocated (using `Qty`) + pub(crate) fn claim_qty(&mut self, want: Qty) -> crate::Result<()> { + if let Some(got) = self.cache.split_off(want) { + return got.claim_return_to_participant(); + } + + find_in_tracker! { + self.tracker => + tracker, state; + self.aid => arecord; + *self.pid => precord; + ?Error + }; + + let mut claim = |want| -> Result<ClaimedQty, _> { + state + .global + .total_used + .claim(precord, want, &state.global.config) + }; + let got = claim(want)?; + + if want <= TARGET_CACHE_CLAIMING { + // While we're here, fill the cache to TARGET_CACHE_CLAIMING. + // Cannot underflow: cache < want (since we failed at `got` earlier + // and we've just checked want <= TARGET_CACHE_CLAIMING. + let want_more_cache = Qty(*TARGET_CACHE_CLAIMING - *self.cache.as_raw()); + if let Ok(add_cache) = claim(want_more_cache) { + // On error, just don't do this; presumably the error will show up later + // (we mustn't early exit here, because we've got the claim in our hand). + self.cache.merge_into(add_cache); + } + } + got.claim_return_to_participant() + } + + /// Record that some memory has been (or will be) freed by a participant + pub fn release(&mut self, have: usize) // infallible + { + self.release_qty(Qty(have)); + } + + /// Record that some memory has been (or will be) freed by a participant (using `Qty`) + pub(crate) fn release_qty(&mut self, have: Qty) // infallible + { + let have = ClaimedQty::release_got_from_participant(have); + self.cache.merge_into(have); + if self.cache > MAX_CACHE { + match (|| { + find_in_tracker! { + self.tracker => + tracker, state; + self.aid => arecord; + *self.pid => precord; + ?None + } + let return_from_cache = Qty(*self.cache.as_raw() - *TARGET_CACHE_RELEASING); + let from_cache = self.cache.split_off(return_from_cache).expect("impossible"); + state.global.total_used.release(precord, from_cache); + Some(()) + })() { + Some(()) => {} // we've given our cache back to the tracker + None => { + // account (or whole tracker!) is gone + // throw away the cache so that we don't take this path again for a bit + self.cache.take().dispose_participant_destroyed(); + } + } + } + } + + /// Obtain a handle onto the account + /// + /// The returned handle is weak, and needs to be upgraded before use, + /// since a [`Participation`] doesn't keep its Account alive. + /// + /// The returned `WeakAccount` is equivalent to + /// all the other account handles for the same account. + pub fn account(&self) -> WeakAccount { + WeakAccount { + aid: self.aid, + tracker: self.tracker.clone(), + } + } +} + +impl Clone for Participation { + fn clone(&self) -> Participation { + let aid = self.aid; + let cache = ClaimedQty::ZERO; + let tracker: Weak<_> = self.tracker.clone(); + let pid = (|| { + let pid = *self.pid; + find_in_tracker! { + self.tracker => + tracker_strong, state; + aid => _arecord; + pid => precord; + ?None + } + let pid = refcount::Ref::new(pid, &mut precord.refcount).ok()?; + // commitment point + Some(pid) + })() + .unwrap_or_else(|| { + // The account has been closed, the participant torn down, or the refcount + // overflowed. We can a busted `Participation`. + // + // We *haven't* incremented the refcount, so we mustn't return pid as a strong + // reference. We aren't supposed to count towards PRecord.refcount, we we *can* + // return the weak reference aid. (`refcount` type-fu assures this is correct.) + // + // If the problem was refcount overflow, we're technically violating the + // documented behaviour. This is OK; see comment in `<Account as Clone>::clone`. + refcount::Ref::null() + }); + Participation { + aid, + pid, + cache, + tracker, + } + } +} + +impl Drop for Participation { + fn drop(&mut self) { + (|| { + find_in_tracker! { + self.tracker => + tracker_strong, state; + self.aid => arecord; + *self.pid => precord; + ?None + } + // release the cached claim + let from_cache = self.cache.take(); + state.global.total_used.release(precord, from_cache); + + if let Some(refcount::Garbage(mut removed)) = + slotmap_dec_ref!(&mut arecord.ps, self.pid.take(), &mut precord.refcount) + { + // We might not have called `release` on everything, so we do that here. + removed.auto_release(&mut state.global); + } + Some(()) + })() + .unwrap_or_else(|| { + // Account or Participation or tracker destroyed. + // (This has no effect except in cfg(test), when it defuses the drop bombs) + self.pid.take().dispose_container_destroyed(); + self.cache.take().dispose_participant_destroyed(); + }); + } +} + +//========== impls on internal types ========== + +impl State { + /// Obtain all of the descendants of `parent_aid` according to the Child relation + /// + /// The returned `HashSet` includes `parent_aid`, its children, + /// their children, and so on. + /// + /// Used in the reclaimation algorithm in [`reclaim`]. + fn get_aid_and_children_recursively(&self, parent_aid: AId) -> HashSet<AId> { + let mut out = HashSet::<AId>::new(); + let mut queue: Vec<AId> = vec![parent_aid]; + while let Some(aid) = queue.pop() { + let Some(arecord) = self.accounts.get(aid) else { + // shouldn't happen but no need to panic + continue; + }; + if out.insert(aid) { + queue.extend(arecord.children.iter().cloned()); + } + } + out + } +} + +impl ARecord { + /// Release all memory that this account's participants claimed + fn auto_release(&mut self, global: &mut Global) { + for (_pid, mut precord) in self.ps.drain() { + precord.auto_release(global); + } + } +} + +impl PRecord { + /// Release all memory that this participant claimed + fn auto_release(&mut self, global: &mut Global) { + let for_teardown = self.used.for_participant_teardown(); + global.total_used.release(self, for_teardown); + } +} diff --git a/crates/tor-memquota/src/mtracker/bookkeeping.rs b/crates/tor-memquota/src/mtracker/bookkeeping.rs new file mode 100644 index 000000000..ec1241196 --- /dev/null +++ b/crates/tor-memquota/src/mtracker/bookkeeping.rs @@ -0,0 +1,308 @@ +//! Quantity bookkeeping +//! +//! Newtypes which wrap up a `Qty` (an amount of memory), +//! and which assure proper accounting. +//! +//! Methods are provided for the specific transactions which are correct, +//! in the accounting scheme in [`tracker`](super). +//! So these types embody the data structure (fields and invariants) from `tracker`. +//! +//! # Panics +//! +//! In tests, these types panic if they are dropped when nonzero, +//! if that's against the rules. + +use super::*; + +define_derive_deftly! { + /// Implement [`BookkeptQty`] and its supertraits + /// + /// By default, dropping when nonzero is forbidden, + /// and you must have a field `bomb: `[`DropBombCondition`]. + /// `#[deftly(allow_nonzero_drop)]` suppresses this. + BookkeptQty = + + ${defcond BOMB not(tmeta(allow_nonzero_drop))} + + impl BookkeepableQty for $ttype { + const ZERO: $ttype = $ttype { + raw: Qty(0), + ${if BOMB { + bomb: DropBombCondition::new_armed(), + }} + }; + + fn as_raw(&self) -> Qty { + self.raw + } + } + + impl<Rhs: BookkeepableQty> PartialEq<Rhs> for $ttype { + fn eq(&self, other: &Rhs) -> bool { + self.as_raw().eq(&other.as_raw()) + } + } + impl<Rhs: BookkeepableQty> PartialOrd<Rhs> for $ttype { + fn partial_cmp(&self, other: &Rhs) -> Option<Ordering> { + self.as_raw().partial_cmp(&other.as_raw()) + } + } + + impl DefaultExtTake for $ttype {} + + impl BookkeptQty for $ttype { + fn from_raw(q: Qty) -> Self { + $ttype { + raw: q, + ${if BOMB { + bomb: DropBombCondition::new_armed(), + }} + } + } + fn into_raw(mut self) -> Qty { + mem::replace(&mut self.raw, Qty(0)) + } + } + + assert_not_impl_any!($ttype: Clone, Into<Qty>, From<Qty>); + + ${if BOMB { + #[cfg(test)] + impl Drop for $ttype { + fn drop(&mut self) { + drop_bomb_disarm_assert!(self.bomb, self.raw == Qty(0)); + } + } + }} +} + +/// Memory quantities that can work with bookkept quantities +/// +/// This trait doesn't imply any invariants; +/// it merely provides read-only access to the underlying value, +/// and ways to make a zero. +/// +/// Used by the derived `PartialEq` and `PartialOrd` impls on bookkept quantities. +/// +/// Implemented by hand for `Qty`. +/// +/// Implemented for bookkept types, along with `BookkeptQty`, by +/// [`#[derive_deftly(BookKept)]`](derive_deftly_template_BookkeptQty). +pub(super) trait BookkeepableQty: Default { + /// Zero (default value) + const ZERO: Self; + + /// Inspect as a raw untracked Qty + fn as_raw(&self) -> Qty; +} + +/// Bookkept memory quantities +/// +/// Each bookkept quantity implements this trait, +/// and has a single field `raw` of type `Qty`. +/// +/// Should be Implemented by +/// [`#[derive_deftly(BookKept)]`](derive_deftly_template_BookkeptQty) +/// and for raw `Qty`. +/// +/// # CORRECTNESS +/// +/// All accesses to `raw`, or calls to `from_raw` or `into_raw`, +/// should be made from transaction functions, +/// which modify one or more bookkept quantities together, +/// preserving the invariants. +/// +/// `raw` may be accessed mutably by such functions, but a bookkept quantity type +/// should be constructed only with `from_raw` and should not be moved out of. +trait BookkeptQty: BookkeepableQty + DefaultExtTake { + /// Make a new bookkept quantity from a raw untracked Qty + /// + fn from_raw(q: Qty) -> Self; + + /// Unwrap into a raw untracked Qty + fn into_raw(self) -> Qty; +} + +impl BookkeepableQty for Qty { + const ZERO: Qty = Qty(0); + + fn as_raw(&self) -> Qty { + *self + } +} + +/// Total used, [`TotalQtyNotifier`].`total_used`, found in [`State`].`total_used`. +/// +/// Can be "poisoned", preventing further claims. +/// (We mark it poisoned if the reclamation task crashes, +/// since in that situation we don't want to continue to use memory, unboundedly.) +// +// Poisoned is indicated by We setting to `MAX`. +#[derive(Default, Debug, Deftly, derive_more::Display)] +#[derive_deftly(BookkeptQty)] +#[deftly(allow_nonzero_drop)] // Dropped only when the whole tracker is dropped +pub(super) struct TotalQty { + /// See [`BookkeptQty`] + raw: Qty, +} + +/// Qty used by a participant, found in [`PRecord`].`used`. +/// +/// The tracker data structure has one of these for each Participant. +/// +/// This is the total amount `claim`ed, plus the caches in each `Participation`. +#[derive(Default, Debug, Deftly, derive_more::Display)] +#[derive_deftly(BookkeptQty)] +#[display(fmt = "{raw}")] +pub(super) struct ParticipQty { + /// See [`BookkeptQty`] + raw: Qty, + + /// See [`BookkeptQty`] + bomb: DropBombCondition, +} + +/// "Cached" claim, on behalf of a Participant +/// +/// Found in [`Participation`].`cache`, +/// and accounted to the Participant (ie, included in `ParticipQty`). +/// +/// Also used as a temporary variable in `claim()` and `release()` functions. +/// When we return to the participant, outside the tracker, we +/// essentially throw this away, since we don't give the caller any representation +/// to store. The participant is supposed to track this separately somehow. +#[derive(Default, Debug, Deftly, derive_more::Display)] +#[derive_deftly(BookkeptQty)] +#[display(fmt = "{raw}")] +#[must_use] +pub(super) struct ClaimedQty { + /// See [`BookkeptQty`] + raw: Qty, + + /// See [`BookkeptQty`] + bomb: DropBombCondition, +} + +impl TotalQty { + /// Claim a quantity, increasing the tracked amounts + /// + /// This module doesn't know anything about the memory quota, + /// so this doesn't do the quota check. + /// + /// The only caller is [`Participation::claim`]. + pub(super) fn claim(&mut self, p_used: &mut ParticipQty, want: Qty) -> Option<ClaimedQty> { + // If poisoned, this add will fail (unless want is 0) + let new_self = self.raw.checked_add(*want)?; + if new_self == usize::MAX { + // This would poison us. If this happens, someone has gone mad, since + // we can't have allocated usize::MAX in total. We'll be reclaiming already. + // We don't want to poison ourselves in this situation. Hopefully the reclaim + // will collapse the errant participants. + return None; + } + let new_p_used = p_used.raw.checked_add(*want)?; + // commit + self.raw = Qty(new_self); + p_used.raw = Qty(new_p_used); + Some(ClaimedQty::from_raw(want)) + } + + /// Release a quantity, decreasing the tracked amounts + /// + /// (Handles underflow by saturating; returning an error is not going to be useful.) + pub(super) fn release(&mut self, p_used: &mut ParticipQty, have: ClaimedQty) { + let have = have.into_raw(); + *p_used.raw = p_used.raw.saturating_sub(*have); + + if self.raw != Qty::MAX { + // Don't unpoison + *self.raw = self.raw.saturating_sub(*have); + } + } + + /// Declare this poisoned, and prevent further claims + pub(super) fn set_poisoned(&mut self) { + self.raw = Qty::MAX; + } +} + +impl ClaimedQty { + /// Split a `ClaimedQty` into two `ClaimedQty`s + pub(super) fn split_off(&mut self, want: Qty) -> Option<ClaimedQty> { + let new_self = self.raw.checked_sub(*want)?; + // commit + *self.raw = new_self; + Some(ClaimedQty::from_raw(want)) + } + + /// Merge two `ClaimedQty`s + /// + /// (Handles overflow by saturating; returning an error is not going to be useful.) + pub(super) fn merge_into(&mut self, have: ClaimedQty) { + let have = have.into_raw(); + *self.raw = self.raw.saturating_add(*have); + } + + /// Obtain result for the participant, after having successfully recorded the amount claimed + /// + /// # CORRECTNESS + /// + /// This must be called only on a successful return path from [`Participation::claim`]. + #[allow(clippy::unnecessary_wraps)] // returns Result; proves it's used on success path + pub(super) fn claim_return_to_participant(self) -> crate::Result<()> { + let _: Qty = self.into_raw(); + Ok(()) + } + + /// When the participant indicates a release, enrol the amount in our bookkeping scheme + /// + /// Handles the quantity argument to [`Participation::release`]. + /// + /// # CORRECTNESS + /// + /// This must be called only on entry to [`Participation::release`]. + pub(super) fn release_got_from_participant(got: Qty) -> Self { + ClaimedQty::from_raw(got) + } + + /// Dispose of a quantity that was claimed by a now-destroyed participant + /// + /// # CORRECTNESS + /// + /// The `ParticipQty` this was claimed from must also have been destroyed. + /// + /// So, + /// [`ParticipQty::for_participant_teardown`] and the corresponding + /// [`release`](TotalQty::release) + /// must have been called earlier - possibly, much earlier. + pub(super) fn dispose_participant_destroyed(mut self) { + let _: Qty = mem::take(&mut self).into_raw(); + } +} + +impl ParticipQty { + /// Prepare to destroy the `ParticipQty` in a participant that's being destroyed + /// + /// When the records of a participant that is being torn down are being destroyed, + /// we must remove our records of the memory that it allocated. + /// + /// This function is for that situation. + /// The returned `ClaimedQty` should then be passed to `release`. + /// + /// # CORRECTNESS + /// + /// The data structure where this `ParticipQty` resides + /// must be torn down (after we return). + /// + /// The `ClaimedQty` must be passed to [`TotalQty::release`], + /// passing the same `p_used`. + // + // We could provide this as a single transaction function, rather than requiring + // two calls. But the main code doesn't have a `TotalQty`, only a `TotalQtyNotifier`, + // so we'd need to add an additional passthrough method to `TotalQtyNotifier`, + // which doesn't seem worth it given that there's only one call site for this fn. + pub(super) fn for_participant_teardown(&self) -> ClaimedQty { + // We imagine that the Participant said it was releasing everything + ClaimedQty::from_raw(self.as_raw()) + } +} diff --git a/crates/tor-memquota/src/mtracker/reclaim.rs b/crates/tor-memquota/src/mtracker/reclaim.rs new file mode 100644 index 000000000..6e45b1f3c --- /dev/null +++ b/crates/tor-memquota/src/mtracker/reclaim.rs @@ -0,0 +1,396 @@ +//! Reclamation algorithm +//! +//! Implementation the of long-running [`task`] function, +//! (which is the only export from here, the wider `mtracker` module). + +use super::*; + +mod deferred_drop; + +use deferred_drop::{DeferredDrop, GuardWithDeferredDrop}; + +/// Total number of participants +/// +/// Used in reporting and in calculations of various edge cases. +/// On 64-bit systems, bigger than the refcounts, which are all `u32` +type NumParticips = usize; + +//========== candiate victim analysis ========== + +/// The nominal data age of a participant +#[derive(Ord, PartialOrd, Eq, PartialEq)] +enum Age { + /// Treat this participant as having very old data + TreatAsVeryOld, + /// Data age value from the [`IsParticipant`] + Actual(CoarseInstant), +} + +/// Participant status, as a candidate victim +enum PStatus { + /// Treat participant as having data of age OldestData + Candidate(Age), + /// Tear this participant down right away + TearDown, + /// Treat participant as not having any data; don't reclaim from it + NoData, +} + +/// Outcome of a completed reclamation run +/// +/// This is used only within `choose_victim`, and only for logging +#[derive(Debug, derive_more::Display)] +enum Outcome { + /// We reached the low water mark + #[display(fmt = "complete")] + TargetReached, + + /// We didn't, but we have so many participants that that's possibly expected + /// + /// (Can only happen on 32-bit platforms.) + #[display(fmt = "{} participants, good enough - stopping", n_particips)] + GoodEnough { + /// The number of participants + n_particips: NumParticips, + }, +} + +/// Figure out whether a participant is a candidate victim, and obtain its data age +fn analyse_particip(precord: &PRecord, defer_drop: &mut DeferredDrop) -> PStatus { + let Some(particip) = precord.particip.upgrade() else { + // Oh! This participant has vanished! + // We can't reclaim from it. It may already be reclaiming. + // Delete it from our data structure. + return PStatus::TearDown; + }; + + let got_oldest = catch_unwind(AssertUnwindSafe(|| particip.get_oldest())); + defer_drop.push(particip); + + match got_oldest { + Ok(Some(age)) => return PStatus::Candidate(Age::Actual(age)), + Ok(None) => {} + Err(_panicked) => { + // _panicked is of a useless type + error!("memory tracker: call to get_oldest panicked!"); + return PStatus::TearDown; + } + } + + // The participant claims not to have any memory + // There might be some cached, let's check + + let Some(max_cached) = precord + .refcount + .as_usize() + .checked_mul(MAX_CACHE.as_usize()) + else { + // WTF! So many Participation clones that the max usage has + // overflowed. (This can only happen on 32-bit platforms + // since refcount is a u32.) Probably we should reclaim + // from this participant. + log_ratelim!( + "memtrack: participant with many clones claims to have no data"; + Err::<Void, _>(internal!("{} Participation clones", *precord.refcount)); + ); + return PStatus::Candidate(Age::TreatAsVeryOld); + }; + + if precord.used.as_raw() > Qty(max_cached) { + // This participant is lying to us somehow. + log_ratelim!( + "memtrack: participant claims to have no data, but our accounting disagrees"; + Err::<Void, _>(internal!("{} used (by {} clones)", precord.used, *precord.refcount)); + ); + return PStatus::Candidate(Age::TreatAsVeryOld); + } + + // Participant plausibly does have no data + PStatus::NoData +} + +//========== reclamation algorith, the main pieces ========== + +/// State while reclamation is active +struct Reclaiming { + /// The heap of candidates, oldest at top of heap + heap: BinaryHeap<Reverse<(Age, AId)>>, +} + +/// A victim we have selected for reclamation +/// +/// This designates a specific Participant. +/// +/// But, note that we always reclaim from an Account, so if we are reclaiming +/// from one `Victim`, we may be reclaiming from other `Victim`s with the same +/// `AId` and different `IsParticipant`s. And because of inheritance, we might +/// be reclaiming from other Accounts too. +type Victim = (AId, drop_reentrancy::ProtectedArc<dyn IsParticipant>); + +/// Marker indicating that the victim's reclaim function panicked +struct VictimPanicked; + +/// Set of responses from the victims, after they have all finished reclaiming. +type VictimResponses = Vec<(AId, Result<Reclaimed, VictimPanicked>)>; + +impl Reclaiming { + /// Check to see if we should start reclaiming, and if so return a `Reclaiming` + /// + /// 1. Checks to see if usage is above `max`; if not, returns `None` + /// 2. Logs that we're starting reclamation + /// 3. Calculates the heap of data ages + fn maybe_start(state: &mut GuardWithDeferredDrop) -> Option<Self> { + let (state, deferred_drop) = state.deref_mut_both(); + + if *state.total_used <= state.global.config.max { + return None; + } + + info!( + "memory tracking: {} > {}, reclamation started (target {})", + *state.total_used, state.config.max, state.config.low_water, + ); + + // `BinaryHeap` is a max heap, so use Rev + let mut heap = BinaryHeap::new(); + + // Build heap of participants we might want to reclaim from + // (and, while we're at it, tear down broken participants) + for (aid, arecord) in &mut state.accounts { + arecord.ps.retain(|_pid, precord| { + match analyse_particip(precord, deferred_drop) { + PStatus::Candidate(age) => { + heap.push(Reverse((age, aid))); + true // retain + } + PStatus::NoData => { + true // retain + } + PStatus::TearDown => { + precord.auto_release(&mut state.global); + false // remove + } + } + }); + } + + Some(Reclaiming { heap }) + } + + /// If we're reclaiming, choose the next victim(s) to reclaim + /// + /// This is the account whose participant has the oldest data age, + /// and all of that account's children. + /// + /// We might discover that we didn't want to continue reclamation after all: + /// this function is responsible for checking our progress + /// against the low water mark. + /// + /// If reclamation should stop, this function logs, and returns `None`. + fn choose_victims(&mut self, state: &mut State) -> Result<Option<Vec<Victim>>, ReclaimCrashed> { + let stop = |state: &mut State, outcome| { + info!( + "memory tracking reclamation reached: {} (target {}): {}", + *state.total_used, state.config.low_water, outcome, + ); + Ok(None) + }; + + if *state.total_used <= state.config.low_water { + return stop(state, Outcome::TargetReached); + } + let Some(Reverse((_, oldest_aid))) = self.heap.pop() else { + // All our remaining participants are NoData. + let n_particips: usize = state + .accounts + .values() + .map(|ar| { + ar.ps + .values() + .map( + |pr| *pr.refcount as NumParticips, // refcount is u32, so this is fine + ) + .sum::<NumParticips>() + }) + .sum::<NumParticips>(); + + if state.total_used.as_raw().as_usize() / n_particips < usize::from(MAX_CACHE) { + // On 32-bit, this could happen due to the cache, if we have + // 2^32 / MAX_CACHE participants. + return stop(state, Outcome::GoodEnough { n_particips }); + } + + // Oh dear. + return Err(internal!( + "memory accounting state corrupted: used={} n_particips={} all NoData", + *state.total_used, + n_particips, + ) + .into()); + }; + + // When we do partial reclamation, rather than just Collapsing: + // + // fudge next_oldest by something to do with number of loop iterations, + // to avoid one-allocation-each-time ping pong between multiple caches + // + // (this match statement will fail to compile when we add a non-Collapsing variant) + // + // let next_oldest = heap.peek_lowest(); + match None { + None | Some(Reclaimed::Collapsing) => {} + } + + let victim_aids = state.get_aid_and_children_recursively(oldest_aid); + + let victims: Vec<Victim> = { + let mut particips = vec![]; + for aid in victim_aids { + let Some(arecord) = state.accounts.get_mut(aid) else { + // shouldn't happen but no need to panic + continue; + }; + arecord.ps.retain(|_pid, precord| { + let Some(particip) = precord.particip.upgrade() else { + // tear this down! + precord.auto_release(&mut state.global); + return false; + }; + particips.push((aid, particip)); + true + }); + } + particips + }; + + Ok(Some(victims)) + } + + /// Notify the chosen victims and obtain their responses + /// + /// This is the async part, and is done with the state unlocked. + // Doesn't actually need `self`, only `victims`, but we take it for form's sake + async fn notify_victims(&mut self, victims: Vec<Victim>) -> VictimResponses { + futures::future::join_all( + // + victims.into_iter().map(|(aid, particip)| async move { + let particip = particip.promise_dropping_is_ok(); + // We run the `.reclaim()` calls within the same task (since that's what + // `join_all` does). So they all run on whatever executor thread is polling + // the reclamation task. + let reclaimed = AssertUnwindSafe(particip.reclaim()) + .catch_unwind() + .await + .map_err(|_panicked| VictimPanicked); + // We drop the `ProtectedArc<dyn IsParticipant>` here, which is OK + // because we don't hold the lock. Since drop isn't async, and + // `join_all` doesn't spawn tasks, we drop them sequentially. + (aid, reclaimed) + }), + ) + .await + } + + /// Process the victim's responses and update `state` accordingly + // Doesn't actually need `self`, only `state`, but we take it for form's sake + fn handle_victim_responses(&mut self, state: &mut State, responses: VictimResponses) { + for (aid, reclaimed) in responses { + match reclaimed { + Ok(Reclaimed::Collapsing) | Err(VictimPanicked) => { + let Some(mut arecord) = state.accounts.remove(aid) else { + // Account is gone, fair enough + continue; + }; + arecord.auto_release(&mut state.global); + // Account is definitely gone now + } + } + } + } +} + +//========== the reclamation task, in terms of the pieces ==========- + +/// Return value from the task, when it finishes due to the tracker being shut down +struct TaskFinished; + +/// Reclaim memory until we reach low water, if necessary +/// +/// Looks to see if we're above `config.max`. +/// If so, constructs a list of victims, and starts reclaiming from them, +/// until we reach low water. +async fn inner_loop(tracker: &Arc<MemoryQuotaTracker>) -> Result<(), ReclaimCrashed> { + let mut reclaiming; + let mut victims; + { + let mut state_guard = GuardWithDeferredDrop::new(tracker.lock()?); + + let Some(r) = Reclaiming::maybe_start(&mut state_guard) else { + return Ok(()); + }; + reclaiming = r; + + // Duplicating this call to reclaiming.choose_victims means we don't + // release the lock between `maybe_start` and `choose_victims` (here) + // and between `handle_victim_responses` and `choose_victims` (bellw). + // (Releasing the lock would not be a bug, but it's not desirable.) + let Some(v) = reclaiming.choose_victims(&mut state_guard)? else { + return Ok(()); + }; + victims = v; + } + + loop { + let responses = reclaiming.notify_victims(mem::take(&mut victims)).await; + let mut state_guard = tracker.lock()?; + reclaiming.handle_victim_responses(&mut state_guard, responses); + let Some(v) = reclaiming.choose_victims(&mut state_guard)? else { + return Ok(()); + }; + victims = v; + } +} + +/// Internal long-running task, handling reclamation - main loop +/// +/// Handles routine logging, but not termination +async fn task_loop( + tracker: &Weak<MemoryQuotaTracker>, + mut wakeup: mpsc::Receiver<()>, +) -> Result<TaskFinished, ReclaimCrashed> { + loop { + // We don't hold a strong reference while we loop around, so we detect + // last drop of an actual client handle. + { + let Some(tracker) = tracker.upgrade() else { + return Ok(TaskFinished); + }; + + inner_loop(&tracker).await?; + } + + let Some(()) = wakeup.next().await else { + // Sender dropped + return Ok(TaskFinished); + }; + } +} + +/// Internal long-running task, handling reclamation +/// +/// This is the entrypoint used by the rest of the `tracker`. +/// It handles logging of crashes. +pub(super) async fn task(tracker: Weak<MemoryQuotaTracker>, wakeup: mpsc::Receiver<()>) { + match task_loop(&tracker, wakeup).await { + Ok(TaskFinished) => {} + Err(bug) => { + let _: Option<()> = (|| { + let tracker = tracker.upgrade()?; + let mut state = tracker.state.lock().ok()?; + state.total_used.set_poisoned(); + Some(()) + })(); + error_report!(bug, "memory tracker task failed"); + } + } +} diff --git a/crates/tor-memquota/src/mtracker/reclaim/deferred_drop.rs b/crates/tor-memquota/src/mtracker/reclaim/deferred_drop.rs new file mode 100644 index 000000000..3b0742a82 --- /dev/null +++ b/crates/tor-memquota/src/mtracker/reclaim/deferred_drop.rs @@ -0,0 +1,81 @@ +//! Deferred drop handling. +//! +//! We sometimes have `Arc<dyn Participant>`s we have obtained but don't want to drop yet +//! +//! See the top-level docs for context. +//! +//! When we drop the `Arc`, the refcount might become zero. +//! Then the inner type would be dropped. +//! The inner type is allowed to call back into us (for example, it may drop an `Account`). +//! We must therefore not drop a caller's `Participant` with our own state lock held. +//! +//! This module has a helper type for assuring that we do defer drops. +// +// There are no separate tests for this module. Drop bombs are hard to test for, +// and the rest of the code is just wrappers. + +use super::*; + +/// `MutexGuard<State>` but also a list of `Arc<dyn Partcipant>` to drop when we unlock +#[derive(Debug, Default)] +pub(super) struct GuardWithDeferredDrop<'m> { + /// The mutex guard + /// + /// Always `Some`; just an `Option` so we can move out during drop + guard: Option<MutexGuard<'m, State>>, + + /// The participants we've acquired and which we want to drop later + deferred_drop: DeferredDrop, +} + +/// Participants we've acquired and which we want to drop later, convenience alias +pub(super) type DeferredDrop = Vec<drop_reentrancy::ProtectedArc<dyn IsParticipant>>; + +impl<'m> GuardWithDeferredDrop<'m> { + /// Prepare for handling deferred drops + pub(super) fn new(guard: MutexGuard<'m, State>) -> Self { + GuardWithDeferredDrop { + guard: Some(guard), + deferred_drop: vec![], + } + } + + /// Obtain mutable borrows of the two components + pub(super) fn deref_mut_both(&mut self) -> (&mut State, &mut DeferredDrop) { + ( + self.guard.as_mut().expect("deref_mut after drop"), + &mut self.deferred_drop, + ) + } +} + +impl Deref for GuardWithDeferredDrop<'_> { + type Target = State; + fn deref(&self) -> &State { + self.guard.as_ref().expect("deref after drop") + } +} +impl DerefMut for GuardWithDeferredDrop<'_> { + fn deref_mut(&mut self) -> &mut State { + self.deref_mut_both().0 + } +} + +// We use ProtectedArc. In tests, that has a drop bomb which requires us to +// call `.promise_dropping_is_ok()`, on pain of panicking. So we must do that here. +// +// Outside tests, the normal drop order would be precisely correct: +// the guard field comes first, so the compiler would drop it before the Arcs. +// So we could make this `#[cfg(test)]` (and add some comments above about field order). +// However, we prefer to use the same code, so that the correctness of +// *production* GuardWithDeferredDrop is assured by the `ProtectedArc`. +impl Drop for GuardWithDeferredDrop<'_> { + fn drop(&mut self) { + let guard = self.guard.take().expect("dropping twice!"); + drop::<MutexGuard<_>>(guard); + // we just unlocked the guard, so drops that re-enter our code are fine + for p in self.deferred_drop.drain(..) { + p.promise_dropping_is_ok(); + } + } +} diff --git a/crates/tor-memquota/src/mtracker/test.rs b/crates/tor-memquota/src/mtracker/test.rs new file mode 100644 index 000000000..43d3dcd7c --- /dev/null +++ b/crates/tor-memquota/src/mtracker/test.rs @@ -0,0 +1,678 @@ +//! `tor_memtrack::tracker::test` + +// @@ begin test lint list maintained by maint/add_warning @@ +#![allow(clippy::bool_assert_comparison)] +#![allow(clippy::clone_on_copy)] +#![allow(clippy::dbg_macro)] +#![allow(clippy::mixed_attributes_style)] +#![allow(clippy::print_stderr)] +#![allow(clippy::print_stdout)] +#![allow(clippy::single_char_pattern)] +#![allow(clippy::unwrap_used)] +#![allow(clippy::unchecked_duration_subtraction)] +#![allow(clippy::useless_vec)] +#![allow(clippy::needless_pass_by_value)] +//! <!-- @@ end test lint list maintained by maint/add_warning @@ --> +#![allow(clippy::let_and_return)] // TODO this lint is annoying and we should disable it + +use super::*; + +use std::collections::BTreeMap; +use std::fmt::Write as _; +use std::time::Duration; + +use itertools::Itertools; +use rand::Rng; +use slotmap::Key as _; +use tracing_test::traced_test; + +use tor_basic_utils::RngExt as _; +use tor_rtcompat::{CoarseDuration, CoarseTimeProvider as _, Runtime}; +use tor_rtmock::MockRuntime; + +//---------- useful utilities ---------- + +fn secs(s: u64) -> CoarseDuration { + Duration::from_secs(s).into() +} + +fn mby(mib: usize) -> usize { + mib * 1024 * 1024 +} + +fn mk_config() -> Config { + Config::builder() + .max(mby(20)) + .low_water(mby(15)) + .build() + .unwrap() +} + +fn mk_tracker(rt: &impl Runtime) -> Arc<MemoryQuotaTracker> { + MemoryQuotaTracker::new(&rt, mk_config()).unwrap() +} + +fn test_with_various_mocks<F, Fut>(f: F) +where + F: Fn(tor_rtmock::MockRuntime) -> Fut, + Fut: Future<Output = ()>, +{ + MockRuntime::test_with_various(|rt| async { + // Make sure we can talk about times at least 1000s in the past + // TODO maybe this should be a feature of MockRuntime but what value to pick? + rt.advance_by(Duration::from_secs(1000)).await; + f(rt).await; + }); +} + +//---------- consistency check (test invariants against outside view) ---------- + +use consistency::*; +mod consistency { + use super::*; + + #[derive(Default)] + pub(super) struct CallerInfoCollector { + g: usize, + acs: BTreeMap<AId, refcount::RawCount>, + pcs: BTreeMap<(AId, PId), (refcount::RawCount, usize)>, + debug_dump: String, + } + + pub(super) trait HasCallerInfo { + fn note_consistency_caller_info(&self, collector: &mut CallerInfoCollector); + } + + impl CallerInfoCollector { + pub(super) fn note_account(&mut self, acct: &Account, reclaimed: ReclaimedOrOk) { + writeln!(self.debug_dump, "acct {acct:?} {reclaimed:?}").unwrap(); + if acct.aid.is_null() || reclaimed.is_err() { + return; + } + let ac = self.acs.entry(*acct.aid).or_default(); + *ac += 1; + } + pub(super) fn note_particip( + &mut self, + p: &Participation, + reclaimed: ReclaimedOrOk, + used: usize, + ) { + writeln!(self.debug_dump, "particip {p:?} {reclaimed:?} {used:?}").unwrap(); + if p.pid.is_null() || p.aid.is_null() || reclaimed.is_err() { + return; + } + self.note_partn_core(p, used); + } + pub(super) fn note_partn_clone(&mut self, p: &Participation) { + writeln!(self.debug_dump, "partn {p:?}").unwrap(); + if p.pid.is_null() { + return; + } + self.note_partn_core(p, 0); + } + fn note_partn_core(&mut self, p: &Participation, x_used: usize) { + let pc = self.pcs.entry((p.aid, *p.pid)).or_default(); + let used = *p.cache.as_raw() + x_used; + pc.0 += 1; + pc.1 += used; + self.g += used; + } + } + + pub(super) fn check_consistency_general( + trk: &Arc<MemoryQuotaTracker>, + collect_caller_info: impl FnOnce(&mut CallerInfoCollector), + ) { + let state = trk.lock().unwrap(); + + let (expected, debug_dump) = { + let mut c = CallerInfoCollector::default(); + collect_caller_info(&mut c); + ((c.g, c.acs, c.pcs), c.debug_dump) + }; + + let got = { + let mut gc = 0; + let mut acs = BTreeMap::new(); + let mut pcs = BTreeMap::new(); + for (aid, arecord) in &state.accounts { + acs.insert(aid, *arecord.refcount); + for (pid, precord) in &arecord.ps { + let used = *precord.used.as_raw(); + gc += used; + pcs.insert((aid, pid), (*precord.refcount, used)); + } + } + + (gc, acs, pcs) + }; + + assert_eq!( + expected, got, + "\n----- dump (start) -----\n{debug_dump}----- dump (end) -----", + ); + } +} + +//---------- common test participant (state) ---------- + +#[derive(Debug)] +struct PartnState { + partn: Participation, + age: Option<CoarseInstant>, + used: usize, + reclaimed: ReclaimedOrOk, + show: String, +} + +#[derive(Debug)] +struct TestPartn { + state: Mutex<PartnState>, +} + +impl TestPartn { + fn lock(&self) -> MutexGuard<PartnState> { + self.state.lock().unwrap() + } +} + +impl From<PartnState> for TestPartn { + fn from(state: PartnState) -> TestPartn { + TestPartn { + state: Mutex::new(state), + } + } +} + +impl TestPartn { + fn get_oldest(&self) -> Option<CoarseInstant> { + self.lock().age + } + fn reclaim(&self) -> ReclaimFuture { + let () = mem::replace(&mut self.lock().reclaimed, Err(())).expect("reclaimed twice!"); + Box::pin(async { Reclaimed::Collapsing }) + } + fn is_reclaimed(&self) -> Result<(), ()> { + self.lock().reclaimed + } +} + +impl IsParticipant for TestPartn { + fn get_oldest(&self) -> Option<CoarseInstant> { + self.get_oldest() + } + fn reclaim(self: Arc<Self>) -> ReclaimFuture { + (*self.clone()).reclaim() + } +} + +impl PartnState { + fn claim(&mut self, qty: usize) -> Result<(), crate::Error> { + claim_via(&mut self.partn, &self.show, &mut self.used, qty) + } + + fn release(&mut self, qty: usize) { + release_via(&mut self.partn, &self.show, &mut self.used, qty); + } +} + +fn claim_via( + via: &mut Participation, + show: impl Display, + used: &mut usize, + qty: usize, +) -> Result<(), crate::Error> { + eprintln!("{show} claim {qty} {qty:#x}"); + via.claim(qty)?; + *used += qty; + Ok(()) +} + +fn release_via(via: &mut Participation, show: impl Display, used: &mut usize, qty: usize) { + eprintln!("{show} release {qty} {qty:#x}"); + via.release(qty); + *used -= qty; +} + +impl HasCallerInfo for PartnState { + fn note_consistency_caller_info(&self, collector: &mut CallerInfoCollector) { + collector.note_particip(&self.partn, self.reclaimed, self.used); + } +} + +//---------- test participant which is directly the accountholder ---------- + +#[derive(Debug, Deref)] +struct UnifiedP { + acct: Account, + #[deref] + state: TestPartn, +} + +type ReclaimedOrOk = Result<(), ()>; + +impl IsParticipant for UnifiedP { + fn get_oldest(&self) -> Option<CoarseInstant> { + self.state.get_oldest() + } + fn reclaim(self: Arc<Self>) -> ReclaimFuture { + self.state.reclaim() + } +} + +impl UnifiedP { + fn new( + rt: &impl Runtime, + trk: &Arc<MemoryQuotaTracker>, + parent: Option<&Account>, + age: CoarseDuration, + show: impl Display, + ) -> Arc<Self> { + let acct = trk.new_account(parent).unwrap(); + + let now = rt.now_coarse(); + + acct.register_participant_with(now, |partn| { + Ok::<_, Void>(Arc::new(UnifiedP { + acct: acct.clone(), + state: PartnState { + partn, + age: Some(now - age), + show: show.to_string(), + used: 0, + reclaimed: Ok(()), + } + .into(), + })) + }) + .unwrap() + .void_unwrap() + } + + async fn settle_check_consistency<'i>( + rt: &'i MockRuntime, + trk: &'i Arc<MemoryQuotaTracker>, + ups: impl IntoIterator<Item = &'i Arc<Self>> + 'i, + ) { + rt.advance_until_stalled().await; + + check_consistency_general(trk, |collector| { + for up in ups { + up.note_consistency_caller_info(collector); + } + }); + } +} + +impl HasCallerInfo for UnifiedP { + fn note_consistency_caller_info(&self, collector: &mut CallerInfoCollector) { + let state = self.lock(); + collector.note_account(&self.acct, state.reclaimed); + state.note_consistency_caller_info(collector); + } +} + +//---------- test cases with unified accountholder/participant ---------- + +#[traced_test] +#[test] +fn basic() { + test_with_various_mocks(|rt| async move { + let trk = mk_tracker(&rt); + + let ps: Vec<Arc<UnifiedP>> = (0..21) + .map(|i| UnifiedP::new(&rt, &trk, None, secs(i), i)) + .collect(); + + for p in &ps[0..19] { + p.lock().claim(mby(1)).unwrap(); + UnifiedP::settle_check_consistency(&rt, &trk, &ps).await; + } + + let count_uncollapsed = || ps.iter().filter(|p| p.is_reclaimed().is_ok()).count(); + + assert_eq!(count_uncollapsed(), 21); + + for p in &ps[20..] { + // check that we are exercising a situation with nonzero cached + // (this is set up by register_participant + assert_ne!(p.lock().partn.cache, Qty(0)); + + p.lock() + .claim(mby(1)) + .expect("allocation rejected, during collapse, but collapse is async"); + } + + UnifiedP::settle_check_consistency(&rt, &trk, &ps).await; + + assert_eq!(count_uncollapsed(), 14); + + // Now we drop everything. This exercises much of the teardown! + }); +} + +#[traced_test] +#[test] +fn parent() { + test_with_various_mocks(|rt| async move { + for ages in [[10, 20], [20, 10]] { + eprintln!("ages: {ages:?}"); + let [parent_age, child_age] = ages.map(secs); + + let trk = mk_tracker(&rt); + + let mk_p = |parent, age, show| UnifiedP::new(&rt, &trk, parent, age, show); + + let parent = mk_p(None, parent_age, "parent"); + parent.lock().claim(mby(7)).unwrap(); + rt.advance_until_stalled().await; + assert!(parent.is_reclaimed().is_ok()); + + let child = mk_p(Some(&parent.acct), child_age, "child"); + child.lock().claim(mby(7)).unwrap(); + assert!(parent.is_reclaimed().is_ok()); + assert!(child.is_reclaimed().is_ok()); + + let trigger = mk_p(None, secs(0), "trigger"); + trigger.lock().claim(mby(7)).unwrap(); + assert!(trigger.is_reclaimed().is_ok()); + + rt.advance_until_stalled().await; + + if parent_age > child_age { + // parent is older than child, we're supposed to have reclaimed + // from the parent, causing reclamation of the child. + assert!(parent.is_reclaimed().is_err()); + assert!(child.is_reclaimed().is_err()); + } else { + // supposed to have reclaimed from child only + assert!(parent.is_reclaimed().is_ok()); + assert!(child.is_reclaimed().is_err()); + } + } + }); +} + +#[traced_test] +#[test] +fn cache() { + test_with_various_mocks(|rt| async move { + let seq = [ + 1, + 1000, + *MAX_CACHE - 2000, + 3000, + *MAX_CACHE, + *MAX_CACHE - 1, + *MAX_CACHE + 1, + ]; + + let trk = mk_tracker(&rt); + let p = UnifiedP::new(&rt, &trk, None, secs(0), "p"); + + for qty in seq { + p.lock().claim(qty).unwrap(); + UnifiedP::settle_check_consistency(&rt, &trk, [&p]).await; + } + + for qty in seq { + p.lock().release(qty); + UnifiedP::settle_check_consistency(&rt, &trk, [&p]).await; + } + + let mut p2 = p.lock().partn.clone(); + + let mut rng = tor_basic_utils::test_rng::Config::Deterministic.into_rng(); + for _iter in 0..10_000 { + let qty = rng.gen_range_checked(0..=*MAX_CACHE).unwrap(); + let p_use_i = rng.gen_range_checked(1..=3).unwrap(); + { + let mut state = p.lock(); + let state = &mut *state; + + let mut p_use_buf; + let p_use = match p_use_i { + 1 => &mut state.partn, + 2 => &mut p2, + 3 => { + p_use_buf = p2.clone(); + &mut p_use_buf + } + x => panic!("{}", x), + }; + + if rng.gen() || qty > state.used { + claim_via(p_use, p_use_i, &mut state.used, qty).unwrap(); + } else { + release_via(p_use, p_use_i, &mut state.used, qty); + } + } + + rt.advance_until_stalled().await; + check_consistency_general(&trk, |collector| { + p.note_consistency_caller_info(collector); + collector.note_partn_clone(&p2); + }); + } + }); +} + +//---------- test client with multiple participants per account ---------- + +#[derive(Debug)] +struct ComplexAH { + acct: Account, + ps: Vec<Arc<TestPartn>>, +} + +impl HasCallerInfo for ComplexAH { + fn note_consistency_caller_info(&self, collector: &mut CallerInfoCollector) { + let reclaimed = self + .ps + .iter() + .map(|p| p.lock().reclaimed) + .dedup() + .exactly_one() + .unwrap(); + + collector.note_account(&self.acct, reclaimed); + for p in &self.ps { + p.lock().note_consistency_caller_info(collector); + } + } +} + +impl ComplexAH { + fn new(trk: &Arc<MemoryQuotaTracker>) -> Self { + ComplexAH { + acct: trk.new_account(None).unwrap(), + ps: vec![], + } + } + + fn add_p(&mut self, now: CoarseInstant, age: CoarseDuration, show: impl Display) -> usize { + let cp = self + .acct + .register_participant_with(now, |partn| { + Ok::<_, Void>(Arc::new(TestPartn::from(PartnState { + partn, + age: Some(now - age), + show: show.to_string(), + used: 0, + reclaimed: Ok(()), + }))) + }) + .unwrap() + .void_unwrap(); + + let i = self.ps.len(); + self.ps.push(cp); + i + } +} + +#[traced_test] +#[test] +fn complex() { + test_with_various_mocks(|rt| async move { + let trk = mk_tracker(&rt); + + let up = UnifiedP::new(&rt, &trk, None, secs(0), "U"); + let mut ah = ComplexAH::new(&trk); + let now = rt.now_coarse(); + + for age in [5, 9] { + ah.add_p(now, secs(age), age); + } + + let settle_check_consistency = || async { + rt.advance_until_stalled().await; + + check_consistency_general(&trk, |collector| { + up.note_consistency_caller_info(collector); + ah.note_consistency_caller_info(collector); + }); + }; + + up.lock().claim(mby(1)).unwrap(); + ah.ps[0].lock().claim(mby(11)).unwrap(); + + settle_check_consistency().await; + + assert!(up.is_reclaimed().is_ok()); + for p in &ah.ps { + assert!(p.is_reclaimed().is_ok()); + } + + ah.ps[1].lock().claim(mby(11)).unwrap(); + + settle_check_consistency().await; + assert!(up.is_reclaimed().is_ok()); + for p in &ah.ps { + assert!(p.is_reclaimed().is_err()); + } + }); +} + +//---------- various error cases ---------- + +#[derive(Debug)] +struct DummyParticipant; + +impl IsParticipant for DummyParticipant { + fn get_oldest(&self) -> Option<CoarseInstant> { + None + } + fn reclaim(self: Arc<Self>) -> ReclaimFuture { + Box::pin(async { Reclaimed::Collapsing }) + } +} + +#[traced_test] +#[test] +fn errors() { + test_with_various_mocks(|rt| async move { + let trk = mk_tracker(&rt); + let now = rt.now_coarse(); + + let mk_ah = || { + let mut ah = ComplexAH::new(&trk); + ah.add_p(now, secs(5), "p"); + ah + }; + + const CLAIM: usize = MAX_CACHE.as_usize() + 1; + + let dummy_dangling = || { + let p = Arc::new(DummyParticipant); + Arc::downgrade(&p) + // p dropped here + }; + assert!(dummy_dangling().upgrade().is_none()); + + macro_rules! assert_error { { $error:ident, $r:expr } => { + let r = $r; + assert!(matches!(r, Err(Error::$error)), "unexpected: {:?} => {:?}", stringify!($r), &r); + } } + + // Dropped account + { + let mut ah = mk_ah(); + let wa1: WeakAccount = ah.acct.downgrade(); + let p = ah.ps.pop().unwrap(); + let wa2: WeakAccount = p.lock().partn.account(); + drop(ah.acct); + + rt.advance_until_stalled().await; + check_consistency_general(&trk, |_collector| ()); + + // account should be dead now + assert!(p.lock().claim(1).is_ok()); // from cache! + assert_error!(AccountClosed, p.lock().claim(CLAIM)); + assert_error!(AccountClosed, wa1.upgrade()); + assert_error!(AccountClosed, wa2.upgrade()); + + // but we can still release + p.lock().release(1); + } + + // Dropped IsParticipant + { + let mut ah = mk_ah(); + let p = ah.ps.pop().unwrap(); + let mut state = Arc::into_inner(p).unwrap().state.into_inner().unwrap(); + + state.claim(mby(30)).unwrap(); // will trigger reclaim, which discovers the loss + + rt.advance_until_stalled().await; + check_consistency_general(&trk, |collector| { + let reclaimed = Ok(()); // didn't manage to make the callback! + collector.note_account(&ah.acct, reclaimed); + }); + + assert_error!(ParticipantShutdown, state.claim(CLAIM)); + } + + // Reclaimed account + { + let ah = mk_ah(); + ah.ps[0].lock().claim(mby(30)).unwrap(); + + rt.advance_until_stalled().await; + check_consistency_general(&trk, |_collector| ()); + + let p = &ah.ps[0]; + + assert!(p.lock().reclaimed.is_err()); + assert_error!(AccountClosed, p.lock().claim(CLAIM)); + + let cloned = ah.acct.clone(); + assert!(cloned.aid.is_null()); + assert_error!( + AccountClosed, + ah.acct.register_participant(dummy_dangling()) + ); + + let mut cloned = p.lock().partn.clone(); + assert!(cloned.pid.is_null()); + assert_error!(AccountClosed, cloned.claim(CLAIM)); + + // but we can still release + p.lock().release(1); + } + + // Dropped tracker + { + let mut ah = mk_ah(); + let p = ah.ps.pop().unwrap(); + let wa = ah.acct.downgrade(); + drop(ah.acct); + let _: MemoryQuotaTracker = Arc::into_inner(trk).unwrap(); + + assert_error!(TrackerShutdown, wa.upgrade()); + assert_error!(TrackerShutdown, p.lock().partn.account().upgrade()); + assert_error!(TrackerShutdown, p.lock().claim(CLAIM)); + } + }); +} diff --git a/crates/tor-memquota/src/mtracker/total_qty_notifier.rs b/crates/tor-memquota/src/mtracker/total_qty_notifier.rs new file mode 100644 index 000000000..4c877ba9a --- /dev/null +++ b/crates/tor-memquota/src/mtracker/total_qty_notifier.rs @@ -0,0 +1,70 @@ +//! `TotalQtyNotifier` +//! +//! This newtype assures that we wake up the reclamation task when nceessary + +use super::*; + +/// Wrapper for `TotalQty` +#[derive(Deref, Debug)] +pub(super) struct TotalQtyNotifier { + /// Total memory usage + /// + /// Invariant: equal to + /// ```text + /// Σ Σ PRecord.used + /// ARecord PRecord + /// ``` + #[deref] + total_used: TotalQty, + + /// Condvar to wake up the reclamation task + /// + /// The reclamation task has another clone of this + reclamation_task_wakeup: mpsc::Sender<()>, +} + +impl TotalQtyNotifier { + /// Make a new `TotalQtyNotifier`, which will notify a specified condvar + pub(super) fn new_zero(reclamation_task_wakeup: mpsc::Sender<()>) -> Self { + TotalQtyNotifier { + total_used: TotalQty::ZERO, + reclamation_task_wakeup, + } + } + + /// Record that some memory has been (or will be) allocated by a participant + /// + /// Signals the wakeup task if we need to. + pub(super) fn claim( + &mut self, + precord: &mut PRecord, + want: Qty, + config: &Config, + ) -> crate::Result<ClaimedQty> { + let got = self + .total_used + .claim(&mut precord.used, want) + .ok_or_else(|| internal!("integer overflow attempting to add claim {}", want))?; + if self.total_used > config.max { + match self.reclamation_task_wakeup.try_send(()) { + Ok(()) => Ok(()), + Err(e) if e.is_full() => Ok(()), + Err(e) => Err(into_internal!("could not notify reclamation task!")(e)), + }?; + } + Ok(got) + } + + /// Declare this poisoned, and prevent further claims + pub(super) fn set_poisoned(&mut self) { + self.total_used.set_poisoned(); + } + + /// Record that some memory has been (or will be) freed by a participant + pub(super) fn release(&mut self, precord: &mut PRecord, have: ClaimedQty) // infallible + { + // TODO if the participant's usage underflows, tell it to reclaim + // (and log some kind of internal error) + self.total_used.release(&mut precord.used, have); + } +} diff --git a/crates/tor-memquota/src/refcount.rs b/crates/tor-memquota/src/refcount.rs new file mode 100644 index 000000000..cbe17b7f0 --- /dev/null +++ b/crates/tor-memquota/src/refcount.rs @@ -0,0 +1,315 @@ +//! Helpers for reference counting +//! +//! Two main purposes: +//! +//! * Consistent handling of overflow and underflow +//! * Assurance of incrementing/decrementing as appropriate, +//! including in combination with a slotmap containing the referenced data. +//! +//! The caller is responsible for making sure that the *right instance*'s +//! [`Count`] is passed to the methods on [`Ref`]. +// +// There are no separate tests for this module. Many of the tests would want to +// exercise the `Ref`s drop bomb, which is troublesome since it's panic in drop, +// which they're making Rust treat as an abort upstream. +// (This scheme did detect a bug or two during development testing, +// so the drop bomb is known to work.) +// +// Anyway, these functions are mostly newtype veneers over existing functionality. +// They're tested by the MemoryQuotaTracker's tests. + +use crate::internal_prelude::*; + +/// Local alias for the counter type +pub(crate) type RawCount = u32; + +/// Decrement a refcount and maybe remove a corresponding slotmap entry +/// +/// ```rust,ignore +/// fn slotmap_dec_ref!<K, V>( +/// slotmap: &mut SlotMap<K, V>, +/// ref_: Ref<K>, +/// refcount: &mut Count<K>, // (typically) borrows from slotmap +/// ) +/// ``` +// +// This macro is a bit out-of-position, up here, because we want to be able to link +// to it in our rustdocs. +macro_rules! slotmap_dec_ref { { $slotmap:expr, $ref_:expr, $refcount:expr } => { { { + use $crate::refcount::*; + let key: Ref<_> = $ref_; + let refcount: &mut Count<_> = $refcount; + if let Some(Garbage(key)) = key.dispose(refcount) { + let slotmap: &mut SlotMap<_, _> = $slotmap; + let removed = slotmap.remove(key).expect("entry vanished or wrong key passed?!"); + Some(Garbage(removed)) + } else { + None + } +} } } } + +/// A reference count, counting references with id type `K` +#[derive(Default, Educe, Ord, PartialOrd, Eq, PartialEq, Deref)] +#[educe(Debug)] +pub(crate) struct Count<K> { + /// Actual count of references + #[deref] + count: RawCount, + /// Bind to the specific key type + // K is generally Send + Sync + 'static so we don't care about variance etc. + #[educe(Debug(ignore))] + marker: PhantomData<K>, +} + +/// An copy of a [`slotmap::Key`] `K`, which is counted by a `RefCount` +/// +/// Ie, a key of type `K` with the property that it +/// keeps the refcounted data structure alive. +/// +/// Must always be deleted using [`dispose`](Ref::dispose), not dropped. +/// In tests, dropping a `RefCounted` will panic. +/// +/// The `Default` value does *not* contribute to a reference count, +/// and is fine to drop. +#[derive(Deref, Educe)] +#[educe(Debug, Default, Ord, Eq, PartialEq)] +pub(crate) struct Ref<K: slotmap::Key> { + /// Actual key (without generics) + #[deref] + raw_key: K, + /// Bind to the specific key type + #[educe(Debug(ignore))] + marker: PhantomData<K>, + /// Drop bomb + /// + /// Also forces `Ref` not to be Clone + #[educe(Debug(ignore), Ord(ignore), Eq(ignore), PartialEq(ignore))] + #[allow(dead_code)] + bomb: DropBombCondition, +} + +// educe's Ord is open-coded and triggers clippy::non_canonical_partial_ord_impl +impl<K: slotmap::Key> PartialOrd for Ref<K> { + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + Some(self.cmp(other)) + } +} + +// Ideally we'd assert_not_impl on Ref but it has generics +assert_not_impl_any!(DropBombCondition: Clone); + +/// Error: refcount overflowed +#[derive(Debug, Clone, Error, Eq, PartialEq)] +#[error("memory tracking refcount overflowed")] +pub(crate) struct Overflow; + +/// Something which has become garbage +/// +/// Often used within `Option`, for clarity. Examples: +/// +/// * Key whose reference count has reached zero - see [`Ref::dispose`] +/// * Value removed from a SlotMap - see [`slotmap_dec_ref!`] +#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)] +pub(crate) struct Garbage<K>(pub(crate) K); + +impl<K> Count<K> { + /// Make a new refcount with a specified value + const fn new_raw(count: RawCount) -> Self { + Count { + count, + marker: PhantomData, + } + } + + /// Obtain this counter as a `usize` + /// + /// (Reference counts are `u32`, so this might be a conversion.) + pub(crate) fn as_usize(&self) -> usize { + // On a 16-bit platform this could theoretically overflow, + // but there would have to be >2^16 clones, which would be impossible. + let r: u32 = **self; + r as usize + } +} + +/// Increment this refcount, but don't care about any [`Ref`]s +fn inc_raw(c: &mut RawCount) -> Result<(), Overflow> { + *c = c.checked_add(1).ok_or(Overflow)?; + Ok(()) +} + +/// Decrement this refcount, but don't care about any [`Ref`]s +/// +/// Returns [`Some(Garbage(()))`] if the count reached zero +fn dec_raw(c: &mut RawCount) -> Option<Garbage<()>> { + *c = c + .checked_sub(1) + // if this happens, our data structure is corrupted, very bad + .expect("refcount underflow"); + (*c == 0).then_some(Garbage(())) +} + +impl<K: slotmap::Key> Ref<K> { + /// Create a refcounted reference `Ref` from an un-counted key, incrementing the count + pub(crate) fn new(key: K, count: &mut Count<K>) -> Result<Self, Overflow> { + inc_raw(&mut count.count)?; + Ok(Ref::from_raw(key)) + } + + /// Creates a null `Ref`, which doesn't refer to any slot (lookups always give `None`) + pub(crate) fn null() -> Self { + Ref::from_raw(K::null()) + } + + /// Internal function for creating a `Ref` + fn from_raw(raw_key: K) -> Self { + Ref { + raw_key, + marker: PhantomData, + bomb: DropBombCondition::new_armed(), + } + } + + /// Dispose of a refcounted reference `Ref`, decrementing the count + /// + /// If the count reaches zero, the raw key is returned; + /// the caller should remove the corresponding data from the data structure. + pub(crate) fn dispose(mut self, refcount: &mut Count<K>) -> Option<Garbage<K>> { + let was = mem::take(&mut self.raw_key); + assert!(!was.is_null()); + dec_raw(&mut refcount.count).map(|_: Garbage<()>| Garbage(was)) + } + + /// Dispose of a refcounted reference whose container no longer exists + /// + /// # CORRECTNESS + /// + /// This just forgets the reference, without decrementing any refcount. + /// If the container *does* still exist, a ref count ref will be leaked. + pub(crate) fn dispose_container_destroyed(mut self) { + let _: K = mem::take(&mut self.raw_key); + } +} + +impl<K: slotmap::Key> DefaultExtTake for Ref<K> {} + +/// Insert a new entry into a slotmap using refcounted keys +/// +/// `value_maker` should take the provided `Count`, +/// and incorporate it into a new value. +/// +/// On return, the entry will be in the map, and there will be one reference, +/// which is returned. +/// +/// There is no corresponding `slotmap_remove` in this module. +/// Use [`Ref::dispose`] and handle any [`Garbage`] it returns. +pub(crate) fn slotmap_insert<K: slotmap::Key, V>( + slotmap: &mut SlotMap<K, V>, + value_maker: impl FnOnce(Count<K>) -> V, +) -> Ref<K> { + let (ref_, ()) = slotmap_try_insert(slotmap, move |refcount| { + Ok::<_, Void>((value_maker(refcount), ())) + }) + .void_unwrap(); + ref_ +} + +/// Insert a new entry into a slotmap using refcounted keys, fallibly and with extra data +/// +/// Like [`slotmap_insert`] but: +/// * `value_maker` can also return extra return data `RD` to the caller +/// * `value_maker` is allowed to fail. +/// +/// On successful return, the entry will be in the map, and +/// the new `Ref` is returned along with the data `D`. +pub(crate) fn slotmap_try_insert<K: slotmap::Key, V, E, RD>( + slotmap: &mut SlotMap<K, V>, + value_maker: impl FnOnce(Count<K>) -> Result<(V, RD), E>, +) -> Result<(Ref<K>, RD), E> { + let refcount = Count::new_raw(1); + let (value, data) = value_maker(refcount)?; + let raw_key = slotmap.insert(value); + let ref_ = Ref { + raw_key, + marker: PhantomData, + bomb: DropBombCondition::new_armed(), + }; + Ok((ref_, data)) +} + +#[cfg(test)] +impl<K: slotmap::Key> Drop for Ref<K> { + fn drop(&mut self) { + drop_bomb_disarm_assert!(self.bomb, self.raw_key.is_null(),); + } +} + +impl From<Overflow> for Error { + fn from(_overflow: Overflow) -> Error { + internal!("reference count overflow in memory tracking (out-of-control subsystem?)").into() + } +} + +#[cfg(test)] +mod test { + // @@ begin test lint list maintained by maint/add_warning @@ + #![allow(clippy::bool_assert_comparison)] + #![allow(clippy::clone_on_copy)] + #![allow(clippy::dbg_macro)] + #![allow(clippy::mixed_attributes_style)] + #![allow(clippy::print_stderr)] + #![allow(clippy::print_stdout)] + #![allow(clippy::single_char_pattern)] + #![allow(clippy::unwrap_used)] + #![allow(clippy::unchecked_duration_subtraction)] + #![allow(clippy::useless_vec)] + #![allow(clippy::needless_pass_by_value)] + //! <!-- @@ end test lint list maintained by maint/add_warning @@ --> + #![allow(clippy::let_and_return)] // TODO this lint is annoying and we should disable it + + use super::*; + + slotmap::new_key_type! { + struct Id; + } + #[derive(Eq, PartialEq, Debug)] + struct Record { + refcount: Count<Id>, + } + type Map = SlotMap<Id, Record>; + + fn setup() -> (Map, Ref<Id>) { + let mut map = Map::default(); + let ref_ = slotmap_insert(&mut map, |refcount| Record { refcount }); + (map, ref_) + } + + #[test] + fn good() { + let (mut map, ref1) = setup(); + + let ent = map.get_mut(*ref1).unwrap(); + let ref2 = Ref::new(*ref1, &mut ent.refcount).unwrap(); + + let g1: Option<Garbage<Record>> = slotmap_dec_ref!(&mut map, ref1, &mut ent.refcount); + assert_eq!(g1, None); + + let ent = map.get_mut(*ref2).unwrap(); + let g2: Option<Garbage<Record>> = slotmap_dec_ref!(&mut map, ref2, &mut ent.refcount); + assert!(g2.is_some()); + } + + #[test] + fn try_insert_fail() { + let mut map = Map::default(); + let () = slotmap_try_insert::<_, _, _, String>(&mut map, |_refcount| Err(())).unwrap_err(); + } + + #[test] + fn drop_ref_without_decrement() { + let (_map, mut ref1) = setup(); + let h = ref1.bomb.make_simulated(); + drop(ref1); + h.expect_exploded(); + } +} diff --git a/crates/tor-memquota/src/utils.rs b/crates/tor-memquota/src/utils.rs new file mode 100644 index 000000000..4540df03d --- /dev/null +++ b/crates/tor-memquota/src/utils.rs @@ -0,0 +1,69 @@ +//! Miscellanous internal utilities + +use crate::internal_prelude::*; + +/// Quantity of memory used, measured in bytes. +/// +/// Like `usize` but `Display`s in a more friendly and less precise way +#[derive(Debug, Clone, Copy, Hash, Default, Eq, PartialEq, Ord, PartialOrd)] // +#[derive(From, Into, Deref, DerefMut, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct Qty(pub(crate) usize); + +impl Qty { + /// Maximum for the type + pub(crate) const MAX: Qty = Qty(usize::MAX); + + /// Return the value as a plain number, a `usize` + /// + /// Provided so call sites don't need to write an opaque `.0` everywhere, + /// even though that would be fine. + pub(crate) const fn as_usize(self) -> usize { + self.0 + } +} + +impl Display for Qty { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let mb = self.0 as f32 / (1024. * 1024.); + write!(f, "{:.2}MiB", mb) + } +} + +/// Convenience extension trait to provide `.take()` +/// +/// Convenient way to provide `.take()` on some of our types. +pub(crate) trait DefaultExtTake: Default { + /// Returns `*self`, replacing it with the default value. + fn take(&mut self) -> Self { + mem::take(self) + } +} + +#[cfg(test)] +mod test { + // @@ begin test lint list maintained by maint/add_warning @@ + #![allow(clippy::bool_assert_comparison)] + #![allow(clippy::clone_on_copy)] + #![allow(clippy::dbg_macro)] + #![allow(clippy::mixed_attributes_style)] + #![allow(clippy::print_stderr)] + #![allow(clippy::print_stdout)] + #![allow(clippy::single_char_pattern)] + #![allow(clippy::unwrap_used)] + #![allow(clippy::unchecked_duration_subtraction)] + #![allow(clippy::useless_vec)] + #![allow(clippy::needless_pass_by_value)] + //! <!-- @@ end test lint list maintained by maint/add_warning @@ --> + + use super::*; + + #[test] + fn display_qty() { + let chk = |by, s| assert_eq!(Qty(by).to_string(), s); + + chk(10 * 1024, "0.01MiB"); + chk(1024 * 1024, "1.00MiB"); + chk(1000 * 1024 * 1024, "1000.00MiB"); + } +} |