Rust Singletons in a ServiceWorker

August 16, 2024 —

10 Minute Read

Background

Recently I was working on the archival editor (https://editor.archival.dev), which is an application in which all the business logic happens in a Rust library. This setup is nice because the UI is a pure function of the state produced by Rust, which is infinitely more reliable than javascript and much faster than round trips to a server.

One of the things that causes a bit of friction when working in this environment is that the singleton pattern isn’t really friendly with Rust’s memory model.

Rust’s design includes the concept of “Fearless Concurrency”, which in practice means that unless any type is marked as Send, the compiler will disallow storing references to it in structures that could potentially be used across threads. For the most part, when working in Rust you'll use types from the stdlib (e.g. Mutex or RwLock) which implement Send and Sync - usually, you won't need to implement these traits yourself. And importantly, implementing these traits must be done outside of the borrow checker (in an unsafe block) because the checker can't guarantee ownership across threads.

This sometimes bumps up against the reality that a service worker is always a single thread. You could go to all the work to make sure every one of your data structures are thread safe, but this would ratchet up the already quite high friction in working with a WASM codebase up to 11 (and in my opinion push it beyond what is reasonable for a web application in terms of complexity). In addition, multithreading doesn't work out of the box in a WASM environment - subthreads will silently not spawn which will often result in an effective deadlock. It makes a lot more sense to build some tooling around helping rust with the single threaded assumption, and then handling any necessary parallel execution via spawning workers.

Why Use Rust At All?

In modern frontend application architecture, a common goal is to make the frontend entirely stateless. This is a strategy long explored in the frontend community, and a big reason why React and other reactive patterns are so popular. However, when working in real javascript applications, the abstraction always leaks. React and other libraries explicitly support storing state in components, and the lines often blur when using patterns like context. Ideally, the only memory that needs to be preserved should be entirely outside of the rendering code.

In my app, this is achieved via using a long-lived service worker to hold all stateful memory - even the javascript code running in this service worker is stateless outside of holding a reference to an instance of my rust library. When it receives a message from the frontend or when the rust lib calls a FFI in the worker, it sends those along to the correct receiver. (It also manages caching & proxying but that’s for a different post).

The Rust application on the other hand becomes extremely stateful. It manages worker lifecycle by caching any long-lived application data (using web APIs and/or FFIs to invoke native Cache APIs).

This gives you the safety of Rust where it matters - in the process of producing a State object - but doesn’t require you to do anything fancy when it comes to rendering the UI. In fact, I just use DOM apis via upd8 and skip a reactive javascript library entirely.

Managing State With a Singleton

Ok, so we want to put our state in Rust. But Rust doesn’t have any stdlib singleton pattern, so we’ll need to use some static memory to track our main library code.

My first attempt was to just write a huge comment warning the reader and to bail out memory safety via unsafe impl Sync/Send on a newType struct:

// HERE BE DRAGONS

/*
 * Editor is "effectively" single threaded, in that it makes heavy use of
 * interior mutability and locking to deal with event orchestration. These
 * locks are generally achieved through RefCells, which are not threadsafe
 * but faster than mutexes. Similarly, ref counting is achieved via Rc
 * instead of Arc for performance. While it may be possible to make Editor
 * threadsafe, the use of stored callbacks make it an even larger challenge,
 * and due to the statefulness of it, there are little upsides. Instead,
 * it is the implementors responsibility to make sure that there is only one
 * Editor instance. In FFI contexts, the implementor can just hold a pointer.
 * But in WASM, it is non-trivial and already unsafe to send a pointer in and
 * out of javascript. So instead, we use a singleton in the WASM thread, which
 * is (as of 2023) already guaranteed to be single threaded.
 */

use crate::editor::Editor;
use once_cell::sync::OnceCell;
use std::ops::{Deref, DerefMut};
use std::sync::{Mutex, MutexGuard};
use tracing::{debug, trace_span, warn};
use wasm_bindgen::JsValue;

pub struct NotActuallyThreadSafeEditor(Editor);
unsafe impl Sync for NotActuallyThreadSafeEditor {}
unsafe impl Send for NotActuallyThreadSafeEditor {}
impl<'a> Deref for NotActuallyThreadSafeEditor {
    type Target = Editor;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
impl DerefMut for NotActuallyThreadSafeEditor {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}
fn wrapped_editor<'a>() -> &'static Mutex<NotActuallyThreadSafeEditor> {
    static INSTANCE: OnceCell<Mutex<NotActuallyThreadSafeEditor>> = OnceCell::new();
    INSTANCE.get_or_init(|| {
        debug!("♻️ Create Editor Singleton");
        Mutex::new(NotActuallyThreadSafeEditor(Editor::init()))
    })
}

pub static LOCK_DESC: Mutex<String> = Mutex::new(String::new());

pub fn editor_guard(
    name: &str,
) -> Result<MutexGuard<'static, NotActuallyThreadSafeEditor>, JsValue> {
    if let Ok(editor) = wrapped_editor().try_lock() {
        debug!("locking: {}", name);
        let s = trace_span!("locking", operation = name);
        let _lt = s.enter();
        *LOCK_DESC.lock().expect(&format!("double lock: {}", name)) = String::from(name);
        Ok(editor)
    } else {
        {
            let current_lock = LOCK_DESC.lock().unwrap();
            warn!(
                "conflict: attempted to lock ({}) while locked ({})",
                name, current_lock
            );
        }
        Err(JsValue::from_str("locked"))
    }
}

