Sharing State in Rust Closures

Obligatory photo to capture attention in thumbnail

When I began using futures in Rust I faced strange issues regarding mutability of state. For instance, the stream type is a future that never resolves, but when using it I still needed results from the closures inside. Closures are a very common and powerful design pattern, but unless you can describe your entire application as a functional chain of monadic futures (in which case why are you even using Rust?..) it is likely you have needed to modify some state variable inside a closure and still have access to it in an outer scope. This article walks through the basic steps to explain this problem and how it’s solvable, both in single and multithreaded environments.

Starting from basics, we want both print statements in the code below to say that c=1. This would mean that the closure was able to increment ‘c’, and access it afterward in a print statement.

fn main() {
    let mut c = 0;
    let mut closure = || {
        c += 1;
        println!("in the closure, c={}", c);
    };
    closure();
    println!("out of the closure, c={}", c);
}

Alas, the borrow checker halts compilation with a very reasonable complaint. The closure borrows ‘c’ as mutable until it leaves scope at the end of main. Therefore the second println cannot borrow ‘c’ at all because a variable can only have one reference if it is mutable.

error[E0502]: cannot borrow `c` as immutable because it is also borrowed as mutable
  --> test.rs:23:19
   |
14 |     let mut closure = || {
   |                       -- mutable borrow occurs here
15 |         c += 1;
   |         - previous borrow occurs due to use of `c` in closure
...
23 |     println!("{}",c);
   |                   ^ immutable borrow occurs here
24 | }
   | - mutable borrow ends here

Understanding how closures are implemented in Rust has drastically helped me debug and use them. It’s also pretty cool so I’ll show it very briefly with an example. Closures are structs that implement a trait that makes them callable (here’s a good article). The above code would be roughly rewritten as follows by the Rust compiler.

struct ClosureType<'a> {
    c: &'a mut i32,
}
impl<'a> ClosureType<'a> {
    fn call(&mut self) {
        *self.c += 1;
        println!("in the closure, c={}",self.c);
    }
}
fn main() {
    let mut c = 0;
    let mut closure = ClosureType { c:&mut c };
    closure.call();
    println!("out of the closure, c={}", c);
}

The lifetimes make this example a little ugly so ignore them if you hate them and haven’t learned to love them, or haven’t even learned to hate them… The important thing to note here is that a closure is a struct! Specifically, a struct with a field for each variable it captures, which in this example is ‘c’. You can see in the main function that c is borrowed mutably. Now the compiler error is much easier to understand. ClosureType borrows mutably so println can’t use it.

Interestingly enough, a “move” on the closure would actually compile. A move causes the closure to take ownership of the variables it captures instead of just borrowing them. In a similar format to above, the closure struct no longer needs a litfetime because it now owns ‘c’.

struct ClosureType {
    c: i32,
}
impl ClosureType {
    fn call(&mut self) {
        self.c += 1;
        println!("in the closure, c={}", self.c);
    }
}
fn main() {
    let mut c = 0;
    let mut closure = ClosureType { c:c };
    closure.call();
    println!("out of the closure, c={}", c);
}

But do you notice the problem here? The program prints

in the closure, c=1
out of the closure, c=0

That’s because c is copied into the closure while the original still exists and is unaffected by the closure operation. If you change ‘c’ to a type that does not implement the Copy trait, like a vector, then the borrow checker will throw an error on compilation. In that situation the type is not copyable so the closure takes ownership of the one an only variable ‘c’, and the second println is left referencing a moved value. See how that was easier to understand when viewing closures as their underlying structs?

Now for a real solution to the problem of referencing a value after a closure has changed it. The fact is that in order to achieve this, we’ll need multiple references to the same value, one of which is mutable. That breaks Rust’s single mutable referencing rule, but luckily Rust is equpped for these situations and has built in primitives to handle common cases. We’ll use RefCell to allow ‘c’ to be referenced while seemingly simultaneously being borrowed mutably. Surprisingly this is not as dangerous as it sounds. The link provided is a great dive into RefCell, but in short, it is a way to enforce that a value is not being accessed unsafely at runtime instead of at compile time where Rust normally does the checking. Trippy. Here it is in action.

fn main() {
    let c = RefCell::new(0);
    let mut closure = || {
        *c.borrow_mut() += 1;
        println!("in the closure, {}", c.borrow());
    };
    closure();
    println!("out of the closure, {}", c.borrow());
}

That prints c=1 both times! It works! This is a way to mutate a captured value in a closure and see the change outside! Now I actually recommend skimming through the whole section of the Rust book on smart pointers. It’s something I didn’t pay attention to on my first read but if I had it would have saved me a lot of time and searching through StackOverflow.

If you have made it this far in the article, I’m going to guess you faced a problem similar enough that this has been relevant information for you. On the other hand, if you just enjoy reading articles on Rust, here’s a real world use case of this design pattern.

I originally ran into this issue when implementing a game server. The server handles client connections asynchronously using the futures pattern with the tokio runtime. Futures are a closure, and there is one future for every client connection in order to easily distribute tasks among multiple threads. When a client sends a player command, the closure must update the corresponding entity in a shared entity database among all futures.

The following is a basic TCP server that on connecting to, will add a new entity to the “game” state. You will notice the state being modified is enclosed in an Arc and a Mutex instead of a RefCell like we used above. In essence a Mutex is the multithreading version of RefCell. Arc is a way to safely access a Mutex among multiple threads (similar to Rc for a single thread). Now I’d be surprised if that made any sense so again I’d recommend the book which does a good job of discerning between Rc/RefCell and Arc/Mutex and why you sometimes need both.

To test this code just create a Cargo.toml file that declares tokio = “0.1” as a dependency. I used this curl command to simulate connections.

curl 127.0.0.1:12345 --resolve "127.0.0.1:12345"