How to Rewrite a Collector in Rust

Bee is a monitoring data collection engine, and it is the second half of a story that starts with the Kotlin collector I wrote about previously. The interesting part is not that it exists in Rust. The interesting part is that it was never a port.

What this post is about. It is about the Rust rewrite of that collector: how the collector talks to a separate program that owns the database connections, how a metric is written down and shipped, and how one concurrency bug was answered three different ways in three commits.

What you need to know first. Four words do most of the explaining. A schema service is the Java program that owns the database connections. A datasource is a named function you can call on a remote machine. A task is one metric you want collected. A wire protocol is the agreed format the two sides talk in.

The last one is where the story is: the same bug — “how many scripts can run at once?” — got three different answers in three commits.

The situation

The product is the same problem as before: pull metrics from hundreds of databases and hosts on a schedule, ship them to Riemann — the time-series database the fleet reports to — and never let one slow instance stall the fleet. The README states the motivation for the rewrite in one sentence:

Bee 是一组小巧，异步的网络监控数据采集服务…… Bee 本身采用 Rust 语言开发，所以其内存占用和 CPU 等系统资源不会占用太大。

(“Bee is a set of small, asynchronous network monitoring collection services… Bee is developed in Rust, so its memory footprint and CPU usage do not take up too much.”)

That sentence is the whole justification, and it is a deployment argument: this collector runs 1:1 on every monitored host and database. A JVM means a 40 MB jar, a heap to tune, and GC pauses (garbage collection stops every thread while it cleans up) on a machine already running the thing being monitored.

How the Rust collector is put together

Two implementations, born together

The repository’s second commit — feat: first commit, 2020-10-20 — is 174 files and 12,391 insertions, and it contains the Rust workspace and the Gradle/Java modules at the same time (Gradle is the standard Java build tool). There was no Java program that was later translated; there were two programs from day one.

That was a deliberate strategy, and the git history shows it being executed:

  • c6a59f2 fix: thread pool (2020-12-23) — the Rust collector and the Java SchemaService: the same bug fixed in both implementations in one commit.
  • 917fbb1 fix: thread pool (2020-12-24) — the Java BeeServer plus the Lua task scripts.
  • 2b3449e fix: postgresql (2021-03-09) — folds five *Schema.java classes into *Connector.java.
  • 735ebe9 fix: suport AES (2021-03-10) — moves password decryption from Java into Rust.

The Java side was never bulk-deleted; it was trimmed file by file as the Rust side took over. The two implementations share a wire protocol, so they can be A/B’d — run side by side and compared — and moved over feature by feature. The Java database and host schemas still exist at HEAD (the latest commit), and the Rust collector talks to them over the same protocol.

A custom binary protocol, because JSON would not do

The wire format between collector and schema service is a hand-rolled binary protocol on top of tokio-util‘s FramedRead/FramedWrite (tokio is Rust’s async runtime — the library that runs async code on a pool of threads — and tokio-util is its bag of helpers). The two helpers read and write length-prefixed frames: chunks of bytes with a header saying how long each one is.

Each packet is a small header followed by a payload of basic values only — nil, integer, number, boolean, string, bytes. The protocol document (doc/src/protocol/data.md) rules out arrays and key-value structures, so anything complicated travels as a byte stream instead.

The protocol has around sixty stable error codes, documented in doc/src/error.md. The reason: “the schema service went away” and “your script threw a Lua error” need to be distinguishable by a collector running on someone else’s database server.

Connecting is where the timeout lives. The collector reads connection_timeout off the instance, defaulting to 10 seconds, adds a second of slack, and wraps the connect in a hard deadline. The + 1 is deliberate: the deadline is one second later than the configured timeout, so the collector’s own timeout fires first and the error names the real problem. set_nodelay(true) turns off Nagle’s algorithm — the trick of merging small writes into fewer packets.

A datasource is just an async function

The plugin model is a trait — Rust’s way of writing “an interface”: a list of methods an implementor must provide. The #[datasource] attribute macro writes the boilerplate for you:

1
2
3
4
5
6
#[async_trait]
pub trait DataSource: Send + Sync {
    fn name(&self) -> &str;
    fn args(&self) -> Columns;
    async fn collect(&self, request: &mut Request) -> Result<()>;
}

libcore/rust/src/datasource.rs:4. Send + Sync are Rust’s marker traits: Send means the value may be moved to another thread, Sync that it may be shared between threads. #[async_trait] lets a trait method be async without boxing every future. The #[datasource] macro derives the input Columns — the named, typed parameters the function takes — from the signature and injects the Promise, Request and Arc<Register> parameters, so the wire representation is generated rather than hand-maintained. A datasource is a named, typed, remotely callable function: columns in, rows out.

A task is a script file

There is no task table. A task is a file in scripts/task/:

1
2
3
4
5
6
pub fn from_file<T: AsRef<Path>>(path: T) -> Result<Task> {
    let id = path.as_ref().file_stem()...;
    let script_type = path.as_ref().extension()...;
    Task { id: id.to_string(), script: read_to_string(path)?,
        timeout: Duration::from_secs(10), ... }
}

bee/src/features/task.rs:58. The id is the filename stem — mysql_qps.lua becomes the task mysql_qps — and the extension decides the action. The ... stands for the error handling that turns the two OsStr values (the raw path components) into strings.

The repository carries 91 task files — 4,313 lines of Lua and SQL, plus 3,057 lines of Perl for the host collection — so adding a metric means adding a file, not a release.

How it works

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
┌──────────────────────────────────────────────────────────────┐
│ bee (collector)                                              │
│                                                              │
│  Schedule ──► Action::SQL ─────────────────┐                │
│      │        Action::Lua ──┐              │                │
│      │                      ▼              ▼                │
│      │              ┌─────────────┐   ┌──────────────┐      │
│      │              │  call rt    │   │  SchemaClient│      │
│      │              │ (tokio MT)  │   │  (wire proto)│      │
│      │              └──────┬──────┘   └──────┬───────┘      │
│      │                     │                 │              │
│      │              ┌──────▼──────┐          ▼              │
│      │              │  Lua VM     │   ┌──────────────┐      │
│      │              │ (mlua)      │   │  agent /     │      │
│      │              └──────┬──────┘   │  db schema   │      │
│      │                     │          │  (java)      │      │
│      │              ┌──────▼──────┐   └──────────────┘      │
│      │              │  State mpsc │─────────┐                │
│      │              │  rows       │         ▼                │
│      │              └─────────────┘   ┌──────────┐          │
│      └─────────────────────────────────►│ Riemann  │          │
│                                        └──────────┘          │
└──────────────────────────────────────────────────────────────┘

Everything lives inside one box: the bee collector process; the arrows show work flowing between the boxes. The schedule on the left fires a task that is either SQL (run against a database) or Lua (a small script). SQL tasks go through SchemaClient, which speaks the binary wire protocol to a Java database schema or the native Rust agent on the host. Lua tasks go into a Lua VM — the interpreter that runs the script — owned by a tokio runtime; “call rt (tokio MT)” is the multi-threaded runtime hosting those interpreters. Script rows flow down the State mpsc box — an mpsc channel is a many-producer, single-consumer queue, so the collector keeps reading while the script keeps writing — and the bottom arrow carries the metrics to Riemann.

Two action types implement the same trait. SQLAction pushes the script into the request arguments and calls query on the remote schema. LuaAction builds an mlua VM, injects the globals — values the script can read by name: logger, task, instance, session, request, cache — runs the common scripts, then the task’s own script. Rows stream back as a small state machine — Ready(columns), Process(row), Err, Ok — and the first row establishes the schema: type information travels once, before the first row, later rows are typed by position.

The configuration is dotted paths

bee.toml — TOML is a simple key = value config format — is read through a tiny helper that walks dotted keys. Call sites read config.get_value("collector.call_max_threads"). The helper splits the key at the . and recurses into the nested table. It is not a general TOML binding library, just enough to keep the config file flat and readable.

The build is optimised for size

The release profile in Cargo.toml sets opt-level = 'z' (optimise for size, not speed), lto = true (link-time optimisation: the linker can inline across crate boundaries — a crate is one compiled Rust library or binary) and codegen-units = 1 (compile everything as one unit). For a binary deployed on every monitored host, artefact size matters more than peak throughput. Cross-compilation targets included x86_64-unknown-linux-musl (a fully static Linux binary) and windows-gnu, so the collector shipped as one static file with no runtime dependency.

The thread pool: three attempts

The git history contains three commits with the same message — fix: thread pool — and each is a different answer to the same question.

Attempt one, in Java: the schema service ran connections on Netty event loops (Netty is the Java network library it was built on) and handed the actual work to Executors.newWorkStealingPool(config.getWorkThreadSize()) — a pool where idle workers pull tasks from their neighbours’ queues.

Attempt two, in Rust: one fresh OS thread per script execution, via std::thread::Builder::new().name(uid).spawn(...). Every Lua run got its own thread, never reused.

Attempt three, which is what the last commit looks like, replaced that with a shared rayon pool (rayon’s threads steal work from each other too) and gave each script its own single-threaded tokio runtime.

Attempt four is still uncommitted in the working tree, and it is the interesting one. The rayon block is commented out and replaced by a tokio LocalSet:

1
2
3
4
5
6
let local = LocalSet::new();
local.spawn_local(async move {
    let rs = exec_lua_script(lua, script, depends).await;
    let outcome = if let Err(err) = rs { Err(Error::from(err)) } else { Ok(()) };
    let _ = signal_tx.send(outcome).await;
});

bee/src/features/collect/mod.rs:259 (the original writes the two branches out longhand, and the commented-out rayon line sits just above this block). A LocalSet is a tokio container for tasks that must stay on the thread they were started on. The motivation is thread affinity: the rayon version had to move a !Send Lua VM — a type Rust has proven must not cross threads — onto a worker thread, which is where Lua::into_static() and Lua::from_static() come from. A LocalSet keeps the script on the caller’s thread, where a !Send VM belongs, and the unsafe impl Send/Sync for Collector can eventually go away. The cost: the rayon pool was what bounded how many scripts ran at once, and a LocalSet runs the script on the caller’s current thread, so the bound must come from the “call” runtime instead.

The lesson is not “use rayon”. It is: in async Rust, “spawn a thread per task” silently reintroduces the exact thread-pool problem you were trying to escape. The Rust attempt-two version had the same unbounded-thread pathology as a naive Java implementation, and it took two commits to notice.

Blocking Lua inside async Rust

The harder problem is that mlua’s Lua is thread-bound and !Send, while the collector wants to be Send + Sync so it can live in a shared Arc — a reference-counted smart pointer you can share between threads. The file answers the compiler with two empty-looking lines, unsafe impl Send for Collector {} and unsafe impl Sync for Collector {}, which tell Rust to accept, without checking, that this type may be moved and shared across threads.

The consequences cascade through the file: a second multi-thread tokio runtime for “call” work, a current-thread runtime per script, async_channel plus LocalSet to await a local task from an async context, and Lua::into_static() / Lua::from_static() to move a VM between threads.

Trade-offs

The Rust collector has no database connectors. The Rust side talks to the Java database schemas over the shared wire protocol. What the port actually replaced was SSH-based host collection: the new native Rust agent runs on the host and streams metrics locally instead of shelling over SSH.

It is not buildable today. All git dependencies point at a Chinese npm-style mirror (github.com.cnpmjs.org/smitea/Rocket.git, forks of mlua, heim, notify, log4rs). bee/src/lib.rs starts with #![feature(decl_macro)], which only compiles on the experimental nightly Rust. The workspace lists libcore/codegen twice. cargo metadata --offline still resolves, so the manifests are valid — but the sources are gone.

There is a real use-after-free in the agent — the program uses memory after handing it back to the allocator. A hand-rolled buffer avoids a Vec (Rust’s normal growable list) allocation per file read, and it frees the memory it just handed out:

1
2
3
4
5
// inside an `unsafe { ... }` block, in new_vec<'a>(size: usize) -> &'a mut [u8]
let layout = Layout::from_size_align_unchecked(size * size_of::<u8>(), size_of::<u8>());
let ptr = std::alloc::alloc(layout);
let buffer = std::slice::from_raw_parts_mut(ptr, size);
std::alloc::dealloc(ptr, layout);   // <-- frees the memory it just handed out

schemas/agent/src/agent/mod.rs:57. The first line builds a raw memory layout without checking it. The next two allocate memory and wrap it in a byte slice. The fourth hands that memory back to the allocator while the slice still points at it — and the function then returns that slice, so any later read touches freed memory. The Vec it was avoiding would have been correct and free. The whole body sits in an unsafe block, and the signature (&'a mut [u8]) promises a slice that outlives the call.

A clock that can panic. SystemTime::duration_since(...).unwrap() on a machine whose clock steps backwards panics the collector.

What I would do differently

Decide the port strategy before writing the second implementation. Dual implementation from commit two is defensible when the wire protocol is stable — and here it worked, because the Java side stayed alive as a reference until the Rust side took over. But it doubles the surface area of every bug fix, and the “same bug fixed in both implementations” commits are a tax paid forever. Better: write the protocol and the Java side, freeze it, then port.

Solve the blocking-Lua problem at the boundary, not with unsafe. The whole unsafe impl Send/Sync for Collector block exists because a Lua VM is thread-bound. The clean answer is a dedicated actor per script — a small task that owns the VM and answers requests over a channel, a queue one task fills and another drains — so the thread-affinity rule lives in the type system, not in a comment.

Keep the size-optimised release profile, drop the rest of the compromises. opt-level = 'z' and lto = true are right for a per-host binary. But the nightly feature gate, the dead mirrors and the duplicated workspace member are neglect, and they are why a five-year-old project cannot be built.

And keep the interesting part. The good decisions came from the domain, not the language: a task is a file, a datasource is a function whose wire format is generated, the collector holds no state, and the protocol has sixty distinguishable error codes because the collector runs on someone else’s database server. Those survive a rewrite; the thread pool did not.

One last thing. In the working tree there is an untracked file, test.metric, dated about a year after the last commit — a sketch of a typed metric-pipeline DSL (a domain-specific language: a small custom language for one job):

1
2
3
4
5
input host;
metric cpu_usage(ctx: @Context, deplay_time: u64) => (host: string, ip: string, idle: f64, ...);
view cpu_usage(ctx: @Context) => (resp: Response);
#[os = 'linux', version = '1.0.0']
cpu_usage_linux(ctx: @Context, deplay_time: u64) as cpu_usage { ... }

input declares a source, metric a measurement and the shape of its output, view how it is presented, and yield hands the values out. It is the ancestor of the pipeline idea in the 2026 bee workspace — the one where a pipeline is a callback node in a DAG (a directed acyclic graph: steps with arrows between them and no loops) and the scripting language is Lua because the people writing the rules are not the people running the cluster. Same idea, five years earlier, in a notation that never shipped. Picking this codebase up again, I would turn test.metric into a real DSL: the pipeline was always the good idea; the thread pool was just the bug I kept fixing.

References

  • tokio — the async runtime, its thread pools and LocalSet
  • mlua — the Lua bindings the task scripts run inside
  • rayon — the work-stealing pool used in the third attempt
  • mlua on docs.rs — why a Lua value cannot cross threads
  • redb — the embedded store used for state in the later bee workspace
  • RocksDB — the storage engine behind that state store
  • RFC 4253 — the protocol the old SSH-based host collection spoke