tor_memquota/mtracker/
total_qty_notifier.rs

1//! `TotalQtyNotifier`
2//!
3//! This newtype assures that we wake up the reclamation task when nceessary
4
5use super::*;
6
7/// Wrapper for `TotalQty`
8#[derive(Deref, Debug)]
9pub(super) struct TotalQtyNotifier {
10    /// Total memory usage
11    ///
12    /// Invariant: equal to
13    /// ```text
14    ///    Σ        Σ         PRecord.used
15    ///     ARecord  PRecord
16    /// ```
17    #[deref]
18    total_used: TotalQty,
19
20    /// Condvar to wake up the reclamation task
21    ///
22    /// The reclamation task has another clone of this
23    reclamation_task_wakeup: mpsc::Sender<()>,
24}
25
26impl TotalQtyNotifier {
27    /// Make a new `TotalQtyNotifier`, which will notify a specified condvar
28    pub(super) fn new_zero(reclamation_task_wakeup: mpsc::Sender<()>) -> Self {
29        TotalQtyNotifier {
30            total_used: TotalQty::ZERO,
31            reclamation_task_wakeup,
32        }
33    }
34
35    /// Record that some memory has been (or will be) allocated by a participant
36    ///
37    /// Signals the wakeup task if we need to.
38    pub(super) fn claim(
39        &mut self,
40        precord: &mut PRecord,
41        want: Qty,
42        config: &ConfigInner,
43    ) -> crate::Result<ClaimedQty> {
44        let got = self
45            .total_used
46            .claim(&mut precord.used, want)
47            .ok_or_else(|| internal!("integer overflow attempting to add claim {}", want))?;
48        self.maybe_wakeup(config);
49        Ok(got)
50    }
51
52    /// Check to see if we need to wake up the reclamation task, and if so, do so
53    pub(super) fn maybe_wakeup(&mut self, config: &ConfigInner) {
54        if self.total_used > config.max {
55            match self.reclamation_task_wakeup.try_send(()) {
56                Ok(()) => {}
57                Err(e) if e.is_full() => {}
58                // reactor shutting down, having dropped reclamation task?
59                Err(e) => debug!("could not notify reclamation task: {e}"),
60            };
61        }
62    }
63
64    /// Declare this poisoned, and prevent further claims
65    pub(super) fn set_poisoned(&mut self) {
66        self.total_used.set_poisoned();
67    }
68
69    /// Record that some memory has been (or will be) freed by a participant
70    pub(super) fn release(&mut self, precord: &mut PRecord, have: ClaimedQty) // infallible
71    {
72        // TODO if the participant's usage underflows, tell it to reclaim
73        // (and log some kind of internal error)
74        self.total_used.release(&mut precord.used, have);
75    }
76}