Tokio, channels and poisoned locks

TL;DR: my trading bot keeps the order books of several exchanges in memory, with one long-lived Tokio task per WebSocket connection. In this post I go through that architecture pattern and the select! loop inside each task. Then I show how each task builds the order book of a market: first a full copy of the book, then apply only the changes, with a check that no message was lost. To wrap up, an illustrative bug: a latent poisoned lock that could have left the order book stale without any warning. A small crate abstracts the example to reproduce it in case you want the hands-on experience.

Disclosure: This post was written by me. I used an LLM for editorial feedback on structure and phrasing, for the figures and for the lab crate. The trading bot, its architecture and the decisions behind it are my own.

Contents

1. The actor model: one task per connection
2. Zooming into one task
3. The poisoned lock
4. Hands-on

1. The actor model: one task per connection

The bot trades on several exchanges, and for each decision it needs the best bid and ask of many markets at the same time. I don't run strategies that need to win the latency race, but I do need the books as fresh as the feeds allow, because each trade opens two hedged positions (one long and one short) and both are decided at the same moment. If the bot asked each exchange over REST at the moment of a decision, it would add a network round trip to the critical path. The decision would be slower, and the prices it uses would be older, so the books need to live in memory. One task per exchange keeps them up to date from the WebSocket, in the background, and the strategy reads them in place. That means one writer and several readers for each exchange's books. On most exchanges they sit behind an RwLock, and on the exchanges that send too many frames, behind arc-swap.

The pattern belongs to the actor model family. An actor is a unit that owns some state and runs on its own. In the strict model, the rest of the program never touches that state and only sends the actor messages. In Tokio the unit is a task, and my feeds follow the model loosely: each task is the only writer of its books, but the strategy reads them in place.

At start-up, the main task spawns one task per enabled exchange:

let book_a = BookFeed::new();
if cfg.exchange_a.enabled {
    tokio::spawn(run_book_feed(cfg.exchange_a.ws_url.clone(), markets, book_a.clone()));
}
router.set_book_feed(book_a.clone());

The main task does not keep the JoinHandle, because nothing needs to wait on the feed. Dropping the handle does not cancel the task, it only detaches it, so each feed runs for the life of the process. Each task owns its WebSocket, its subscriptions and the queue of messages it has to send.

With the std::sync::RwLock (Arc<RwLock<HashMap<MarketId, Book>>>), the task takes the write lock for each update, and the strategy takes a short read lock. With arc-swap there is no lock: after each frame, the task publishes an immutable view of the books, and readers load it without waiting. The second one is closer to the actor idea, because only the task holds the mutable state.

With arc-swap the pointer swap is atomic, so the reader never gets a half-written frame. However, a reader may get the view from one frame earlier, so it is a trade-off. That's why I use the lock on most exchanges, and arc-swap only where the socket sends too many frames. With that many writes, the writer and the readers keep waiting on each other, the RwLock behaves almost like a Mutex, and the lock slows everything down.

2. Zooming into one task

Each feed task has three levels:

1. The reconnect loop runs for the life of the process. Each turn opens one WebSocket connection and runs a session on it.
2. The select! loop is the session. It runs while the connection is up, and it waits on the socket and on three timers.
3. The book update runs for each frame that brings book data. It applies the frame to the order book of that market and checks that no message was lost.

The reconnect loop

The outermost level is a loop that calls session again and again. A session opens one connection to the exchange and returns when that connection ends. A connection can end in three ways:

• the exchange sends a close frame
• the socket returns an error
• the liveness timer below finds that no frame has arrived for 10 s

Regardless of the reason, the loop will mark each book invalid, because it may not be current anymore. It will then wait and reconnect with exponential backoff. The wait starts at 1 s and doubles on each reconnect, up to 30 s. If the session that just ended stayed up for more than 30 s, the wait goes back to 1 s.

pub async fn run_book_feed(ws_url: String, markets: Vec<MarketId>, feed: BookFeed) {
    let mut backoff = MIN_BACKOFF; // 1 s
    loop {
        let started = Instant::now();
        if let Err(e) = session(&ws_url, &markets, &feed).await {
            log::warn!("stream failed: {e:#}, reconnecting");
        }
        invalidate_all(&feed.book);
        if started.elapsed() > MAX_BACKOFF {
            backoff = MIN_BACKOFF;
        }
        tokio::time::sleep(backoff).await;
        backoff = (backoff * 2).min(MAX_BACKOFF); // 30 s
    }
}

The select! loop

Inside a session, the task waits on four things. One is to handle incoming data, and the other three are for keeping the connection healthy:

1. ws.next(), the next frame from the socket. This is the branch that carries the book data.
2. liveness.tick(), a timer every 5 s. It ends the session if no frame of any kind has arrived for 10 s. Since a socket can die silently, with no close frame and no error, this check is the only way to see it.
3. ping.tick(), a timer every 60 s. It marks a ping as due, because the exchange closes a socket that sends nothing for two minutes.
4. sender.tick(), a timer every 200 ms. It sends the pending ping and the queued subscribe and unsubscribe messages, as many as the exchange's rate limit allows.

So tokio::select! polls all four branches, runs the first one that is ready and drops the others until the next turn of the loop.

loop {
    tokio::select! {
        frame = ws.next() => {
            let Some(frame) = frame else { return Ok(()) };
            last_frame = Instant::now();
            handle_frame(frame?, &feed.book, &mut to_send)?;
        }
        _ = liveness.tick() => {
            if last_frame.elapsed() >= SILENCE_LIMIT {
                anyhow::bail!("no frame for {:?}, treating the socket as dead", last_frame.elapsed());
            }
        }
        _ = ping.tick() => { ping_due = true; }
        _ = sender.tick() => drain_queue(&mut ws, &mut to_send, &mut ping_due).await?,
    }
}

Note: a branch that select! drops has to be safe to drop halfway (cancel safe). ws.next() and interval.tick() are, so they can sit in the loop as they are. These feeds have no shutdown branch, because they live as long as the process. A service that has to stop cleanly would add one more branch, on a watch channel or a cancellation token (the crate at the end does).

Building the order book

An order book is the list of open orders on one market, grouped by price. Each price that has orders on it is a level, with the total size waiting at that price. Bids (buy orders) are on one side and asks (sell orders) on the other. The task keeps each side in a BTreeMap keyed by price. A BTreeMap keeps its keys sorted, so the best bid is the last key on the bid side and the best ask is the first key on the ask side.

The exchange does not send the whole book on every change. When the task subscribes to a market, the exchange first sends a snapshot (the whole book at that moment). After that it only sends deltas (the levels that changed since its previous message). A delta gives the new size of each level, not the difference, and a size of zero removes the level. Each delta also carries its own sequence number and the sequence number of the message before it. So the task can check that each delta follows the last one it applied.

A book is valid if and only if the two conditions hold:

• it was built from a snapshot
• every delta since then followed the one before it

That is the only test for the strategy to know it can read the book. The age of the book doesn't matter: on a live connection with no missed delta, a book that hasn't changed for a minute is still considered current.

When a delta does not follow the last one, there is a gap, so we know a message was lost. The task marks that market's book invalid and asks for a fresh snapshot. On this exchange, subscribing again to a market that is already subscribed does nothing, so the repair is an unsubscribe followed by a subscribe. The other markets on the same socket are untouched.

To get an idea of how often this happens, I measured it from the production server: 0 gaps in 3,365 updates across 5 markets.

fn apply(&mut self, d: Delta) -> Result<(), BookError> {
    let Some(seq) = self.seq else { return Err(BookError::NotSynced) };
    if d.prev_seq != seq {
        self.invalidate();
        return Err(BookError::Gap { expected: seq, got: d.prev_seq });
    }
    for level in d.bids { set_level(&mut self.bids, level) }
    for level in d.asks { set_level(&mut self.asks, level) }
    self.seq = Some(d.seq);
    Ok(())
}

3. The poisoned lock

std::sync::Mutex and RwLock remember when a thread panicked while holding the guard. From then on lock() (or write(), read()) returns Err(PoisonError), because the panic may have left the data half written and the lock has no way to know. You can do three things with that error:

• unwrap() and panic as well;
• into_inner(), which hands you the guard anyway, for when you can repair the data or are about to throw it away;
• clear_poison(), once the data is repaired.

Note: tokio::sync::Mutex and parking_lot don't poison.

The latent bug: in my feed the write path was if let Ok(mut b) = book.write() { ... }. If a panic in the middle of a write poisoned the lock, every later write would have been skipped with no log line while the socket was still up, so the books would stop updating and nothing would indicate why. I found it with a test harness, not in production: one of the tests takes the write guard in a thread and panics, the way a bug in the update code would.

The fix makes a poisoned lock an error that ends the session, and the reconnect path, which invalidates everything anyway, takes the guard back and clears the poison. Then the next session can write fresh snapshots:

let mut b = book
    .write()
    .map_err(|_| anyhow::anyhow!("book lock poisoned, rebuilding from fresh snapshots"))?;
 
fn invalidate_all(book: &SharedBook) {
    let mut b = book.write().unwrap_or_else(|poisoned| {
        book.clear_poison();
        poisoned.into_inner()
    });
    for s in b.values_mut() {
        s.invalidate();
    }
}

Recovering with into_inner is right here because we immediately throw away the data.

4. Hands-on

I created a small crate that abstracts this architecture and these errors, so you can run them yourself: tokio-actor-orderbook-lab. It has no exchange code. The exchange is a fake socket made of two channels, and it can drop deltas or lose the first snapshot on purpose.

git clone https://github.com/reymom/tokio-actor-orderbook-lab
cd tokio-actor-orderbook-lab
cargo test            # the book, the gap and its repair, shutdown, the poisoned lock
cargo run --bin demo  # three exchanges for two seconds, then a clean shutdown

The crate has five modules:

• book: the order book from section 2, with the snapshot, the deltas and the sequence check.
• venue: the fake exchange for one market, as a Tokio task. Like the real one, it ignores a second subscribe, so a repair is an unsubscribe followed by a subscribe.
• actor: one task per exchange. Its select! waits on the next frame, a heartbeat and a shutdown signal on a watch channel (the branch my feeds don't have).
• aggregator: one task that receives the best bid and ask of every exchange over a single mpsc channel.
• poison: the same book behind a std::sync::Mutex, with the bug from section 3 and its fix.

In the demo, exchange A is clean, B drops every 25th delta and C loses its first snapshot. A run prints something like this:

venue A  snapshots=1  deltas=201  dropped_early=0  gaps=0  snapshot_repairs=0
venue B  snapshots=9  deltas=185  dropped_early=0  gaps=8  snapshot_repairs=0
venue C  snapshots=1  deltas=150  dropped_early=51  gaps=0  snapshot_repairs=1
aggregator  updates=555  best_bid=99 (C)  best_ask=102 (A)

B repairs each gap with a fresh snapshot, so it ends with one snapshot more than gaps. C receives deltas before any snapshot, so it drops them until the heartbeat asks for the snapshot again.

The poisoned lock is two tests. The first one shows the bug: after the panic, the write is skipped and nothing says so.

#[test]
fn a_poisoned_lock_makes_the_buggy_writer_go_quiet() {
    let book = Arc::new(Mutex::new(Book::default()));
    assert!(write_if_ok(&book, 1, &[lv(99, 1)], &[lv(101, 1)])); // lv(price, size) is one level
    poison_on_purpose(&book);
    assert!(book.is_poisoned());
    assert!(!write_if_ok(&book, 2, &[lv(98, 1)], &[lv(102, 1)]));
}

The second one runs the fix: the write fails with an error, the reconnect path clears the poison and throws the book away, and the next snapshot is written again.