Add SharedCell helper for wrapping types used in Python (#391)

Summary:
Pull Request resolved: #391

Adds a helper class, implemented as a thin wrapper around `tokio::sync::RwLock`,
to help manage Rust structs that get shared via Python wrappers.  Access to the wrapped
struct is supported through `.borrow()` and the object can be consumed via `.take()`
when there are no outstanding borrows, leaving the cell in an unusable state.

Reviewed By: mariusae, shayne-fletcher

Differential Revision: D77244960

fbshipit-source-id: a7479b2df6189a328722c25be3d97d30887c05cf

Author: andrewjcg

hyperactor_mesh/Cargo.toml

@@ -40,6 +40,7 @@ libc = "0.2.139"
 mockall = "0.13.1"
 ndslice = { version = "0.0.0", path = "../ndslice" }
 nix = { version = "0.29.0", features = ["dir", "event", "hostname", "inotify", "ioctl", "mman", "mount", "net", "poll", "ptrace", "reboot", "resource", "sched", "signal", "term", "time", "user", "zerocopy"] }
+preempt_rwlock = { version = "0.0.0", path = "../preempt_rwlock" }
 rand = { version = "0.8", features = ["small_rng"] }
 serde = { version = "1.0.185", features = ["derive", "rc"] }
 serde_bytes = "0.11"

hyperactor_mesh/src/lib.rs

@@ -24,6 +24,7 @@ pub mod mesh_selection;
 mod metrics;
 pub mod proc_mesh;
 pub mod reference;
+pub mod shared_cell;
 pub mod shortuuid;
 pub mod test_utils;
 

hyperactor_mesh/src/shared_cell.rs

@@ -0,0 +1,263 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::fmt::Debug;
+use std::ops::Deref;
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use async_trait::async_trait;
+use dashmap::DashMap;
+use futures::future::try_join_all;
+use preempt_rwlock::OwnedPreemptibleRwLockReadGuard;
+use preempt_rwlock::PreemptibleRwLock;
+use tokio::sync::TryLockError;
+
+#[derive(thiserror::Error, Debug)]
+pub struct EmptyCellError {}
+
+impl From<TryLockError> for EmptyCellError {
+    fn from(_err: TryLockError) -> Self {
+        Self {}
+    }
+}
+
+impl std::fmt::Display for EmptyCellError {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "already taken")
+    }
+}
+
+#[derive(thiserror::Error, Debug)]
+pub enum TryTakeError {
+    #[error("already taken")]
+    Empty,
+    #[error("cannot lock: {0}")]
+    TryLockError(#[from] TryLockError),
+}
+
+struct PoolRef {
+    map: Arc<DashMap<usize, Arc<dyn SharedCellDiscard + Send + Sync>>>,
+    key: usize,
+}
+
+impl Debug for PoolRef {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("PoolRef").field("key", &self.key).finish()
+    }
+}
+
+#[derive(Debug)]
+struct Inner<T> {
+    value: Option<T>,
+    pool: Option<PoolRef>,
+}
+
+impl<T> Drop for Inner<T> {
+    fn drop(&mut self) {
+        if let Some(pool) = &self.pool {
+            pool.map.remove(&pool.key);
+        }
+    }
+}
+
+/// A wrapper class that facilitates sharing an item across different users, supporting:
+/// - Ability grab a reference-counted reference to the item
+/// - Ability to consume the item, leaving the cell in an unusable state
+#[derive(Debug)]
+pub struct SharedCell<T> {
+    inner: Arc<PreemptibleRwLock<Inner<T>>>,
+}
+
+impl<T> Clone for SharedCell<T> {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+        }
+    }
+}
+
+impl<T> From<T> for SharedCell<T> {
+    fn from(value: T) -> Self {
+        Self {
+            inner: Arc::new(PreemptibleRwLock::new(Inner {
+                value: Some(value),
+                pool: None,
+            })),
+        }
+    }
+}
+
+impl<T> SharedCell<T> {
+    fn with_pool(value: T, pool: PoolRef) -> Self {
+        Self {
+            inner: Arc::new(PreemptibleRwLock::new(Inner {
+                value: Some(value),
+                pool: Some(pool),
+            })),
+        }
+    }
+}
+
+pub struct SharedCellRef<T, U = T> {
+    guard: OwnedPreemptibleRwLockReadGuard<Inner<T>, U>,
+}
+
+impl<T> SharedCellRef<T> {
+    fn from(guard: OwnedPreemptibleRwLockReadGuard<Inner<T>>) -> Result<Self, EmptyCellError> {
+        if guard.value.is_none() {
+            return Err(EmptyCellError {});
+        }
+        Ok(Self {
+            guard: OwnedPreemptibleRwLockReadGuard::map(guard, |guard| {
+                guard.value.as_ref().unwrap()
+            }),
+        })
+    }
+
+    pub fn map<F, U>(self, f: F) -> SharedCellRef<T, U>
+    where
+        F: FnOnce(&T) -> &U,
+    {
+        SharedCellRef {
+            guard: OwnedPreemptibleRwLockReadGuard::map(self.guard, f),
+        }
+    }
+
+    pub async fn preempted(&self) {
+        self.guard.preempted().await
+    }
+}
+
+impl<T, U: Debug> Debug for SharedCellRef<T, U> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        Debug::fmt(&**self, f)
+    }
+}
+
+impl<T, U> Deref for SharedCellRef<T, U> {
+    type Target = U;
+
+    fn deref(&self) -> &Self::Target {
+        &self.guard
+    }
+}
+
+impl<T> SharedCell<T> {
+    /// Borrow the cell, returning a reference to the item. If the cell is empty, returns an error.
+    /// While references are held, the cell cannot be taken below.
+    pub fn borrow(&self) -> Result<SharedCellRef<T>, EmptyCellError> {
+        SharedCellRef::from(self.inner.clone().try_read_owned()?)
+    }
+
+    /// Take the item out of the cell, leaving it in an unusable state.
+    pub async fn take(&self) -> Result<T, EmptyCellError> {
+        let mut inner = self.inner.write(true).await;
+        inner.value.take().ok_or(EmptyCellError {})
+    }
+
+    pub fn blocking_take(&self) -> Result<T, EmptyCellError> {
+        let mut inner = self.inner.blocking_write(true);
+        inner.value.take().ok_or(EmptyCellError {})
+    }
+
+    pub fn try_take(&self) -> Result<T, TryTakeError> {
+        let mut inner = self.inner.try_write(true)?;
+        inner.value.take().ok_or(TryTakeError::Empty)
+    }
+}
+
+/// A pool of `SharedCell`s which can be used to mass `take()` and discard them all at once.
+pub struct SharedCellPool {
+    map: Arc<DashMap<usize, Arc<dyn SharedCellDiscard + Send + Sync>>>,
+    token: AtomicUsize,
+}
+
+impl SharedCellPool {
+    pub fn new() -> Self {
+        Self {
+            map: Arc::new(DashMap::new()),
+            token: AtomicUsize::new(0),
+        }
+    }
+
+    pub fn insert<T>(&self, value: T) -> SharedCell<T>
+    where
+        T: Send + Sync + 'static,
+    {
+        let map = self.map.clone();
+        let key = self.token.fetch_add(1, Ordering::Relaxed);
+        let pool = PoolRef { map, key };
+        let value: SharedCell<_> = SharedCell::with_pool(value, pool);
+        self.map.entry(key).insert(Arc::new(value.clone()));
+        value
+    }
+
+    /// Run `take` on all cells in the pool and immediately drop them.
+    pub async fn discard_all(self) -> Result<(), EmptyCellError> {
+        try_join_all(
+            self.map
+                .iter()
+                .map(|r| async move { r.value().discard().await }),
+        )
+        .await?;
+        Ok(())
+    }
+}
+
+/// Trait to facilitate storing `SharedCell`s of different types in a single pool.
+#[async_trait]
+pub trait SharedCellDiscard {
+    async fn discard(&self) -> Result<(), EmptyCellError>;
+    fn blocking_discard(&self) -> Result<(), EmptyCellError>;
+    fn try_discard(&self) -> Result<(), TryTakeError>;
+}
+
+#[async_trait]
+impl<T: Send + Sync> SharedCellDiscard for SharedCell<T> {
+    fn try_discard(&self) -> Result<(), TryTakeError> {
+        self.try_take()?;
+        Ok(())
+    }
+
+    async fn discard(&self) -> Result<(), EmptyCellError> {
+        self.take().await?;
+        Ok(())
+    }
+
+    fn blocking_discard(&self) -> Result<(), EmptyCellError> {
+        self.blocking_take()?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn borrow_after_take() -> Result<()> {
+        let cell = SharedCell::from(0);
+        let _ = cell.take().await;
+        assert!(cell.borrow().is_err());
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn take_after_borrow() -> Result<()> {
+        let cell = SharedCell::from(0);
+        let b = cell.borrow()?;
+        assert!(cell.try_take().is_err());
+        std::mem::drop(b);
+        cell.try_take()?;
+        Ok(())
+    }
+}