Implement TypeSize for Cache

Cargo.toml

@@ -17,7 +17,8 @@ exclude = [".circleci", ".devcontainer", ".github", ".gitpod.yml", ".vscode"]
 [features]
 default = ["sync"]
 
-sync = ["dashmap"]
+sync = ["dep:dashmap"]
+typesize = ["dep:typesize", "dashmap?/typesize"]
 
 [dependencies]
 crossbeam-channel = "0.5.5"
@@ -31,6 +32,10 @@ triomphe = { version = "0.1.13", default-features = false }
 
 # Optional dependencies (enabled by default)
 dashmap = { version = "6.1", optional = true }
+typesize = { version = "0.1.9", optional = true }
+
+[patch.crates-io.typesize]
+git = "https://github.com/GnomedDev/typesize"
 
 [dev-dependencies]
 anyhow = "1.0.19"

src/common.rs

@@ -7,6 +7,7 @@ pub(crate) mod builder_utils;
 pub(crate) mod deque;
 pub(crate) mod frequency_sketch;
 pub(crate) mod time;
+pub(crate) mod typesize_helpers;
 
 // Note: `CacheRegion` cannot have more than four enum variants. This is because
 // `crate::{sync,unsync}::DeqNodes` uses a `tagptr::TagNonNull<DeqNode<T>, 2>`

src/common/concurrent.rs

@@ -13,6 +13,8 @@ pub(crate) mod atomic_time;
 
 use self::entry_info::EntryInfo;
 
+use super::typesize_helpers::MaybeOwnedArc;
+
 pub(crate) type Weigher<K, V> = Arc<dyn Fn(&K, &V) -> u32 + Send + Sync + 'static>;
 
 pub(crate) trait AccessTime {
@@ -60,8 +62,9 @@ impl<K> KeyDate<K> {
     }
 }
 
+#[cfg(feature = "typesize")]
 pub(crate) struct KeyHashDate<K> {
-    key: Arc<K>,
+    key: MaybeOwnedArc<K>,
     hash: u64,
     entry_info: TrioArc<EntryInfo<K>>,
 }

src/common/concurrent/atomic_time.rs

@@ -2,6 +2,7 @@ use super::Instant;
 
 use std::sync::RwLock;
 
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 pub(crate) struct AtomicInstant {
     instant: RwLock<Option<Instant>>,
 }

src/common/concurrent/deques.rs

