diff --git a/lading/src/bin/payloadtool.rs b/lading/src/bin/payloadtool.rs index 0e90d2a9b..5596b6915 100644 --- a/lading/src/bin/payloadtool.rs +++ b/lading/src/bin/payloadtool.rs @@ -486,6 +486,12 @@ fn check_generator(config: &generator::Config, args: &Args) -> Result { + if args.fingerprint { + return Ok(None); + } + unimplemented!("TcpCrr not supported") + } generator::Inner::TcpRr(_) => { if args.fingerprint { return Ok(None); diff --git a/lading/src/blackhole.rs b/lading/src/blackhole.rs index e23db4c69..1d130598f 100644 --- a/lading/src/blackhole.rs +++ b/lading/src/blackhole.rs @@ -15,6 +15,7 @@ pub mod otlp; pub mod splunk_hec; pub mod sqs; pub mod tcp; +pub mod tcp_crr; pub mod tcp_rr; pub mod udp; pub mod unix_datagram; @@ -26,6 +27,9 @@ pub enum Error { /// See [`crate::blackhole::tcp::Error`] for details. #[error(transparent)] Tcp(tcp::Error), + /// See [`crate::blackhole::tcp_crr::Error`] for details. + #[error(transparent)] + TcpCrr(tcp_crr::Error), /// See [`crate::blackhole::tcp_rr::Error`] for details. #[error(transparent)] TcpRr(tcp_rr::Error), @@ -87,6 +91,8 @@ pub struct General { pub enum Inner { /// See [`crate::blackhole::tcp::Config`] for details. Tcp(tcp::Config), + /// See [`crate::blackhole::tcp_crr::Config`] for details. + TcpCrr(tcp_crr::Config), /// See [`crate::blackhole::tcp_rr::Config`] for details. TcpRr(tcp_rr::Config), /// See [`crate::blackhole::datadog::Config`] for details. @@ -117,6 +123,8 @@ pub enum Inner { pub enum Server { /// See [`crate::blackhole::tcp::Tcp`] for details. Tcp(tcp::Tcp), + /// See [`crate::blackhole::tcp_crr::TcpCrr`] for details. + TcpCrr(tcp_crr::TcpCrr), /// See [`crate::blackhole::tcp_rr::TcpRr`] for details. TcpRr(tcp_rr::TcpRr), /// See [`crate::blackhole::datadog::Datadog`] for details. @@ -152,6 +160,9 @@ impl Server { pub fn new(config: Config, shutdown: lading_signal::Watcher) -> Result { let server = match config.inner { Inner::Tcp(conf) => Self::Tcp(tcp::Tcp::new(config.general, &conf, shutdown)), + Inner::TcpCrr(conf) => { + Self::TcpCrr(tcp_crr::TcpCrr::new(config.general, &conf, shutdown)) + } Inner::TcpRr(conf) => Self::TcpRr(tcp_rr::TcpRr::new(config.general, &conf, shutdown)), Inner::Datadog(conf) => { Self::Datadog(datadog::Datadog::new(config.general, conf, shutdown)) @@ -194,6 +205,7 @@ impl Server { pub async fn run(self) -> Result<(), Error> { match self { Server::Tcp(inner) => inner.run().await.map_err(Error::Tcp), + Server::TcpCrr(inner) => inner.run().await.map_err(Error::TcpCrr), Server::TcpRr(inner) => inner.run().await.map_err(Error::TcpRr), Server::Datadog(inner) => inner.run().await.map_err(Error::Datadog), Server::DatadogStatefulLogs(inner) => { diff --git a/lading/src/blackhole/tcp_crr.rs b/lading/src/blackhole/tcp_crr.rs new file mode 100644 index 000000000..fb437df7c --- /dev/null +++ b/lading/src/blackhole/tcp_crr.rs @@ -0,0 +1,144 @@ +//! TCP connect/request/response (`tcp_crr`) blackhole - the server side. +//! Based on +//! +//! Listens for incoming connections and, for each flow, reads a fixed-size +//! request then writes a fixed-size response. The CRR client closes the +//! connection after each response; the server side is identical to `tcp_rr` +//! and delegates to the same shared machinery. +//! +//! The event-loop machinery lives in [`crate::neper::rr`]; this module is a +//! thin wrapper that supplies configuration. +//! +//! ## Metrics +//! +//! `connections_accepted`: Incoming connections accepted +//! `requests_received`: Completed request reads +//! `responses_sent`: Completed response writes +//! `bytes_received`: Request bytes read +//! `bytes_written`: Response bytes sent +//! `connections_closed`: Flow removals (client close + I/O errors) + +use std::net::{IpAddr, SocketAddr}; +use std::num::{NonZeroU16, NonZeroUsize}; + +use serde::{Deserialize, Serialize}; + +use super::General; +use crate::neper::rr::{self, Mode, ServerParams}; + +fn default_nonzero_u16() -> NonZeroU16 { + NonZeroU16::new(1).expect("1 is nonzero") +} + +fn default_nonzero_usize() -> NonZeroUsize { + NonZeroUsize::new(1).expect("1 is nonzero") +} + +fn default_control_port() -> u16 { + 12866 +} + +fn default_data_port() -> u16 { + 12867 +} + +fn default_backlog() -> i32 { + 1024 +} + +const fn default_true() -> bool { + true +} + +#[derive(Debug, Deserialize, Serialize, Clone, Copy, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +/// Configuration for the `tcp_crr` blackhole. +pub struct Config { + /// IP address to bind on. + pub addr: IpAddr, + /// Data port for flow connections. Default 12867. + #[serde(default = "default_data_port")] + pub data_port: u16, + /// Control port for startup synchronization with the generator. Default 12866. + #[serde(default = "default_control_port")] + pub control_port: u16, + /// Number of OS server threads. Default 1. When > 1, uses `SO_REUSEPORT` + /// with an eBPF program for load balancing. + #[serde(default = "default_nonzero_u16")] + pub threads: NonZeroU16, + /// Total number of TCP flows the generator should open (neper `-F`). + /// Default 1. Sent to the generator over the control connection at + /// startup; the generator does not configure this independently. + #[serde(default = "default_nonzero_u16")] + pub flows: NonZeroU16, + /// Bytes to read per request. Default 1. + #[serde(default = "default_nonzero_usize")] + pub request_size: NonZeroUsize, + /// Bytes to send per response. Default 1. + #[serde(default = "default_nonzero_usize")] + pub response_size: NonZeroUsize, + /// Whether to set `TCP_NODELAY` on accepted connections. Default true. + #[serde(default = "default_true")] + pub no_delay: bool, + /// Listener backlog (pending-connection queue length) passed to `listen(2)`. + /// Default 1024. CRR workloads benefit from a larger backlog to absorb + /// connect bursts. + #[serde(default = "default_backlog")] + pub backlog: i32, +} + +#[derive(thiserror::Error, Debug)] +/// Errors produced by [`TcpCrr`]. +pub enum Error { + /// Shared neper-style request/response error. + #[error(transparent)] + Rr(#[from] rr::Error), +} + +#[derive(Debug)] +/// The `tcp_crr` blackhole (server side). +pub struct TcpCrr { + config: Config, + metric_labels: Vec<(String, String)>, + shutdown: lading_signal::Watcher, +} + +impl TcpCrr { + /// Create a new [`TcpCrr`] blackhole instance. + #[must_use] + pub fn new(general: General, config: &Config, shutdown: lading_signal::Watcher) -> Self { + let mut metric_labels = vec![ + ("component".to_string(), "blackhole".to_string()), + ("component_name".to_string(), "tcp_crr".to_string()), + ]; + if let Some(id) = general.id { + metric_labels.push(("id".to_string(), id)); + } + Self { + config: *config, + metric_labels, + shutdown, + } + } + + /// Run the blackhole to completion or until a shutdown signal is received. + /// + /// # Errors + /// + /// Returns an error if binding fails or a worker thread panics. + pub async fn run(self) -> Result<(), Error> { + let params = ServerParams { + data_addr: SocketAddr::new(self.config.addr, self.config.data_port), + control_addr: SocketAddr::new(self.config.addr, self.config.control_port), + threads: self.config.threads.get(), + flows: self.config.flows.get(), + request_size: self.config.request_size.get(), + response_size: self.config.response_size.get(), + no_delay: self.config.no_delay, + backlog: self.config.backlog, + mode: Mode::Crr, + }; + rr::run_server(params, self.metric_labels, self.shutdown, "tcp_crr").await?; + Ok(()) + } +} diff --git a/lading/src/blackhole/tcp_rr.rs b/lading/src/blackhole/tcp_rr.rs index b81fc513c..9da929ace 100644 --- a/lading/src/blackhole/tcp_rr.rs +++ b/lading/src/blackhole/tcp_rr.rs @@ -1,10 +1,13 @@ -//! TCP request/response (`tcp_rr`) blackhole — the server side. +//! TCP request/response (`tcp_rr`) blackhole - the server side. //! Based on //! //! Listens for incoming connections and, for each flow, reads a fixed-size //! request then writes a fixed-size response, repeating until the flow closes //! or lading shuts down. //! +//! The event-loop machinery lives in [`crate::neper::rr`]; this module is a +//! thin wrapper that supplies configuration. +//! //! ## Metrics //! //! `connections_accepted`: Incoming connections accepted @@ -13,25 +16,13 @@ //! `bytes_received`: Request bytes read //! `bytes_written`: Response bytes sent -use std::io::{ErrorKind, Read, Write}; -use std::net::{self, IpAddr, SocketAddr}; +use std::net::{IpAddr, SocketAddr}; use std::num::{NonZeroU16, NonZeroUsize}; -use std::os::fd::AsRawFd; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering::Relaxed}; -use std::time::Duration; -use mio::net::{TcpListener, TcpStream}; -use mio::{Events, Interest, Poll, Token}; use serde::{Deserialize, Serialize}; -use tokio::sync::mpsc; -use tracing::{info, trace, warn}; use super::General; -use crate::neper::bpf; -use crate::neper::flow::{self, Action, Flow, FlowMap}; -use crate::neper::metrics::{self, ThreadMetrics}; -use crate::neper::thread; +use crate::neper::rr::{self, Mode, ServerParams}; fn default_nonzero_u16() -> NonZeroU16 { NonZeroU16::new(1).expect("1 is nonzero") @@ -70,9 +61,14 @@ pub struct Config { #[serde(default = "default_control_port")] pub control_port: u16, /// Number of OS server threads. Default 1. When > 1, uses `SO_REUSEPORT` - /// with an eBPF program for load balancing + /// with an eBPF program for load balancing. #[serde(default = "default_nonzero_u16")] pub threads: NonZeroU16, + /// Total number of TCP flows the generator should open. + /// Default 1. Sent to the generator over the control connection at + /// startup; the generator does not configure this independently. + #[serde(default = "default_nonzero_u16")] + pub flows: NonZeroU16, /// Bytes to read per request. Default 1. #[serde(default = "default_nonzero_usize")] pub request_size: NonZeroUsize, @@ -91,21 +87,9 @@ pub struct Config { #[derive(thiserror::Error, Debug)] /// Errors produced by [`TcpRr`]. pub enum Error { - /// IO error + /// Shared neper-style request/response error. #[error(transparent)] - Io(#[from] std::io::Error), - /// Error binding TCP listener - #[error("Failed to bind TCP listener to {addr}: {source}")] - Bind { - /// Binding address - addr: SocketAddr, - /// Underlying IO error - #[source] - source: Box, - }, - /// Worker thread panicked - #[error("Worker thread panicked")] - ThreadPanicked, + Rr(#[from] rr::Error), } #[derive(Debug)] @@ -116,13 +100,6 @@ pub struct TcpRr { shutdown: lading_signal::Watcher, } -enum ServerState { - RecvRequest, - SendResponse, -} - -const LISTENER_TOKEN: Token = Token(0); - impl TcpRr { /// Create a new [`TcpRr`] blackhole instance. #[must_use] @@ -146,354 +123,19 @@ impl TcpRr { /// # Errors /// /// Returns an error if binding fails or a worker thread panics. - /// - /// # Panics - /// - /// Panics if the ready-barrier tokio task is cancelled. - #[allow(clippy::too_many_lines)] pub async fn run(self) -> Result<(), Error> { - let shutdown_flag = thread::new_shutdown_flag(); - let num_threads = self.config.threads.get(); - - let thread_metrics = Arc::new( - (0..num_threads) - .map(|_| ThreadMetrics::new()) - .collect::>(), - ); - - let metrics_handle = { - let tm = Arc::clone(&thread_metrics); - let labels = self.metric_labels.clone(); - let flag = Arc::clone(&shutdown_flag); - thread::spawn_named("tcp_rr-bh-metrics", move || { - metrics::run_metrics_thread(&tm, &labels, &flag); - }) - }; - - // Pre-build thread 0's listener here so the BPF program is attached - // to the reuseport group before any other thread calls bind(). This - // removes the need for a cross-thread BPF barrier — if the bind fails - // or panics, it propagates as an error directly from this task. - let binding_addr = SocketAddr::new(self.config.addr, self.config.data_port); - let thread0_listener = if num_threads > 1 { - Some(create_listener( - 0, - num_threads, - binding_addr, - self.config.backlog, - )) - } else { - None + let params = ServerParams { + data_addr: SocketAddr::new(self.config.addr, self.config.data_port), + control_addr: SocketAddr::new(self.config.addr, self.config.control_port), + threads: self.config.threads.get(), + flows: self.config.flows.get(), + request_size: self.config.request_size.get(), + response_size: self.config.response_size.get(), + no_delay: self.config.no_delay, + backlog: self.config.backlog, + mode: Mode::Rr, }; - - // Each thread sends a ready signal via this channel after binding. - // If a thread panics before signaling, its sender drops; once all - // senders are gone, recv() returns None and we detect the failure - // instead of hanging forever. - let (ready_tx, mut ready_rx) = mpsc::unbounded_channel::<()>(); - - let mut handles = Vec::with_capacity(num_threads as usize); - let mut thread0_listener = thread0_listener; - for i in 0..num_threads { - let request_size = self.config.request_size.get(); - let response_size = self.config.response_size.get(); - let no_delay = self.config.no_delay; - let backlog = self.config.backlog; - let flag = Arc::clone(&shutdown_flag); - let tm = Arc::clone(&thread_metrics); - let prebuilt = if i == 0 { - thread0_listener.take() - } else { - None - }; - let tx = ready_tx.clone(); - let handle = thread::spawn_named(&format!("tcp_rr-server-{i}"), move || { - server_thread_main( - i, - num_threads, - binding_addr, - prebuilt, - backlog, - request_size, - response_size, - no_delay, - &flag, - &tm[i as usize], - tx, - ); - }); - handles.push(handle); - } - // Drop our own copy so the channel closes when all worker threads exit. - drop(ready_tx); - - // Wait for each thread to signal ready. If a sender drops without - // signaling (thread panicked), recv() eventually returns None. - for _ in 0..num_threads { - if ready_rx.recv().await.is_none() { - shutdown_flag.store(true, Relaxed); - thread::join_all(handles).map_err(|()| Error::ThreadPanicked)?; - return Err(Error::ThreadPanicked); - } - } - - // All data listeners are up. Open control port so the generator - // can connect and know we're ready. - let control_addr = SocketAddr::new(self.config.addr, self.config.control_port); - let control_listener = - net::TcpListener::bind(control_addr).map_err(|source| Error::Bind { - addr: control_addr, - source: Box::new(source), - })?; - control_listener - .set_nonblocking(true) - .expect("failed to set control listener nonblocking"); - info!("control port listening on {control_addr}, waiting for generator"); - - handles.push(metrics_handle); - - // Accept with shutdown awareness: poll accept in a loop. - let flag = Arc::clone(&shutdown_flag); - let shutdown_clone = self.shutdown.clone(); - tokio::spawn(async move { - shutdown_clone.recv().await; - flag.store(true, Relaxed); - }); - let mut generator_connected = false; - loop { - if shutdown_flag.load(Relaxed) { - info!("shutdown before generator connected"); - break; - } - match control_listener.accept() { - Ok((_conn, peer)) => { - info!("generator connected from {peer}, data threads running"); - generator_connected = true; - break; - } - Err(ref e) if e.kind() == ErrorKind::WouldBlock => { - tokio::time::sleep(Duration::from_millis(100)).await; - } - Err(e) => { - return Err(Error::Bind { - addr: control_addr, - source: Box::new(e), - }); - } - } - } - drop(control_listener); - - if generator_connected { - self.shutdown.recv().await; - info!("shutdown signal received"); - } - shutdown_flag.store(true, Relaxed); - - thread::join_all(handles).map_err(|()| Error::ThreadPanicked)?; - + rr::run_server(params, self.metric_labels, self.shutdown, "tcp_rr").await?; Ok(()) } } - -/// Create a listener socket. When `num_threads` > 1, sets `SO_REUSEPORT` -/// and (for thread 0) attaches the reuseport eBPF program. -fn create_listener( - thread_index: u16, - num_threads: u16, - binding_addr: SocketAddr, - backlog: i32, -) -> net::TcpListener { - let domain = if binding_addr.is_ipv4() { - socket2::Domain::IPV4 - } else { - socket2::Domain::IPV6 - }; - let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP)) - .expect("failed to create socket"); - socket - .set_nonblocking(true) - .expect("failed to set nonblocking"); - socket - .set_cloexec(true) - .expect("failed to set close-on-exec"); - socket - .set_reuse_address(true) - .expect("failed to set SO_REUSEADDR"); - - if num_threads > 1 { - socket - .set_reuse_port(true) - .expect("failed to set SO_REUSEPORT"); - - if thread_index == 0 { - match bpf::load_reuseport_ebpf(u32::from(num_threads)) { - Ok(prog) => { - if let Err(e) = bpf::attach_reuseport_ebpf(socket.as_raw_fd(), &prog) { - warn!("failed to attach reuseport eBPF: {e}, falling back to kernel hash"); - } - } - Err(e) => { - warn!("failed to load reuseport eBPF: {e}, falling back to kernel hash"); - } - } - } - } - - socket - .bind(&binding_addr.into()) - .unwrap_or_else(|e| panic!("failed to bind to {binding_addr}: {e}")); - socket.listen(backlog).expect("failed to listen"); - - socket.into() -} - -#[allow(clippy::too_many_arguments)] -fn server_thread_main( - thread_index: u16, - num_threads: u16, - binding_addr: SocketAddr, - prebuilt_listener: Option, - backlog: i32, - request_size: usize, - response_size: usize, - no_delay: bool, - shutdown_flag: &AtomicBool, - metrics: &ThreadMetrics, - ready_tx: mpsc::UnboundedSender<()>, -) { - // Thread 0 uses the pre-built listener (with BPF already attached); - // others bind their own sockets that join the existing reuseport group. - let std_listener = prebuilt_listener - .unwrap_or_else(|| create_listener(thread_index, num_threads, binding_addr, backlog)); - - // Signal that this thread's listener is bound and ready. If this send - // fails the receiver has gone away (blackhole is shutting down). - let _ = ready_tx.send(()); - drop(ready_tx); - - let mut listener = TcpListener::from_std(std_listener); - let mut poll = Poll::new().expect("failed to create mio::Poll"); - let mut events = Events::with_capacity(256); - - poll.registry() - .register(&mut listener, LISTENER_TOKEN, Interest::READABLE) - .expect("failed to register listener"); - - let mut request_buf = vec![0u8; request_size]; - let response_buf = vec![0u8; response_size]; - let mut flows: FlowMap = FlowMap::new(); - let mut next_token: usize = 1; - - loop { - let _ = poll.poll(&mut events, Some(Duration::from_millis(100))); - if shutdown_flag.load(Relaxed) { - break; - } - - let mut attempts = 0; - for event in &events { - if event.token() == LISTENER_TOKEN { - loop { - match listener.accept() { - Ok((stream, _addr)) => { - set_nodelay_mio(&stream, no_delay); - let token = Token(next_token); - next_token += 1; - let mut mio_stream = stream; - poll.registry() - .register(&mut mio_stream, token, Interest::READABLE) - .expect("failed to register flow"); - flows.insert(Flow { - stream: mio_stream, - token, - state: ServerState::RecvRequest, - xfer: request_size, - }); - metrics.connections_accepted.add(1); - } - Err(ref e) if e.kind() == ErrorKind::WouldBlock => break, - Err(e) => { - if attempts > 2 { - break; - } - warn!("accept error: {e}"); - attempts += 1; - std::thread::sleep(Duration::from_millis(1000)); - } - } - } - } else { - let token = event.token(); - let Some(fl) = flows.get_mut(token) else { - continue; - }; - let action = handle_server_event(fl, &mut request_buf, &response_buf, metrics); - flow::apply_action(action, token, &mut flows, poll.registry()); - } - } - } -} - -/// Set `TCP_NODELAY` on a mio [`TcpStream`] via a borrowed `socket2::SockRef`. -fn set_nodelay_mio(stream: &TcpStream, no_delay: bool) { - let sock = socket2::SockRef::from(stream); - if let Err(e) = sock.set_tcp_nodelay(no_delay) { - trace!("failed to set TCP_NODELAY: {e}"); - } -} - -fn handle_server_event( - flow: &mut Flow, - request_buf: &mut [u8], - response_buf: &[u8], - metrics: &ThreadMetrics, -) -> Action { - match flow.state { - ServerState::RecvRequest => { - let offset = request_buf.len() - flow.xfer; - match flow.stream.read(&mut request_buf[offset..]) { - Ok(0) => Action::Remove, - Ok(n) => { - flow.xfer -= n; - if flow.xfer == 0 { - flow.xfer = response_buf.len(); - flow.state = ServerState::SendResponse; - metrics.requests_received.add(1); - metrics.bytes_received.add(request_buf.len() as u64); - Action::Reregister(Interest::WRITABLE) - } else { - Action::Continue - } - } - Err(e) if e.kind() == ErrorKind::WouldBlock => Action::Continue, - Err(e) => { - trace!("read error: {e}"); - Action::Remove - } - } - } - ServerState::SendResponse => { - let offset = response_buf.len() - flow.xfer; - match flow.stream.write(&response_buf[offset..]) { - Ok(n) => { - flow.xfer -= n; - if flow.xfer == 0 { - flow.xfer = request_buf.len(); - flow.state = ServerState::RecvRequest; - metrics.responses_sent.add(1); - metrics.bytes_written.add(response_buf.len() as u64); - Action::Reregister(Interest::READABLE) - } else { - Action::Continue - } - } - Err(e) if e.kind() == ErrorKind::WouldBlock => Action::Continue, - Err(e) => { - trace!("write error: {e}"); - Action::Remove - } - } - } - } -} diff --git a/lading/src/generator.rs b/lading/src/generator.rs index 0d2744680..f2ea9490f 100644 --- a/lading/src/generator.rs +++ b/lading/src/generator.rs @@ -27,6 +27,7 @@ pub mod process_tree; pub mod procfs; pub mod splunk_hec; pub mod tcp; +pub mod tcp_crr; pub mod tcp_rr; pub mod trace_agent; pub mod udp; @@ -39,6 +40,9 @@ pub enum Error { /// See [`crate::generator::tcp::Error`] for details. #[error(transparent)] Tcp(#[from] tcp::Error), + /// See [`crate::generator::tcp_crr::Error`] for details. + #[error(transparent)] + TcpCrr(#[from] tcp_crr::Error), /// See [`crate::generator::tcp_rr::Error`] for details. #[error(transparent)] TcpRr(#[from] tcp_rr::Error), @@ -115,6 +119,8 @@ pub struct General { pub enum Inner { /// See [`crate::generator::tcp::Config`] for details. Tcp(tcp::Config), + /// See [`crate::generator::tcp_crr::Config`] for details. + TcpCrr(tcp_crr::Config), /// See [`crate::generator::tcp_rr::Config`] for details. TcpRr(tcp_rr::Config), /// See [`crate::generator::udp::Config`] for details. @@ -156,6 +162,8 @@ pub enum Inner { pub enum Server { /// See [`crate::generator::tcp::Tcp`] for details. Tcp(tcp::Tcp), + /// See [`crate::generator::tcp_crr::TcpCrr`] for details. + TcpCrr(tcp_crr::TcpCrr), /// See [`crate::generator::tcp_rr::TcpRr`] for details. TcpRr(tcp_rr::TcpRr), /// See [`crate::generator::udp::Udp`] for details. @@ -201,6 +209,9 @@ impl Server { pub fn new(config: Config, shutdown: lading_signal::Watcher) -> Result { let srv = match config.inner { Inner::Tcp(conf) => Self::Tcp(tcp::Tcp::new(config.general, &conf, shutdown)?), + Inner::TcpCrr(conf) => { + Self::TcpCrr(tcp_crr::TcpCrr::new(config.general, &conf, shutdown)) + } Inner::TcpRr(conf) => Self::TcpRr(tcp_rr::TcpRr::new(config.general, &conf, shutdown)), Inner::Udp(conf) => Self::Udp(udp::Udp::new(config.general, &conf, shutdown)?), Inner::Http(conf) => Self::Http(http::Http::new(config.general, conf, shutdown)?), @@ -276,6 +287,7 @@ impl Server { match self { Server::Tcp(inner) => inner.spin().await?, + Server::TcpCrr(inner) => inner.spin().await?, Server::TcpRr(inner) => inner.spin().await?, Server::Udp(inner) => inner.spin().await?, Server::Http(inner) => inner.spin().await?, diff --git a/lading/src/generator/tcp_crr.rs b/lading/src/generator/tcp_crr.rs new file mode 100644 index 000000000..ceb5e2888 --- /dev/null +++ b/lading/src/generator/tcp_crr.rs @@ -0,0 +1,131 @@ +//! TCP connect/request/response (`tcp_crr`) generator — the client side. +//! Based on +//! +//! Implements neper's `tcp_crr` protocol: each flow connects, sends a +//! fixed-size request, reads a fixed-size response, closes the connection, +//! then reconnects and repeats. This measures connection-establishment rate +//! end-to-end, including kernel and TCP-handshake overhead. +//! +//! The event-loop machinery lives in [`crate::neper::rr`]; this module is a +//! thin wrapper that supplies configuration and selects [`Mode::Crr`]. +//! +//! ## Metrics +//! +//! `connections_initiated`: Successful client-side connect completions +//! `requests_sent`: Completed request writes +//! `responses_received`: Completed response reads +//! `bytes_written`: Request bytes sent +//! `bytes_read`: Response bytes received +//! `connections_failed`: Failed connection attempts + +use std::net::{IpAddr, SocketAddr}; +use std::num::{NonZeroU16, NonZeroUsize}; + +use serde::{Deserialize, Serialize}; + +use super::General; +use crate::generator::common::MetricsBuilder; +use crate::neper::rr::{self, ClientParams, Mode}; + +fn default_nonzero_u16() -> NonZeroU16 { + NonZeroU16::new(1).expect("1 is nonzero") +} + +fn default_nonzero_usize() -> NonZeroUsize { + NonZeroUsize::new(1).expect("1 is nonzero") +} + +const fn default_true() -> bool { + true +} + +fn default_control_port() -> u16 { + 12866 +} + +fn default_data_port() -> u16 { + 12867 +} + +#[derive(Debug, Deserialize, Serialize, PartialEq, Clone)] +#[serde(deny_unknown_fields)] +/// Configuration for the `tcp_crr` generator. +/// +/// Flow count (neper `-F`) is *not* configured here — it is owned by the +/// `tcp_crr` blackhole and communicated to the generator over the control +/// port during startup. +pub struct Config { + /// The IP address of the `tcp_crr` server. + pub addr: String, + /// Data port for flow connections. Default 12867. + #[serde(default = "default_data_port")] + pub data_port: u16, + /// Control port for startup synchronization with the blackhole. Default 12866. + #[serde(default = "default_control_port")] + pub control_port: u16, + /// Number of OS threads (neper -T). Default 1. + #[serde(default = "default_nonzero_u16")] + pub threads: NonZeroU16, + /// Bytes per request. Default 1. + #[serde(default = "default_nonzero_usize")] + pub request_size: NonZeroUsize, + /// Bytes per response to read back. Default 1. + #[serde(default = "default_nonzero_usize")] + pub response_size: NonZeroUsize, + /// Whether to set `TCP_NODELAY` on connections. Default true. + #[serde(default = "default_true")] + pub no_delay: bool, +} + +#[derive(thiserror::Error, Debug)] +/// Errors produced by [`TcpCrr`]. +pub enum Error { + /// Shared neper-style request/response error. + #[error(transparent)] + Rr(#[from] rr::Error), +} + +#[derive(Debug)] +/// The `tcp_crr` generator (client side). +pub struct TcpCrr { + config: Config, + metric_labels: Vec<(String, String)>, + shutdown: lading_signal::Watcher, +} + +impl TcpCrr { + /// Create a new [`TcpCrr`] generator instance. + #[must_use] + pub fn new(general: General, config: &Config, shutdown: lading_signal::Watcher) -> Self { + let metric_labels = MetricsBuilder::new("tcp_crr").with_id(general.id).build(); + Self { + config: config.clone(), + metric_labels, + shutdown, + } + } + + /// Run the generator to completion or until a shutdown signal is received. + /// + /// # Errors + /// + /// Returns an error if a worker thread panics or configuration is invalid. + /// + /// # Panics + /// + /// Panics if `addr` cannot be parsed as an IP address. + pub async fn spin(self) -> Result<(), Error> { + let ip: IpAddr = self.config.addr.parse().expect("invalid addr"); + let params = ClientParams { + data_addr: SocketAddr::new(ip, self.config.data_port), + control_addr: SocketAddr::new(ip, self.config.control_port), + threads: self.config.threads.get(), + request_size: self.config.request_size.get(), + response_size: self.config.response_size.get(), + no_delay: self.config.no_delay, + mode: Mode::Crr, + }; + rr::run_client(params, self.metric_labels, self.shutdown, "tcp_crr").await?; + Ok(()) + } +} diff --git a/lading/src/generator/tcp_rr.rs b/lading/src/generator/tcp_rr.rs index 5a7a3394d..2185e751d 100644 --- a/lading/src/generator/tcp_rr.rs +++ b/lading/src/generator/tcp_rr.rs @@ -5,6 +5,9 @@ //! waits for a fixed-size response, and repeats. Flows are distributed across //! OS threads and multiplexed via mio. //! +//! The event-loop machinery lives in [`crate::neper::rr`]; this module is a +//! thin wrapper that supplies configuration. +//! //! ## Metrics //! //! `requests_sent`: Completed request writes @@ -13,23 +16,14 @@ //! `bytes_read`: Response bytes received //! `connections_failed`: Failed connection attempts -use std::io::{self, ErrorKind, Read, Write}; -use std::net::{self, IpAddr, SocketAddr}; +use std::net::{IpAddr, SocketAddr}; use std::num::{NonZeroU16, NonZeroUsize}; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering::Relaxed}; -use std::time::{Duration, Instant}; -use mio::net::TcpStream; -use mio::{Events, Interest, Poll, Token}; use serde::{Deserialize, Serialize}; -use tracing::{info, trace}; use super::General; use crate::generator::common::MetricsBuilder; -use crate::neper::flow::{self, Action, Flow, FlowMap}; -use crate::neper::metrics::{self, ThreadMetrics}; -use crate::neper::thread; +use crate::neper::rr::{self, ClientParams, Mode}; fn default_nonzero_u16() -> NonZeroU16 { NonZeroU16::new(1).expect("1 is nonzero") @@ -54,6 +48,10 @@ fn default_data_port() -> u16 { #[derive(Debug, Deserialize, Serialize, PartialEq, Clone)] #[serde(deny_unknown_fields)] /// Configuration for the `tcp_rr` generator. +/// +/// Flow count is *not* configured here — it is owned by the +/// `tcp_rr` blackhole and communicated to the generator over the control port +/// during startup. pub struct Config { /// The IP address of the `tcp_rr` server. pub addr: String, @@ -66,9 +64,6 @@ pub struct Config { /// Number of OS threads (neper -T). Default 1. #[serde(default = "default_nonzero_u16")] pub threads: NonZeroU16, - /// Total number of TCP flows/connections (neper -F). Default 1. - #[serde(default = "default_nonzero_u16")] - pub flows: NonZeroU16, /// Bytes per request. Default 1. #[serde(default = "default_nonzero_usize")] pub request_size: NonZeroUsize, @@ -83,15 +78,9 @@ pub struct Config { #[derive(thiserror::Error, Debug)] /// Errors produced by [`TcpRr`]. pub enum Error { - /// IO error + /// Shared neper-style request/response error. #[error(transparent)] - Io(#[from] std::io::Error), - /// Worker thread panicked - #[error("Worker thread panicked")] - ThreadPanicked, - /// Invalid configuration. - #[error("invalid config: {0}")] - Config(String), + Rr(#[from] rr::Error), } #[derive(Debug)] @@ -102,11 +91,6 @@ pub struct TcpRr { shutdown: lading_signal::Watcher, } -enum ClientState { - SendRequest, - RecvResponse, -} - impl TcpRr { /// Create a new [`TcpRr`] generator instance. #[must_use] @@ -123,226 +107,23 @@ impl TcpRr { /// /// # Errors /// - /// Returns an error if a worker thread panics. + /// Returns an error if a worker thread panics or configuration is invalid. /// /// # Panics /// - /// Panics if `addr` cannot be resolved to a socket address. + /// Panics if `addr` cannot be parsed as an IP address. pub async fn spin(self) -> Result<(), Error> { - if self.config.threads > self.config.flows { - return Err(Error::Config(format!( - "threads ({}) must be <= flows ({})", - self.config.threads, self.config.flows - ))); - } - let ip: IpAddr = self.config.addr.parse().expect("invalid addr"); - let data_addr = SocketAddr::new(ip, self.config.data_port); - let control_addr = SocketAddr::new(ip, self.config.control_port); - - let shutdown_flag = thread::new_shutdown_flag(); - - // Wait for the blackhole to be ready by connecting to its control port. - info!("waiting for blackhole control port at {control_addr}"); - let deadline = Instant::now() + Duration::from_secs(300); - { - let flag = Arc::clone(&shutdown_flag); - let shutdown = self.shutdown.clone(); - tokio::spawn(async move { - shutdown.recv().await; - flag.store(true, Relaxed); - }); - } - loop { - if shutdown_flag.load(Relaxed) { - return Err(Error::Io(io::Error::new( - ErrorKind::ConnectionRefused, - format!( - "shutdown before blackhole control port {control_addr} became reachable" - ), - ))); - } - match net::TcpStream::connect(control_addr) { - Ok(_conn) => { - info!("blackhole ready, starting flows"); - break; - } - Err(e) => { - if Instant::now() >= deadline { - return Err(Error::Io(io::Error::new( - ErrorKind::TimedOut, - format!( - "blackhole control port {control_addr} not reachable after 5 minutes: {e}" - ), - ))); - } - std::thread::sleep(Duration::from_millis(100)); - } - } - } - let num_threads = self.config.threads.get(); - let num_flows = self.config.flows.get(); - let request_size = self.config.request_size.get(); - let response_size = self.config.response_size.get(); - - let flow_dist = thread::distribute_flows(num_flows, num_threads); - - let thread_metrics = Arc::new( - (0..num_threads) - .map(|_| ThreadMetrics::new()) - .collect::>(), - ); - - let metrics_handle = { - let tm = Arc::clone(&thread_metrics); - let labels = self.metric_labels.clone(); - let flag = Arc::clone(&shutdown_flag); - thread::spawn_named("tcp_rr-metrics", move || { - metrics::run_metrics_thread(&tm, &labels, &flag); - }) + let params = ClientParams { + data_addr: SocketAddr::new(ip, self.config.data_port), + control_addr: SocketAddr::new(ip, self.config.control_port), + threads: self.config.threads.get(), + request_size: self.config.request_size.get(), + response_size: self.config.response_size.get(), + no_delay: self.config.no_delay, + mode: Mode::Rr, }; - - let mut worker_handles = Vec::with_capacity(num_threads as usize); - for i in 0..num_threads { - let thread_flows = flow_dist[i as usize]; - let flag = Arc::clone(&shutdown_flag); - let tm = Arc::clone(&thread_metrics); - let no_delay = self.config.no_delay; - let handle = thread::spawn_named(&format!("tcp_rr-client-{i}"), move || { - client_thread_main( - data_addr, - thread_flows, - request_size, - response_size, - no_delay, - &flag, - &tm[i as usize], - ); - }); - worker_handles.push(handle); - } - - self.shutdown.recv().await; - info!("shutdown signal received"); - shutdown_flag.store(true, Relaxed); - - worker_handles.push(metrics_handle); - thread::join_all(worker_handles).map_err(|()| Error::ThreadPanicked)?; - + rr::run_client(params, self.metric_labels, self.shutdown, "tcp_rr").await?; Ok(()) } } - -fn client_thread_main( - addr: SocketAddr, - num_flows: u16, - request_size: usize, - response_size: usize, - no_delay: bool, - shutdown_flag: &AtomicBool, - metrics: &ThreadMetrics, -) { - let mut poll = Poll::new().expect("failed to create mio::Poll"); - let mut events = Events::with_capacity(num_flows as usize); - let request_buf = vec![0u8; request_size]; - let mut response_buf = vec![0u8; response_size]; - let mut flows: FlowMap = FlowMap::new(); - let mut next_token: usize = 0; - - for _ in 0..num_flows { - match net::TcpStream::connect(addr) { - Ok(std_stream) => { - let _ = std_stream.set_nodelay(no_delay); - std_stream - .set_nonblocking(true) - .expect("failed to set nonblocking"); - let mut stream = TcpStream::from_std(std_stream); - let token = Token(next_token); - next_token += 1; - poll.registry() - .register(&mut stream, token, Interest::WRITABLE) - .expect("failed to register flow"); - flows.insert(Flow { - stream, - token, - state: ClientState::SendRequest, - xfer: request_size, - }); - } - Err(e) => { - trace!("connection to {addr} failed: {e}"); - metrics.connections_failed.add(1); - } - } - } - - loop { - let _ = poll.poll(&mut events, Some(Duration::from_millis(100))); - if shutdown_flag.load(Relaxed) { - break; - } - for event in &events { - let token = event.token(); - let Some(fl) = flows.get_mut(token) else { - continue; - }; - let action = handle_client_event(fl, &request_buf, &mut response_buf, metrics); - flow::apply_action(action, token, &mut flows, poll.registry()); - } - } -} - -fn handle_client_event( - flow: &mut Flow, - request_buf: &[u8], - response_buf: &mut [u8], - metrics: &ThreadMetrics, -) -> Action { - match flow.state { - ClientState::SendRequest => { - let offset = request_buf.len() - flow.xfer; - match flow.stream.write(&request_buf[offset..]) { - Ok(n) => { - flow.xfer -= n; - if flow.xfer == 0 { - flow.xfer = response_buf.len(); - flow.state = ClientState::RecvResponse; - metrics.requests_sent.add(1); - metrics.bytes_written.add(request_buf.len() as u64); - Action::Reregister(Interest::READABLE) - } else { - Action::Continue - } - } - Err(e) if e.kind() == ErrorKind::WouldBlock => Action::Continue, - Err(e) => { - trace!("write error: {e}"); - Action::Remove - } - } - } - ClientState::RecvResponse => { - let offset = response_buf.len() - flow.xfer; - match flow.stream.read(&mut response_buf[offset..]) { - Ok(0) => Action::Remove, - Ok(n) => { - flow.xfer -= n; - if flow.xfer == 0 { - flow.xfer = request_buf.len(); - flow.state = ClientState::SendRequest; - metrics.responses_received.add(1); - metrics.bytes_read.add(response_buf.len() as u64); - Action::Reregister(Interest::WRITABLE) - } else { - Action::Continue - } - } - Err(e) if e.kind() == ErrorKind::WouldBlock => Action::Continue, - Err(e) => { - trace!("read error: {e}"); - Action::Remove - } - } - } - } -} diff --git a/lading/src/neper.rs b/lading/src/neper.rs index 1a07396ea..befe81fa0 100644 --- a/lading/src/neper.rs +++ b/lading/src/neper.rs @@ -11,4 +11,5 @@ pub(crate) mod bpf; pub(crate) mod bpf; pub(crate) mod flow; pub(crate) mod metrics; +pub(crate) mod rr; pub(crate) mod thread; diff --git a/lading/src/neper/flow.rs b/lading/src/neper/flow.rs index 102a53121..ed8af24df 100644 --- a/lading/src/neper/flow.rs +++ b/lading/src/neper/flow.rs @@ -31,34 +31,67 @@ pub(crate) enum Action { /// Token-indexed flow storage. /// /// Flows are stored in a `Vec` indexed by token value. Removed slots become -/// `None` and are not reused — tokens are monotonically increasing, matching +/// `None` and are not reused - tokens are monotonically increasing, matching /// neper's behavior. pub(crate) struct FlowMap { inner: Vec>>, } +/// Errors produced by `FlowMap`. +#[derive(thiserror::Error, Debug)] +pub(crate) enum FlowMapError { + /// No capacity + #[error("Server flow map is at capacity: {0}")] + NoCapacity(usize), +} + impl FlowMap { - pub(crate) fn new() -> Self { - Self { inner: Vec::new() } + pub(crate) fn new(flows: usize) -> Self { + Self { + inner: Vec::with_capacity(flows * 2), + } } - /// Insert a flow. Grows the backing vec if needed. - pub(crate) fn insert(&mut self, flow: Flow) { - let idx = flow.token.0; - if idx >= self.inner.len() { + /// Insert a flow. + pub(crate) fn insert(&mut self, flow: Flow) -> Result<(), FlowMapError> { + let idx = flow.token.0 % self.inner.capacity(); + if self.inner.len() <= idx { self.inner.resize_with(idx + 1, || None); } - self.inner[idx] = Some(flow); + if self.inner[idx].is_none() { + self.inner[idx] = Some(flow); + return Ok(()); + } + + Err(FlowMapError::NoCapacity(self.inner.capacity())) } /// Get a mutable reference to the flow at the given token. + /// + /// Returns `None` for an empty slot. A slot occupied by a flow with a + /// different token indicates a token-collision bug (tokens congruent + /// modulo capacity share a slot), so it is asserted against. pub(crate) fn get_mut(&mut self, token: Token) -> Option<&mut Flow> { - self.inner.get_mut(token.0).and_then(|slot| slot.as_mut()) + let idx = token.0 % self.inner.capacity(); + let flow = self.inner.get_mut(idx).and_then(|slot| slot.as_mut())?; + assert_eq!( + flow.token, token, + "FlowMap slot {idx} holds a flow with a mismatched token" + ); + Some(flow) } /// Remove and return the flow at the given token. + /// + /// Returns `None` for an empty slot; asserts the occupant's token matches. pub(crate) fn remove(&mut self, token: Token) -> Option> { - self.inner.get_mut(token.0).and_then(Option::take) + let idx = token.0 % self.inner.capacity(); + let flow = self.inner.get_mut(idx).and_then(Option::take)?; + assert_eq!( + flow.token, token, + "FlowMap slot {idx} holds a flow with a mismatched token" + ); + Some(flow) } } @@ -73,12 +106,16 @@ pub(crate) fn apply_action( Action::Continue => {} Action::Reregister(interest) => { if let Some(flow) = flows.get_mut(token) { - let _ = registry.reregister(&mut flow.stream, flow.token, interest); + registry + .reregister(&mut flow.stream, flow.token, interest) + .expect("reregister of a live, owned flow must succeed"); } } Action::Remove => { if let Some(mut flow) = flows.remove(token) { - let _ = registry.deregister(&mut flow.stream); + registry + .deregister(&mut flow.stream) + .expect("deregister of a registered, owned flow must succeed"); } } } diff --git a/lading/src/neper/metrics.rs b/lading/src/neper/metrics.rs index f449c4535..8c3fdab06 100644 --- a/lading/src/neper/metrics.rs +++ b/lading/src/neper/metrics.rs @@ -69,8 +69,10 @@ define_thread_metrics! { responses_received, bytes_written, bytes_read, + connections_initiated, connections_failed, connections_accepted, + connections_closed, requests_received, responses_sent, bytes_received, diff --git a/lading/src/neper/rr.rs b/lading/src/neper/rr.rs new file mode 100644 index 000000000..99c328042 --- /dev/null +++ b/lading/src/neper/rr.rs @@ -0,0 +1,1006 @@ +//! Shared client/server machinery for neper-style request/response workloads. +//! +//! Provides [`run_client`] and [`run_server`] entry points used by `tcp_rr` +//! (and forthcoming `tcp_crr`). The shared code owns the mio event loops, flow +//! lifecycle, control-port synchronization, and per-thread metrics plumbing; +//! per-variant modules build the [`ClientParams`] / [`ServerParams`] and +//! call in. + +use std::io::{self, ErrorKind, Read, Write}; +use std::net::{self, SocketAddr}; +use std::os::fd::AsRawFd; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering::Relaxed}; +use std::time::{Duration, Instant}; + +use mio::net::{TcpListener, TcpStream}; +use mio::{Events, Interest, Poll, Token}; +use tokio::sync::mpsc; +use tracing::{info, warn}; + +use crate::neper::bpf; +use crate::neper::flow::{self, Action, Flow, FlowMap}; +use crate::neper::metrics::{self, ThreadMetrics}; +use crate::neper::thread; + +/// Errors produced by [`run_client`] and [`run_server`]. +#[derive(thiserror::Error, Debug)] +pub enum Error { + /// IO error. + #[error(transparent)] + Io(#[from] std::io::Error), + /// Failed to bind a listener. + #[error("Failed to bind TCP listener to {addr}: {source}")] + Bind { + /// Binding address. + addr: SocketAddr, + /// Underlying IO error. + #[source] + source: Box, + }, + /// A worker thread panicked. + #[error("Worker thread panicked")] + ThreadPanicked, + /// Invalid configuration. + #[error("invalid config: {0}")] + Config(String), +} + +/// Which neper-style protocol the client is driving. +#[derive(Clone, Copy, Debug)] +pub(crate) enum Mode { + /// `tcp_rr`: persistent connection, request/response loop forever. + Rr, + /// `tcp_crr`: connect, request/response, close, reconnect, repeat. + Crr, +} + +/// Parameters for [`run_client`]. +/// +/// Flow count is *not* a client parameter - it is owned by the server and +/// communicated to the client over the control connection during startup. +pub(crate) struct ClientParams { + /// Address of the server's data port. + pub(crate) data_addr: SocketAddr, + /// Address of the server's control port. + pub(crate) control_addr: SocketAddr, + /// Number of OS threads. + pub(crate) threads: u16, + /// Bytes per request. + pub(crate) request_size: usize, + /// Bytes per response. + pub(crate) response_size: usize, + /// Whether to set `TCP_NODELAY`. + pub(crate) no_delay: bool, + /// RR or CRR. + pub(crate) mode: Mode, +} + +/// Parameters for [`run_server`]. +pub(crate) struct ServerParams { + /// Address to bind the data listener on. + pub(crate) data_addr: SocketAddr, + /// Address to bind the control listener on. + pub(crate) control_addr: SocketAddr, + /// Number of OS server threads. + pub(crate) threads: u16, + /// Total number of TCP flows the client should open. Sent to the client + /// over the control connection during startup. + pub(crate) flows: u16, + /// Bytes to read per request. + pub(crate) request_size: usize, + /// Bytes to send per response. + pub(crate) response_size: usize, + /// Whether to set `TCP_NODELAY` on accepted connections. + pub(crate) no_delay: bool, + /// Listener backlog. + pub(crate) backlog: i32, + /// RR or CRR. + pub(crate) mode: Mode, +} + +enum ClientState { + /// CRR only: waiting for a non-blocking `connect(2)` to complete. The + /// `WRITABLE` readiness event signals connect completion; `take_error()` + /// then tells us if it succeeded. + Connecting, + SendRequest, + RecvResponse, +} + +/// Actions returned by [`handle_client_event`]. Supersets [`flow::Action`] +/// with [`ClientAction::Reconnect`] for CRR's per-transaction reconnect. +#[derive(Clone, Copy)] +enum ClientAction { + Continue, + Reregister(Interest), + /// CRR: response complete - close this socket and open a new one with the + /// same `Token`. Handled by [`apply_client_action`]. + Reconnect, + Remove, +} + +enum ServerState { + RecvRequest, + SendResponse, + CloseStream, +} + +const LISTENER_TOKEN: Token = Token(0); + +/// Control-channel handshake: server writes `flows` to the accepted control +/// connection as a 2-byte big-endian `u16` and closes; client reads the same +/// 2 bytes after connecting. Internal protocol - no magic / version byte. +const HANDSHAKE_LEN: usize = 2; +const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(5); + +/// Run the neper-style client (generator side). +/// +/// Connects `flows` TCP flows distributed across `threads` OS threads, then +/// runs a request/response loop on each until `shutdown` fires. +/// +/// `thread_prefix` is used to name OS threads (`{prefix}-metrics`, +/// `{prefix}-client-{i}`) so multiple variants can coexist in `top -H`. +/// +/// # Errors +/// +/// Returns an error if configuration is invalid, the blackhole control port +/// is never reachable, or a worker thread panics. +#[allow(clippy::too_many_lines)] +pub(crate) async fn run_client( + params: ClientParams, + metric_labels: Vec<(String, String)>, + shutdown: lading_signal::Watcher, + thread_prefix: &'static str, +) -> Result<(), Error> { + let shutdown_flag = thread::new_shutdown_flag(); + + // Wait for the blackhole to be ready by connecting to its control port, + // then read the flow count over that connection. + info!( + "waiting for blackhole control port at {}", + params.control_addr + ); + let deadline = Instant::now() + Duration::from_secs(300); + { + let flag = Arc::clone(&shutdown_flag); + let shutdown = shutdown.clone(); + tokio::spawn(async move { + shutdown.recv().await; + flag.store(true, Relaxed); + }); + } + let flows: u16 = loop { + if shutdown_flag.load(Relaxed) { + return Err(Error::Io(io::Error::new( + ErrorKind::ConnectionRefused, + format!( + "shutdown before blackhole control port {} became reachable", + params.control_addr + ), + ))); + } + match generator_connect_blocking(params.control_addr) { + Ok(mut conn) => { + conn.set_read_timeout(Some(HANDSHAKE_TIMEOUT)) + .expect("set_read_timeout on connected TcpStream must succeed"); + let mut buf = [0u8; HANDSHAKE_LEN]; + conn.read_exact(&mut buf)?; + let received = u16::from_be_bytes(buf); + info!("blackhole ready, {received} flows to open"); + break received; + } + Err(e) => { + if Instant::now() >= deadline { + return Err(Error::Io(io::Error::new( + ErrorKind::TimedOut, + format!( + "blackhole control port {} not reachable after 5 minutes: {e}", + params.control_addr + ), + ))); + } + std::thread::sleep(Duration::from_millis(100)); + } + } + }; + + if params.threads > flows { + return Err(Error::Config(format!( + "threads ({}) must be <= flows received from blackhole ({flows})", + params.threads + ))); + } + + let flow_dist = thread::distribute_flows(flows, params.threads); + + let thread_metrics = Arc::new( + (0..params.threads) + .map(|_| ThreadMetrics::new()) + .collect::>(), + ); + + let metrics_handle = { + let tm = Arc::clone(&thread_metrics); + let labels = metric_labels.clone(); + let flag = Arc::clone(&shutdown_flag); + thread::spawn_named(&format!("{thread_prefix}-metrics"), move || { + metrics::run_metrics_thread(&tm, &labels, &flag); + }) + }; + + let data_addr = params.data_addr; + let request_size = params.request_size; + let response_size = params.response_size; + let no_delay = params.no_delay; + let mode = params.mode; + let mut worker_handles = Vec::with_capacity(params.threads as usize); + for i in 0..params.threads { + let thread_flows = flow_dist[i as usize]; + let flag = Arc::clone(&shutdown_flag); + let tm = Arc::clone(&thread_metrics); + let handle = thread::spawn_named(&format!("{thread_prefix}-client-{i}"), move || { + client_thread_main( + data_addr, + thread_flows, + request_size, + response_size, + no_delay, + mode, + &flag, + &tm[i as usize], + ); + }); + worker_handles.push(handle); + } + + shutdown.recv().await; + info!("shutdown signal received"); + shutdown_flag.store(true, Relaxed); + + worker_handles.push(metrics_handle); + thread::join_all(worker_handles).map_err(|()| Error::ThreadPanicked)?; + + Ok(()) +} + +/// `IP_LOCAL_PORT_RANGE` socket option (Linux >= 6.3). Not yet exposed by the +/// `libc` crate, so it is defined here from ``. +const IP_LOCAL_PORT_RANGE: libc::c_int = 51; +/// Lowest ephemeral local port the generator may use for its sockets. +const LOCAL_PORT_LOW: u16 = 1024; +/// Highest ephemeral local port the generator may use for its sockets. +const LOCAL_PORT_HIGH: u16 = 60999; + +/// Set once the kernel is found not to support `IP_LOCAL_PORT_RANGE`, so the +/// "unsupported, falling back" warning is logged a single time rather than on +/// every socket the generator opens. +static PORT_RANGE_UNSUPPORTED: AtomicBool = AtomicBool::new(false); + +/// Increase `socket`'s automatic source-port selection to +/// `[LOCAL_PORT_LOW, LOCAL_PORT_HIGH]` via `IP_LOCAL_PORT_RANGE`. The option +/// value packs the high port in the upper 16 bits and the low port in the +/// lower 16 bits. +/// +/// This is done to reduce `EADDRNOTAVAIL` errors when a large number of flows are +/// created especially for `tcp_crr` workload. +/// Since port ranges are specific to network namespaces, this should not cause issues +/// for other daemons coming online on lower port ranges when lading is launched in its own +/// namespace. +fn set_local_port_range(socket: &socket2::Socket) -> io::Result<()> { + let value: u32 = (u32::from(LOCAL_PORT_HIGH) << 16) | u32::from(LOCAL_PORT_LOW); + // SAFETY: `socket` owns a valid fd for the duration of the borrow, and we + // pass a pointer to a correctly sized `u32` as the option value, exactly as + // `IP_LOCAL_PORT_RANGE` expects. + let ret = unsafe { + libc::setsockopt( + socket.as_raw_fd(), + libc::IPPROTO_IP, + IP_LOCAL_PORT_RANGE, + std::ptr::addr_of!(value).cast::(), + std::mem::size_of::() + .try_into() + .expect("u32 size fits in socklen_t"), + ) + }; + if ret != 0 { + let err = io::Error::last_os_error(); + // ENOPROTOOPT / EOPNOTSUPP means the running kernel predates + // IP_LOCAL_PORT_RANGE (< 6.3). Degrade gracefully: fall back to the + // system-wide ephemeral range rather than failing the connection. + if matches!( + err.raw_os_error(), + Some(libc::ENOPROTOOPT | libc::EOPNOTSUPP) + ) { + if !PORT_RANGE_UNSUPPORTED.swap(true, Relaxed) { + warn!( + "IP_LOCAL_PORT_RANGE not supported by this kernel; \ + falling back to the system ephemeral port range" + ); + } + return Ok(()); + } + return Err(err); + } + Ok(()) +} + +/// Create a TCP socket for the generator with its local port range constrained +/// to `[LOCAL_PORT_LOW, LOCAL_PORT_HIGH]`. +fn new_generator_socket(addr: SocketAddr) -> io::Result { + let socket = socket2::Socket::new( + socket2::Domain::for_address(addr), + socket2::Type::STREAM, + Some(socket2::Protocol::TCP), + )?; + set_local_port_range(&socket)?; + Ok(socket) +} + +/// Blocking connect to `addr` using a port-range-constrained generator socket. +/// +/// The kernel picks the source port from the constrained ephemeral range (see +/// [`set_local_port_range`]); the generator does not manage source ports itself. +fn generator_connect_blocking(addr: SocketAddr) -> io::Result { + let socket = new_generator_socket(addr)?; + socket.connect(&addr.into())?; + Ok(net::TcpStream::from(socket)) +} + +/// Non-blocking connect to `addr` using a port-range-constrained generator +/// socket, returning a mio stream whose connect is in progress (completion is +/// signalled by a `WRITABLE` readiness event). +/// +/// The kernel picks the source port from the constrained ephemeral range (see +/// [`set_local_port_range`]); the generator does not manage source ports itself. +fn generator_connect_nonblocking(addr: SocketAddr) -> io::Result { + let socket = new_generator_socket(addr)?; + socket.set_nonblocking(true)?; + // A non-blocking connect reports in-progress as EINPROGRESS / WouldBlock; + // that is expected and not an error. + match socket.connect(&addr.into()) { + Ok(()) => {} + Err(e) if e.raw_os_error() == Some(libc::EINPROGRESS) => {} + Err(e) if e.kind() == ErrorKind::WouldBlock => {} + Err(e) => return Err(e), + } + Ok(TcpStream::from_std(net::TcpStream::from(socket))) +} + +#[allow(clippy::too_many_arguments)] +fn client_thread_main( + addr: SocketAddr, + num_flows: u16, + request_size: usize, + response_size: usize, + no_delay: bool, + mode: Mode, + shutdown_flag: &AtomicBool, + metrics: &ThreadMetrics, +) { + let mut poll = Poll::new().expect("failed to create mio::Poll"); + let mut events = Events::with_capacity(num_flows as usize); + let request_buf = vec![0u8; request_size]; + let mut response_buf = vec![0u8; response_size]; + let mut flows: FlowMap = FlowMap::new(num_flows as usize); + let mut next_token: usize = 0; + + for _ in 0..num_flows { + let token = Token(next_token); + match generator_connect_blocking(addr) { + Ok(std_stream) => { + let _ = std_stream.set_nodelay(no_delay); + std_stream + .set_nonblocking(true) + .expect("failed to set nonblocking"); + let mut stream = TcpStream::from_std(std_stream); + next_token += 1; + poll.registry() + .register(&mut stream, token, Interest::WRITABLE) + .expect("failed to register flow"); + flows + .insert(Flow { + stream, + token, + state: ClientState::SendRequest, + xfer: request_size, + }) + .expect("client should never be able to exceed FlowMap capacity"); + metrics.connections_initiated.add(1); + } + Err(e) => { + warn!("connection to {addr} failed: {e}"); + metrics.connections_failed.add(1); + } + } + } + + loop { + let _ = poll.poll(&mut events, Some(Duration::from_millis(100))); + if shutdown_flag.load(Relaxed) { + break; + } + for event in &events { + let token = event.token(); + let Some(fl) = flows.get_mut(token) else { + continue; + }; + let action = handle_client_event(fl, mode, &request_buf, &mut response_buf, metrics); + apply_client_action(action, token, &mut flows, &poll, addr, no_delay, metrics); + } + } +} + +/// Apply a [`ClientAction`] to the flow map. Handles the CRR +/// reconnect transition (deregister old stream, open a new +/// non-blocking connect, reregister with the same token). +fn apply_client_action( + action: ClientAction, + token: Token, + flows: &mut FlowMap, + poll: &Poll, + addr: SocketAddr, + no_delay: bool, + metrics: &ThreadMetrics, +) { + let registry = poll.registry(); + match action { + ClientAction::Continue => {} + ClientAction::Reregister(interest) => { + if let Some(flow) = flows.get_mut(token) { + let _ = registry.reregister(&mut flow.stream, flow.token, interest); + } + } + ClientAction::Reconnect => { + // Take the flow out of the map so the old socket is fully closed + // before the new connection is opened. + let Some(mut flow) = flows.remove(token) else { + return; + }; + let _ = registry.deregister(&mut flow.stream); + { + // Abortive close: SO_LINGER with a zero timeout makes the drop + // below emit a RST instead of a FIN, so the socket skips + // TIME_WAIT and its source port returns to the constrained + // ephemeral range immediately rather than lingering 2*MSL. + let sock = socket2::SockRef::from(&flow.stream); + if let Err(e) = sock.set_linger(Some(Duration::from_secs(0))) { + warn!("failed to set SO_LINGER for abortive close on reconnect: {e}"); + } + } + // Drop the old stream now (sends RST) before the new connect. + drop(flow); + match generator_connect_nonblocking(addr) { + Ok(mut new_stream) => { + { + let sock = socket2::SockRef::from(&new_stream); + if let Err(e) = sock.set_tcp_nodelay(no_delay) { + warn!("failed to set TCP_NODELAY on reconnect: {e}"); + } + } + if let Err(e) = registry.register(&mut new_stream, token, Interest::WRITABLE) { + warn!("reconnect register failed: {e}"); + metrics.connections_failed.add(1); + } else if let Err(err) = flows.insert(Flow { + stream: new_stream, + token, + state: ClientState::Connecting, + xfer: 0, + }) { + warn!("failed to reinsert reconnected flow: {err}"); + metrics.connections_failed.add(1); + } + } + Err(e) => { + warn!("reconnect to {addr} failed: {e}"); + metrics.connections_failed.add(1); + } + } + } + ClientAction::Remove => { + metrics.connections_closed.add(1); + if let Some(mut flow) = flows.remove(token) { + let _ = registry.deregister(&mut flow.stream); + // Abortive close (RST) so the source port skips TIME_WAIT and + // is reclaimed immediately rather than lingering 2*MSL. + let sock = socket2::SockRef::from(&flow.stream); + if let Err(e) = sock.set_linger(Some(Duration::from_secs(0))) { + warn!("failed to set SO_LINGER for abortive close: {e}"); + } + // `flow` (and its stream fd) dropped here -> RST sent. + } + } + } +} + +fn handle_client_event( + flow: &mut Flow, + mode: Mode, + request_buf: &[u8], + response_buf: &mut [u8], + metrics: &ThreadMetrics, +) -> ClientAction { + // Connecting -> SendRequest transition: mio is edge-triggered, so the + // single WRITABLE event that signaled connect completion is also the + // event that must drive the first write. Transition state and fall + // through to SendRequest in the same call. + if matches!(flow.state, ClientState::Connecting) { + match flow.stream.take_error() { + Ok(None) => { + flow.state = ClientState::SendRequest; + flow.xfer = request_buf.len(); + metrics.connections_initiated.add(1); + // fall through + } + Ok(Some(e)) => { + warn!("connect failed: {e}"); + metrics.connections_failed.add(1); + return ClientAction::Reconnect; + } + Err(e) => { + warn!("take_error failed: {e}"); + metrics.connections_failed.add(1); + return ClientAction::Reconnect; + } + } + } + + match flow.state { + ClientState::Connecting => unreachable!("transitioned out of Connecting above"), + ClientState::SendRequest => { + let offset = request_buf.len() - flow.xfer; + match flow.stream.write(&request_buf[offset..]) { + Ok(n) => { + flow.xfer -= n; + if flow.xfer == 0 { + flow.xfer = response_buf.len(); + flow.state = ClientState::RecvResponse; + metrics.requests_sent.add(1); + metrics.bytes_written.add(request_buf.len() as u64); + ClientAction::Reregister(Interest::READABLE) + } else { + ClientAction::Continue + } + } + Err(e) if e.kind() == ErrorKind::WouldBlock => ClientAction::Continue, + Err(e) => { + warn!("write error: {e}"); + ClientAction::Remove + } + } + } + ClientState::RecvResponse => { + let offset = response_buf.len() - flow.xfer; + match flow.stream.read(&mut response_buf[offset..]) { + Ok(0) => ClientAction::Remove, + Ok(n) => { + flow.xfer -= n; + if flow.xfer == 0 { + flow.xfer = request_buf.len(); + flow.state = ClientState::SendRequest; + metrics.responses_received.add(1); + metrics.bytes_read.add(response_buf.len() as u64); + match mode { + Mode::Rr => { + flow.xfer = request_buf.len(); + flow.state = ClientState::SendRequest; + ClientAction::Reregister(Interest::WRITABLE) + } + Mode::Crr => ClientAction::Reconnect, + } + } else { + ClientAction::Continue + } + } + Err(e) if e.kind() == ErrorKind::WouldBlock => ClientAction::Continue, + Err(e) => { + warn!("read error: {e}"); + ClientAction::Remove + } + } + } + } +} + +/// Run the neper-style server (blackhole side). +/// +/// Binds a data listener (with `SO_REUSEPORT` + reuseport eBPF when +/// `threads > 1`), then accepts and services request/response flows until +/// `shutdown` fires. +/// +/// `thread_prefix` is used to name OS threads (`{prefix}-bh-metrics`, +/// `{prefix}-server-{i}`). +/// +/// # Errors +/// +/// Returns an error if binding fails or a worker thread panics. +/// +/// # Panics +/// +/// Panics if the ready-barrier tokio task is cancelled. +#[allow(clippy::too_many_lines)] +pub(crate) async fn run_server( + params: ServerParams, + metric_labels: Vec<(String, String)>, + shutdown: lading_signal::Watcher, + thread_prefix: &'static str, +) -> Result<(), Error> { + let shutdown_flag = thread::new_shutdown_flag(); + let num_threads = params.threads; + + let thread_metrics = Arc::new( + (0..num_threads) + .map(|_| ThreadMetrics::new()) + .collect::>(), + ); + + let metrics_handle = { + let tm = Arc::clone(&thread_metrics); + let labels = metric_labels.clone(); + let flag = Arc::clone(&shutdown_flag); + thread::spawn_named(&format!("{thread_prefix}-bh-metrics"), move || { + metrics::run_metrics_thread(&tm, &labels, &flag); + }) + }; + + // Pre-build thread 0's listener here so the BPF program is attached to the + // reuseport group before any other thread calls bind(). This removes the + // need for a cross-thread BPF barrier - if bind fails or panics, it + // propagates as an error directly from this task. + let binding_addr = params.data_addr; + let thread0_listener = if num_threads > 1 { + Some(create_listener( + 0, + num_threads, + binding_addr, + params.backlog, + )) + } else { + None + }; + + // Each thread sends a ready signal via this channel after binding. If a + // thread panics before signaling, its sender drops; once all senders are + // gone, recv() returns None and we detect the failure instead of hanging + // forever. + let (ready_tx, mut ready_rx) = mpsc::unbounded_channel::<()>(); + + let mut handles = Vec::with_capacity(num_threads as usize); + let mut thread0_listener = thread0_listener; + let flows = params.flows; + for i in 0..num_threads { + let request_size = params.request_size; + let response_size = params.response_size; + let no_delay = params.no_delay; + let backlog = params.backlog; + let flag = Arc::clone(&shutdown_flag); + let tm = Arc::clone(&thread_metrics); + let prebuilt = if i == 0 { + thread0_listener.take() + } else { + None + }; + let tx = ready_tx.clone(); + let handle = thread::spawn_named(&format!("{thread_prefix}-server-{i}"), move || { + server_thread_main( + i, + num_threads, + binding_addr, + prebuilt, + backlog, + flows, + request_size, + response_size, + no_delay, + &flag, + &tm[i as usize], + tx, + params.mode, + ); + }); + handles.push(handle); + } + // Drop our own copy so the channel closes when all worker threads exit. + drop(ready_tx); + + // Wait for each thread to signal ready. If a sender drops without + // signaling (thread panicked), recv() eventually returns None. + for _ in 0..num_threads { + if ready_rx.recv().await.is_none() { + shutdown_flag.store(true, Relaxed); + thread::join_all(handles).map_err(|()| Error::ThreadPanicked)?; + return Err(Error::ThreadPanicked); + } + } + + // All data listeners are up. Open control port so the generator can + // connect and know we're ready. + let control_addr = params.control_addr; + let control_listener = net::TcpListener::bind(control_addr).map_err(|source| Error::Bind { + addr: control_addr, + source: Box::new(source), + })?; + control_listener + .set_nonblocking(true) + .expect("failed to set control listener nonblocking"); + info!("control port listening on {control_addr}, waiting for generator"); + + handles.push(metrics_handle); + + let flag = Arc::clone(&shutdown_flag); + let shutdown_clone = shutdown.clone(); + tokio::spawn(async move { + shutdown_clone.recv().await; + flag.store(true, Relaxed); + }); + let mut generator_connected = false; + let flows_bytes = params.flows.to_be_bytes(); + loop { + if shutdown_flag.load(Relaxed) { + info!("shutdown before generator connected"); + break; + } + match control_listener.accept() { + Ok((mut conn, peer)) => { + // accept(2) on Linux returns a blocking socket regardless of + // the listener's O_NONBLOCK; a small write_timeout guards + // against a generator that connects but never reads. + conn.set_write_timeout(Some(HANDSHAKE_TIMEOUT)) + .expect("set_write_timeout on accepted TcpStream must succeed"); + conn.write_all(&flows_bytes)?; + info!( + "generator connected from {peer}, sent flows={}, data threads running", + params.flows + ); + generator_connected = true; + break; + } + Err(ref e) if e.kind() == ErrorKind::WouldBlock => { + tokio::time::sleep(Duration::from_millis(100)).await; + } + Err(e) => { + return Err(Error::Bind { + addr: control_addr, + source: Box::new(e), + }); + } + } + } + drop(control_listener); + + if generator_connected { + shutdown.recv().await; + info!("shutdown signal received"); + } + shutdown_flag.store(true, Relaxed); + + thread::join_all(handles).map_err(|()| Error::ThreadPanicked)?; + + Ok(()) +} + +/// Create a listener socket. When `num_threads` > 1, sets `SO_REUSEPORT` +/// and (for thread 0) attaches the reuseport eBPF program. +fn create_listener( + thread_index: u16, + num_threads: u16, + binding_addr: SocketAddr, + backlog: i32, +) -> net::TcpListener { + let domain = if binding_addr.is_ipv4() { + socket2::Domain::IPV4 + } else { + socket2::Domain::IPV6 + }; + let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP)) + .expect("failed to create socket"); + socket + .set_nonblocking(true) + .expect("failed to set nonblocking"); + socket + .set_cloexec(true) + .expect("failed to set close-on-exec"); + socket + .set_reuse_address(true) + .expect("failed to set SO_REUSEADDR"); + + if num_threads > 1 { + socket + .set_reuse_port(true) + .expect("failed to set SO_REUSEPORT"); + + if thread_index == 0 { + match bpf::load_reuseport_ebpf(u32::from(num_threads)) { + Ok(prog) => { + if let Err(e) = bpf::attach_reuseport_ebpf(socket.as_raw_fd(), &prog) { + warn!("failed to attach reuseport eBPF: {e}, falling back to kernel hash"); + } + } + Err(e) => { + warn!("failed to load reuseport eBPF: {e}, falling back to kernel hash"); + } + } + } + } + + socket + .bind(&binding_addr.into()) + .unwrap_or_else(|e| panic!("failed to bind to {binding_addr}: {e}")); + socket.listen(backlog).expect("failed to listen"); + + socket.into() +} + +#[allow(clippy::too_many_arguments)] +fn server_thread_main( + thread_index: u16, + num_threads: u16, + binding_addr: SocketAddr, + prebuilt_listener: Option, + backlog: i32, + num_flows: u16, + request_size: usize, + response_size: usize, + no_delay: bool, + shutdown_flag: &AtomicBool, + metrics: &ThreadMetrics, + ready_tx: mpsc::UnboundedSender<()>, + mode: Mode, +) { + // Thread 0 uses the pre-built listener (with BPF already attached); others + // bind their own sockets that join the existing reuseport group. + let std_listener = prebuilt_listener + .unwrap_or_else(|| create_listener(thread_index, num_threads, binding_addr, backlog)); + + // Signal that this thread's listener is bound and ready. If this send + // fails the receiver has gone away (blackhole is shutting down). + let _ = ready_tx.send(()); + drop(ready_tx); + + let mut listener = TcpListener::from_std(std_listener); + let mut poll = Poll::new().expect("failed to create mio::Poll"); + // Worst case under SO_REUSEPORT: every flow lands on this thread, so size + // for the total flow count plus the listener token. + let mut events = Events::with_capacity(num_flows as usize + 1); + + poll.registry() + .register(&mut listener, LISTENER_TOKEN, Interest::READABLE) + .expect("failed to register listener"); + + let mut request_buf = vec![0u8; request_size]; + let response_buf = vec![0u8; response_size]; + let mut flows: FlowMap = FlowMap::new(num_flows as usize); + let mut next_token: usize = 1; + + loop { + let _ = poll.poll(&mut events, Some(Duration::from_millis(100))); + if shutdown_flag.load(Relaxed) { + break; + } + + let mut attempts = 0; + for event in &events { + if event.token() == LISTENER_TOKEN { + loop { + match listener.accept() { + Ok((stream, _addr)) => { + set_nodelay_mio(&stream, no_delay); + let token = Token(next_token); + next_token += 1; + // Insert first: `insert` takes the flow by value and + // drops it (closing the fd) on failure, so there is + // nothing registered to clean up on the error path. + if let Err(err) = flows.insert(Flow { + stream, + token, + state: ServerState::RecvRequest, + xfer: request_size, + }) { + warn!( + "failed to insert flow in server FlowMap: {err} {0}", + token.0 + ); + break; + } + let flow = flows.get_mut(token).expect("flow was just inserted"); + poll.registry() + .register(&mut flow.stream, token, Interest::READABLE) + .expect("failed to register flow"); + metrics.connections_accepted.add(1); + } + Err(ref e) if e.kind() == ErrorKind::WouldBlock => break, + Err(e) => { + if attempts > 2 { + break; + } + warn!("accept error: {e}"); + attempts += 1; + std::thread::sleep(Duration::from_millis(1000)); + } + } + } + } else { + let token = event.token(); + let Some(fl) = flows.get_mut(token) else { + continue; + }; + let action = + handle_server_event(fl, &mut request_buf, &response_buf, metrics, mode); + flow::apply_action(action, token, &mut flows, poll.registry()); + } + } + } +} + +/// Set `TCP_NODELAY` on a mio [`TcpStream`] via a borrowed `socket2::SockRef`. +fn set_nodelay_mio(stream: &TcpStream, no_delay: bool) { + let sock = socket2::SockRef::from(stream); + if let Err(e) = sock.set_tcp_nodelay(no_delay) { + warn!("failed to set TCP_NODELAY: {e}"); + } +} + +fn handle_server_event( + flow: &mut Flow, + request_buf: &mut [u8], + response_buf: &[u8], + metrics: &ThreadMetrics, + mode: Mode, +) -> Action { + match flow.state { + ServerState::RecvRequest => { + let offset = request_buf.len() - flow.xfer; + match flow.stream.read(&mut request_buf[offset..]) { + Ok(0) => { + metrics.connections_closed.add(1); + Action::Remove + } + Ok(n) => { + flow.xfer -= n; + if flow.xfer == 0 { + flow.xfer = response_buf.len(); + flow.state = ServerState::SendResponse; + metrics.requests_received.add(1); + metrics.bytes_received.add(request_buf.len() as u64); + Action::Reregister(Interest::WRITABLE) + } else { + Action::Continue + } + } + Err(e) if e.kind() == ErrorKind::WouldBlock => Action::Continue, + Err(e) => { + warn!("read error: {e}"); + metrics.connections_closed.add(1); + Action::Remove + } + } + } + ServerState::SendResponse => { + let offset = response_buf.len() - flow.xfer; + match flow.stream.write(&response_buf[offset..]) { + Ok(n) => { + flow.xfer -= n; + if flow.xfer == 0 { + flow.xfer = request_buf.len(); + match mode { + Mode::Rr => flow.state = ServerState::RecvRequest, + Mode::Crr => flow.state = ServerState::CloseStream, + } + metrics.responses_sent.add(1); + metrics.bytes_written.add(response_buf.len() as u64); + Action::Reregister(Interest::READABLE) + } else { + Action::Continue + } + } + Err(e) if e.kind() == ErrorKind::WouldBlock => Action::Continue, + Err(e) => { + warn!("write error: {e}"); + metrics.connections_closed.add(1); + Action::Remove + } + } + } + ServerState::CloseStream => Action::Remove, + } +}