Thread-safe way to either take or borrow

I came up with this using lazycell:

use std::sync::Mutex;

use lazycell::AtomicLazyCell;

#[derive(Eq, PartialEq, Debug)]
pub enum Error {
    Taken,
    TakenOrBorrowed,
}

pub struct Container<T> {
    initial: Mutex<Option<T>>,
    borrowed: AtomicLazyCell<T>,
}

impl<T> Container<T> {
    pub fn new(value: T) -> Self {
        Self {
            initial: Mutex::new(Some(value)),
            borrowed: AtomicLazyCell::new(),
        }
    }

    pub fn try_take(&self) -> Result<T, Error> {
        let mut lock = self.initial.lock().unwrap();
        lock.take().ok_or(Error::TakenOrBorrowed)
    }

    pub fn try_borrow(&self) -> Result<&T, Error> {
        {
            let mut lock = self.initial.lock().unwrap();
            if let Some(value) = lock.take() {
                if self.borrowed.fill(value).is_err() {
                    panic!("Tried to store borrowed value twice");
                }
            }
        }

        self.borrowed.borrow().ok_or(Error::Taken)
    }
}

fn assert_sync<T: Sync>(_: &T) {}

#[derive(Eq, PartialEq, Debug)]
struct MyValue(i32);

#[test]
fn can_take() {
    let c = Container::new(MyValue(42));
    assert_sync(&c);

    assert_eq!(Ok(MyValue(42)), c.try_take());
    assert_eq!(Err(Error::TakenOrBorrowed), c.try_take());
    assert_eq!(Err(Error::Taken), c.try_borrow());
}

#[test]
fn can_borrow() {
    let c = Container::new(MyValue(42));

    assert_eq!(Ok(&MyValue(42)), c.try_borrow());
    assert_eq!(Ok(&MyValue(42)), c.try_borrow());
    assert_eq!(Err(Error::TakenOrBorrowed), c.try_take());
}
3 Likes

nice, you did it cleanly and without using unsafe!

1 Like

This seems very similar to my atomic-take crate. It uses an atomic and a bit of unsafe, and this could be implemented using the same strategy.

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.