diff --git a/.cargo/config.toml b/.cargo/config.toml index 94a019606..bdd5800c8 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -10,3 +10,4 @@ rustflags = ["-C", "target-cpu=x86-64-v3", "-C", "force-frame-pointers=yes"] # Force frame pointers to make performance and memory profiling much easier [target.aarch64-unknown-linux-gnu] rustflags = ["-C", "target-cpu=neoverse-n1", "-C", "force-frame-pointers=yes"] + diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 508d1a3dd..189e06f44 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -109,7 +109,7 @@ jobs: steps: - uses: actions/checkout@v4 - name: Install Rust components - run: rustup component add rustfmt clippy + run: rustup component add rustfmt clippy && rustup target add wasm32-unknown-unknown - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 with: cache-bin: false @@ -120,6 +120,8 @@ jobs: run: cargo clippy --workspace --all-targets --locked -- -D warnings - name: cargo clippy (prefill-router) run: cargo clippy -p switchyard-server --all-targets --features prefill-router --locked -- -D warnings + - name: cargo clippy (libsy, wasm32) + run: cargo clippy -p switchyard-libsy --target wasm32-unknown-unknown --locked -- -D warnings - name: cargo test run: cargo test --workspace --locked - name: cargo test (prefill-router) diff --git a/Cargo.lock b/Cargo.lock index b1cbce1be..98639abca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2325,6 +2325,7 @@ version = "0.3.0" dependencies = [ "async-trait", "futures", + "getrandom 0.4.3", "http", "jsonptr", "jsonschema", @@ -2343,6 +2344,8 @@ dependencies = [ "tracing-opentelemetry", "tracing-subscriber", "uuid", + "wasm-bindgen-futures", + "web-time", ] [[package]] diff --git a/crates/libsy/Cargo.toml b/crates/libsy/Cargo.toml index fa92e3094..08b9109ec 100644 --- a/crates/libsy/Cargo.toml +++ b/crates/libsy/Cargo.toml @@ -33,12 +33,25 @@ rand.workspace = true regex.workspace = true switchyard-protocol.workspace = true thiserror.workspace = true -tokio.workspace = true -tokio-stream = "0.1" +tokio-stream = { version = "0.1", default-features = false } tracing.workspace = true tracing-opentelemetry.workspace = true uuid.workspace = true +[target.'cfg(not(all(target_arch = "wasm32", target_os = "unknown")))'.dependencies] +tokio.workspace = true + +[target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies] +# Workers/browser targets have no OS threads or epoll: only tokio's +# runtime-agnostic pieces (channels, mutexes, macros) are linked. +tokio = { version = "1", default-features = false, features = ["sync", "macros"] } +wasm-bindgen-futures = "0.4" +web-time = "1" +# `rand` 0.10 pulls getrandom 0.4; route it through the JS crypto API. +# (`jsonschema` already does the same for the 0.3 line that `ahash` uses.) +getrandom = { version = "0.4", default-features = false, features = ["wasm_js"] } +uuid = { workspace = true, features = ["js"] } + [dev-dependencies] # SDK + in-memory exporter to assert what the observability layer records. opentelemetry_sdk = { version = "0.32", features = ["metrics", "testing", "trace"] } diff --git a/crates/libsy/src/algorithms/advisor_gate.rs b/crates/libsy/src/algorithms/advisor_gate.rs index d638864f8..5cc843b04 100644 --- a/crates/libsy/src/algorithms/advisor_gate.rs +++ b/crates/libsy/src/algorithms/advisor_gate.rs @@ -31,7 +31,8 @@ //! holds the only mutable state. use std::sync::Arc; -use std::time::Instant; + +use crate::rt::Instant; use switchyard_protocol::{ Category, ContentBlock, InstructionBlock, LlmRequest, Message, ModelId, OutputParams, Request, diff --git a/crates/libsy/src/algorithms/fall_through.rs b/crates/libsy/src/algorithms/fall_through.rs index 8d6f5e71d..967dc4fc2 100644 --- a/crates/libsy/src/algorithms/fall_through.rs +++ b/crates/libsy/src/algorithms/fall_through.rs @@ -19,14 +19,19 @@ use std::{ collections::HashMap, - sync::{Arc, Once, Weak}, - time::{Duration, Instant}, + sync::{Arc, Once}, + time::Duration, }; +#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))] +use std::sync::Weak; + use async_trait::async_trait; use parking_lot::Mutex; use tokio::sync::Mutex as AsyncMutex; +use crate::rt::Instant; + use crate::core::algorithm::{self, Algorithm, Driver}; use crate::core::classifier::{Classifier, Score}; use crate::core::processor::{Event, Processor}; @@ -58,6 +63,8 @@ pub struct FallThrough { classifiers: Vec>>, session_states: Option>>, cleanup_started: Once, + #[cfg(all(target_arch = "wasm32", target_os = "unknown"))] + next_cleanup: Mutex, } impl FallThrough<()> { @@ -69,6 +76,8 @@ impl FallThrough<()> { classifiers: Vec::new(), session_states: None, cleanup_started: Once::new(), + #[cfg(all(target_arch = "wasm32", target_os = "unknown"))] + next_cleanup: Mutex::new(Instant::now() + SESSION_CLEANUP_INTERVAL), } } } @@ -85,6 +94,8 @@ where classifiers: Vec::new(), session_states: Some(Arc::new(Mutex::new(HashMap::new()))), cleanup_started: Once::new(), + #[cfg(all(target_arch = "wasm32", target_os = "unknown"))] + next_cleanup: Mutex::new(Instant::now() + SESSION_CLEANUP_INTERVAL), } } @@ -128,7 +139,12 @@ where }; let states = Arc::downgrade(states); self.cleanup_started.call_once(move || { + // Timer-driven background cleanup needs a Tokio runtime; wasm hosts + // sweep expired sessions inline in `session_state` instead. + #[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))] drop(tokio::spawn(cleanup_inactive_sessions(states))); + #[cfg(all(target_arch = "wasm32", target_os = "unknown"))] + drop(states); }); } @@ -184,8 +200,16 @@ where fn session_state(&self, request: &Request) -> Option>> { let states = self.session_states.as_ref()?; let session_id = session_id(request)?; - let mut states = states.lock(); let now = Instant::now(); + #[cfg(all(target_arch = "wasm32", target_os = "unknown"))] + { + let mut next_cleanup = self.next_cleanup.lock(); + if now >= *next_cleanup { + remove_inactive_sessions(states, now, SESSION_STATE_TTL); + *next_cleanup = now + SESSION_CLEANUP_INTERVAL; + } + } + let mut states = states.lock(); let session = states.entry(session_id).or_insert_with(|| SessionState { state: Arc::new(AsyncMutex::new(S::default())), last_accessed: now, @@ -246,6 +270,7 @@ where } } +#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))] async fn cleanup_inactive_sessions(states: Weak>) where S: Send + 'static, diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 119809278..6548ee025 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -4,10 +4,7 @@ //! The [`Algorithm`] trait and its [`Driver`] — the orchestration contract every //! algorithm implements and the offload channel it uses for routing-time model calls. -use std::{ - collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc, - time::Instant, -}; +use std::{collections::HashMap, future::Future, panic::AssertUnwindSafe, pin::Pin, sync::Arc}; use async_trait::async_trait; use futures::{FutureExt, Stream, StreamExt}; @@ -26,6 +23,7 @@ use tracing::Instrument; /// [`switchyard_protocol::LlmResponseStream`] or the terminal aggregate. use switchyard_protocol::{Category, ModelId, Request, Response}; +use crate::rt::Instant; use crate::{DriverError, LibsyError, Result, observability}; /// A boxed, `Send` stream of [`Step`]s — the output of @@ -419,7 +417,7 @@ fn panic_message(payload: &(dyn std::any::Any + Send)) -> String { } /// Abort guard -struct AbortOnDrop(tokio::task::AbortHandle); +struct AbortOnDrop(crate::rt::AbortHandle); impl Drop for AbortOnDrop { fn drop(&mut self) { @@ -526,7 +524,7 @@ pub trait Algorithm: Send + Sync + 'static { fn run_stream(self: Arc, request: Request, models: Arc) -> StepStream { let (driver, step_rx) = Driver::new(self.name(), models); let span = observability::run_span(self.name(), &request); - let handle = tokio::spawn( + let handle = crate::rt::spawn_abortable( async move { let algorithm = self.name().to_string(); // Catch a panicking algorithm so the run still publishes a terminal step. @@ -548,7 +546,7 @@ pub trait Algorithm: Send + Sync + 'static { .instrument(span), ); // Dropping the stream aborts the algorithm task when its consumer goes away. - let abort_guard = AbortOnDrop(handle.abort_handle()); + let abort_guard = AbortOnDrop(handle); Box::pin(ReceiverStream::new(step_rx).map(move |step| { // link abort guard to stream let _keep_alive = &abort_guard; diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 9d07aeff9..d4426dd53 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -14,6 +14,7 @@ pub use core::processor::{Event, Processor}; pub use core::state::{State, StateValue}; mod error; +mod rt; pub use error::{DriverError, LibsyError, Result}; mod algorithms; diff --git a/crates/libsy/src/observability.rs b/crates/libsy/src/observability.rs index 69c6d9295..ce721814d 100644 --- a/crates/libsy/src/observability.rs +++ b/crates/libsy/src/observability.rs @@ -33,7 +33,9 @@ //! negligible next to a model call. use std::future::Future; -use std::time::{Duration, Instant}; +use std::time::Duration; + +use crate::rt::Instant; use opentelemetry::metrics::Meter; use opentelemetry::{KeyValue, global}; diff --git a/crates/libsy/src/rt.rs b/crates/libsy/src/rt.rs new file mode 100644 index 000000000..25a118e02 --- /dev/null +++ b/crates/libsy/src/rt.rs @@ -0,0 +1,41 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Runtime abstraction over task spawning and monotonic clocks, so the crate +//! runs on native Tokio hosts and on single-threaded wasm32 hosts (browsers, +//! Cloudflare Workers) alike. + +/// Monotonic instant. On wasm32 `std::time::Instant::now()` aborts, so a +/// JS-clock-backed drop-in replacement is used there. +#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))] +pub(crate) use std::time::Instant; +#[cfg(all(target_arch = "wasm32", target_os = "unknown"))] +pub(crate) use web_time::Instant; + +/// Handle that aborts the task returned by [`spawn_abortable`]. +#[cfg(all(target_arch = "wasm32", target_os = "unknown"))] +pub(crate) use futures::future::AbortHandle; +#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))] +pub(crate) use tokio::task::AbortHandle; + +/// Spawns a future on the host runtime and returns a handle that aborts it. +#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))] +pub(crate) fn spawn_abortable(future: F) -> AbortHandle +where + F: std::future::Future + Send + 'static, +{ + tokio::spawn(future).abort_handle() +} + +/// Spawns a future on the JS microtask queue and returns a handle that aborts it. +#[cfg(all(target_arch = "wasm32", target_os = "unknown"))] +pub(crate) fn spawn_abortable(future: F) -> AbortHandle +where + F: std::future::Future + 'static, +{ + let (handle, registration) = AbortHandle::new_pair(); + wasm_bindgen_futures::spawn_local(async move { + let _ = futures::future::Abortable::new(future, registration).await; + }); + handle +}