Mutexes, RwLock, Panics - oh my!

This all brings me to the main issue - I generate bindings from Rust to JS using wasm-bindgen - this saves me a lot of headache around instantiation and types, and works out of the box. However something that wasn’t immediately clear to me when getting started with this toolkit is that while Rust applications usually terminate when they panic, this environment has no “crash mode” - exceptions can be thrown, but that just ends the current tick and prints something to the console. The next tick, the event loop will continue, with your memory intact. This totally breaks a core assumption in Rust that memory will be dropped after a panic, and often results in unreachable statements being hit as Rust code is not expected to continue running after a panic.

What does this have to do with singletons? Because, when the above code panics, the locks are never poisoned. Instead, wasm-bindgen wraps panics and throws exceptions. In addition, the guard isn’t dropped - which makes sense, JS can’t guarantee that we won’t reuse that memory. This results in a permanently locked Mutex which isn't marked as poisoned. The result for end users is an infinite load, or otherwise non-modifiable state, since the Mutex will never unlock.

One Solution

I initially looked at “manual poisoning”, but this reasonably doesn't exist, the implementation details of poisoning and unlocking are private. So instead I wrote my own manually tracked lock, which would allow me to poison it in a panic hook. Then, I used a static variable to track a global instance that is lazily initialized. When the lock is poisoned, we just recreate the instance the next time it is asked for. The result looks roughly like this:


use crate::editor::Editor;
use std::fmt::Debug;
use std::ops::{Deref, DerefMut};
use tracing::{debug, warn};
use wasm_bindgen::JsValue;

pub static mut INSTANCE: Option<UnsafeAssumedSingleThreadLock<Editor>> = None;

unsafe fn guarantee_instance() {
    if INSTANCE.is_none() {
        debug!("♻️ Create Editor Singleton");
        INSTANCE = Some(UnsafeAssumedSingleThreadLock::new(Editor::init()));
    }
}

pub fn set_panicked() {
    unsafe {
        if let Some(inst) = INSTANCE.as_mut() {
            inst.poison();
        }
    }
}

pub fn editor_guard(name: &str) -> Result<UASTGuard<Editor>, JsValue> {
    unsafe {
        guarantee_instance();
        match INSTANCE.as_mut().unwrap().lock(name) {
            UnsafeAssumedSingleThreadResult::Locked(existing) => Err(format!(
                "lock conflict: Attempted to acquire editor lock for {} while it was already locked for {}",
                name, existing
            )
            .into()),
            UnsafeAssumedSingleThreadResult::Poisoned => {
                INSTANCE = None;
                editor_guard(name)
            }
            UnsafeAssumedSingleThreadResult::Unlocked(guard) => Ok(guard),
        }
    }
}