@@ -7,6 +7,8 @@ use crate::common::{
 use std::ptr::NonNull;
 use tagptr::TagNonNull;
 use triomphe::Arc as TrioArc;
+
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 pub(crate) struct Deques<K> {
     pub(crate) window: Deque<KeyHashDate<K>>, //    Not used yet.
     pub(crate) probation: Deque<KeyHashDate<K>>,

src/common/concurrent/housekeeper.rs

@@ -18,6 +18,7 @@ pub(crate) trait InnerSync {
     fn now(&self) -> Instant;
 }
 
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 pub(crate) struct Housekeeper {
     is_sync_running: AtomicBool,
     sync_after: AtomicInstant,

src/common/deque.rs

@@ -14,6 +14,8 @@
 
 use std::{marker::PhantomData, ptr::NonNull};
 
+use typesize::TypeSize;
+
 use super::CacheRegion;
 
 // `crate::{sync,unsync}::DeqNodes` uses a `tagptr::TagNonNull<DeqNode<T>, 2>`
@@ -328,6 +330,11 @@ impl<T> Deque<T> {
     }
 }
 
+#[cfg(feature = "typesize")]
+impl<T: TypeSize> typesize::TypeSize for Deque<T> {
+    // TODO: Implement extra_size properly
+}
+
 #[cfg(test)]
 mod tests {
     use super::{CacheRegion::MainProbation, DeqNode, Deque};

src/common/time.rs

@@ -6,6 +6,7 @@ pub(crate) use clock::Clock;
 
 /// a wrapper type over Instant to force checked additions and prevent
 /// unintentional overflow. The type preserve the Copy semantics for the wrapped
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 #[derive(PartialEq, PartialOrd, Clone, Copy)]
 pub(crate) struct Instant(clock::Instant);
 

src/common/time/clock.rs

@@ -6,18 +6,21 @@ use std::{
 #[cfg(test)]
 use std::time::Duration;
 
+use crate::common::typesize_helpers::MaybeOwnedArc;
+
 pub(crate) type Instant = StdInstant;
 
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 pub(crate) struct Clock {
-    mock: Option<Arc<Mock>>,
+    mock: Option<MaybeOwnedArc<Mock>>,
 }
 
 impl Clock {
     #[cfg(test)]
-    pub(crate) fn mock() -> (Clock, Arc<Mock>) {
-        let mock = Arc::new(Mock::default());
+    pub(crate) fn mock() -> (Clock, MaybeOwnedArc<Mock>) {
+        let mock = MaybeOwnedArc::new(Mock::default());
         let clock = Clock {
-            mock: Some(Arc::clone(&mock)),
+            mock: Some(MaybeOwnedArc::clone(&mock)),
         };
         (clock, mock)
     }
@@ -31,6 +34,7 @@ impl Clock {
     }
 }
 
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 pub(crate) struct Mock {
     now: RwLock<Instant>,
 }

src/common/typesize_helpers.rs

@@ -0,0 +1,52 @@
+use std::sync::Arc;
+
+/// Wrapper around `SizableArc<T, Owned>`` with support for disabling typesize.
+///
+/// This denotes an Arc where T's size should be considered when calling `TypeSize::get_size`
+#[derive(Debug)]
+pub(crate) struct MaybeOwnedArc<T>(
+    #[cfg(feature = "typesize")] typesize::ptr::SizableArc<T, typesize::ptr::Owned>,
+    #[cfg(not(feature = "typesize"))] Arc<T>,
+);
+
+impl<T> MaybeOwnedArc<T> {
+    pub(crate) fn new(inner: T) -> Self {
+        Self(Arc::new(inner).into())
+    }
+
+    pub(crate) fn get_inner(self) -> Arc<T> {
+        #[cfg(feature = "typesize")]
+        let inner = self.0 .0;
+        #[cfg(not(feature = "typesize"))]
+        let inner = self.0;
+
+        inner
+    }
+}
+
+#[cfg(feature = "typesize")]
+impl<T: typesize::TypeSize> typesize::TypeSize for MaybeOwnedArc<T> {
+    fn extra_size(&self) -> usize {
+        self.0.extra_size()
+    }
+
+    typesize::if_typesize_details! {
+        fn get_collection_item_count(&self) -> Option<usize> {
+            self.0.get_collection_item_count()
+        }
+    }
+}
+
+impl<T> std::ops::Deref for MaybeOwnedArc<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<T> Clone for MaybeOwnedArc<T> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone().into())
+    }
+}

src/sync/base_cache.rs

@@ -16,6 +16,7 @@ use crate::{
         deque::{DeqNode, Deque},
         frequency_sketch::FrequencySketch,
         time::{CheckedTimeOps, Clock, Instant},
+        typesize_helpers::MaybeOwnedArc,
         CacheRegion,
     },
     Policy,
@@ -38,11 +39,14 @@ use std::{
 };
 use triomphe::Arc as TrioArc;
 
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 pub(crate) struct BaseCache<K, V, S = RandomState> {
-    pub(crate) inner: Arc<Inner<K, V, S>>,
+    pub(crate) inner: MaybeOwnedArc<Inner<K, V, S>>,
+    #[typesize(skip)]
     read_op_ch: Sender<ReadOp<K, V>>,
+    #[typesize(skip)]
     pub(crate) write_op_ch: Sender<WriteOp<K, V>>,
-    pub(crate) housekeeper: Option<Arc<Housekeeper>>,
+    pub(crate) housekeeper: Option<MaybeOwnedArc<Housekeeper>>,
 }
 
 impl<K, V, S> Clone for BaseCache<K, V, S> {
@@ -52,7 +56,7 @@ impl<K, V, S> Clone for BaseCache<K, V, S> {
     /// pointers to the shared internal data structures.
     fn clone(&self) -> Self {
         Self {
-            inner: Arc::clone(&self.inner),
+            inner: MaybeOwnedArc::clone(&self.inner),
             read_op_ch: self.read_op_ch.clone(),
             write_op_ch: self.write_op_ch.clone(),
             housekeeper: self.housekeeper.clone(),
@@ -110,10 +114,10 @@ where
         );
         Self {
             #[cfg_attr(beta_clippy, allow(clippy::arc_with_non_send_sync))]
-            inner: Arc::new(inner),
+            inner: MaybeOwnedArc::new(inner),
             read_op_ch: r_snd,
             write_op_ch: w_snd,
-            housekeeper: Some(Arc::new(Housekeeper::default())),
+            housekeeper: Some(MaybeOwnedArc::new(Housekeeper::default())),
         }
     }
 
@@ -202,7 +206,7 @@ where
         inner: &impl InnerSync,
         ch: &Sender<WriteOp<K, V>>,
         now: Instant,
-        housekeeper: Option<&Arc<Housekeeper>>,
+        housekeeper: Option<&Housekeeper>,
     ) {
         let w_len = ch.len();
 
@@ -255,7 +259,7 @@ where
         op: ReadOp<K, V>,
         now: Instant,
     ) -> Result<(), TrySendError<ReadOp<K, V>>> {
-        self.apply_reads_if_needed(self.inner.as_ref(), now);
+        self.apply_reads_if_needed(&*self.inner, now);
         let ch = &self.read_op_ch;
         match ch.try_send(op) {
             // Discard the ReadOp when the channel is full.
@@ -452,9 +456,12 @@ type CacheStore<K, V, S> = dashmap::DashMap<Arc<K>, TrioArc<ValueEntry<K, V>>, S
 
 type CacheEntryRef<'a, K, V> = DashMapRef<'a, Arc<K>, TrioArc<ValueEntry<K, V>>>;
 
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 pub(crate) struct Inner<K, V, S> {
     max_capacity: Option<u64>,
+    #[typesize(skip)]
     entry_count: AtomicCell<u64>,
+    #[typesize(skip)]
     weighted_size: AtomicCell<u64>,
     cache: CacheStore<K, V, S>,
     build_hasher: S,

src/sync/cache.rs

@@ -224,6 +224,7 @@ use std::{
 /// [build-with-hasher-method]: ./struct.CacheBuilder.html#method.build_with_hasher
 /// [ahash-crate]: https://crates.io/crates/ahash
 ///
+#[cfg_attr(feature = "typesize", derive(typesize::derive::TypeSize))]
 pub struct Cache<K, V, S = RandomState> {
     base: BaseCache<K, V, S>,
 }
@@ -443,15 +444,9 @@ where
 
     pub(crate) fn insert_with_hash(&self, key: Arc<K>, hash: u64, value: V) {
         let (op, now) = self.base.do_insert_with_hash(key, hash, value);
-        let hk = self.base.housekeeper.as_ref();
-        Self::schedule_write_op(
-            self.base.inner.as_ref(),
-            &self.base.write_op_ch,
-            op,
-            now,
-            hk,
-        )
-        .expect("Failed to insert");
+        let hk = self.base.housekeeper.as_deref();
+        Self::schedule_write_op(&*self.base.inner, &self.base.write_op_ch, op, now, hk)
+            .expect("Failed to insert");
     }
 
     /// Discards any cached value for the key.
@@ -466,15 +461,9 @@ where
         if let Some(kv) = self.base.remove_entry(key) {
             let op = WriteOp::Remove(kv);
             let now = self.base.current_time_from_expiration_clock();
-            let hk = self.base.housekeeper.as_ref();
-            Self::schedule_write_op(
-                self.base.inner.as_ref(),
-                &self.base.write_op_ch,
-                op,
-                now,
-                hk,
-            )
-            .expect("Failed to remove");
+            let hk = self.base.housekeeper.as_deref();
+            Self::schedule_write_op(&*self.base.inner, &self.base.write_op_ch, op, now, hk)
+                .expect("Failed to remove");
         }
     }
 
@@ -576,7 +565,7 @@ where
         ch: &Sender<WriteOp<K, V>>,
         op: WriteOp<K, V>,
         now: Instant,
-        housekeeper: Option<&Arc<Housekeeper>>,
+        housekeeper: Option<&Housekeeper>,
     ) -> Result<(), TrySendError<WriteOp<K, V>>> {
         let mut op = op;