rust-async-tokio

name	rust-async-tokio
description	Tokio,

Async Rust in Rocky

Rocky is async end-to-end. Tokio is the only runtime; there's a single #[tokio::main] entrypoint at engine/rocky/src/main.rs:471. Adapter I/O (Databricks REST, Snowflake REST, Fivetran REST, Valkey, webhooks) all runs on the same runtime.

Runtime setup

• Workspace dep — tokio = { version = "1", features = ["full"] } (from engine/Cargo.toml). The "full" feature is deliberate: we use macros, rt-multi-thread, time, fs, net, sync, and process across the workspace, and pinning a narrower feature set per crate creates friction when async code moves between crates.
• Entrypoint — engine/rocky/src/main.rs:471 has #[tokio::main] on async fn main() -> anyhow::Result<()>. Library crates never spawn their own runtime — they take &self on async methods and trust the binary to drive the reactor.
• Library crates depend on the workspace tokio re-export. Don't cargo add tokio in a sub-crate; inherit from [workspace.dependencies] in engine/Cargo.toml.

Async trait rule

Every async trait uses #[async_trait] from the async-trait crate (async-trait = "0.1" in [workspace.dependencies]). This is a hard rule from engine/CLAUDE.md.

// DO — matches the pattern used throughout rocky-adapter-sdk
#[async_trait::async_trait]
pub trait WarehouseAdapter: Send + Sync {
    async fn execute(&self, sql: &str) -> Result<StatementResult, AdapterError>;
    async fn describe_table(&self, qname: &QualifiedName) -> Result<TableSchema, AdapterError>;
}

// DON'T — native async-in-traits still has object-safety gaps in 2024 edition
pub trait WarehouseAdapter: Send + Sync {
    async fn execute(&self, sql: &str) -> Result<StatementResult, AdapterError>;
}

Concrete examples: crates/rocky-adapter-sdk/src/traits.rs has four #[async_trait] trait definitions at lines 219, 333, 346, 372 — those are the shapes adapters must match.

When in doubt, read how rocky-databricks or rocky-fivetran implement the trait — both are full end-to-end examples.

Backpressure: the AIMD pattern

Rocky uses adaptive concurrency on remote calls that can be rate-limited. The canonical implementation is crates/rocky-databricks/src/throttle.rs::AdaptiveThrottle:

• Slow start — below half of max, increase by 2 every increase_interval successes.
• Congestion avoidance — above half of max, increase by 1.
• Multiplicative decrease — on on_rate_limit() (triggered by 429 / 503 / "TEMPORARILY_UNAVAILABLE"), halve the current concurrency down to min_concurrency.
• Clamped — never goes below min_concurrency (≥ 1) or above max_concurrency.
• Atomic shared state — the inner counters are AtomicUsize + AtomicU64 inside an Arc<ThrottleInner>, so .clone() shares state across tasks. No locks on the hot path.

When to reach for this pattern: any new adapter that talks to a rate-limited remote API. Don't re-invent it; either reuse AdaptiveThrottle directly (if it's a Databricks-family API) or copy its shape. The tests at crates/rocky-databricks/src/throttle.rs:122 cover the invariants you'd want to preserve (starts at max, halves on rate limit, never below min, never above max, clone shares state).

How to wire it:

1. Create one AdaptiveThrottle per warehouse/endpoint at adapter construction.
2. Acquire a permit via a tokio::sync::Semaphore whose available_permits() is refreshed from throttle.current() before each batch.
3. On the result of each request:
   • 2xx → throttle.on_success()
   • 429 / 503 → throttle.on_rate_limit() and retry after a small delay
   • anything else → neither (it's a user error, not a concurrency signal)

Structured logging: tracing, not println!

Hard rule from engine/CLAUDE.md: use tracing, not println! or eprintln!. The subscriber is initialized in crates/rocky-observe/src/tracing_setup.rs and emits structured JSON lines when RUST_LOG / tracing_subscriber::EnvFilter is set.

use tracing::{info, warn, error, debug};

// DO — structured fields, no interpolation in the message
info!(
    connector_id = %connector.id,
    table_count = tables.len(),
    "discover completed"
);

warn!(
    from = old,
    to = new,
    "adaptive throttle: rate limit detected, reducing concurrency"
);

// DON'T — stringify values into the message
info!("discover completed: connector={}, tables={}", connector.id, tables.len());

• %value → uses Display
• ?value → uses Debug
• Bare field = literal → treats as JSON literal

warn! is the level rocky-databricks/src/throttle.rs uses for rate-limit detection — follow that pattern for any adaptive-concurrency event you add.

Timeouts

Any network call needs a timeout. Use tokio::time::timeout:

use std::time::Duration;
use tokio::time::timeout;

let result = timeout(Duration::from_secs(30), client.execute(sql))
    .await
    .context("databricks execute timed out after 30s")??;
//                                                      ^^ one ? for Elapsed → anyhow, one ? for the inner Result

Rule of thumb: if a function makes an HTTP/SQL/network call and does not have an outer timeout, it's a bug waiting to happen. The anyhow .context layer is what surfaces the timeout reason into the Dagster event log — don't skip it.

Cancellation and select!

tokio::select! is the right tool when you want to race two futures (e.g. "wait for this statement to finish or for the user to ctrl-C"). Guidelines:

• Every branch must be cancel-safe. Read the tokio::select! docs on which tokio primitives are cancel-safe and which aren't. (Hint: AsyncRead::read_buf is not cancel-safe.)
• Prefer tokio::sync::oneshot for "signal this future to stop" over manual flags.
• Don't hold a MutexGuard across an .await in any select! branch.

Spawn vs. join

Situation	Use
Fan out N tasks and wait for all	futures::future::try_join_all
Fan out N tasks, take first result	futures::future::select_ok
Fire-and-forget background worker	tokio::spawn(...) — but the spawned future must own its data ('static), and you must handle its JoinHandle if it can fail
Parallel CPU-bound chunks	tokio::task::spawn_blocking — not spawn. Blocking the reactor starves other adapters.

duckdb and sqlparser operations are CPU-bound and should be wrapped in spawn_blocking if they're called from an async context on a hot path.

What NOT to do

Anti-pattern	Why
std::thread::spawn in async code	Bypasses the runtime; task never wakes correctly. Use tokio::spawn or spawn_blocking.
std::sync::Mutex held across .await	Can deadlock the reactor. Use tokio::sync::Mutex or (better) restructure to avoid holding the lock.
Per-crate #[tokio::main] or nested runtimes	There's exactly one runtime, driven by engine/rocky/src/main.rs. Libraries don't own the runtime.
futures::executor::block_on inside async code	Nested block_on will panic under tokio::main.
async fn foo(...) -> Box<dyn Future<...>>	Use #[async_trait] for trait methods or impl Future for free functions.
Spinning on throttle.current() in a busy loop	Drive concurrency off a Semaphore; the throttle is a signal, not a gate.

Related skills

• rust-error-handling — async errors follow the same two-tier model (library → thiserror, CLI → anyhow with .context).
• rust-clippy-triage — the async-family clippy lints (e.g. clippy::unused_async, clippy::await_holding_lock) fire in this surface area.
• rust-unsafe — duckdb calls are sync/FFI and need spawn_blocking, not direct await.