DEV Community: Raine-Yang-UofT

Resolving Memory Fragmentation for Linkedlist Heap Allocator

Raine-Yang-UofT — Sun, 11 Feb 2024 21:37:07 +0000

This blog is based on the chapter Allocator Design of tutorial Writing an OS in Rust (https://os.phil-opp.com/).

Heap Allocator

A heap is a memory region for a program to store dynamically-sized data at runtime. The data stored in heap can have varying lifecycles. A heap allocator is responsible for allocating a heap region of appropriate size to store new data, and deallocate the region when the data is released. In many languages, the deallocation of heap is handled by garbage collection system. However, in Rust, with the ownership and lifecycle system, the deallocation of heap is called directly by the compiler once the object ends its lifecycle.

In our OS kernel, we want to implement a heap allocator that can be used by the Rust compiler as a global allocator, which requires the implementation of GlobalAlloc trait:

pub unsafe trait GlobalAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8;
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { ... }
    unsafe fn realloc(
        &self,
        ptr: *mut u8,
        layout: Layout,
        new_size: usize
    ) -> *mut u8 { ... }
}

We need to implement alloc and dealloc methods, and the other two methods have default implementations. For the alloc method, we need to find an available heap region with size and memory alignment described in layout parameter, or returns a null pointer if the heap is out of memory. For the dealloc method, we need to mark the region to be freed as unused so it can be allocated again.

Linkedlist Heap Allocator Implementation

The heap allocator should be able to keep track of all the unused memory regions in the heap. Since the number of unused regions is dynamic at runtime, we need to store the list of regions (also known as free list) in a data structure with dynamic size. However, we cannot use collection types like Vectors or Hashmaps since their implementations are based on a heap, which leaves as the choice to use linkedlist.

A linkedlist allocator stores all the free regions in a linkedlist, whose each node stores the starting address and size of the region, and a pointer to the next node. When we need to allocate a region, we traverse the linkedlist to find a node that contains a region that is large enough, and remove and return it. During deallocation, we append the released region as a new node into the free list.

The implementation of a node:

struct ListNode {
    size: usize,
    next: Option<&'static mut ListNode>,
}

impl ListNode {
    const fn new(size: usize) -> Self {
        ListNode { size, next: None }
    }

    fn start_addr(&self) -> usize {
        self as *const Self as usize
    }

    fn end_addr(&self) -> usize {
        self.start_addr() + self.size
    }
}

The implementation of LinkedlistAllocator:

pub struct LinkedListAllocator {
    head: ListNode,
}

impl LinkedListAllocator {
    /// Creates an empty LinkedListAllocator.
    pub const fn new() -> Self {
        Self {
            head: ListNode::new(0),
        }
    }

    /// Initialize the allocator with the given heap bounds.
    ///
    /// This function is unsafe because the caller must guarantee that the given
    /// heap bounds are valid and that the heap is unused. This method must be
    /// called only once.
    pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
        self.add_free_region(heap_start, heap_size);
    }
}

Note that in this code we make the head of a list an empty ListNode, and head.next represent the first valid node in free list. The init method creates one node that records the whole heap region. The implementation of add_free_region is as below:

impl LinkedListAllocator {
    /// Adds the given memory region to the front of the list.
    unsafe fn add_free_region(&mut self, addr: usize, size: usize) {
        // ensure that the freed region is capable of holding ListNode
        assert_eq!(align_up(addr, mem::align_of::<ListNode>()), addr);
        assert!(size >= mem::size_of::<ListNode>());

        // create a new list node and append it at the start of the list
        let mut node = ListNode::new(size);
        node.next = self.head.next.take();
        let node_ptr = addr as *mut ListNode;
        node_ptr.write(node);
        self.head.next = Some(&mut *node_ptr)
    }
}

In the add_free_region method, we first check that whether the memory region is properly aligned and is greater than the size of a ListNode (to keep track of the region with a ListNode, we need to first ensure the region is large enough the store the node!) Then we append the node to the head of the list (self.head.next). Note that it is only during node_ptr.write(node) the node is written to the memory and has a 'static lifecycle. Before that the node is a local variable stored on stack. We will come to this important distinction later.

This diagram shows the process of appending the node to the head of the list:

Next, we need a method alloc_from_region that checks whether a list node is enough to allocate a memory region with a given size. We first need to check that the node's end address is greater than the end of allocation (the region recorded by the node is larger than the region). In addition, since we want to keep track of the remaining region of the node after allocation, we need to make sure the remaining region after allocation is large enough to store another ListNode.

impl LinkedListAllocator {
    /// Try to use the given region for an allocation with given size and
    /// alignment.
    ///
    /// Returns the allocation start address on success.
    fn alloc_from_region(region: &ListNode, size: usize, align: usize)
        -> Result<usize, ()>
    {
        let alloc_start = align_up(region.start_addr(), align);
        let alloc_end = alloc_start.checked_add(size).ok_or(())?;

        if alloc_end > region.end_addr() {
            // region too small
            return Err(());
        }

        let excess_size = region.end_addr() - alloc_end;
        if excess_size > 0 && excess_size < mem::size_of::<ListNode>() {
            // rest of region too small to hold a ListNode (required because the
            // allocation splits the region in a used and a free part)
            return Err(());
        }

        // region suitable for allocation
        Ok(alloc_start)
    }
}

The find_region method is used to traverse the linkedlist to find a memory region that is large enough for the allocation. If such a node is found, it returns the node and removes it from the linkedlist. Otherwise, None is returned.

impl LinkedListAllocator {
    /// Looks for a free region with the given size and alignment and removes
    /// it from the list.
    ///
    /// Returns a tuple of the list node and the start address of the allocation.
    fn find_region(&mut self, size: usize, align: usize)
        -> Option<(&'static mut ListNode, usize)>
    {
        // reference to current list node, updated for each iteration
        let mut current = &mut self.head;
        // look for a large enough memory region in linked list
        while let Some(ref mut region) = current.next {
            if let Ok(alloc_start) = Self::alloc_from_region(&region, size, align) {
                // region suitable for allocation -> remove node from list
                let next = region.next.take();
                let ret = Some((current.next.take().unwrap(), alloc_start));
                current.next = next;
                return ret;
            } else {
                // region not suitable -> continue with next region
                current = current.next.as_mut().unwrap();
            }
        }

        // no suitable region found
        None
    }
}

With the methods we implemented above, we are able to implement the alloc and dealloc method for GlobalAlloc trait:

unsafe impl GlobalAlloc for Locked<LinkedListAllocator> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // perform layout adjustments
        let (size, align) = LinkedListAllocator::size_align(layout);
        let mut allocator = self.lock();

        if let Some((region, alloc_start)) = allocator.find_region(size, align) {
            let alloc_end = alloc_start.checked_add(size).expect("overflow");
            let excess_size = region.end_addr() - alloc_end;
            if excess_size > 0 {
                allocator.add_free_region(alloc_end, excess_size);
            }
            alloc_start as *mut u8
        } else {
            ptr::null_mut()
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        // perform layout adjustments
        let (size, _) = LinkedListAllocator::size_align(layout);

        self.lock().add_free_region(ptr as usize, size)
    }
}

In the alloc method, we first find a memory region to allocate using find_region method. Then, if there is remaining memory on the region after allocation, we store the excess region back to the linkedlist. In the dealloc method, we only need to append the freed region into our linkedlist.

Memory Fragmentation Issue

The above implementation of linkedlist allocator works in general. However, it has a serious problem: after repeated allocations and deallocations, the memory will be fragmented into smaller and smaller list nodes and it will be harder to find regions to allocate large objects. The diagram below shows how this problem happens:

Assume we have three regions being allocated. When these regions are freed, new list nodes are created to keep track of each freed regions during dealloc. However, even through these freed regions are consecutive, they are tracked by different nodes, and in out current implementation we do not have any ways to merge the nodes together. As a result, in repeated allocations, the memory will be fragmented into smaller and smaller nodes, which would compromises the efficiency of traversing the linkedlist and makes it unable to allocate large regions even if there is enough space.

Solution to the Memory Fragmentation Issue

In order to resolve the issue above, we need to find a way to merge consecutive memory regions to a single node. To implement this, we first need to sort the nodes by their starting address in add_free_region method instead of appending new nodes directly at the head.

impl LinkedListAllocator {
    // apppend a free region with given starting address and size to the linkedlist
    unsafe fn add_free_region(&mut self, addr: usize, size: usize) {
        // ensure the memory is aligned 
        assert_eq!(align_up(addr, mem::align_of::<ListNode>()), addr);
        // ensure the memory is large enough to hold the linkedlist
        assert!(size >= mem::size_of::<ListNode>());

        // the new node
        let mut new_node = ListNode::new(size); // the node is now on stack
        // find the proper location of node in linkedlist

        // add node to the head of list
        let mut current = &mut self.head;

        if current.next.is_none() ||
            addr < current.next.as_ref().unwrap().start_addr() {    // the head is None or new node has smaller start address than first node
            new_node.next = current.next.take();
        } else {
            // traverse the linkedlist
            while let Some(ref mut next) = current.next {
                if next.start_addr() >= addr {  
                    break;
                }
                current = current.next.as_mut().unwrap()
            }
            // assign the node's next node
            // we can only assign its previous node when the node is written to memory (having a 'static lifecycle)
            new_node.next = current.next.take();
        }


        let node_ptr = addr as *mut ListNode;
        node_ptr.write(new_node);   // write the new node in memory
        current.next = Some(&mut *node_ptr);    // connect the node with the previous node

    }
}

In general, the code is similar to the method in find_region that traverses the linkedlist. However, there are a few points to note about:

We need to modify self.head.next (the first node) if self.head.next is None or if the start address of new node is smaller than the start address of current self.head.next.
When we compare the address of nodes in the linkedlist with the newly appended node, we have to use addr instead of new_node.start_address(), this is because before new_node is written to memory in node_ptr.write(new_node), new_node is a local variable stored on stack. Comparing its start address on stack with addresses in the heap will very likely lead to confusing results.
The next attribute of ListNode need to ensure that the reference has a 'static lifecycle, meaning that the ListNode being refered to must be stored in heap. We can assign the next of new_node in the while let loop (since it points to an existing node in heap), but we have to assign the next of the node before new_node (current.next) only after new_node is being written to heap. (We need to thank Rust's lifecycle system on this matter, otherwise it would be very easy to accidentally assign the reference to a local variable, which gets cleared after the method).

After maintaining the nodes sorted by start address, we are able to merge consecutive memory by checking whether the end address of one node is equal to the start address of its next node. If so, we creates a single node in the linkedlist that overrides the two nodes.

impl LinkedlistAllocator {
// merge consecutive free memory regions into larger regions
    fn merge_region(&mut self) {
        let mut current = &mut self.head;

        while let Some(ref mut node) = current.next {
            // get the start and end address of current node (prevent borrowing issue)
            let start_addr = node.start_addr();
            let end_addr = node.end_addr();
            // check whether it is adjacent to the next node
            if let Some(ref mut next) = node.next {
                if end_addr == next.start_addr() {
                    // combine two memory regions
                    node.size = next.end_addr() - start_addr;
                    node.next = next.next.take();
                }
            }

            current = current.next.as_mut().unwrap();
        }
    }
}

Finally, we call merge_region method in dealloc to merge consecutive regions after every deallocation:

unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
    let (size, _) = LinkedListAllocator::size_align(layout);
    // add the freed region to free list
    self.lock().add_free_region(ptr as usize, size);;
    // merge unused regions
    self.lock().merge_region();
}

