Avoid locking on single threaded async

Question

Let's say I have the following struct:

struct Time {
    hour: u8,
    minute: u8
}

impl Time {
    pub fn set_time(&mut self, hour: u8, minute: u8) {
        self.hour = hour;
        self.minute = minute;
    }
}

On a multithreaded program, having a mutable reference to it shared across multiple threads could cause race conditions, but on a single-threaded one that can't happen (there is no way for the task to yield inside set_time).

Is there any way to avoid having to use locks in such a situation?

Here is an example where two tasks are running on a single thread and could share the mutable reference without a problem:

use tokio::join;

struct Time {
  hour: u8,
  minute: u8
}

impl Time {
  pub fn set_time(&mut self, hour: u8, minute: u8) {
      self.hour = hour;
      self.minute = minute;
  }
}

fn main() {
  let mut runtime_builder = tokio::runtime::Builder::new_current_thread();
  runtime_builder.enable_time();
  let runtime = runtime_builder.build().unwrap();
  runtime.block_on(async_main());
}

async fn async_main() {
  let mut time = Time {hour: 0, minute: 0};

  join!(
    task_1(&mut time),
    task_2(&mut time) // <- Rust wont allow this
  );
}

async fn task_1(time: &mut Time) {
  loop {
    // Do something
    tokio::task::yield_now();
  }
}

async fn task_2(time: &mut Time) {
  loop {
    // Do something
    tokio::task::yield_now();
  }
}

Answer 1

Mutable reference aliasing is never allowed by the borrow checker. This is one of conditions that the borrow checker enforces to guarantee memory safety at compile time.

This is the reason why the following code doesn't compile. There is 2 mutable references at the same time.

  join!(
    task_1(&mut time), // <- first mutable reference
    task_2(&mut time) // <- second mutable reference ERROR
  );

This rule is not specific to multithreaded context.

In the other hand, immutable reference aliasing is allowed by the borrow checker. but the following code would not compile because we can't mutate a field behind a shared reference.

use tokio::join;

struct Time {
  hour: u8,
  minute: u8
}

impl Time {
  pub fn set_time(&self, hour: u8, minute: u8) {
      self.hour = hour; // Error
      self.minute = minute; // Error
  }
}

[...]

async fn async_main() {
  let mut time = Time {hour: 0, minute: 0};

  join!(
    task_1(&time),
    task_2(&time)
  );
}

[...]

}

To have both aliasing and mutation capabilities, you need to use the interior mutability pattern.

Interior mutability:
"A type has interior mutability if its internal state can be changed through a shared reference to it. This goes against the usual requirement that the value pointed to by a shared reference is not mutated." (see rust reference)

There is several types which implement interior mutability pattern:

• Cell
• RefCell
• Atomics kind
• Mutex
• RwLock

with single threaded runtime, you can use either RefCell or Cell

With RefCell<T>, The rules still applies but at runtime instead of compile time. With a single threaded runtime, RefCell allows us to mutate shared reference in separate task, since there is no parallelism involved.

use std::sync::Arc;
use std::cell::RefCell;
use tokio::join;

#[derive(Debug)]
struct Time {
    hour: RefCell<u8>,
    minute: RefCell<u8>,
}

impl Time {
    pub fn set_time(&self, hour: u8, minute: u8) {
        *self.hour.borrow_mut() = hour; 
        *self.minute.borrow_mut() = minute;
    }
}

async fn task_1(time: &Time) {
    time.set_time(11, 54);
    println!("Task 1: {:?}", time);
    tokio::task::yield_now().await;
}

async fn task_2(time: &Time) {
    time.set_time(8, 12);
    println!("Task 2: {:?}", time);
    tokio::task::yield_now().await;
}

fn main() {
    let mut runtime_builder = tokio::runtime::Builder::new_current_thread();
    runtime_builder.enable_time();
    let runtime = runtime_builder.build().unwrap();
    runtime.block_on(async {
        let time = Time {
            hour: RefCell::new(0),
            minute: RefCell::new(0),
        };

        let _ = join!(task_1(&time), task_2(&time));
    });
}

[based on trentcl suggestion]

If you want to avoid the cost of runtime check, you can use Cell<T>. The API is pretty convenient when T is Copy.

use std::cell::Cell;
use tokio::join;

#[derive(Debug)]
struct Time {
    hour: Cell<u8>,
    minute: Cell<u8>,
}

impl Time {
    pub fn set_time(&self, hour: u8, minute: u8) {
        self.hour.replace(hour); 
        self.minute.replace(minute);
    }
}

async fn task_1(time: &Time) {
    time.set_time(11, 54);
    println!("Task 1: {:?}", time);
    tokio::task::yield_now().await;
}

async fn task_2(time: &Time) {
    time.set_time(8, 12);
    println!("Task 2: {:?}", time);
    tokio::task::yield_now().await;
}

fn main() {
    let mut runtime_builder = tokio::runtime::Builder::new_current_thread();
    runtime_builder.enable_time();
    let runtime = runtime_builder.build().unwrap();
    runtime.block_on(async {
        let time = Time {
            hour: Cell::new(0),
            minute: Cell::new(0),
        };

        let _ = join!(task_1(&time), task_2(&time));
    });
}

With multithreaded runtime, you can use atomics kind, Mutex, RwLock, and channel messaging.

If we take the specific case of AtomicU8, this type provides interior mutability and its store method is lock-free (at least on x86).

by using AtomicU8, we can aliased and mutate our Time struct with no data race, and no blocking.

AtomicU8 is Sync and Send, but we need to satisfy 'static bound for tokio::spawn, so taking a shared ref is not an option. We need to wrap the structure into an Arc.

use std::sync::atomic::{AtomicU8, Ordering};
use tokio::join;
use std::sync::Arc;

#[derive(Debug)]
struct Time {
    hour: AtomicU8,
    minute: AtomicU8,
}

impl Time {
    pub fn set_time(&self, hour: u8, minute: u8) {
        self.hour.store(hour, Ordering::SeqCst); // <-- mutation OK
        self.minute.store(minute, Ordering::SeqCst); // <-- mutation OK
    }
}

async fn task_1(time: Arc<Time>) {
    time.set_time(11, 54);
    println!("Task 1: {:?}", time);
    tokio::task::yield_now().await;
    
}

async fn task_2(time: Arc<Time>) {
    time.set_time(8, 12);
    println!("Task 2: {:?}", time);
    tokio::task::yield_now().await;
}

fn main() {
    let mut runtime_builder = tokio::runtime::Builder::new_multi_thread();
    runtime_builder.enable_time();
    let runtime = runtime_builder.build().unwrap();
    runtime.block_on(async {
        let time = Arc::new(Time {
            hour: AtomicU8::new(0),
            minute: AtomicU8::new(0),
        });

        let h1 = tokio::spawn(task_1(time.clone()));
        let h2 = tokio::spawn(task_2(time));
        

        let _ = join!(
            h1, // <-- aliasing Ok
            h2 // <-- aliasing Ok
        );
    });
}