Line data    Source code

       1              : use std::{collections::HashMap, sync::Arc};
       2              : 
       3              : /// A map of locks covering some arbitrary identifiers. Useful if you have a collection of objects but don't
       4              : /// want to embed a lock in each one, or if your locking granularity is different to your object granularity.
       5              : /// For example, used in the storage controller where the objects are tenant shards, but sometimes locking
       6              : /// is needed at a tenant-wide granularity.
       7              : pub(crate) struct IdLockMap<T>
       8              : where
       9              :     T: Eq + PartialEq + std::hash::Hash,
      10              : {
      11              :     /// A synchronous lock for getting/setting the async locks that our callers will wait on.
      12              :     entities: std::sync::Mutex<std::collections::HashMap<T, Arc<tokio::sync::RwLock<()>>>>,
      13              : }
      14              : 
      15              : impl<T> IdLockMap<T>
      16              : where
      17              :     T: Eq + PartialEq + std::hash::Hash,
      18              : {
      19            0 :     pub(crate) fn shared(
      20            0 :         &self,
      21            0 :         key: T,
      22            0 :     ) -> impl std::future::Future<Output = tokio::sync::OwnedRwLockReadGuard<()>> {
      23            0 :         let mut locked = self.entities.lock().unwrap();
      24            0 :         let entry = locked.entry(key).or_default();
      25            0 :         entry.clone().read_owned()
      26            0 :     }
      27              : 
      28            0 :     pub(crate) fn exclusive(
      29            0 :         &self,
      30            0 :         key: T,
      31            0 :     ) -> impl std::future::Future<Output = tokio::sync::OwnedRwLockWriteGuard<()>> {
      32            0 :         let mut locked = self.entities.lock().unwrap();
      33            0 :         let entry = locked.entry(key).or_default();
      34            0 :         entry.clone().write_owned()
      35            0 :     }
      36              : 
      37              :     /// Rather than building a lock guard that re-takes the [`Self::entities`] lock, we just do
      38              :     /// periodic housekeeping to avoid the map growing indefinitely
      39            0 :     pub(crate) fn housekeeping(&self) {
      40            0 :         let mut locked = self.entities.lock().unwrap();
      41            0 :         locked.retain(|_k, lock| lock.try_write().is_err())
      42            0 :     }
      43              : }
      44              : 
      45              : impl<T> Default for IdLockMap<T>
      46              : where
      47              :     T: Eq + PartialEq + std::hash::Hash,
      48              : {
      49            0 :     fn default() -> Self {
      50            0 :         Self {
      51            0 :             entities: std::sync::Mutex::new(HashMap::new()),
      52            0 :         }
      53            0 :     }
      54              : }