futex.rs

// A different version of chapter 8-01 on futex.
//! https://github.com/m-ou-se/rust-atomics-and-locks/blob/main/examples/ch8-01-futex.rs
//!
//! This version demonstrates how we can create an anonymous
//! memory-backed file (through memfd_create syscall), map it into the
//! program's memory (through mmap, memmap2 crate) and store a futex
//! on it. The futex is awakened later after the thread resumes from sleep.
//!
//! Futexes can be useful for having a finely grained concurrency
//! control on memory.

#![cfg(target_os = "linux")]

use memmap2::MmapMut;
use rustix::thread::{futex, FutexFlags, FutexOperation};
use std::mem::size_of;
use std::os::fd::AsRawFd;
use std::ptr::{null_mut, slice_from_raw_parts};
use std::sync::atomic::{AtomicU32, Ordering};
use std::thread;
use std::time::Duration;

const STD_OFFSET: isize = size_of::<AtomicU32>() as isize;

pub trait ToMemmap {
    fn to_mem(&self) -> MmapMut;
}

impl<T> ToMemmap for T
where
    T: AsRawFd,
{
    fn to_mem(&self) -> MmapMut {
        unsafe { memmap2::MmapOptions::new().map_mut(self) }.unwrap()
    }
}

pub unsafe fn wait(mut mmapped_futex: MmapMut, expected: u32) {

    let atomic = mmapped_futex
        .as_mut_ptr()
        .byte_offset(STD_OFFSET)
        .cast::<AtomicU32>()
        .as_mut()
        .unwrap();
  
    loop {

        if atomic.load(Ordering::Relaxed) == expected {

            let res = futex(
                atomic.as_ptr(),
                FutexOperation::Wait,
                FutexFlags::empty(),
                expected,
                null_mut(),
                null_mut(),
                0,
            )
            .unwrap();

            break;
        }
    }
}

pub unsafe fn wake_one(mut mmapped: MmapMut) {

    // retrieve futex from shared memory
    let atomic = mmapped
        .as_mut_ptr()
        .byte_offset(STD_OFFSET)
        .cast::<AtomicU32>()
        .as_mut()
        .unwrap();
  
    atomic.store(20, Ordering::Relaxed);

    futex(
        atomic.as_ptr(),
        FutexOperation::Wake,
        FutexFlags::empty(),
        1, // number of waiters to wake up, generally one or all
        null_mut(),
        null_mut(),
        0,
    )
    .unwrap();

}

fn main() {
    let file = memfile::CreateOptions::new().create("futexes").unwrap();
    file.set_len(1024).unwrap();
    let mmapped_futex = unsafe { memmap2::MmapOptions::new().map_mut(&file) }.unwrap();
    let wait_value = 10;

    unsafe {
        let mut mmaped = file.to_mem();
        let mem_ptr = mmaped.as_mut_ptr().byte_offset(STD_OFFSET);
        let atomic = mem_ptr.cast::<AtomicU32>().as_mut().unwrap();
        atomic.store(wait_value, Ordering::Relaxed);
        mmaped.flush().unwrap();
    };

    let slc = slice_from_raw_parts(mmapped_futex.as_ptr(), 50);
    let a = unsafe { slc.as_ref() }.unwrap();
    println!("before thread {:?}", a);

    thread::scope(|s| unsafe {
        s.spawn(|| {
            thread::sleep(Duration::from_secs(3));
            wake_one(file.to_mem());
        });

        println!("Waiting...");
        wait(file.to_mem(), wait_value);
        println!("Done!");
    });
}