Zero-copy is a way to elide CPU copies between the kernel and user space buffers that is particularly useful in high throughput applications like database engines. It makes a huge difference in performance under high load, particularly when your working set is no longer cache resident.

What Is Zero-Copy

Here is what a typical database engine looks like. For this post, focus on two copy boundaries: the OS boundary, and the path from the buffer pool into the layers above it.

  ┌─────────────────────────────────────────────────────────┐
  │                      Query Layer                        │
  └────────────────────────┬────────────────────────────────┘
                           │
  ┌────────────────────────▼────────────────────────────────┐
  │                    Execution Engine                     │
  └────────────────────────┬────────────────────────────────┘
                           │
  ┌────────────────────────▼────────────────────────────────┐
  │                    Transaction Manager                  │
  └──────────┬─────────────┴─────────────────┬──────────────┘
             │                               │
  ┌──────────▼──────────┐       ┌────────────▼────────────┐
  │    Lock Manager     │       │      Log Manager        │
  └─────────────────────┘       └─────────────────────────┘
                           │  fresh copies into higher layers
  ┌────────────────────────▼────────────────────────────────┐
  │                    Buffer Pool                          │
  └────────────────────────┬────────────────────────────────┘
                           │  copy at OS boundary
  ┌────────────────────────▼────────────────────────────────┐
  │                    Disk                                 │
  └─────────────────────────────────────────────────────────┘

Trying to build a high performance engine requires eliding any non-useful work as far as possible and copying data falls squarely in this category.

Think of each copy operation as an equivalent of memcpy()memcpy can actually cause pipeline stalls which is something you want to avoid in high perf applications. which requires the CPU to copy data from a source and put it into a destination. You’re spending cycles on non-essential work and this can cause eviction of hot data from CPU caches.

Here’s a great example of the lifecycle that a typical read or write operation goes through. All those CPU copies are useless work that are burning cycles.

Image taken from https://www.linuxjournal.com/article/6345

Now, let’s focus on eliminating copies at the layer between the buffer pool and disk first.

The Buffer Pool And Direct IO

⊕Here’s a famous tirade from Linus Torvalds on the direct IO interface. He doesn’t like database developers.

The buffer pool opens and stores file descriptors with the open() syscall. When we call read() and write() on those file descriptors it goes through the whole cycle you saw earlier with copies between userspace, kernel and DMA.

An easy win here is to use direct IO with the O_DIRECT flagA large number of modern databases use this approach, although there are holdouts like Postgres.. This will force the application to bypass the OS page cacheThis assumes 64-bit DMA capable hardware. On systems with 32-bit DMA devices or confidential computing VMs (AMD SEV, Intel TDX), the kernel may silently introduce a SWIOTLB bounce buffer, reintroducing a CPU copy. See the Linux kernel docs on swiotlb..

O_DIRECT requires that the buffers submitted are pointer aligned, along with I/O length and file offset. In Rust, we guarantee the former with #[repr(align(4096))] on the buffer holding our page, and 4 KiB page-sized reads and writes at page-aligned offsets satisfy the rest. Without this, O_DIRECT reads or writes would often fail with EINVALHere’s a gist showing this in C — the first program uses malloc (not 4096-aligned) and the write fails, the second uses posix_memalign and succeeds..

Since we’re bypassing the kernel page cache we don’t get useful boosts like readahead or write coalescing but this is exactly why a buffer pool is so important in a database.

The buffer pool is a replacement for the OS page cache designed with specific workloads in mind. It’s always helpful to think of this in terms of mechanism + policy.

Mechanism: A fixed size page table which serves page requests for the layers above and evicts some pages to make room for others.

Policy: A way to decide which page to evict (eviction policy)

  ┌─────────────────────────────────────────────────────────┐
  │                      Query Layer                        │
  │                                                         │
  │   TableScan / IndexScan / BTreeScan                     │
  │   (iterates rows, calls next(), get_value())            │
  └────────────────────────┬────────────────────────────────┘
                           │ uses
  ┌────────────────────────▼────────────────────────────────┐
  │                    Transaction                          │
  │                                                         │
  │   tx_id | ConcurrencyManager | RecoveryManager          │
  └────────────────────────┬────────────────────────────────┘
                           │ pin() / unpin()
  ┌────────────────────────▼────────────────────────────────┐
  │                    Buffer Pool                          │
  │                                                         │
  │  ┌──────────┐  ┌──────────┐  ┌──────────┐               │
  │  │ Frame 0  │  │ Frame 1  │  │ Frame 2  │  ...          │
  │  │          │  │          │  │          │               │
  │  │ file:3   │  │ file:7   │  │  empty   │               │
  │  │ pins: 1  │  │ pins: 0  │  │  pins: 0 │               │
  │  │          │  │          │  │          │               │
  │  │ 4KB data │  │ 4KB data │  │          │               │
  │  └──────────┘  └──────────┘  └──────────┘               │
  │                                                         │
  │  policy: LRU | CLOCK | SIEVE                            │
  └────────────────────────┬────────────────────────────────┘
                           │ 
                           │ 
  ┌────────────────────────▼────────────────────────────────┐
  │                       Disk                              │
  │                                                         │
  │    [ file:1 ]  [ file:3 ]  [ file:7 ]  [ file:9 ]       │
  └─────────────────────────────────────────────────────────┘

Choosing the right policy for your system depends on the characteristics of your workload but most systems typically go with CLOCK which is an LRU approximation. Here’s a non-exhaustive list of replacement policies used in buffer pools.

O_DIRECT removes the copy at the OS boundary. The next problem is avoiding fresh copies once the page is already in the buffer pool.

Eliminating Copies From The Read Path

So far, zero-copy has meant removing copies between the kernel and the buffer pool. From here on, I’m going to broaden it slightly to mean removing redundant copies inside the engine too.

Rust has a great and terrible way to avoid dealing with copies of data - references. It’s great because it’s a single character (&), it’s terrible because now we have to learn to deal with lifetimes.

The simplest way to think about lifetimes is that you are proving to the compiler that any reference held by type A will not outlive the data it points to.

Let’s start with defining the raw bytes for a single page like this

pub struct PageBytes {
    bytes: [u8; PAGE_SIZE_BYTES as usize],
}

Now, we’ll define the data that is held within a single buffer pool frame. The RwLock<T> type here is our page latch.

#[derive(Debug)]
pub struct BufferFrame {
    page: RwLock<PageBytes>,
}

Now, we’re going to store this frame’s data inside a PageReadGuard.

// To keep the example small, the next type is schematic rather than literal.
// I'm using it to show the ownership tradeoff, not the exact implementation
// details of `RwLockReadGuard`.
/// Read guard providing shared access to a pinned page.
pub struct PageReadGuard {
    page: PageBytes,
}

This version is simple to model, but it bakes copying into the design. If every higher-level page object owns its own PageBytes, then constructing those objects from buffer-pool storage means materializing fresh owned values.

What we actually want is not ownership, but a borrowed view into bytes that already live somewhere else. We can model that by introducing a lifetime.

pub struct PageReadGuard<'a> {
    page: &'a PageBytes,
}

With this lifetime annotation, we are proving to the compiler that PageReadGuard will not outlive PageBytes, which means higher-level page objects can become views into existing bytes rather than owned copies.

In the real implementation, the field is RwLockReadGuard<'a, PageBytes> rather than &'a PageBytes, but the ownership story is the same: the guard borrows the page bytes instead of owning them, and our wrapper carries that borrow forward.

pub struct PageReadGuard<'a> {
    page: RwLockReadGuard<'a, PageBytes>
}

Typical database engines have two core page types - heap and btree pages. So let’s focus on the former. The structure of the page is going to be a standard slotted page layout.

At this point, the bytes are already borrowed through PageReadGuard<'a>. Now, the question is where the guard should live and where the parsed references into the page should live.

The most natural thing to try is to keep everything in one struct.

struct HeapPage<'a> {
    guard: PageReadGuard<'a>,
    header: &'a [u8],
    line_pointers: &'a [u8],
    record_space: &'a [u8],
}

This leads us to the classic self-referential struct issue in Rust, which makes pointer invalidation very hard. Imagine for a moment, you have a struct and it has two fields - A and B, with B pointing to A. Now, your struct A moves. What does B point to? It’s going to continue pointing to where A was but that’s invalid and would lead to UB. There are ways around this in Rust with Pin, unsafe raw pointers, Arc pointers and external crates like ouroboros. But, all of these have overhead associated with them.

So, the next thing to try is to separate ownership of the guard from the parsed view into the page.

pub struct HeapPage<'a> {
    guard: PageReadGuard<'a>,
}

pub struct HeapPageView<'a> {
    header: &'a [u8],
    line_pointers: &'a [u8],
    record_space: &'a [u8],
    layout: &'a Layout,
}

let page = HeapPage::new(guard);
let view = HeapPageView::new(&'a page, layout); // borrows from page and page stays alive on stack

This works but runs into an issue when we want to mutate the bytes. Imagine that we want to insert a record into the heap page. This requires mutating record_space and also mutating line_pointers. Now, the bounds of both might have changed which means we have stale references in our struct. We need to drop the struct and re-create it again. While it is cheap, it’s a leaky abstraction.

A better form would be to do the following

pub struct HeapPage<'a> {
    guard: PageReadGuard<'a>,
}

pub struct HeapPageView<'a> {
    header: &'a [u8],
    body_bytes: &'a [u8],
    layout: &'a Layout,
}

let page = HeapPage::new(guard);
let view = HeapPageView::new(&'a page, layout); // borrows from page and page stays alive on stack

In slotted pages, the header size is fixed and anytime we want to perform some operation we re-parse the body bytes using information from the header and get an accurate view into the bytes.

But, the above requires keeping the page and view alive on the stack and leaks implementation details up to the query layer. Once again, a leaky abstraction.

The version I went with flips this around.The old computer science move: “Any problem in computer science can be solved with another level of indirection”. To make that concrete, a slotted page is divided into the header, the line pointers and the record space, so we’ll store those parsed references in HeapPage and keep the PageReadGuard in HeapPageView.

pub struct HeapHeaderRef<'a> {
    bytes: &'a [u8],
}

struct LinePtrBytes<'a> {
    bytes: &'a [u8],
}

struct LinePtrArray<'a> {
    bytes: LinePtrBytes<'a>,
    len: usize,
    capacity: usize,
}

struct HeapRecordSpace<'a> {
    bytes: &'a [u8],
    base_offset: usize,
}

struct HeapPage<'a> {
    header: HeapHeaderRef<'a>,
    line_pointers: LinePtrArray<'a>,
    record_space: HeapRecordSpace<'a>,
}

pub struct HeapPageView<'a> {
    guard: PageReadGuard<'a>,
    layout: &'a Layout,
}

All of them share the exact same lifetime of 'a, which means that any of these references held by another type will not outlive the type.

And all of these types are references into the exact same set of bytes held by PageBytes in the BufferFrame.

PageBytes (owned by BufferFrame)
┌─────────────────────────────────────────────────────┐
│ Header (34 bytes)                                   │
│ page_type | slot_count | free_lower | free_upper .. │
├─────────────────────────────────────────────────────┤
│ Line Pointers (grows →)                             │
│ [ slot 0 ] [ slot 1 ] [ slot 2 ] ...                │
├─────────────────────────────────────────────────────┤
│ Free Space                                          │
├─────────────────────────────────────────────────────┤
│ Record Space (grows ←)                              │
│ ... [ tuple 2 ] [ tuple 1 ] [ tuple 0 ]             │
└─────────────────────────────────────────────────────┘
      │               │                    │
HeapHeaderRef<'a>  LinePtrArray<'a>  HeapRecordSpace<'a>
      │               │                    │
      └───────────────┴────────────────────┘
                      │
                 HeapPage<'a>  ←─────────────────────────┐
                                                         │
                                                  built from
                                                         │
                                            HeapPageView<'a>
                                          ┌──────────────────────────┐
                                          │ guard: PageReadGuard<'a> │
                                          │ layout: &'a Layout       │
                                          └──────────────────────────┘
                                                   │
                                            owns the lock on
                                                   │
                                            PageBytes in BufferFrame

The link between them is a method on HeapPageView that constructs a HeapPage whenever some operation needs an interpreted view of the page bytes.

impl<'a> HeapPageView<'a> {
    fn build_page(&'a self) -> HeapPage<'a> {
        HeapPage::new(self.guard.bytes()).unwrap()
    }

    pub fn row(&self, slot: SlotId) -> Option<LogicalRow<'_>> {
        let view = self.build_page();
        let mut current = slot;
        loop {
            match view.tuple_ref(current)? {
                //  code elided for simplicity
            }
        }
    }
}

impl<'a> HeapPage<'a> {
    fn new(bytes: &'a [u8]) -> SimpleDBResult<Self> {
        // Use shared parsing logic from PageKind trait
        let layout = Self::parse_layout(bytes)?;

        let header = HeapHeaderRef::new(layout.header);

        // Additional heap-specific validation
        let free_upper = header.free_upper() as usize;
        let page_size = PAGE_SIZE_BYTES as usize;
        if free_upper < header.free_lower() as usize || free_upper > page_size {
            return Err("heap page free_upper out of bounds".into());
        }

        let page = Self::from_parts(header, layout.line_ptrs, layout.records, layout.base_offset);
        assert_eq!(
            page.slot_count(),
            header.slot_count() as usize,
            "slot directory length must match header slot_count"
        );
        Ok(page)
    }
}

When the query layer wants to read something from the data pages, it will get a HeapPageView and any operation on the HeapPageView that requires access to some logical segment of data will construct a HeapPage which understands how the bytes on the page are laid out.

Now, you’re probably wondering about the cost of reconstructing the HeapPage each time. It’s actually really cheap because it’s composed entirely of arithmetic operations sprinkled with a few panics in case some invariants aren’t met. And arithmetic operations are extremely cheap for the CPU to perform, especially when compared to memcpy() operations.

Not all CPU operations are created equal. Source: ithare.com, via Andrew Kelley’s talk on Data Oriented Design

Eliminating Copies From The Write Path

In the previous section we saw PageReadGuard, HeapPage and HeapPageView which collectively constitute the read side path of the page. However, Rust has the principle of aliasing XOR mutability and this means that we either get multiple &T or a single &mut T. Everything we saw above is on the &T path and we need a &mut T path.

/// Write guard providing exclusive access to a pinned page.
pub struct PageWriteGuard<'a> {
    page: RwLockWriteGuard<'a, PageBytes>,
}

pub struct HeapPageMut<'a> {
    header: HeapHeaderMut<'a>,
    body_bytes: &'a mut [u8],
}

pub struct HeapPageViewMut<'a> {
    guard: PageWriteGuard<'a>,
    layout: &'a Layout,
}

impl<'a> HeapPageViewMut<'a> {
    fn build_mut_page(&mut self) -> SimpleDBResult<HeapPageMut<'_>> {
        HeapPageMut::new(self.guard.bytes_mut())
    }

    pub fn insert_tuple(&mut self, tuple: &[u8]) -> SimpleDBResult<SlotId> {
        let mut page = self.build_mut_page()?;
        page.insert(tuple)
    }
}

Now, we abide by Rust’s aliasing rules and if we have a BufferFrame, we can either acquire the read latch multiple times and build the read path or acquire the write latch and build the write path.

                    BufferFrame
                    RwLock<PageBytes>
                         │
              ┌──────────┴──────────┐
              │                     │
         read_page()           write_page()
              │                     │
    RwLockReadGuard          RwLockWriteGuard
    (shared, N readers)      (exclusive, 1 writer)
              │                     │
       PageReadGuard           PageWriteGuard
           &'a [u8]               &'a mut [u8]
              │                     │
         HeapPage              HeapPageMut
       (borrows &'a [u8])        (borrows &'a mut [u8])
              │                     │
       HeapPageView<'a>         HeapPageViewMut<'a>

And, of course, Rust’s borrow checker makes use-after-unpin a compile error, not a runtime hazard.

There’s one more asymmetry worth noting. HeapPage splits the bytes into three fields (header, line pointers, record space) because reads are idempotent and those boundaries never shift.

HeapPageMut can’t do the same: a single insert moves both free_lower and free_upper, making any pre-split reference immediately stale. So HeapPageMut keeps a single body_bytes: &mut [u8] and re-derives sub-regions from the header on each operation. Rust prevents aliased mutable access, but the deeper invariant that split points must always match the header has to be enforced through design.

Nested Borrows

If you’ve noticed, so far we’ve been using only a single lifetime of 'a, which makes sense because we’ve been borrowing everything from the same set of underlying bytes.

But, we’ve actually been imprecise and let the compiler handle some of the drudgery for us. The chain of borrows so far has been:

PageBytes  →  RwLock{Read,Write}Guard<'a>  →  Page{Read,Write}Guard<'a>  →  HeapPage[Mut]<'a>  →  HeapPageView[Mut]<'a>

Each successive borrow nests within the previous one but the compiler allows us to do all of this with a single lifetime because of lifetime variance.

The crux of lifetime variance is that shared references are covariant and exclusive references are invariant. Another way this is frequently expressed is that &T is covariant over 'a and &mut T is invariant over 'a.

For covariant lifetimes, Rust can shorten a longer-lived &T into a shorter-lived &T when needed, which allows us to elide the nested lifetimes and let the compiler infer the intermediate lifetimes for us.

Mutable borrows are less forgiving because mutation makes those coercions much stricter. With &mut T, Rust can’t be as flexible about the inner lifetime without risking that a shorter-lived reference gets written into a place that promised to hold a longer-lived one.

Here’s a concrete example that demonstrates the asymmetry. We define a struct with two lifetimes and try to shorten one in a function:

struct Inner<'a> {
    data: &'a [u8],
}

struct Outer<'short, 'long> {
    inner: &'short Inner<'long>,
}

fn shorten<'long, 'short>(outer: Outer<'short, 'long>) -> Outer<'short, 'short> { outer }

This compiles because &T is covariant over 'long. The function can coerce a reference where the data outlives the borrow.

Now, change the outer reference to mutable:

struct Outer<'short, 'long> {
    inner: &'short mut Inner<'long>,
}

fn shorten<'long, 'short>(outer: Outer<'short, 'long>) -> Outer<'short, 'short> { outer }

This fails to compile. &mut T is invariant over its lifetime parameter — it cannot be shortened even when the data outlives the borrow. The lifetimes must match exactly.

This is exactly what happens in HeapPageViewMut::row_mut() and the construction of LogicalRowMut:

pub struct LogicalRowMut<'row, 'page: 'row> {
    view: &'row mut HeapPageViewMut<'page>,
}

impl<'a> HeapPageViewMut<'a> {
    /// Decodes the live tuple at `slot` into a `LogicalRowMut` for editing.
    /// Changes are written back to the page automatically when the returned value is dropped.
    pub fn row_mut<'row>(
        &'row mut self,
        slot: SlotId,
    ) -> SimpleDBResult<Option<LogicalRowMut<'row, 'a>>> {
        //  code elided for simplicity
        Ok(Some(LogicalRowMut {
            view: self,
            slot,
            values,
            layout,
            dirty: false,
        }))
    }
}

'page is the lifetime of the underlying page view, which already borrows the pinned page bytes. 'row is the shorter lifetime of one exclusive edit session on top of that view. The relation 'page: 'row says that the page view has to stay valid for at least as long as the mutable row editor borrowing it.

Page bytes in BufferFrame
┌─────────────────────────────────────────────────────────────┐
│ header │ line pointers │ free space │ tuple 2 │ tuple 1 │… │
└─────────────────────────────────────────────────────────────┘
<---------------------- borrowed for 'page ---------------------->

                                 one call to row_mut()
                                            │
                                            ▼

                       LogicalRowMut<'row, 'page>
                       ┌─────────────────────────┐
                       │ edits one logical row   │
                       └─────────────────────────┘
                       <---- borrowed for 'row --->

Constraint: 'page : 'row

The Cost Of Safe Abstractions

The split into separate read and write types has a real ergonomic cost. HeapPageViewMut doesn’t automatically get HeapPageView’s read methods. In Rust’s standard library, &mut Vec<T> coerces to &Vec<T> automatically via Deref, so mutable references get all immutable methods for free.

This works because Vec<T> is backed by a single raw pointer internally. It implements Deref<Target=[T]> and DerefMut<Target=[T]> to give you either a &[T] or &mut [T]. From one pointer, the borrow checker decides at the call site whether you get a shared or exclusive slice based on how you borrowed it. The unsafe is inside Vec, audited once and invisible to callers.

pub struct Vec<T, #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global> {
    buf: RawVec<T, A>,
    len: usize,
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<T, A: Allocator> ops::Deref for Vec<T, A> {
    type Target = [T];

    #[inline]
    fn deref(&self) -> &[T] {
        self.as_slice()
    }
}

pub const fn as_slice(&self) -> &[T] {
    unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
}

Our design can’t do this. PageReadGuard and PageWriteGuard hold fundamentally different types — RwLockReadGuard and RwLockWriteGuard — that can’t be unified into a single raw pointer. The RwLock enforces the read/write distinction at runtime, so it has to be reflected as two distinct types at compile time. Any read method you want on HeapPageViewMut has to be written explicitly.

This is the tradeoff in Rust API design. There is no unsafe, and there is a clear separation of capabilities, but the ergonomics aren’t as nice as the ones unsafe-backed types get for free.

Conclusion

By the end, the path looks more like this:

  ┌─────────────────────────────────────────────────────────┐
  │                      Query Layer                        │
  └────────────────────────┬────────────────────────────────┘
                           │
  ┌────────────────────────▼────────────────────────────────┐
  │                    Execution Engine                     │
  └────────────────────────┬────────────────────────────────┘
                           │
  ┌────────────────────────▼────────────────────────────────┐
  │                    Transaction Manager                  │
  └──────────┬─────────────┴─────────────────┬──────────────┘
             │                               │
  ┌──────────▼──────────┐       ┌────────────▼────────────┐
  │    Lock Manager     │       │      Log Manager        │
  └─────────────────────┘       └─────────────────────────┘
                           │  borrowed views over page bytes
  ┌────────────────────────▼────────────────────────────────┐
  │                    Buffer Pool                          │
  └────────────────────────┬────────────────────────────────┘
                           │  `O_DIRECT` removes this copy
  ┌────────────────────────▼────────────────────────────────┐
  │                    Disk                                 │
  └─────────────────────────────────────────────────────────┘

O_DIRECT removes the copy between disk and the buffer pool, and the page/view design removes fresh copies above the buffer pool by turning higher-level page objects into borrowed views over bytes that are already pinned in memory.

To me, the interesting part of this design is that it moves ownership into one place, the buffer pool. Everything above that is about views over the same set of bytes. This comes at the price of API ergonomics though. I had to structure the zero-copy page access so that the compiler can see the same invariants I care about. While it litters the code with lifetime annotations it eliminates a certain class of bugs and avoids redundant data movement.

On February 2nd, Conviva’s streaming analytics platform suddenly ground to a crawl but only for one customer. P99 latency spiked without clear reason, pushing our DAG engine to its limits. What started as a puzzling slowdown soon became a deep dive into concurrency pitfalls.

Conviva’s platform is built to handle 5 trillion daily events, powered by a DAG (directed acyclic graph) based analytics engine. Each customer’s logic is compiled into a DAG, running concurrently on a custom actor model built atop Tokio.

This post unpacks how a seemingly innocuous atomic counter in a shared type registry became the bottleneck and what we learned about concurrency, cache lines, and the right data structures for the job. If you use Rust at scale, or plan to, you’ll enjoy this.

Setting The Stage

We intially tried debugging the issue by eliminating the obvious causes - watermarking, inaccurate metrics etc.

Traffic from gateway showing the P99 latency spike

There was some spirited discussion around whether the way the Tokio runtime was scheduling its tasks across physical threads was causing issues but that seemed improbable given that we use an actor system and each DAG processing task runs independently on a specific actor, and it was unlikely that multiple actors were being scheduled onto the same underlying physical thread.

There were additional lines of inquiry around whether HDFS writes were what was causing the lag to build up and eventually causing a backpressure throughout the system. More analysis of more graphs showed increased context switching during the incident but still with no clear evidence of the cause.

Analyzing The Evidence

We were able to reproduce the issue by saving the event data to GCS buckets and replaying this in an environment enabled with perf. This was a relief because at least the issue wasn’t tied to the prod environment, which would have been a nightmare to debug.

We track active sessions across our system, so we have a reasonable measure of how much load our system is under. However, further analysis in the perf environment revealed that while there was a spike in the number of active sessions, those gradually dropped off while the DAG processing time continued to stay high.

Session count tracker and processing time

While this was still puzzling, at least we had a clear indication of where to look - inside our DAG compiler/engine. All the clues pointed to this as being the source of the issue for the P99 latency spike and the backpressure we were seeing throughout the system.

While we knew where to look, this investigation had already taken weeks and things took a turn for the worse when we hit the issue again on February 23rd. However, there was more evidence coming our way about where to look. All Grafana metrics pointed to DAG processing actually being the cause of slowdown. Another interesting graph that came up was this one displaying the jump in context switches during the incident. While it didn’t lead us directly to the root cause at that point, it became important later on as we identified the issue and resolved it because it tied in neatly with our analysis.

Context switches

Recreating The Crime Scene

Thanks to the earlier work in recreating the issue in a perf environment, we were able to generate these flamegraphs that highlighted hot paths in the code. The first one displays the flamegraph during normal traffic and the second one displays the flamegraph during the incident.

Normal traffic flamegraph

Incident traffic flamegraph

In the incident flamegraph, you can clearly see the dreaded wide bars which indicate longer processing time.

Looking carefully at the flamegraph generated during the incident, you can see a very high load for call paths involving AtomicUsize::fetch_sub which was being called from creating and dropping a ReadGuard in flashmap, which we were using as a concurrent hash map. This concurrent hash map was being used as a type registry which was globally shared amongst all DAG’s across our system.

use flashmap::{ReadHandle, WriteHandle};

pub struct TypeRegistry<C: Clone, const N: usize = DEFAULT_N> {
    writer: Mutex<WriteHandle<ShortTypeId, TypeMetadata<C, N>, BuildNoHashHasher<ShortTypeId>>>,
    pub(crate) reader: ReadHandle<ShortTypeId, TypeMetadata<C, N>, BuildNoHashHasher<ShortTypeId>>
}

In the context of this, the earlier graph about the context switches spiking during the incident makes some sense. The ReadGuard in the hot path of the flamegraph was responsible for handling reads from various threads and each thread would increment and decrement the counter.

Now, one important thing about the hashmap in the type registry is that it is almost read-only. It is initialized with some types on start-up and then only updated when a new type is seen but that rarely ever happened. However, on a critical path, it would check to see if the type was already registered which is where the atomic increments and decrements were occurring.

Now, the question was what to do about this. In the flashmap documentation, the performance comparison shows that in the read-heavy scenario dashmap performed better in terms of both latency & throughput. Unfortunately, replacing flashmap with dashmap did nothing to fix the performance problems. In fact, the flamegraphs turned out to be worse in the same situation with dashmap.

Dashmap flamegraph

Finally, we implemented an ArcSwap based solution and the flamegraph improved and the CPU load dropped to 40% in the perf environment.

use arc_swap::ArcSwap;

pub struct TypeRegistry<C: Clone, const N: usize = DEFAULT_N> {
    pub(crate) types:
        ArcSwap<HashMap<ShortTypeId, TypeMetadata<C, N>, BuildNoHashHasher<ShortTypeId>>>,
}

Post Mortem

So, ArcSwap fixed the problem but let’s look at why it fixed the problem.

First, let’s dig into how concurrent hash maps typically operate. Many designs involve mechanisms like counters to track readers and writers, though the specifics can vary. For example, some implementations use a single, shared counter while others employ sharded designs or multiple counters to reduce contention.

For instance, Dashmap uses a sharded design where each shard is a separate HashMap guarded by a RWLock

pub struct DashMap<K, V, S = RandomState> {
    shift: usize,
    shards: Box<[CachePadded<RwLock<HashMap<K, V>>>]>,
    hasher: S,
}

  [Core 1]             [Core 2]             [Core 3]
     |                   |                   |
     | read()            | read()            | read()
     |                   |                   |
     v                   v                   v
