diff --git a/Cargo.lock b/Cargo.lock index 13b670f55..487045792 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1618,6 +1618,24 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "fake-driver-rs" +version = "0.0.0" +dependencies = [ + "clap", + "futures", + "miette", + "openshell-core", + "serde", + "thiserror 2.0.18", + "tokio", + "tokio-stream", + "toml", + "tonic", + "tracing", + "tracing-subscriber", +] + [[package]] name = "fastrand" version = "2.4.1" diff --git a/Cargo.toml b/Cargo.toml index f450cd5c8..506b88108 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,7 @@ [workspace] resolver = "2" -members = ["crates/*"] +members = ["crates/*", "examples/fake-driver-rs"] [workspace.package] version = "0.0.0" diff --git a/examples/fake-driver-rs/Cargo.toml b/examples/fake-driver-rs/Cargo.toml new file mode 100644 index 000000000..043763396 --- /dev/null +++ b/examples/fake-driver-rs/Cargo.toml @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "fake-driver-rs" +description = "Scriptable fake compute driver for testing gateway integration" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +repository.workspace = true + +[[bin]] +name = "fake-driver-rs" +path = "src/main.rs" + +[dependencies] +openshell-core = { path = "../../crates/openshell-core", default-features = false } + +tokio = { workspace = true } +tokio-stream = { workspace = true } +tonic = { workspace = true, features = ["transport"] } +futures = { workspace = true } +serde = { workspace = true } +toml = { workspace = true } +clap = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } +thiserror = { workspace = true } +miette = { workspace = true } + +[lints] +workspace = true diff --git a/examples/fake-driver-rs/scenarios/happy-path.toml b/examples/fake-driver-rs/scenarios/happy-path.toml new file mode 100644 index 000000000..e986fb77c --- /dev/null +++ b/examples/fake-driver-rs/scenarios/happy-path.toml @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Happy-path scenario: all RPCs succeed. +# +# Each [[rpc_name]] block is one entry in the scripted sequence. The driver +# advances through the list on each call; when the list is exhausted the last +# entry repeats (with a warning logged). +# +# To return a gRPC error instead of a success response, add: +# +# [rpc_name.error] +# code = "NOT_FOUND" +# message = "sandbox not found" + +[[validate_sandbox_create]] +# success — no error field + +[[create_sandbox]] +# success — no error field + +[[get_sandbox]] +[get_sandbox.sandbox] +id = "sandbox-abc123" +name = "test-sandbox" +namespace = "default" + +[get_sandbox.sandbox.status] +sandbox_name = "test-sandbox" +instance_id = "instance-1" +agent_fd = "127.0.0.1:12345" +sandbox_fd = "127.0.0.1:12346" + +[[list_sandboxes]] +# Empty list on first call. +sandboxes = [] + +[[list_sandboxes]] +# Populated list on subsequent calls (last entry repeats). +[[list_sandboxes.sandboxes]] +id = "sandbox-abc123" +name = "test-sandbox" +namespace = "default" + +[list_sandboxes.sandboxes.status] +sandbox_name = "test-sandbox" +instance_id = "instance-1" +agent_fd = "127.0.0.1:12345" +sandbox_fd = "127.0.0.1:12346" + +[[stop_sandbox]] +# success — no error field + +[[delete_sandbox]] +deleted = true + +# WatchSandboxes streams these events in order, then closes. +[watch_sandboxes] +delay_ms = 100 + +[[watch_sandboxes.events]] +[watch_sandboxes.events.sandbox] +id = "sandbox-abc123" +name = "test-sandbox" +namespace = "default" + +[watch_sandboxes.events.sandbox.status] +sandbox_name = "test-sandbox" +instance_id = "instance-1" +agent_fd = "127.0.0.1:12345" +sandbox_fd = "127.0.0.1:12346" + +[[watch_sandboxes.events]] +[watch_sandboxes.events.deleted] +sandbox_id = "sandbox-abc123" diff --git a/examples/fake-driver-rs/src/config.rs b/examples/fake-driver-rs/src/config.rs new file mode 100644 index 000000000..7a9191eb6 --- /dev/null +++ b/examples/fake-driver-rs/src/config.rs @@ -0,0 +1,174 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! TOML configuration types for the fake driver's scripted response sequences. +//! +//! Each RPC has an ordered list of entries. The driver advances through the list +//! on each call. When the sequence is exhausted the last entry is repeated and a +//! warning is logged. + +use serde::Deserialize; + +/// Top-level configuration loaded from the `--config` file. +#[derive(Debug, Default, Deserialize)] +pub struct Config { + #[serde(default)] + pub validate_sandbox_create: Vec, + #[serde(default)] + pub create_sandbox: Vec, + #[serde(default)] + pub get_sandbox: Vec, + #[serde(default)] + pub list_sandboxes: Vec, + #[serde(default)] + pub stop_sandbox: Vec, + #[serde(default)] + pub delete_sandbox: Vec, + #[serde(default)] + pub watch_sandboxes: WatchSandboxesConfig, +} + +/// A gRPC error to return in place of a success response. +#[derive(Debug, Deserialize)] +pub struct GrpcError { + /// gRPC status code name, e.g. `"NOT_FOUND"`, `"INVALID_ARGUMENT"`. + pub code: String, + pub message: String, +} + +impl GrpcError { + pub fn to_status(&self) -> tonic::Status { + let code = match self.code.to_uppercase().as_str() { + "CANCELLED" => tonic::Code::Cancelled, + "INVALID_ARGUMENT" => tonic::Code::InvalidArgument, + "DEADLINE_EXCEEDED" => tonic::Code::DeadlineExceeded, + "NOT_FOUND" => tonic::Code::NotFound, + "ALREADY_EXISTS" => tonic::Code::AlreadyExists, + "PERMISSION_DENIED" => tonic::Code::PermissionDenied, + "RESOURCE_EXHAUSTED" => tonic::Code::ResourceExhausted, + "FAILED_PRECONDITION" => tonic::Code::FailedPrecondition, + "ABORTED" => tonic::Code::Aborted, + "OUT_OF_RANGE" => tonic::Code::OutOfRange, + "UNIMPLEMENTED" => tonic::Code::Unimplemented, + "INTERNAL" => tonic::Code::Internal, + "UNAVAILABLE" => tonic::Code::Unavailable, + "DATA_LOSS" => tonic::Code::DataLoss, + "UNAUTHENTICATED" => tonic::Code::Unauthenticated, + _ => tonic::Code::Unknown, + }; + tonic::Status::new(code, &self.message) + } +} + +/// Entry for RPCs that return an empty success response. +/// +/// Omit `error` for success; set `error` to return a gRPC status error. +#[derive(Debug, Default, Deserialize)] +pub struct SimpleEntry { + pub error: Option, +} + +/// Entry for `ValidateSandboxCreate`. Same shape as `SimpleEntry` but kept +/// distinct so the config section name is self-documenting. +#[derive(Debug, Default, Deserialize)] +pub struct ValidateSandboxCreateEntry { + pub error: Option, +} + +/// Entry for `GetSandbox`. +#[derive(Debug, Default, Deserialize)] +pub struct GetSandboxEntry { + pub error: Option, + /// The sandbox to return on success. If absent the RPC returns `NOT_FOUND`. + pub sandbox: Option, +} + +/// Entry for `ListSandboxes`. +#[derive(Debug, Default, Deserialize)] +pub struct ListSandboxesEntry { + pub error: Option, + #[serde(default)] + pub sandboxes: Vec, +} + +/// Entry for `DeleteSandbox`. +#[derive(Debug, Default, Deserialize)] +pub struct DeleteSandboxEntry { + pub error: Option, + /// Whether a resource was deleted. Defaults to `true`. + #[serde(default = "default_deleted")] + pub deleted: bool, +} + +fn default_deleted() -> bool { + true +} + +/// Simplified sandbox descriptor used in scripted responses. +#[derive(Debug, Default, Deserialize)] +pub struct SandboxConfig { + #[serde(default)] + pub id: String, + #[serde(default)] + pub name: String, + #[serde(default)] + pub namespace: String, + pub status: Option, +} + +/// Simplified sandbox status used in scripted responses. +#[derive(Debug, Default, Deserialize)] +pub struct SandboxStatusConfig { + #[serde(default)] + pub sandbox_name: String, + #[serde(default)] + pub instance_id: String, + /// Address the agent service is reachable at (e.g. `"127.0.0.1:12345"`). + #[serde(default)] + pub agent_fd: String, + /// Address the sandbox supervisor service is reachable at. + #[serde(default)] + pub sandbox_fd: String, +} + +/// Configuration for the `WatchSandboxes` streaming RPC. +/// +/// A single flat sequence of events is emitted in order with `delay_ms` +/// between each. The stream closes after the last event. +#[derive(Debug, Deserialize)] +pub struct WatchSandboxesConfig { + /// Milliseconds to wait between emitting consecutive events. + #[serde(default = "default_delay_ms")] + pub delay_ms: u64, + #[serde(default)] + pub events: Vec, +} + +fn default_delay_ms() -> u64 { + 100 +} + +impl Default for WatchSandboxesConfig { + fn default() -> Self { + Self { + delay_ms: default_delay_ms(), + events: Vec::new(), + } + } +} + +/// A single event emitted on the `WatchSandboxes` stream. +/// +/// Exactly one of `sandbox` or `deleted` should be set. +#[derive(Debug, Deserialize)] +pub struct WatchEventConfig { + /// Emit a sandbox-updated event with the given snapshot. + pub sandbox: Option, + /// Emit a sandbox-deleted event with the given sandbox ID. + pub deleted: Option, +} + +#[derive(Debug, Deserialize)] +pub struct DeletedEventConfig { + pub sandbox_id: String, +} diff --git a/examples/fake-driver-rs/src/driver.rs b/examples/fake-driver-rs/src/driver.rs new file mode 100644 index 000000000..e67b54e62 --- /dev/null +++ b/examples/fake-driver-rs/src/driver.rs @@ -0,0 +1,282 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![allow(clippy::result_large_err)] // gRPC handlers return Result<_, tonic::Status> + +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use futures::Stream; +use openshell_core::proto::compute::v1::{ + CreateSandboxRequest, CreateSandboxResponse, DeleteSandboxRequest, DeleteSandboxResponse, + DriverSandbox, DriverSandboxStatus, GetCapabilitiesRequest, GetCapabilitiesResponse, + GetSandboxRequest, GetSandboxResponse, ListSandboxesRequest, ListSandboxesResponse, + StopSandboxRequest, StopSandboxResponse, ValidateSandboxCreateRequest, + ValidateSandboxCreateResponse, WatchSandboxesDeletedEvent, WatchSandboxesEvent, + WatchSandboxesRequest, WatchSandboxesSandboxEvent, compute_driver_server::ComputeDriver, + watch_sandboxes_event::Payload, +}; +use tokio_stream::wrappers::ReceiverStream; +use tonic::{Request, Response, Status}; + +use crate::config::{Config, SandboxConfig, WatchSandboxesConfig}; +use crate::sequence::Sequence; + +pub struct FakeDriver { + state: Arc>, + watch_config: WatchSandboxesConfig, +} + +struct DriverState { + validate_sandbox_create: Sequence, + create_sandbox: Sequence, + get_sandbox: Sequence, + list_sandboxes: Sequence, + stop_sandbox: Sequence, + delete_sandbox: Sequence, +} + +impl FakeDriver { + pub fn new(config: Config) -> Self { + let Config { + validate_sandbox_create, + create_sandbox, + get_sandbox, + list_sandboxes, + stop_sandbox, + delete_sandbox, + watch_sandboxes, + } = config; + + Self { + state: Arc::new(Mutex::new(DriverState { + validate_sandbox_create: Sequence::new( + "validate_sandbox_create", + validate_sandbox_create, + ), + create_sandbox: Sequence::new("create_sandbox", create_sandbox), + get_sandbox: Sequence::new("get_sandbox", get_sandbox), + list_sandboxes: Sequence::new("list_sandboxes", list_sandboxes), + stop_sandbox: Sequence::new("stop_sandbox", stop_sandbox), + delete_sandbox: Sequence::new("delete_sandbox", delete_sandbox), + })), + watch_config: watch_sandboxes, + } + } +} + +fn sandbox_to_proto(cfg: &SandboxConfig) -> DriverSandbox { + DriverSandbox { + id: cfg.id.clone(), + name: cfg.name.clone(), + namespace: cfg.namespace.clone(), + status: cfg.status.as_ref().map(|s| DriverSandboxStatus { + sandbox_name: s.sandbox_name.clone(), + instance_id: s.instance_id.clone(), + agent_fd: s.agent_fd.clone(), + sandbox_fd: s.sandbox_fd.clone(), + ..Default::default() + }), + ..Default::default() + } +} + +#[tonic::async_trait] +impl ComputeDriver for FakeDriver { + async fn get_capabilities( + &self, + _request: Request, + ) -> Result, Status> { + Ok(Response::new(GetCapabilitiesResponse { + driver_name: "fake-driver-rs".to_string(), + driver_version: env!("CARGO_PKG_VERSION").to_string(), + default_image: "ghcr.io/nvidia/openshell-sandbox:latest".to_string(), + })) + } + + async fn validate_sandbox_create( + &self, + request: Request, + ) -> Result, Status> { + let sandbox_id = request + .into_inner() + .sandbox + .as_ref() + .map_or_else(String::new, |s| s.id.clone()); + let mut state = self.state.lock().expect("state lock poisoned"); + let entry = state.validate_sandbox_create.next().ok_or_else(|| { + Status::unimplemented("no validate_sandbox_create entries configured") + })?; + tracing::info!( + sandbox_id, + rpc = "validate_sandbox_create", + "handling request" + ); + if let Some(err) = &entry.error { + return Err(err.to_status()); + } + Ok(Response::new(ValidateSandboxCreateResponse {})) + } + + async fn create_sandbox( + &self, + request: Request, + ) -> Result, Status> { + let sandbox_id = request + .into_inner() + .sandbox + .as_ref() + .map_or_else(String::new, |s| s.id.clone()); + let mut state = self.state.lock().expect("state lock poisoned"); + let entry = state + .create_sandbox + .next() + .ok_or_else(|| Status::unimplemented("no create_sandbox entries configured"))?; + tracing::info!(sandbox_id, rpc = "create_sandbox", "handling request"); + if let Some(err) = &entry.error { + return Err(err.to_status()); + } + Ok(Response::new(CreateSandboxResponse {})) + } + + async fn get_sandbox( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + tracing::info!( + sandbox_id = req.sandbox_id, + sandbox_name = req.sandbox_name, + rpc = "get_sandbox", + "handling request" + ); + let mut state = self.state.lock().expect("state lock poisoned"); + let entry = state + .get_sandbox + .next() + .ok_or_else(|| Status::unimplemented("no get_sandbox entries configured"))?; + if let Some(err) = &entry.error { + return Err(err.to_status()); + } + let sandbox = entry + .sandbox + .as_ref() + .map(sandbox_to_proto) + .ok_or_else(|| Status::not_found("no sandbox in scripted entry"))?; + Ok(Response::new(GetSandboxResponse { + sandbox: Some(sandbox), + })) + } + + async fn list_sandboxes( + &self, + _request: Request, + ) -> Result, Status> { + tracing::info!(rpc = "list_sandboxes", "handling request"); + let mut state = self.state.lock().expect("state lock poisoned"); + let entry = state + .list_sandboxes + .next() + .ok_or_else(|| Status::unimplemented("no list_sandboxes entries configured"))?; + if let Some(err) = &entry.error { + return Err(err.to_status()); + } + let sandboxes = entry.sandboxes.iter().map(sandbox_to_proto).collect(); + Ok(Response::new(ListSandboxesResponse { sandboxes })) + } + + async fn stop_sandbox( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + tracing::info!( + sandbox_id = req.sandbox_id, + sandbox_name = req.sandbox_name, + rpc = "stop_sandbox", + "handling request" + ); + let mut state = self.state.lock().expect("state lock poisoned"); + let entry = state + .stop_sandbox + .next() + .ok_or_else(|| Status::unimplemented("no stop_sandbox entries configured"))?; + if let Some(err) = &entry.error { + return Err(err.to_status()); + } + Ok(Response::new(StopSandboxResponse {})) + } + + async fn delete_sandbox( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + tracing::info!( + sandbox_id = req.sandbox_id, + sandbox_name = req.sandbox_name, + rpc = "delete_sandbox", + "handling request" + ); + let mut state = self.state.lock().expect("state lock poisoned"); + let entry = state + .delete_sandbox + .next() + .ok_or_else(|| Status::unimplemented("no delete_sandbox entries configured"))?; + if let Some(err) = &entry.error { + return Err(err.to_status()); + } + Ok(Response::new(DeleteSandboxResponse { + deleted: entry.deleted, + })) + } + + type WatchSandboxesStream = + Pin> + Send + 'static>>; + + async fn watch_sandboxes( + &self, + _request: Request, + ) -> Result, Status> { + tracing::info!(rpc = "watch_sandboxes", "starting stream"); + + // Convert config events to proto eagerly so config errors are caught before streaming. + let mut proto_events: Vec = Vec::new(); + for e in &self.watch_config.events { + let event = if let Some(sandbox) = &e.sandbox { + WatchSandboxesEvent { + payload: Some(Payload::Sandbox(WatchSandboxesSandboxEvent { + sandbox: Some(sandbox_to_proto(sandbox)), + })), + } + } else if let Some(deleted) = &e.deleted { + WatchSandboxesEvent { + payload: Some(Payload::Deleted(WatchSandboxesDeletedEvent { + sandbox_id: deleted.sandbox_id.clone(), + })), + } + } else { + return Err(Status::internal( + "watch_sandboxes event has neither `sandbox` nor `deleted` payload", + )); + }; + proto_events.push(event); + } + + let delay = Duration::from_millis(self.watch_config.delay_ms); + let (tx, rx) = tokio::sync::mpsc::channel(16); + + tokio::spawn(async move { + for event in proto_events { + tokio::time::sleep(delay).await; + if tx.send(Ok(event)).await.is_err() { + return; // receiver dropped; gateway disconnected + } + } + tracing::info!(rpc = "watch_sandboxes", "stream exhausted; closing"); + }); + + Ok(Response::new(Box::pin(ReceiverStream::new(rx)))) + } +} diff --git a/examples/fake-driver-rs/src/main.rs b/examples/fake-driver-rs/src/main.rs new file mode 100644 index 000000000..f8c6c7295 --- /dev/null +++ b/examples/fake-driver-rs/src/main.rs @@ -0,0 +1,105 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Scriptable fake compute driver for testing gateway integration. +//! +//! Implements the `ComputeDriver` gRPC service with responses driven by a +//! TOML config file. Each RPC has an ordered sequence of scripted responses; +//! when the sequence is exhausted the last entry repeats. +//! +//! # Usage +//! +//! ```text +//! fake-driver-rs --socket /tmp/fake-driver.sock --config scenario.toml +//! ``` +//! +//! Point the gateway at it via its remote driver config: +//! +//! ```toml +//! [compute] +//! driver = "remote" +//! +//! [compute.remote] +//! socket = "/tmp/fake-driver.sock" +//! ``` + +use std::path::PathBuf; + +use clap::Parser; +use miette::{IntoDiagnostic, Result, WrapErr}; +use openshell_core::proto::compute::v1::compute_driver_server::ComputeDriverServer; +use tokio::net::UnixListener; +use tokio::signal::unix::{SignalKind, signal}; +use tokio_stream::wrappers::UnixListenerStream; +use tracing::info; +use tracing_subscriber::EnvFilter; + +mod config; +mod driver; +mod sequence; + +use config::Config; +use driver::FakeDriver; + +#[derive(Parser)] +#[command(name = "fake-driver-rs")] +#[command(about = "Scriptable fake compute driver for testing gateway integration")] +struct Args { + /// Unix socket path the driver will listen on. + #[arg(long, env = "FAKE_DRIVER_SOCKET")] + socket: PathBuf, + + /// Path to the TOML scenario config file. + #[arg(long, env = "FAKE_DRIVER_CONFIG")] + config: PathBuf, + + #[arg(long, env = "FAKE_DRIVER_LOG_LEVEL", default_value = "info")] + log_level: String, +} + +#[tokio::main] +async fn main() -> Result<()> { + let args = Args::parse(); + + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(&args.log_level)), + ) + .init(); + + let raw = std::fs::read_to_string(&args.config) + .into_diagnostic() + .wrap_err_with(|| format!("failed to read config from {}", args.config.display()))?; + let config: Config = toml::from_str(&raw) + .into_diagnostic() + .wrap_err("failed to parse config")?; + + // Remove a stale socket file if present so bind succeeds. + if args.socket.exists() { + std::fs::remove_file(&args.socket) + .into_diagnostic() + .wrap_err("failed to remove stale socket")?; + } + + let listener = UnixListener::bind(&args.socket) + .into_diagnostic() + .wrap_err_with(|| format!("failed to bind socket {}", args.socket.display()))?; + + info!(socket = %args.socket.display(), "fake-driver-rs listening"); + + let service = ComputeDriverServer::new(FakeDriver::new(config)); + + let mut sigterm = signal(SignalKind::terminate()).into_diagnostic()?; + + tonic::transport::Server::builder() + .add_service(service) + .serve_with_incoming_shutdown(UnixListenerStream::new(listener), async move { + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + _ = sigterm.recv() => {} + } + info!("received shutdown signal"); + }) + .await + .into_diagnostic() +} diff --git a/examples/fake-driver-rs/src/sequence.rs b/examples/fake-driver-rs/src/sequence.rs new file mode 100644 index 000000000..555fc711c --- /dev/null +++ b/examples/fake-driver-rs/src/sequence.rs @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +/// Advances through a scripted sequence of responses. +/// +/// When the sequence is exhausted the last entry is repeated indefinitely and a +/// warning is logged on each repeated call. +pub struct Sequence { + rpc: &'static str, + entries: Vec, + index: usize, +} + +impl Sequence { + pub fn new(rpc: &'static str, entries: Vec) -> Self { + Self { + rpc, + entries, + index: 0, + } + } + + /// Return the next entry in the sequence, repeating the last when exhausted. + /// + /// Returns `None` only when the sequence is empty. + pub fn next(&mut self) -> Option<&T> { + if self.entries.is_empty() { + return None; + } + if self.index < self.entries.len() { + let entry = &self.entries[self.index]; + self.index += 1; + Some(entry) + } else { + tracing::warn!( + rpc = self.rpc, + "scripted sequence exhausted; repeating last entry" + ); + self.entries.last() + } + } +}