Eli Bendersky's website - Rust

Concurrent Servers: Part 7 - Rust

2026-08-15T09:41:00-07:00

This is part 7 in a series of posts on writing concurrent network servers. In this part, we discuss how the challenges described in earlier parts are tackled in the Rust programming language.

All posts in the series:

Part 1 - Introduction
Part 2 - Threads
Part 3 - Event-driven
Part 4 - libuv
Part 5 - Redis case study
Part 6 - Callbacks, Promises and async/await
Part 7 - Rust (this part)
Part 8 - Go

Several years have passed since the previous parts were published. I've recently went over them to make sure the information presented is still relevant and all the code samples build and run using modern toolchains. I strongly recommend reviewing the previous parts before reading this one.

This post assumes a basic familiarity with the Rust programming language. It will only explain Rust constructs when we encounter code that wouldn't appear in an introductory book or tutorial.

Setting the baseline - a sequential state machine server

The first few parts in the series focused on a socket server that implements a simple state machine protocol. See part 1 for a complete description of the protocol. Let's start by showing how this protocol is implemented in a basic sequential Rust server:

use async_socket_server::serve_connection;
use std::net::TcpListener;

fn main() -> std::io::Result<()> {
    let port = match std::env::args().nth(1) {
        Some(s) => s,
        None => "9090".to_string(),
    };
    let addr = format!("127.0.0.1:{port}");

    let listener = TcpListener::bind(addr)?;
    println!("Serving on port {port}");

    loop {
        let (stream, addr) = listener.accept()?;
        println!("connection received from {}", addr);
        if let Err(e) = serve_connection(stream) {
            eprintln!("error serving connection: {}", e);
        } else {
            println!("peer done {addr}");
        }
    }
}

With the function serve_connection defined as:

pub enum ProcessingState {
    WaitForMsg,
    InMsg,
}

pub fn serve_connection(mut stream: TcpStream) -> std::io::Result<()> {
    stream.write_all(b"*")?;

    let mut state = ProcessingState::WaitForMsg;
    let mut buf = [0u8; 1024];
    loop {
        let n = stream.read(&mut buf)?;
        if n == 0 {
            // Connection closed by the client.
            break;
        }

        for byte in &buf[..n] {
            match state {
                ProcessingState::WaitForMsg => {
                    if *byte == b'^' {
                        state = ProcessingState::InMsg;
                    }
                }
                ProcessingState::InMsg => {
                    if *byte == b'$' {
                        state = ProcessingState::WaitForMsg;
                    } else {
                        let newbyte = byte.wrapping_add(1);
                        stream.write_all(&[newbyte])?;
                    }
                }
            }
        }
    }

    Ok(())
}

As a reminder, this server version is sequential because it accepts clients one by one; the main loop blocks on serve_connection until it's done (the client closes the connection), and only then goes back to accept the next client.

One thread per client

Clearly, handling clients one by one won't do. In part 2, we've discussed approaches that use OS threads to handle clients concurrently. Let's start with the unbounded one-thread-per-client solution in Rust:

use async_socket_server::serve_connection;
use std::{net::TcpListener, thread};

fn main() -> std::io::Result<()> {
    let port = match std::env::args().nth(1) {
        Some(s) => s,
        None => "9090".to_string(),
    };
    let addr = format!("127.0.0.1:{port}");

    let listener = TcpListener::bind(addr)?;
    println!("Serving on port {port}");

    loop {
        let (stream, addr) = listener.accept()?;
        println!("connection received from {}", addr);

        let res = thread::Builder::new().spawn(move || {
            if let Err(e) = serve_connection(stream) {
                eprintln!("error serving connection: {}", e);
            } else {
                println!("peer done {addr}");
            }
        });

        if let Err(e) = res {
            eprintln!("error spawning thread: {}", e);
        }
    }
}

The spawn method returns a Result<JoinHandle<T>>; on success, we allow the handle to be dropped at the end of the loop iteration. In Rust, this detaches the thread; we don't actually wait for it to complete. This is reasonable for our code sample, because the loop is infinite; it never terminates anyway. The potential for runaway threads is just one of the issues with the unbounded threads approach discussed in part 2. The solution is to use a fixed thread pool.

Thread pool

Before diving into the code, a quick note on the design: the thread pool is a fixed set of threads that await "jobs" and handle them to completion. In our case a "job" is serve_connection for a specific client. There are many ways to implement a thread pool; for our use case, I went with a set of threads that all get a shared channel to which the main thread sends jobs. A worker thread picks up the next job from the channel, serves it to completion, and goes back to waiting for the next job. Here's how this looks in code:

struct Job {
    stream: TcpStream,
    addr: SocketAddr,
}

fn worker(receiver: Receiver<Job>) {
    while let Ok(job) = receiver.recv() {
        if let Err(e) = serve_connection(job.stream) {
            eprintln!("error serving connection from {}: {}", job.addr, e);
        } else {
            println!("peer done {}", job.addr);
        }
    }
}

What is Receiver? It's a type from the crossbeam_channel crate:

use crossbeam_channel::{Receiver, bounded};

Rust's builtin channels in std are mpsc - multi producer, single consumer, but what we need for our job queue is a channel that supports multiple consumers (the worker threads). While std does have mpmc, this is an experimental API only available in nightly versions at the time of writing. Therefore, I've opted to include the crossbeam_channel crate that provides well-tested mpmc channels for this sample [1].

And here's the main function:

use async_socket_server::serve_connection;
use crossbeam_channel::{Receiver, bounded};
use std::{
    io,
    net::{SocketAddr, TcpListener, TcpStream},
    thread,
};

const NUM_WORKERS: usize = 256;
const JOB_QUEUE_CAPACITY: usize = NUM_WORKERS;

fn main() -> io::Result<()> {
    let port = match std::env::args().nth(1) {
        Some(s) => s,
        None => "9090".to_string(),
    };
    let addr = format!("127.0.0.1:{port}");

    let listener = TcpListener::bind(addr)?;
    println!("Serving on port {port}");

    let (tx, rx) = bounded::<Job>(JOB_QUEUE_CAPACITY);
    for _ in 0..NUM_WORKERS {
        let receiver = rx.clone();
        thread::Builder::new().spawn(move || worker(receiver))?;
    }
    drop(rx);

    loop {
        let (stream, addr) = listener.accept()?;
        println!("connection received from {addr}");

        // A full queue blocks this loop, preventing it from accepting more
        // connections until a worker becomes available.
        if tx.send(Job { stream, addr }).is_err() {
            return Err(io::Error::other("all connection workers stopped"));
        }
    }
}

Note that our job channel is bounded - it has a fixed size. This helps naturally implement a backpressure mechanism - if too many clients connect, the following clients will have to wait - the main loop blocks on tx.send and won't accept additional clients on the socket until jobs are cleared from the channel.

Asynchronous, event-driven server

In parts 4, 5 and 6 of the series we've discussed event-driven, or asynchronous servers. Let's see how it's done in Rust. Specifically, part 6 presented a gradation from callbacks to promises to async/await mechanisms; Rust supports all of these and - as you'd expect - modern code is usually written with async/await while hiding all the details of promises (called futures in Rust) underneath.

Without further ado, here's our simple state machine protocol in asynchronous Rust:

use async_socket_server::ProcessingState;
use tokio::io::{self, AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;

#[tokio::main]
async fn main() -> io::Result<()> {
    let port = match std::env::args().nth(1) {
        Some(s) => s,
        None => "9090".to_string(),
    };
    let addr = format!("127.0.0.1:{port}");

    let listener = TcpListener::bind(addr).await?;
    println!("Serving on port {port}");

    loop {
        let (socket, addr) = listener.accept().await?;
        println!("connection received from {addr:?}");

        tokio::spawn(async move {
            if let Err(e) = serve_connection_async(socket).await {
                eprintln!("error serving connection from {addr:?}: {e}");
            } else {
                println!("peer done {addr:?}");
            }
        });
    }
}

Rust takes an interesting approach to async programming: it supports some of its fundamental building blocks (like futures and the async and await keywords) in the core language, but leaves the actual async engine implementation (the thing that implements the event loop) to external crates. By far the most popular crate for async programming in Rust is is tokio, so that's what we're using here.

After reading the JS code in part 6, the Rust snippet above should appear fairly familiar, except perhaps the explicit tokio task "spawn". Instead of enqueuing a callback on the connection returned by listener.accept, the code spawns a tokio task, which can be seen as a green thread, and hence uses similar terminology [2]. These tasks must not issue blocking calls; therefore, they are supposed to use tokio's I/O utilities instead of the usual, blocking std utilities. In fact, we have to implement an async version of serve_connection to make this work:

async fn serve_connection_async(mut stream: tokio::net::TcpStream) -> std::io::Result<()> {
    stream.write_all(b"*").await?;

    let mut state = ProcessingState::WaitForMsg;
    let mut buf = [0u8; 1024];
    loop {
        let n = stream.read(&mut buf).await?;
        if n == 0 {
            // Connection closed by the client.
            break;
        }

        for byte in &buf[..n] {
            match state {
                ProcessingState::WaitForMsg => {
                    if *byte == b'^' {
                        state = ProcessingState::InMsg;
                    }
                }
                ProcessingState::InMsg => {
                    if *byte == b'$' {
                        state = ProcessingState::WaitForMsg;
                    } else {
                        let newbyte = byte.wrapping_add(1);
                        stream.write_all(&[newbyte]).await?;
                    }
                }
            }
        }
    }

    Ok(())
}

Note how similar this code is to serve_connection from earlier; the only real differences are the await calls on socket reads and writes [3], and the types involved. For example, instead of a std::net::TcpStream used in the synchronous samples, here we're using tokio::net::TcpStream. Tokio has an underlying dependency called mio to handle non-blocking APIs for all kinds of I/O. It wraps OS-specific event loops like epoll to do so efficiently.

Asynchronous primality testing servers

While most of the series has been using a simple state machine server as the driving example, part 6 switched focus to a server for primality testing which simulates long compute tasks. Let's see how this is done in Rust with tokio:

use tokio::io::{self, AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;

#[tokio::main]
async fn main() -> io::Result<()> {
    let port = match std::env::args().nth(1) {
        Some(s) => s,
        None => "8070".to_string(),
    };
    let addr = format!("127.0.0.1:{port}");

    let listener = TcpListener::bind(addr).await?;
    println!("Serving on port {port}");

    loop {
        let (socket, addr) = listener.accept().await?;
        println!("connection received from {addr:?}");

        tokio::spawn(async move {
            if let Err(e) = serve_client(socket).await {
                eprintln!("error serving connection from {addr:?}: {e}");
            } else {
                println!("peer done {addr:?}");
            }
        });
    }
}

async fn serve_client(mut stream: tokio::net::TcpStream) -> std::io::Result<()> {
    let mut buf = [0u8; 1024];

    loop {
        let n = stream.read(&mut buf).await?;
        if n == 0 {
            // Connection closed by the client.
            return Ok(());
        }

        // Parse read buf into u64.
        let num = std::str::from_utf8(&buf[..n])
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?
            .trim()
            .parse::<u64>()
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;

        let answer = if isprime(num, true) {
            "prime"
        } else {
            "composite"
        };
        stream
            .write_all((answer.to_string() + "\n").as_bytes())
            .await?;
    }
}

This code is very similar to the previous snippet conceptually; isprime is:

// Check if n is prime, returning a boolean. The delay parameter is optional;
// if true, the function will block for n milliseconds before computing the
// answer. This is useful for simulating a long-running computation.
fn isprime(n: u64, delay: bool) -> bool {
    if delay {
        std::thread::sleep(std::time::Duration::from_millis(n));
    }

    if n < 2 {
        return false;
    }
    if n % 2 == 0 {
        return n == 2;
    }

    let mut r = 3;
    while r * r <= n {
        if n % r == 0 {
            return false;
        }
        r += 2;
    }
    true
}

Note that this sample demonstrates a job that can block (simulated with a sleep in this case). This can be problematic in an async context, as the tokio documentation explains.

One potential solution would be to dispatch a blocking task to a separate thread pool and use tokio channels to communicate with it; this is similar to the approach we've taken in the thread pool sample above.

Part 6 also included a version of this server that caches data on a local Redis instance; the goal was to demonstrate the complexity of event-driven code when additional layers of callbacks are added and how async/await can help mitigate that. Here's our Rust version of this server, using the redis crate (that has a tokio component enabled explicitly to support async calls):

use redis::AsyncTypedCommands;
use redis::aio::MultiplexedConnection;
use tokio::io::{self, AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;

#[derive(Clone)]
struct AppState {
    redis_connection: MultiplexedConnection,
}

const REDIS_URL: &str = "redis://127.0.0.1";

#[tokio::main]
async fn main() -> io::Result<()> {
    let port = match std::env::args().nth(1) {
        Some(s) => s,
        None => "8070".to_string(),
    };
    let addr = format!("127.0.0.1:{port}");

    let app_state = AppState {
        redis_connection: {
            redis::Client::open(REDIS_URL)
                .map_err(io::Error::other)?
                .get_multiplexed_async_connection()
                .await
                .map_err(io::Error::other)?
        },
    };

    let listener = TcpListener::bind(addr).await?;
    println!("Serving on port {port}");

    loop {
        let (socket, addr) = listener.accept().await?;
        println!("connection received from {addr:?}");
        let app_state = app_state.clone();

        tokio::spawn(async move {
            if let Err(e) = serve_client(socket, app_state).await {
                eprintln!("error serving connection from {addr:?}: {e}");
            } else {
                println!("peer done {addr:?}");
            }
        });
    }
}

async fn serve_client(
    mut stream: tokio::net::TcpStream,
    mut app_state: AppState,
) -> std::io::Result<()> {
    let mut buf = [0u8; 1024];

    loop {
        let n = stream.read(&mut buf).await?;
        if n == 0 {
            // Connection closed by the client.
            return Ok(());
        }

        // Parse read buf into u64.
        let num = std::str::from_utf8(&buf[..n])
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?
            .trim()
            .parse::<u64>()
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;

        // Try the redis cache first. If found in cache, send the cached
        // answer.
        let cachekey = format!("primecache:{num}");
        match app_state.redis_connection.get(cachekey.clone()).await {
            Ok(Some(cached)) => {
                stream.write_all((cached.clone() + "\n").as_bytes()).await?;
                println!("cached num {num} is {cached}");
                continue;
            }
            Ok(None) => {
                // Not found in cache, continue to compute.
            }
            Err(error) => {
                return Err(io::Error::other(error));
            }
        }

        let answer = if isprime(num, true) {
            "prime"
        } else {
            "composite"
        };

        // Save answer in cache and send it to the client.
        app_state
            .redis_connection
            .set(cachekey, answer)
            .await
            .map_err(io::Error::other)?;
        stream
            .write_all((answer.to_string() + "\n").as_bytes())
            .await?;
    }
}

Notes:

Because of the function color problem, the redis crate has a connection constructor specifically for async: get_multiplexed_async_connection.
Here we have an example of shared state between tokio tasks - the Redis connection. Note that we don't require any particular synchronization because MultiplexedConnection is Clone; cloning it to different tasks is safe - and in fact that's what we do for each new task. There's no magic here; if you look inside MultiplexedConnection, you'll see that it already has all the synchronization mechanisms implemented internally, as needed.
Due to the magic of async/await, the code in serve_client is nice and linear. We simply await on the Redis call, and once it's back we continue with the rest of the handler. Since we're using an async Redis connection, in case waiting is required, control will be ceded to some other task that's not currently blocked on I/O.

In conclusion, while Rust provides excellent support for async programming, it doesn't solve its inherent issues like function colors and the need for careful separation between blocking and non-blocking tasks. These issues are typically surmountable with some extra care, and async programming with Tokio in Rust is very popular due to its performance benefits.

Code

All the code for this post is available on GitHub.

[1]

Rust makes a conscious design choice to keep its standard library minimal. In part 2, we've used Python's builtin stdlib thread pools; Rust also has several crates that implement thread pools, but I decided against using them because they obscure the underlying working of the code too much. Our sample implements a simple thread pool, but unfortunately still has to reach for an external crate to use a well-supported multi-consumer channel mechanism.

[2] tokio is quite sophisticated; it manages a thread pool with an efficient work-stealing mechanism, onto which new tasks are spawned. Therefore, all shared data must be protected with mutexes or other synchronization primitives, and Rust's traits like Send and Sync play an important role, where applicable.

[3]	Having to code this alternative just for the sake of async invocation is a classical example of the function color problem.

Plugins case study: mdBook preprocessors

2025-12-17T18:11:00-08:00

mdBook is a tool for easily creating books out of Markdown files. It's very popular in the Rust ecosystem, where it's used (among other things) to publish the official Rust book.

mdBook has a simple yet effective plugin mechanism that can be used to modify the book output in arbitrary ways, using any programming language or tool. This post describes the mechanism and how it aligns with the fundamental concepts of plugin infrastructures.

mdBook preprocessors

mdBook's architecture is pretty simple: your contents go into a directory tree of Markdown files. mdBook then renders these into a book, with one file per chapter. The book's output is HTML by default, but mdBook supports other outputs like PDF.

The preprocessor mechanism lets us register an arbitrary program that runs on the book's source after it's loaded from Markdown files; this program can modify the book's contents in any way it wishes before it all gets sent to the renderer for generating output.

The official documentation explains this process very well.

Sample plugin

I rewrote my classical "nacrissist" plugin for mdBook; the code is available here.

In fact, there are two renditions of the same plugin there:

One in Python, to demonstrate how mdBook can invoke preprocessors written in any programming language.
One in Rust, to demonstrate how mdBook exposes an application API to plugins written in Rust (since mdBook is itself written in Rust).

Fundamental plugin concepts in this case study

Let's see how this case study of mdBook preprocessors measures against the Fundamental plugin concepts that were covered several times on this blog.

Discovery

Discovery in mdBook is very explicit. For every plugin we want mdBook to use, it has to be listed in the project's book.toml configuration file. For example, in the code sample for this post, the Python narcissist plugin is noted in book.toml as follows:

[preprocessor.narcissistpy]
command = "python3 ../preprocessor-python-narcissist/narcissist.py"

Each preprocessor is a command for mdBook to execute in a sub-process. Here it uses Python, but it can be anything else that can be validly executed.

Registration

For the purpose of registration, mdBook actually invokes the plugin command twice. The first time, it passes the arguments supports <renderer> where <renderer> is the name of the renderer (e.g. html). If the command returns 0, it means the preprocessor supports this renderer; otherwise, it doesn't.

In the second invocation, mdBook passes some metadata plus the entire book in JSON format to the preprocessor through stdin, and expects the preprocessor to return the modified book as JSON to stdout (using the same schema).

Hooks

In terms of hooks, mdBook takes a very coarse-grained approach. The preprocessor gets the entire book in a single JSON object (along with a context object that contains metadata), and is expected to emit the entire modified book in a single JSON object. It's up to the preprocessor to figure out which parts of the book to read and which parts to modify.

Given that books and other documentation typically have limited sizes, this is a reasonable design choice. Even tens of MiB of JSON-encoded data are very quick to pass between sub-processes via stdout and marshal/unmarshal. But we wouldn't be able to implement Wikipedia using this design.

Exposing an application API to plugins

This is tricky, given that the preprocessor mechanism is language-agnostic. Here, mdBook only offers additional utilities to preprocessors implemented in Rust. These get access to mdBook's API to unmarshal the JSON representing the context metadata and book's contents. mdBook offers the Preprocessor trait Rust preprocessors can implement, which makes it easier to wrangle the book's contents. See my Rust version of the narcissist preprocessor for a basic example of this.

Renderers / backends

Actually, mdBook has another plugin mechanism, but it's very similar conceptually to preprocessors. A renderer (also called a backend in some of mdBook's own doc pages) takes the same input as a preprocessor, but is free to do whatever it wants with it. The default renderer emits the HTML for the book; other renderers can do other things.

The idea is that the book can go through multiple preprocessors, but at the end a single renderer.

The data a renderer receives is exactly the same as a preprocessor - JSON encoded book contents. Due to this similarity, there's no real point getting deeper into renderers in this post.

FAAS in Go with WASM, WASI and Rust

2023-05-06T06:44:00-07:00

This post is best described as a technology demonstration; it melds together web servers, plugins, WebAssembly, Go, Rust and ABIs. Here's what it shows:

How to load WASM code with WASI in a Go environment and hook it up to a web server.
How to implement web server plugins in any language that can be compiled to WASM.
How to translate Go programs into WASM that uses WASI.
How to translate Rust programs into WASM that uses WASI.
How to write WAT (WebAssembly Text) code that uses WASI to interact with a non-JS environment.

We're going to build a simple FAAS (Function as a Service) server in Go that lets us write modules in any language that has a WASM target. Comparing to existing technologies, it's something between GCP's Cloud Functions, Cloud Run and good old CGI.

Design

Let's start with a high-level diagram describing how the system works:

The steps numbered in the diagram are:

The FAAS server receives an HTTP GET request, with a path consisting of a module name (func in the example in the diagram) and an arbitrary query string.
The FAAS server finds and loads the WASM module corresponding to the module name it was provided, and invokes it with a description of the HTTP request.
The module emits output to its stdout, which is captured by the FAAS server.
The FAAS server uses the module's stdout as the contents of an HTTP Response to the request it received.

The FAAS server

We'll start our deep dive with the FAAS server itself (full code here). The HTTP handling part is straightforward:

func httpHandler(w http.ResponseWriter, req *http.Request) {
  parts := strings.Split(strings.Trim(req.URL.Path, "/"), "/")
  if len(parts) < 1 {
    http.Error(w, "want /{modulename} prefix", http.StatusBadRequest)
    return
  }
  mod := parts[0]
  log.Printf("module %v requested with query %v", mod, req.URL.Query())

  env := map[string]string{
    "http_path":   req.URL.Path,
    "http_method": req.Method,
    "http_host":   req.Host,
    "http_query":  req.URL.Query().Encode(),
    "remote_addr": req.RemoteAddr,
  }

  modpath := fmt.Sprintf("target/%v.wasm", mod)
  log.Printf("loading module %v", modpath)
  out, err := invokeWasmModule(mod, modpath, env)
  if err != nil {
    log.Printf("error loading module %v", modpath)
    http.Error(w, "unable to find module "+modpath, http.StatusNotFound)
    return
  }

  // The module's stdout is written into the response.
  fmt.Fprint(w, out)
}

func main() {
  mux := http.NewServeMux()
  mux.HandleFunc("/", httpHandler)
  log.Fatal(http.ListenAndServe(":8080", mux))
}

This server listens on port 8080 (feel free to change this or make it more configurable), and registers a catch-all handler for the root path. The handler parses the actual request URL to find the module name. It then stores some information to pass to the loaded module in the env map.

The loaded module is found with a filesystem lookup in the target directory relative to the FAAS server binary. All of this is just for demonstration purposes and can be easily changed, of course. The handler then calls invokeWasmModule, which we'll get to shortly. This function returns the invoked module's stdout, which the handler prints out into the HTTP response.

Running WASM code in Go

Given a WASM module, how do we run it programmatically in Go? There are several high-quality WASM runtimes that work outside the browser environment, and many of them have Go bindings; for example wasmtime-go. The one I like most, however, is wazero; it's a zero-dependency, pure Go runtime that doesn't have any prerequisites except running a go get. Our FAAS server is using wazero to load and run WASM modules.

Here's invokeWasmModule:

// invokeWasmModule invokes the given WASM module (given as a file path),
// setting its env vars according to env. Returns the module's stdout.
func invokeWasmModule(modname string, wasmPath string, env map[string]string) (string, error) {
  ctx := context.Background()

  r := wazero.NewRuntime(ctx)
  defer r.Close(ctx)
  wasi_snapshot_preview1.MustInstantiate(ctx, r)

  // Instantiate the wasm runtime, setting up exported functions from the host
  // that the wasm module can use for logging purposes.
  _, err := r.NewHostModuleBuilder("env").
    NewFunctionBuilder().
    WithFunc(func(v uint32) {
      log.Printf("[%v]: %v", modname, v)
    }).
    Export("log_i32").
    NewFunctionBuilder().
    WithFunc(func(ctx context.Context, mod api.Module, ptr uint32, len uint32) {
      // Read the string from the module's exported memory.
      if bytes, ok := mod.Memory().Read(ptr, len); ok {
        log.Printf("[%v]: %v", modname, string(bytes))
      } else {
        log.Printf("[%v]: log_string: unable to read wasm memory", modname)
      }
    }).
    Export("log_string").
    Instantiate(ctx)
  if err != nil {
    return "", err
  }

  wasmObj, err := os.ReadFile(wasmPath)
  if err != nil {
    return "", err
  }

  // Set up stdout redirection and env vars for the module.
  var stdoutBuf bytes.Buffer
  config := wazero.NewModuleConfig().WithStdout(&stdoutBuf)

  for k, v := range env {
    config = config.WithEnv(k, v)
  }

  // Instantiate the module. This invokes the _start function by default.
  _, err = r.InstantiateWithConfig(ctx, wasmObj, config)
  if err != nil {
    return "", err
  }

  return stdoutBuf.String(), nil
}

Interesting things to note about this code:

wazero supports WASI, which has to be instantiated explicitly to be usable by the loaded modules.
A lot of the code deals with exporting logging functions from the host (the Go code of the FAAS server) to the WASM module.
We set up the loaded module's stdout to be redirected to a buffer, and set up its environment variables to match the env map passed in.

There are several way for host code to interact with WASM modules using only the WASI API and ABI. Here, we opt for using environment variables for input and stdout for output, but there are other options (see the Other resources section in the bottom for some pointers).

This is it - the whole FAAS server, about 100 LOC of commented Go code. Now let's move on to see some WASM modules this thing can load and run.

Writing modules in Go

We can compile Go code to WASM that uses WASI. Here's a basic Go program that emits a greeting and a listing of its environment variables to stdout:

package main

import (
  "fmt"
  "os"
)

func main() {
  fmt.Println("goenv environment:")

  for _, e := range os.Environ() {
    fmt.Println(" ", e)
  }
}

To build using the standard Go toolchain:

$ GOOS=wasip1 GOARCH=wasm gotip build -o target/goenv.wasm examples/goenv/goenv.go

(the wasip1 target name refers to "WASI Preview 1")

Sharp-eyed readers will recall that the target/ directory is precisely where the FAAS server looks for *.wasm files to load as modules. Now that we've placed a module named goenv.wasm there, we're ready to launch our server with go run . in the root directory. We can issue a HTTP request to its goenv module in a separate terminal:

$ curl "localhost:8080/goenv?foo=bar&id=1234"
goenv environment:
  http_method=GET
  http_host=localhost:8080
  http_query=foo=bar&id=1234
  remote_addr=127.0.0.1:59268
  http_path=/goenv

And looking at the terminal where the FAAS server runs we'll see some logging like:

2023/04/29 06:35:59 module goenv requested with query map[foo:[bar] id:[1234]]
2023/04/29 06:35:59 loading module target/goenv.wasm

Another option to build this is using the TinyGo compiler. In our FAAS project structure, the invocation from the root directory is:

$ tinygo build -o target/goenv.wasm -target=wasi examples/goenv/goenv.go

Before Go 1.21, TinyGo has been the only way to target WASI, but this support has been part of the standard toolchain for a few releases now.

Writing modules in Rust

Rust is another language that has good support for WASM and WASI in the build system. After adding the wasm32-wasi target with rustup, it's as simple as passing the target name to cargo:

$ cargo build --target wasm32-wasi --release

The code is straightforward, similarly to the Go version:

use std::env;

fn main() {
    println!("rustenv environment:");

    for (key, value) in env::vars() {
        println!("  {key}: {value}");
    }
}

Writing modules in WebAssembly Text (WAT)

As we've seen, compiling Go and Rust code to WASM is fairly easy; looking for a challenge, let's write a module in WAT! As I've written before, I enjoy writing directly in WAT; it's educational, and produces remarkably compact binaries.

The "educational" aspect quickly becomes apparent when thinking about our task. How exactly am I supposed to write to stdout or read environment variables using WASM? This is where WASI comes in. WASI defines both an API and ABI, both of which will be visible in our sample. The following shows some code snippets with explanations; for the full code check out the sample repository.

First, I want to show how output to stdout is done; we start by importing the fd_write WASI system call:

(import "wasi_snapshot_preview1" "fd_write" (func $fd_write (param i32 i32 i32 i32) (result i32)))

Apparently, it has four i32 parameters and returns an i32; what do all of these mean? Unfortunately, WASI documentation could use a lot of work; the resources I found useful are [1]:

With this in hand, I was able to concoct a useful println equivalent in WAT that uses fd_write under the hood:

;; println prints a string to stdout using WASI.
;; It takes the string's address and length as parameters.
(func $println (param $strptr i32) (param $len i32)
    ;; Print the string pointed to by $strptr first.
    ;;   fd=1
    ;;   data vector with the pointer and length
    (i32.store (global.get $datavec_addr) (local.get $strptr))
    (i32.store (global.get $datavec_len) (local.get $len))
    (call $fd_write
        (i32.const 1)
        (global.get $datavec_addr)
        (i32.const 1)
        (global.get $fdwrite_ret)
    )
    drop

    ;; Print out a newline.
    (i32.store (global.get $datavec_addr) (i32.const 850))
    (i32.store (global.get $datavec_len) (i32.const 1))
    (call $fd_write
        (i32.const 1)
        (global.get $datavec_addr)
        (i32.const 1)
        (global.get $fdwrite_ret)
    )
    drop
)

This uses some globals that you'll have to look up in the full code sample if you're interested. Here's another helper function that prints out a zero-terminated string to stdout:

;; show_env emits a single env var pair to stdout. envptr points to it,
;; and it's 0-terminated.
(func $show_env (param $envptr i32)
    (local $i i32)
    (local.set $i (i32.const 0))

    ;; for i = 0; envptr[i] != 0; i++
    (loop $count_loop (block $break_count_loop
        (i32.eqz (i32.load8_u (i32.add (local.get $envptr) (local.get $i))))
        br_if $break_count_loop

        (local.set $i (i32.add (local.get $i) (i32.const 1)))
        br $count_loop
    ))

    (call $println (local.get $envptr) (local.get $i))
)

The fun part about writing assembly is that there are no abstractions. Everything is out in the open. You know how strings are typically represented using either zero termination (like in C) or a (start, len) pair? In manual WAT code that uses WASI we have the pleasure of using both approaches in the same program :-)

Finally, our main function:

(func $main (export "_start")
    (local $i i32)
    (local $num_of_envs i32)
    (local $next_env_ptr i32)

    (call $log_string (i32.const 750) (i32.const 19))

    ;; Find out the number of env vars.
    (call $environ_sizes_get (global.get $env_count) (global.get $env_len))
    drop

    ;; Get the env vars themselves into memory.
    (call $environ_get (global.get $env_ptrs) (global.get $env_buf))
    drop

    ;; Print out the preamble
    (call $println (i32.const 800) (i32.const 19))

    ;; for i = 0; i != *env_count; i++
    ;;   show env var i
    (local.set $num_of_envs (i32.load (global.get $env_count)))
    (local.set $i (i32.const 0))
    (loop $envvar_loop (block $break_envvar_loop
        (i32.eq (local.get $i) (local.get $num_of_envs))
        (br_if $break_envvar_loop)

        ;; next_env_ptr <- env_ptrs[i*4]
        (local.set
            $next_env_ptr
            (i32.load (i32.add  (global.get $env_ptrs)
                                (i32.mul (local.get $i) (i32.const 4)))))

        ;; print out this env var
        (call $show_env (local.get $next_env_ptr))

        (local.set $i (i32.add (local.get $i) (i32.const 1)))
        (br $envvar_loop)
    ))
)

We can now compile this WAT code into a FAAS module and re-run the server:

$ wat2wasm examples/watenv.wat -o target/watenv.wasm
$ go run .

Let's try it:

$ curl "localhost:8080/watenv?foo=bar&id=1234"
watenv environment:
http_host=localhost:8080
http_query=foo=bar&id=1234
remote_addr=127.0.0.1:43868
http_path=/watenv
http_method=GET

WASI: API and ABI

I've mentioned the WASI API and ABI earlier; now it's a good time to explain what that means. An API is a set of functions that programs using WASI have access to; one can think of it as a standard library of sorts. Go programmers have access to the fmt package and the Println function within it. Programs targeting WASI have access to the fd_write system call in the wasi_snapshow_preview1 module, and so on. The API of fd_write also defines how this function takes parameters and what it returns. Our sample uses three WASI functions: fd_write, environ_sizes_get and environ_get.

An ABI is a little bit less familiar to most programmers; it's the run-time contract between a program and its environment. The WASI ABI is currently unstable and is described here. In our program, the ABI manifests in two ways:

The main entry point we export is the _start function. This is automatically called by a WASI-supporting host after setup.
Our WASM code exports its linear memory to the host with (memory (export "memory") 1). Since WASI APIs require passing pointers to memory, both the host and the WASM module need a shared understanding of how to access this memory.

Naturally, both the Go and Rust implementations of FAAS modules comply to the WASI API and ABI, but this is hidden by the compiler from programmers. In the Go program, for example, all we need to do is write a main function as usual and therein emit to stdout using Println. The Go compiler will properly export _start and memory:

$ wasm-objdump -x target/goenv.wasm

... snip

Export[2]:
 - func[1028] <_rt0_wasm_wasip1> -> "_start"
 - memory[0] -> "memory"

... snip

And will properly hook things up to call our code from _start, etc.

WASI and plugins

The FAAS server presented in this post is clearly an example of developing plugins using WASM and WASI. This is an emerging and exciting area in programming and lots of progress is being made on multiple fronts. Right now, WASI modules are limited to interacting with the environment via means like environment variables and stdin/stdout; while this is fine for interacting with the outside world, for host-to-module communication it's not amazing, in my experience. Therefore the WASM standards committee is working of further improvements to WASI that may include sockets and other means of passing data between hosts and modules.

In the meantime, projects are making do with what they have. For example, the sqlc Go package supports WASM plugins. The communication with plugins happens as follows: the host encodes a command into a protobuf and emits it to the plugin's stdin; it then reads the plugin's stdout for a protobuf-encoded response.

Other projects are taking more maverick approaches; for example, the Envoy proxy supports WASM plugins by defining a custom API and ABI between the host and WASM modules. I'll probably write more about this in a later post.

Other resources

Here are some additional resources on the same topic as this post:

Building a WebAssembly-powered serverless platform - the post that inspired this one. Other WASM-related articles from that blog are also recommended.
wagi - a more featured approch with the same idea, implemented in Rust.
WCGI - an even more CGI-conformant approach.

[1]	The main Go compiler supported WASM for a while, but only as a target for browsers & JS (I wrote a bit about it recently). The news in 1.21 is the `GOOS=wasip1` support, which hooks up Go's internal syscalls to WASI APIs, sets up the ABI etc.

[2]	These resources tend to disappear and move around without notice; if you notice a dead link, please drop me a note and I'll look for a more up-to-date source.

How I went about learning Rust

2022-07-15T20:23:00-07:00

Rust has been on my radar for a long time now, and about a year ago I finally began allocating some time every week to learning it.

In this post I'll provide details on the learning path I've followed for Rust, in the hope that this may prove useful for others. You'll note that this isn't exactly a "Learn X in 24 hours" kind of journey, as it stretches over a long time period and covers lots of different materials. YMMV.

This learning journey alternates consuming information (mostly books) with coding activities (some guided by books, some not). In my experience, both are critical for really learning a language or any other tool deeply, and the combination of the two is powerful.

Consuming information

These are the Rust-specific books I've read, roughly in chronological order. There are additional books in the "writing code" section, but those are not Rust-specific.

Programming Rust (link to review) is the first book I read for an initial introduction to the language.

The Rust Programming Language - I haven't read this one cover to cover, but have skimmed major parts of it, implementing the projects it lists. I've also used it a bunch as a reference since it usually comes up high in Google rankings; it's a better reference than Programming Rust, IMHO.

Rust in Action (link to review).

Rust for Rustaceans (link to review).

In addition to these books, I've watched a few of John Gjengset's videos on YouTube. These are very interesting, and I'll probably watch more to learn about specific advanced concepts of Rust, as needed.

Finally, over this year I've read a bunch of blog posts about Rust, mostly stuff that made it to the front pages of HN or r/rust from time to time.

Writing code

rustlings - small exercises for reading and writing snippets of Rust code. Nice concept, though not very comprehensive. It's still quite useful when just getting started.

Advent of Code - I love AOC, and the 2021 edition was a great opportunity to practice some Rust. I've implemented the solutions for days 1-18 in Rust, for a total of ~2 KLOC. Solving AOC problems is a terrific way to learn and practice programming languages. I'm very likely to tackle AOC 2022 in Rust again, to keep my skills sharp.

The Ray Tracer Challenge book (link to review). This book is language agnostic, and I picked Rust to implement a simple ray tracer. Another great way to practice a new language! Almost 4 KLOC for this one.

Crafting Interpreters (link to review). Reimplemented the majority of the clox compiler and VM in Rust, ~2 KLOC.

In addition to these, I've done quite a bit of experimentation, writing code to explore various areas of Rust. Some of this was mentioned in blog posts, but most wasn't. I estimate there's a total of ~6 KLOC of Rust there.

Overall, over these past months I've written about 14 KLOC of Rust in the process of learning and practicing the language.

What's next

I wouldn't call myself a Rust expert, but over this last year I believe I've developed a good understanding and intuition for the language, to the extent that I would feel confident using it for a production project.

Overall, I like Rust and I'm happy to have added it to my toolbox. It has issues, and I certainly have opinions about it, but that will have to wait for a separate post :-)

Rewriting the lexer benchmark in Rust

2022-05-30T05:53:00-07:00

I've reimplemented my lexical analyzer (lexer) for the the TableGen language in Python, JS and Go so far. The latest post in the series discussed how several years of new Go versions improved my lexer's performance by roughly 2x, and how several additional optimizations won another 37%; the final result is a lexer that churns through 1 MiB of source code in just 5.6 milliseconds.

Since I've also been playing with Rust recently, I thought it would be interesting to re-implement this lexer once again, this time in Rust. Rust is certainly the lowest-level language among those I've used so far for this task, so I expect to see some top-notch performance.

The full code for this post is available on GitHub.

Designing an API

I find that Rust's strict ownership rules makes one think carefully about API design from very early on. If we want to create a lexer and pass it a string as input - who owns the string? In Rust, the answer to this question is not an implicit contract (like it always is in C and sometimes in C++), and it cannot be deferred to the runtime either (like one would do in Python, JS or Go). The answer has to be explicitly encoded into the types of the program.

Since one of the goals of this series of posts is performance, I decided to go with a zero-copy API at first, where the user of the lexer owns the input string; the lexer, in turn, returns tokens that contain references into this input string - so the user ends up owning these too. Rust's lifetime specifiers make this fairly natural; here's the type:

pub struct Lexer<'source> {
  input: &'source str,
  iter: Peekable<CharIndices<'source>>,

  // c is the last char taken from iter, and ci is its offset in the input.
  c: char,
  ci: usize,

  // error is true iff the lexer encountered and error.
  error: bool,
}

