tor_memquota/
refcount.rs

Help
1//! Helpers for reference counting
2//!
3//! Two main purposes:
4//!
5//!  * Consistent handling of overflow and underflow
6//!  * Assurance of incrementing/decrementing as appropriate,
7//!    including in combination with a slotmap containing the referenced data.
8//!
9//! The caller is responsible for making sure that the *right instance*'s
10//! [`Count`] is passed to the methods on [`Ref`].
11//
12// There are no separate tests for this module.  Many of the tests would want to
13// exercise the `Ref`s drop bomb, which is troublesome since it's panic in drop,
14// which they're making Rust treat as an abort upstream.
15// (This scheme did detect a bug or two during development testing,
16// so the drop bomb is known to work.)
17//
18// Anyway, these functions are mostly newtype veneers over existing functionality.
19// They're tested by the MemoryQuotaTracker's tests.
20
21use crate::internal_prelude::*;
22
23/// Local alias for the counter type
24pub(crate) type RawCount = u32;
25
26/// Decrement a refcount and maybe remove a corresponding slotmap entry
27///
28/// ```rust,ignore
29/// fn slotmap_dec_ref!<K, V>(
30///    slotmap: &mut SlotMap<K, V>,
31///    ref_: Ref<K>,
32///    refcount: &mut Count<K>, // (typically) borrows from slotmap
33/// )
34/// ```
35//
36// This macro is a bit out-of-position, up here, because we want to be able to link
37// to it in our rustdocs.
38macro_rules! slotmap_dec_ref { { $slotmap:expr, $ref_:expr, $refcount:expr } => { { {
39    use $crate::refcount::*;
40    let key: Ref<_> = $ref_;
41    let refcount: &mut Count<_> = $refcount;
42    if let Some(Garbage(key)) = key.dispose(refcount) {
43        let slotmap: &mut SlotMap<_, _> = $slotmap;
44        let removed = slotmap.remove(key).expect("entry vanished or wrong key passed?!");
45        Some(Garbage(removed))
46    } else {
47        None
48    }
49} } } }
50
51/// A reference count, counting references with id type `K`
52#[derive(Default, Educe, Ord, PartialOrd, Eq, PartialEq, Deref)]
53#[educe(Debug)]
54pub(crate) struct Count<K> {
55    /// Actual count of references
56    #[deref]
57    count: RawCount,
58    /// Bind to the specific key type
59    // K is generally Send + Sync + 'static so we don't care about variance etc.
60    #[educe(Debug(ignore))]
61    marker: PhantomData<K>,
62}
63
64/// An copy of a [`slotmap_careful::Key`] `K`, which is counted by a `RefCount`
65///
66/// Ie, a key of type `K` with the property that it
67/// keeps the refcounted data structure alive.
68///
69/// Must always be deleted using [`dispose`](Ref::dispose), not dropped.
70/// In tests, dropping a `RefCounted` will panic.
71///
72/// The `Default` value does *not* contribute to a reference count,
73/// and is fine to drop.
74#[derive(Deref, Educe)]
75#[educe(Debug, Default, Ord, Eq, PartialEq)]
76pub(crate) struct Ref<K: slotmap_careful::Key> {
77    /// Actual key (without generics)
78    #[deref]
79    raw_key: K,
80    /// Bind to the specific key type
81    #[educe(Debug(ignore))]
82    marker: PhantomData<K>,
83    /// Drop bomb
84    ///
85    /// Also forces `Ref` not to be Clone
86    #[educe(Debug(ignore), Ord(ignore), Eq(ignore), PartialEq(ignore))]
87    #[allow(dead_code)]
88    bomb: DropBombCondition,
89}
90
91// educe's Ord is open-coded and triggers clippy::non_canonical_partial_ord_impl
92impl<K: slotmap_careful::Key> PartialOrd for Ref<K> {
93    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
94        Some(self.cmp(other))
95    }
96}
97
98// Ideally we'd assert_not_impl on Ref but it has generics
99assert_not_impl_any!(DropBombCondition: Clone);
100
101/// Error: refcount overflowed
102#[derive(Debug, Clone, Error, Eq, PartialEq)]
103#[error("memory tracking refcount overflowed")]
104pub(crate) struct Overflow;
105
106/// Something which has become garbage
107///
108/// Often used within `Option`, for clarity.  Examples:
109///
110///  * Key whose reference count has reached zero - see [`Ref::dispose`]
111///  * Value removed from a SlotMap - see [`slotmap_dec_ref!`]
112#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
113pub(crate) struct Garbage<K>(pub(crate) K);
114
115impl<K> Count<K> {
116    /// Make a new refcount with a specified value
117    const fn new_raw(count: RawCount) -> Self {
118        Count {
119            count,
120            marker: PhantomData,
121        }
122    }
123
124    /// Obtain this counter as a `usize`
125    ///
126    /// (Reference counts are `u32`, so this might be a conversion.)
127    pub(crate) fn as_usize(&self) -> usize {
128        // On a 16-bit platform this could theoretically overflow,
129        // but there would have to be >2^16 clones, which would be impossible.
130        let r: u32 = **self;
131        r as usize
132    }
133}
134
135/// Increment this refcount, but don't care about any [`Ref`]s
136fn inc_raw(c: &mut RawCount) -> Result<(), Overflow> {
137    *c = c.checked_add(1).ok_or(Overflow)?;
138    Ok(())
139}
140
141/// Decrement this refcount, but don't care about any [`Ref`]s
142///
143/// Returns [`Some(Garbage(()))`] if the count reached zero
144fn dec_raw(c: &mut RawCount) -> Option<Garbage<()>> {
145    *c = c
146        .checked_sub(1)
147        // if this happens, our data structure is corrupted, very bad
148        .expect("refcount underflow");
149    (*c == 0).then_some(Garbage(()))
150}
151
152impl<K: slotmap_careful::Key> Ref<K> {
153    /// Create a refcounted reference `Ref` from an un-counted key, incrementing the count
154    pub(crate) fn new(key: K, count: &mut Count<K>) -> Result<Self, Overflow> {
155        inc_raw(&mut count.count)?;
156        Ok(Ref::from_raw(key))
157    }
158
159    /// Creates a null `Ref`, which doesn't refer to any slot (lookups always give `None`)
160    pub(crate) fn null() -> Self {
161        Ref::from_raw(K::null())
162    }
163
164    /// Internal function for creating a `Ref`
165    fn from_raw(raw_key: K) -> Self {
166        Ref {
167            raw_key,
168            marker: PhantomData,
169            bomb: DropBombCondition::new_armed(),
170        }
171    }
172
173    /// Dispose of a refcounted reference `Ref`, decrementing the count
174    ///
175    /// If the count reaches zero, the raw key is returned;
176    /// the caller should remove the corresponding data from the data structure.
177    pub(crate) fn dispose(mut self, refcount: &mut Count<K>) -> Option<Garbage<K>> {
178        let was = mem::take(&mut self.raw_key);
179        assert!(!was.is_null());
180        dec_raw(&mut refcount.count).map(|_: Garbage<()>| Garbage(was))
181    }
182
183    /// Dispose of a refcounted reference whose container no longer exists
184    ///
185    /// # CORRECTNESS
186    ///
187    /// This just forgets the reference, without decrementing any refcount.
188    /// If the container *does* still exist, a ref count ref will be leaked.
189    pub(crate) fn dispose_container_destroyed(mut self) {
190        let _: K = mem::take(&mut self.raw_key);
191    }
192}
193
194impl<K: slotmap_careful::Key> DefaultExtTake for Ref<K> {}
195
196/// Insert a new entry into a slotmap using refcounted keys
197///
198/// `value_maker` should take the provided `Count`,
199/// and incorporate it into a new value.
200///
201/// On return, the entry will be in the map, and there will be one reference,
202/// which is returned.
203///
204/// There is no corresponding `slotmap_remove` in this module.
205/// Use [`Ref::dispose`] and handle any [`Garbage`] it returns.
206pub(crate) fn slotmap_insert<K: slotmap_careful::Key, V>(
207    slotmap: &mut SlotMap<K, V>,
208    value_maker: impl FnOnce(Count<K>) -> V,
209) -> Ref<K> {
210    let (ref_, ()) = slotmap_try_insert(slotmap, move |refcount| {
211        Ok::<_, Void>((value_maker(refcount), ()))
212    })
213    .void_unwrap();
214    ref_
215}
216
217/// Insert a new entry into a slotmap using refcounted keys, fallibly and with extra data
218///
219/// Like [`slotmap_insert`] but:
220///  * `value_maker` can also return extra return data `RD` to the caller
221///  * `value_maker` is allowed to fail.
222///
223/// On successful return, the entry will be in the map, and
224/// the new `Ref` is returned along with the data `D`.
225pub(crate) fn slotmap_try_insert<K: slotmap_careful::Key, V, E, RD>(
226    slotmap: &mut SlotMap<K, V>,
227    value_maker: impl FnOnce(Count<K>) -> Result<(V, RD), E>,
228) -> Result<(Ref<K>, RD), E> {
229    let refcount = Count::new_raw(1);
230    let (value, data) = value_maker(refcount)?;
231    let raw_key = slotmap.insert(value);
232    let ref_ = Ref {
233        raw_key,
234        marker: PhantomData,
235        bomb: DropBombCondition::new_armed(),
236    };
237    Ok((ref_, data))
238}
239
240/// Unconditionally remove en entry from the slotmap, given a strong ref
241///
242/// Other references to this entry will become dangling.
243pub(crate) fn slotmap_remove_early<K: slotmap_careful::Key, V>(
244    slotmap: &mut SlotMap<K, V>,
245    key: Ref<K>,
246) -> Option<V> {
247    let r = slotmap.remove(*key);
248    // Correctness: this becomes a reference to a missing entry
249    key.dispose_container_destroyed();
250    r
251}
252
253#[cfg(test)]
254impl<K: slotmap_careful::Key> Drop for Ref<K> {
255    fn drop(&mut self) {
256        drop_bomb_disarm_assert!(self.bomb, self.raw_key.is_null(),);
257    }
258}
259
260impl From<Overflow> for Error {
261    fn from(_overflow: Overflow) -> Error {
262        internal!("reference count overflow in memory tracking (out-of-control subsystem?)").into()
263    }
264}
265
266#[cfg(test)]
267mod test {
268    // @@ begin test lint list maintained by maint/add_warning @@
269    #![allow(clippy::bool_assert_comparison)]
270    #![allow(clippy::clone_on_copy)]
271    #![allow(clippy::dbg_macro)]
272    #![allow(clippy::mixed_attributes_style)]
273    #![allow(clippy::print_stderr)]
274    #![allow(clippy::print_stdout)]
275    #![allow(clippy::single_char_pattern)]
276    #![allow(clippy::unwrap_used)]
277    #![allow(clippy::unchecked_duration_subtraction)]
278    #![allow(clippy::useless_vec)]
279    #![allow(clippy::needless_pass_by_value)]
280    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
281    #![allow(clippy::let_and_return)] // TODO this lint is annoying and we should disable it
282
283    use super::*;
284
285    slotmap_careful::new_key_type! {
286        struct Id;
287    }
288    #[derive(Eq, PartialEq, Debug)]
289    struct Record {
290        refcount: Count<Id>,
291    }
292    type Map = SlotMap<Id, Record>;
293
294    fn setup() -> (Map, Ref<Id>) {
295        let mut map = Map::default();
296        let ref_ = slotmap_insert(&mut map, |refcount| Record { refcount });
297        (map, ref_)
298    }
299
300    #[test]
301    fn good() {
302        let (mut map, ref1) = setup();
303
304        let ent = map.get_mut(*ref1).unwrap();
305        let ref2 = Ref::new(*ref1, &mut ent.refcount).unwrap();
306
307        let g1: Option<Garbage<Record>> = slotmap_dec_ref!(&mut map, ref1, &mut ent.refcount);
308        assert_eq!(g1, None);
309
310        let ent = map.get_mut(*ref2).unwrap();
311        let g2: Option<Garbage<Record>> = slotmap_dec_ref!(&mut map, ref2, &mut ent.refcount);
312        assert!(g2.is_some());
313    }
314
315    #[test]
316    fn try_insert_fail() {
317        let mut map = Map::default();
318        let () = slotmap_try_insert::<_, _, _, String>(&mut map, |_refcount| Err(())).unwrap_err();
319    }
320
321    #[test]
322    fn drop_ref_without_decrement() {
323        let (_map, mut ref1) = setup();
324        let h = ref1.bomb.make_simulated();
325        drop(ref1);
326        h.expect_exploded();
327    }
328}
tor_memquota/refcount.rs

tor_memquota/
refcount.rs
