diff --git a/Cargo.lock b/Cargo.lock index 8d65115..b148e3e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -590,6 +590,7 @@ dependencies = [ "tokio-rustls", "tokio-stream", "tonic", + "tonic-reflection", ] [[package]] @@ -2502,6 +2503,20 @@ dependencies = [ "tonic-build", ] +[[package]] +name = "tonic-reflection" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acccd136a4bf19810a1fde9c74edc6129b42a66b44d0c1c8aaa67aeb49a146a7" +dependencies = [ + "prost", + "prost-types", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", +] + [[package]] name = "tower" version = "0.5.2" diff --git a/README.md b/README.md index 40b0ee2..4497bc8 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,41 @@ dwd udp --no-ui --generator profile.yaml --api-addr 127.0.0.1:8080 curl http://127.0.0.1:8080/api/v1/metrics ``` +Alternatively (or additionally), expose the gRPC control API with `--grpc-addr` and drive the run remotely — see below. + +## Remote control (gRPC API) + +The interactive TUI talks to the load-generating core through a gRPC API (`dwd.v1.Dwd`, see `dwd-proto/dwdpb/dwd.proto`). Passing the global `--grpc-addr` flag exposes that same API on a network address, giving remote clients full control over the run — everything the TUI can do: + +- **`Control`** — set a fixed RPS, suspend/resume the load profile, or stop the run entirely (the engine drains and the process shuts down gracefully, exactly like a TUI exit or `SIGTERM`); +- **`StreamStats`** — a live stream of cumulative statistic snapshots (requests, bytes, response codes, latency histogram buckets); +- **`Describe`** — the engine kind and which statistic groups it provides. + +The flag works with and without the TUI; combined with `--no-ui` it makes a fully API-driven headless run: + +```bash +dwd udp --no-ui --grpc-addr 127.0.0.1:8081 +``` + +Server reflection is enabled, so `grpcurl` works without any proto files: + +```bash +# Set the load to 5000 RPS: +grpcurl -plaintext -d '{"set":{"rps":"5000"}}' 127.0.0.1:8081 dwd.v1.Dwd/Control + +# Suspend / resume the load profile: +grpcurl -plaintext -d '{"suspend":{}}' 127.0.0.1:8081 dwd.v1.Dwd/Control +grpcurl -plaintext -d '{"resume":{}}' 127.0.0.1:8081 dwd.v1.Dwd/Control + +# Watch live statistics: +grpcurl -plaintext -d '{}' 127.0.0.1:8081 dwd.v1.Dwd/StreamStats + +# Stop the run gracefully: +grpcurl -plaintext -d '{"stop":{}}' 127.0.0.1:8081 dwd.v1.Dwd/Control +``` + +Note: `set` overrides the profile with a fixed value (implicitly suspending it); `resume` returns control to the profile. The API is unauthenticated plaintext HTTP/2 — bind it to localhost or a trusted network only. + ## Load profiles A load profile is a declarative description of how the load should be generated and for how long. diff --git a/dwd-core/Cargo.toml b/dwd-core/Cargo.toml index c63593f..47ecd39 100644 --- a/dwd-core/Cargo.toml +++ b/dwd-core/Cargo.toml @@ -51,10 +51,14 @@ rand = "0.9" axum = "0.8" prometheus-client = "0.24" -# gRPC service (server half of the in-process API seam). +# gRPC service (server half of the API seam; in-process for the TUI and +# optionally exposed on a network address for remote control). dwd-proto = { path = "../dwd-proto" } tonic = "0.14" -tokio-stream = "0.1" +# Reflection lets `grpcurl` & co discover the API without local proto files. +tonic-reflection = "0.14" +# `net` provides `TcpListenerStream` for the network gRPC endpoint. +tokio-stream = { version = "0.1", features = ["net"] } # Optional DPDK FFI backend (Linux only, needs DPDK 24.11 installed). dpdk = { version = "0.1", path = "../dpdk", package = "dwd-dpdk", optional = true } diff --git a/dwd-core/src/grpc/mod.rs b/dwd-core/src/grpc/mod.rs index e7a98a8..56906d2 100644 --- a/dwd-core/src/grpc/mod.rs +++ b/dwd-core/src/grpc/mod.rs @@ -1,8 +1,10 @@ -//! Server half of the in-process gRPC API seam. +//! Server half of the gRPC API seam. //! //! [`server`] implements the [`Dwd`](dwd_proto::dwd_server::Dwd) service over the //! running core; [`snapshot`] maps the core statistics traits onto the wire -//! snapshot. The client half (TUI) lives in the `dwd` binary. +//! snapshot. The service is hosted in two ways: [`serve`] over an in-memory pipe +//! for the built-in TUI, and [`serve_tcp`] on a network address for remote +//! clients. The client half (TUI) lives in the `dwd` binary. pub mod server; pub mod snapshot; @@ -12,8 +14,10 @@ pub use self::{ snapshot::{EngineDescriptor, SnapshotSource}, }; -use tokio::{io::DuplexStream, task::JoinHandle}; -use tokio_stream::StreamExt as _; +use core::net::SocketAddr; + +use tokio::{io::DuplexStream, net::TcpListener, task::JoinHandle}; +use tokio_stream::{wrappers::TcpListenerStream, StreamExt as _}; use tonic::transport::Server; use dwd_proto::dwd_server::DwdServer; @@ -35,3 +39,38 @@ pub fn serve(service: DwdService, server_io: DuplexStream) -> JoinHandle Result<(SocketAddr, JoinHandle>), anyhow::Error> { + let listener = TcpListener::bind(addr).await?; + let addr = listener.local_addr()?; + log::info!("gRPC API server listening on {addr}"); + + let reflection_v1 = tonic_reflection::server::Builder::configure() + .register_encoded_file_descriptor_set(dwd_proto::FILE_DESCRIPTOR_SET) + .build_v1()?; + let reflection_v1alpha = tonic_reflection::server::Builder::configure() + .register_encoded_file_descriptor_set(dwd_proto::FILE_DESCRIPTOR_SET) + .build_v1alpha()?; + + let handle = tokio::spawn(async move { + Server::builder() + .add_service(reflection_v1) + .add_service(reflection_v1alpha) + .add_service(DwdServer::new(service)) + .serve_with_incoming(TcpListenerStream::new(listener)) + .await + }); + + Ok((addr, handle)) +} diff --git a/dwd-core/src/grpc/server.rs b/dwd-core/src/grpc/server.rs index af82c64..b9271e5 100644 --- a/dwd-core/src/grpc/server.rs +++ b/dwd-core/src/grpc/server.rs @@ -25,6 +25,10 @@ const STATS_INTERVAL: Duration = Duration::from_millis(25); const STATS_CHANNEL_CAP: usize = 4; /// The [`Dwd`] service implementation over the in-process core. +/// +/// Cloning is cheap (channels and `Arc`s), so the same service can back both +/// the in-memory TUI seam and the network endpoint at once. +#[derive(Clone)] pub struct DwdService { /// Forwards control RPCs to `run_generator` (mirrors the old UI mpsc). control_tx: Sender, @@ -65,6 +69,15 @@ impl Dwd for DwdService { Kind::Suspend(_) => GeneratorEvent::Suspend, Kind::Resume(_) => GeneratorEvent::Resume, Kind::Set(set) => GeneratorEvent::Set(set.rps), + Kind::Stop(_) => { + // Stop bypasses the generator channel: flipping the shared + // run flag drains the engine and the generator loop exactly + // as a TUI exit or SIGTERM would. + log::info!("stop requested via the API"); + self.is_running.store(false, Ordering::SeqCst); + + return Ok(Response::new(ControlResponse {})); + } }; // Non-blocking, coalescing send — mirrors the TUI's `try_send`: if the @@ -103,3 +116,72 @@ impl Dwd for DwdService { Ok(Response::new(ReceiverStream::new(rx))) } } + +#[cfg(test)] +mod tests { + use dwd_proto::{SetControl, StatGroup, StopControl, SuspendControl}; + + use super::*; + use crate::grpc::snapshot::EngineDescriptor; + + struct EmptySource; + + impl SnapshotSource for EmptySource { + fn snapshot(&self) -> StatsSnapshot { + StatsSnapshot::default() + } + } + + fn service() -> (DwdService, mpsc::Receiver, Arc) { + let (control_tx, control_rx) = mpsc::channel(1); + let descriptor = EngineDescriptor { + engine_kind: "udp", + groups: vec![StatGroup::Common], + } + .into_response(); + let is_running = Arc::new(AtomicBool::new(true)); + let service = DwdService::new(control_tx, Arc::new(EmptySource), descriptor, is_running.clone()); + + (service, control_rx, is_running) + } + + #[tokio::test] + async fn control_forwards_generator_events() { + let (service, mut control_rx, is_running) = service(); + + service + .control(Request::new(ControlRequest { + kind: Some(Kind::Set(SetControl { rps: 1234 })), + })) + .await + .expect("set control"); + assert!(matches!(control_rx.try_recv(), Ok(GeneratorEvent::Set(1234)))); + + service + .control(Request::new(ControlRequest { + kind: Some(Kind::Suspend(SuspendControl {})), + })) + .await + .expect("suspend control"); + assert!(matches!(control_rx.try_recv(), Ok(GeneratorEvent::Suspend))); + + // Generator control never touches the run flag. + assert!(is_running.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn control_stop_flips_run_flag() { + let (service, mut control_rx, is_running) = service(); + + service + .control(Request::new(ControlRequest { + kind: Some(Kind::Stop(StopControl {})), + })) + .await + .expect("stop control"); + + // Stop is not a generator event: it acts on the run flag directly. + assert!(!is_running.load(Ordering::SeqCst)); + assert!(control_rx.try_recv().is_err()); + } +} diff --git a/dwd-proto/build.rs b/dwd-proto/build.rs index 4afdc6c..13ddd2e 100644 --- a/dwd-proto/build.rs +++ b/dwd-proto/build.rs @@ -3,9 +3,14 @@ fn main() -> Result<(), Box> { // protobuf-compiler (keeps CI and the Docker builds self-contained). std::env::set_var("PROTOC", protoc_bin_vendored::protoc_bin_path()?); + let out_dir = std::path::PathBuf::from(std::env::var("OUT_DIR")?); + tonic_prost_build::configure() .build_client(true) .build_server(true) + // Emitted for the gRPC reflection service, which lets tools like + // `grpcurl` discover the API without a local copy of the proto file. + .file_descriptor_set_path(out_dir.join("dwd_descriptor.bin")) .compile_protos(&["dwdpb/dwd.proto"], &["dwdpb"])?; Ok(()) } diff --git a/dwd-proto/dwdpb/dwd.proto b/dwd-proto/dwdpb/dwd.proto index 1f23419..957f575 100644 --- a/dwd-proto/dwdpb/dwd.proto +++ b/dwd-proto/dwdpb/dwd.proto @@ -5,9 +5,10 @@ package dwd.v1; // The control + observability API seam between the DWD CLI/TUI (client) and the // load-generating core (server). // -// The service is consumed in-process over an in-memory transport (no sockets), -// but the contract is a regular gRPC service so the CLI talks to the core -// exclusively through this API. +// The service is consumed in-process over an in-memory transport (no sockets) +// by the built-in TUI, and can additionally be exposed on a network address +// (`--grpc-addr`) for remote control — the contract is a regular gRPC service +// either way, so every client talks to the core exclusively through this API. service Dwd { // Describes the running engine so the client knows which statistics groups // (i.e. which TUI widgets) to render. Replaces the previous compile-time @@ -48,6 +49,7 @@ message ControlRequest { SuspendControl suspend = 1; ResumeControl resume = 2; SetControl set = 3; + StopControl stop = 4; } } @@ -59,6 +61,9 @@ message ResumeControl {} message SetControl { uint64 rps = 1; } +// Stops the run entirely: the engine drains and the process shuts down +// gracefully, exactly as if the TUI had exited or SIGTERM was received. +message StopControl {} message ControlResponse {} diff --git a/dwd-proto/src/lib.rs b/dwd-proto/src/lib.rs index ad8049a..4ef46d3 100644 --- a/dwd-proto/src/lib.rs +++ b/dwd-proto/src/lib.rs @@ -10,3 +10,7 @@ pub mod pb { } pub use pb::*; + +/// Encoded `FileDescriptorSet` of the API, consumed by the gRPC reflection +/// service so tools like `grpcurl` can discover the API without proto files. +pub const FILE_DESCRIPTOR_SET: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dwd_descriptor.bin")); diff --git a/dwd/src/cfg.rs b/dwd/src/cfg.rs index ab4a541..d636818 100644 --- a/dwd/src/cfg.rs +++ b/dwd/src/cfg.rs @@ -22,6 +22,8 @@ pub struct Config { pub generator_fn: BoxedGeneratorNew, /// Address to expose the Prometheus API on. pub api_addr: Option, + /// Address to expose the gRPC control API on. + pub grpc_addr: Option, /// Run headless, without the interactive TUI. pub no_ui: bool, } @@ -32,6 +34,7 @@ impl TryFrom for Config { fn try_from(v: Cmd) -> Result { let mode = v.mode.into_config()?; let api_addr = v.api_addr; + let grpc_addr = v.grpc_addr; let no_ui = v.no_ui; let generator_fn = { let path = v.generator.clone(); @@ -49,6 +52,12 @@ impl TryFrom for Config { })) }; - Ok(Self { mode, generator_fn, api_addr, no_ui }) + Ok(Self { + mode, + generator_fn, + api_addr, + grpc_addr, + no_ui, + }) } } diff --git a/dwd/src/cmd.rs b/dwd/src/cmd.rs index a64d41b..bbccd8d 100644 --- a/dwd/src/cmd.rs +++ b/dwd/src/cmd.rs @@ -40,6 +40,15 @@ pub struct Cmd { /// to observe the generator state and metrics. #[clap(long, global = true, value_name = "IP:PORT")] pub api_addr: Option, + /// Address to expose the gRPC control API on. + /// + /// When specified, serves the same `dwd.v1.Dwd` gRPC service the built-in + /// TUI uses, so remote clients get full control over the run: set the RPS, + /// suspend/resume the generator, stop the run and stream live statistics. + /// Server reflection is enabled, so `grpcurl` works without proto files. + /// Combine with `--no-ui` to drive a headless run entirely over the API. + #[clap(long, global = true, value_name = "IP:PORT")] + pub grpc_addr: Option, /// Run headless, without the interactive TUI. /// /// The load runs until the profile is exhausted or the process receives diff --git a/dwd/src/grpc/client.rs b/dwd/src/grpc/client.rs index 02ad4aa..c1592ad 100644 --- a/dwd/src/grpc/client.rs +++ b/dwd/src/grpc/client.rs @@ -379,4 +379,71 @@ mod tests { is_running.store(false, Ordering::SeqCst); server.abort(); } + + /// Exercises the network endpoint (`--grpc-addr`) over a real TCP socket: + /// remote control (set + stop) must behave exactly like the in-process + /// seam, including flipping the shared run flag on stop. + #[tokio::test] + async fn network_endpoint_roundtrip() { + use core::sync::atomic::{AtomicBool, Ordering}; + + use dwd_core::grpc::{DwdService, EngineDescriptor, SnapshotSource}; + use dwd_proto::{control_request::Kind, ControlRequest, DescribeRequest, SetControl, StopControl}; + + struct FixedSource; + impl SnapshotSource for FixedSource { + fn snapshot(&self) -> StatsSnapshot { + StatsSnapshot::default() + } + } + + let (control_tx, mut control_rx) = tokio::sync::mpsc::channel::(1); + let descriptor = EngineDescriptor { + engine_kind: "http", + groups: vec![StatGroup::Common, StatGroup::Http], + } + .into_response(); + let is_running = Arc::new(AtomicBool::new(true)); + let service = DwdService::new(control_tx, Arc::new(FixedSource), descriptor, is_running.clone()); + + let (addr, server) = dwd_core::grpc::serve_tcp(service, "127.0.0.1:0".parse().expect("loopback")) + .await + .expect("bind the network endpoint"); + + let channel = Endpoint::try_from(format!("http://{addr}")) + .expect("endpoint uri") + .connect() + .await + .expect("connect over TCP"); + let mut client = DwdClient::new(channel); + + // The remote client sees the same engine description as the TUI. + let described = client + .describe(DescribeRequest {}) + .await + .expect("describe") + .into_inner(); + assert_eq!(described.engine_kind, "http"); + + // Remote load control reaches the generator's control channel. + client + .control(ControlRequest { + kind: Some(Kind::Set(SetControl { rps: 9000 })), + }) + .await + .expect("set control"); + let event = control_rx.recv().await.expect("control event delivered"); + assert!(matches!(event, GeneratorEvent::Set(9000))); + + // Remote stop shuts the run down gracefully via the shared flag. + client + .control(ControlRequest { + kind: Some(Kind::Stop(StopControl {})), + }) + .await + .expect("stop control"); + assert!(!is_running.load(Ordering::SeqCst)); + + server.abort(); + } } diff --git a/dwd/src/runtime.rs b/dwd/src/runtime.rs index c733fea..642d82c 100644 --- a/dwd/src/runtime.rs +++ b/dwd/src/runtime.rs @@ -1,9 +1,11 @@ -//! The composition root: wires the core engine, the in-process gRPC seam, and -//! the TUI client together into the single bundled process. +//! The composition root: wires the core engine, the gRPC seam (in-process for +//! the TUI, optionally on a network address for remote clients), and the TUI +//! client together into the single bundled process. //! //! This is the only place that knows about *both* the core (server) and the TUI //! (client). The boundary between them is the [`dwd_proto`] gRPC contract: the -//! TUI never touches the engine directly. +//! TUI never touches the engine directly, and remote clients get the exact same +//! API over `--grpc-addr`. use core::{ future, @@ -69,22 +71,44 @@ impl Runtime { }; // The control channel that `run_generator` polls; the gRPC `Control` - // handler is now the only writer (the TUI reaches it through the API). + // handler is now the only writer (the TUI and remote clients both reach + // it through the API). let (control_tx, control_rx) = mpsc::channel::(1); - // The client half of the process exists only when the TUI runs. Headless - // (`--no-ui`) there is no client to serve, so the in-process gRPC seam is - // not stood up at all, and shutdown comes from process signals instead of - // the UI loop. - let seam = if self.cfg.no_ui { + // The single service instance behind every transport: the in-memory TUI + // seam and the network endpoint are the same API over the same core. + // Fully headless (no TUI, no `--grpc-addr`) there is no client at all, + // so no service is stood up and the control channel closes. + let service = if self.cfg.no_ui && self.cfg.grpc_addr.is_none() { drop(control_tx); + None + } else { + Some(DwdService::new( + control_tx, + snapshot_source, + descriptor, + self.is_running.clone(), + )) + }; + + // The network gRPC endpoint gives remote clients full control over the + // run (set / suspend / resume / stop + the stats stream). Bound eagerly + // so a bad address fails the run instead of a background task. + let grpc_handle = match (self.cfg.grpc_addr, &service) { + (Some(addr), Some(service)) => Some(dwd_core::grpc::serve_tcp(service.clone(), addr).await?.1), + _ => None, + }; + + // Headless runs have no UI loop to translate key presses into shutdown, + // so process signals take that role instead. + let seam = if self.cfg.no_ui { self.spawn_signal_handler(); None } else { // Stand up the in-process gRPC server over the core and connect a client // over the same in-memory pipe — no OS socket is opened. - let service = DwdService::new(control_tx, snapshot_source, descriptor, self.is_running.clone()); + let service = service.clone().expect("service exists whenever the TUI runs"); let (client_io, server_io) = tokio::io::duplex(DUPLEX_BUFFER); let server_handle = dwd_core::grpc::serve(service, server_io); let grpc_client = client::connect(client_io)?; @@ -148,6 +172,9 @@ impl Runtime { // lands on a clean stderr. summary.finish().log(); + if let Some(handle) = grpc_handle { + handle.abort(); + } if let Some(handle) = api_handle { handle.abort(); }