|-------------------Shared Read Counter-------------------|
                    (on one cache line)

                           CPU Caches
   ┌────────────┐     ┌────────────┐     ┌────────────┐
   │  Core 1    │     │  Core 2    │     │  Core 3    │
   │  Cache     │<==> │  Cache     │<==> │  Cache     │
   └────────────┘     └────────────┘     └────────────┘

       ↑ Cache line invalidated each time counter is written
       ↑ "Ping-pong" as cache line bounces across cores

⚠️ Every reader updates the same atomic/shared counter
⚠️ Constant inter-core cache line transfers = degraded perf

In cases where the data is guarded by a single, shared counter or resides on the same shard, contention can arise under high loads. This is because every CPU core attempting to increment or decrement the counter causes cache invalidation due to cache coherence. Each modification forces the cache line containing the counter to “ping-pong” between cores, leading to degraded performance. To understand this better, look at this section below from a great PDF titled What every systems programmer should know about concurrency by Matt Kline.

This also ties in with the context switching graph we saw earlier, which showed a spike in context switches during the incident.

Note: If you’re interested in understanding more about hardware caches and their implications, look at this post

Now, let’s contrast this with the approach that ArcSwap uses. ArcSwap follows the read-copy-update (RCU) methodology where:

• readers access the data without locking
• writers create a new copy of the data
• writers atomically swap in the new data
• old data is reclaimed later during a reclamation phase

The ArcSwap repo even has a method called rcu.

ArcSwap avoids cache contention issues for readers that typically crop up when updating a shared read counter with a thread-local epoch counter to track “debt”.

A new version of the data is swapped in using the standard cmp_xchg operation. This marks the beginning of a new epoch, but the data associated with the old epoch isn’t cleaned up until all “debt” is paid off, that is until all readers of the previous epoch have finished.

  [Core 1]             [Core 2]             [Core 3]
     |                   |                   |
     | load()            | load()            | load()
     |                   |                   |
     v                   v                   v

[Thread-Local Epoch 1] [Thread-Local Epoch 2] [Thread-Local Epoch 3]
      (read guard)           (read guard)         (read guard)

      ┌─────────────────────────────┐
      │        ArcSwap<T>          │
      │ ┌────────────────────────┐ │
      │ │ Arc<T>: Current value  │ │ <── atomic ptr (no cache bouncing)
      │ └────────────────────────┘ │
      └─────────────────────────────┘

     Writer swaps in new Arc<T> using atomic store()
     └── Old Arc<T> placed into deferred queue
         └── Only dropped when all read guards released

✅ No shared counters for read
✅ No cache line bouncing
✅ Readers are wait-free and isolated

The big difference between a concurrent hash map and ArcSwap is that ArcSwap requires swapping out the entirety of the underlying data with every write but trades this off with very cheap reads. Writers don’t even have to wait for all readers to finish since a new epoch is created with the new version of data.

Hash maps on the other hand allow updating invidual portions of data in the hash map but this is where it becomes important that we have an almost read-only scenario with a small dataset because the additional overhead of writes with ArcSwap is worth paying here since reads are faster.

Conclusion

Given that we had a situation which was almost read-only with a small dataset, the overhead of a concurrent hash map was not suitable since we had no use case for frequent, granular updates. Trading that for ArcSwap, which is a specialized AtomicRef, something that is designed for occasional swaps where the entire ref is updated, turned out to be a much better fit.

I’ve been trying to understand profiling and performance tooling better, and there didn’t seem to be a better way than to try to write a cache aware hash map.

The hard part here is getting familiar with the tooling and semantics of tools on different platforms (hint: avoid OSX)

Setting Up The Tooling

My only recommendation in this section is to avoid using any platform other than Linux for profiling. perf is a godsend, and Apple is a locked down nightmare of an OS. The closest thing to perf on OSX is dtrace but you need to disable System Integrity Protection (SIP) on Mac for that. This is a very involved process that involves rebooting your Mac. Even then, I don’t think you get access to hardware counters on Mac.

Contrast this with Linux, where I managed to capture all the requirements to get up and running with all profiling metrics in this GitHub issue. Additionally, there is great documentation for both strace and perf on Linux whereas documentation for Instruments is inadequate on OSX.

Building A HashMap

Let’s dig into building the hashmap itself now. Everyone is familiar with the external API surface of a hashmap - 4 methods. There’s a few things we need to track internally too.

struct HashMap {
  data: Vec<T>  //  not specifying what T is here
  size: usize,
  capacity: usize
}

pub fn new(capacity: usize) -> Self
pub fn get(key: K) -> Result<Option<V>>
pub fn insert(key: K, value: V) -> Result<()>
pub fn delete(key: K) -> Result<()>
fn get_load_factor() -> f64
fn resize() -> Result<()>

The load factor determines when the hash map needs to be resized. I’m setting it to 0.7 for all the examples here.

Building a hash map itself isn’t too hard, it’s figuring out how to handle collisions that’s challenging. There’s broadly 2 ways to do this:

• chaining
• open addressing

Chaining

The chained hash map is the most intuitive way to think about handling collisions - convert each entry into a linked list.

Now, you can append onto each entry any time you encounter the same index.

For chaining, the hashmap struct would look something like this

struct Node<K, V> {
  key: K,
  value: V,
  next: Option<Box<Node<K, V>>>
}

struct LinkedList<K,V> {
  head: Option<Box<Node<K, V>>>
}

struct HashMap<K, V> {
  data: Vec<LinkedList<K, V>>,
  size: usize,
  capacity: usize
}

Open Addressing

The open addressing hash map, instead, tries to find the next empty slot in the vector. There’s a couple of different variations here:

• linear probing
• quadratic probing
• double hashing

We’ll use linear probing in this example to keep things simple. The important thing to note is that with open addressing we expect better cache performance, and therefore better overall performance. If you want a detailed explanation of this, look at this SO answer. The tl;dr is that following the pointers in a linked list leads to cache trashing since the nodes are not “clustered” together in memory.

For open addressing the hash map struct would look something like this

enum Entry<K, V> {
  Empty,
  Deleted(K),
  Occupied(K, V)
}

struct HashMap<K, V> {
  data: Vec<Entry<K, V>>,
  size: usize,
  capacity: usize,
}

Cache Hardware

Now, let’s try and understand why we expect better performance from our open addressing implementation vs chaining implementation.

CPU caches come in 3 flavours - L1, L2 and a shared L3 cache. L3 connects out to RAM.

A CPU will first check its L1 cache for some data. Upon failing to find it in L1, it will search L2, then L3 and eventually RAM.

We want the CPU to stay as far up this hierarchy as possible because each cache miss means an extra fetch which wastes cycles.

Each cache has a unit called cache lines - these are the actual storage units. For instance, an L1 cache of size 128 KB with a 64 byte cache line, has 2,048 cache lines. We want our data to fit in these cache lines to prevent additional fetches.

Now, when the program first loads and a cache miss is encountered, the load will reach RAM and attempt to read the data. However, when the memory is read, the principle of spatial locality is followed and some surrounding memory is read as well.

If we are storing our data in a Vec with as few memory hops as possible, this is great because it is unlikely we’ll encounter a cache miss soon. However, if our data structure has a fragmented memory approach, like in Vec<LinkedLinked<T>>, then it is more likely we’ll end up encountering a cache miss again in the near future.

So, what we should expect to see is that the cache performance of the open addressing approach is better than the chaining approach.

Cache Performance

Let’s run our actual program and collect some performance metrics. I wrote some basic benchmarks and a test harness that you can see here and here.

Below are the cache counter stats, followed by CPU stats. For the cache stats, think of the first 2 lines as showing a rough aggregate. The LLC below stands for last level cache which is the L3 cache.

  Performance counter stats for './target/release/hashmap -w load_factor -i chaining':

       762,544,185      cache-references                                                        (66.66%)
       597,421,905      cache-misses                     #   78.35% of all cache refs           (66.66%)
     8,145,732,959      L1-dcache-loads                                                         (66.67%)
       279,075,767      L1-dcache-load-misses            #    3.43% of all L1-dcache accesses   (66.68%)
       144,218,848      LLC-loads                                                               (66.67%)
       106,241,540      LLC-load-misses                  #   73.67% of all LL-cache accesses    (66.66%)

      14.427876578 seconds time elapsed

      13.234905000 seconds user
       1.191811000 seconds sys
  
  Performance counter stats for './target/release/hashmap -w load_factor -i open_addressing':

       654,770,201      cache-references                                                        (66.67%)
       509,578,087      cache-misses                     #   77.83% of all cache refs           (66.68%)
     5,511,667,961      L1-dcache-loads                                                         (66.67%)
       272,838,534      L1-dcache-load-misses            #    4.95% of all L1-dcache accesses   (66.67%)
        91,460,491      LLC-loads                                                               (66.66%)
        59,470,563      LLC-load-misses                  #   65.02% of all LL-cache accesses    (66.65%)

       9.068350211 seconds time elapsed

       7.816879000 seconds user
       1.250820000 seconds sys

  Performance counter stats for './target/release/hashmap -w load_factor -i chaining':

                61      context-switches                 #    4.226 /sec
                 0      cpu-migrations                   #    0.000 /sec
    33,334,941,421      cycles                           #    2.309 GHz
    34,737,785,689      instructions                     #    1.04  insn per cycle
     6,377,780,492      branches                         #  441.853 M/sec
        37,015,167      branch-misses                    #    0.58% of all branches
         14,434.16 msec cpu-clock                        #    1.000 CPUs utilized
         14,434.20 msec task-clock                       #    1.000 CPUs utilized

      14.436026853 seconds time elapsed

      13.282750000 seconds user
       1.151891000 seconds sys

  Performance counter stats for './target/release/hashmap -w load_factor -i open_addressing':

                24      context-switches                 #    2.635 /sec
                 1      cpu-migrations                   #    0.110 /sec
    24,113,035,970      cycles                           #    2.647 GHz
    26,051,667,680      instructions                     #    1.08  insn per cycle
     4,472,654,897      branches                         #  491.004 M/sec
        30,880,291      branch-misses                    #    0.69% of all branches
          9,109.20 msec cpu-clock                        #    1.000 CPUs utilized
          9,109.21 msec task-clock                       #    1.000 CPUs utilized

       9.110056061 seconds time elapsed

       7.864538000 seconds user
       1.244926000 seconds sys

You can find a summary of the important metrics in the table below

The results are surprising to say the least. Open addressing performs better overall as expected, but not for the reasons we expected. We expected better cache performance from open addressing but don’t see that at all.

In fact, the majority of the gains come from having to execute fewer instructions for open addressing. There’s no walking the chain and loading each node from memory in open addressing.

But, why is the cache performance so poor in open addressing? In fact, the cache performance of open addressing is actually worse at the L1 level than chaining and only slightly better at the L3 cache level.

To figure out why we are getting such poor cache performance, we need to understand the memory layout of our data structures and the size of our cache. This will help us figure out how many elements can be placed in a single cache line.

Cache Sizing

Figuring out the size of the cache on Linux is straightforward. You run lscpu | grep "cache" and it gives you the information

zaid-humayun@zaid-humayun-XPS-13-9370:~$ lscpu | grep "cache"
L1d cache:                            128 KiB (4 instances)
L1i cache:                            128 KiB (4 instances)
L2 cache:                             1 MiB (4 instances)
L3 cache:                             8 MiB (1 instance)

This tells me that I have 4 cores, each with an L1 cache (128KiB) and an L2 cache (1MiB) and finally one shared L3 cache (8MiB).

Next, we need to figure out the size of the cache line and it’s easy to do that with cat /proc/cpuinfo | grep "cache_alignment" which shows 64 bytes. So, we have 2,048 cache lines, since 64 bytes * 2048 = 128KiB

Memory Layout

Now, this is where the real meat of the problem is since this is what we control. Let’s look back at our data structures we were using again.

Before we get into this, a quick reminder that a String in Rust is represented by 24 bytes (8 for the heap pointer, 8 for the size & 8 for the capacity).

For chaining, we are using a linked list to store the data. We can figure out the size of the LinkedList struct using std::mem::size_of::<LinkedList>() and when K and V are of type String, this turns out to be 8 bytes. With a cache line of 64 bytes, we can fit 8 entries per cache line. Here, LinkedList itself is an Option<Box<Node<K, V>>>

Here’s a visualisation of how that is being laid out in the cache.

This is good in terms of layout but each entry ends up being a pointer into the heap, so that’s not great because that means we always end up doing an extra fetch for each entry.

For open addressing, we are using an Enum where the memory layout is 48 bytes when K and V are of type String. It’s 48 bytes because the largest variant is a tuple of String and each String type is 24 bytes.

Once we have padding alignment, we can only fit 1 element of our open addressing variant per cache line. This is significantly worse than the chaining layout but each entry here isn’t a pointer into heap memory (well, it is because the type is a String but we’ll get to that).

In fact, since we can fit 8 entries per cache line with the chaining approach, we see a slightly better performance in the L1 cache there (3.53% vs 4.95%). However, you can see the cost of pointer chasing with the total number of loads in the L1 cache (8.1B vs 5.5B). Even though chaining is more compact and cache efficient, the overall cost of pointer chasing ends up dominating the whole operation and open addressing performs better.

Types And Pointer Chasing

So, we know that memory layout of our data structures can have a major performance impact, but so could the actual types we use in the hash map.

Strings inherently require pointer chasing. Regardless of how efficient we are with our memory layout for our struct, having a String in the key and value significantly impacts the cache performance.

However, if we replace the key-value pairs with u64 then for open addressing each entry only takes 24 bytes now instead of 48, which means 2 entries per cache line instead of 1. So, I tried running the same profiling with u64 key-value pairs.

Performance counter stats for './target/release/hashmap -w load_factor -i chaining':

       372,807,778      cache-references                                                        (66.65%)
       252,436,979      cache-misses                     #   67.71% of all cache refs           (66.65%)
     2,380,114,906      L1-dcache-loads                                                         (66.67%)
       140,770,087      L1-dcache-load-misses            #    5.91% of all L1-dcache accesses   (66.69%)
        72,878,214      LLC-loads                                                               (66.68%)
        41,523,145      LLC-load-misses                  #   56.98% of all LL-cache accesses    (66.66%)

       6.273809754 seconds time elapsed

       5.943481000 seconds user
       0.330026000 seconds sys

 Performance counter stats for './target/release/hashmap -w load_factor -i open_addressing':

       217,178,237      cache-references                                                        (66.63%)
       167,946,747      cache-misses                     #   77.33% of all cache refs           (66.64%)
       880,710,536      L1-dcache-loads                                                         (66.68%)
        81,384,415      L1-dcache-load-misses            #    9.24% of all L1-dcache accesses   (66.71%)
        23,916,052      LLC-loads                                                               (66.70%)
        14,262,560      LLC-load-misses                  #   59.64% of all LL-cache accesses    (66.65%)

       2.421592557 seconds time elapsed

       2.128437000 seconds user
       0.293060000 seconds sys

Overall performance is considerably better, which validates our theory about u64 being better for cache performance but we still need to improve on our cache performance with open addressing.

Compact Memory Layout

Let’s design a more compact memory layout, something that will fit more easily into our cache. We’ll call this our open addressing compact variation, since it builds on top of open addressing.

We can try splitting up the state (whether the entry is free, occupied or tombstoned) and the actual data into separate vectors.

pub struct HashMap<K, V>
where
    K: Key,
    V: Value,
{
    status_bits: Vec<u8>,
    entries: Vec<(K, V)>,
    capacity: usize,
    size: usize,
}

Now, when we have K and V as u64, each entry in the entries vector will be a tuple of size 16 bytes, which means 4 of them should fit in one cache line. To denote whether an entry is occupied, deleted or empty we’ll use 2 bits from the status_bits vector for each entry.

This means that we have to load two separate vectors into our cache but we still end up being more efficient. Each byte in the status_bits vector can hold the status of 4 entries, so a single entry of 8 bytes in the cache can hold the status of 32 entries.

Here’s some code that shows the indexing scheme. If you’re interesting in seeing the code in full context, look here.

// 2 bits per entry: 00 = empty, 01 = deleted, 11 = occupied
const EMPTY: u8 = 0b00;
const DELETED: u8 = 0b01;
const OCCUPIED: u8 = 0b11;

fn get_status(&self, index: usize) -> u8 {
    let byte_idx = index / 4;
    let bit_offset = (index % 4) * 2;
    (self.status_bits[byte_idx] >> bit_offset) & 0b11
}

fn set_status(&mut self, index: usize, status: u8) {
    let byte_idx = index / 4;
    let bit_offset = (index % 4) * 2;
    // Clear the two bits
    self.status_bits[byte_idx] &= !(0b11 << bit_offset);
    // Set the new status
    self.status_bits[byte_idx] |= (status & 0b11) << bit_offset;
}

Now, let’s try running our profiling again and see how we fare.

Performance counter stats for './target/release/hashmap -w load_factor -i chaining':

       372,807,778      cache-references                                                        (66.65%)
       252,436,979      cache-misses                     #   67.71% of all cache refs           (66.65%)
     2,380,114,906      L1-dcache-loads                                                         (66.67%)
       140,770,087      L1-dcache-load-misses            #    5.91% of all L1-dcache accesses   (66.69%)
        72,878,214      LLC-loads                                                               (66.68%)
        41,523,145      LLC-load-misses                  #   56.98% of all LL-cache accesses    (66.66%)

       6.273809754 seconds time elapsed

       5.943481000 seconds user
       0.330026000 seconds sys

 Performance counter stats for './target/release/hashmap -w load_factor -i open_addressing':

       217,178,237      cache-references                                                        (66.63%)
       167,946,747      cache-misses                     #   77.33% of all cache refs           (66.64%)
       880,710,536      L1-dcache-loads                                                         (66.68%)
        81,384,415      L1-dcache-load-misses            #    9.24% of all L1-dcache accesses   (66.71%)
        23,916,052      LLC-loads                                                               (66.70%)
        14,262,560      LLC-load-misses                  #   59.64% of all LL-cache accesses    (66.65%)

       2.421592557 seconds time elapsed

       2.128437000 seconds user
       0.293060000 seconds sys

 Performance counter stats for './target/release/hashmap -w load_factor -i open_addressing_compact':

       171,399,933      cache-references                                                        (66.61%)
        86,106,253      cache-misses                     #   50.24% of all cache refs           (66.61%)
       955,677,380      L1-dcache-loads                                                         (66.66%)
        64,703,719      L1-dcache-load-misses            #    6.77% of all L1-dcache accesses   (66.73%)
        24,137,327      LLC-loads                                                               (66.73%)
         6,734,891      LLC-load-misses                  #   27.90% of all LL-cache accesses    (66.66%)

       1.617498666 seconds time elapsed

       1.411503000 seconds user
       0.205927000 seconds sys

The overall cache misses have dropped to 50.24% for the open addressing compact variation, which is a ~50% improvement over open addressing. It also improves on the cache miss rate from the chaining implementation.

The more stark drop is in the LLC case where it has dropped from the high 50% range to 27.90%. That’s a change of over 50%! And, we can see this change reflected in the overall processing time of the load test too, which is a 1.5x improvement!

Conclusion

Being performant with a cache line isn’t just about using a contiguous data structure like a Vector over a non-contiguous data structure like a LinkedList. It matters far more how many elements you can fit into a cache line and how many pointer chases you go through to get to the actual data.

In the case of Strings, it seems we are always doomed to an additional pointer chase which is why u64 as keys is more cache performant. The best you could do here is to inline smaller strings onto the stack with something like smol_str, but this only works for strings upto 23 bytes long.

Thank you to Debasish Ghosh for reviewing a draft of this post.

This post is the third in a series about exploring what exactly an async runtime is, and what async I/O really means. Here is a link to part I where I built a basic future in Rust and polled it to completion with an executor, and here is part II, where I built a simple event loop which uses the kqueue async io interface.

In this post, I’m going to combine learnings from both posts to build a simple, single-threaded async runtime in ~900 lines of Rust code.

All of this started with some fundamental questions I had about what exactly the term “async” means, what fibers, coroutines, green threads etc. really are.

Async Words

I want to first explore the terminology used in async because having a shared vocabulary makes it much easier to have abstract conversations. While it might just sound like jargon or yak-shaving, being able to clearly differentiate things in designs or conversations is critical because abstractions quickly pile on.

The image above gives a high-level overview of the state of async terminology. The broadest classifier of things are coroutines, of which there are two - stackful & stackless.

Stackful coroutines are usually referred to by other names such as fibers or green threads. Stackless coroutines are just state machines under the hood, and they are sometimes called tasks. Both styles of coroutines are sometimes referred to as using the M:N threading model, where M user-space threads or tasks are multiplexed onto N threads of the underlying host system (thanks to King Protty for pointing this out on Twitter).

The primary difference between the two types of coroutines is in the name - stackful have call stacks allocated (similar to OS thread stacks), stackless have no call stacks allocated to them.

Now, there are broadly two classes of schedulers(sometimes referred to as executors or runtimes):

• Pre-emptive
• Co-operative

Pre-emptive schedulers mean that the scheduler is capable of making a coroutine pause at any point during it’s execution. Co-operative schedulers mean that the scheduler is incapable of making a coroutine pause at any point during it’s execution, the coroutine is responsible for pausing at certain points so that it doesn’t block the scheduler.

With stackful coroutines, it is possible to have either a pre-emptive or co-operative scheduler (this is more of a spectrum, not binary). But, with stackless coroutines, you can only ever have a co-operative scheduler. The reason is that stackless coroutines compile down to state machines which don’t have a call stack allocated to them that stores execution information, so you can’t pause them at any point. Stackful coroutines on the other hand do have a call stack allocated to them which stores execution information. This allows the scheduler to stop & resume them at any point.

We’ll dig deeper into stackless coroutines and co-operative schedulers in this post, but if you’re interested in learning more about stackful coroutines, I highly recommend chapter 4 of Asynchronous Programming in Rust.

Stackful & stackless coroutines have a famous allegory associated with them thanks to Bob Nystrom’s What Colour Is Your Function blog post from back in 2015. In the post, he presents the case for why he thinks stackless coroutines are the wrong abstraction to represent concurrency in a language.

Let’s look at some examples - consider the Go code below

package main

import "fmt"

func f(n int) {
  for i := 0; i < 10; i++ {
    fmt.Println(n, ":", i)
  }
}

func main() {
  go f(0)
  var input string
  fmt.Scanln(&input)
}

It kicks off a goroutine in the background and waits for some input from the user. Notice that you don’t have to do anything explicit in terms of making this asynchronous apart from using the go keyword. The functions have no “colour”.

Now, here is the equivalent JavaScript code (I’m using JS here because it’s syntactically easier to parse than Rust). Now, you use the async keyword to denote that some function can run in the background and it needs to have await called against it.

const readline = require('readline').promises;

async function f(n) {
    for (let i = 0; i < 10; i++) {
        console.log(n, ":", i);
        // Simulate some asynchronous work
        await new Promise(resolve => setTimeout(resolve, 100));
    }
}

async function main() {
    // Create a promise that resolves when f(0) completes
    const task = f(0);

    // Set up readline interface
    const rl = readline.createInterface({
        input: process.stdin,
        output: process.stdout
    });

    // Wait for user input
    await rl.question('Press Enter to exit...');

    // Close the readline interface
    rl.close();

    // Ensure f(0) has completed
    await task;
}

(async () => {
    try {
        await main();
    } catch (error) {
        console.error(error);
    }
})();

The biggest downside is that you also need to denote your main function as async now because you can never call an async function from a regular function. This is where your code gets “coloured”. This might not seem like a big deal but consider a situation where you are trying to do simple iteration in your code but need to call an async function there now. Now, your iteration code also requires async against it even if it doesn’t actually do any I/O waiting. It becomes hard to differentiate which parts of your code are actually doing I/O operations, which voids doing function colouring to begin with. Eventually, your code just ends up becoming the colour of async.

In the JS code, you’ll also notice await, which you can think of as the points at which this coroutine is yielding back to the scheduler (it’s similar to the yield keyword used in generators, and generators are similar to async functions).

This is why async/await is considered co-operative - the scheduler has no way of stopping a future/promise in the midst of it’s execution because there is nowhere to save it’s execution information so it can resume later. In the equivalent Go code, there was no yield or await keyword - the language’s compiler allows each thread that is spun up to execute for some fixed amount of time before stopping it, which as of Go 1.19.1 is 10ms(source).

Because of these details, there is a lot more implicit “magic” happening with Go’s runtime but it’s probably safer because it’s harder to shoot yourself in the foot. Conversely with the async/await situation, things are more explicit, but one poorly misplaced synchronous operation between two await operations could block your runtime. Rust uses async/await but doesn’t bundle a runtime with the language and depends on the ecosystem to provide a runtime. This choice providers users more power since they can choose a runtime based on the workload but adds a lot of mental overhead (a recurring theme with Rust).

In a language without “coloured” functions, since each coroutine is allocated a stack to keep track of it’s execution status, there is more overhead which isn’t present for stackless coroutines. Therefore, you should be able to spin up a greater number of stackless coroutines thus providing greater concurrency (this statement is caveated by implementation details of the application itlsef, though. Just because you can spin up more coroutines does not necessarily mean you are actually providing more concurrency and therefore greater throughput).

There is a certain irony to the terminology here - runtimes without “coloured” functions spin up “green” threads, whereas runtimes with “coloured” functions have no colour to their tasks.

Now that we’ve explored the differences between stackful & stackless coroutines, let’s get down to building a simple runtime for stackless coroutines.

State Machines

The most important thing to understand when it comes to stackless coroutines is that they are typically compiled down to state machines which are then “run” to completion.

The important thing to note here is that the moment you use the async keyword in Rust, the compiler transforms that into code that resembles a state machine. Look at this Rust Playground link as an example. It just has the following Rust code

fn main() {
    async {
        println!("Hello world");
    };
}

If you click the view MIR option in the playground, you’ll see the code generated by the compiler for the async block above.

The “run” to completion bit is important above because these state machines are lazy - they need to actually be run, sometimes repeatedly. This is the job of the executor/scheduler.

Taking a more complete example involving Rust futures, consider the code below

async fn example() {
    let a = async_fn_1().await;
    let r = sync_fn();
    let b = async_fn_2().await;
}

which will compile down to something resembling the following

enum ExampleState {
    Start,
    AwaitingFutureA,
    AwaitingFutureB,
    Done
}

struct Example {
    state: ExampleState,
    output: Option<()>,
}

impl Future for Example {
    //  the logic to move the state machine through it's various states
}

Rememeber that Rust futures are stackless coroutines that cannot be pre-empted by the scheduler. If your sync_fn() above takes too long, you are actually blocking this task and the thread of the executor this task is running on.

Building A Runtime

Let’s start writing some code to build a basic single-threaded runtime. Here are the major components we are going to be building:

• tasks (wrappers around futures)
• waker
• reactor
• sync thread pool
• executor

If you’re unsure what any of these terms mean, refer back to parts i and ii of this series

Here’s a great visual representation of what we’re trying to build using a TCP socket as an example.

We’ll tackle this component by component starting with actually representing our futures first.

Tasks

Representing tasks is fairly simple, they are just wrappers around futures with some additional metadata:

pub struct Task {
    pub id: usize,
    pub future: RefCell<Pin<Box<dyn Future<Output = ()> + 'static>>>,
}

impl Display for Task {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "task {}", self.id)
    }
}

impl std::fmt::Debug for Task {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "task {}, ", self.id)
    }
}

If you’re interested in understanding the Pin<Box<T>> verbosity, checkout fasterthanlime’s post. Put very simply, it’s to ensure that the future is not moved from the memory region it’s stored in.

Waker

The waker is probably the component with the most complicated code, mostly because I chose to implement it in unsafe Rust. If you’re following along and trying to build this on your own, you can always use the simpler ArcWake implementation. I chose not to do that mainly because I never intended to make this runtime multi-threaded and I like the masochism of unsafe Rust.

Before I jump into the code for the waker, I need to explain a fat pointer. A pointer is typically one word size since it only holds a memory address. However, a fat pointer is more than one word size since it holds a memory address and some additional data (in this case, a vtable).

Note: If you want to look at a better implementation of the code below, look at the std lib’s implementation

Here’s the code for the waker, and I’ll walk through it after

pub struct MyWaker {
    task: Rc<Task>,
    sender: Sender<Rc<Task>>,
}

impl MyWaker {
    const VTABLE: RawWakerVTable =
        RawWakerVTable::new(Self::clone, Self::wake, Self::wake_by_ref, Self::drop);

    pub fn new(task: Rc<Task>, sender: Sender<Rc<Task>>) -> Waker {
        let pointer = Rc::into_raw(Rc::new(MyWaker { task, sender })) as *const ();
        let vtable = &MyWaker::VTABLE;
        unsafe { Waker::from_raw(RawWaker::new(pointer, vtable)) }
    }

    unsafe fn clone(ptr: *const ()) -> RawWaker {
        let waker = std::mem::ManuallyDrop::new(Rc::from_raw(ptr as *const MyWaker));
        let cloned_waker = Rc::clone(&waker);
        let raw_pointer = Rc::into_raw(cloned_waker);
        RawWaker::new(raw_pointer as *const (), &Self::VTABLE)
    }

    unsafe fn wake(ptr: *const ()) {
        let waker = Rc::from_raw(ptr as *const MyWaker);
        waker.sender.send(Rc::clone(&waker.task)).unwrap();
    }

    unsafe fn wake_by_ref(ptr: *const ()) {
        let waker = &*(ptr as *const MyWaker);
        waker.sender.send(Rc::clone(&waker.task)).unwrap();
    }

    unsafe fn drop(ptr: *const ()) {
        drop(Rc::from_raw(ptr as *const MyWaker));
    }
}

I have a custom MyWaker struct which holds the task and a sender to a channel. Apart from that, there is the implementation of the dynamic virtual table.

The clone method above uses std::mem::ManuallyDrop to ensure that the default destructor for Rc doesn’t run when the scope ends (I spent 6 hours chasing down that bug).

The wake & wake_by_ref methods send the task onto a channel so that it can be polled by the executor.

Task Queue

I didn’t explicitly mention this component in the list above because it’s tiny and could typically be rolled into the executor. But, it’s where tasks are sent to & then read from. Fairly self-explanatory code.

pub struct TaskQueue {
    pub tasks: Vec<Rc<Task>>,
    sender: Sender<Rc<Task>>,
    receiver: Receiver<Rc<Task>>,
}

impl TaskQueue {
    pub fn new() -> Self {
        let (sender, recv) = mpsc::channel();
        Self {
            tasks: Vec::new(),
            sender,
            receiver: recv,
        }
    }

    pub fn sender(&self) -> Sender<Rc<Task>> {
        self.sender.clone()
    }

    pub fn receive(&mut self) {
        while let Ok(task) = self.receiver.try_recv() {
            self.tasks.push(task);
        }
    }

    pub fn pop(&mut self) -> Option<Rc<Task>> {
        self.tasks.pop()
    }

    pub fn len(&self) -> usize {
        self.tasks.len()
    }

    pub fn is_empty(&self) -> bool {
        self.tasks.len() == 0
    }
}

Reactor

I already covered reactors in detail in part ii, so refer to that if you want more detail. The big update to this component is that I now store the wakers that are built for each task, so that the task can be enqueued when the reactor receives an event.

#[derive(Debug)]
pub struct Event {
    pub fd: usize,
    pub readable: bool,
    pub writable: bool,
}

impl Event {
    pub fn none(fd: usize) -> Event {
        Event {
            fd,
            readable: false,
            writable: false,
        }
    }

    pub fn readable(fd: usize) -> Event {
        Event {
            fd,
            readable: true,
            writable: false,
        }
    }

    pub fn writable(fd: usize) -> Event {
        Event {
            fd,
            readable: false,
            writable: true,
        }
    }

    pub fn all(fd: usize) -> Event {
        Event {
            fd,
            readable: true,
            writable: true,
        }
    }
}

#[derive(Debug)]
pub enum InterestType {
    Read,
    Write,
}

pub struct Reactor {
    kqueue_fd: RawFd,
    notifier: (UnixStream, UnixStream),
    readable: HashMap<usize, Vec<Waker>>,
    writable: HashMap<usize, Vec<Waker>>,
}

impl Reactor {
    /// Create a reactor instance
    pub fn new() -> std::io::Result<Self> {
        let kq = unsafe { libc::kqueue() };
        if kq < 0 {
            return Err(std::io::Error::last_os_error());
        }
        let (reader, writer) = UnixStream::pair()?;
        reader.set_nonblocking(true)?;
        writer.set_nonblocking(true)?;
        let reactor = Reactor {
            kqueue_fd: kq.as_raw_fd(),
            notifier: (reader, writer),
            readable: HashMap::new(),
            writable: HashMap::new(),
        };

        reactor.modify(
            reactor.notifier.0.as_raw_fd(),
            Event::readable(reactor.notifier.0.as_raw_fd().try_into().unwrap()),
        )?;
        Ok(reactor)
    }

    /// Function to determine what interests this source has
    fn get_interest(&self, source: usize) -> Event {
        match (
            self.readable.contains_key(&source),
            self.writable.contains_key(&source),
        ) {
            (false, false) => Event::none(source),
            (true, false) => Event::readable(source),
            (false, true) => Event::writable(source),
            (true, true) => Event::all(source),
        }
    }

    /// Function to register interest for a specific source
    pub fn register_interest(
        &mut self,
        source: i32,
        interest: InterestType,
        context: &mut Context,
    ) {
        match interest {
            InterestType::Read => {
                self.readable
                    .entry(source as usize)
                    .and_modify(|v| v.push(context.waker().clone()))
                    .or_insert(vec![context.waker().clone()]);
                self.modify(source, Event::readable(source as usize))
                    .unwrap();
            }
            InterestType::Write => {
                self.writable
                    .entry(source as usize)
                    .and_modify(|v| v.push(context.waker().clone()))
                    .or_insert(vec![context.waker().clone()]);
                self.modify(source, Event::writable(source as usize))
                    .unwrap();
            }
        }
    }

    pub fn get_wakers(&mut self, events: Vec<Event>) -> Vec<Waker> {
        let mut wakers = Vec::new();
        for event in events {
            if event.readable {
                if let Some(readable_wakers) = self.readable.remove(&event.fd) {
                    wakers.extend(readable_wakers);
                }
            } else if event.writable {
                if let Some(writable_wakers) = self.writable.remove(&event.fd) {
                    wakers.extend(writable_wakers);
                }
            }
        }
        wakers
    }

    pub fn waiting_on_events(&self) -> bool {
        if self.readable.is_empty() && self.writable.is_empty() {
            return false;
        }
        true
    }

    /// Function to accept the source to register an interest in and the type of interest
    pub fn add(&mut self, source: RawFd) -> std::io::Result<()> {
        self.modify(source, self.get_interest(source as usize))
    }

    /// A helper notify method to unblock the scheduler
    pub fn notify(&mut self) -> std::io::Result<usize> {
        self.notifier.1.write(&[1])
    }

    /// The function that removes interest for a file descriptor with the actual underlying syscall
    pub fn remove(&mut self, fd: RawFd, ev: Event) -> std::io::Result<()> {
        self.readable.remove(&(fd as usize));
        self.writable.remove(&(fd as usize));
        let registered_interest = self.get_interest(fd as _);
        let mut changelist = Vec::new();
        if ev.readable && registered_interest.readable {
            changelist.push(kevent {
                ident: fd as _,
                filter: EVFILT_READ,
                flags: EV_DELETE,
                fflags: 0,
                data: 0,
                udata: ev.fd as *mut c_void,
            });
        }
        if ev.writable && registered_interest.writable {
            changelist.push(kevent {
                ident: fd as _,
                filter: EVFILT_WRITE,
                flags: EV_DELETE,
                fflags: 0,
                data: 0,
                udata: ev.fd as *mut c_void,
            });
        }

        if changelist.is_empty() {
            return Ok(());
        }

        let result = unsafe {
            kevent(
                self.kqueue_fd,
                changelist.as_mut_ptr(),
                changelist.len() as i32,
                std::ptr::null_mut(),
                0,
                std::ptr::null(),
            )
        };
        if result < 0 {
            error!(
              "There was an error while attempting to modify the kqueue list for {} for event {:?}",
              fd, ev
            );
            return Err(std::io::Error::last_os_error());
        }
        Ok(())
    }

    /// The function that registers interest with the actual underlying syscall
    fn modify(&self, fd: RawFd, ev: Event) -> std::io::Result<()> {
        debug!("Adding file {} for event {:?} to reactor", fd, ev);
        let mut changelist = Vec::new();
        if ev.readable {
            changelist.push(kevent {
                ident: fd as _,
                filter: EVFILT_READ,
                flags: EV_ADD | EV_ONESHOT,
                fflags: 0,
                data: 0,
                udata: ev.fd as *mut c_void,
            });
        }

        if ev.writable {
            changelist.push(kevent {
                ident: fd as _,
                filter: EVFILT_WRITE,
                flags: EV_ADD | EV_ONESHOT,
                fflags: 0,
                data: 0,
                udata: ev.fd as *mut c_void,
            });
        }

        if changelist.is_empty() {
            return Ok(());
        }

        let result = unsafe {
            kevent(
                self.kqueue_fd,
                changelist.as_ptr(),
                changelist.len() as i32,
                std::ptr::null_mut(),
                0,
                std::ptr::null(),
            )
        };

        if result < 0 {
            return Err(std::io::Error::last_os_error());
        }
        Ok(())
    }

    /// Blocking poll function to get events
    pub fn poll(&mut self) -> std::io::Result<Vec<Event>> {
        let mut events: Vec<libc::kevent> = Vec::new();
        let result = unsafe {
            events.resize(1, std::mem::zeroed());
            kevent(
                self.kqueue_fd,
                std::ptr::null(),
                0,
                events.as_mut_ptr(),
                1,
                std::ptr::null(),
            )
        };

        if result < 0 {
            return Err(std::io::Error::last_os_error());
        }

        let mapped_events: std::io::Result<Vec<_>> = events
            .iter()
            .map(|event| {
                let ident = event.ident;
                let filter = event.filter;

                if ident == self.notifier.0.as_raw_fd().try_into().unwrap() {
                    let mut buf = [0; 8];
                    self.notifier.0.read(&mut buf)?;
                    self.modify(
                        self.notifier.0.as_raw_fd(),
                        Event::readable(self.notifier.0.as_raw_fd().try_into().unwrap()),
                    )?;
                }

                let event = Event {
                    fd: ident,
                    readable: filter == EVFILT_READ,
                    writable: filter == EVFILT_WRITE,
                };
                Ok(event)
            })
            .collect();
        info!("Received events {:?}", mapped_events);

        mapped_events
    }
}

The code above creates a kqueue file descriptor and then starts listening for events in one-shot mode on specific file descriptors. kqueue is only available on OSX & BSD distributions, the equivalent for Linux would be epoll but they provide a similar API.

The code also uses a Unix pipe hack to unblock the scheduler (I’ll show that below) since I just block the scheduler while listening for events from kqueue. I imagine that a production grade scheduler would typically provide a timeout.

Pay attention to the register_interest and get_interest methods here that update the readable & writable hash maps, this is where the wakers are stored.

Sync Thread Pool

I mentioned earlier that we typically don’t want to block within our async code, but sometimes you need to perform a blocking operation, like reading from a file when io_uring isn’t available.

In these situations, it’s helpful to have a separate thread pool which can be used to run these blocking tasks in a non-blocking manner. This frees up the executor to run other tasks and check with this task periodically.

use super::reactor::Reactor;

pub struct FileIOPool {
    sender: Sender<FileReaderTask>,
}

pub struct FileReaderTask {
    pub path: PathBuf,
    pub responder: Sender<std::io::Result<Vec<u8>>>,
}

impl FileIOPool {
    pub fn new(num_threads: usize, shutdown_rx: Receiver<()>) -> Self {
        let (sender, receiver) = channel::<FileReaderTask>();
        let recv = Arc::new(Mutex::new(receiver));
        let shutdown_rx = Arc::new(Mutex::new(shutdown_rx));
        for _ in 0..num_threads {
            let recv_clone = Arc::clone(&recv);
            let shutdown_rx_clone = Arc::clone(&shutdown_rx);
            std::thread::spawn(move || loop {
                let task = recv_clone.lock().unwrap().try_recv();
                let shutdown_signal = shutdown_rx_clone.lock().unwrap().try_recv();
                match (task, shutdown_signal) {
                    (Ok(task), _) => {
                        let result = std::fs::read(task.path);
                        let _ = task.responder.send(result);
                    }
                    (_, Ok(())) => {
                        debug!("File io pool received shutdown signal, shutting down");
                        break;
                    }
                    (Err(_), Err(_)) => {}
                }
            });
        }
        Self { sender }
    }

    pub fn read_file(&self, path: PathBuf, reactor: Rc<RefCell<Reactor>>) -> ReadFileFuture {
        let (file_completion_sender, file_completion_recv) =
            std::sync::mpsc::channel::<std::io::Result<Vec<u8>>>();
        let file_reader_task = FileReaderTask {
            path,
            responder: file_completion_sender,
        };
        self.sender
            .send(file_reader_task)
            .expect("Error while sending the file reader task to io pool");
        ReadFileFuture {
            reactor,
            receiver: file_completion_recv,
        }
    }
}

pub struct ReadFileFuture {
    reactor: Rc<RefCell<Reactor>>,
    receiver: Receiver<std::io::Result<Vec<u8>>>,
}

impl Future for ReadFileFuture {
    type Output = std::io::Result<Vec<u8>>;
    fn poll(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        match self.receiver.try_recv() {
            Ok(data) => return std::task::Poll::Ready(data),
            Err(e) => match e {
                TryRecvError::Empty => {
                    debug!("received empty from the file reader task receiver");
                    self.reactor.borrow_mut().notify()?; //  force the reactor to wake up so scheduler can continue
                    cx.waker().wake_by_ref();
                    return std::task::Poll::Pending;
                }
                TryRecvError::Disconnected => {
                    return std::task::Poll::Ready(Err(std::io::Error::new(
                        std::io::ErrorKind::Other,
                        "Received Disconnected while waiting for file read to complete",
                    )))
                }
            },
        }
    }
}

Every time the thread pool receives a task to read a file, it returns a future which can be polled to check whether the file reading has completed.

Executor

Finally, the center piece of all of this - the executor. There’s a lot of code here but it’s a simple component to understand. It’s simpler to think of this as a combination of 2 components:

• event loop
• futures poller

In the “hot” loop of this component, it listens for events from the reactor and also checks if the task queue has any tasks on it to poll, and it does this combination every iteration.

pub struct Executor {
    task_queue: Rc<RefCell<TaskQueue>>,
    reactor: Rc<RefCell<Reactor>>,
    monotonic_clock: Mutex<usize>,
}

impl Executor {
    pub fn new(
        task_queue: Rc<RefCell<TaskQueue>>,
        reactor: Rc<RefCell<Reactor>>,
    ) -> std::io::Result<Self> {
        Ok(Self {
            task_queue,
            reactor,
            monotonic_clock: Mutex::new(0),
        })
    }

    pub fn block_on<F>(&self, future: F) -> F::Output
    where
        F: Future<Output = ()> + 'static,
    {
        let task = Task {
            id: *self.monotonic_clock.lock().unwrap(),
            future: RefCell::new(Box::pin(future)),
        };
        *self.monotonic_clock.lock().unwrap() += 1;
        self.task_queue
            .borrow()
            .sender()
            .send(Rc::new(task))
            .unwrap();
        self.run();
    }

    pub fn spawn<F>(&self, future: F) -> F::Output
    where
        F: Future<Output = ()> + 'static,
    {
        let task = Task {
            id: *self.monotonic_clock.lock().unwrap(),
            future: RefCell::new(Box::pin(future)),
        };
        *self.monotonic_clock.lock().unwrap() += 1;
        self.task_queue
            .borrow()
            .sender()
            .send(Rc::new(task))
            .unwrap();
        self.reactor.borrow_mut().notify().unwrap();
    }

    fn run(&self) {
        loop {
            self.task_queue.borrow_mut().receive();
            loop {
                let task = {
                    if let Some(task) = self.task_queue.borrow_mut().pop() {
                        task
                    } else {
                        break;
                    }
                };

                let waker = MyWaker::new(Rc::clone(&task), self.task_queue.borrow().sender());
                let mut context = Context::from_waker(&waker);
                match task.future.borrow_mut().as_mut().poll(&mut context) {
                    std::task::Poll::Ready(_output) => {
                        debug!(
                            "The future for task {} has completed and returned on thread {:?}",
                            task.id,
                            std::thread::current().id()
                        );
                    }
                    std::task::Poll::Pending => {
                        debug!(
                            "The future for task {} is pending on thread {:?}",
                            task.id,
                            std::thread::current().id()
                        );
                    }
                };
            }

            self.task_queue.borrow_mut().receive();
            if !self.reactor.borrow().waiting_on_events() && self.task_queue.borrow().is_empty() {
                debug!("no events to wait on and no events in the queue, so breaking out");
                break;
            }

            if self.reactor.borrow().waiting_on_events() {
                debug!("waiting on events from the reactor");
                match self.wait_for_io() {
                    Ok(events) => self.wake_futures_on_io(events),
                    Err(e) => {
                        if e.kind() == std::io::ErrorKind::Interrupted {
                            break;
                        }
                        eprintln!("Error while waiting for IO events :{}", e);
                    }
                }
            }
        }
    }

    fn wait_for_io(&self) -> std::io::Result<Vec<Event>> {
        self.reactor.borrow_mut().poll()
    }

    fn wake_futures_on_io(&self, events: Vec<Event>) {
        let wakers = self.reactor.borrow_mut().get_wakers(events);
        let _ = wakers
            .into_iter()
            .map(|waker| waker.wake())
            .collect::<Vec<_>>();
    }
}

It has two public API methods - block_on and spawn. These are the equivalent of the Tokio methods and do roughly the same thing. block_on is used for the top level future and blocks the main thread waiting for this future to complete, and spawn is used for nested futures.

You’ll notice that all my futures have an output of the unit type and have a static lifetime. The former is done to simplify the code since this runtime is only going to service a web server, which typically write their results back via the TCP stream.

I did the latter because it was the easiest way to set up everything without having to worry about lifetimes. If you’re interested in a way that futures can be built without any lifetimes at all, read this great post by Evan Schwartz.

And with that, we’re done with the main internal components of our system.

Custom Futures

Now, that we have internal components mainly set up, let’s switch focus over to writing some custom futures which can be run on the executor.

We’re going to be building a basic web server on top of this, so let’s buid a simple TCP listener & client.

TCP Listener

Here’s the code for the listener. Notice that it registers interest with the reactor upon recognizing that there are no clients attempting to connect right now. Also, notice the drop functionality where it unregisters itself from the reactor.

The listener.set_nonblocking(true) is important here because otherwise the connect function call will block on the async task, rendering the entire runtime pointless. When the listener notices that there are no clients trying to connect, it registers interest with the reactor and unblocks the scheduler.

pub struct TcpListener {
    listener: std::net::TcpListener,
    reactor: Rc<RefCell<Reactor>>,
}

impl TcpListener {
    pub fn bind(addr: &str, reactor: Rc<RefCell<Reactor>>) -> std::io::Result<TcpListener> {
        let listener = std::net::TcpListener::bind(addr)?;
        listener.set_nonblocking(true)?;
        Ok(Self { listener, reactor })
    }

    pub fn accept(&self) -> std::io::Result<ListenerFuture> {
        Ok(ListenerFuture {
            listener: &self.listener,
            reactor: Rc::clone(&self.reactor),
        })
    }
}

impl Drop for TcpListener {
    fn drop(&mut self) {
        self.reactor
            .borrow_mut()
            .remove(
                self.listener.as_raw_fd(),
                Event::all(self.listener.as_raw_fd() as _),
            )
            .unwrap();
    }
}

pub struct ListenerFuture<'listener> {
    listener: &'listener std::net::TcpListener,
    reactor: Rc<RefCell<Reactor>>,
}

impl Future for ListenerFuture<'_> {
    type Output = std::io::Result<(std::net::TcpStream, std::net::SocketAddr)>;
    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        debug!("Received a poll call on the ListenerFuture");
        match self.listener.accept() {
            Ok((stream, addr)) => std::task::Poll::Ready(Ok((stream, addr))),
            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                debug!("Attempting to accept a connection on the listener WouldBlock, so registering with reactor and yielding control back to the executor");
                let fd = self.listener.as_raw_fd();
                self.as_mut().reactor.borrow_mut().register_interest(
                    fd,
                    super::reactor::InterestType::Read,
                    cx,
                );
                std::task::Poll::Pending
            }
            Err(e) => {
                eprintln!("received an error in the ListenerFuture {}", e);
                std::task::Poll::Ready(Err(e))
            }
        }
    }
}

TCP Client

The TCP client is a little more involved but it’s useful to break this down into a set of smaller components as well. While handling a request:

• the client reads bytes from the stream
• reads a response file from the local fs using the threadpool
• writes a response back to the stream

Here’s the code

pub struct TcpClient {
    client: TcpStream,
    _addr: SocketAddr,
    reactor: Rc<RefCell<Reactor>>,
    file_io_pool: Rc<RefCell<FileIOPool>>,
}

impl TcpClient {
    pub fn new(
        client: TcpStream,
        addr: SocketAddr,
        reactor: Rc<RefCell<Reactor>>,
        file_io_pool: Rc<RefCell<FileIOPool>>,
    ) -> Self {
        client.set_nonblocking(true).unwrap();
        Self {
            client,
            _addr: addr,
            reactor,
            file_io_pool,
        }
    }

    pub async fn handle_request(&mut self) -> std::io::Result<()> {
        debug!("handling client request");
        self.read().await?;
        let file_path = PathBuf::from("hello.html");
        let bytes = self
            .file_io_pool
            .borrow()
            .read_file(file_path, Rc::clone(&self.reactor))
            .await
            .expect("Error while reading the response file");
        let mut response = Vec::new();
        let headers = format!(
            "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
            bytes.len()
        );
        response.extend_from_slice(headers.as_bytes());
        response.extend_from_slice(&bytes);
        self.write(response).await?;
        self.client.shutdown(std::net::Shutdown::Write)?;
        Ok(())
    }

    fn read(&self) -> AsyncTcpReader {
        AsyncTcpReader {
            client: &self.client,
            reactor: Rc::clone(&self.reactor),
            buffer: Vec::with_capacity(1024),
            total_read: 0,
        }
    }

    fn write(&self, buffer: Vec<u8>) -> AsyncTcpWriter {
        AsyncTcpWriter {
            client: &self.client,
            reactor: Rc::clone(&self.reactor),
            buffer,
            bytes_written: 0,
        }
    }
}

impl Drop for TcpClient {
    fn drop(&mut self) {
        self.reactor
            .borrow_mut()
            .remove(
                self.client.as_raw_fd(),
                Event::all(self.client.as_raw_fd().try_into().unwrap()),
            )
            .unwrap();
    }
}

pub struct AsyncTcpReader<'read_stream> {
    client: &'read_stream TcpStream,
    reactor: Rc<RefCell<Reactor>>,
    buffer: Vec<u8>,
    total_read: usize,
}

impl<'read_stream> Future for AsyncTcpReader<'read_stream> {
    type Output = std::io::Result<usize>;
    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        loop {
            let mut chunk = [0u8; 1024];
            match self.client.read(&mut chunk) {
                Ok(0) => {
                    // received EOF
                    self.buffer.clear();
                    return std::task::Poll::Ready(Ok(self.total_read));
                }
                Ok(n) => {
                    self.buffer.extend_from_slice(&chunk);
                    self.total_read += n;
                    let headers = std::str::from_utf8(&chunk[..n]).unwrap();
                    if headers.ends_with("\r\n\r\n") {
                        self.buffer.clear();
                        return std::task::Poll::Ready(Ok(self.total_read));
                    }
                    return std::task::Poll::Pending;
                }
                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                    self.reactor.borrow_mut().register_interest(
                        self.client.as_raw_fd(),
                        InterestType::Read,
                        cx,
                    );
                    return std::task::Poll::Pending;
                }
                Err(e) => return std::task::Poll::Ready(Err(e)),
            }
        }
    }
}

pub struct AsyncTcpWriter<'write_stream> {
    client: &'write_stream TcpStream,
    reactor: Rc<RefCell<Reactor>>,
    buffer: Vec<u8>,
    bytes_written: usize,
}

impl<'write_stream> Future for AsyncTcpWriter<'write_stream> {
    type Output = std::io::Result<usize>;
    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        let this = self.as_mut().get_mut();
        loop {
            match this.client.write(&this.buffer[this.bytes_written..]) {
                Ok(n) => {
                    this.bytes_written += n;
                    if this.bytes_written >= this.buffer.len() {
                        return std::task::Poll::Ready(Ok(this.bytes_written));
                    }
                }
                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                    this.reactor.borrow_mut().register_interest(
                        this.client.as_raw_fd(),
                        InterestType::Write,
                        cx,
                    );
                    return std::task::Poll::Pending;
                }
                Err(e) => {
                    return std::task::Poll::Ready(Err(e));
                }
            }
        }
    }
}

Notice that because we have a threadpool via which to read files, all of the 3 tasks above become asynchronous.

Wiring It All Up

Now, that we have all the different components, we need to wire everything up together.

Here’s the main function where we’ll build the top-level future and then pass it along to the executor. It’s quite ugly because of the excessive Rc syntax, but that’s just Rust I suppose.

use std::{cell::RefCell, rc::Rc};

use log::debug;
use timer_future::async_runtime::{
    client::TcpClient, executor::Executor, file_io_pool::FileIOPool, listener::TcpListener,
    reactor::Reactor, task_queue::TaskQueue,
};

fn main() {
    env_logger::init();
    let task_queue = Rc::new(RefCell::new(TaskQueue::new()));
    let reactor = Rc::new(RefCell::new(Reactor::new().unwrap()));
    let (shutdown_tx, shutdown_rx) = std::sync::mpsc::channel();
    setup_ctrlc_handler(shutdown_tx);
    let file_io_pool = Rc::new(RefCell::new(FileIOPool::new(5, shutdown_rx)));
    let runtime = Rc::new(Executor::new(task_queue, Rc::clone(&reactor)).unwrap());
    let runtime_clone = Rc::clone(&runtime);
    runtime.block_on(async move {
        let listener = TcpListener::bind("localhost:8000", Rc::clone(&reactor)).unwrap();
        while let Ok((client, addr)) = listener.accept().unwrap().await {
            debug!("Received a client connection from client {:?}", client);
            let reactor_clone = Rc::clone(&reactor);
            let file_io_pool_clone = Rc::clone(&file_io_pool);
            runtime_clone.spawn(async move {
                let mut tcp_client =
                    TcpClient::new(client, addr, reactor_clone, Rc::clone(&file_io_pool_clone));
                tcp_client
                    .handle_request()
                    .await
                    .expect("Error occurred while handling the tcp request");
            });
            debug!("Handed off client connection to executor");
        }
    });
    println!("Done executing the top level future");
}

fn setup_ctrlc_handler(shutdown_tx: std::sync::mpsc::Sender<()>) {
    ctrlc::set_handler(move || {
        shutdown_tx
            .send(())
            .expect("failed to send shutdown signal");
    })
    .unwrap();
}

There’s some additional functionality around shutting down the entire runtime by listening for a SIGINT, but the meat of the code is in the block_on function. The connect loop listens for connections, and upon receiving one immediately spawns a separate task for the client.

That was a lot of code, so let me leave you with another diagram illustrating what we’ve built. It’s very similar to the diagram I shared earlier with some more detail attached.

It looks a lot more complicated but it’s just showing more of the detail in the system, the core components remain the same. I’m showing multiple cores being utilized in the diagram, whereas we just utilized one core.

Conclusion

So, there we have it - a single-threaded async runtime which shows how to use async IO to wait on events and poll tasks via a scheduler in ~900 lines of Rust code. Building this prototype was a great way for me to intuit the core ideas behind an async runtime for myself.

Of course, there’s a lot we skipped over here for the sake of simplicity:

• most production schedulers are multi-threaded with tasks being multiplexed across these threads (for our example, we’d end up having to use Arc instead of Rc)
• timers in an async runtime & cancellable tasks
• work stealing across threads, something Tokio famously does
• thread-per-core model with tasks pinned to specific threads

I hope this post gives you a better idea of what’s going on under the hood of an async runtime.

References

1. Async I/O in Depth video series
2. Async Rust Without Send, Sync Or Static
3. Asynchronous Programming In Rust

This post is a follow up to a previous post, where I built a basic custom Future and hooked that up to an Executor and polled the future to completion.

This time I’m going to build a single-threaded event loop that uses async I/O interfaces in ~500 lines of Rust code.

Reactors

Async runtimes have an important component called reactors. These components sit on top of the OS and listen for events from the underlying file descriptors. This is what allows an application to offload a lot of the work.

For instance, consider that your application is listening on a TCP socket. Without an async runtime, you’d typically block on the main thread and listen for connections. Then, as you receive each connection, you’d offload the incoming request onto a separate thread or pass it to a thread pool.

This works fine, but you can make use of certain syscalls to improve your throughput with fewer number of threads (even a single thread).

Let’s start with building a simple reactor.

I’ve used the libc crate to allow a FFI to C so that syscalls can be made. It’s a simple design, just stores the events and has a notification mechanism built on Unix streams.

pub struct Reactor {
    kq: RawFd, //  the kqueue fd
    events: Vec<libc::kevent>,
    capacity: usize,
    notifier: (UnixStream, UnixStream),
}

The reactor exposes the following API:

• add to start watching new file descriptors
• delete to stop watching file descriptors
• notify to unblock the reactor
• poll to ask for events from the OS in a blocking manner

The poll is the most important method because that’s where the blocking routine comes in. To implement this we’re going to have to make use I/O multiplexing systems like kqueue (macOS), epoll(Linux) and IOCP(Windows). This is all part of the “async i/o” umbrella that I’m trying to demystify for myself.

The basic idea behind these interfaces is that you can use a single application level thread to monitor all the I/O sources you are interested in for readiness.

The implementation for the Reactor is provided below.

impl Reactor {
    pub fn new() -> std::io::Result<Self> {
        let kq = unsafe { libc::kqueue() };
        if kq < 0 {
            return Err(std::io::Error::last_os_error());
        }
        let (read_stream, write_stream) = UnixStream::pair()?;
        read_stream.set_nonblocking(true)?;
        write_stream.set_nonblocking(true)?;
        let reactor = Self {
            kq,
            events: Vec::new(),
            capacity: 1,
            notifier: (read_stream, write_stream),
        };
        reactor.modify(
            reactor.notifier.0.as_raw_fd(),
            Event::readable(reactor.notifier.0.as_raw_fd()),
        )?;
        Ok(reactor)
    }

    pub fn add(&mut self, fd: RawFd, ev: Event) -> std::io::Result<()> {
        self.modify(fd, ev)
    }

    pub fn delete(&mut self, fd: RawFd) -> std::io::Result<()> {
        self.modify(fd, Event::none(fd))
    }

    pub fn notify(&mut self) -> std::io::Result<()> {
        self.notifier.1.write(&[1])?;
        Ok(())
    }

    fn modify(&self, fd: RawFd, ev: Event) -> std::io::Result<()> {
        let read_flags = if ev.readable {
            EV_ADD | EV_ONESHOT
        } else {
            EV_DELETE
        };
        let changes = [kevent {
            ident: fd as usize,
            filter: EVFILT_READ,
            flags: read_flags,
            fflags: 0,
            data: 0,
            udata: std::ptr::null_mut(),
        }];
        let result = unsafe {
            kevent(
                self.kq,
                changes.as_ptr(),
                1,
                std::ptr::null_mut(),
                0,
                std::ptr::null_mut(),
            )
        };
        if result < 0 {
            return Err(std::io::Error::last_os_error());
        }
        Ok(())
    }

    pub fn poll(&mut self) -> std::io::Result<Vec<Event>> {
        let max_capacity = self.capacity as c_int;
        self.events.clear();
        let result = unsafe {
            self.events
                .resize(max_capacity as usize, std::mem::zeroed());
            let result = kevent(
                self.kq,
                std::ptr::null(),
                0,
                self.events.as_mut_ptr(),
                max_capacity,
                std::ptr::null(),
            );
            result
        };
        if result < 0 {
            //  return the last OS error
            return Err(std::io::Error::last_os_error());
        }

        let mut mapped_events = Vec::new();
        for i in 0..result as usize {
            let kevent = &self.events[i];
            let ident = kevent.ident;
            let filter = kevent.filter;

            let mut buf = [0; 8];
            if ident == self.notifier.0.as_raw_fd() as usize {
                self.notifier.0.read(&mut buf)?;
                self.modify(
                    self.notifier.0.as_raw_fd(),
                    Event::readable(self.notifier.0.as_raw_fd()),
                )?;
            }

            let mut event = Event {
                fd: ident,
                readable: false,
                writable: false,
            };

            match filter {
                EVFILT_READ => event.readable = true,
                EVFILT_WRITE => event.writable = true,
                _ => {}
            };
            self.modify(event.fd as i32, Event::readable(event.fd as i32))?;
            mapped_events.push(event);
        }
        Ok(mapped_events)
    }
}

The notify method is an interesting little hack because it writes dummy data into one end of the Unix socket, and because we tell kqueue we are interested in the other end of the stream, we have a way to unblock the kevent syscall in the poll method.

Event Loops

In the previous post, we discussed Executors which were a component responsible for polling futures to completion. In the context of asynchronous I/O, you need an analogous event loop which is constantly monitoring for I/O events and then reacting based on those events.

Let’s build a dead simple event loop which will wait on I/O events and also handle a small set of tasks. To handle the tasks, we’ll create a task queue component and feed that into our event loop to read from.

We’ll have 3 types of tasks - registration tasks, unregistration tasks and scheduled tasks.

The scheduled tasks are a way of some object telling the event loop that it wants to be polled, which is where the Unix pipe hack comes in.

pub struct RegistrationTask {
    pub fd: usize,
    pub reference: Box<dyn EventHandler>,
}

pub struct UnregistrationTask {
    pub fd: usize,
}

pub struct ScheduledTask {
    pub fd: usize,
}

pub enum Task {
    RegistrationTask(RegistrationTask),
    UnregistrationTask(UnregistrationTask),
    ScheduledTask(ScheduledTask),
}

impl Display for Task {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Task::RegistrationTask(task) => write!(f, "RegistrationTask: fd {}", task.fd),
            Task::UnregistrationTask(task) => write!(f, "UnregistrationTask: fd {}", task.fd),
            Task::ScheduledTask(task) => write!(f, "ScheduledTask: fd {}", task.fd),
        }
    }
}

pub struct TaskQueue {
    pub queue: Vec<Task>,
}

impl TaskQueue {
    pub fn new() -> Self {
        Self { queue: Vec::new() }
    }

    pub fn add_task(&mut self, task: Task) {
        self.queue.push(task);
    }
}

Now, for the event loop itself.

struct EventLoop {
    reactor: Arc<Mutex<Reactor>>,
    task_queue: Arc<Mutex<TaskQueue>>,
    references: HashMap<usize, Box<dyn EventHandler>>,
}

I’m using the Box<dyn EventHandler> to hold a reference to the object backing the file descriptor (more on that below).

There is the overhead of dynamic dispatch here but we’re not trying to build a performant system, so that’s okay.

impl EventLoop {
    fn new(reactor: Arc<Mutex<Reactor>>, task_queue: Arc<Mutex<TaskQueue>>) -> Self {
        Self {
            reactor,
            task_queue,
            references: HashMap::new(),
        }
    }

    /// Add a reference to the object backing the file descriptor
    fn register(&mut self, fd: usize, reference: Box<dyn EventHandler>) {
        self.references.insert(fd, reference);
    }

    /// Remove the reference backing the file descriptor
    fn unregister(&mut self, fd: usize) {
        self.references.remove(&fd);
    }

    fn process_tasks(&mut self) {
        let mut tasks_to_process = Vec::new();

        {
            // Collect tasks to process
            let mut task_queue = self.task_queue.lock().unwrap();
            while let Some(task) = task_queue.queue.pop() {
                tasks_to_process.push(task);
            }
        }

        // Process collected tasks
        for task in tasks_to_process {
            match task {
                Task::RegistrationTask(registration_task) => {
                    self.register(registration_task.fd, registration_task.reference);
                }
                Task::UnregistrationTask(unregistration_task) => {
                    self.unregister(unregistration_task.fd);
                }
                Task::ScheduledTask(scheduled_task) => {
                    if let Some(reference) = self.references.get_mut(&scheduled_task.fd) {
                        reference.poll();
                    }
                }
            }
        }
    }

    fn handle_events(&mut self, events: Vec<Event>) {
        for event in events {
            if let Some(reference) = self.references.get_mut(&event.fd) {
                reference.event(event);
            }
        }
    }

    fn run(&mut self) {
        loop {
            self.process_tasks();
            let events = self
                .reactor
                .lock()
                .unwrap()
                .poll()
                .expect("Error polling the reactor");

            self.handle_events(events);
        }
    }
}

The most important part of the EventLoop is the run function at the bottom. We do three things here - get the list of tasks on the task queue and take action on them, poll the Reactor for any new events and handle any events that are returned.

When an event is returned, the EventLoop calls the event function on the object backing the file descriptor.

It’s worth it to stop and paint a picture here of what is going on since we’ve added a lot to our runtime in this section.

This is mostly the same as the previous image, except that I’ve expanded on what the application really contains. But there’s still that EventHandler box I haven’t really explained.

EventHandler

This is a really simple part of the system. It’s just some object that handles events. That’s all it is!

The object is expected to conform to some interface like the one below

trait EventHandler {
    fn event(&mut self, event: Event);
    fn poll(&mut self);
}

Now, what objects do we need to conform to this trait? If we’re setting up a TCP listener, we need the object serving as a listener to conform to that trait, so let’s do that first.

enum AsyncTcpListenerState {
    WaitingForConnection,
    Accepting(TcpStream),
}

struct AsyncTcpListener {
    listener: TcpListener,
    fd: usize,
    reactor: Arc<Mutex<Reactor>>,
    task_queue: Arc<Mutex<TaskQueue>>,
    state: Option<AsyncTcpListenerState>,
}

The AsyncTcpListener is the object backing the TCP server file descriptor. It maintains some internal state to make it easier to encode logic, so let’s look at the implementation.

impl AsyncTcpListener {
    fn new(
        listener: TcpListener,
        reactor: Arc<Mutex<Reactor>>,
        task_queue: Arc<Mutex<TaskQueue>>,
    ) -> std::io::Result<Self> {
        let fd = listener.as_raw_fd();
        reactor.lock().unwrap().add(fd, Event::readable(fd))?;
        Ok(AsyncTcpListener {
            listener,
            fd: fd as usize,
            reactor,
            task_queue,
            state: Some(AsyncTcpListenerState::WaitingForConnection),
        })
    }
}

impl EventHandler for AsyncTcpListener {
    fn event(&mut self, event: Event) {
        match event.readable {
            true => match self.listener.accept() {
                Ok((client, addr)) => {
                    self.state.replace(AsyncTcpListenerState::Accepting(client));
                    self.task_queue
                        .lock()
                        .unwrap()
                        .add_task(Task::ScheduledTask(ScheduledTask { fd: self.fd }));
                }
                Err(e) => eprintln!("Error accepting connection: {}", e),
            },
            false => {
                panic!("AsyncTcpListener received an event that is not readable")
            }
        }
    }

    fn poll(&mut self) {
        match self.state.take() {
            Some(AsyncTcpListenerState::Accepting(client)) => {
                let client = AsyncTcpClient::new(
                    client,
                    Arc::clone(&self.reactor),
                    Arc::clone(&self.task_queue),
                )
                .unwrap();
                self.task_queue
                    .lock()
                    .unwrap()
                    .add_task(Task::RegistrationTask(RegistrationTask {
                        fd: client.fd,
                        reference: Box::new(client),
                    }));
            }
            Some(AsyncTcpListenerState::WaitingForConnection) => {
                panic!("The WaitingForConnection state should not be reached in the poll fn for listener")
            }
            None => {
                panic!("No state found in the poll fn for listener")
            }
        }
    }
}

Now, when we receive a connection from the internet, we need to do something similar for each client because we need an object backing the file descriptors for clients as well.

#[derive(Debug)]
enum AsyncTcpClientState {
    Waiting,
    Reading,
    Writing,
    Close,
    Closed,
}

struct AsyncTcpClient {
    client: TcpStream,
    fd: usize,
    reactor: Arc<Mutex<Reactor>>,
    task_queue: Arc<Mutex<TaskQueue>>,
    state: Option<AsyncTcpClientState>,
}

There’s a lot more states but that’s entirely implementation dependent. You could make do with just 2 or 3 states too. It all comes down to how you want to encapsulate the logic.

Now, here’s the implementation for the client

impl AsyncTcpClient {
    fn new(
        client: TcpStream,
        reactor: Arc<Mutex<Reactor>>,
        task_queue: Arc<Mutex<TaskQueue>>,
    ) -> std::io::Result<Self> {
        let fd = client.as_raw_fd();
        reactor.lock().unwrap().add(fd, Event::readable(fd))?;
        Ok(Self {
            client,
            fd: fd as usize,
            reactor,
            task_queue,
            state: Some(AsyncTcpClientState::Waiting),
        })
    }
}

impl EventHandler for AsyncTcpClient {
    fn event(&mut self, event: Event) {
        match self.state.take() {
            Some(AsyncTcpClientState::Waiting) => {
                if event.readable {
                    self.state.replace(AsyncTcpClientState::Reading);
                    self.task_queue
                        .lock()
                        .unwrap()
                        .add_task(Task::ScheduledTask(ScheduledTask { fd: self.fd }));
                }
            }
            Some(s) => {
                self.state.replace(s);
            }
            None => {
                panic!("state was none");
            }
        }
    }

    fn poll(&mut self) {
        match self.state.take() {
            None => {}
            Some(AsyncTcpClientState::Waiting) => {
                panic!("The Waiting state should not be reached in the poll fn for client")
            }
            Some(AsyncTcpClientState::Reading) => {
                let reader = BufReader::new(&self.client);
                let http_request: Vec<_> = reader
                    .lines()
                    .map(|line| line.unwrap())
                    .take_while(|line| !line.is_empty())
                    .collect();
                if http_request
                    .iter()
                    .next()
                    .unwrap()
                    .contains("GET / HTTP/1.1")
                {
                    self.state.replace(AsyncTcpClientState::Writing);
                } else {
                    eprintln!("received invalid request, closing the socket connection");
                    self.state.replace(AsyncTcpClientState::Close);
                }
                self.state.replace(AsyncTcpClientState::Writing);
                self.task_queue
                    .lock()
                    .unwrap()
                    .add_task(Task::ScheduledTask(ScheduledTask { fd: self.fd }));
                self.reactor.lock().unwrap().notify().unwrap();
            }
            Some(AsyncTcpClientState::Writing) => {
                let path = Path::new("hello.html");
                let content = std::fs::read(path).unwrap();
                let response = format!(
                    "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
                    content.len(),
                    String::from_utf8_lossy(&content)
                );
                self.client.write_all(response.as_bytes()).unwrap();
                self.state.replace(AsyncTcpClientState::Close);
                self.task_queue
                    .lock()
                    .unwrap()
                    .add_task(Task::ScheduledTask(ScheduledTask { fd: self.fd }));
                self.reactor.lock().unwrap().notify().unwrap();
            }
            Some(AsyncTcpClientState::Close) => {
                //  remove the client fd from the reactor and unregister from the event loop
                self.reactor
                    .lock()
                    .unwrap()
                    .delete(self.fd.try_into().unwrap())
                    .unwrap();
                self.task_queue
                    .lock()
                    .unwrap()
                    .add_task(Task::UnregistrationTask(UnregistrationTask { fd: self.fd }));
                self.client.shutdown(std::net::Shutdown::Both).unwrap();
                self.state.replace(AsyncTcpClientState::Closed);
            }
            Some(AsyncTcpClientState::Closed) => {}
        }
    }
}

This looks like a lot of code, but both of these objects are just state machines backing a file descriptor. There’s nothing more to it. The client encapsulates logic in it’s state machine like reading from a file and then writing back to the socket, and the listener only has 2 states - either accepting connections or waiting for new connections.

You could do fancy things with the client like spawning threads when reading from a file to ensure that you don’t block the event loop in any way, but for our experimental purposes this should be good enough.

Each of the objects above implement the EventHandler trait - which means they both have a poll && event method.

The former is called by the event loop when it is scheduled to do so - that is, it gets a task from the task queue telling it to poll a specific object.

The latter is called when the event loop receives a new event for a specific file descriptor. It then calls the object backing that file descriptor.

Wiring It All Up

Okay, we understand each of the components on their own but we need to wire up everything to see how it works together.

fn main() {
    let reactor = Arc::new(Mutex::new(Reactor::new().unwrap()));
    let task_queue = Arc::new(Mutex::new(TaskQueue::new()));
    let mut event_loop = EventLoop::new(Arc::clone(&reactor), Arc::clone(&task_queue));

    //  start listener
    let tcp_listener = TcpListener::bind("127.0.0.1:8000").unwrap();
    let listener =
        AsyncTcpListener::new(tcp_listener, Arc::clone(&reactor), Arc::clone(&task_queue)).unwrap();
    task_queue
        .lock()
        .unwrap()
        .add_task(Task::RegistrationTask(RegistrationTask {
            fd: listener.fd,
            reference: Box::new(listener),
        }));

    //  start the event loop
    event_loop.run();
}

We create all our objects and pass them into the event loop. We then create a TcpListener and construct a backing object on top of it. Next, we push a task into the task queue to register the new file descriptor and it’s backing object with the event loop and we fire off the event loop.

I’m placing the run function of the event loop below again because it’s the central piece of this entire system. When we call process_tasks, the event loop will pick up the RegistrationTask and add the file descriptor and the backing object to it’s references.

Next, it will listen for events from the reactor and act on those events. When it receives an event, it gets it for a specific file descriptor, so it will look up the file descriptor in it’s references, find the backing object and call the object’s event method.

fn run(&mut self) {
  loop {
      self.process_tasks();
      let events = self
          .reactor
          .lock()
          .unwrap()
          .poll()
          .expect("Error polling the reactor");

      self.handle_events(events);
  }
}

I mentioned earlier that the scheduled tasks are the way of some object telling the event loop that it wants to be polled, and that’s what the event loop does in the process_tasks method. It just calls the poll method on the underlying object.

fn process_tasks(&mut self) {
    let mut tasks_to_process = Vec::new();

    {
        // Collect tasks to process
        let mut task_queue = self.task_queue.lock().unwrap();
        while let Some(task) = task_queue.queue.pop() {
            tasks_to_process.push(task);
        }
    }

    // Process collected tasks
    for task in tasks_to_process {
        match task {
            Task::RegistrationTask(registration_task) => {
                self.register(registration_task.fd, registration_task.reference);
            }
            Task::UnregistrationTask(unregistration_task) => {
                self.unregister(unregistration_task.fd);
            }
            Task::ScheduledTask(scheduled_task) => {
                if let Some(reference) = self.references.get_mut(&scheduled_task.fd) {
                    reference.poll();
                }
            }
        }
    }
}

Now, if you go back and look at the code for the AsyncTcpClient, you’ll see calls to notify on the reactor placed at the end of each state transition. The purpose of that call is for the object to tell the event loop that it needs to be polled because it’s ready to move forward (this is where the Unix pipe hack comes in).

Well, there you have it - a single-threaded async runtime in ~500 lines of code.

I intend to continue the series with another post by trying to hook up my futures and executor with the event loop and reactor. Should be interesting!

References

1. YouTube Playlist On Async I/O In Rust
2. Polling Crate
3. MIO Crate
4. Unix Streams
5. Async Rust By Evan Schwartz
6. Async I/O In Rust Git Repo
7. King Protty’s Talk On Zig I/O Concurrency

I’m trying to understand async runtimes better, specifically in Rust. This post is a short attempt to build the most basic example of a future and execute it.

A better and more detailed version of this can be found in this book.

Futures

Creating a future in Rust is very straightforward. You just need to implement the Future trait, which requires one method, poll, to be implemented.

Let’s build a simple future.

use std::{
    future::Future,
    sync::{
        mpsc::{sync_channel, Receiver, SyncSender},
        Arc, Mutex,
    },
    task::{Context, Poll, Waker},
    thread,
    time::Duration,
};

use futures::{
    future::{BoxFuture, FutureExt},
    task::{waker_ref, ArcWake},
};

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

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

The Waker referenced above is a way for the task to inform it’s executor that it should be polled again. You can read more from the docs here.

Now, let’s make it a future and construct an instance of it.

Note: If you want to understand more about what Pin is and why it’s required, look here. Long story short - it’s a way for Rust to ensure that the data being pointed at isn’t moved around in memory

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

impl TimerFuture {
    fn new(duration: Duration) -> Self {
        let shared_state = Arc::new(Mutex::new(SharedState {
            completed: false,
            waker: None,
        }));
        let thread_shared_state = Arc::clone(&shared_state);
        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 }
    }
}

There we go, we now have a future. Now, we need to run this future to completion.

Executing A Future

Let’s get some conceptual modeling out of the way first. We have the following terms - Spawner, Executor, Task, Future & Waker. Let’s build the mental model bottom up.

A Future is something that will complete at some point in time.

A Waker is a way for the future to ensure that it is polled again (typically by placing the future back onto the executor queue)

A Task is a wrapper around a Future and a Waker.

A Spawner is a component that constructs a Task and provides it to the Executor.

An Executor constantly checks it’s list of Tasks and polls their futures for completion

A picture is worth a thousand words, so here’s one

Let’s write some code to represent this.

Overall Structure

Here’s the basic structs for Task, Spawner & Executor.

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

struct Spawner {
    task_sender: SyncSender<Arc<Task>>,
}

struct Executor {
    task_queue: Receiver<Arc<Task>>,
}

Task and Spawner keep putting tasks onto the channel and the Executor has the receiving end of the channel and keeps polling the futures within these tasks.

Task

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

The BoxFuture is an aliased type for pub type BoxFuture<'a, T> = Pin<alloc::boxed::Box<dyn Future<Output = T> + Send + 'a>>.

It’s a complicated type but it essentially means that I have a Box container around a dynamic Future and I do not want the data within the Box container to be moved around in memory so I wrap it in a Pin.

The easiest way to understand Pin is to work through a simple example. Let’s create a simple Pin wrapper and try to move it. The example below compiles with an error because String does not have Copy semantics, so when it tries to move the value, it cannot because the container has been pinned in place.

#[derive(Debug)]
struct MyData {
    value: String,
}
let my_data = Pin::new(Box::new(MyData {
    value: String::from("hello"),
}));
let moved_data = Box::new(MyData {
    value: my_data.value,
});
error[E0507]: cannot move out of dereference of `Pin<Box<MyData>>`
  --> src/main.rs:40:16
   |
40 |         value: my_data.value,
   |                ^^^^^^^^^^^^^ move occurs because value has type `String`, which does not implement the `Copy` trait
   |
help: consider cloning the value if the performance cost is acceptable
   |
40 |         value: my_data.value.clone(),
   |                             ++++++++

Now, we said earlier that Task is a wrapper for a Future and a Wake object. By implementing the trait below, we allow a Waker object to be constructed from a Task. More on that below.

impl ArcWake for Task {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        let arc_clone = Arc::clone(&arc_self);
        arc_self.task_sender.send(arc_clone).unwrap();
    }
}

Spawner

struct Spawner {
    task_sender: SyncSender<Arc<Task>>,
}

This component pushes new tasks onto the channel.

impl Spawner {
    fn spawn(&self, future: impl Future<Output = ()> + Send + 'static) {
        let box_future = future.boxed();
        let task = Arc::new(Task {
            future: Mutex::new(box_future),
            task_sender: self.task_sender.clone(),
        });
        self.task_sender.send(task).unwrap();
    }
}

Executor

struct Executor {
    task_queue: Receiver<Arc<Task>>,
}

This component blocks on the receiver until it receives new tasks. Once it gets these tasks, it polls them to determine if they are complete.

There is also a convenience method for constructing an Executor and Spawner.

impl Executor {
    fn run(&self) {
        while let Ok(task) = self.task_queue.recv() {
            let mut fut = task.future.lock().unwrap();
            let waker = waker_ref(&task);
            let context = &mut Context::from_waker(&waker);
            if fut.as_mut().poll(context).is_ready() {
                println!("The future is done running");
            }
        }
    }

    fn executor_and_spawner() -> (Executor, Spawner) {
        let (sync_sender, receiver) = sync_channel(10000);
        let executor = Executor {
            task_queue: receiver,
        };
        let spawner = Spawner {
            task_sender: sync_sender,
        };
        (executor, spawner)
    }
}

In the run method, you can see the Waker object being constructed from the Task. This is possible because we implemented the ArcWake trait for the Task above, and we are providing it in the Context object to the poll method.

Now, when the future is ready and it calls the wake method on the Wake object, the method in the ArcWake trait is executed.

Putting It Together

Here’s a simple test to demonstrate how this works together

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn spawn_tasks() {
    let (executor, spawner) = Executor::executor_and_spawner();
    spawner.spawn(async {
        println!("Hello");
        TimerFuture::new(Duration::from_secs(2)).await;
        println!("World");
    });
    drop(spawner);
    executor.run();
  }
}

The output should be Hello, followed by a 2 second pause and then a World.

Here’s a flow diagram for what is going on

-> Spawner creates task 
-> Spawner puts task into execution queue 
-> Executor polls the future in the task and finds it pending 
-> The future will keep track of the Waker object 
-> Duration elapses 
-> The future calls the `.wake` method 
-> The `ArcWake` implementation on Task is triggered 
-> The task is again placed on the queue 
-> Executor polls the future and finds it completed

Conclusion

So, there you have it - a very, very simple implementation of a single threaded executor.

References

• Cloudflare Blog On Pin And Unpin
• Asynchronous Programming In Rust
• GitHub Repo
• A Guided Tour Of Streams In Rust
• Pin And Suffering
• Scheduling Internals

I stumbled upon the talent plan course by TiDB and decided to use it to understand distributed transactions (a.k.a Percolator) better.

The point of the Percolator course is to read the paper and implement it in Rust, so I’ll be diving into bits and pieces of the paper as well.

Overview

Percolator came out of Google in 2010 as a way to run distributed transactions on top of Bigtable, Google’s distributed storage system. At the time Google wanted to run transactions on a regular DBMS but nothing came close to handling the volume Google had, so they built a distributed transaction model on top of their own distributed data source.

Mechanics

The mechanics of Percolator are deceptively simple. It’s a straightforward algorithm to understand, albeit with a lot of bookkeeping involved. Let’s look at how it would work for a simple key-value store.

Every key has 3 columns attached to it - a Data column, a Lock column and a Write column. Every transaction commit has 2 phases - a pre-write phase and a commit phase.

The pre-write phase involves writing to the Data & Lock columns - that is, whatever the new value the transaction wants to write and an associated lock. Now, the transaction hasn’t committed yet, so this newly written value is not visible yet.

The transaction commit phase involves some validation checks and assuming those validation checks pass, the lock is removed and a pointer to the Data column is placed in the Write column.

Let’s run through an example transaction. (Look here for a detailed example from TiKV).

I’ll work through an example that uses markers more pertinent to the codebase for the course. The examples provided in the paper and the blog post from TiKV use slightly different markers in the rows, but they’re functionally the same.

Initial State

Let’s say we start out with a simple key value store that represents the accounts of individuals. There are 2 account holders Bob and Joe with an initial balance of 10$ and 2$ respectively. We will run a transaction that will transfer 7$ from Bob to Joe.

The image below shows the initial state we will start from.

The key is a tuple representing the (key, timestamp) of the transaction that wrote that data into it. The value can either be a set of bytes representing the data or a timestamp serving as a pointer (The value can also have a third variant, which is wall clock timestamp. I’ll get to that next). The write column has a pointer pointing to the timestamp in the data column where the data actually resides as a value.

You’ll notice that the timestamps in the data column and the write column are different because they represent the start and commit timestamps of the transaction that wrote this data and committed it. An invariant that must always be held is that commit timestamp > start timestamp. This is ensured by a timestamp oracle, which is a fancy term for a server that hands out strictly monotonically increasing values. The strictly is doing a lot of work here because it means that you are guaranteed to never see a timestamp more than once.

Pre-Write

Now, let’s say the transaction begins that will transfer the amount of 7$ from Bob to Joe. This transaction will begin with a timestamp of t2. The first thing this transaction does it to acquire a lock on both keys that are involved in the transaction, and write down the data that should be the end result of this transaction.

The entries in the two columns are different - one of them is selected as the primary lock (the key with a # symbol) and the other (or others, depending on the number of keys involved in the transaction) are selected as secondaries. The primary lock is primarily used as a synchronization point in case of any cleanup required. This cleanup is usually done if a transaction has crashed or failed to remove all of it’s locks.

Commit

As part of the commit phase, the transaction will perform some validation checks (this is implementation dependent but the main check is that the primary lock hasn’t been removed). Assuming these validation checks pass, the transaction will erase the locks placed on all the rows and place a new entry in the write column serving as a pointer to the entry in the data column that was created in the pre-commit phase. It marks the new entry in the write column with a commit timestamp.

Implementation

Let’s dive into an implementation of the Percolator algorithm. The base for this was taken from the TiKV course. If you just want to jump straight to my full implementation, go here.

The course gives you a client and server implementation with the functionality stubbed out. The goal is for you to write that functionality. Very helpfully, the course gives you a network implementation between the client and the server which simulates faults like dropped requests which help you test your implementation. (testing distributed systems is more than half the problem).

Write Path

Let’s cover the write path first. A transaction client write’s all data into a private buffer until it is time to actually perform the commit.

#[derive(Clone)]
pub struct Client {
    // Your definitions here.
    tso_client: TSOClient,
    txn_client: TransactionClient,
    transaction: Option<Transaction>,
}

#[derive(Clone, Debug)]
pub struct KVPair {
    key: Vec<u8>,
    value: Vec<u8>,
}

#[derive(Clone)]
pub struct Transaction {
    pub start_ts: u64,
    pub write_buffer: Vec<KVPair>,
}

impl Transaction {
    pub fn new(start_ts: u64) -> Self {
        Transaction {
            start_ts,
            write_buffer: Vec::new(),
        }
    }
}

impl Client {
    /// Creates a new Client.
    pub fn new(tso_client: TSOClient, txn_client: TransactionClient) -> Client {
        // Your code here.
        Client {
            tso_client,
            txn_client,
            transaction: None,
        }
    }

    /// Gets a timestamp from a TSO.
    pub fn get_timestamp(&self) -> Result<u64> {
        let rpc = || self.tso_client.get_timestamp(&TimestampRequest {});
        match executor::block_on(self.call_with_retry(rpc)) {
            Ok(ts) => Ok(ts.timestamp),
            Err(e) => Err(e),
        }
    }

    /// Begins a new transaction.
    pub fn begin(&mut self) {
        let ts = self
            .get_timestamp()
            .expect("unable to get a timestamp from the oracle");
        let transaction = Transaction::new(ts);
        self.transaction = Some(transaction);
    }

    /// Sets keys in a buffer until commit time.
    pub fn set(&mut self, key: Vec<u8>, value: Vec<u8>) {
        // Your code here.
        if let Some(transaction) = self.transaction.as_mut() {
            transaction.write_buffer.push(KVPair { key, value });
            return;
        }
        panic!("attempting to set a key value pair without a txn");
    }
}

Now, when the transaction performs the commit it runs through the two phases - pre-write and commit per key-value pair on the client.

/// Commits a transaction.
    pub fn commit(&self) -> Result<bool> {
        //  PRE-WRITE PHASE
        let transaction = self.transaction.as_ref().expect("transaction not found");
        let kv_pair = &transaction.write_buffer;
        let primary = kv_pair
            .first()
            .expect("cannot find the first key value pair");
        let secondaries = &kv_pair[1..];
        //  acquire a lock on the primary first
        let args = PrewriteRequest {
            timestamp: transaction.start_ts,
            kv_pair: Some(KvPair {
                key: primary.key.clone(),
                value: primary.value.clone(),
            }),
            primary: Some(KvPair {
                key: primary.key.clone(),
                value: primary.value.clone(),
            }),
        };
        let rpc = || self.txn_client.prewrite(&args);
        if executor::block_on(self.call_with_retry(rpc))?.res == false {
            return Ok(false);
        }
        //  acquire locks on the secondaries now
        for kv_pair in secondaries {
            let args = PrewriteRequest {
                timestamp: transaction.start_ts,
                kv_pair: Some(KvPair {
                    key: kv_pair.key.clone(),
                    value: kv_pair.value.clone(),
                }),
                primary: Some(KvPair {
                    key: primary.key.clone(),
                    value: primary.value.clone(),
                }),
            };
            let rpc = || self.txn_client.prewrite(&args);
            if executor::block_on(self.call_with_retry(rpc))?.res == false {
                return Ok(false);
            }
        }
        //  END PRE-WRITE PHASE

        //  COMMIT PHASE
        let commit_ts = self.get_timestamp()?;
        assert!(
            commit_ts > transaction.start_ts,
            "panic because the commit ts is not strictly greater than the start ts of the txn"
        );
        let args = CommitRequest {
            start_ts: transaction.start_ts,
            commit_ts,
            is_primary: true,
            kv_pair: Some(KvPair {
                key: primary.key.clone(),
                value: primary.value.clone(),
            }),
        };
        let rpc = || self.txn_client.commit(&args);
        match executor::block_on(self.call_with_retry(rpc)) {
            Ok(response) => {
                return Ok(response.res);
            }
            Err(e) => match e {
                labrpc::Error::Other(e_string) if e_string == "reqhook" => {
                    return Ok(false);
                }
                _ => return Err(e),
            },
        }
        for kv_pair in secondaries {
            let args = CommitRequest {
                start_ts: transaction.start_ts,
                commit_ts,
                is_primary: false,
                kv_pair: Some(KvPair {
                    key: primary.key.clone(),
                    value: primary.value.clone(),
                }),
            };
            let rpc = || self.txn_client.commit(&args);
            match executor::block_on(self.call_with_retry(rpc)) {
                Ok(response) => return Ok(response.res),
                Err(e) => match e {
                    labrpc::Error::Other(e_string) => {
                        if e_string == "reqhook" {
                            return Ok(true);
                        }
                    }
                    _ => return Err(e),
                },
            }
        }
        //  END COMMIT PHASE
        Ok(true)
    }

There’s a bunch of additional code ensuring that the RPC call is retried x number of times, but the gist is the same as that of the paper.

Now, for the server side, we have 3 main components - a memory storage server, a key value table and the timestamp oracle. The key value table holds the actual data and the memory storage server manages the business logic of the transaction.

// KvTable is used to simulate Google's Bigtable.
// It provides three columns: Write, Data, and Lock.
#[derive(Clone, Default)]
pub struct KvTable {
    write: BTreeMap<Key, Value>,
    data: BTreeMap<Key, Value>,
    lock: BTreeMap<Key, Value>,
}

impl KvTable {
    // Reads the latest key-value record from a specified column
    // in MemoryStorage with a given key and a timestamp range.
    #[inline]
    fn read(
        &self,
        key: Vec<u8>,
        column: Column,
        ts_start_inclusive: Option<u64>,
        ts_end_inclusive: Option<u64>,
    ) -> Option<(Key, Value)> {
        let col = match column {
            Column::Data => &self.data,
            Column::Lock => &self.lock,
            Column::Write => &self.write,
        };
        let mut res = None;
        let mut max_timestamp_seen = 0;

        for ((k, ts), value) in col.iter() {
            if k == &key
                && ts_start_inclusive.map_or(true, |start| *ts >= start)
                && ts_end_inclusive.map_or(true, |end| *ts <= end)
                && *ts >= max_timestamp_seen
            {
                max_timestamp_seen = *ts;
                res = Some(((k.clone(), *ts), value.clone()));
            }
        }
        res
    }

    // Writes a record to a specified column in MemoryStorage.
    #[inline]
    fn write(&mut self, key: Vec<u8>, column: Column, ts: u64, value: Value) {
        let col = match column {
            Column::Data => &mut self.data,
            Column::Lock => &mut self.lock,
            Column::Write => &mut self.write,
        };
        col.insert((key, ts), value);
    }

    #[inline]
    // Erases a record from a specified column in MemoryStorage.
    fn erase(&mut self, key: Vec<u8>, column: Column, commit_ts: u64) {
        let col = match column {
            Column::Data => &mut self.data,
            Column::Lock => &mut self.lock,
            Column::Write => &mut self.write,
        };
        let mut keys_to_remove = Vec::new();
        for ((k, ts), _) in col.iter() {
            if k == &key && *ts == commit_ts {
                keys_to_remove.push((k.clone(), *ts));
            }
        }
        for key in keys_to_remove {
            let value = col.remove(&key);
            assert!(value.is_some());
        }
    }
}

// MemoryStorage is used to wrap a KvTable.
// You may need to get a snapshot from it.
#[derive(Clone, Default)]
pub struct MemoryStorage {
    data: Arc<Mutex<KvTable>>,
}

#[async_trait::async_trait]
impl transaction::Service for MemoryStorage {
  // example prewrite RPC handler.
    async fn prewrite(&self, req: PrewriteRequest) -> labrpc::Result<PrewriteResponse> {
        let primary = req.primary.ok_or_else(|| {
            labrpc::Error::Other("primary kv_pair is missing in the prewrite request".to_string())
        })?;
        let kv_pair = req.kv_pair.ok_or_else(|| {
            labrpc::Error::Other("kv_pair is missing in the prewrite request".to_string())
        })?;
        let mut storage = self.data.lock().unwrap();
        match storage.read(
            kv_pair.key.clone(),
            Column::Write,
            Some(req.timestamp),
            None,
        ) {
            Some(_) => return Ok(PrewriteResponse { res: false }),
            None => (),
        };
        match storage.read(kv_pair.key.clone(), Column::Lock, Some(0), None) {
            Some(_) => return Ok(PrewriteResponse { res: false }),
            None => (),
        };
        //  all checks completed, place data and lock
        storage.write(
            kv_pair.key.clone(),
            Column::Data,
            req.timestamp,
            Value::Vector(kv_pair.value.clone()),
        );
        if primary == kv_pair {
            storage.write(
                kv_pair.key.clone(),
                Column::Lock,
                req.timestamp,
                Value::LockPlacedAt(SystemTime::now()),
            );
        } else {
            storage.write(
                kv_pair.key.clone(),
                Column::Lock,
                req.timestamp,
                Value::Vector(primary.key),
            );
        }
        Ok(PrewriteResponse { res: true })
    }

    // example commit RPC handler.
    async fn commit(&self, req: CommitRequest) -> labrpc::Result<CommitResponse> {
        let mut storage = self.data.lock().unwrap();
        let kv_pair = req
            .kv_pair
            .expect("kv_pair is missing in the commit request");
        if req.is_primary {
            //  check lock on primary still holds
            match storage.read(
                kv_pair.key.clone(),
                Column::Lock,
                Some(req.start_ts),
                Some(req.start_ts),
            ) {
                Some(_) => (),
                None => {
                    return Ok(CommitResponse { res: false });
                }
            };
        }

        //  create write and remove lock
        storage.write(
            kv_pair.key.clone(),
            Column::Write,
            req.commit_ts,
            Value::Timestamp(req.start_ts),
        );
        storage.erase(kv_pair.key, Column::Lock, req.start_ts);
        Ok(CommitResponse { res: true })
    }
}

While doing the prewrite phase, you’ll notice that there is a check to determine what the primary lock is supposed to be. The value written into the primary lock is the wall clock time when it is being placed. This is important during the read path.

The commit handler is a lot simpler - the only validation it does is to determine that if the primary is being committed that the lock is still valid. It doesn’t need to do that for the secondaries.

Read Path

Let’s deal with a straightforward read path first - one where it doesn’t encounter any locks. This post from TiKV has a great explanation of the read path, so I’m going to use the text from there as an explanation.

The basic idea is that you first check if there are any pending locks. Assuming there are none, fetch the latest write record within range and use the pointer there to the data column to get the actual value.

#[async_trait::async_trait]
impl transaction::Service for MemoryStorage {
  // example get RPC handler.
    async fn get(&self, req: GetRequest) -> labrpc::Result<GetResponse> {
        loop {
            let mut storage = self.data.lock().unwrap();
            let is_row_locked =
                storage.read(req.key.clone(), Column::Lock, Some(0), Some(req.timestamp));
            if is_row_locked.is_some() {
                drop(storage);
                self.back_off_maybe_clean_up_lock(req.timestamp, req.key.clone());
                std::thread::sleep(Duration::from_millis(100));
                continue;
            }

            let start_ts =
                match storage.read(req.key.clone(), Column::Write, Some(0), Some(req.timestamp)) {
                    Some(((_, commit_ts), value)) => match value {
                        Value::Timestamp(start_ts) => start_ts,
                        _ => {
                            return Err(labrpc::Error::Other(format!(
                                "unexpected value found in write column for key {:?} at ts {}",
                                req.key, commit_ts
                            )))
                        }
                    },
                    None => {
                        return Ok(GetResponse {
                            success: false,
                            value: Vec::new(),
                        });
                    }
                };

            let data = match storage.read(
                req.key.clone(),
                Column::Data,
                Some(start_ts),
                Some(start_ts),
            ) {
                Some(((_, _), value)) => match value {
                    Value::Vector(bytes) => bytes,
                    _ => {
                        return Err(labrpc::Error::Other(format!(
                            "unexpected value found in data column for key {:?} at ts {}",
                            req.key, start_ts
                        )));
                    }
                },
                None => {
                    return Err(labrpc::Error::Other(format!(
                        "No value found in data column for key {:?} at timestamp {}",
                        req.key, start_ts
                    )));
                }
            };

            return Ok(GetResponse {
                success: true,
                value: data,
            });
        }
    }
}

The code is verbose because of all the match guards but hopefully straighforward to follow. (On another note, I find myself writing very verbose Rust code which is what the language seems to lend itself to.)

Now, let’s dive into dealing with pending locks and potentially failed transactions on the read path. First, let’s look at what the paper says about dealing with these transactions.

So Percolator distinguishes between crashes by determining when a client has crashed.

1. If it crashes after the pre-commit phase but before the commit phase, the transaction must be rolled back. The way to determine this is to check if the TTL on the wall clock time for the primary lock has elapsed.

2. If the client crashed after committing the primary lock in the commit phase but before committing the secondary locks, the transaction must be rolled forward since the only synchronization point is the primary lock itself.

Here’s all the code related to that. It’s slightly complicated by the fact that a get request might come across a secondary lock, in which case it needs handle the indirection. I’ve tried to document this function as clearly as possible so it’s easier to follow.

impl MemoryStorage {
    fn back_off_maybe_clean_up_lock(&self, start_ts: u64, key: Vec<u8>) {
        //  STEPS:
        //  1. Recheck the condition that prompted this call by re-acquiring lock. Things might have changed
        //  2. Check if the lock is the primary lock. If secondary lock, get primary lock
        //  3. If primary lock present and
        //      a. has expired, roll-back the txn
        //      b. has not expired, do nothing and retry after some time
        //  4. If primary lock not present and
        //      a. Data found in Write column, roll-forward the txn
        //      b. No data found in Write column, remove stale lock

        let mut storage = self.data.lock().unwrap();
        let ((key, start_ts), value) =
            match storage.read(key.clone(), Column::Lock, Some(0), Some(start_ts)) {
                Some((key, value)) => (key, value),
                None => return,
            };

        let is_primary_lock = match value {
            Value::LockPlacedAt(creation_time) => true,
            Value::Vector(ref data) => false,
            Value::Timestamp(_) => panic!(
                "unexpected value of bytes found in lock column, expected SystemTime or Vec<u8>"
            ),
        };

        if is_primary_lock {
            if self.check_if_primary_lock_expired(value) {
                self.remove_lock_and_rollback(&mut storage, key, start_ts);
            }
            return;
        }

        //  handle the secondary lock here
        let primary_key = value
            .as_vector()
            .expect("unexpected value in lock column, expected Vec<u8>");
        match storage.read(primary_key.clone(), Column::Lock, Some(0), Some(start_ts)) {
            Some(((_, conflicting_start_ts), value)) => {
                if self.check_if_primary_lock_expired(value) {
                    self.remove_lock_and_rollback(&mut storage, primary_key, conflicting_start_ts);
                    self.remove_lock_and_rollback(&mut storage, key, start_ts);
                    return;
                }
            }
            None => {
                //  the primary lock is gone, check for data
                match storage.read(primary_key, Column::Write, None, None) {
                    None => {
                        self.remove_lock_and_rollback(&mut storage, key, start_ts);
                    }
                    Some(((_, commit_ts), value)) => {
                        let start_ts = value
                            .as_timestamp()
                            .expect("unexpected value in write column, expected ts");
                        self.remove_lock_and_roll_forward(&mut storage, key, start_ts, commit_ts);
                    }
                }
            }
        }
    }

    fn check_if_primary_lock_expired(&self, value: Value) -> bool {
        let lock_creation_time = value
            .as_lock_placed_at()
            .expect("unexpected value in lock column, expected SystemTime");
        let ttl_duration = Duration::from_nanos(TTL);
        let future_time = lock_creation_time + ttl_duration;
        future_time < SystemTime::now()
    }

    fn remove_lock_and_rollback(
        &self,
        storage: &mut std::sync::MutexGuard<KvTable>,
        key: Vec<u8>,
        timestamp: u64,
    ) {
        storage.erase(key.clone(), Column::Lock, timestamp);
        storage.erase(key.clone(), Column::Data, timestamp);
    }

    fn remove_lock_and_roll_forward(
        &self,
        storage: &mut std::sync::MutexGuard<KvTable>,
        key: Vec<u8>,
        start_ts: u64,
        commit_ts: u64,
    ) {
        storage.erase(key.clone(), Column::Lock, start_ts);
        storage.write(key, Column::Write, commit_ts, Value::Timestamp(start_ts));
    }
}

Aside

As an aside, something that bothered me while doing this course was that failed transactions are only cleaned up on the read path. What happens if we have a set of transactions that do only writes with no reads and they all keep failing? Theoretically, we could have an entire database locked up that would require a read request to come in and start freeing up locked rows.

I’m basing all this on the pseudo-code that the paper presents, part of which I have included below.

1  class Transaction {
2  struct Write { Row row; Column col; string value; };
3  vector<Write> writes ;
4  int start ts ;
5 
6  Transaction() : start ts (oracle.GetTimestamp()) {}
7  void Set(Write w) { writes .push back(w); }
8  bool Get(Row row, Column c, string* value) {
9  while (true) {
10 bigtable::Txn T = bigtable::StartRowTransaction(row);
11 // Check for locks that signal concurrent writes.
12 if (T.Read(row, c+"lock", [0, start ts ])) {
13 // There is a pending lock; try to clean it and wait
14 BackoffAndMaybeCleanupLock(row, c);
15 continue;
16 }
Figure 6: Pseudocode for Percolator transaction protocol.

I still don’t have a satisfactory answer to this question. I imagine different implementations of Percolator include cleanup logic on the write path as well.

References

1. The TiKV courses
2. The Percolator paper
3. TiKV’s blog post about Percolator
4. My solution
5. Alternate solution 1
6. Alternate solution 2

This is a series of posts around building an LSM-based key-value store. It will include flushing memtables to disk, compaction and a very minimal implementation of serializable snapshot isolation (SSI) using Write Snapshot Isolation (WSI). If you want to see the complete code for all sections of this series, check it out here. If you want to see the complete code for only this post, check out this commit. I’m going to follow the pattern of the course, so one post for every week of the course. The commits on my repo follow the same pattern, so you should be able to easily tell which commit to checkout.

Goals

By the end of the post we will have covered the following 3 goals:

1. Writing to an in-memory buffer
2. Flushing the in-memory buffer to disk
3. Running point and range queries on the sources

Introduction

If you want an overview of what LSM trees are and how different components work together without diving into too much code, I recommend looking at this post by Justin Jaffrey and this one by Garren Smith.

Typically, LSM engines have 3 main components:

1. The in memory component (memtable)
2. Flushing data to disk (sorted string table)
3. Combining data from multiple sstables together (compaction)

We’ll cover compaction in a future post.

The write path looks like this. Client put request -> add to memtable -> eventually flush to sstable. The eventually flush to SST is usually implementation defined where there is a limit to the memtable buffer. When this buffer is exceeded, the memtable is frozen and then converted to a sorted string table at some later point. Let’s jump into the code.

The Database Engine

We need a database struct that can hold the reference to the memtable and everything else we need. So, let’s define those structs first.

/// A thin wrapper for `LsmStorageInner` and the user interface for MiniLSM.
pub struct MiniLsm {
    pub(crate) inner: Arc<LsmStorageInner>,
    /// Notifies the L0 flush thread to stop working. (In week 1 day 6)
    flush_notifier: crossbeam_channel::Sender<()>,
    /// The handle for the compaction thread. (In week 1 day 6)
    flush_thread: Mutex<Option<std::thread::JoinHandle<()>>>,
    /// Notifies the compaction thread to stop working. (In week 2)
    compaction_notifier: crossbeam_channel::Sender<()>,
    /// The handle for the compaction thread. (In week 2)
    compaction_thread: Mutex<Option<std::thread::JoinHandle<()>>>,
}

impl Drop for MiniLsm {
    fn drop(&mut self) {
        self.compaction_notifier.send(()).ok();
        self.flush_notifier.send(()).ok();
    }
}

impl MiniLsm {
    pub fn close(&self) -> Result<()> {
        let mut flush_thread = self.flush_thread.lock();
        if let Some(flush_thread) = flush_thread.take() {
            match flush_thread.join() {
                std::result::Result::Ok(_) => (),
                Err(e) => eprintln!("flush thread panicked: {:?}", e),
            }
        }
        Ok(())
    }

    /// Start the storage engine by either loading an existing directory or creating a new one if the directory does
    /// not exist.
    pub fn open(path: impl AsRef<Path>, options: LsmStorageOptions) -> Result<Arc<Self>> {
        let inner = Arc::new(LsmStorageInner::open(path, options)?);
        let (tx1, rx) = crossbeam_channel::unbounded();
        let compaction_thread = inner.spawn_compaction_thread(rx)?;
        let (tx2, rx) = crossbeam_channel::unbounded();
        let flush_thread = inner.spawn_flush_thread(rx)?;
        Ok(Arc::new(Self {
            inner,
            flush_notifier: tx2,
            flush_thread: Mutex::new(flush_thread),
            compaction_notifier: tx1,
            compaction_thread: Mutex::new(compaction_thread),
        }))
    }
}

Now, let’s define the structs for the actual state.

/// Represents the state of the storage engine.
#[derive(Clone)]
pub struct LsmStorageState {
    /// The current memtable.
    pub memtable: Arc<MemTable>,
    /// Immutable memtables, from latest to earliest.
    pub imm_memtables: Vec<Arc<MemTable>>,
    /// L0 SSTs, from latest to earliest.
    pub l0_sstables: Vec<usize>,
    /// SsTables sorted by key range; L1 - L_max for leveled compaction, or tiers for tiered
    /// compaction.
    pub levels: Vec<(usize, Vec<usize>)>,
    /// SST objects.
    pub sstables: HashMap<usize, Arc<SsTable>>,
}

pub(crate) struct LsmStorageInner {
    pub(crate) state: Arc<RwLock<Arc<LsmStorageState>>>,
    pub(crate) state_lock: Mutex<()>,
    path: PathBuf,
    pub(crate) block_cache: Arc<BlockCache>,
    next_sst_id: AtomicUsize,
    pub(crate) options: Arc<LsmStorageOptions>,
    pub(crate) compaction_controller: CompactionController,
    pub(crate) manifest: Option<Manifest>,
    pub(crate) mvcc: Option<LsmMvccInner>,
    pub(crate) compaction_filters: Arc<Mutex<Vec<CompactionFilter>>>,
}

You’ll notice that there are two separate locks here - one for the accessing the LSM state and another which is a simple mutex. The idea is, to increase the write throughput (avoiding time waiting for disk I/O), any time you want to modify the actual state of the LSM engine (think either flushing a memtable or compacting SST’s), you acquire the state_lock mutex in conjunction with the write lock to make the actual changes. Any other time where you just want to write to the memtable, you can do a read lock on state.

If you depended on the write lock alone, you would acquire it and then be stuck waiting for disk I/O while preventing any writes to the memtable resulting in a latency spike.

There are also a few channels set up for cross thread communication - namely for flushing the frozen memtable and for running a compaction.

Source: https://www.creativcoder.dev/blog/what-is-a-lsm-tree

The MemTable

Here’s the definition for the MemTable.

pub struct MemTable {
    id: usize,
    map: Arc<SkipMap<KeyBytes, Bytes>>,
    approximate_size: Arc<AtomicUsize>,
    wal: Option<Wal>,
}

I’m using the crossbeam crate’s implementation of a SkipList. It supports concurrency out of the box, so no need to wrap it in a mutex. Also, all key value pairs will be written as bytes and the bytes crate provides an implementation which stores references to the underlying bytes on disk, so cloning is cheap.

KeyBytes is a wrapper around bytes::Bytes, which you can understand better by looking at this file in the repo.

The implementation for put and delete on the memtable is relatively straightforward. The delete implementation is LSM engines requires writing a tombstone value for the key, so it just uses the same put interface.

impl Memtable {
  pub fn put(&self, _key: &[u8], _value: &[u8]) -> Result<()> {
    self.map
        .insert(Bytes::from(_key.to_vec()), Bytes::from(_value.to_vec()));
    let current_size = self
        .approximate_size
        .load(std::sync::atomic::Ordering::Relaxed);
    let key_value_size = _key.len() + _value.len();
    let new_size = current_size + key_value_size;
    self.approximate_size
        .store(new_size, std::sync::atomic::Ordering::Relaxed);
    Ok(())
  }
}

Now, the more interesting part here is freezing the memtable, which is handled in LsmStorageInner.

impl LsmStorageInner {
    /// Put a key-value pair into the storage by writing into the current memtable.
    pub fn put(&self, _key: &[u8], _value: &[u8]) -> Result<()> {
        let state = self.state.read();
        let result = state.memtable.put(_key, _value);
        if state.memtable.approximate_size() >= self.options.target_sst_size {
            let state_lock = self.state_lock.lock();
            if state.memtable.approximate_size() >= self.options.target_sst_size {
                drop(state);
                self.force_freeze_memtable(&state_lock)?
            }
        }
        result
    }

    /// Remove a key from the storage by writing an empty value.
    pub fn delete(&self, _key: &[u8]) -> Result<()> {
        let state = self.state.read();
        let result = state.memtable.put(_key, &[]);
        if state.memtable.approximate_size() >= self.options.target_sst_size {
            let state_lock = self.state_lock.lock();
            if state.memtable.approximate_size() >= self.options.target_sst_size {
                drop(state);
                self.force_freeze_memtable(&state_lock)?
            }
        }
        result
    }

    /// Force freeze the current memtable to an immutable memtable
    pub fn force_freeze_memtable(&self, _state_lock_observer: &MutexGuard<'_, ()>) -> Result<()> {
        let state_read = self.state.read();
        let current_memtable = Arc::clone(&state_read.memtable);
        drop(state_read);

        let new_memtable = Arc::new(MemTable::create(self.next_sst_id()));

        let mut state_guard = self.state.write();
        let state = Arc::make_mut(&mut state_guard);
        state.imm_memtables.insert(0, current_memtable);
        state.memtable = new_memtable;

        Ok(())
    }
}

We make a simple check to determine if the memtable has exceeded the size limit, and if so freeze the memtable and place it inside a vector. More importantly, we need to create a new memtable to accept future writes. For that tiny duration, we obviously can’t accept any new writes.

This is great, so far we can write to the memtable. Now, what if you want to read the value back from the memtable. For now, since our write path consists of only the memtables, we can just iterate over them and check if any of them have the key we are looking for. Currently, we can support point queries but not range queries, so let’s go ahead and support that.

/// Get a key from the storage. In day 7, this can be further optimized by using a bloom filter.
impl LsmStorageInner {
    pub fn get(&self, key: &[u8]) -> Result<Option<Bytes>> {
        //  first probe the memtables
        let state_guard = self.state.read();
        let mut memtables = Vec::new();
        memtables.push(Arc::clone(&state_guard.memtable));
        memtables.extend(
            state_guard
                .imm_memtables
                .iter()
                .map(|memtable| Arc::clone(memtable)),
        );
        for memtable in memtables {
            if let Some(value) = memtable.get(key) {
                if value.is_empty() {
                    return Ok(None);
                }
                return Ok(Some(value));
            }
        }
    }
}

This works because, by default, the SkipMap stores data in a sorted order.

Awesome! We have in-memory buffer we can write to and read from. If only we had unlimited memory, but alas!

Sorted String Table

Without infinite memory, let’s turn our attention to on-disk representations for our data. Let’s start building this out visually and then represent that in code. We can put all our data into blocks and then use those blocks as part of a file.

We will also need to store some metadata for the file so that we can quickly tell what portion of the file we want to read. We don’t want to be loading the entire file into memory, especially when we can get away with reading just a specific block within the file.

Here’s the encoding format we will use. We include a bloom filter with each SST to quickly tell us whether a key does not exist in the SST file. This is a great addition because bloom filters can tell us with a 100% certainty whether something does not exist, saving us the I/O of loading any blocks from this file.

If you want to understand bloom filters better, this is a great post by Sam Who

-----------------------------------------------------------------------------------------------------
|         Block Section         |                            Meta Section                           |
-----------------------------------------------------------------------------------------------------
| data block | ... | data block | metadata | meta block offset | bloom filter | bloom filter offset |
|                               |  varlen  |         u32       |    varlen    |        u32          |
-----------------------------------------------------------------------------------------------------

Now, let’s dive a little deeper and look at what each individual block looks like.

----------------------------------------------------------------------------------------------------
|             Data Section             |              Offset Section             |      Extra      |
----------------------------------------------------------------------------------------------------
| Entry #1 | Entry #2 | ... | Entry #N | Offset #1 | Offset #2 | ... | Offset #N | num_of_elements |
----------------------------------------------------------------------------------------------------

And below is what each entry within a block looks like

-----------------------------------------------------------------------
|                           Entry #1                            | ... |
-----------------------------------------------------------------------
| key_len (2B) | key (keylen) | value_len (2B) | value (varlen) | ... |
-----------------------------------------------------------------------

Let’s represent this in code. First thing is a struct that can help us quickly build individual blocks out. There’s a simple encoding scheme in the block which uses key overlap with the first key to compress the size of the data.

use crate::key::{Key, KeySlice, KeyVec};

use super::Block;

/// Builds a block.
pub struct BlockBuilder {
    /// Offsets of each key-value entries.
    offsets: Vec<u16>,
    /// All serialized key-value pairs in the block.
    data: Vec<u8>,
    /// The expected block size.
    block_size: usize,
    /// The first key in the block
    first_key: KeyVec,
}

impl BlockBuilder {
    /// Creates a new block builder.
    pub fn new(block_size: usize) -> Self {
        BlockBuilder {
            offsets: Vec::new(),
            data: Vec::new(),
            block_size,
            first_key: Key::new(),
        }
    }

    /// Adds a key-value pair to the block. Returns false when the block is full.
    #[must_use]
    pub fn add(&mut self, key: KeySlice, value: &[u8]) -> bool {
        //  get the overlap of the key with the first key
        let key_overlap = key
            .into_inner()
            .iter()
            .zip(self.first_key.as_key_slice().into_inner().iter())
            .take_while(|(a, b)| a == b)
            .count() as u16;
        let key_overlap_bytes = key_overlap.to_le_bytes();
        let rest_of_key = &(key.into_inner())[key_overlap as usize..];
        let rest_of_key_len = (rest_of_key.len() as u16).to_le_bytes();

        let value_length = value.len();
        let value_length_bytes = (value_length as u16).to_le_bytes();
        let entry_size = 2 + rest_of_key.len() + 2 + value_length + 2;

        if self.data.len() + self.offsets.len() + entry_size > self.block_size
            && self.first_key.raw_ref().len() > 0
        {
            return false;
        }

        self.offsets.push(self.data.len() as u16);

        self.data.extend_from_slice(&key_overlap_bytes);
        self.data.extend_from_slice(&rest_of_key_len);
        self.data.extend_from_slice(rest_of_key);
        self.data.extend_from_slice(&value_length_bytes);
        self.data.extend_from_slice(value);

        if self.first_key.raw_ref().len() == 0 {
            let mut new_key = Key::new();
            new_key.set_from_slice(key);
            self.first_key = new_key;
        }
        true
    }

    /// Check if there is no key-value pair in the block.
    pub fn is_empty(&self) -> bool {
        self.offsets.is_empty()
    }

    /// Finalize the block.
    pub fn build(self) -> Block {
        Block {
            data: self.data,
            offsets: self.offsets,
        }
    }

    pub fn size(&self) -> usize {
        self.offsets.len()
    }
}

Next, we need a struct that can help us build out an SSTable in a similar fashion by using the block builder

use std::path::Path;
use std::sync::Arc;

use anyhow::Result;

use super::{bloom::Bloom, BlockMeta, FileObject, SsTable};
use crate::{
    block::{Block, BlockBuilder},
    key::{Key, KeySlice},
    lsm_storage::BlockCache,
};

/// Builds an SSTable from key-value pairs.
pub struct SsTableBuilder {
    builder: BlockBuilder,
    first_key: Vec<u8>,
    last_key: Vec<u8>,
    data: Vec<u8>,
    pub(crate) meta: Vec<BlockMeta>,
    block_size: usize,
    key_hashes: Vec<u32>,
}

impl SsTableBuilder {
    /// Create a builder based on target block size.
    pub fn new(block_size: usize) -> Self {
        SsTableBuilder {
            builder: BlockBuilder::new(block_size),
            first_key: Vec::new(),
            last_key: Vec::new(),
            data: Vec::new(),
            meta: Vec::new(),
            block_size,
            key_hashes: Vec::new(),
        }
    }

    /// Adds a key-value pair to SSTable.
    ///
    /// Note: You should split a new block when the current block is full.(`std::mem::replace` may
    /// be helpful here)
    pub fn add(&mut self, key: KeySlice, value: &[u8]) {
        if self.first_key.is_empty() {
            self.first_key = key.to_key_vec().raw_ref().to_vec();
        }

        self.key_hashes
            .push(farmhash::fingerprint32(key.into_inner()));

        if self.builder.add(key, value) {
            self.last_key = key.to_key_vec().raw_ref().to_vec();
            return;
        }

        self.freeze_block();

        assert!(self.builder.add(key, value));
        self.first_key = key.to_key_vec().raw_ref().to_vec();
        self.last_key = key.to_key_vec().raw_ref().to_vec();
    }

    /// This function will take current block builder, build it and replace it with a fresh block builder
    /// It will add the block to the SSTable data and then create and store the metadata for this block
    fn freeze_block(&mut self) {
        //  the block is full, split block and replace older builder
        let builder = std::mem::replace(&mut self.builder, BlockBuilder::new(self.block_size));
        let block = builder.build();
        let encoded_block = Block::encode(&block);

        //  get metadata for split block
        let block_meta = BlockMeta {
            offset: self.data.len(),
            first_key: Key::from_vec(self.first_key.clone()).into_key_bytes(),
            last_key: Key::from_vec(self.last_key.clone()).into_key_bytes(),
        };

        self.data.extend_from_slice(&encoded_block);
        self.meta.push(block_meta);
    }

    /// Get the estimated size of the SSTable.
    ///
    /// Since the data blocks contain much more data than meta blocks, just return the size of data
    /// blocks here.
    pub fn estimated_size(&self) -> usize {
        self.data.len()
    }

    /// Builds the SSTable and writes it to the given path. Use the `FileObject` structure to manipulate the disk objects.
    pub fn build(
        mut self,
        id: usize,
        block_cache: Option<Arc<BlockCache>>,
        path: impl AsRef<Path>,
    ) -> Result<SsTable> {
        self.freeze_block();

        //  create the bloom filter
        let bits_per_key = Bloom::bloom_bits_per_key(self.key_hashes.len(), 0.01);
        let bloom_filter = Bloom::build_from_key_hashes(&self.key_hashes, bits_per_key);

        let mut encoded_sst: Vec<u8> = Vec::new();
        encoded_sst.extend_from_slice(&self.data);

        //  encode meta section for each block and add it to encoding
        let mut encoded_meta: Vec<u8> = Vec::new();
        BlockMeta::encode_block_meta(&self.meta, &mut encoded_meta);
        encoded_sst.extend_from_slice(&encoded_meta);

        //  encode the meta block offset in the next 4 bytes
        let data_len = (self.data.len() as u32).to_le_bytes();
        encoded_sst.extend_from_slice(&data_len);

        //  encode the bloom filter and add it to encoded table
        let bloom_filter_offset = encoded_sst.len() as u32;
        bloom_filter.encode(&mut encoded_sst);
        encoded_sst.extend_from_slice(&bloom_filter_offset.to_le_bytes());

        //  write the entire encoding to disk
        let file = FileObject::create(path.as_ref(), encoded_sst)?;
        Ok(SsTable {
            file,
            block_meta_offset: self.data.len(),
            id,
            block_cache: None,
            first_key: self.meta.first().unwrap().first_key.clone(),
            last_key: self.meta.last().unwrap().last_key.clone(),
            block_meta: self.meta,
            bloom: Some(bloom_filter),
            max_ts: 0,
        })
    }

    #[cfg(test)]
    pub(crate) fn build_for_test(self, path: impl AsRef<Path>) -> Result<SsTable> {
        self.build(0, None, path)
    }
}

Now, there’s a couple of missing pieces regarding the encoding & decoding of the table and the blocks, so let’s fill those out. Here’s the encoding and decoding for the block.

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct BlockMeta {
    /// Offset of this data block.
    pub offset: usize,
    /// The first key of the data block.
    pub first_key: KeyBytes,
    /// The last key of the data block.
    pub last_key: KeyBytes,
}

impl BlockMeta {
    /// Encode block meta to a buffer.
    /// You may add extra fields to the buffer,
    /// in order to help keep track of `first_key` when decoding from the same buffer in the future.
    pub fn encode_block_meta(block_meta: &[BlockMeta], buf: &mut Vec<u8>) {
        for ind_block_meta in block_meta {
            let offset_bytes = (ind_block_meta.offset as u16).to_le_bytes();
            buf.extend_from_slice(&offset_bytes);

            let first_key_length = ind_block_meta.first_key.len() as u16;
            let first_key_length_bytes = first_key_length.to_le_bytes();
            buf.extend_from_slice(&first_key_length_bytes);
            buf.extend_from_slice(ind_block_meta.first_key.raw_ref());

            let last_key_length = ind_block_meta.last_key.len() as u16;
            let last_key_length_bytes = last_key_length.to_le_bytes();
            buf.extend_from_slice(&last_key_length_bytes);
            buf.extend_from_slice(ind_block_meta.last_key.raw_ref());
        }
    }

    /// Decode block meta from a buffer.
    pub fn decode_block_meta(mut buf: impl Buf) -> Vec<BlockMeta> {
        let mut block_metas: Vec<BlockMeta> = Vec::new();

        while buf.remaining() > 0 {
            let offset = buf.get_u16_le() as usize;

            let first_key_length = buf.get_u16_le() as usize;
            let mut first_key = vec![0; first_key_length];
            buf.copy_to_slice(&mut first_key);

            let last_key_length = buf.get_u16_le() as usize;
            let mut last_key = vec![0; last_key_length];
            buf.copy_to_slice(&mut last_key);

            block_metas.push(BlockMeta {
                offset,
                first_key: Key::from_vec(first_key).into_key_bytes(),
                last_key: Key::from_vec(last_key).into_key_bytes(),
            });
        }

        block_metas
    }
}

Finally, we need the ability to open a file from disk and decode the bytes. So, let’s add that in as well.

pub(crate) mod bloom;
mod builder;
mod iterator;

use std::fs::File;
use std::path::Path;
use std::sync::Arc;

use anyhow::Result;
pub use builder::SsTableBuilder;
use bytes::Buf;
pub use iterator::SsTableIterator;

use crate::block::Block;
use crate::key::{Key, KeyBytes, KeySlice};
use crate::lsm_storage::BlockCache;

use self::bloom::Bloom;

/// A file object.
pub struct FileObject(Option<File>, u64);

impl FileObject {
    pub fn read(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
        use std::os::unix::fs::FileExt;
        let mut data = vec![0; len as usize];
        self.0
            .as_ref()
            .unwrap()
            .read_exact_at(&mut data[..], offset)?;
        Ok(data)
    }

    pub fn size(&self) -> u64 {
        self.1
    }

    /// Create a new file object (day 2) and write the file to the disk (day 4).
    pub fn create(path: &Path, data: Vec<u8>) -> Result<Self> {
        std::fs::write(path, &data)?;
        File::open(path)?.sync_all()?;
        Ok(FileObject(
            Some(File::options().read(true).write(false).open(path)?),
            data.len() as u64,
        ))
    }

    pub fn open(path: &Path) -> Result<Self> {
        let file = File::options().read(true).write(false).open(path)?;
        let size = file.metadata()?.len();
        Ok(FileObject(Some(file), size))
    }
}

/// An SSTable.
pub struct SsTable {
    /// The actual storage unit of SsTable, the format is as above.
    pub(crate) file: FileObject,
    /// The meta blocks that hold info for data blocks.
    pub(crate) block_meta: Vec<BlockMeta>,
    /// The offset that indicates the start point of meta blocks in `file`.
    pub(crate) block_meta_offset: usize,
    id: usize,
    block_cache: Option<Arc<BlockCache>>,
    first_key: KeyBytes,
    last_key: KeyBytes,
    pub(crate) bloom: Option<Bloom>,
    /// The maximum timestamp stored in this SST, implemented in week 3.
    max_ts: u64,
}

impl SsTable {
    #[cfg(test)]
    pub(crate) fn open_for_test(file: FileObject) -> Result<Self> {
        Self::open(0, None, file)
    }

    /// Open SSTable from a file.
    pub fn open(id: usize, block_cache: Option<Arc<BlockCache>>, file: FileObject) -> Result<Self> {
        let len = file.size();
        //  read the last 4 bytes to get the bloom filter offset
        let bloom_filter_offset_raw = file.read(len - 4, 4)?;
        let bloom_filter_offset = (&bloom_filter_offset_raw[..]).get_u32_le() as u64;

        //  use the bloom filter offset to read the data starting
        let raw_bloom_filter = file.read(bloom_filter_offset, len - bloom_filter_offset - 4)?;
        let bloom_filter = Bloom::decode(&raw_bloom_filter)?;

        //  read the 4 bytes before the bloom offset to get the meta offset
        let raw_meta_offset = file.read(bloom_filter_offset - 4, 4)?;
        let meta_offset = (&raw_meta_offset[..]).get_u32_le() as u64;

        //  use the meta offset to read the metadata from the file
        let raw_meta = file.read(meta_offset as u64, bloom_filter_offset - 4 - meta_offset)?;
        let meta = BlockMeta::decode_block_meta(raw_meta.as_slice());

        Ok(SsTable {
            file,
            block_meta_offset: meta_offset as usize,
            id,
            block_cache,
            first_key: meta.first().unwrap().first_key.clone(),
            last_key: meta.last().unwrap().last_key.clone(),
            block_meta: meta,
            bloom: Some(bloom_filter),
            max_ts: 0,
        })
    }

    /// Create a mock SST with only first key + last key metadata
    pub fn create_meta_only(
        id: usize,
        file_size: u64,
        first_key: KeyBytes,
        last_key: KeyBytes,
    ) -> Self {
        Self {
            file: FileObject(None, file_size),
            block_meta: vec![],
            block_meta_offset: 0,
            id,
            block_cache: None,
            first_key,
            last_key,
            bloom: None,
            max_ts: 0,
        }
    }

    /// Read a block from the disk.
    pub fn read_block(&self, block_idx: usize) -> Result<Arc<Block>> {
        let block_offset_start = self.block_meta[block_idx].offset;
        let block_offset_end = if block_idx + 1 < self.block_meta.len() {
            self.block_meta[block_idx + 1].offset
        } else {
            self.block_meta_offset
        };
        let block_len = block_offset_end - block_offset_start;
        let block_data_raw = self
            .file
            .read(block_offset_start as u64, block_len as u64)?;
        let block_data = Block::decode(&block_data_raw);
        Ok(Arc::new(block_data))
    }

    /// Read a block from disk, with block cache. (Day 4)
    pub fn read_block_cached(&self, block_idx: usize) -> Result<Arc<Block>> {
        if let Some(block_cache) = &self.block_cache {
            let block = block_cache
                .try_get_with((self.id, block_idx), || self.read_block(block_idx))
                .map_err(|e| {
                    println!("Error: {:?}", e);
                    anyhow::anyhow!(e)
                })?;
            Ok(block)
        } else {
            self.read_block(block_idx)
        }
    }

    /// Find the block that may contain `key`.
    /// Note: You may want to make use of the `first_key` stored in `BlockMeta`.
    /// You may also assume the key-value pairs stored in each consecutive block are sorted.
    pub fn find_block_idx(&self, key: KeySlice) -> usize {
        self.block_meta
            .partition_point(|meta| meta.first_key.as_key_slice() <= key) //  parition_point does binary search
            .saturating_sub(1)
    }

    /// Get number of data blocks.
    pub fn num_of_blocks(&self) -> usize {
        self.block_meta.len()
    }

    pub fn first_key(&self) -> &KeyBytes {
        &self.first_key
    }

    pub fn last_key(&self) -> &KeyBytes {
        &self.last_key
    }

    pub fn table_size(&self) -> u64 {
        self.file.1
    }

    pub fn sst_id(&self) -> usize {
        self.id
    }

    pub fn max_ts(&self) -> u64 {
        self.max_ts
    }
}

Now, we need to wire this all up back to our storage engine. We will have a flush thread running which executes every X number of milliseconds and it will pick the oldest frozen memtable and attempt to flush that.

Don’t worry about the fact that the sstable is pushed into something called L0, we’ll get to that later.

impl LsmStorageInner {
    /// Force flush the earliest-created immutable memtable to disk
    pub fn force_flush_next_imm_memtable(&self) -> Result<()> {
        let _state_lock = self.state_lock.lock();
        let oldest_memtable = {
            let state_guard = self.state.write();
            let oldest_memtable = state_guard
                .imm_memtables
                .last()
                .expect("No memtable found")
                .clone();
            oldest_memtable
        };

        let mut sst_builder = SsTableBuilder::new(self.options.block_size);
        oldest_memtable.flush(&mut sst_builder)?;
        let sst = Arc::new(sst_builder.build(
            oldest_memtable.id(),
            Some(Arc::clone(&self.block_cache)),
            self.path_of_sst(oldest_memtable.id()),
        )?);

        {
            let mut state_guard = self.state.write();
            let mut snapshot = state_guard.as_ref().clone();
            let oldest_memtable = snapshot.imm_memtables.pop().expect("No memtable found");
            snapshot.l0_sstables.insert(0, oldest_memtable.id());
            snapshot.sstables.insert(oldest_memtable.id(), sst);
            *state_guard = Arc::new(snapshot);
        }

        Ok(())
    }

    fn trigger_flush(&self) -> Result<()> {
        {
            let state_guard = self.state.read();
            if state_guard.imm_memtables.len() < self.options.num_memtable_limit {
                return Ok(());
            }
        }

        self.force_flush_next_imm_memtable()?;
        Ok(())
    }

    pub(crate) fn spawn_flush_thread(
        self: &Arc<Self>,
        rx: crossbeam_channel::Receiver<()>,
    ) -> Result<Option<std::thread::JoinHandle<()>>> {
        let this = self.clone();
        let handle = std::thread::spawn(move || {
            let ticker = crossbeam_channel::tick(Duration::from_millis(50));
            loop {
                crossbeam_channel::select! {
                    recv(ticker) -> _ => if let Err(e) = this.trigger_flush() {
                        eprintln!("flush failed: {}", e);
                    },
                    recv(rx) -> _ => return
                }
            }
        });
        Ok(Some(handle))
    }
}

Great! We now have the mechanism to take our frozen memtable, write that to a disk friendly format, encode it and then read it back if we want to later.

Now, we already have a mechanism to read back from the memtable, so we also need a way to read back from an SSTable. For the memtable we just iterated over all the memtables in reverse chronological order and checked if we could find the key. The first instance we found was the key we were looking for.

We could do the exact same thing with the SSTables as well. First check memtables for the key and if it can’t be found there start checking the SSTables one by one, and that would actually work fine. However, let’s make this a bit more efficient in a way we can also support range queries.

It’s Iterators All The Way Down

Assume that we have a database engine that ends up in the following state.

put("x", 5)
put("y", 3)
put("x", Tombstone)
//  freeze and flush occurs
put("x", 9)
put("y", 3)
put("y", 1)
//  freeze and flush occurs
put("a", 12)
put("b", Tombstone)

Now, at this point the state looks like

memtable ("a" -> 12, "b" -> Tombstone)
l0 -> sst1 ("x" -> 9, "y" -> 3, "y" -> 1), sst2 ("x" -> 5, "y" -> 3, "x" -> Tombstone)

We have two different data sources here and if a user wanted to get all the key-value pairs in storage, we would have to iterate over every source simultaneously and figure out what the correct value for it is.

So, let’s define a trait that an iterator over any data source would need to stick to. This makes it easier to define new iterators as and when required.

pub trait StorageIterator {
    type KeyType<'a>: PartialEq + Eq + PartialOrd + Ord
    where
        Self: 'a;

    /// Get the current value.
    fn value(&self) -> &[u8];

    /// Get the current key.
    fn key(&self) -> Self::KeyType<'_>;

    /// Check if the current iterator is valid.
    fn is_valid(&self) -> bool;

    /// Move to the next position.
    fn next(&mut self) -> anyhow::Result<()>;

    /// Number of underlying active iterators for this iterator.
    fn num_active_iterators(&self) -> usize {
        1
    }
}

The trait has an associated type for the type of key it returns.

MemTable Scan

So, let’s go ahead and define the memtable iterator first so that we can do a scan across a memtable.

type SkipMapRangeIter<'a> =
    crossbeam_skiplist::map::Range<'a, Bytes, (Bound<Bytes>, Bound<Bytes>), Bytes, Bytes>;

#[self_referencing]
pub struct MemTableIterator {
    /// Stores a reference to the skipmap.
    map: Arc<SkipMap<Bytes, Bytes>>,
    /// Stores a skipmap iterator that refers to the lifetime of `MemTableIterator` itself.
    #[borrows(map)]
    #[not_covariant]
    iter: SkipMapRangeIter<'this>,
    /// Stores the current key-value pair.
    item: (Bytes, Bytes),
}

impl StorageIterator for MemTableIterator {
    type KeyType<'a> = KeySlice<'a>;

    fn value(&self) -> &[u8] {
        self.with_item(|item| &item.1)
    }

    fn key(&self) -> KeySlice {
        let key = self.with_item(|item| &item.0);
        Key::from_slice(key)
    }

    fn is_valid(&self) -> bool {
        !self.with_item(|item| item.0.is_empty())
    }

    fn next(&mut self) -> Result<()> {
        let next_entry = self.with_iter_mut(|iter| {
            iter.next()
                .map(|entry| (entry.key().clone(), entry.value().clone()))
                .unwrap_or_else(|| (Bytes::from_static(&[]), Bytes::from_static(&[])))
        });
        self.with_item_mut(|item| *item = next_entry);
        Ok(())
    }
}

impl MemTable {
    /// Get an iterator over a range of keys.
    pub fn scan(&self, _lower: Bound<&[u8]>, _upper: Bound<&[u8]>) -> MemTableIterator {
        let mut iterator = MemTableIteratorBuilder {
            map: self.map.clone(),
            iter_builder: |map| map.range((map_bound(_lower), map_bound(_upper))),
            item: (Bytes::new(), Bytes::new()),
        }
        .build();
        iterator.next().unwrap();
        iterator
    }
}

So, this was relatively simple to do because we can build an iterator using the Range operator provided by crossbeam.

SSTable Scan

Okay, now we need something that can scan the sorted string table. Now, when scanning an SSTable, each SSTable is composed of individual blocks and those blocks are composed of key-value pairs, so we actually need 2 iterators - one for a block and one for a table itself.

Let’s define the one for the block first. There’s nothing fancy in this iterator, at it’s heart it uses binary search to look through the list of keys. However, since we’re dealing with raw bytes, there’s a fair amount of decoding that goes on here.

use std::sync::Arc;

use bytes::Buf;

use crate::key::{KeySlice, KeyVec};

use super::Block;

/// Iterates on a block.
pub struct BlockIterator {
    /// The internal `Block`, wrapped by an `Arc`
    block: Arc<Block>,
    /// The current key, empty represents the iterator is invalid
    key: KeyVec,
    /// the value range from the block
    value_range: (usize, usize),
    /// Current index of the key-value pair, should be in range of [0, num_of_elements)
    idx: usize,
    /// The first key in the block
    first_key: KeyVec,
}

impl BlockIterator {
    fn new(block: Arc<Block>) -> Self {
        Self {
            first_key: BlockIterator::decode_first_key(&block),
            block,
            key: KeyVec::new(),
            value_range: (0, 0),
            idx: 0,
        }
    }

    fn decode_first_key(block: &Arc<Block>) -> KeyVec {
        let mut buf = &block.data[..];
        let overlap_length = buf.get_u16_le(); //  read the overlap length (should be 0)
        assert_eq!(overlap_length, 0);
        let key_length = buf.get_u16_le(); //  read the key length
        let key = &buf[..key_length as usize];
        KeyVec::from_vec(key.to_vec())
    }

    /// Creates a block iterator and seek to the first entry.
    pub fn create_and_seek_to_first(block: Arc<Block>) -> Self {
        let mut iter = BlockIterator::new(block);
        iter.seek_to_first();
        iter
    }

    /// Creates a block iterator and seek to the first key that >= `key`.
    pub fn create_and_seek_to_key(block: Arc<Block>, key: KeySlice) -> Self {
        let mut iter = BlockIterator::new(block);
        iter.seek_to_key(key);
        iter
    }

    /// Returns the key of the current entry.
    pub fn key(&self) -> KeySlice {
        self.key.as_key_slice()
    }

    /// Returns the value of the current entry.
    pub fn value(&self) -> &[u8] {
        let value_range = self.value_range;
        let value_raw = &self.block.data[value_range.0..value_range.1];
        value_raw
    }

    /// Returns true if the iterator is valid.
    /// Note: You may want to make use of `key`
    pub fn is_valid(&self) -> bool {
        if self.key.is_empty() {
            return false;
        }
        true
    }

    /// Seeks to the first key in the block.
    pub fn seek_to_first(&mut self) {
        self.seek_to(0);
        self.idx = 0;
    }

    fn seek_to(&mut self, index: usize) {
        let offset = self.block.offsets[index] as usize;
        let data_to_consider = &self.block.data[offset..];

        let (key_overlap_length_raw, rest) = data_to_consider.split_at(2);
        let key_overlap_length = u16::from_le_bytes(key_overlap_length_raw.try_into().unwrap());

        let (key_length_raw, rest) = rest.split_at(2);
        let key_length = u16::from_le_bytes(key_length_raw.try_into().unwrap());

        let (key, rest) = rest.split_at(key_length as usize);
        let key_overlap = &(self.first_key.clone().into_inner())[..key_overlap_length as usize];
        let mut full_key = Vec::new();
        full_key.extend_from_slice(&key_overlap);
        full_key.extend_from_slice(&key);
        self.key = KeyVec::from_vec(full_key);

        let (value_length_raw, rest) = rest.split_at(2);
        let value_length = u16::from_le_bytes(value_length_raw.try_into().unwrap());

        let (_, _) = rest.split_at(value_length as usize);
        let new_value_start = offset + 2 + 2 + key_length as usize;
        self.value_range = (
            new_value_start + 2,
            new_value_start + 2 + value_length as usize,
        );
    }

    /// Move to the next key in the block.
    pub fn next(&mut self) {
        self.idx += 1;

        if self.idx >= self.block.offsets.len() {
            self.key.clear();
            self.value_range = (0, 0);
            return;
        }

        self.seek_to(self.idx);
    }

    /// Seek to the first key that >= `key`.
    /// Note: You should assume the key-value pairs in the block are sorted when being added by
    /// callers.
    pub fn seek_to_key(&mut self, key: KeySlice) {
        let mut low = 0;
        let mut high = self.block.offsets.len() - 1;

        while low <= high {
            let mid = low + (high - low) / 2;
            self.seek_to(mid);
            self.idx = mid;
            let mid_key = self.key.as_key_slice();

            match mid_key.cmp(&key) {
                std::cmp::Ordering::Less => low = mid + 1,
                std::cmp::Ordering::Greater => {
                    if mid == 0 {
                        break;
                    }
                    high = mid - 1;
                }
                std::cmp::Ordering::Equal => return,
            }
        }

        if low >= self.block.offsets.len() {
            self.key.clear();
            self.value_range = (0, 0);
            return;
        }
        self.idx = low;
        self.seek_to(self.idx);
    }
}

Okay, now let’s focus on the table iterator. The table iterator will use the block iterator from above and will determine which block needs to be read. If the iterator just needs to go to the first key, it will read the first block and if the iterator needs to look for a specific key, it will do a binary search on the metadata of the blocks to determine the key.

use std::sync::Arc;

use anyhow::Result;

use super::SsTable;
use crate::{block::BlockIterator, iterators::StorageIterator, key::KeySlice};

/// An iterator over the contents of an SSTable.
pub struct SsTableIterator {
    table: Arc<SsTable>,
    blk_iter: BlockIterator,
    blk_idx: usize,
}

impl SsTableIterator {
    /// Create a new iterator and seek to the first key-value pair in the first data block.
    pub fn create_and_seek_to_first(table: Arc<SsTable>) -> Result<Self> {
        //  get the first block from the sstable and build an iterator on top of it
        let block = table.read_block_cached(0)?;
        let block_iterator = BlockIterator::create_and_seek_to_first(block);
        let iter = SsTableIterator {
            table,
            blk_iter: block_iterator,
            blk_idx: 0,
        };
        Ok(iter)
    }

    /// Seek to the first key-value pair in the first data block.
    pub fn seek_to_first(&mut self) -> Result<()> {
        //  get metata for first block index
        let block = self.table.read_block_cached(0)?;
        let block_iterator = BlockIterator::create_and_seek_to_first(block);
        self.blk_idx = 0;
        self.blk_iter = block_iterator;
        Ok(())
    }

    fn seek_to(table: &Arc<SsTable>, key: KeySlice) -> Result<(usize, BlockIterator)> {
        let mut block_index = table.find_block_idx(key);
        let block = table.read_block_cached(block_index).unwrap();
        let mut block_iter = BlockIterator::create_and_seek_to_key(block, key);
        if !block_iter.is_valid() {
            block_index += 1;
            if block_index < table.num_of_blocks() {
                block_iter =
                    BlockIterator::create_and_seek_to_first(table.read_block_cached(block_index)?);
            }
        }
        Ok((block_index, block_iter))
    }

    /// Create a new iterator and seek to the first key-value pair which >= `key`.
    pub fn create_and_seek_to_key(table: Arc<SsTable>, key: KeySlice) -> Result<Self> {
        let (block_index, block_iter) = Self::seek_to(&table, key)?;
        let iter = SsTableIterator {
            table,
            blk_iter: block_iter,
            blk_idx: block_index,
        };
        Ok(iter)
    }

    /// Seek to the first key-value pair which >= `key`.
    /// Note: You probably want to review the handout for detailed explanation when implementing
    /// this function.
    pub fn seek_to_key(&mut self, key: KeySlice) -> Result<()> {
        let (block_index, block_iter) = Self::seek_to(&self.table, key)?;
        self.blk_iter = block_iter;
        self.blk_idx = block_index;
        Ok(())
    }
}

impl StorageIterator for SsTableIterator {
    type KeyType<'a> = KeySlice<'a>;

    /// Return the `key` that's held by the underlying block iterator.
    fn key(&self) -> KeySlice {
        self.blk_iter.key()
    }

    /// Return the `value` that's held by the underlying block iterator.
    fn value(&self) -> &[u8] {
        self.blk_iter.value()
    }

    /// Return whether the current block iterator is valid or not.
    fn is_valid(&self) -> bool {
        self.blk_iter.is_valid()
    }

    /// Move to the next `key` in the block.
    /// Note: You may want to check if the current block iterator is valid after the move.
    fn next(&mut self) -> Result<()> {
        if self.is_valid() {
            self.blk_iter.next();
            if !self.blk_iter.is_valid() {
                self.blk_idx += 1;
                if self.blk_idx < self.table.num_of_blocks() {
                    let new_block = self.table.read_block_cached(self.blk_idx)?;
                    let new_block_iter = BlockIterator::create_and_seek_to_first(new_block);
                    self.blk_iter = new_block_iter;
                }
            }
        }
        Ok(())
    }
}

This is great, we now have two iterators for 2 different types of data - one for the memtable and one for the sstable. But we have multiple iterators for the memtable and multiple iterators for the sstable. We need a way to combine them.

Merging Iterators

Now, we need a way to represent multiple iterators of the same type as one iterator. This is the classic k-way merge algorithm implemented using a min heap.

First, let’s define a custom heap that will compare our keys to determine which key should be the first one.

struct HeapWrapper<I: StorageIterator>(pub usize, pub Box<I>);

impl<I: StorageIterator> PartialEq for HeapWrapper<I> {
    fn eq(&self, other: &Self) -> bool {
        self.partial_cmp(other).unwrap() == cmp::Ordering::Equal
    }
}

impl<I: StorageIterator> Eq for HeapWrapper<I> {}

impl<I: StorageIterator> PartialOrd for HeapWrapper<I> {
    #[allow(clippy::non_canonical_partial_ord_impl)]
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        match self.1.key().cmp(&other.1.key()) {
            cmp::Ordering::Greater => Some(cmp::Ordering::Greater),
            cmp::Ordering::Less => Some(cmp::Ordering::Less),
            cmp::Ordering::Equal => self.0.partial_cmp(&other.0),
        }
        .map(|x| x.reverse())
    }
}

impl<I: StorageIterator> Ord for HeapWrapper<I> {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.partial_cmp(other).unwrap()
    }
}

Now, we can define the actual merge iterator.

pub struct MergeIterator<I: StorageIterator> {
    iters: BinaryHeap<HeapWrapper<I>>,
    current: Option<HeapWrapper<I>>,
}

impl<I: StorageIterator> MergeIterator<I> {
    pub fn create(iters: Vec<Box<I>>) -> Self {
        if iters.len() == 0 {
            return MergeIterator {
                iters: BinaryHeap::new(),
                current: None,
            };
        }

        let mut heap: BinaryHeap<HeapWrapper<I>> = BinaryHeap::new();

        //  if none of the iterators are valid, just pick the last one as current
        if iters.iter().all(|iter| !iter.is_valid()) {
            let mut iters = iters;
            return MergeIterator {
                iters: heap,
                current: Some(HeapWrapper(0, iters.pop().unwrap())),
            };
        }

        for (index, iter) in iters.into_iter().enumerate() {
            if iter.is_valid() {
                let heap_wrapper = HeapWrapper(index, iter);
                heap.push(heap_wrapper);
            }
        }

        let current = heap.pop().unwrap();
        MergeIterator {
            iters: heap,
            current: Some(current),
        }
    }
}

impl<I: 'static + for<'a> StorageIterator<KeyType<'a> = KeySlice<'a>>> StorageIterator
    for MergeIterator<I>
{
    type KeyType<'a> = KeySlice<'a>;

    fn key(&self) -> KeySlice {
        self.current.as_ref().unwrap().1.key()
    }

    fn value(&self) -> &[u8] {
        self.current.as_ref().unwrap().1.value()
    }

    fn is_valid(&self) -> bool {
        self.current
            .as_ref()
            .map(|heap_wrapper| heap_wrapper.1.is_valid())
            .unwrap_or(false)
    }

    fn next(&mut self) -> Result<()> {
        let current = self.current.as_mut().unwrap();

        //  Check if there are any keys that are identical - advance the lower ranked iterators in that case
        while let Some(mut heap_wrapper) = self.iters.peek_mut() {
            if heap_wrapper.1.key() == current.1.key() {
                //  The current and the heap top have the same key. Ignore the heap top key because we organised by reverse
                //  chronological order when building the heap. The value in current should be what's upheld. Advance the top
                if let Err(e) = heap_wrapper.1.next() {
                    PeekMut::pop(heap_wrapper);
                    return Err(e);
                }

                if !heap_wrapper.1.is_valid() {
                    PeekMut::pop(heap_wrapper);
                }
            } else {
                break;
            }
        }

        //  advance the current iterator
        current.1.next()?;

        //  check if the current iterator continues to be valid - if not, replace with the top
        if !current.1.is_valid() {
            if let Some(heap_wrapper) = self.iters.pop() {
                self.current = Some(heap_wrapper);
            }
            return Ok(());
        }

        //  check if the current iterator should be replaced by the top value in the heap
        if let Some(mut heap_wrapper) = self.iters.peek_mut() {
            if current < &mut heap_wrapper {
                std::mem::swap(current, &mut *heap_wrapper);
            }
        }
        Ok(())
    }

    fn num_active_iterators(&self) -> usize {
        let heap_active_iters: usize = self
            .iters
            .iter()
            .map(|iter| iter.1.num_active_iterators())
            .sum();
        let current_active_iters: usize = self
            .current
            .iter()
            .map(|iter| iter.1.num_active_iterators())
            .sum();
        heap_active_iters + current_active_iters
    }
}

Let’s take a quick break and look at where we’re at in terms of scanning the data right now.

So we have the memtable iterator, the sst iterator & the merge iterator defined. It’s important to note that the order of iterators matters because we always want to check the latest sources of data first - that is the memtables in reverse chronological order and then the sstables in reverse chronological order.

Now, we need one more iterator which can combine two merge iterators into one so that we can search across multiple sources of data in the scan interface - this is the two merge iterator which is defined below.

use anyhow::Result;

use super::StorageIterator;

/// Merges two iterators of different types into one. If the two iterators have the same key, only
/// produce the key once and prefer the entry from A.
pub struct TwoMergeIterator<A: StorageIterator, B: StorageIterator> {
    a: A,
    b: B,
    // Add fields as need
    use_iterator: u8, // this can be 0 (use a), 1 (use b), 2 (use both)
}

impl<
        A: 'static + StorageIterator,
        B: 'static + for<'a> StorageIterator<KeyType<'a> = A::KeyType<'a>>,
    > TwoMergeIterator<A, B>
{
    pub fn create(a: A, b: B) -> Result<Self> {
        let use_iterator = TwoMergeIterator::decide_which_iter_to_use(&a, &b);
        Ok(TwoMergeIterator { a, b, use_iterator })
    }

    fn decide_which_iter_to_use(a: &A, b: &B) -> u8 {
        if !a.is_valid() && b.is_valid() {
            return 1;
        }
        if a.is_valid() && !b.is_valid() {
            return 0;
        }
        if !a.is_valid() && !b.is_valid() {
            return u8::MAX;
        }
        if a.key() < b.key() {
            0
        } else if a.key() > b.key() {
            1
        } else {
            2
        }
    }
}

impl<
        A: 'static + StorageIterator,
        B: 'static + for<'a> StorageIterator<KeyType<'a> = A::KeyType<'a>>,
    > StorageIterator for TwoMergeIterator<A, B>
{
    type KeyType<'a> = A::KeyType<'a>;

    fn key(&self) -> Self::KeyType<'_> {
        if self.use_iterator == 0 || self.use_iterator == 2 {
            return self.a.key();
        }
        self.b.key()
    }

    fn value(&self) -> &[u8] {
        if self.use_iterator == 0 || self.use_iterator == 2 {
            return self.a.value();
        }
        self.b.value()
    }

    fn is_valid(&self) -> bool {
        if self.use_iterator == u8::MAX {
            false
        } else if self.use_iterator == 0 {
            self.a.is_valid()
        } else {
            self.b.is_valid()
        }
    }

    fn next(&mut self) -> Result<()> {
        if self.use_iterator == 0 {
            //  advance the first iterator because the second one wasn't used
            if self.a.is_valid() {
                self.a.next()?;
            }
        } else if self.use_iterator == 1 {
            //  advance the second iterator because the first one wasn't used
            if self.b.is_valid() {
                self.b.next()?;
            }
        } else if self.use_iterator == 2 {
            //  advance both
            if self.a.is_valid() {
                self.a.next()?;
            }
            if self.b.is_valid() {
                self.b.next()?;
            }
        }
        self.use_iterator = TwoMergeIterator::decide_which_iter_to_use(&self.a, &self.b);
        Ok(())
    }

    fn num_active_iterators(&self) -> usize {
        self.a.num_active_iterators() + self.b.num_active_iterators()
    }
}

Now, this would have actually been enough at this point. However, the course defines two additional types of iterators to make changing things easier, so let’s just add those. However, the hierarchy of the iterators is the same as we saw in the above diagram. We will add 2 new layers but nothing fundamentally changed.

use std::{
    io::{self, ErrorKind},
    ops::Bound,
};

use anyhow::Result;
use bytes::Bytes;

use crate::{
    iterators::{
        merge_iterator::MergeIterator, two_merge_iterator::TwoMergeIterator, StorageIterator,
    },
    mem_table::MemTableIterator,
    table::SsTableIterator,
};

/// Represents the internal type for an LSM iterator. This type will be changed across the tutorial for multiple times.
type LsmIteratorInner =
    TwoMergeIterator<MergeIterator<MemTableIterator>, MergeIterator<SsTableIterator>>;

pub struct LsmIterator {
    inner: LsmIteratorInner,
    upper_bound: Bound<Bytes>,
    is_valid: bool,
}

impl LsmIterator {
    pub(crate) fn new(iter: LsmIteratorInner, upper_bound: Bound<Bytes>) -> Result<Self> {
        let mut lsm_iter = Self {
            is_valid: iter.is_valid(),
            inner: iter,
            upper_bound,
        };

        //  when an iterator is first created, there is the possibility that
        //  the very first key-value pair is a tombstone so need to account for that
        lsm_iter.skip_deleted_values()?;

        Ok(lsm_iter)
    }

    //  if the value associated with the key after calling next is an empty string
    //  this marks a tombstone. this key should be skipped so that the consumer
    //  of this iterator does not see it
    fn skip_deleted_values(&mut self) -> Result<()> {
        while self.inner.is_valid() && self.inner.value().is_empty() {
            self.inner.next()?;
            if !self.inner.is_valid() {
                self.is_valid = false;
                return Ok(());
            }
            match self.upper_bound.as_ref() {
                Bound::Included(key) => {
                    if self.inner.key().raw_ref() > key {
                        //invalidate the iterator
                        self.is_valid = false;
                    }
                }
                Bound::Excluded(key) => {
                    if self.inner.key().raw_ref() >= key {
                        //  invalidate the iterator
                        self.is_valid = false;
                    }
                }
                Bound::Unbounded => {}
            }
        }
        Ok(())
    }
}

impl StorageIterator for LsmIterator {
    type KeyType<'a> = &'a [u8];

    fn is_valid(&self) -> bool {
        self.is_valid
    }

    fn key(&self) -> &[u8] {
        self.inner.key().into_inner()
    }

    fn value(&self) -> &[u8] {
        self.inner.value()
    }

    fn next(&mut self) -> Result<()> {
        self.inner.next()?;
        if !self.inner.is_valid() {
            self.is_valid = false;
            return Ok(());
        }
        match self.upper_bound.as_ref() {
            Bound::Included(key) => {
                if self.inner.key().raw_ref() > key {
                    //invalidate the iterator
                    self.is_valid = false;
                }
            }
            Bound::Excluded(key) => {
                if self.inner.key().raw_ref() >= key {
                    //  invalidate the iterator
                    self.is_valid = false;
                }
            }
            Bound::Unbounded => {}
        }
        self.skip_deleted_values()?;
        Ok(())
    }

    fn num_active_iterators(&self) -> usize {
        self.inner.num_active_iterators()
    }
}

/// A wrapper around existing iterator, will prevent users from calling `next` when the iterator is
/// invalid. If an iterator is already invalid, `next` does not do anything. If `next` returns an error,
/// `is_valid` should return false, and `next` should always return an error.
pub struct FusedIterator<I: StorageIterator> {
    iter: I,
    has_errored: bool,
}

impl<I: StorageIterator> FusedIterator<I> {
    pub fn new(iter: I) -> Self {
        Self {
            iter,
            has_errored: false,
        }
    }
}

impl<I: StorageIterator> StorageIterator for FusedIterator<I> {
    type KeyType<'a> = I::KeyType<'a> where Self: 'a;

    fn is_valid(&self) -> bool {
        if self.has_errored {
            return false;
        }
        self.iter.is_valid()
    }

    fn key(&self) -> Self::KeyType<'_> {
        self.iter.key()
    }

    fn value(&self) -> &[u8] {
        self.iter.value()
    }

    fn next(&mut self) -> Result<()> {
        if self.has_errored {
            return Err(io::Error::new(ErrorKind::Other, "The iterator has errored").into());
        }
        if !self.is_valid() {
            return Ok(());
        }
        match self.iter.next() {
            Ok(_) => Ok(()),
            Err(e) => {
                self.has_errored = true;
                Err(e)
            }
        }
    }

    fn num_active_iterators(&self) -> usize {
        self.iter.num_active_iterators()
    }
}

We are finally at a point where we can run a scan on the entire data store and get accurate results. So, let’s define the interface and implement it for that.

/// Create an iterator over a range of keys.
    pub fn scan(
        &self,
        lower: Bound<&[u8]>,
        upper: Bound<&[u8]>,
    ) -> Result<FusedIterator<LsmIterator>> {
        //  create the merge iterator for the memtables here
        let state_guard = self.state.read();
        let mut memtables = Vec::new();
        memtables.push(Arc::clone(&state_guard.memtable));
        memtables.extend(
            state_guard
                .imm_memtables
                .iter()
                .map(|memtable| Arc::clone(memtable)),
        );
        let mut memtable_iterators = Vec::new();
        for memtable in memtables {
            //  create a memtable iterator for each memtable
            let iterator = memtable.scan(lower, upper);
            memtable_iterators.push(Box::new(iterator));
        }
        let memtable_merge_iterator = MergeIterator::create(memtable_iterators);
        drop(state_guard);

        //  create the merge iterator for the SsTables here
        let snapshot = {
            let state_guard = self.state.read();
            Arc::clone(&state_guard)
        };

        //  retrieve the ids of the sstables
        let sstable_ids = &*snapshot.l0_sstables;
        let mut sstable_iterators: Vec<Box<SsTableIterator>> = Vec::new();
        for sstable_id in sstable_ids {
            let sstable = snapshot.sstables.get(&sstable_id).unwrap();
            //  need to skip building any iterators that cannot contain the key
            if !self.range_overlap(lower, upper, sstable.first_key(), sstable.last_key()) {
                continue;
            }

            let sstable_iter = match lower {
                Bound::Included(key) => SsTableIterator::create_and_seek_to_key(
                    Arc::clone(sstable),
                    KeySlice::from_slice(key),
                )?,
                Bound::Excluded(key) => {
                    let mut iterator = SsTableIterator::create_and_seek_to_key(
                        Arc::clone(sstable),
                        KeySlice::from_slice(key),
                    )?;
                    if iterator.is_valid() && iterator.key().raw_ref() == key {
                        iterator.next()?;
                    }
                    iterator
                }
                Bound::Unbounded => SsTableIterator::create_and_seek_to_first(Arc::clone(sstable))?,
            };
            sstable_iterators.push(Box::new(sstable_iter));
        }
        let sstable_merge_iterator = MergeIterator::create(sstable_iterators);

        let lsm_iterator = LsmIterator::new(
            TwoMergeIterator::create(memtable_merge_iterator, sstable_merge_iterator)?,
            map_bound(upper),
        )?;

        Ok(FusedIterator::new(lsm_iterator))
    }

And here’s how we’d use the above interface by utilising Rust’s bounds.

lsm.put(b"a", "1");
lsm.put(b"b", "2");
let mut iter = lsm.scan(std::ops::Bound::Unbounded, std::ops::Bound::Unbounded)?;
let mut cnt = 0;
while iter.is_valid() {
    println!(
        "{:?}={:?}",
        Bytes::copy_from_slice(iter.key()),
        Bytes::copy_from_slice(iter.value()),
    );
    iter.next()?;
    cnt += 1;
}

Recap

Congrats on making it this far in the post. To recap, here’s what we’ve covered:

1. Writing user data to an in-memory buffer called the memtable
2. Flushing data from the memtable to a sorted string table
3. Servicing point and range queries via iterators

This was quite a lengthy post and involved a ton of code. If you want to play around with the engine to get a feel for it, clone the repo and run it locally with

cargo run --bin mini-lsm-cli -- --compaction none

You can run commands like

fill 1000 3000 -> put values within that range

get 1001 -> get a specific value

scan 1001 1005 -> scan a specific range

References

1. The repo
2. Three places for data
3. Log structured merge tree
4. What is an LSM tree
5. Bloom filters

Race Conditions In Rust

Here’s some code written in Rust that has a race condition in it because of application logic. It uses the classic example of tranferring money from one account to another but does it in an extremely silly way

struct Account {
    balance: Mutex<i32>,
}

fn transfer2(amount: i32, account_from: Arc<Account>, account_to: Arc<Account>) -> &'static str {
    // First atomic block
    let bal;
    {
        let from_balance = account_from.balance.lock().unwrap();
        bal = *from_balance;
    }

    // Check balance
    if bal < amount {
        return "NOPE";
    }

    // Second atomic block
    {
        let mut to_balance = account_to.balance.lock().unwrap();
        *to_balance += amount;
    }

    // Third atomic block
    {
        let mut from_balance = account_from.balance.lock().unwrap();
        *from_balance -= amount;
    }

    "YEP"
}

fn main() {
    let account_from = Arc::new(Account {
        balance: Mutex::new(100),
    });
    let account_to = Arc::new(Account {
        balance: Mutex::new(50),
    });

    let account_from_clone = Arc::clone(&account_from);
    let account_to_clone = Arc::clone(&account_to);

    let handle = thread::spawn(move || {
        let result = transfer2(30, account_from_clone, account_to_clone);
        println!("Transfer result: {}", result);
    });

    // Simulate another transfer in the main thread
    let result = transfer2(80, Arc::clone(&account_from), Arc::clone(&account_to));
    println!("Transfer result: {}", result);

    handle.join().unwrap();
}

Here’s a simple execution trace that triggers the race condition in the code

T1 -> amount = 30, from_balance = 100, to_balance = 50
T2 -> amount = 80, from_balance = 100, to_balance = 50
Both T1 and T2 pass the check of bal >= amount
T1 -> amount = 30, from_balance = 100, to_balance = 80
T2 -> amount = 80, from_balance = 100, to_balance = 160
T1 -> amount = 30, from_balance = 70, to_balance = 160
T2 -> amount = 80, from_balance = -10, to_balance = 160

The execution results in the bank account being overdrawn. However, the Rust compiler allows this code to be compiled because there is no verifiable way to prevent every possible logical race condition.

It is impossible to prevent all classes of race errors unless you control the scheduler.

Deadlocks in Rust

Now, we’ve seen that Rust’s concurrency patterns can’t prevent application logic race conditions. Rust also cannot prevent your code from deadlocking, which is a specific case of a race condition where your program halts execution. This is also an application logic error.

Assume that you have a program that is attempting to acquire locks in a specific order. We’ll have two locks - A and B and two programs, one of which attempts to acquire A first and then B, and the other program which does the reverse.

Here’s some sample code which demonstrates that

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

struct Locks {
    lock_a: Mutex<u64>,
    lock_b: Mutex<u64>,
}

impl Locks {
    fn acquire(&self) {
        println!("running acquire");
        let mut a = self.lock_a.lock().unwrap();
        println!("got lock a");
        thread::sleep(std::time::Duration::from_millis(1000));
        let mut b = self.lock_b.lock().unwrap();
        println!("got lock b");
        *a += 1;
        *b += 1;
        println!("The sum is {}", *a + *b);
    }

    fn acquire_rev(&self) {
        println!("running acquire rev");
        thread::sleep(std::time::Duration::from_millis(100));
        let mut b = self.lock_b.lock().unwrap();
        println!("got lock b");
        let mut a = self.lock_a.lock().unwrap();
        println!("got lock a");
        *b += 1;
        *a += 1;
        println!("The reverse sum is {}", *b + *a);
    }
}

fn main() {
    let mut handles = vec![];

    let locks = Arc::new(Locks {
        lock_a: Mutex::new(1),
        lock_b: Mutex::new(2),
    });

    for _ in 0..100 {
        let locks = Arc::clone(&locks);
        let handle = thread::spawn(move || {
            locks.acquire();
            locks.acquire_rev();
        });
        handles.push(handle);
    }

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

I’ve introduced a couple of sleeps in the code and added some logs to increase the likelihood of the code deadlocking. Also, this code is extremely silly (for illustrative purposes, obviously) but the point remains that this code compiles perfectly and deadlocks.

It’s impossible for the Rust compiler to catch the issue here because this is an application logic error. As far as the compiler is concerned, everything is wrapped behind an appropriate mutex and the data behind the mutex is being read correctly.

Now, let’s look at a subset of a race condition called a data race - something that the Rust compiler can definitely catch.

Data Races

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

struct Counter {
    counter: Mutex<u64>,
}

impl Counter {
    fn increment_1(&self) {
        let mut counter = self.counter.lock().unwrap();
        *counter += 1;
        println!("The new value of counter in increment_1 {}", counter);
    }

    fn increment_2(&self) {
        let mut counter = self.counter.lock().unwrap();
        *counter += 1;
        println!("The new value of counter in increment_2 {}", counter);
    }
}

fn main() {
    let mut handles = vec![];

    let counter = Arc::new(Counter { counter: 0.into() });

    for _ in 0..100 {
        let counter_1 = Arc::clone(&counter);
        let counter_2 = Arc::clone(&counter);
        let handle_1 = thread::spawn(move || {
            counter_1.increment_1();
        });
        let handle_2 = thread::spawn(move || {
            counter_2.increment_2();
        });
        handles.push(handle_1);
        handles.push(handle_2);
    }

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

There is no way to write this code in Rust such that the memory location represented by counter could potentially be updated by two threads at once. First, I need to wrap the Counter object in an Arc because I need to send it across threads. Arc gives me a thread safe read counter.

However, if I tried to remove the Mutex around counter, I would need to make the methods increment_1 and increment_2 use the signature &mut self. Now, if I were to try to call a mutable method on an object wrapped in an Arc like the code below, I get an error

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

struct Counter {
    counter: u64,
}

impl Counter {
    fn increment_1(&mut self) {
        let mut counter = self.counter;
        counter += 1;
        self.counter = counter;
        println!("The new value of counter in increment_1 {}", counter);
    }

    fn increment_2(&mut self) {
        let mut counter = self.counter;
        counter += 1;
        self.counter = counter;
        println!("The new value of counter in increment_2 {}", counter);
    }
}

fn main() {
    let mut handles = vec![];

    let counter = Arc::new(Counter { counter: 0 });

    for _ in 0..100 {
        let counter_1 = Arc::clone(&counter);
        let counter_2 = Arc::clone(&counter);
        let handle_1 = thread::spawn(move || {
            counter_1.increment_1();
        });
        let handle_2 = thread::spawn(move || {
            counter_2.increment_2();
        });
        handles.push(handle_1);
        handles.push(handle_2);
    }

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

The error states that

error[E0596]: cannot borrow data in an `Arc` as mutable
  --> src/main.rs:33:13
   |
33 |             counter_1.increment_1();
   |             ^^^^^^^^^ cannot borrow as mutable
   |
   = help: trait `DerefMut` is required to modify through a dereference, but it is not implemented for `Arc<Counter>`

So, if I want to share some state across threads, I am forced to wrap that data in a Mutex, preventing any kind of data race.

There is no way to write this code with a race condition which the Rust compiler will allow.

Let’s look at a simpler case of how Rust prevents data races by ignoring threads.

One Mutable Reference Or Multiple Immutable References

If you’ve read the Rust Book, you’ve probably heard about how you can have either one mutable reference to an object or multiple immutable references. This is similar to the idea of a RWLock, where you can have multiple readers or a single writer at any given point of time.

Here’s a small example in Rust showing how the compiler prevents this class of errors

struct Data {
    var_a: u64,
    var_b: u64,
}

fn main() {
    let mut data = Data { var_a: 1, var_b: 2 };
    let a = &data.var_a;
    data.var_a += 1;
    println!("var a {}", data.var_a);
    println!("a {}", a);
}

If you try to compile the above code, you see the following error

error[E0506]: cannot assign to `data.var_a` because it is borrowed
  --> src/main.rs:61:5
   |
60 |     let a = &data.var_a;
   |             ----------- `data.var_a` is borrowed here
61 |     data.var_a += 1;
   |     ^^^^^^^^^^^^^^^ `data.var_a` is assigned to here but it was already borrowed
62 |     println!("var a {}", data.var_a);
63 |     println!("a {}", a);
   |                      - borrow later used here

The error is essentially saying that because the variable a borrows data.var_a, data.var_a cannot later be changed later since this invalidates the reference that a holds. This is preventing a data race because the memory location that a points to cannot be changed if a is going to be used.

Here’s a different version of the above code with methods implemented on the Data struct but demonstrating the same principle regarding data race prevention.

struct Data {
    var_a: u64,
    var_b: u64,
}

impl Data {
    fn get_a(&self) -> &u64 {
        &self.var_a
    }

    fn increment_a(&mut self) {
        self.var_a += 1;
    }
}

fn main() {
    let mut data = Data { var_a: 1, var_b: 2 };
    let a_ref = data.get_a();
    data.increment_a();
    println!("The ref {}", a_ref);
}

This code gives the following error

error[E0502]: cannot borrow `data` as mutable because it is also borrowed as immutable
   --> src/main.rs:109:5
    |
108 |     let a_ref = data.get_a();
    |                 ---- immutable borrow occurs here
109 |     data.increment_a();
    |     ^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
110 |     println!("The ref {}", a_ref);
    |                            ----- immutable borrow later used here

The basic underlying principle behind Rust preventing data races is quite simple: you can either have a single mutable reference to some location in memory and mutate that data in any way you wish, or you can hold multiple immutable references to some location in memory and read from it via as many threads as you’d like.

References

1. Race Condition vs Data Race
2. Rustnomicon on data races and race conditions

In this post, I’m going to focus more on how to go about testing the cluster. While building the Raft implementation, I realised that half the battle with distributed systems is in building a useful test harness. This has become such a problem in the distributed systems space that companies like FoundationDB and TigerBeetle have written something called a deterministic simulation testing (DST) engine. It’s a fancy term for a more advanced form of fuzz testing (atleast from what I understand).

The founder of FoundationDB gave a talk about why they went about building a DST simulator and he later went on to found Antithesis whose entire business is around trying to provide a generalizable DST engine to other companies to test their products!

Anyway, back to writing a much simpler test cluster in Raft!

Mocking

While I was trying to write my Raft implementation, there was a very helpful Twitter reply about mocking away the network and the clock so that you’re essentially testing a deterministic state machine (Raft itself, however, isn’t deterministic).

In the previous post we saw an RPCManager to facilitate communication and a tick method to logically advance the time in the cluster. These methods are defined for each node.

This tick function which I’ve reproduced below calls an advance_time_by method, which in turn calls clock.advance on a node. Each node has a clock which is injected into it.

struct RaftNode<T: RaftTypeTrait, S: StateMachine<T>, F: RaftFileOps<T>> {
    id: ServerId,
    state: Mutex<RaftNodeState<T>>,
    state_machine: S,
    config: ServerConfig,
    peers: Vec<ServerId>,
    to_node_sender: mpsc::Sender<MessageWrapper<T>>,
    from_rpc_receiver: mpsc::Receiver<MessageWrapper<T>>,
    rpc_manager: CommunicationLayer<T>,
    persistence_manager: F,
    clock: RaftNodeClock,
}

impl RaftNodeClock {
    fn advance(&mut self, duration: Duration) {
        match self {
            RaftNodeClock::RealClock(_) => (), //  this method should do nothing for a real clock
            RaftNodeClock::MockClock(clock) => clock.advance(duration),
        }
    }
}

fn advance_time_by(&mut self, duration: Duration) {
    self.clock.advance(duration);
}

fn tick(&mut self) {
    let mut state_guard = self.state.lock().unwrap();
    match state_guard.status {
        RaftNodeStatus::Leader => {
            //  the leader should send heartbeats to all followers
            self.send_heartbeat(&state_guard);
        }
        RaftNodeStatus::Candidate | RaftNodeStatus::Follower => {
            //  the candidate or follower should check if the election timeout has elapsed
            self.check_election_timeout(&mut state_guard);
        }
    }
    drop(state_guard);

    //  listen to incoming messages from the RPC manager
    self.listen_for_messages();
    self.advance_time_by(Duration::from_millis(1));
}

Clocks

I defined two variants of a clock - a mock and a real implementation to use for my cluster (another area I found where Rust’s enum pattern matching really shines).

use std::time::{Duration, Instant};

pub trait Clock {
    fn now(&self) -> Instant;
}

pub struct RealClock;
impl Clock for RealClock {
    fn now(&self) -> Instant {
        Instant::now()
    }
}

pub struct MockClock {
    pub current_time: Instant,
}

impl MockClock {
    pub fn advance(&mut self, duration: Duration) {
        self.current_time += duration;
    }
}

impl Clock for MockClock {
    fn now(&self) -> Instant {
        self.current_time
    }
}

So any time I was creating a cluster for testing, I would create nodes with a MockClock and for actual implementations (I never created those) I would use the RealClock. You’ll notice that the RealClock cannot have its time advanced.

Network

Similar to the clocks, for the network layer I created a MockRPCManager which contained a message queue indicating what messages were sent by the node for every logical tick of the cluster. And when a node wanted to send a message, it would simply get pushed into the sent_messages queue.

There are also two additional helper methods called get_messages_in_queue and replay_messages_in_queue which allow me to either drain the messages or inspect them without mutating them. The get_messages method drains the messages from a node’s queue and sends them across the network and replay_messages keeps the messages as they are but allows for inspection.

Note: I picked up these ideas from an implementation on GitHub which you can find here

#[derive(Debug, Clone)]
struct MessageWrapper<T: RaftTypeTrait> {
    from_node_id: ServerId,
    to_node_id: ServerId,
    message: RPCMessage<T>,
}

struct MockRPCManager<T: RaftTypeTrait> {
    server_id: ServerId,
    to_node_sender: mpsc::Sender<MessageWrapper<T>>,
    sent_messages: RefCell<Vec<MessageWrapper<T>>>,
}

impl<T: RaftTypeTrait> MockRPCManager<T> {
    fn new(server_id: ServerId, to_node_sender: mpsc::Sender<MessageWrapper<T>>) -> Self {
        MockRPCManager {
            server_id,
            to_node_sender,
            sent_messages: RefCell::new(vec![]),
        }
    }
}

impl<T: RaftTypeTrait> MockRPCManager<T> {
    fn start(&self) {}

    fn stop(&self) {}

    fn send_message(&self, _to_address: String, message: MessageWrapper<T>) {
        self.sent_messages.borrow_mut().push(message);
    }

    fn get_messages_in_queue(&mut self) -> Vec<MessageWrapper<T>> {
        let mut mock_messages_vector: Vec<MessageWrapper<T>> = Vec::new();
        for message in self.sent_messages.borrow_mut().drain(..) {
            mock_messages_vector.push(message.clone());
        }
        mock_messages_vector
    }

    fn replay_messages_in_queue(&self) -> Ref<Vec<MessageWrapper<T>>> {
        self.sent_messages.borrow()
    }
}

The CommunicationLayer trait defines what methods are available to a node to consume via its network layer.

trait Communication<T: RaftTypeTrait> {
    fn start(&self);

    fn stop(&self);

    fn send_message(&self, to_address: String, message: MessageWrapper<T>);
}

enum CommunicationLayer<T: RaftTypeTrait> {
    MockRPCManager(MockRPCManager<T>),
    RPCManager(RPCManager<T>),
}

impl<T: RaftTypeTrait> Communication<T> for CommunicationLayer<T> {
    fn start(&self) {
        match self {
            CommunicationLayer::MockRPCManager(manager) => manager.start(),
            CommunicationLayer::RPCManager(manager) => manager.start(),
        }
    }

    fn stop(&self) {
        match self {
            CommunicationLayer::MockRPCManager(manager) => manager.stop(),
            CommunicationLayer::RPCManager(manager) => manager.stop(),
        }
    }

    fn send_message(&self, to_address: String, message: MessageWrapper<T>) {
        match self {
            CommunicationLayer::MockRPCManager(manager) => {
                manager.send_message(to_address, message)
            }
            CommunicationLayer::RPCManager(manager) => manager.send_message(to_address, message),
        }
    }
}

impl<T: RaftTypeTrait> CommunicationLayer<T> {
    fn get_messages(&mut self) -> Vec<MessageWrapper<T>> {
        match self {
            CommunicationLayer::MockRPCManager(manager) => {
                return manager.get_messages_in_queue();
            }
            CommunicationLayer::RPCManager(_) => {
                panic!("This method is not supported for the RPCManager");
            }
        }
    }

    fn replay_messages(&self) -> Ref<Vec<MessageWrapper<T>>> {
        match self {
            CommunicationLayer::MockRPCManager(manager) => {
                return manager.replay_messages_in_queue();
            }
            CommunicationLayer::RPCManager(_) => {
                panic!("This method is not supported for the RPCManager");
            }
        }
    }
}

Test Cluster

Okay, now that we’ve got the mocking out of the way, let’s jump into defining the test cluster itself.

Here’s some basic definitions for my TestCluster.

#[derive(Clone)]
pub struct ClusterConfig {
    pub election_timeout: Duration,
    pub heartbeat_interval: Duration,
    pub ports: Vec<u64>,
}

pub struct TestCluster {
    pub nodes: Vec<RaftNode<i32, KeyValueStore<i32>, DirectFileOpsWriter>>,
    pub nodes_map:
        BTreeMap<ServerId, RaftNode<i32, KeyValueStore<i32>, DirectFileOpsWriter>>,
    pub message_queue: Vec<MessageWrapper<i32>>,
    pub connectivity: HashMap<ServerId, HashSet<ServerId>>,
    pub config: ClusterConfig,
}

I have a bunch of helper methods on this TestCluster but it would be too much to go into all of them so I’ll just go over the tick method and the time based methods here. If you want to see all the methods, check out the repo here.

The tick method on the cluster calls the get_messages method we saw earlier for each node. It collects all these messages in a central queue and then dispatches them in the same logical tick. But, it dispatches the messages only after each node has gone through its own tick function. There’s no special reason for the ordering of events here, it’s just how I decided to do it.

impl TestCluster {
pub fn tick(&mut self) {
    //  Collect all messages from nodes and store them in the central queue
    self.nodes.iter_mut().for_each(|node| {
        let mut messages_from_node = node.rpc_manager.get_messages();
        self.message_queue.append(&mut messages_from_node);
    });

    //  allow each node to tick
    self.nodes.iter_mut().for_each(|node| {
        node.tick();
    });

    //  deliver all messages from the central queue
    self.message_queue
        .drain(..)
        .into_iter()
        .for_each(|message| {
            let node = self
                .nodes
                .iter()
                .find(|node| node.id == message.to_node_id)
                .unwrap();
            //  check if these pair of nodes are partitioned
            if self
                .connectivity
                .get_mut(&message.from_node_id)
                .unwrap()
                .contains(&message.to_node_id)
            {
                match node.to_node_sender.send(message) {
                    Ok(_) => (),
                    Err(e) => {
                        panic!(
                            "There was an error while sending the message to the node: {}",
                            e
                        )
                    }
                };
            }
        });
}

The methods below allow me to advance logical time in my cluster for a bunch of different configurations. Either all nodes can advance by the same amount, a single node can advance by some amount, or all nodes can advance by some variable amount.

pub fn tick_by(&mut self, tick_interval: u64) {
    for _ in 0..tick_interval {
        self.tick();
    }
}

pub fn advance_time_by_for_node(&mut self, node_id: ServerId, duration: Duration) {
    let node = self
        .nodes
        .iter_mut()
        .find(|node| node.id == node_id)
        .expect(&format!("Could not find node with id: {}", node_id));
    node.advance_time_by(duration);
}

pub fn advance_time_by_variably(&mut self, duration: Duration) {
    //  for each node in the cluster, advance it's mock clock by the duration + some random variation
    for node in &mut self.nodes {
        let jitter = rand::thread_rng().gen_range(0..50);
        let new_duration = duration + Duration::from_millis(jitter);
        node.advance_time_by(new_duration);
    }
}

pub fn advance_time_by(&mut self, duration: Duration) {
    //  for each node in the cluster, advance it's mock clock by the duration
    for node in &mut self.nodes {
        node.advance_time_by(duration);
    }
}

There are two more methods I found super helpful while doing the testing which are below - the partition and heal_partition methods. These allow me to simulate a network partition by stopping the flow of messages from one subset of nodes to another and later fixing that.

/// Separates the cluster into two smaller clusters where only the nodes within
/// each cluster can communicate between themselves
pub fn partition(&mut self, group1: &[ServerId], group2: &[ServerId]) {
    for &node_id in group1 {
        self.connectivity
            .get_mut(&node_id)
            .unwrap()
            .retain(|&id| group1.contains(&id));
    }
    for &node_id in group2 {
        self.connectivity
            .get_mut(&node_id)
            .unwrap()
            .retain(|&id| group2.contains(&id));
    }
}

/// Removes any partition in the cluster and restores full connectivity among all nodes
pub fn heal_partition(&mut self) {
    let all_node_ids = self
        .nodes
        .iter()
        .map(|node| node.id)
        .collect::<HashSet<ServerId>>();
    for node_id in all_node_ids.iter() {
        self.connectivity.insert(*node_id, all_node_ids.clone());
    }
}

And here’s what the code for creating, starting & stopping a cluster looks like.

pub fn new(number_of_nodes: u64, config: ClusterConfig) -> Self {
    let mut nodes: Vec<RaftNode<i32, KeyValueStore<i32>, DirectFileOpsWriter>> =
        Vec::new();
    let nodes_map: BTreeMap<
        ServerId,
        RaftNode<i32, KeyValueStore<i32>, DirectFileOpsWriter>,
    > = BTreeMap::new();

    let node_ids: Vec<ServerId> = (0..number_of_nodes).collect();
    let addresses: Vec<String> = config
        .ports
        .iter()
        .map(|port| format!("127.0.0.1:{}", port))
        .collect();
    let mut id_to_address_mapping: HashMap<ServerId, String> = HashMap::new();
    for (node_id, address) in node_ids.iter().zip(addresses.iter()) {
        id_to_address_mapping.insert(*node_id, address.clone());
    }

    let mut counter = 0;
    for (node_id, address) in node_ids.iter().zip(addresses.iter()) {
        let server_config = ServerConfig {
            election_timeout: config.election_timeout,
            heartbeat_interval: config.heartbeat_interval,
            address: address.clone(),
            port: config.ports[counter],
            cluster_nodes: node_ids.clone(),
            id_to_address_mapping: id_to_address_mapping.clone(),
        };

        let state_machine = KeyValueStore::<i32>::new();
        let persistence_manager = DirectFileOpsWriter::new("data", *node_id).unwrap();
        let (to_node_sender, from_rpc_receiver) =
            mpsc::channel::<MessageWrapper<i32>>();
        let rpc_manager = CommunicationLayer::MockRPCManager(MockRPCManager::new(
            *node_id,
            to_node_sender.clone(),
        ));
        let mock_clock = RaftNodeClock::MockClock(MockClock {
            current_time: Instant::now(),
        });

        let node = RaftNode::new(
            *node_id,
            state_machine,
            server_config,
            node_ids.clone(),
            persistence_manager,
            rpc_manager,
            to_node_sender,
            from_rpc_receiver,
            mock_clock,
        );
        nodes.push(node);
        counter += 1;
    }
    let message_queue = Vec::new();

    let mut connectivity_hm: HashMap<ServerId, HashSet<ServerId>> = HashMap::new();
    for node_id in &node_ids {
        connectivity_hm.insert(*node_id, node_ids.clone().into_iter().collect());
    }

    TestCluster {
        nodes,
        nodes_map,
        message_queue,
        connectivity: connectivity_hm,
        config,
    }
}

pub fn start(&mut self) {
    for node in &mut self.nodes {
        node.start();
    }
}

pub fn stop(&self) {
    for node in &self.nodes {
        node.stop();
    }
}

Actual Tests

Now that we have a reasonable test harness set up, I’m going to dive into some specific test scenarios to explain how I use the test harness to move the cluster into a specific configuration and then test out certain scenarios.

Note: After spending time doing this, I understand why people prefer property based testing or fuzz testing as a methodology. Creating scenario based tests is a very time consuming process and doesn’t give you enough coverage to justify the time spent on it.

My tests aren’t exhaustive (and might even be buggy) and I genuinely don’t think scenario based testing is the way to get exhaustive coverage anyway. I have more tests in the repo if you want to look at those.

We’ll start with some simple tests for a single node cluster.

const ELECTION_TIMEOUT: Duration = Duration::from_millis(150);
const HEARTBEAT_INTERVAL: Duration = Duration::from_millis(50);
const MAX_TICKS: u64 = 100;

/// This test checks whether a node in a single cluster will become leader as soon as the election timeout is reached
#[test]
fn leader_election() {
    let _ = env_logger::builder().is_test(true).try_init();
    //  create a cluster with a single node first
    let cluster_config = ClusterConfig {
        election_timeout: ELECTION_TIMEOUT,
        heartbeat_interval: HEARTBEAT_INTERVAL,
        ports: vec![8000],
    };
    let mut cluster = TestCluster::new(1, cluster_config);
    cluster.start();
    cluster.advance_time_by(ELECTION_TIMEOUT + Duration::from_millis(100 + 5)); //  picking 255 here because 150 + a max jitter of 100 guarantees that election has timed out
    cluster.wait_for_stable_leader(MAX_TICKS);
    cluster.stop();
    assert_eq!(cluster.has_leader(), true);
}

You can see that setting up the cluster earlier makes testing significantly easier here. I just call my helper methods on the cluster to move the cluster into a specific state and then assert whatever conditions I want against the cluster.

As an aside, you could convert this to a property based test using the proptest crate. You could change the number of nodes in the cluster to determine that a leader gets elected regardless of the number of nodes in the cluster. Something like the below code

extern crate proptest;
use proptest::prelude::*;
use std::time::Duration;

proptest! {
    #![proptest_config(ProptestConfig::with_cases(10))]

    fn leader_election_property_based_test(node_count in 1usize..5) {
        let _ = env_logger::builder().is_test(true).try_init();
        
        let cluster_config = ClusterConfig {
            election_timeout: ELECTION_TIMEOUT,
            heartbeat_interval: HEARTBEAT_INTERVAL,
            ports: (8000..8000 + node_count as u16).collect(),
        };
        let mut cluster = TestCluster::new(node_count, cluster_config);
        cluster.start();
        
        cluster.advance_time_by(ELECTION_TIMEOUT + Duration::from_millis(100 + 5));
        cluster.wait_for_stable_leader(MAX_TICKS * node_count as u64); // Adjusted wait time
        cluster.stop();
        
        assert_eq!(cluster.has_leader(), true);
    }
}

Here’s a slightly more complicated test with 3 nodes that simulates a network partition in the cluster and asserts that a leader still gets elected in the majority partitioned cluster.

/// This test models the scenario where the leader in a cluster is network partitioned from the
/// rest of the cluster and a new leader is elected
#[test]
fn network_partition_new_leader() {
    let _ = env_logger::builder().is_test(true).try_init();
    let cluster_config = ClusterConfig {
        election_timeout: ELECTION_TIMEOUT,
        heartbeat_interval: HEARTBEAT_INTERVAL,
        ports: vec![8000, 8001, 8002],
    };
    let mut cluster = TestCluster::new(3, cluster_config);
    cluster.start();
    cluster.advance_time_by_for_node(0, ELECTION_TIMEOUT + Duration::from_millis(50));
    cluster.wait_for_stable_leader(MAX_TICKS);

    //  partition the leader from the rest of the group
    let group1 = &[cluster.get_leader().unwrap().id];
    let group2 = &cluster
        .get_all_followers()
        .iter()
        .map(|node| node.id)
        .collect::<Vec<ServerId>>();
    cluster.partition(group1, group2);
    cluster.advance_time_by_for_node(1, ELECTION_TIMEOUT + Duration::from_millis(100));
    cluster.wait_for_stable_leader_partition(MAX_TICKS, group2);
    cluster.stop();
    assert_eq!(cluster.has_leader_in_partition(group2), true);
}

Again, you can see how useful the helper methods turn out to be because it’s much easier to separate one node from the rest by simulating a network partition and then checking certain properties on the individual partitions.

Now, here’s a much more complicated scenario-based test which models a network partition occurring between the leader and the rest of the cluster, the network partition healing and the old leader rejoining the rest of the cluster and having to catch up.

/// The test models the following scenario
/// 1. A leader is elected
/// 2. Client requests are received and replicated
/// 3. A partition occurs and a new leader is elected
/// 4. Clients requests are processed by the new leader
/// 5. The partition heals and the old leader rejoins the cluster
/// 6. The old leader must recognize its a follower and get caught up with the new leader
#[test]
fn network_partition_log_healing() {
    let _ = env_logger::builder().is_test(true).try_init();
    let cluster_config = ClusterConfig {
        election_timeout: ELECTION_TIMEOUT,
        heartbeat_interval: HEARTBEAT_INTERVAL,
        ports: vec![8000, 8001, 8002],
    };
    let mut cluster = TestCluster::new(3, cluster_config);

    //  cluster starts and a leader is elected
    cluster.start();
    cluster.advance_time_by_for_node(0, ELECTION_TIMEOUT + Duration::from_millis(50));
    cluster.wait_for_stable_leader(MAX_TICKS);
    assert_eq!(cluster.has_leader(), true);

    //  client requests are received and replicated across cluster
    let current_leader_term = cluster
        .get_leader()
        .unwrap()
        .state
        .lock()
        .unwrap()
        .current_term;
    let mut last_log_index = cluster
        .get_leader()
        .unwrap()
        .state
        .lock()
        .unwrap()
        .log
        .last()
        .map_or(0, |entry| entry.index);
    last_log_index += 1;
    let log_entry_1 = LogEntry {
        term: current_leader_term,
        index: last_log_index,
        command: LogEntryCommand::Set,
        key: "a".to_string(),
        value: 1,
    };
    last_log_index += 1;
    let log_entry_2 = LogEntry {
        term: current_leader_term,
        index: last_log_index,
        command: LogEntryCommand::Set,
        key: "b".to_string(),
        value: 2,
    };
    cluster.apply_entries_across_cluster(vec![&log_entry_1, &log_entry_2], MAX_TICKS);
    cluster.verify_logs_across_cluster_for(vec![&log_entry_1, &log_entry_2], MAX_TICKS);

    //  partition and new leader election here
    let group1 = &[cluster.get_leader().unwrap().id];
    let group2 = &cluster
        .get_all_followers()
        .iter()
        .map(|node| node.id)
        .collect::<Vec<ServerId>>();
    cluster.partition(group1, group2);
    cluster.advance_time_by_for_node(1, ELECTION_TIMEOUT + Duration::from_millis(100));
    cluster.wait_for_stable_leader_partition(MAX_TICKS, group2);
    assert_eq!(cluster.has_leader_in_partition(group2), true);

    //  send requests to group2 leader
    let log_entry_3 = {
        let current_leader_term = cluster
            .get_leader_in_cluster(group2)
            .unwrap()
            .state
            .lock()
            .unwrap()
            .current_term;
        let mut last_log_index = cluster
            .get_leader()
            .unwrap()
            .state
            .lock()
            .unwrap()
            .log
            .last()
            .map_or(0, |entry| entry.index);
        last_log_index += 1;
        let log_entry_3 = LogEntry {
            term: current_leader_term,
            index: last_log_index,
            command: LogEntryCommand::Set,
            key: "c".to_string(),
            value: 3,
        };
        cluster.apply_entries_across_cluster_partition(
            vec![&log_entry_3],
            group2,
            MAX_TICKS,
        );
        log_entry_3
    };

    //  partition heals and old leader rejoins the cluster
    //  cluster needs to be verified to ensure all logs are up to date
    let leader_id = cluster.get_leader_in_cluster(group2).unwrap().id;
    cluster.heal_partition();
    cluster.tick_by(MAX_TICKS);
    cluster.wait_for_stable_leader(MAX_TICKS);
    assert_eq!(cluster.get_leader().unwrap().id, leader_id);
    cluster.verify_logs_across_cluster_for(
        vec![&log_entry_1, &log_entry_2, &log_entry_3],
        MAX_TICKS,
    );
}

I think this is one of those situations where property-based testing would really shine not because it would make setting up the test simpler but because it would allow me to use the same test case and test the cluster under different scenarios.

For instance, I could use proptest to generate clusters of different sizes, generate different variations of partitions and add a different number of log entries based on some generated value. This would give me more coverage.

So, in a sense, you could almost think of a single property test combining multiple scenario based tests.

A Little Bit About DST

I mentioned earlier that I’ve been digging into DST a little bit and with the help of some folks online, I’ve come to understand that DST is essentially about using a random seed value to generate property based tests while maintaining the ability to re-create a specific scenario using nothing but the seed value.

If we take the last test case we were discussing about generating clusters of different sizes, different variations of partitions and adding a different number of log entries, imagine if we did that with a randomly generated number which we can call a seed value.

This seed value would then be used to generate all the other values in a deterministic way. For instance, the number of nodes in the cluster could be the seed value itself, the number of log entries written could be seed_value * 2 and the the partition configuration could be defined by (seed_value % num_of_nodes) - 1. This way, if the test case fails and the seed value is logged, I can re-create the test scenario using nothing but the seed value itself.

Since the seed value changes on every test run, I can test the system under a variety of different scenarios using property based testing. And in the case of a failure in a specific scenario, I am able to immediately re-create the scenario to debug the error.

I was under the impression that DST was a magic pill in the sense that if you wrote code that passed the simulator, you have bug free code. This is obviously incorrect. You’ve only reduced the probability of bugs in production by a significant factor (hopefully). In the video I linked above, Will Wilson explains how they considered writing a second DST simulator at FoundationDB because the engineers were getting really good at writing code that passed the simulator but still had bugs in it.

Conclusion

So there you have it. Writing a test harness for a distributed system is significantly more challenging than writing the system itself I think. I can see now why FoundationDB spent years building their simulator before writing their engine, although I don’t know if that’s a realistic approach for most people.

Still, it’s fun to dig into this stuff.