The merging of memory region:

Rust Learning Note: Writing a Simple Tokio

Raine-Yang-UofT — Sat, 06 Jan 2024 16:52:31 +0000

This article is a summary of Chapter 6.8 of Rust Course (course.rs/)

Future Trait

In async/await programming, an async function will not be executed until it is invoked by .await. This delayed execution is implemented through Future trait, and the return value of any async function is a trait object of Future. Here is the definition of Future trait:

pub trait Future {
    type Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output>;
}

Future trait defines a series of asynchronous computations. poll() method is used to progress in the computation. It has an associated type Ouput, which stands for the return type when Future completes the computation. poll() returns an Poll enum type, including Poll::Ready, indicating that the computation is finished, and Poll::Pending, indicating that the computation has not been completed.

Here we have a struct Delay that implements Future. Its poll method waits for a certain time interval before returning Poll::Ready("done").

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

struct Delay {
    when: Instant
}

impl Future for Delay {
    type Output = &'static str;
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<&'static str> {
        if Instant::now() >= self.when {
            println!("Hello world");
            Poll::Ready("done");
        } else {
            // ignore this code by now
            cx.waker().wake_by_ref();
            Poll::Pending
        }
    }
}

#[tokio::main]
async fn main() {
    let when = Instant::now() + Duration::from_millis(10);
    let future = Delay { when };
    let out = future.await;
    assert_eq!(out, "done");
}

Executor

In the code above, we need to create an async main function from tokio to run the Future. As we mentioned above, only .await can invoke a Future. However, .await can only be declared in async functions, whose return types themselves are Future. As a result, we need to implement an Executor to process the Future in the outmost layer. Here is an example code:

use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use futures::task;

struct MiniTokio {
    tasks: VecDeque<Task>
}

type Task = Pin<Box<dyn Future<Output = ()> + Send>>;

impl MiniTokio {
    fn new() -> MiniTokio {
        MiniTokio {
            tasks: VecDeque::new()
        }
    }

    fn spawn<F>(&mut self, future: F)
    where F: Future<Output = ()> + Send + 'static {
        self.tasks.push_back(Box::pin(future));
    }

    fn run(&mut self) {
        let waker = task::noop_waker();
        let mut cx = Context::from_waker(&waker);
        while let Some(mut task) = self.tasks.pop_front() {
            if task.as_mut().poll(&mut cx).is_pending() {
                self.tasks.push_back(task);
            }
        }
    }
}

In this code, the run method continuously retrieves Future from the queue calls the poll() method. If the Future returns Poll::Pending, it puts the Future back to the queue. However, this method is inefficient. The Executor would poll all Future continuously, and very likely most polling are meaningless as the Future simply returns Poll::Pending.

A more ideal solution is that Future could notify the Executor when to poll. Only when the Future is able to continue its processing (such as when it finally gets a result from another calculation or I/O) would it be polled again. This approach is implemented through Waker.

Waker

fn poll(self: Pin<&mut Self>, cx: &mut Context)
    -> Poll<Self::Output>;

In the poll() method, Waker is contained in parameter Context and can be retrieved by calling cx.waker(). The wake() method defined in Waker is used to notify the executor that the Future is ready to be polled again.

Here we implement Waker for Delay:

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};

struct Delay {
    when: Instant
}

impl Future for Delay {
    type Output = &'static str;
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<&'static str> {
        if Instant::now() >= self.when {
            println!("Hello world");
            Poll::Ready("done");
        } else {
            let waker = cx.waker().clone();
            let when = self.when;
            thread::spawn(move || {
                let now = Instant::now();
                if now < when {
                    thread::sleep(when - now);
                }
                waker.wake();
            });
            Poll::Pending
        }
    }
}

In this updated code, we rewrite the else branch in which Delay returns Poll::Pending. We spawn a new thread and sleeps for the given time interval, and calls waker.wake() when the thread wakes up. Note that when a Future returns Poll::Pending, wake() must be invoked, otherwise the Future will be suspended indefinitely and will never be polled again.

Processing wake()

We will now rewrite MiniTokio to allow it receives wake notification. We will use a message channel to store all Future waiting to be polled. Once a Future calls wake(), it will be added into the message channel. To implement this, we will decorate Future with a new struct Task, which includes the Future object and the Sender of message channel. Once wake() is invoked, the sender will send the Task into the message channel, which will be polled by the Executor.

Since the message receiver and sender may be in different threads, the message channel need to implement Sync and Send. However, the message channel in standard library does not implement Sync, so we need to use library crossbeam

// in Cargo.toml [dependencies]
crossbeam = "0.8"

struct MiniTokio {
    scheduled: channel::Receiver<Arc<Task>>,
    sender: channel::Sender<Arc<Task>>
}

impl MiniTokio {
    fn new() -> MiniTokio {
        let (sender, scheduled) = channel::unbounded();
        MiniTokio { scheduled, sender }
    }
}

struct Task {
    future: Mutex<Pin<Box<dyn Future<Output = ()> + Send>>>,
    executor: channel::Sender<Arc<Task>>
}

impl Task {
    fn schedule(self: &Arc<Self>) {
        self.executor.send(self.clone());
    }
}

In the code, MiniTokio holds the receiver (scheduled) and sender (sender) of the message channel. Task contains the Future trait object and the sender of message channel. Method schedule() is used to send a message to MiniTokio.

In order to convert Task in to a Waker, we can make it implement ArcWake trait in futures module. First we need to add futures to dependencies in Cargo.toml

futures = "0.3"

And then implement ArcWake for Task

impl ArcWake for Task {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self.schedule();
    }
}

When the Future calls waker.wake(), the schedule() method in Task will be invoked. After that, we can implement two other methods for Task: poll() for passing the waker to the poll() method in Future and calling poll() in Future, and spawn() for creating a new Task from a Future and sending the Task to MiniTokio for initial polling.

impl Task {
    fn poll(self: Arc<Self>) {
        let waker = task::waker(self.clone());
        let mut cx = Context::from_waker(&waker);
        let mut future = self.future.try_lock().unwrap();
        let _ = future.as_mut().poll(&mut cx);
    }

    fn spawn<F>(future: F, sender: &channel::Sender<Arc<Task>>)
    where F: Future<Output = ()> + Send + 'static {
        let task = Arc::new(Task {
            future: Mutex::new(Box::pin(future)),
            executor: sender.clone() 
        });

        let _ = sender.send(task);
    }
}

Finally, we need to implement the run() method for MiniTokio that retrieves Task from message channel and calls poll() method of the Task.

impl MiniTokio {
    fn run(&self) {
        while let Ok(task) = self.scheduled.recv() {
            task.poll();
        }
    }

    fn spawn<F>(&self, future: F)
    where F: Future<Output = ()> + Send + 'static {
        Task::spawn(future, &self.sender);
    }
}

Here is the complete code for MiniTokio:

use std::future::Future;
use futures::task::{self, ArcWake};
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use std::thread;
use crossbeam::channel;
use std::sync::{Arc, Mutex};

struct Delay {
    when: Instant
}

impl Future for Delay {
    type Output = &'static str;

    fn poll(self:Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<&'static str> {
        if Instant::now() >= self.when {
            println!("Hello world");
            Poll::Ready("done")
        } else {
            let waker = cx.waker().clone();
            let when = self.when;

            thread::spawn(move || {
                let now = Instant::now();

                if now < when {
                    thread::sleep(when - now);
                }

                waker.wake();
            });

            Poll::Pending
        }
    }
}


struct MiniTokio {
    scheduled: channel::Receiver<Arc<Task>>,
    sender: channel::Sender<Arc<Task>>
}

impl MiniTokio {
    fn run(&self) {
        while let Ok(task) = self.scheduled.recv() {
            task.poll();
        }
    }

    fn new() -> MiniTokio {
        let (sender, scheduled) = channel::unbounded();
        MiniTokio {scheduled, sender}
    }

    fn spawn<F>(&self, future: F)
    where F: Future<Output = ()> + Send + 'static {
        Task::spawn(future, &self.sender);
    }
}


struct Task {
    future: Mutex<Pin<Box<dyn Future<Output = ()> + Send>>>,
    executor: channel::Sender<Arc<Task>>
}

impl Task {
    fn schedule(self: &Arc<Self>) {
        self.executor.send(self.clone());
    }

    fn poll(self: Arc<Self>) {
        let waker = task::waker(self.clone());
        let mut cx = Context::from_waker(&waker);
        let mut future = self.future.try_lock().unwrap();
        let _ = future.as_mut().poll(&mut cx);
    }

    fn spawn<F>(future: F, sender: &channel::Sender<Arc<Task>>)
    where F: Future<Output = ()> + Send + 'static {
        let task = Arc::new(Task {
            future: Mutex::new(Box::pin(future)),
            executor: sender.clone()
        });

        let _ = sender.send(task);
    }
}

impl ArcWake for Task {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self.schedule();
    }
}

Rust Learning Note: Running Multiple Future Simultaneously

Raine-Yang-UofT — Tue, 02 Jan 2024 19:35:34 +0000

This article is a summary of Chapter 4.11.5 of Rust Course (course.rs/)

join!

.await does not allow us to run multiple Future simultaneously. We have to wait for one Future to finish processing before another. Sometimes, we want to process a few Future concurrently, which requires the use of join! macro:

use futures::join;

async fn enjoy_book_and_music() -> (Book, Music) {
    let book_fut = enjoy_book();
    let music_fut = enjoy_music();
    join!(book_fut, music_fut)
}

join! takes a tuple of two Future, and returns a tuple containing the results of the two Future after they are both completed. If we want to concurrently run multiple Future stored in an array, we can use **futures::future::join_all **method

use futures::future::join_all;

async fn async_task(id: u32) -> u32 {
    async_std::task::sleep(std::time::Duration::from_secs(1)).await;
    println!("Task {} completed", id);
    id
}

#[tokio::main]
async fn main() {
    let tasks = vec![
        async_task(1),
        async_task(2),
        async_task(3)
    ];

    let results = join_all(tasks).await;
    println!("All tasks completed");

    for result in results {
        println!("Result: {}", result);
    }
}

try_join!

join! is completed executing only when all the Future inside finish. If we want to terminal all Future if any Future throws an error, we can use try_join!

use futures::try_join;

async fn get_book() -> Result<Book, String> {Ok(Book)};
async fn get_music() -> Result<Music, String> {Ok(Music)};

async fn get_book_and_music() -> Result<(Book, Music), String> {
    let book_fut = get_book();
    let music_fut = get_music();
    try_join!(book_fut, music_fut)
}

Note that the Future in try_join! must have the same error type. If the functions return different errors, we can use map_err and err_into in futures::future::TryFutureExt for error conversion.

use futures::{
    future::TryFutureExt,
    try_join
};

async fn get_book() -> Result<Book, ()> {Ok(Book)};
async fn get_music() -> Result<Music, String> {Ok(Music)};

async fn get_book_and_music() -> Result<(Book, Music), String> {
    let book_fut = get_book().map_err(|()|, "unable to get book".to_string());
    let music_fut = get_music();
    try_join!(book_fut, music_fut)
}

select!

join! only returns when all Future complete. Sometimes, we want to wait for a few Future, and returns whenever one Future completes and does not wait for the rest of the Future. In this case, we can use select! macro.

use futures::{
    future::FutureExt,
    pin_mut,
    select
};

async fn task_one() {};
async fn task_two() {};

