feat(providers): make the async runtime pluggable (tokio | compio) by x-at-01 · Pull Request #1435 · microsoft/DiskANN

x@01 (x-at-01) · 2026-09-23T08:58:23Z

Problem

diskann, diskann-providers and diskann-inmem hard-bind to tokio: diskann and diskann-providers declare a non-optional tokio dependency (rt, rt-multi-thread). Our production services run on compio, a thread-per-core runtime, and therefore could not consume the async provider surface (diskann-providers::index::diskann_async, the async graph index, and the synchronous wrapped_async wrappers) even though the crates only actually require task spawn + join semantics from the runtime — no tokio I/O, timers, or sync primitives are used in library code.

Solution

Make the async runtime backend selectable via cargo features while keeping upstream behavior identical by default:

New tokio feature (on by default) and new compio feature (optional dependency, default-features = false, features = ["runtime"]). With both enabled, tokio takes precedence; with neither, a compile_error! in the crate root explains the requirement.
A small internal runtime facade in diskann (spawn, JoinHandle, JoinError, JoinSet) replaces the five direct tokio::spawn call sites and the JoinSet/JoinError uses in the graph index. ExecutionContext::wrap_spawn is unchanged.
- The compio backend spawns on the calling thread's current runtime (compio::runtime::spawn), matching compio's thread-per-core execution model (one runtime per worker thread). Its JoinSet joins in FIFO order, which satisfies every call site since all task results are consumed.
diskann-providers' synchronous wrappers (wrapped_async) move from raw tokio::runtime::Runtime/Handle onto a blocking facade. Under compio the facade is stateless: the wrapper stores no runtime (staying Send + Sync, so one index can be shared across worker threads) and every blocking call resolves the current thread's runtime via Runtime::try_current(), panicking with an actionable message when called outside any compio runtime context. The external-handle constructors (new_with_handle, load_with_handle) remain tokio-only, since a compio runtime cannot be entered from a foreign thread.
Member manifests that consume diskann with default-features disabled forward the backend explicitly (features = ["tokio"]), reproducing the previous effective dependency set.

Behavior compatibility

Default build (diskann-providers default features, or plain diskann): feature set and behavior are unchanged — the tokio backend is compiled in exactly as before, all public signatures are identical, and the tokio-form test suites pass unchanged (diskann lib 324 tests, inmem lib 103 tests, providers integration tests).
The only intentional semantic note: under compio, dropping the last handle to a spawned task cancels it (tokio detaches). All current call sites join every spawned handle, so this is not observable.

Verification

Feature matrix: cargo check -p <crate> --all-features and cargo check -p <crate> --no-default-features --features compio --all-targets for diskann, diskann-providers, diskann-inmem — both clean (0 warnings) on stable 1.97.1 (the pinned toolchain) and nightly 1.100.
Tokio form: existing suites pass unchanged; compio form: the full providers lib suite (287 tests) passes with tokio compiled out.
New compio end-to-end tests: full index lifecycle (parallel multi_insert build, save, reload, search) driven under compio::runtime::Runtime::block_on/enter with tokio compiled out; a Send + Sync assertion for the wrapper types under both backends; a should_panic guard covering calls outside a compio runtime context; nested rt.block_on(async { .. }) coverage for the downstream worker-thread shape.
Tokio-only tests are gated behind feature = "tokio".
cargo fmt --check passes for the touched crates; clippy is clean apart from pre-existing test-only unwrap_used warnings.

The nightly-1.100 trait solver also stopped inferring &Arc -> &B deref coercions at generic call sites; the three affected call sites now deref explicitly (&*batch_clone / &*vectors_clone), which both solvers accept.

