Search code examples
multithreadingrustconcurrencyautomatic-ref-counting

Lock-free concurrent field usage


Say I have the following struct:

struct MyStruct {
   field1: Box<MyField>,
   field2: Box<MyField>,
}

and I'd like to be able to do the following:

  1. Use the fields in separate threads in a mutable fashion.
  2. Wait until complete
  3. Then start using the struct as a whole.

In theory, a reference counted type is sufficient (i.e. no locks).

To illustrate, here's some broken code that demonstrates the problem:

// Function to call when everything's done.
fn all_done(s: MyStruct) {
    ...
}

// Holds MyStruct until all references are dropped and then calls `all_done`
struct Guard(Option<MyStruct>);

impl Guard {
    fn new(s: MyStruct) -> Self {
        Self(Some(s))
    }
}

impl Drop for Guard {
    fn drop(&mut self) {
        all_done(self.0.take().unwrap());
    }
}

fn main() {
    let mut s: MyStruct = ...;
    let f1 = &mut s.field1;
    let f2 = &mut s.field2;
    let s = Arc::new(Guard::new(s));

    let s_copy = s.clone();
    thread::spawn(move || {
        do_some_mutating_with(f1);
        drop(s_copy);
    });

    thread::spawn(move || {
        do_some_mutating_with(f2);
        drop(s);
    });
}

In this contrived example, a few things are worth noting:

  1. The references don't outlive the value. This is ensured because the threads hold an Arc instance for that value.
  2. No part of the struct is used concurrently in multiple threads.

The compiler complains because:

  1. Arc::new(Guard::new(s)) moves something that's borrowed. On the other hand, if we created the thing in the Arc from the get-go, we'd have a different problem: "cannot borrow data in an Arc as mutable". Certainly, the references shouldn't be to the value on the stack as they are above.
  2. It's not convinced the references live long enough.

Is there any way to restructure the code to accomplish what it's trying to do? Absent that, are there any libraries that enable something equivalent?


Solution

  • The answer by @Jmb does answer the question, but is limited to spawning new threads. The original question did spawn new threads, but more for explanatory purposes. The goal was to work in async Rust and other contexts in which doing that each time is not an option.

    Turns out the key to making this work is "interior mutability", which replaces compile-time guarantees for runtime checks. Instead of references, I can use something like crossbeam::atomic::AtomicCell.

    Here's a rough sketch of a solution:

    use crossbeam::atomic::AtomicCell;
    use std::{sync::Arc, thread};
    
    struct MyStruct {
        field1: AtomicCell<Option<Box<i32>>>,
        field2: AtomicCell<Option<Box<i32>>>,
    }
    
    fn all_done(s: MyStruct) {
        println!(
            "Done: field1={}, field2={}",
            s.field1.swap(None).unwrap(),
            s.field2.swap(None).unwrap(),
        );
    }
    
    struct Guard(Option<MyStruct>);
    
    impl Guard {
        fn new(s: MyStruct) -> Self {
            Self(Some(s))
        }
    }
    
    impl Drop for Guard {
        fn drop(&mut self) {
            all_done(self.0.take().unwrap());
        }
    }
    
    fn main() {
        let s = MyStruct {
            field1: AtomicCell::new(Some(Box::new(0))),
            field2: AtomicCell::new(Some(Box::new(0))),
        };
        let s = Arc::new(Guard::new(s));
    
        let s_copy = s.clone();
        let jh1 = thread::spawn(move || {
            let mut f1 = s_copy.0.as_ref().unwrap().field1.swap(None).unwrap();
            *f1 = 1;
            s_copy.0.as_ref().unwrap().field1.swap(Some(f1));
        });
    
        let jh2 = thread::spawn(move || {
            let mut f2 = s.0.as_ref().unwrap().field2.swap(None).unwrap();
            *f2 = 2;
            s.0.as_ref().unwrap().field2.swap(Some(f2));
        });
    
        jh1.join().unwrap();
        jh2.join().unwrap();
    }
    

    Rust Playground link