async fn race_tasks() {
    let t1 = task_one().fuse();
    let t2 = task_two().fuse();
    pin_mut!(t1, t2);
    select! {
        () = t1 => println!("task 1 finishes first"),
        () = t2 => println!("task 2 finishes first")
    }
}

In the code above, we use fuse() to implement FusedFuture for Future, which makes select unable to poll a Future once it's finished. pin_mut!() is used to implement Unpin trait, which allows select to obtain a mutable reference of Future. These two traits FusedFuture and Unpin are required for select to work properly.

select! also supports default and complete branch. complete branch is called when all Future and Stream are completed. It is commonly used together with loop. default branch is called when none of the Future or Stream are at Ready state.

use futures::future;
use futures::select;
pub fn main() {
    let mut a_fut = future::ready(4);
    let mut b_fut = future::ready(6);
    let mut total = 0;

    loop {
        select!{
            a = a_fut => total += a,
            b = b_fut => total += b,
            complete => break,
            default => panic!()
        };
    }
    assert_eq!(total, 10);
}

With Stream, we need to implement FusedStream with fuse(). When a Stream implements FusedStream, the Future it returns with next() will also implement FusedFuture.

use futures::{
    stream::{Stream, StreamExt, FusedStream},
    select
}

async fn add_two_streams(
    mut s1: impl Stream<Item = u8> + FusedStream + Unpin,
    mut s2: impl Stream<Item = u8> + FusedStream + Unpin
) -> u8 {
    let mut total = 0;

    loop {
        let item = select! {
            x = s1.next() => x,
            x = s2.next() => x,
            complete => break
        };
        if let Some(next_num) = item {
            total += next_num;
        }
    }
    total
}

Rust Learning Note: async/await and Stream

Raine-Yang-UofT — Tue, 02 Jan 2024 16:37:24 +0000

This article is a summary of Chapter 4.11.4 of Rust Course (course.rs/)

The Lifecycle of async

If an async function has reference type parameters, the parameter it refers to must life longer than the Future the function returns. For instance, the code below would throw an error since the variable x only lives inside function bad(), but the async function borrow_x, after being returned, lives in a larger scope than x.

use std::future::Future;
fn bad() -> impl Future<Output = u8> {
    let x = 5;
    borrow_x(&x)
}

async fn borrow_x(x: &u8) -> u8 {*x}

One solution to this problem is to put the variables the function refer to along with the function inside a single async block. In this way, the function and the variables it refers to always exist in the same scope, and thus have the same lifecycle.

use std::future::Future;

async fn borrow_x(x: &u8) -> u8 { *x }

fn good() -> impl Future<Output = u8> {
    async {
        let x = 5;
        borrow_x(&x).await
    }
}

Another solution is to use move keyword to transfer the ownership of the variable into the async function, which is similar to move in closures. However, this will no longer allow us to use the moved variable anywhere else.

fn move_block() -> impl Future<Output = ()> {
    let my_string = "foo".to_string();
    async move {
        println!("{my_string}");
    }
}

.await in Multithreading Executor

Variable in async blocks need to be passed among multiple threads (since async/await is a M:N threading model), so data types that do not implement Send and Sync traits, like Rc, RefCell, cannot be used in async blocks. Also, the Mutex in std::sync::Mutex is not safe to use in async/await, since it is possible that one .await function that acquires a lock is suspended before releasing the lock, and another function currently processed by the thread tries the acquire the lock, leading to a deadlock. We should use the Mutex in futures:🔒:Mutex instead.

Stream Processing

Stream trait is similar to Future trait, except that it can return multiple Future objects. Stream is similar to Iterator.

trait Stream {
    type Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>)
        -> Poll<Option<Self::Item>>;
}

One use of Stream is the Receiver of message channel. The Stream receives a Some(val) when the sender sends a message, and None when the channel is closed. In the code below, Stream is implicitly used in the receiver (rx). The method call rx.next() is invoking the Stream trait implementation for the receiver.

async fn send_recv() {
    const BUFFER_SIZE: usize = 10;
    let (mut tx, mut rx) = mpsc::channel::<i32>(BUFFER_SIZE);

    tx.send(1).await.unwrap();
    tx.send(2).await.unwrap();
    drop(tx);

    assert_eq!(Some(1), rx.next().await);
    assert_eq!(Some(2), rx.next().await);
    assert_eq!(None, rx.next().await);   
}

Similar to an Iterator, we can also iterate a Stream and use methods like map, filter, and fold. However, we cannot use a for loop to loop through a Stream. Instead, we can use while let for looping.

async fn sum_with_next(mut stream: Pin<&mut dyn Stream<Item = i32>>) -> i32 {
    use futures::stream::streamExt;
    let mut sum = 0;
    while let Some(item) = stream.next().await {
        sum += item;
    }
    sum
}

async fn sum_with_try_next(
    mut stream: Pin<&mut dyn Stream<Item = Result<i32, io::Error>>>
) -> Result<i32, io::Error> {
    use futures::stream::TryStreamExt;
    let mut sum = 0;
    while let Some(item) = stream.try_next().await? {
        sum += item;
    }
    Ok(sum)
}

However, the approach above processes only one value each time, and blocks the thread with await when waiting for the next value, which loses the meaning of concurrent programming. We can use for_each_concurrent or try_for_each_concurrent to process multiple values from Stream concurrently.

async fn jump_around(
    mut stream: Pin<&mut dyn Stream<Item = Result<u8, io::Error>>>
) -> Result<(), io::Error> {
    use futures::stream::TryStreamExt;
    const MAX_CONCURRENT_JUMPERS: usize = 100;

    stream.try_for_each_concurrent(MAX_CONCURRENT_JUMPERS, |num| async move {
        jump_n_times(num).await?;
        report_n_jumps(num).await?;
        Ok(())
    }).await?;

    Ok(())
}

In this example, try_for_each_concurrent method is used to apply an asynchronous closure to each item in the stream concurrently. Inside the closure, we defines two custome functions and uses ? to propagate error.

Rust Learning Note: Pin and Unpin

Raine-Yang-UofT — Sun, 31 Dec 2023 22:20:50 +0000

This article is a summary of Chapter 4.11.3 of Rust Course (course.rs/) and blog (folyd.com/blog/rust-pin-unpin/)

The Use of Pin

Pin is a smart pointer that encapsulates another pointer. It is represented as Pin

, in which P is a pointer pointing to an object of type T. Pin makes the object T unable to move its location in memory if T does not implement Unpin trait (that is, T is type !Unpin). If T implements Unpin, Pin has no effect on P at all.

Pin is mainly used in resolving self-reference type, a types whose some attribute is a reference to its other attributes, as the example below:

#[derive(Debug)]
struct Test {
    a: String,
    b: *const String
}

impl Test {
    fn new(txt: &str) -> Self {
        Test {
            a: String::from(txt),
            b: std::ptr::null()
        }
    }

    fn init(&mut self) {
        let self_ref: *const String = &self.a;
        self.b = self_ref;
    }

    fn a(&self) -> &str {
        &self.a
    }

    fn b(&self) -> &String {
        assert!(!self.b.is_null());
        unsafe { &*(self.b) }
    }
}

In this code, we create a raw pointer b and points it to the memory address of a in init() method. An error would occur if the memory address of a is moved yet b is still pointing to the old address, as in this example:

fn main() {
    let mut test1 = Test::new("test1");
    test1.init();
    let mut test2 = Test::new("test2");
    test2.init();

    println!("a: {} b: {}", test1.a(), test1.b());
    std::mem::swap(&mut test1, &mut test2);
    println!("a: {} b: {}", test2.a(), test2.b());
}

In the code, we expect the second println statement would print "a: test1, b: test1". However, the output is actually "a: test1, b: test2". This is because b is still pointing to the previous memory address which is now in test1.

Fig 1. The raw pointer pointing to a moved address. Image reproduced from (Rust Course)

Fix the Value on Stack

With Pin, we can solve the problem concerning moved data above. If an object does not implement Unpin trait, when it is put inside smart pointer Pin, we would be unable to get its mutable reference, and thus cannot invoke functions like swap. However, almost all Rust types automatically implement the Unpin trait, with only a few exceptions

PhantomPinned
Future in async/await

Since a composite data type would only automatically implement a trait only if all its attribute implement the trait, if we include at least one !Unpin attribute to the composite type, it would not implement Unpin. Normally, we include PhantomPinned into a type to make it !Unpin.

use std::pin::Pin;
use std::marker::PhantomPinned;

#[derive(Debug)]
struct Test {
    a: String,
    b: *const String,
    _marker: PhantomPinned
}

impl Test {
    fn new(txt: &str) -> Self {
        Test {
            a: String::from(txt),
            b: std::ptr::null(),
            _marker: PhantomPinned
        }
    }

    fn init(self: Pin<&mut Self>) {
        let self_ptr: *const String = &self.a;
        let this = unsafe { self.get_unchecked_mut() };
        this.b = self_ptr;
    }

    fn a(self: Pin<&Self>) -> &str {
        &self.get_ref().a
    }

    fn b(self: Pin<&Self>) -> &String {
        assert!(!self.b.is_null());
        unsafe { &*(self.b) }
    }
}

In this code, we will not be able to get a mutable reference of Test object in safe Rust. Note that in init we used a method get_unchecked_mut to get the mutable reference in the unsafe block. get_unchecked_mut returns a mutable reference of the pointer even if the type is !Unpin. Using this method requires us to follow the rules of Pin manually:

If P implements Unpin, P must be unpinned throughout its lifecycle
If P implements !Unpin, P must be pinned throughout its lifecycle

Fix the Value on Heap

Using smart pointer Box and method Box::pin, we can allocation an !Unpin type on the heap and fix its location with Pin.

use std::pin::Pin;
use std::marker::PhantomPinned;

#[derive(Debug)]
struct Test {
    a: String,
    b: *const String,
    _marker: PhantomPinned
}

impl Test {
    fn new(txt: &str) -> Pin<Box<Self>> {
        let t = Test {
            a: String::from(txt),
            b: std::ptr::null(),
            _marker: PhantomPinned
        };
        let mut boxed = Box::pin(t);
        let self_ptr: *const String = &boxed.as_ref().a;
        unsafe { boxed.as_mut().get_unchecked_mut().b = self_ptr };

        boxed
    }

    fn a(self: Pin<&Self>) -> &str {
        &self.get_ref().a
    }

    fn b(self: Pin<&Self>) -> &String {
        unsafe { &*(self.b) }
    }
}

pub fn main() {
    let test1 = Test::new("test1");
    let test2 = Test::new("test2");
    println!("a: {}, b: {}", test1.as_ref().a(), test1.as_ref().b());
    println!("a: {}, b: {}", test2.as_ref().a(), test2.as_ref().b());
}

The Use of Pin in Future

The underlying implementation of async/await programming is the Future trait. Future is one of the few traits that implement !Unpin. In fact, the introduction of Pin in Rust is largely due to issues related to Future. This is the API of Future trait:

trait Future {
    type Output;
    fn poll(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>
    ) -> Poll<Self::Output>;
}

pub enum Poll<T> {
    Ready(T),
    Pending
}

Here's a standard use of async/await:

let fut_one = /*... */ // Future 1
let fut_two = /* ... */ // Future 2
async move {
    fut_one.await;
    fut_two.await;
}

The code above is actually a syntatic suger. Its underlying implementation is a Future type with poll method:

struct AsyncFuxture {
    fut_one: FutOne,
    fut_two: FutTwo,
    state: State
}