…or compio The crate hard-depends on tokio (features rt/rt-multi-thread), which excludes thread-per-core runtimes such as compio even though the crate only needs `spawn` + join semantics. - Introduce an optional `compio` dependency (workspace, version 0.19.2, `default-features = false, features = ["runtime"]`) and make `tokio` optional behind a `tokio` feature that is enabled by default, so upstream behavior is unchanged. - Add an internal `runtime` facade with `spawn`, `JoinHandle`, `JoinError` and `JoinSet`: - `tokio` backend delegates to `tokio::spawn` / `tokio::task::JoinSet`. - `compio` backend builds on `compio::runtime::spawn` + `block_on` semantics (thread-per-core, no detached task registry); its `JoinSet` joins in FIFO order, which is sufficient for all call sites since they consume every task result. - With both features enabled, `tokio` wins; with neither, a `compile_error!` in the crate root explains the requirement. - Route the five `tokio::spawn` call sites and the `JoinSet`/`JoinError` uses in the graph index through the facade; the `ExecutionContext:: wrap_spawn` customization point is unchanged. - Gate runtime-dependent tests behind `feature = "tokio"` and add compio end-to-end tests (spawn, JoinSet join order, join-error plumbing) that run under `compio::runtime::Runtime::block_on`. - Workspace: declare compio once and make `diskann` default-features opt-in so member crates forward the backend explicitly (`tokio` + `tracing` reproduces the previous effective defaults). Verification: `cargo check -p diskann --all-features`, `cargo check -p diskann --no-default-features --features compio` (and the same with `--all-targets`) pass warning-free; the tokio-form unit suite passes unchanged, and the new compio tests pass with tokio compiled out.

`wrapped_async` owns or borrows a `tokio::runtime::Runtime`/`Handle` pair, hard-binding the synchronous API to tokio. - Add an internal `runtime` facade: an owned `Runtime` and a cloneable blocking `Handle` with `block_on`. - `tokio` backend preserves the existing multi-thread / current-thread flavors and the external-handle constructors (`new_with_handle`, `load_with_handle`, tokio-only by nature). - `compio` backend provides the same blocking semantics on top of `compio::runtime::Runtime` (single flavor, bound to its creating thread; the wrapper types are `!Send` there and external-handle constructors are compiled out). - Switch `DiskANNIndex` and `PagedSearch` to the facade; all public tokio-form signatures and behavior are unchanged. - Feature wiring: `default = ["tokio", "tracing"]` reproduces the previous effective dependency behavior (tokio via the former non-optional dep, tracing via the former diskann defaults); `compio` forwards `diskann/compio`. When both are enabled, `tokio` wins. - Gate the tokio-only tests behind `feature = "tokio"` and add a compio end-to-end test (build through the parallel `multi_insert` path, save, reload, search) that runs with tokio compiled out. Verification: `cargo check -p diskann-providers --all-features` and `--no-default-features --features compio` (with `--all-targets`) pass warning-free; tokio-form tests pass unchanged; the full compio-form lib suite (279 tests) passes, including the new end-to-end test.

The crate already treated tokio as optional; wire it into an explicit feature and add the compio counterpart. - Features: `default = ["tokio", "tracing"]` reproduces the previous effective behavior (tokio via former diskann defaults); `compio` forwards `diskann/compio`. When both are enabled, `tokio` wins. - Gate the tokio-only smoke test behind `feature = "tokio"` and add a compio twin that drives the provider/index flow under `compio::runtime::Runtime::block_on` with tokio compiled out. Verification: default and `--no-default-features --features compio` forms (with `--all-targets`) pass; the compio smoke test passes.

…fests With `diskann`'s tokio made optional, member crates that consume it (or diskann-providers/diskann-inmem) with default features disabled must name the backend explicitly to keep their previous behavior: - diskann-disk / diskann-label-filter / diskann-tools: add `["tracing", "tokio"]` to their diskann / diskann-providers edges. - diskann-bftree / diskann-benchmark-core: same for their diskann edges.

Resolve conflicts with the post-0.59.0 refactor: - diskann/src/provider.rs: keep the tokio feature gate and adopt the upstream #[expect] style. - diskann-providers wrapped_async: our version already replaces the runtime constructors touched upstream; port the removal of the type_complexity allow and align the facade with #[expect]. - diskann-inmem: take the reworked provider (store extraction) and re-express the backend selection as thin tokio/compio wrappers around a shared `smoke_flow`, keeping the compio end-to-end proof intact.