pub fn read_editor() -> &'static Editor {
    unsafe {
        guarantee_instance();
        INSTANCE.as_ref().unwrap()
    }
}

#[derive(Clone, Debug)]
enum UnsafeAssumedSingleThreadLockStatus {
    Unlocked,
    Locked(String),
    Poisoned,
}

#[derive(Clone, Debug)]
struct UnsafeAssumedSingleThreadLock<T> {
    status: UnsafeAssumedSingleThreadLockStatus,
    value: T,
}

enum UnsafeAssumedSingleThreadResult<'a, T> {
    Unlocked(UASTGuard<'a, T>),
    Locked(&'a str),
    Poisoned,
}

pub struct UASTGuard<'a, T> {
    lock: &'a mut UnsafeAssumedSingleThreadLock<T>,
}
impl<'a, T> UASTGuard<'a, T> {
    fn new(lock: &'a mut UnsafeAssumedSingleThreadLock<T>) -> Self {
        Self { lock }
    }
}
impl<'a, T> Drop for UASTGuard<'a, T> {
    fn drop(&mut self) {
        self.lock.status = UnsafeAssumedSingleThreadLockStatus::Unlocked;
    }
}
impl<'a, T> Deref for UASTGuard<'a, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        &self.lock
    }
}
impl<'a, T> DerefMut for UASTGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.lock.value
    }
}

impl<T> UnsafeAssumedSingleThreadLock<T> {
    pub fn new(thing: T) -> Self {
        Self {
            status: UnsafeAssumedSingleThreadLockStatus::Unlocked,
            value: thing,
        }
    }
    pub fn lock(&mut self, description: &str) -> UnsafeAssumedSingleThreadResult<T> {
        match &self.status.clone() {
            UnsafeAssumedSingleThreadLockStatus::Locked(_) => {
                if let UnsafeAssumedSingleThreadLockStatus::Locked(name) = &self.status {
                    UnsafeAssumedSingleThreadResult::Locked(name)
                } else {
                    panic!("invariant")
                }
            }
            UnsafeAssumedSingleThreadLockStatus::Poisoned => {
                UnsafeAssumedSingleThreadResult::Poisoned
            }
            UnsafeAssumedSingleThreadLockStatus::Unlocked => {
                self.status = UnsafeAssumedSingleThreadLockStatus::Locked(description.to_string());
                UnsafeAssumedSingleThreadResult::Unlocked(UASTGuard::new(self))
            }
        }
    }
    pub fn poison(&mut self) {
        self.status = UnsafeAssumedSingleThreadLockStatus::Poisoned;
    }
}

impl<'a, T> Deref for UnsafeAssumedSingleThreadLock<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

Now, in lib.rs I can implement the panic hook:

#[wasm_bindgen(start)]
pub fn start() -> Result<(), JsValue> {
    panic::set_hook(Box::new(|info| {
        editor_singleton::set_panicked();
        // Also give ourselves some readable stack traces, of course.
        console_error_panic_hook::hook(info);
    }));
    Ok(())
}

All together, this results in what I want - our instance is lazily initialized whenever accessed, allowed to be read without locking, and allowed to be written by obtaining a UASTGuard, which when dropped will release the lock. If set_panicked() is called, then the next time the instance is accessed, it will be recreated.

Now our WASM code can run statelessly, and if a panic happens, the editor will be transparently recreated without the user noticing more than the cost of initialization, which in Rust is rarely noticeable.

Conclusion

After working through this, I thought it was a pretty good test case for understanding some of the subtleties when bridging the assumptions in Rust and the assumptions in ServiceWorkers and Javascript threads. This solution feels so far like the best of both worlds. I may at some point publish some library code related to this, but in its current state it feels best managed inside application codebases, especially due to the dangerous nature of the unsafe blocks - you really must only use this type of pattern when you can guarantee a single threaded environment - something I can’t do from inside a crate.