enum State {
    AwaitingFutOne,
    AwaitingFutTwo,
    Done
}

impl Future for AsyncFuture {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
        loop {
            match self.state {
                State::AwaitingFutOne => match self.fut_one.poll(..) {
                    Poll::Ready(()) => self.state = State::AwaitingFutTwo,
                    Poll::Pending => return Poll::Pending
                State::AwaitingFutTwo => match self.fut_two.poll(..) {
                    Poll::Ready(()) => self.state = State::Done,
                    Poll::Pending => return Poll::Pending
                }
                State::Done => return Poll::Ready(())
            }
        }
    }
}

If we include reference types in async block, issues may occur:

async {
    let mut x = [0; 128];
    let read_into_buf_fut = read_into_buf(&mut x);
    read_into_buf_fut.await;
    println!("{:?}", x);
}

This code will be converted to:

struct ReadIntoBuf<'a> {
    buf: &'a mut [u8]
}

struct AsyncFuture {
    x: [u8; 128],
    read_into_buf_fut: ReadIntoBuf
}

In this case, ReadIntoBuf has an attribute buf, which is a pointer to x. As a result, if AsyncFuture is moved, x will be moved as well, leading to an invalid address for buf. The way to prevent that is to fix Future to a specific memory location using Pin. This is the reason the poll method in Future trait uses Pin<&mut Self> as type annotation for self.

Rust Learning Note: Creating a Timer with Async/Await

Raine-Yang-UofT — Sun, 31 Dec 2023 17:36:25 +0000

This article is a summary of Chapter 4.11.2 of Rust Course (course.rs/)

Future Trait

Future trait is the fundamental building block for Rust async/await mechanism. It represents a computation that may produce a value in the future. Future trait defines a poll method that is used to produce the output asynchronously.

trait Future {
    type Output;
    fn poll(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>
    ) -> Poll<Self::Output>;
}

pub enum Poll<T> {
    Ready(T),
    Pending
}

When the poll method is invoked, it attempts to make progress in its computation. If the computation is completed in one poll, Ready is returned indicating the task is completed. If the computation cannot be completed in a poll, Pending is returned. In addition, the Waker obtained from parameter cx is stored in future. Waker is used to wake up a task by notifying the executor that the task can be polled again.

Here is a Timer that uses Future trait. It creates a Future that completes after a specified duration. When the timer is created, a thread a started a sleep for a certain duration. The thread notifies the Future when the sleep ends.

use std::{
    future::Future,
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll, Waker},
    thread,
    time::Duration
};


pub struct TimerFuture {
    shared_state: Arc<Mutex<SharedState>>
}

struct SharedState {
    completed: bool,
    waker: Option<Waker>
}

impl Future for TimerFuture {
    type Output = ();
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut shared_state = self.shared_state.lock().unwrap();
        if shared_state.completed {
            Poll::Ready(())
        } else {
            shared_state.waker = Some(cx.waker().clone());
            Poll::Pending
        }
    }
}

impl TimerFuture {
    pub fn new(duration: Duration) -> Self {
        let shared_state = Arc::new(Mutex::new(SharedState {
            completed: false,
            waker:None
        }));

        let thread_shared_state = shared_state.clone();
        thread::spawn(move || {
            thread::sleep(duration);
            let mut shared_state = thread_shared_state.lock().unwrap();
            shared_state.completed = true;
            if let Some(waker) = shared_state.waker.take() {
                waker.wake()
            }
        });

        TimerFuture { shared_state }
    }
}

Explanation:

SharedState struct is used to sychronize the state between Future and thread. When the thread finishes the sleep, it changes completed to true and wakes up the Future. Since the SharedState struct need to be accessed and mutated by multiple threads, it is inside Arc>
We implement Future trait for TimerFuture. The poll method returns Ready (meaning the task is completed) when shared_state.completed is set to true. If it is false, we store the waker get from the parameter cx and returns Pending (meaning the task is not completed).
We then implement constructor for TimerFuture. The constructor creates a new SharedState, and creates a thread the sleeps for a certain duration. When the thread wakes up from sleep, it sets shared_state.completed to true and notifies the Future to poll again through calling wake() method.

To use the Future timer, we need an Executor to initially poll the timers. After that, the Executor only polls a timer when it is notified by the wake method. Here is a simple implementation of an executor:

use {
    futures::{
        future::{BoxFuture, FutureExt},
        task::{waker_ref, ArcWake}
    },
    std::{
        future::Future,
        sync::mpsc::{sync_channel, Receiver, SyncSender},
        sync::{Arc, Mutex},
        task::{Context, Poll},
        time::Duration
    },
    timer_future::TimerFuture
};


struct Executor {
    ready_queue: Receiver<Arc<Task>>
}

#[derive(Clone)]
struct Spawner {
    task_sender: SyncSender<Arc<Task>>
}

struct Task {
    future: Mutex<Option<BoxFuture<'static, ()>>>,
    task_sender: SyncSender<Arc<Task>>
}

fn new_executor_and_spawner() -> (Executor, Spawner) {
    const MAX_QUEUED_TASKS: usize = 10000;
    let (task_sender, ready_queue) = sync_channel(MAX_QUEUED_TASKS);
    (Executor {ready_queue}, Spawner {task_sender})
}

impl Spawner {
    fn spawn(&self, future: impl Future<Output = ()> + 'static + Send) {
        let future = future.boxed();
        let task = Arc::new(Task {
            future: Mutex::new(Some(future)),
            task_sender: self.task_sender.clone()
        });
        self.task_sender.send(task).expect("The task queue is full");
    }
}

impl ArcWake for Task {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        let cloned = arc_self.clone();
        arc_self.task_sender.send(cloned).expect("The task queue is full");
    }
}

impl Executor {
    fn run(&self) {
        while let Ok(task) = self.ready_queue.recv() {
            let mut future_slot = task.future.lock().unwrap();
            if let Some(mut future) = future_slot.take() {
                let waker = waker_ref(&task);
                let context = &mut Context::from_waker(&*waker);
                if future.as_mut().poll(context).is_pending() {
                    *future_slot = Some(future);
                }
            }
        }
    }
}


fn main() {
    let (executor, spawner) = new_executor_and_spawner();

    spawner.spawn(async {
        println!("howdy!");
        TimerFuture::new(Duration::new(2, 0)).await;
        println!("done!");
    });

    drop(spawner);

    executor.run();
}

Executor is used to receive tasks from the mpsc channel and execute them. Spawner is used to create new Future timers and send them to the mpsc channel. Task contains the Future timer and sender to send itself to Executor.
new_executor_and_spawner() method initializes a sychronous mpsc channel with capacity 10000, and assigns the channel sender and receiver to Spawner and Executor.
The spawn method of Spawner takes a trait object of Future and wraps it in a Box to make its lifecycle 'static. It creates a Task from the Future object and sends it to the mpsc channel.
ArcWake trait enables Task to wake itself up when it is ready to be polled again. This trait contains wake_by_ref method, in which we make a clone of an Arc pointer to the task and send it back to the mpsc channel.
Finally, we implement run method for Executor to contiously polling tasks from the mpsc channel.

Rust Learning Note: Unsafe Rust

Raine-Yang-UofT — Sat, 30 Dec 2023 15:40:43 +0000

This article is a summary of Chapter 4.9.1 of Rust Course (course.rs/)

Code inside unsafe allows behaviors that are not permitted by the compiler in safe Rust, including:

Dereferencing a raw pointer
Invoking an unsafe or external function
Accessing or mutating a mutable static variable
Implementing an unsafe trait
Skipping borrowing checking

These functionalities allows for more convenient and efficient implementations of certain tasks, and prevents misjudgements of the compiler. However, since some of the checkings by the compiler are skipped, the developer is responsible for the potential bugs lead by unsafe Rust. In general, the unsafe block should be as small as possible.

Dereference Raw Pointer

Raw pointers are similar to references in their functions. However, raw pointers have these differences from references and smart pointers:

Raw pointers can bypass borrowing rules in Rust. The coexistence of multiple mutable and immutable pointers is allowed
Raw pointers may point to invalid memory
Raw pointers can be null
Raw pointers will not be dropped automatically

let mut num = 5;

let r1 = &num as *const i32;
let r2 = &mut num as *mut i32;
unsafe {
    println!("r1: {}", *r1);
}

In the code above we created an immutable raw pointer r1 and a mutable raw pointer r2. In this case, it is allowed to use the immutable pointer r1 even when a mutable pointer is defined. Note the the creation of raw pointers is not unsafe, but the dereferencing of raw pointers is.

We can also create raw pointers from memory address. Note that we should first get the address and then use still. Prevent making up a memory address as it will likely lead to unexpected behaviors or segmentation fault error.

use std::{slice::from_raw_parts, str::from_utf8_unchecked};

fn get_memory_location() -> (usize, usize) {
    let string = "Hello World";
    let pointer = string.as_ptr() as usize;
    let length = string.len();
    (pointer, length)
}

fn get_str_at_location(pointer: usize, length: usize) -> &'static str {
    unsafe { from_utf8_unchecked(from_raw_parts(pointer as *const u8, length)) }
}

fn main() {
    let (pointer, length) = get_memory_location();
    let message = get_str_at_location(pointer, length);
    println!("The {} bytes at 0x{:X} stored: {}", length, pointer, message);

    // let message = get_str_at_location(1000, 10);    Dangerous Behavior
}

We can derefence a raw pointer with *. This must be in unsafe block since we cannot guarantee the memory safety of raw pointers.

let a = 1;
let b: *const i32 = &a as *const i32;
unsafe {
    println!("{}", *c);
}

Finally, we can create raw pointers from smart pointers using reference or method into_raw()

let a: Box<i32> = Box::new(10);
let b: *const i32 = &*a;
let c: *consy i32 = Box::into_raw(a);

Invoke Unsafe Functions and Methods

Unsafe function definition is to indicate that the function may contain undefined behaviors. The calling of unsafe functions must also be embedded in unsafe block. There is no need to contain addition unsafe blocks inside an unsafe function.

unsafe fn dangerous() {}
fn main() {
    unsafe { dangerous(); }
}

A trait is unsafe if it contains at least one unsafe method

unsafe trait Foo {
}

unsafe impl Foo for i32 {
}

FFI

FFI (Foreign Function Interface) allows using functions from another language in Rust. Code that uses FFI need to be inside unsafe block since other languages may not implement the rules of Rust compiler.

extern "C" {
    fn abs(input: i32) -> i32;
}

fn main() {
    unsafe {
        println!("{}", abs(-3));
    }
}

In this code, we define the abs function in C in extern keyword. To invoke this function, the code need to be in unsafe block. In extern "C" we list the function signiture of the external function to invoke, and "C" refers to the ABI (Application Binary Interface) of the external function. ABI defines how to make function calls on an assembly level.

Similarly, we can also define rust functions that can be used by other languages:

#[no_mangle]
pub extern "C" fn call_from_c() {
    println!("Call Rust function from C");
}

In the code above, we allow the function call_from_c to be called in C language. The annotation #[no_mangle] tells the compiler not to change the function name at compiler time so the name can be correctly identified by other languages.

Access Static Variables

A static variable is a global variable that can be accessed throughout the program. Accessing and mutating a static variable is unsafe since it may lead to data competing in multithreading. We should only static variables only in single thread.

static mut REQUEST_RECV: usize = 0;
fn main() {
    unsafe {
        REQUEST_RECV += 1;
        assert_eq!(REQUEST_RECV, 1);
    }
}

Access Union

Union is similar to struct except that all of its attributes share the same memory, that is, when one attribute is mutated, the other attributes will also be overwritten. Thus, union is considered as unsafe.

 #[repr(C)]
union MyUnion {
    f1: u32,
    f2: f32
}

In this code, since f1 and f2 share the same memory space. If we assign a value to f1 and access f2, what we get is a f32 representation of value in f1.

Rust Learning Note: Exception Handling

Raine-Yang-UofT — Thu, 28 Dec 2023 17:56:20 +0000

This article is a summary of Chapter 4.8 of Rust Course (course.rs/)

Type Conversions of Option and Result

or() and and()

or() and and() are similar to the or and and in boolean expressions. They evaluate an Option/Result from two Option/Result values.

or(): returns the first expression that is Some/Ok, returns the second None/Err if both expressions are None/Err.
and(): returns the second Some/Ok if both expressions are Some/Ok, returns the first occurance of None/Err2

fn main() {
    let s1 = Some(1);
    let s2 = Some(2);
    let n: Option<i32> = None;

    let o1: Result<i32, i32> = Ok(1);
    let o2: Result<i32, i32> = Ok(2);
    let e1: Result<i32, i32> = Err(1);
    let e2: Result<i32, i32> = Err(2);

    assert_eq!(s1.or(s2), s1); // Some1 or Some2 = Some1
    assert_eq!(s1.or(n), s1);  // Some or None = Some
    assert_eq!(n.or(s1), s1);  // None or Some = Some
    assert_eq!(n.or(n), n);    // None1 or None2 = None2

    assert_eq!(o1.or(o2), o1); // Ok1 or Ok2 = Ok1
    assert_eq!(o1.or(e1), o1); // Ok or Err = Ok
    assert_eq!(e1.or(o1), o1); // Err or Ok = Ok
    assert_eq!(e1.or(e2), e2); // Err1 or Err2 = Err2

    assert_eq!(s1.and(s2), s2); // Some1 and Some2 = Some2
    assert_eq!(s1.and(n), n);   // Some and None = None
    assert_eq!(n.and(s1), n);   // None and Some = None
    assert_eq!(n.and(n), n);    // None1 and None2 = None1

    assert_eq!(o1.and(o2), o2); // Ok1 and Ok2 = Ok2
    assert_eq!(o1.and(e1), e1); // Ok and Err = Err
    assert_eq!(e1.and(o1), e1); // Err and Ok = Err
    assert_eq!(e1.and(e2), e1); // Err1 and Err2 = Err1
}

or_else() and and_then()

or_else() and and_then() follow the same evaluation rules as or() and and(). However, they take a closure as the second expression.

fn main() {
    let s1 = Some(1);
    let s2 = Some(2);
    let fn_some = || Some(2);
    let n: Option<i32> = None;
    let fn_none = || None;

    assert_eq!(s1.or_else(fn_some), s1);  // Some1 or_else Some2 = Some1
    assert_eq!(s1.or_else(fn_none), s1);  // Some or_else None = Some
    assert_eq!(n.or_else(fn_some), s2);   // None or_else Some = Some
    assert_eq!(n.or_else(fn_none), None); // None1 or_else None2 = None2

    let o1: Result<i32, i32> = Ok(1);
    let o2: Result<i32, i32> = Ok(2);
    let fn_ok = || Ok(2);
    let e1: Result<i32, i32> = Err(1);
    let e2: Result<i32, i32> = Err(2);
    let fn_err = || Err(2);

    assert_eq!(o1.and_then(fn_ok), o2);  // Ok1 and_then Ok2 = Ok2
    assert_eq!(o1.and_then(fn_err), e2); // Ok and_then Err = Err
    assert_eq!(e1.and_then(fn_ok), e1);  // Err and_then Ok = Err
    assert_eq!(e1.and_then(fn_err), e1); // Err1 and_then Err2 = Err1
}

filter

filter is used to filter an Option. It takes a closure that returns a boolean value, and keeps the value if the closure returns true.

fn main() {
    let s1 = Some(3);
    let s2 = Some(6);
    let n = None;

    let fn_is_even = |x: &i8| x % 2 == 0;

    assert_eq!(s1.filter(fn_is_even), n);
    assert_eq!(s2.filter(fn_is_even), s2);
    assert_eq!(n.filter(fn_is_even), n);
}

map() and map_err()

map takes a closure as parameter and maps the value in Some or Ok to another value. It cannot be used to change the value in Err.

fn main() {
    let s1 = Some("abcde");
    let o1: Result<&str, &str> = Ok("abcde");
    let fn_character_count = |s: &str| s.chars().count();

    let e1: Result<&str, &str> = Err("abcde");    
    let e2: Result<usize, &str> = Err("abcde");

    assert_eq!(s1.map(fn_character_count), 5);
    assert_eq!(o1.map(fn_character_count), 5);

    assert_eq!(e1.map(fn_character_count), e2);
}

If we want to map the value in Err to another value, we can use method map_err

fn main() {
    let e1: Result<&str, &str> = Err("404");
    let e2: Result<&str, isize> = Err(404);
    let fn_character_count = |s: &str| -> isize {s.parse().unwrap()};
    assert_eq!(e1.map_err(fn_character_count), e2);
}

map_or provides a default value for map. If the Option/Result is None/Err, the default value is returned. map_or_else is similar to map_or, except that it evaluates the default value with a closure.

fn main() {
    const DEFAULT: u32 = 1;

    let s: Result<u32, ()> = Ok(10);
    let n: Option<u32> = None;
    let fn_closure = |v: u32| v + 2;

    assert_eq!(s.map_or(DEFAULT, fn_closure), 12);
    assert_eq!(n.map_or(DEFAULT, fn_closure), DEFAULT);

    let o = Ok(10);
    let e = Err(5);
    let fn_default = |v: i8| v + 1;
    assert_eq!(o.map_or_else(fn_default, fn_closure), 12);
    assert_eq!(e.map_or_else(fn_default, fn_closure), 6);

Note that in the code above, we are able to use the value in Err in the closure for map_or_else.

ok_or() and ok_or_else()

ok_or and ok_or_else are used to convert an Option to a Result. ok_or takes a a parameter as the default Err value, and ok_or_else takes a closure as Err value.

fn main() {
    const DEFAULT: &str = "error message";
    let s = Some("abcde");
    let n: Option<&str> = None;

    assert_eq!(s.ok_or(DEFAULT), Ok("abcde"));
    assert_eq!(n.ok_or(DEFAULT), Err(DEFAULT));

    let fn_err = || "error message";
    assert_eq!(s.ok_or_else(fn_err), Ok("abcde"));
    assert_eq!(n.ok_or_else(fn_err), Err("error message"));

Defining Custome Error Type

std::error::Error trait allows an object to be represented as error values, i.e., values inside Err() for Result enumerator. A custome errror type needs to implement Debug and Display, which are prerequistes of Error.

use std::fmt;

struct AppError {
    code: usize,
    message: String
}

impl fmt::Display for AppError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let err_msg = match self.code {
            404 => "Page not found",
            - => "Something went wrong"
        };

        write!(f, "{}", err_msg)
    }
}

impl fmt::Debug for AppError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "AppError {{ code: {}, message: {} }}",
            self.code, self.message
        )
    }
}

fn produce_error() -> Result<(), AppError> {
    Err(AppError {
        code: 404,
        message: String::from("Page not found")
    })
}

fn main() {
    match produce_error() {
        Err(e) => eprintln!("{}", e),
        _ => println!("No error")
    }
}

Conversion among errors with From trait

From trait allows us to convert one type to another. Using From, we can convert different standard error types to our custome type.

use std::fs::File
use std::io::{self, Read};
use std::num

#[derive(Debug)]
struct AppError {
    kind: String,
    message: String
}

impl From<io::Error> for AppError {
    fn from(error: io:Error) -> Self {
        AppError {
            kind: String::from("io"),
            message: error.to_string()
        }
    }
}

impl From<num::ParseIntError> for AppError {
    fn from(error: num::ParseIntError) -> Self {
        AppError {
            kind: String::from("parse"),
            message: error.to_string()
        }
    }
}

fn main() -> Result<(), AppError> {
    let mut file = File::open("hello world.txt")?;
    let mut content = String::new();
    file.read_to_string(&mut content)?;
    let number: usize;
    number = content.parse()?;

Note that ? is able to implicitly convert the standard error types to our custome type AppError once we implement the From trait for type conversion.

Unifying different error types

Sometimes we want a function to return different types of errors. In this case, we need to have a unified error type for the function type annotation.

One easy way to do that is to use trait object:

use std::fs::read_to_string;
use std::error::Error;
fn main() -> Result<(), Box<dyn Error>> {
    let html = render()?;
    println!("{}", html);
    Ok(())
}

fn render() -> Result<String, Box<dyn Error>> {
    let file = std:;env::var("MARKDOWN")?;
    let source = read_to_string(file)?;
    Ok(source)
}

One limitation of this method is that the trait object would not restrict the specific types of errors. As long as a type implements Debug and Display type, it can be returned as the error. Another way to avoid this is to define custome error types that can be converted from the standard errors:

use std::fs::read_to_string;

fn main() -> Result<(), MyError> {
  let html = render()?;
  println!("{}", html);
  Ok(())
}

fn render() -> Result<String, MyError> {
  let file = std::env::var("MARKDOWN")?;
  let source = read_to_string(file)?;
  Ok(source)
}

#[derive(Debug)]
enum MyError {
  EnvironmentVariableNotFound,
  IOError(std::io::Error),
}

impl From<std::env::VarError> for MyError {
  fn from(_: std::env::VarError) -> Self {
    Self::EnvironmentVariableNotFound
  }
}

impl From<std::io::Error> for MyError {
  fn from(value: std::io::Error) -> Self {
    Self::IOError(value)
  }
}

impl std::error::Error for MyError {}

impl std::fmt::Display for MyError {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    match self {
      MyError::EnvironmentVariableNotFound => write!(f, "Environment variable not found"),
      MyError::IOError(err) => write!(f, "IO Error: {}", err.to_string()),
    }
  }
}

We can utilize third-party libraries to simplify the code above. For example thiserror (https://github.com/dtolnay/thiserror) provides convenien derive macro for Error trait:

use std::fs::read_to_string;

fn main() -> Result<(), MyError> {
  let html = render()?;
  println!("{}", html);
  Ok(())
}

fn render() -> Result<String, MyError> {
  let file = std::env::var("MARKDOWN")?;
  let source = read_to_string(file)?;
  Ok(source)
}

#[derive(thiserror::Error, Debug)]
enum MyError {
    #[error("Environment variable not found")]
    EnvironmentVariableNotFound(#[from] std::env::VarError),
    #[error(transparent)]
    IOError(#[from] std::io::Error)
}

Rust Learning Note: Global Variable

Raine-Yang-UofT — Wed, 27 Dec 2023 15:24:48 +0000

This article is a summary of Chapter 4.7 of Rust Course (course.rs/)

Global Variables Initialized at Compile Time

Static Constant

Static constants are global constants that can be used anywhere in the program, regardless of the scope in which it is defined, since the lifecycle of a constant is always 'static. However, different references to the same constant do not guarantee the same memory address.

const MAX_ID: usize = usize::MAX / 2;
fn main() {
    println!({MAX_ID});
}

Here are a few things to note about static constants:

We use const instead of let to define a constant
Type annotations for constants cannot be omitted
A constant must be assigned to a specific value or an expression that can be evaluated at compiler time.
There cannot be duplicate constant names

Static Variable

Static variables are global variables that can be accessed throughout the program. A static variable must be accessed and mutated in unsafe block, since it is unsafe especially in multithreading. A static variable must be defined as a value known in compile time, and the value must implement Sync trait. Any references to a static variable always yield the same memory address.

static mut REQUEST_RECV: usize = 0;
fn main() {
    unsafe {
        REQUEST_RECV += 1;
    }
}

Atomic Type

If we want to use a static variable safely in multithreading, we can apply atomic types:

use std::sync::atomic::{AtomicUsize, Ordering};
static REQUEST_RECV: AtomicUsize = AtomicUsize::new(0);
fn main() {
    for _ in 0..100 {
        REQUEST_RECV.fetch_add(1, Ordering::Relaxed);
    }
}

Here an example global ID generator with global variables mentioned above:

use std::sync::atomic::{Ordering, AtomicUsize};

struct Factory {
    factory_id: usize,
}

static GLOBAL_ID_COUNTER: AtomicUsize = AtomicUsize::new(0);
const MAX_ID: usize = usize::MAX / 2;

fn generate_id() -> usize {
    let current_val = GLOBAL_ID_COUNTER.load(Ordering::Relaxed);
    if current_val > MAX_ID {
        panic!("Factory ids overflowed");
    }
    GLOBAL_ID_COUNTER.fetch_add(1, Ordering::Relaxed);
    let next_id = GLOBAL_ID_COUNTER.load(Ordering::Relaxed);
    if next_id > MAX_ID {
        panic!("Factory ids overflowed");
    }
    next_id
}

impl Factory {
    fn new() -> Self {
        Self {
            factory_id: generate_id()
        }
    }
}

Global Variable Initialized at Runtime

Lazy_static

The global variables above can only be initialized at compile time. However, sometimes we want to initialize global variables at runtime, such as creating a global Mutex lock. In this case, we can use macro lazy_static:

use std::sync::Mutex
use lazy_static::lazy_static;
lazy_static! {
    static ref NAMES: Mutex<String> = Mutex::new(String::from("Sunface, Jack, Allen"));
}

fn main() {
    let mut v = NAMES.lock().unwrap();
    v.push_str(", Myth");
    println!("{}", v);
}

lazy_static can also be used to load a global configuration that is initialized only after the program starts. Here's an example of implementing a global cache with lazy_static:

use lazy_static::lazy_static;
use std::collections::HashMap;

lazy_static! {
    static ref HASHMAP: HashMap<u32, &'static str> = {
        let mut m = HashMap::new();
        m.insert(0, "foo");
        m.insert(1, "bar");
        m.insert(2, "baz");
        m
    };
}

fn main() {
    println!("The entry for 0 is {}", HASHMAP.get(&0).unwrap());
    println!("The entry for 1 is {}", HASHMAP.get(&1).unwrap());
}

In the code above, HASHMAP is initialized only when it is first invoked at the first line in main method.

Box::leak

Box::leak returns a variable with a 'static lifecycle. It can be used for reassignment of static variable. It is not permitted to reassign a global variable to a value of a local variable, since the local variable has a shorter lifecycle. However, we can achieve this with Box::leak.

#[derive(Debug)]
struct Config {
    a: String,
    b: String
}

static mut CONFIG: Option<&mut Config> = None;

fn main() {
    let c = Box::new(Config {
        a: "A".to_string(),
        b: "B".to_string()
    });

    unsafe {
        CONFIG = Some(Box::leak(c));
        println!("{:?}", CONFIG);
    }
}

Using Box::leak, we can also return a global cariable from a function:

#[derive(Debug)]
struct Config {
    a: String,
    b: String
}

static mut CONFIG: Option<&mut Config> = None;

fn init() -> Option<&'static mut Config> {
    let c = Box::new(Config {
        a: "A".to_string(),
        b: "B".to_string()
    });
    Some(Box::leak(c))
}

fn main() {
    unsafe {
        CONFIG = init();
        println!("{:?}", CONFIG)
    }
}

Rust Learning Note: Thread Sychronization

Raine-Yang-UofT — Mon, 25 Dec 2023 18:39:05 +0000

this article is a summary of Chapter 4.6.3 - 4.6.6 of Rust Course (course.rs/)

Multiple Producer Single Consumer

The Rust standard library provides a channel among multiple threads called std::sync::mpsc. mpsc is short for multiple producer, single consumer, meaning that the channel supports multiple senders but only one receive. Here is a simple example use of mpsc:

use std::sync::mpsc;
use std::thread;

fn main() {
    let (tx, rx) = mpsc::channel();

    thread::spawn(move || {
        tx.send(1).unwrap();
    });

    println!("receive {}", rx.recv().unwrap());
}

In the code above, mpsc::channel() creates a new channel, which returns a tuple containing the sender (tx), and receiver (rx). The spawned thread sends 1 with tx, which is received by the main thread with rx. We need to use move to transfer the ownership of tx into the closure. The send method returns a Result, meaning that it may send the error if an error occurs.

try_recv method

Method recv() would block the thread until it receives the message. Another method called try_recv would try to read the message, returns an error if the message has not been sent, and proceed without blocking the thread.

use std::sync::mpsc;
use std::thread;

fn main() {
    let (tx, rx) = mpsc::channel();

    thread::spawn(move || {
        tx.send(1).unwrap();
    });

    println!("receive {}", rx.try_recv().unwrap());
}

If we replace the rx.recv() previous code with rx.try_recv(), we would likely get an Err(Empty) since the creation of a new thread is much slower than the processing of the main thread.

When sending a variable that has ownership (does not implement Copy), send() would transfer the ownership of the variable from the sender to the receiver.

Multiple Senders

To have multiple senders of a channel, we need to clone the tx in the channel and assign the copy of tx to the other thread.

use std::sync::mpsc;
use std::thread;

fn main() {
    let (tx, rx) = mpsc::channel();
    let tx1 = tx.clone();
    thread::spawn(move || {
        for i in 0..10 {
            tx.send(i).unwrap();
        }
    });

    thread::spawn(move || {
        for i in 0..10 {
            tx1.send(i).unwrap();
        }
    });

    for received in rx {
        println!("Got: {received}");
    }
}

Asychronized and Sychronized Channel

The mpsc::channel we used above is asychronized, meaning that the sender continues executing the following codes after sending the message, regardless of whether the message is received. Using mpsc::sync_channel we create a sychronized channel, in which the sender is blocked until the receiver receives the message.

use std::sync::mpsc;
use std::thread;
use std::time::Duration;

fn main() {
    let (tx, rx) = mpsc::sync_channel(0);

    let handle = thread::spawn(move || {
        println!("before sending");
        tx.send(1).unwrap();
        println!("after sending");
    });

    println!("before sleeping");
    thread::sleep(Duration::from_secs(3));
    println!("after sleeping");

    println!("receive {}", rx.recv().unwrap());
    handle.join().unwrap();
}

In this code, "after sending" by the spawned thread will be printed only after "after sleeping" in the main thread, indicating that the spawned thread is blocked when the main thread does not receive the message. The parameter inside mpsc::sync_channel is the size of the cache. When the cache is set to N, the sender can send at most N messages to the receiver without being blocked, and is blocked when the cache is full.

Lock

Besides the use of message system, we can also sychronize threads with shared memory. In this case, we need to assign locks to the data accessed by a thread to prevent data competing.

Mutex

Mutex stands for mutual exclusion. To access data in Mutex, the thread need to request a lock with lock(), which blocks the thread if the lock is acquired by other threads. This allows multiple threads to access the data one by one. Since Mutex is also a smart pointer, it implements Deref, allowing automatic dereferencing to retrieve the data inside, and Drop, allowing the lock to be released once it leaves the scope.

use std::sync::{Arc, Mutex};
use std::thread;

fn main() {
    let counter = Arc::new(Mutex::new(0));
    let mut handles = vec![];

    for _ in 0..10 {
        let counter = Arc::clone(&counter);
        let handle = thread::spawn(move || {
            let mut num = counter.lock().unwrap();
            *num += 1;
        });
        handles.push(handle);
    }

    for handle in handles {
        handle.join().unwrap();
    }

    println!("Result: {}", *counter.lock().unwrap());
}

In the code above, we use Arc to allow multiple ownership of data, and Mutex to achieve internal mutability. In single thread, Rc/RefCell is a common combination to achieve internal mutability and multiple ownership. Its multithreading version is Arc/Mutex.

Deadlock

Deadlock would happen if a lock is being requested when it is not released yet. Here is a single example in one thread:

use std::sync::Mutex;

fn main() {
    let data = Mutex::new(0);
    let d1 = data.lock();
    let d2 = data.lock();
}

In multithreading, if two threads each acquires a lock and tries to acquire the lock owned by the other thread, deadlock may happen:

use std::{sync::{Mutex, MutexGuard}, thread};
use std::thread::sleep;
use std::time::Duration;

use lazy_static::lazy_static;
lazy_static! {
    static ref MUTEX1: Mutex<i64> = Mutex::new(0);
    static ref MUTEX2: Mutex<i64> = Mutex::new(0);
}

fn main() {
    let mut children = vec![];
    for i_thread in 0..1 {
        children.push(thread::spawn(move || {
            if i_thread % 2 == 0 {
                let guard1: MutexGuard<i64> = MUTEX1.lock().unwrap();
                println!("MUTEX1 locked by {i_thread}");
                let guard1 = MUTEX2.lock().unwrap();
            } else {
                let guard2: MutexGuard<i64> = MUTEX2.lock().unwrap();
                println!("MUTEX2 locked by {i_thread}");
                let guard2 = MUTEX1.lock().unwrap();
            }
        }));
    }

    for child in children {
        let _ = child.join();
    }
}

In this code, thread 1 locks MUTEX1 and tries to acquire MUTEX2, and thread 2 locks MUTEX2 and tries to acquire MUTEX1. Both threads will then be blocked indefinitely.

try_lock

In contrast with lock, try_lock tries to acquire the lock once without blocking the thread, and returns an error if the lock has been acquired.

RwLock

Mutex would assign a lock for both reading and writing operations, which could lead to inefficiency when we have large number of reading operations from multiple threads. RwLock does not set locks for reading operations, and only blocks other threads when writing operations occur. In short, RwLock allows either reading operations from multiple threads or writing operations from one thread at a time.

use std::sync::RwLock;

fn main() {
    let lock = RwLock::new(5);

    {
        // multiple readings allowed
        let r1 = lock.read().unwrap();
        let r2 = lock.read().unwrap();
        assert_eq!(*r1, 5);
    }

    {
        // only one writing allowed
        let mut w = lock.write().unwrap();
        *w += 1;
        assert_eq!(*w, 6);    
    }
}

In general, Mutex is more efficient in its operations. We should use RwLock only when we want multiple threads to read the data simultaneouly and possibly acquire the data for a long time.

Condition Variable

Condition variable can suspend a thread and continue it only after certain contations. It can be used to control the order of data access of multiple threads.

use std::sync::{Arc, Mutex, Condvar};
use std::thread::{spawn, sleep};
use std::time::Duration;

fn main() {
    let flag = Arc::new(Mutex::new(false));
    let cond = Arc::new(Condvar::new());
    let cflag = flag.clone();
    let ccond = cond.clone();

    let hdl = spawn(move || {
        let mut lock = cflag.lock().unwrap();
        let mut counter = 0;

        while counter < 3 {
            while !*lock {
                lock = ccond.wait(lock).unwrap();
            }

            *lock = false;
            counter += 1;
            println!("inner counter: {counter}");
        }
    });

    let mut counter = 0;
    loop {
        sleep(Duration::from_millis(1000));
        *flag.lock().unwrap() = true;
        counter += 1;
        if counter > 3 {
            break;
        }
        println!("outer counter: {counter}");
        cond.notify_one();
    }
    hdl.join().unwrap();
    println!("{:?}", flag);
}

In this case, if the thread hdl reaches the while loop before the main thread sets the flag to true, it would be blocked by wait() until the main thread invokes notify_one(). The result will be alternative printings from the main thread and hdl triggered by the main thread.

outer counter: 1
inner counter: 1
outer counter: 2
inner counter: 2
outer counter: 3
inner counter: 3
Mutex { data: true, poisoned: false, .. }

Semaphore

Semaphore allows us to control the maximum tasks in the program. The code below creates a semaphore with capacity 3. Where are are more than 3 tasks, the rest of them need to wait after the tasks being executed.

use std::sync::Arc;
use tokio::sync::Semaphore;

#[tokio::main]
async fn main() {
    let semaphore = Arc::new(Semaphore::new(3));
    let mut join_handles = Vec::new();

    for _ in 0..5 {
        let permit = semaphore.clone().acquire_owned().await.unwrap();
        join_handles.push(tokio::spawn(async move {
            // process the task
            drop(permit);
        }));
    }

    for handle in join_handles {
        handle.await.unwrap();
    }
}

Atomic Type

The use of atomic type is another way to sychronize among multiple threads. An atomic operation refers to an indivisible operation on shared memory, ensuing that no other operation can happen simultaneously. One typical application of atomic type is the creation of gobal variables:

use std::sync::atomic::{AtomicU64, Ordering};
use std::thread::{self, JoinHandle};

const N_TIMES: u64 = 1000000;
const N_THREADS: usize = 10;

static R: AtomicU64 = AtomicU64::new(0);

fn add_n_times(n: u64) -> JoinHandle<()> {
    thread::spawn(move || {
        for _ in 0..n {
            R.fetch_add(1, Ordering::Relaxed);
        }
    })
}

fn main() {
    let mut threads = Vec::with_capacity(N_THREADS);

    for _ in 0..N_THREADS {
        threads.push(add_n_times(N_TIMES));
    }

    for thread in threads {
        thread.join().unwrap();
    }
    assert_eq!(N_TIMES * N_THREADS as u64, R.load(Ordering::Relaxed));
}

In the code above, we have 10 threads adding 1 continuously to atomic type R. The final result of R is equal to N_TIMES * N_THREADS demonstates the concurrent safety of atomic type. Also, compared with locks such as Mutex, atomic types have much greater efficiency.

One thing to note about is the parameter inside atomic type operation, Odering::Relaxed. This parameter is used to indicate the memory ordering. When the source code is compiled and run, the order of native code does not always strictly follow the order of the source code due to reordering during compiler optimization and CPU cache. There are five orders to choose from:

Relaxed: no restrictions on reordering
Release: set a memory barrier to ensure that all operations before will always come before it at runtime.
Acquire: set a memory barrier to ensure that all operations after will always come after it at runtime.
AcqRel: the combination of Release and Acquire
SeqCst: the strictest memory ordering, ensuring total order of all operations.

Normally, Release is used for writing operations, and Acquire is used for reading operations. A safer approach is to always use SeqCst, even though it may impair efficiency to some extent.

Here is another example that implements a spinning lock with Atomic:

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::{hint, thread};

fn main() {
    let spinlock = Arc::new(AtomicUsize::new(1));
    let spinlock_clone = Arc::clone(&spinlock);
    let thread = thread::spawn(move || {
        spinlock_clone.store(0, Ordering::SeqCst);
    });

    while spinlock.load(Ordering::SeqCst) != 0 {
        hint::spin_loop();
    }

    if let Err(panic) = thread.join() {
        println!("{:?}", panic);
    }
}

Thread Safety Based on Send and Sync Traits

An object must implement Send and Sync traits to be passed among differen threads. These two traits are marker traits, meaning that they do not have concrete methods. Send trait allows a type to be passed among threads via ownership transfer. Sync trait allows a type to be shared among treads through reference. If an object implements Sync, its reference must implement Send (that is, if T implements Sync, &T implements Send)

Most types in Rust implement Send and Sync by default. However, there are a few exceptions:

raw pointer: raw pointer has no safety guarantee, so it implements neither traits.
Cell and RefCell: Cell and RefCell are not Sync since their base implementation UnsafeCell is not Sync.
Rc: Rc is neither Send nor Sync, so it cannot be used in multithreading.

For a compound data type, if all of its attributes are Send and Sync, then the type also implements Send and Sync.

Rust Learning Note: Multithreading

Raine-Yang-UofT — Sat, 23 Dec 2023 19:09:14 +0000

This article is a summary of chapter 4.6.1 and 4.6.2 of Rust Course (course.rs/)

Concurrent and Parallel Programming

Fig 1 Difference between concurrent and parallel. Image reproduced from (Rust Course)

Concurrent refers to having one CPU core dealing with multiple threads. Every there there is only one thread being executed, and the operating system has specific dispatch algorithms to handing the switching among multiple threads, so the threads appear to be running simultaneously.

Parallel refers to having a CPU core handling one thread, so multiple threads are indeed processed in parallel.

In a more formal definition, a system is concurrent if it allows multiple threads to coexist, and it is parallel if it allows multiple threads to be executed simultaneously. From this definition, parallel processing is a subset of concurrent procession.

Different programming languages have different implementations of concurrent programming. Rust uses 1:1 thread model, meaning that it invokes API provided by the operating system to create threads, and the threads in the program are exactly the same as threads created by the operating systems. Some languages like Go use M:N thread model, meaning that the language has its own implementation of threading model, and M threads created in the language are mapped to N system threads based on the model.

Creating Thread

use std::thread;
use std::time::Duration;

fn main() {
    thread::spawn(|| {
        for i in 1..10 {
            println!("number {} from the spawned thread", i);
            thread::sleep(Duration::from_millis(1));
        }
    });

    for i in 1..5 {
        println!("number {} from the main thread", i);
        thread::sleep(Duration::from_millis(1));
    }
}

We use thread::spawn to create a new thread. The code inside a thread is executed using a closure. thread::sleep will make a thread sleep for a certain amount of time, during which other threads will be processed. In the code above, the thread spawned may not be able to finish execution since the whole program will terminate once the main thread stops. To solve this program, we can use handle.join, which blocks the current thread until the thread it waits for finishes processing.

use std::thread;
use std::time::Duration;

fn main() {
    let handle = thread::spawn(|| {
        for i in 1..5 {
            println!("number {} from the spawned thread", i);
            thread::sleep(Duration::from_millis(1));
        }
    });

    handle.join().unwrap();

    for i in 1..5 {
        println!("number {} from the main thread", i);
        thread::sleep(Duration::from_millis(1));
    }
}

In the code above, handle.join().unwrap() waits the tread handle to be executed before the for loop in the main thread. If we put handle.join().unwrap() at the bottom, the for loop in the thread handle and that in the main thread will be executed alternatively until they both finish.

Transfer External Variables with Move

Keyword move is used to take the ownership of a variable when it is used in a closure. It can also be used to tranfer ownership from one thread to another.

use std::thread;

fn main() {
    let v = vec![1, 2, 3];

    let handle = thread::spawn(move || {
        println!("{:?}", v);
    });

    handle.join().unwrap();
}

In this code, move transfers the ownership of v from main thread to the thread handle. An error would occur without move since Rust cannot determine the lifecycle of v compared with the lifecycle of the thread. It is possible that v is released before the thread.

How Threads Terminate

In Rust, once a thread is created, it would not terminate until it is finished processing or the main thread finished, even if the thread that create it is finished. This is to prevent unexpected behavior for a thread to be terminated before it finishes.

use std::thread;
use std::time::Duration;
fn main() {
    let new_thread = thread::spawn(move || {
        thread::spawn(move || {
            loop {
                println!("I am a new thread");
            }
        })
    });

    new_thread.join().unwrap();
    println!("Child thread is finish");
    thread::sleep(Duration::from_millis(100));
}

In this case, the thread created inside the new_thread will continuously print "I am a new thread" until the main thread finishes, even after new_thread finishes.

Thread Barrier

Thread barrier is used to block a thread until multiple threads reach a certain point.

use std::sync::{Arc, Barrier};
use std::thread;

fn main() {
    let mut handles = Vec::with_capacity(6);
    let barrier = Arc::new(Barrier::new(6));

    for _ in 0..6 {
        let b = barrier.clone();
        handles.push(thread::spawn(move|| {
            println!("before wait");
            b.wait();
            println!("after wait");
        )));
    }

    for handle in handles {
        handle.join().unwrap();
    }
}

In this example, we create a new Barrier that blocks the thread until 6 threads reach the barrier. b.wait() blocks a single thread, and when all 6 threads reach wait() the threads will continue executing.

Thread Local Variable

A thread local variable is a variable that has a separate and independent value for in thread. Each thread would get the initial value of the variable and the change of the variable in each thread would not affect that in other threads. In Rust, we use macro thread_local to initialize a thread local variable, and we use with to use the variable inside a thread. A thread local variable has a 'static lifecycle declared with keyword static.

use std::cell::RefCell;
use std::thread;

thread_local!(static FOO: RefCell<u32> = RefCell::new(1));

FOO.with(|F| {
    assert_eq!(*f.borrow(), 1);
    *f.borrow_mut() = 2;
});

let t = thread::spawn(move || {
    FOO.with(|f| {
        assert_eq!(*f.borrow(), 1);
        *f.borrow_mut() = 3;
    });
});

t.join().unwrap();

FOO.with(|f| {
    assert_eq!(*f.borrow(), 2);
});

Conditional Variable

Using conditional variable with Mutex can suspend a thread and continue it after certain condition.

use std::thread;
use std::sync::{Arc, Mutex, Condvar};

fn main() {
    let pair = Arc::new((Mutex::new(false), Condvar::new()));
    let pair2 = pair.clone();

    thread::spawn(move || {
        let (lock, cvar) = &*pair2;
        let mut started = lock.lock().unwrap();
        println!("changing started");
        *started = true;
        cvar.notify_one();
    });

    let (lock, cvar) = &*pair;
    let mut started = lock.lock().unwrap();
    while !*started {
        started = cvar.wait(started).unwrap();
    }

    println!("started changed");
}

In the code above, Mutex provides mutual exclusion, allowing only one thread to access the data at a time. Condvar is the conditional variable that enables the thread to wait for a particular condition to become true.

We create a Arc containing a tuple (Mutex, Condvar) to share the pair across multiple threads. In the thread, the spawned thread acquires a lock on Mutex and changes started to true, then it notify the Condvar that the condition has been met. In the main thread, the main thread acquires a lock on Mutex, and loops until started become true. In the loop, cvar waits the notification of spawned thread that changes started.

call_once Method

Sometimes, we want a certain functions, like initializing global variables, to be only invoked once by one thread and is ignored by threads following it. This can be achieved through call_once function in Once type

use std::thread;
use std::sync::Once;

static mut VAL: usize = 0;
static INIT: Once = Once::new();

fn main() {
    let handle1 = thread::spawn(move || {
        INIT.call_once(|| {
            unsafe {
                VAL = 1;
            }
        });
    });

    let handle2 = thread::spawn(move || {
        INIT.call_once(|| {
            unsafe {
                VAL = 2;
            }
        });
    });

    handle1.join().unwrap();
    handle2.join().unwrap();

    println!("{}", unsafe { VAL });
}

In this case, INIT.call_once will be invoked by the first thread that calls the function. After that, the second thread would not invoke its INIT.call_once method. Since the creation of threads are asynchorous and slow, the code above may produce 1 or 2.

Rust Learning Note: Smart Pointer

Raine-Yang-UofT — Wed, 20 Dec 2023 18:20:42 +0000

This article is a summary of Chapter 4.4 in Rust Course (course.rs/)

Box
Box is used to allocate a value in the heap, and creates a pointer on the stack pointing to the value. It has these common uses:

1 Allocate values on the heap

Using Box, we can intentionally allocate values on the heap even if the value is by default allocated on the stack (such as i32).

fn main() {
    let a = Box::new(3);
    println!("a = {}", a);

At the println! statement, the smart pointer is automatically deferenced with Def, so we don't need to deference the pointer manually.

2 Prevent copying values on the stack

Variable reassignment on the stack are achieved through copying the entire values, and no transfer of ownership happens. However, values in the heap will not be transferred. Instead, a pointer to the value is transferred from one variable to another, which leads to ownership transfer.

fn main() {
    let arr = [0; 1000]    // stored in stack
    let arr1 = arr;    // copy the array, does not take ownership
    println!("{:?}", arr.len());
    println!("{:?}". arr1.len());

    let arr = Box::new([0; 1000]);    // stored in heap
    let arr1 = arr;    // arr1 takes the ownership of arr

3 Invoking Dynamically Sized Type (DST)

A data type is dynamically sized type (DST) if its size is unknown at compiler time. To give a few examples:

1 slicing: slicing types are DST, we can only use reference of slicing.

2 str: different from &str and String, str is a DST, since we cannot expect strings with different lengths to be stored in a data type with uniform size.

3 trait objects: since we do not know the exact type of objects that will be passed as trait objects, trait objects are DST.

4 recursive types: types that invoke itself in definition. For instance:

enum List {
    Cons(i32, List),
    Nil
}

All DST cannot be directly used. They can only be invoked through pointers, which have fixed size at compile time. For example, the recursive type above can be changed as follows:

enum List {
    Cons(i32, Box<List>),
    Nil
}

Memory Allocation with Box
When we store a vector called vec1, the pointer is stored in stack and the elements are in the heap:

Fig 1. Memory model of a vector. Image reproduced from (Rust Course)

If the elements are of type Box, the pointer stored in stack points to a vector of Box in heap. Each Box pointers to its corresponding element.

Fig 2. Memory model of a vector with Box. Image reproduced from (Rust Course)

To retrieve values from vec2, we need to first deference vec2, and then dereference Box.

fn main() {
    let arr = vec![Box::new(1), Box::new(2)];
    let (first, second) = (&arr[0], &arr[1]);
    let sum = **first + **second

Box::leak

Box::leak is an associated function that can consume Box and returns the value inside it. One use of it is to convert the lifecycle of an object to 'static at runtime. For example, in the code below, using Box::leak, we get a &str type with 'static lifecycle from the String type.

fn main() {
    let s = gen_static_str();
    println!("{}", s);
}

fn gen_static_Str() -> &'static str {
    let mut s = String::new();
    s.push_Str("hello, world");
    Box::leak(s.into_boxed_str())
}

Automatic Dereferencing with Deref Trait

When we want to dereference an object with *, the object must implement trait Deref. Here is an example implementation of Deref with MyBox

struct MyBox<T>(T);

impl<T> MyBox<T> {
    fn new(x: T) -> MyBox<T> {
        MyBox(x)
    }
}

impl<T> Deref for MyBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

After MyBox implements Deref, we can use *MyBox to get value x from MyBox(x). This is achieved through internally invoking

*(y.deref())

When an object that implements deref is passed as a function parameter, deref() can be automatically invoked based on the parameter signiture, for example:

fn main() {
    let s = String::from("hello world");
    display(&s)
}

fn display(s: &str) {
    println!("{}", s);
}

In the code above, String implements Deref trait and can be deferenced to &str. When &s is passed into function display, its type &String is automatically deferenced to &str. Note at the parameter must be passed as a reference to trigger the dereferencing.

The automatic dereferencing can be invoked consecutively until finding a suitable data type, as the code below would also work:

fn main() {
    let s = Box::new(String::from("hello world"));
    display(&s)
}

fn display(s: &str) {
    println!("{}", s);
}

Deref and DerefMut

T: Deref can convert &T to &U, or &mut T to &U

T: DerefMut can convert &mut T to &mut U

To implement DerefMut, we need to first implement Deref for the object, as the example below:

struct MyBox<T> {
    v: T
}

impl<T> MyBox<T> {
    fn new(x: T) -> MyBox<T> {
        MyBox {v: x}
    }
}

use std::ops::Deref;

impl<T> Deref for MyBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.v
    }
}

use std::ops::DerefMut;

impl<T> DerefMut for MyBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.v
    }
}

fn main() {
    let mut s = MyBox::new(String::from("hello, "));
    display(&mut s)
}

fn display(s: &mut String) {
    s.push_str("world");
    println!("{}", s);
}

Object Release with Drop

When an object leaves the scope and is to be released, the compiler will automatically call the drop method of the object. Thus, Drop trait can be used to execute certain cleanup task before the object is released.

struct HasDrop1;
struct HasDrop2;
impl Drop for HasDrop1 {
    fn drop(&mut self) {
        println!("Dropping HasDrop1");
    }
}
impl Drop for HasDrop2 {
    fn drop(&mut self) {
        println!("Dropping HasDrop2");
    }
}

struct HasTwoDrops {
    one: HasDrop1,
    two: HasDrop2
}
impl Drop for HasTwoDrops {
    fn drop(&mut self) {
        println!("Dropping HasTwoDrops");
    }
}

struct Foo;
impl Drop for Foo {
    fn drop(&mut self) {
        println!("Dropping Foo")
    }
}

fn main() {
    let x = HasTwoDrops {
        two: HasDrop2,
        one: HasDrop1
    };
    let foo = Foo
    println!("Running");

In the code above, we implement Drop trait for HasDrop1, HasDrop2, HasTwoDrops, and Foo. The result of this code will be:

Running
Dropping Foo
Dropping HasTwoDrops
Dropping HasDrop1
Dropping HasDrop2

The dropping of variables in a scope follows the reverse order, as x is created before foo and foo is dropped before x. The dropping of struct attributes follows the order in which the attributes are defined. Since HasDrop2 is defined after HasDrop1 (even though assigned first), HasDrop1 is dropped before HasDrop2.

Even if we do not define Drop trait for x, the code would still run as follows, since almost every data type in Rust has implemented Drop by default.

When we want to drop an object manually, we can invoke the drop method of the object. However, we have to write in the form of drop(object), instead of object.drop(), since in the latter case the drop method only uses a mutable borrowing of the object, and the variable is still accessible even when the object it refers to is dropped, leading to a compile error. drop(object) takes the ownership of the variable, so both the object and the variable are invalidated.

fn main() {
    let foo = Foo;
    // foo.drop();    ERROR
    drop(foo);

A final thing to note about is that trait Copy and Drop cannot coexist. Copy implies shallow copy semantics where the original and the copy are independent, while Drop indicates the object needs special cleanup when the value is dropped. The coexistence of Copy and Drop can lead to undefined behaviors such as double-free.

Reference Counting with Rc and Arc

The ownership and borrowing mechanism requires that every reference type object can only have one reference. However, in some cases, we may want an object to have mutiple references. For example, in a graph data structure we may have multiple edges pointing to one node, and in multithreading we want multiple threads to hold one data. The solutions to these problems are

Rc (in a single thread), and Arc (in multithreading).

Rc uses reference counting to keep track of references to the object. The counting pluses one when a new variable refers to Rc object and minus one when a variable is released. The Rc object is released when the counter goes to 0. Using Rc::clone() we can assign a variable to an existing Rc object, and with Rc::strong_count() we can know the number of references to Rc

use std::rc::Rc;
fn main() {
    let a = Rc::new(String::from("test ref counting"));
    println!("count after creating a {}", Rc::strong_count(&a));    // 1
    let b = Rc::clone(&a);
    println!("count after creating b {}", Rc::strong_count(&a));    // 2
    {
        let c = Rc::clone(&a);
        println!("count after creating c {}", Rc::strong_count(&c));    // 3
    }
    println!("count after c goes out of scope {}", Rc::strong_count(&a));    // 2
}

Rc pointer is an immutable reference to the object in heap, since we do not allow the coexistence of multiple mutable references to the object.

Arc (Atomic Rc) has the same API as Rc. However, Arc supports sharing data among multiple threads. Arc is much less efficient than Rc, so we should use Rc when possible unless we're dealing with multithreading.

use std::sync::Arc;
use std::thread;

fn main() {
    let s = Arc::new(String::from("test Arc"));
    for _ in 0..10 {
        let s = Arc::clone(&s);
        let handle = thread::spawn(move || {println!("{}", s)});
    }
}

Internal Mutability with Cell and RefCell

Cell is used when T implements Copy. It has method get() to retrieve the element and set() to change the element. Cell can be used as a workaround of Rust borrowing rule, as in this example:

use std::cell::Cell;
fn main() {
    let c = Cell::new("asdf");
    let one = c.get();
    c.set("qwer");
    let two = c.get();
    println!("{}, {}", one, two);    // asdf, qwer
}

If we use reference in the code above, aftering assigning one to be an immutable borrowing, we cannot change c to another value, since that would lead to coexistence of mutable and immutable borrowing. However, with Cell we are able to change the value even when it has an immutable borrowing.

RefCell is used when T does not implement Copy. It allows mutable and immutable reference to coexist at compile time. However, it cannot ignore the rule as it checks the borrowing errors at runtime instead. RefCell is used when the compiler misidentify the code as violating the borrowing rule.

One of such cases is the implementation of internal mutability: the ability to mutate data even when accessed through immutable borrowing. In the example below, we want to implement a message queue that stores value in cache in send method. However, assume that the send method uses an immutable borrowing &self, we need to use RefCell to mutuate self.msg_cache

use std::cell::RefCell;
pub trait Messenger {
    fn send(&self, msg: String);
}

pub struct MsgQueue {
    msg_cache: RefCell<Vec<String>>
}

impl Messenger for MsgQueue {
    fn send(&self, msg: String) {
        self.msg_cache.borrow_mut().push(msg)
    }
}

fn main() {
    let mq = MsgQueue { msg_cache: RefCell::new(Vec::new()) };
    mq.send("hello, world".to_string());
}

Combining Rc and RefCell

Rc allows multiple ownership of an object, and RefCell allows internal mutability. A common combination in Rust is using Rc and RefCell together:

use std::cell::RefCell;
use std::rc::Rc;
fn main() {
    let s = Rc::new(RefCell::new("123".to_string()));
    let s1 = s.clone();
    let s2 = s.clone();
    s2.borrow_mut().push_str("4");
    println!("{:?}\n{:?}\n{:?}", s, s1, s2);
}

The result of the code above would be:

RefCell { value: "1234" }
RefCell { value: "1234" }
RefCell { value: "1234" }

Note that when one of the reference mutates the object with borrow_mut(), all references refer to the updated object.