…; harden against the nightly trait solver Two follow-ups to the runtime-backend work, found in review: 1. rustc's next-generation trait solver (nightly ~1.100) no longer applies deref coercion while unifying `&Arc` against a `&B` parameter, so generic call sites that pass `&Arc` infer `B := Arc` and reject the unimplemented `Arc: Batch` bound. Deref `Arc` to `&B` explicitly at the three affected call sites (`set_chunk` spawned task, `search_and_prune_batch` spawned task and local call). Behavior under the pinned stable toolchain is unchanged. 2. The compio backend of the providers runtime facade stored a `compio::runtime::Runtime` (the `Rc`-shared executor/driver pair) inside the wrapper types, which made the synchronous `DiskANNIndex` wrapper `!Send + !Sync` — breaking its real usage shape: one index shared across worker threads, each thread running its own compio runtime. The facade is now stateless under compio: `Runtime`/`Handle` store a marker, and every blocking call resolves the *current thread's* runtime via `compio::runtime::Runtime::try_current()` (panicking with an actionable message when called outside of any compio runtime context). Task spawning already dispatches through the `diskann::runtime` facade, which spawns on the calling thread's current runtime. The tokio backend (default) is unchanged. Also adds boundary tests (wrapper `Send + Sync` assertion under both backends, `should_panic` guard for calls outside a compio runtime context, nested `rt.block_on(async { .. })` coverage), refreshes the stale compio constructor/load docs, documents the backend difference in `JoinHandle` drop semantics (tokio detaches, compio cancels), and drops the `default-features = false` overrides in member manifests that reference a workspace entry without that field (cargo warns this will become a hard error). Verification: nightly 1.100 and stable 1.97.1, both backend feature combinations, `--all-targets` checks pass; providers compio suite 287 tests and diskann lib 324 tests pass under both toolchains; `cargo fmt --check` and clippy are clean for the touched crates apart from pre-existing test-only `unwrap_used` warnings that predate this change.

Copilot

Copilot review overview

🟡 Changes recommended

Unresolved critical and moderate findings remain in feature forwarding, backend diagnostics, and test-helper gating.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 2 · 2 · 2

Open (6)

Enable diskann's tokio feature for integration tests · New
Gate unused test helpers in compio-only builds · New
Prevent unresolved runtime imports in no-backend builds · New
Prevent unresolved runtime imports in no-backend graph builds · New
Fix broken new_with_handle link in compio-only docs · New
Document differing task behavior between Tokio and compio · New

What changed in this PR

This PR makes DiskANN’s async runtime pluggable between Tokio (default) and Compio.

Changes:

Adds runtime features, dependency forwarding, and runtime facades.
Updates synchronous wrappers for Tokio and Compio.
Gates Tokio-specific tests and adds Compio coverage.