Ignore the fields for now, focusing just on the struct definition. This is the constructor:

impl<'source> Lexer<'source> {
    pub fn new(input: &'source str) -> Self {
      // ...
    }
}

The 'source lifetime specifier is used to explicitly annotate the lifetime of the input string slice, since we want to store it in our Lexer and refer to it later. Once a Lexer is created, the user obtains new tokens by calling next_token:

pub fn next_token(&mut self) -> Token<'source> {
  // ...
}

Note that the returned Token also has the same lifetime annotation. Here's how the type is defined:

#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Token<'source> {
    pub value: TokenValue<'source>,
    pub pos: usize,
}

#[derive(Debug, PartialEq, Clone, Copy)]
pub enum TokenValue<'source> {
    EOF,
    Error,

    Plus,
    Minus,
    Multiply,
    Divide,
    Period,
    Backslash,
    Colon,
    Percent,
    Pipe,
    Exclamation,
    Question,
    Pound,
    Ampersand,
    Semi,
    Comma,
    LeftParen,
    RightParen,
    LeftAng,
    RightAng,
    LeftBrace,
    RightBrace,
    LeftBracket,
    RightBracket,
    Equals,
    Comment(&'source str),
    Identifier(&'source str),
    Number(&'source str),
    Quote(&'source str),
}

Now it should be clear how these lifetimes are tied together. Some tokens hold slices into the user's input, and this is encoded in the explicit lifetimes. The signature of the next_token method says "we return tokens with a lifetime that's tied to the lifetime of the input passed into the constructor".

We can also provide a more natural iteration API for the Lexer by implementing the Iterator trait:

// Lexer is an Iterator; it returns tokens until EOF is encountered, when it
// returns None (the EOF token itself is not returned). Note that errors are
// still returned as tokens with TokenValue::Error.
impl<'source> Iterator for Lexer<'source> {
    type Item = Token<'source>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.error {
            // If an error has already been set before we invoke next_token,
            // it means we've already returned TokenValue::Error once and now
            // we should terminate the iteration.
            return None;
        }

        let tok = self.next_token();
        if tok.value == TokenValue::EOF {
            None
        } else {
            Some(tok)
        }
    }
}

The iterator implementation makes it possible to integrate the lexer with the rest of Rust very elegantly; for example, to obtain all the tokens in a given input as a vector:

pub fn tokenize_all_collect<'source>(data: &'source str) -> Vec<Token<'source>> {
    let lex = Lexer::new(&data);
    lex.collect()
}

Implementation

Generally, the implementation of the lexer in Rust follows the same approach used by my previous hand-written lexers in this series. I'd like to highlight a couple of Rust-specific aspects here.

Our lexer fully supports Unicode, so I decided to use Rust's string iterator support to obtain chars from &str. Rust provides a helpful iterator called CharIndices, which yields chars along with their position in the input - we need the position to report token locations. Furthermore, since our lexer requires a bit of look-ahead [1], the iterator is wrapped in Peekable, which provides the peek method. As we've seen above already, the final definition is:

iter: Peekable<CharIndices<'source>>,

With this in mind, the code of the lexer should be very readable, even for someone not too familiar with Rust.

Another note is how sub-string extraction happens when tokens are returned. As an example, let's look at the scan_number method which is invoked when the lexer's current character is a digit:

fn scan_number(&mut self) -> Token<'source> {
    let startpos = self.ci;
    while self.c.is_digit(10) {
        self.scan_char();
    }
    Token {
        value: TokenValue::Number(&self.input[startpos..self.ci]),
        pos: startpos,
    }
}

Recall that the Number variant of the TokenValue enum is defined as Number(&'source str) - it contains a reference to the input source string. In the code for scan_number, we see how this is actually implemented, by sub-slicing the input slice. This creates a slice with the same lifetime as the input slice (which is encoded in the lifetimes of the types). As in Go, sub-slicing is a very cheap operation in Rust (no heap allocation).

Performance

The performance of this Rust-implemented lexer is very good! I ran it on the same large TableGen file I've used for all the benchmarking, and it finishes tokenizing it in just 3.7 ms; this is about 33% faster than my fastest Go version from the previous post.

Profiling Rust is quite a bit trickier than Go, but I managed to cobble together enough magic flags and perf invocations to ascertain that next_token indeed takes the bulk of the time; this is good - the time is being spent where it's supposed to be spent.

Trying a variant with allocations

As described above, the API for this Rust lexer was designed to be zero-copy. Since my Go lexers experimented with different approaches, I've decided to see how a different API in Rust would look.

There are two aspects to ownership we can change: taking ownership of the input string, and/or returning owned Strings in the tokens.

For taking ownership of the input string, our constructor would have to look like:

pub fn new(input: String) -> Self

Does it make sense for a lexer to own its input? This isn't clear, and in reality it turned out to be tricky to implement due to lifetime issues. Rust is very unhappy when a struct field is a reference to another field of the same struct, because there is no safe way to move instances of such structs. In our case, the iterator (a struct field) needs a reference to the string (another field in the same struct), and this is doesn't pass the Rust compiler's scrutiny. I wanted to be able to implement this without actively fooling the compiler by using opaque indices or unsafe.

When the language fights you this hard, it may be a good sign that the design is wrong - it's better to leave the ownership to the code that creates a Lexer.

Ownership of returned tokens was easier to set up. The code for this is available in the owning.rs file of the repository. In this variant, the constructor doesn't change, but tokens are defined differently:

#[derive(Debug, PartialEq, Clone)]
pub struct Token {
    pub value: TokenValue,
    pub pos: usize,
}

#[derive(Debug, PartialEq, Clone)]
pub enum TokenValue {
    // .. types
    Comment(String),
    Identifier(String),
    Number(String),
    Quote(String),
}

Note that variants like Identifier now hold an owning String instead of a string reference. Therefore, lifetime annotations are no longer necessary on Token.

The scanning code now allocates a new String and reads characters into it from the iterator. We've seen previously how scan_number looks; here it is again, for the owning token variant (with some helper methods):

fn scan_number(&mut self) -> Token {
    let startpos = self.ci;
    Token {
        value: TokenValue::Number(self.scan_while_true(|c| c.is_digit(10))),
        pos: startpos,
    }
}

// Helper to scan chars while `pred(c)` returns true, into the given `s`.
fn scan_while_true_into<F>(&mut self, s: &mut String, pred: F)
where
    F: Fn(char) -> bool,
{
    while pred(self.c) {
        s.push(self.c);
        self.scan_char();
    }
}

// Helper to scan chars while `pred(c)` returns true and return all scanned
// chars in a new String.
fn scan_while_true<F>(&mut self, pred: F) -> String
where
    F: Fn(char) -> bool,
{
    let mut s = String::with_capacity(8);
    self.scan_while_true_into(&mut s, pred);
    s
}

Performance of this variant

When I first tried this variant, its performance wasn't great at all! It was about 30% slower than the string-copying version in Go. I think this has to do with the slight difference in approach - in Go, my lexer figures out the token boundaries using numerical indices and then converts a []byte into string in one fell swoop. The Rust version fetches chars one by one from an iterator and writes them into a String.

In particular, since Rust's String is dynamically allocated, this may incur reallocations (depending on how much is allocated initially).

So my solution was to create these strings with_capacity - as you can see in the previous code sample. This cut down the execution time to be roughly equal to Go's version where strings are copied.

Which API is better?

IMHO there's little doubt that the original "slice" API is better, for multiple reasons:

Performance: the results speak for themselves - the slice API is zero-copy and incurs no additional heap allocations. Like in Go, it deals in slice headers, which are just pointers to parts of a string.
The ownership story is very clear, symmetrical and explicitly documented with lifetime annotations. The 'source lifetime controls everything: it's the lifetime of the &str passed into the lexer's constructor, and it's the lifetime of the tokens. The symmetry feels important - the code that creates a lexer controls the lifetime of the input, as well as the output.
In general, there's a known good practice in API design which is not to force allocations on users, if possible. This is well articulated in this blog post for Go, but it applies equally well in Rust. In our slice API, the user can always call to_owned on the returned slice, if they so wish. But why do it for them? Why should we assume the users want us to return them an owned string? Returning a slice provides more flexibility in using the API.

Performance improvement: avoiding `Peekable`

After this post was published, the reader Utkarsh Kukreti mentioned that using Peekable is not the most optimal approach, since its next is slower than the underlying iterator's next, and calling it is on the hot path. Instead of using Peekable, we could just clone the iterator when we see a / and invoke next on the clone; iterators are small so cloning them is cheap.

Here's a patch with this change. Applying it to my lexer makes it ~11% faster on my machine, finishing the benchmark in about 3.4 ms!

To be honest, this is a little bit surprising and disappointing to me, since I was hoping that Peekable would fall into the zero abstraction promise of Rust. Perhaps this is something that can fixed in the future.

[1]	For properly tokenizing comments; when the lexer sees `/`, it needs to peek forward to see if this is the beginning of a `//`, or else the division operator. This is a common use case in lexers, especially when tokenizing multi-character operators (like `>=`).