Add OnDemandHashMap and TokenSet

src/utils/on_demand_hashmap.rs

@@ -0,0 +1,304 @@
+use std::{
+    collections::HashMap,
+    fmt,
+    hash::Hash,
+    marker::PhantomData,
+    ops::Deref,
+    sync::{Arc, Weak},
+};
+
+use tokio::sync::{mpsc, Mutex, OwnedMutexGuard, RwLock};
+use tracing::{trace, warn, Level};
+
+use crate::observability::METRICS;
+
+/// Data shared between [`OnDemandHashMap`] and the cleanup task
+///
+/// Importantly it does not contain the `cleanup_sender`, since it getting
+/// dropped signals the cleanup task to exit. If the cleanup task had an owned
+/// reference to it, the only way for it to exit would be for every [`Entry`] to
+/// be dropped, we don't want to rely on that.
+struct SharedData<K, V> {
+    name: Arc<str>,
+    /// Values are owned by their [entries][Entry]
+    entries: RwLock<HashMap<K, Weak<V>>>,
+}
+
+impl<K, V> SharedData<K, V>
+where
+    K: Hash + Eq + Clone + fmt::Debug,
+{
+    #[tracing::instrument(
+        level = Level::TRACE,
+        skip(self),
+        fields(name = self.name.as_ref()),
+    )]
+    async fn try_cleanup_entry(&self, key: K) {
+        let mut map = self.entries.write().await;
+
+        let Some(weak) = map.get(&key) else {
+            trace!("Entry has already been cleaned up");
+            return;
+        };
+
+        if weak.strong_count() != 0 {
+            trace!("Entry is in use");
+            return;
+        }
+
+        trace!("Cleaning up unused entry");
+        map.remove(&key);
+        METRICS.record_on_demand_hashmap_size(self.name.clone(), map.len());
+    }
+
+    #[tracing::instrument(level = Level::TRACE, skip(map))]
+    fn try_get_live_value(
+        pass: usize,
+        map: &HashMap<K, Weak<V>>,
+        key: &K,
+    ) -> Option<Arc<V>> {
+        if let Some(value) = map.get(key) {
+            if let Some(value) = value.upgrade() {
+                trace!(pass, "Using existing value");
+                return Some(value);
+            }
+
+            trace!(
+                pass,
+                "Existing value is stale and needs cleanup, creating new"
+            );
+        } else {
+            trace!(pass, "No existing value, creating new");
+        }
+
+        None
+    }
+
+    /// Either returns an existing live value, or creates a new one and inserts
+    /// it into the map.
+    #[tracing::instrument(level = Level::TRACE, skip(self, create))]
+    async fn get_or_insert_with<F>(&self, key: &K, create: F) -> Arc<V>
+    where
+        F: FnOnce() -> V,
+    {
+        {
+            // first, take a read lock and try to get an existing value
+
+            // TODO check if this fast path actually makes it faster, possibly
+            // make it configurable per OnDemandHashMap depending on contention
+            // and how expensive create() is
+            let map = self.entries.read().await;
+            if let Some(v) = Self::try_get_live_value(1, &map, key) {
+                return v;
+            }
+        }
+
+        // no entry or it has died, create a new one
+        let value = Arc::new(create());
+        let weak = Arc::downgrade(&value);
+
+        // take a write lock, try again, otherwise insert our new value
+        let mut map = self.entries.write().await;
+        if let Some(v) = Self::try_get_live_value(2, &map, key) {
+            // another entry showed up while we had let go of the lock,
+            // use that
+            drop(value);
+            drop(weak);
+            return v;
+        }
+
+        map.insert(key.clone(), weak);
+        METRICS.record_on_demand_hashmap_size(self.name.clone(), map.len());
+
+        value
+    }
+}
+
+/// A [`HashMap`] whose entries are automatically removed once they are no
+/// longer referenced.
+pub(crate) struct OnDemandHashMap<K, V> {
+    /// The data shared between the [`OnDemandHashMap`] and the cleanup task.
+    shared: Arc<SharedData<K, V>>,
+    /// This is the only non-[weak][mpsc::WeakUnboundedSender] `Sender`, which
+    /// means that dropping the `OnDemandHashMap` causes the cleanup
+    /// process to exit.
+    cleanup_sender: mpsc::UnboundedSender<K>,
+}
+
+impl<K, V> OnDemandHashMap<K, V>
+where
+    K: Hash + Eq + Clone + fmt::Debug + Send + Sync + 'static,
+    V: Send + Sync + 'static,
+{
+    /// Creates a new `OnDemandHashMap`. The `name` is used for metrics and
+    /// should be unique to this instance.
+    pub(crate) fn new(name: String) -> Self {
+        let (cleanup_sender, mut receiver) = mpsc::unbounded_channel();
+
+        let shared = Arc::new(SharedData {
+            name: name.into(),
+            entries: RwLock::new(HashMap::new()),
+        });
+
+        {
+            let shared = Arc::clone(&shared);
+            tokio::task::spawn(async move {
+                loop {
+                    let Some(key) = receiver.recv().await else {
+                        trace!(
+                            name = shared.name.as_ref(),
+                            "Channel has died, exiting cleanup task"
+                        );
+                        return;
+                    };
+
+                    shared.try_cleanup_entry(key).await;
+                }
+            });
+        }
+
+        Self {
+            shared,
+            cleanup_sender,
+        }
+    }
+
+    #[tracing::instrument(level = Level::TRACE, skip(self, create))]
+    pub(crate) async fn get_or_insert_with<F>(
+        &self,
+        key: K,
+        create: F,
+    ) -> Entry<K, V>
+    where
+        F: FnOnce() -> V,
+    {
+        let value = self.shared.get_or_insert_with(&key, create).await;
+
+        Entry {
+            drop_guard: EntryDropGuard {
+                cleanup_sender: self.cleanup_sender.downgrade(),
+                key: Some(key),
+            },
+            value,
+        }
+    }
+}
+
+struct EntryDropGuard<K> {
+    cleanup_sender: mpsc::WeakUnboundedSender<K>,
+    /// Only `None` during `drop()`
+    key: Option<K>,
+}
+
+impl<K> Drop for EntryDropGuard<K> {
+    fn drop(&mut self) {
+        let Some(cleanup_sender) = self.cleanup_sender.upgrade() else {
+            trace!("Backing map has already been dropped");
+            return;
+        };
+
+        if let Err(error) = cleanup_sender
+            .send(self.key.take().expect("drop should only be called once"))
+        {
+            warn!(%error, "Failed to send cleanup message");
+        };
+    }
+}
+
+/// A wrapper around a key/value pair inside an [`OnDemandHashMap`]
+///
+/// If every `Entry` for a specific key is dropped, the value is removed from
+/// the map.
+pub(crate) struct Entry<K, V> {
+    drop_guard: EntryDropGuard<K>,
+    value: Arc<V>,
+}
+
+impl<K, V> Deref for Entry<K, V> {
+    type Target = V;
+
+    fn deref(&self) -> &Self::Target {
+        self.value.as_ref()
+    }
+}
+
+/// Internal zero-sized type used to swallow the [`TokenSet`]'s marker type
+struct TokenMarker<T>(PhantomData<fn(T) -> T>);
+
+/// A collection of dynamically-created locks, one for each value of `K`.
+///
+/// A given key can be locked using [`TokenSet::lock_key()`], which will either
+/// return an ownership token immediately if the key is not currently locked, or
+/// wait until the previous lock has been released.
+///
+/// The marker type `M` can be used to disambiguate different `TokenSet`
+/// instances to avoid misuse of tokens.
+pub(crate) struct TokenSet<K, M = ()> {
+    inner: OnDemandHashMap<K, Mutex<TokenMarker<M>>>,
+}
+
+impl<K, M> TokenSet<K, M>
+where
+    K: Hash + Eq + Clone + fmt::Debug + Send + Sync + 'static,
+    M: 'static,
+{
+    /// Creates a new `TokenSet`. The `name` is used for metrics and should be
+    /// unique to this instance.
+    pub(crate) fn new(name: String) -> Self {
+        Self {
+            inner: OnDemandHashMap::new(name),
+        }
+    }
+
+    /// Locks this key in the `TokenSet`, returning a token proving
+    /// unique access.
+    #[tracing::instrument(level = Level::TRACE, skip(self))]
+    pub(crate) async fn lock_key(&self, key: K) -> KeyToken<K, M> {
+        let Entry {
+            drop_guard,
+            value,
+        } = self
+            .inner
+            .get_or_insert_with(key, || Mutex::new(TokenMarker(PhantomData)))
+            .await;
+
+        KeyToken {
+            drop_guard,
+            _mutex_guard: value.lock_owned().await,
+        }
+    }
+}
+
+/// Unique token for a given key in a [`TokenSet`].
+///
+/// Ownership of this token proves that no other [`KeyToken`] for this key in
+/// this [`TokenSet`] currently exists.
+///
+/// Access to the underlying key is provided by a [`Deref`] impl.
+pub(crate) struct KeyToken<K, M = ()> {
+    drop_guard: EntryDropGuard<K>,
+    _mutex_guard: OwnedMutexGuard<TokenMarker<M>>,
+}
+
+impl<K, M> Deref for KeyToken<K, M> {
+    type Target = K;
+
+    fn deref(&self) -> &Self::Target {
+        self.drop_guard
+            .key
+            .as_ref()
+            .expect("key should only be None during Drop")
+    }
+}
+
+impl<K: fmt::Debug, M> fmt::Debug for KeyToken<K, M> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{:?}", &**self)
+    }
+}
+
+impl<K: fmt::Display, M> fmt::Display for KeyToken<K, M> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", &**self)
+    }
+}