File	Reviewed changes
`diskann/src/test/mod.rs`	Gates Tokio test helpers.
`diskann/src/runtime.rs`	Adds runtime-agnostic task facade.
`diskann/src/provider.rs`	Updates spawning tests.
`diskann/src/lib.rs`	Adds backend validation and module wiring.
`diskann/src/graph/test/provider.rs`	Gates Tokio-dependent tests.
`diskann/src/graph/test/mod.rs`	Gates Tokio graph tests.
`diskann/src/graph/index.rs`	Uses the runtime facade for spawning.
`diskann/src/graph/ext/labeled.rs`	Gates Tokio filtered-search tests.
`diskann/src/flat/test/harness.rs`	Supports Tokio and Compio test runtimes.
`diskann/src/flat/strategy.rs`	Gates Tokio tests.
`diskann/src/flat/index.rs`	Uses the runtime join facade.
`diskann/Cargo.toml`	Adds runtime features and dependencies.
`diskann-tools/Cargo.toml`	Selects Tokio features explicitly.
`diskann-providers/src/storage/index_storage.rs`	Gates Tokio tests.
`diskann-providers/src/runtime.rs`	Adds the blocking runtime facade.
`diskann-providers/src/model/graph/provider/layers/betafilter.rs`	Gates Tokio tests.
`diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs`	Gates Tokio tests.
`diskann-providers/src/model/graph/provider/async_/memory_vector_provider.rs`	Gates Tokio tests.
`diskann-providers/src/model/graph/provider/async_/memory_quant_vector_provider.rs`	Gates Tokio tests.
`diskann-providers/src/model/graph/provider/async_/inmem/scalar.rs`	Gates Tokio tests.
`diskann-providers/src/model/graph/provider/async_/inmem/provider.rs`	Gates Tokio tests.
`diskann-providers/src/model/graph/provider/async_/fast_memory_vector_provider.rs`	Gates Tokio tests.
`diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs`	Gates Tokio tests.
`diskann-providers/src/lib.rs`	Configures the runtime module.
`diskann-providers/src/index/wrapped_async.rs`	Makes wrappers backend-aware.
`diskann-providers/src/index/diskann_async.rs`	Gates Tokio tests.
`diskann-providers/Cargo.toml`	Adds runtime features and forwarding.
`diskann-label-filter/Cargo.toml`	Selects Tokio explicitly.
`diskann-inmem/src/provider.rs`	Adds Compio coverage.
`diskann-inmem/Cargo.toml`	Adds runtime features and forwarding.
`diskann-disk/Cargo.toml`	Selects Tokio features explicitly.
`diskann-bftree/Cargo.toml`	Updates test feature selection.
`diskann-benchmark-core/Cargo.toml`	Selects Tokio features explicitly.
`Cargo.toml`	Adds the workspace Compio dependency.
`Cargo.lock`	Locks Compio and transitive dependencies.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+default = ["tokio", "tracing"]
+
+# Use the tokio runtime (forwarded to `diskann`); upstream default behavior.
+tokio = ["dep:tokio", "diskann/tokio"]


    /// This needs to be manually `async` so we can aply the `'static` bound. Since it's
    /// recursive, Rust struggles to properly deduce the hidden type for the opaque return
    /// type.
+    #[cfg(feature = "tokio")]


+#[cfg(any(feature = "tokio", feature = "compio"))]
+pub(crate) mod runtime;


+#[cfg(any(feature = "tokio", feature = "compio"))]
+pub(crate) mod runtime;


+///
+/// The runtime backend is selected by crate features: `tokio` (default) or
+/// `compio`. Under `tokio`, external handles to an existing runtime are
+/// supported via [`Self::new_with_handle`]; under `compio` the runtime is


+/// Dropping all handles to a task may cancel it, mirroring the semantics of
+/// both backends.


Mark Hildebrand (hildebrandmw) · 2026-09-23T22:18:09Z

Thanks - this isn't the right way to do this. As identified, we don't use many features of tokio. It would be better to actually make the code runtime agnostic, instead of needing N features for N different runtimes.

x@01 (x-at-01) · 2026-09-24T03:12:39Z

Closing this PR to redesign with a clean runtime-agnostic approach instead of feature multiplexing, as suggested.

webc-fork added 6 commits September 23, 2026 14:14

x@01 (x-at-01) requested review from a team and a lite review from Copilot September 23, 2026 08:58

Copilot started reviewing on behalf of x@01 (x-at-01) September 23, 2026 08:59

View session

Copilot AI reviewed Sep 23, 2026

View reviewed changes

x@01 (x-at-01) closed this Sep 24, 2026

Navigation Menu

Repository navigation

feat(providers): make the async runtime pluggable (tokio | compio) - #1435

feat(providers): make the async runtime pluggable (tokio | compio)#1435
x@01 (x-at-01) wants to merge 6 commits into
microsoft:mainfrom
webc-fork:runtime-agnostic

x@01 (x-at-01) commented Sep 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Mark Hildebrand (hildebrandmw) commented Sep 23, 2026

Uh oh!

x@01 (x-at-01) commented Sep 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		#[cfg(any(feature = "tokio", feature = "compio"))]
		pub(crate) mod runtime;

		/// Dropping all handles to a task may cancel it, mirroring the semantics of
		/// both backends.

Conversation

x@01 (x-at-01) commented Sep 23, 2026

Problem

Solution

Behavior compatibility

Verification

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Mark Hildebrand (hildebrandmw) commented Sep 23, 2026

Uh oh!

x@01 (x-at-01) commented Sep 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants