summaryrefslogtreecommitdiff
path: root/crates/tor-memquota/src/mtracker
diff options
context:
space:
mode:
Diffstat (limited to 'crates/tor-memquota/src/mtracker')
-rw-r--r--crates/tor-memquota/src/mtracker/bookkeeping.rs308
-rw-r--r--crates/tor-memquota/src/mtracker/reclaim.rs396
-rw-r--r--crates/tor-memquota/src/mtracker/reclaim/deferred_drop.rs81
-rw-r--r--crates/tor-memquota/src/mtracker/test.rs678
-rw-r--r--crates/tor-memquota/src/mtracker/total_qty_notifier.rs70
5 files changed, 1533 insertions, 0 deletions
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);
+ }
+}
generated by cgit v1.2.3 (git 2.46.0) at 2026-08-29 07:51:59 +0000