use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use crate::aux_egress::AuxEgress;
use crate::aux_mux::{self, AuxChannel};
use crate::aux_rpc::{self, RpcMethod};
use crate::frame::HEADER_SIZE;
use serde::Serialize;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{UnixListener, UnixStream};
use tokio::sync::{oneshot, Mutex, Notify};
pub const RPC_DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
const IPC_WRITE_TIMEOUT: Duration = Duration::from_millis(200);
const IPC_CONNECT_TIMEOUT: Duration = Duration::from_millis(200);
const RPC_RETRY_DELAYS: [Duration; 4] = [
Duration::from_millis(150),
Duration::from_millis(250),
Duration::from_millis(400),
Duration::from_millis(700),
];
const MAX_PENDING_CALLS: usize = 1024;
const REQUEST_FRAGMENT_PACING: Duration = Duration::from_millis(5);
pub const DEFAULT_RESPONSE_SOCK: &str = "/run/ados/aux-rpc-responses.sock";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RpcResponseOwned {
pub status: u16,
pub body: Vec<u8>,
pub headers: Vec<aux_rpc::ResponseHeader>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RpcError {
TooLarge,
Encode,
Send(String),
Timeout,
Incomplete { received: usize, total: u16 },
ChannelClosed,
Busy,
}
impl std::fmt::Display for RpcError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::TooLarge => write!(f, "request body exceeds the relay limit"),
Self::Encode => write!(f, "request field too long to encode"),
Self::Send(e) => write!(f, "aux egress send failed: {e}"),
Self::Timeout => write!(f, "no response before timeout"),
Self::Incomplete { received, total } => {
write!(f, "relay response incomplete: {received}/{total} fragments")
}
Self::ChannelClosed => write!(f, "aux consumer channel closed before response"),
Self::Busy => write!(f, "relay proxy has too many calls in flight"),
}
}
}
impl std::error::Error for RpcError {}
struct PendingCall {
sender: oneshot::Sender<RpcResponseOwned>,
expected_sender: Option<Box<str>>,
assembly: Option<Assembly>,
response_started: bool,
}
struct Assembly {
status: u16,
total: u16,
received: usize,
accepted: u64,
decoder: aux_rpc::ResponseDecoder,
}
#[derive(Default)]
pub struct RpcProxyCounters {
calls_started: AtomicU64,
calls_ok: AtomicU64,
calls_timeout: AtomicU64,
calls_incomplete: AtomicU64,
calls_send_error: AtomicU64,
calls_channel_closed: AtomicU64,
calls_busy: AtomicU64,
retransmits: AtomicU64,
retransmit_send_errors: AtomicU64,
fragments_received: AtomicU64,
fragments_duplicate: AtomicU64,
fragments_wrong_sender: AtomicU64,
fragments_bad_symbol: AtomicU64,
responses_out_of_range: AtomicU64,
pending_now: AtomicU64,
}
#[derive(Debug, Default, Clone, Copy, Serialize, PartialEq, Eq)]
pub struct RpcProxyStats {
pub calls_started: u64,
pub calls_ok: u64,
pub calls_timeout: u64,
pub calls_incomplete: u64,
pub calls_send_error: u64,
pub calls_channel_closed: u64,
pub calls_busy: u64,
pub retransmits: u64,
pub retransmit_send_errors: u64,
pub fragments_received: u64,
pub fragments_duplicate: u64,
pub fragments_wrong_sender: u64,
pub fragments_bad_symbol: u64,
pub responses_out_of_range: u64,
pub pending_now: u64,
}
#[derive(Clone)]
pub struct AuxRpcProxy {
egress: Arc<AuxEgress>,
pending: Arc<Mutex<HashMap<u32, PendingCall>>>,
next_id: Arc<AtomicU32>,
counters: Arc<RpcProxyCounters>,
timeout: Duration,
}
impl AuxRpcProxy {
pub fn new(egress: AuxEgress) -> Self {
Self::with_timeout(egress, RPC_DEFAULT_TIMEOUT)
}
pub fn with_timeout(egress: AuxEgress, timeout: Duration) -> Self {
let seed = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.subsec_nanos() ^ std::process::id())
.unwrap_or(1)
| 1;
Self {
egress: Arc::new(egress),
pending: Arc::new(Mutex::new(HashMap::new())),
next_id: Arc::new(AtomicU32::new(seed)),
counters: Arc::new(RpcProxyCounters::default()),
timeout,
}
}
pub async fn call(
&self,
target: &[u8],
method: RpcMethod,
path: &[u8],
body: &[u8],
) -> Result<RpcResponseOwned, RpcError> {
self.call_request(
target,
method,
&aux_rpc::RequestParts {
path,
body,
..Default::default()
},
)
.await
}
pub async fn call_request(
&self,
target: &[u8],
method: RpcMethod,
parts: &aux_rpc::RequestParts<'_>,
) -> Result<RpcResponseOwned, RpcError> {
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
self.call_with_id(id, target, method, parts).await
}
async fn send_fragments(&self, fragments: &[Vec<u8>]) -> Result<(), String> {
let last = fragments.len().saturating_sub(1);
for (index, fragment) in fragments.iter().enumerate() {
self.egress
.send(AuxChannel::Request, fragment)
.await
.map_err(|e| e.to_string())?;
if index < last {
tokio::time::sleep(REQUEST_FRAGMENT_PACING).await;
}
}
Ok(())
}
async fn call_with_id(
&self,
id: u32,
target: &[u8],
method: RpcMethod,
parts: &aux_rpc::RequestParts<'_>,
) -> Result<RpcResponseOwned, RpcError> {
let fragments =
aux_rpc::encode_request(method, id, target, parts).map_err(|e| match e {
aux_rpc::RequestEncodeError::BodyTooLarge => RpcError::TooLarge,
aux_rpc::RequestEncodeError::FieldTooLong => RpcError::Encode,
})?;
let expected_sender = (!target.is_empty()).then(|| {
String::from_utf8_lossy(target)
.into_owned()
.into_boxed_str()
});
let (tx, rx) = oneshot::channel::<RpcResponseOwned>();
{
let mut pending = self.pending.lock().await;
if pending.len() >= MAX_PENDING_CALLS {
self.counters.calls_busy.fetch_add(1, Ordering::Relaxed);
return Err(RpcError::Busy);
}
pending.insert(
id,
PendingCall {
sender: tx,
expected_sender,
assembly: None,
response_started: false,
},
);
}
self.counters.pending_now.fetch_add(1, Ordering::Relaxed);
self.counters.calls_started.fetch_add(1, Ordering::Relaxed);
let mut guard = PendingGuard {
pending: Arc::clone(&self.pending),
counters: Arc::clone(&self.counters),
id,
armed: true,
};
if let Err(e) = self.send_fragments(&fragments).await {
guard.disarm();
self.remove_pending(id).await;
self.counters
.calls_send_error
.fetch_add(1, Ordering::Relaxed);
return Err(RpcError::Send(e));
}
let deadline = tokio::time::Instant::now() + self.timeout;
let mut attempt = 0usize;
tokio::pin!(rx);
loop {
let wake = match RPC_RETRY_DELAYS.get(attempt).copied() {
Some(d) => (tokio::time::Instant::now() + d).min(deadline),
None => deadline,
};
tokio::select! {
biased;
res = &mut rx => {
guard.disarm();
return match res {
Ok(resp) => {
self.counters.calls_ok.fetch_add(1, Ordering::Relaxed);
Ok(resp)
}
Err(_) => {
self.remove_pending(id).await;
self.counters
.calls_channel_closed
.fetch_add(1, Ordering::Relaxed);
Err(RpcError::ChannelClosed)
}
};
}
_ = tokio::time::sleep_until(wake) => {
if tokio::time::Instant::now() >= deadline {
guard.disarm();
let partial = self.remove_pending(id).await;
return match partial.and_then(|p| p.assembly) {
Some(a) => {
self.counters
.calls_incomplete
.fetch_add(1, Ordering::Relaxed);
Err(RpcError::Incomplete {
received: a.received,
total: a.total,
})
}
None => {
self.counters.calls_timeout.fetch_add(1, Ordering::Relaxed);
Err(RpcError::Timeout)
}
};
}
attempt += 1;
let started = self
.pending
.lock()
.await
.get(&id)
.map(|p| p.response_started)
.unwrap_or(false);
if !started {
match self.send_fragments(&fragments).await {
Ok(()) => {
self.counters.retransmits.fetch_add(1, Ordering::Relaxed);
}
Err(e) => {
self.counters
.retransmit_send_errors
.fetch_add(1, Ordering::Relaxed);
tracing::warn!(
request_id = id,
attempt,
error = %e,
"aux_rpc_retransmit_send_failed"
);
}
}
}
}
}
}
}
async fn remove_pending(&self, id: u32) -> Option<PendingCall> {
let removed = self.pending.lock().await.remove(&id);
if removed.is_some() {
self.counters.pending_now.fetch_sub(1, Ordering::Relaxed);
}
removed
}
pub async fn dispatch_response(&self, response: &aux_rpc::RpcResponse<'_>) {
self.counters
.fragments_received
.fetch_add(1, Ordering::Relaxed);
let mut pending = self.pending.lock().await;
let Some(mut call) = pending.remove(&response.id) else {
return;
};
if call
.expected_sender
.as_deref()
.is_some_and(|expected| response.sender != expected.as_bytes())
{
self.counters
.fragments_wrong_sender
.fetch_add(1, Ordering::Relaxed);
pending.insert(response.id, call);
return;
}
call.response_started = true;
if call.assembly.is_none() {
let bounded =
response.total > 0 && response.total as usize <= aux_rpc::MAX_RESPONSE_FRAGMENTS;
let seeded = bounded
.then(|| aux_rpc::ResponseDecoder::new(response.oti))
.flatten();
let Some(decoder) = seeded else {
self.counters
.responses_out_of_range
.fetch_add(1, Ordering::Relaxed);
self.counters.pending_now.fetch_sub(1, Ordering::Relaxed);
let _ = call.sender.send(RpcResponseOwned {
status: 502,
body: b"relay response geometry out of range".to_vec(),
headers: Vec::new(),
});
return;
};
call.assembly = Some(Assembly {
status: response.status,
total: response.total,
received: 0,
accepted: 0,
decoder,
});
}
let index = response.index as usize;
let bit = 1u64 << (index % 64);
let assembly = call.assembly.as_mut().expect("seeded above");
let fresh = assembly.total == response.total
&& index < assembly.total as usize
&& assembly.accepted & bit == 0;
let outcome = fresh.then(|| assembly.decoder.push(response.index, response.body));
if matches!(outcome, Some(aux_rpc::FragmentOutcome::Pending)) {
assembly.accepted |= bit;
assembly.received += 1;
}
let status = assembly.status;
match outcome {
None => {
self.counters
.fragments_duplicate
.fetch_add(1, Ordering::Relaxed);
pending.insert(response.id, call);
}
Some(aux_rpc::FragmentOutcome::BadSymbol) => {
self.counters
.fragments_bad_symbol
.fetch_add(1, Ordering::Relaxed);
pending.insert(response.id, call);
}
Some(aux_rpc::FragmentOutcome::Pending) => {
pending.insert(response.id, call);
}
Some(aux_rpc::FragmentOutcome::Complete(object)) => {
self.counters.pending_now.fetch_sub(1, Ordering::Relaxed);
let (status, body, headers) = aux_rpc::unpack_response(status, object);
let _ = call.sender.send(RpcResponseOwned {
status,
body,
headers,
});
}
}
}
pub fn stats(&self) -> RpcProxyStats {
let c = &self.counters;
RpcProxyStats {
calls_started: c.calls_started.load(Ordering::Relaxed),
calls_ok: c.calls_ok.load(Ordering::Relaxed),
calls_timeout: c.calls_timeout.load(Ordering::Relaxed),
calls_incomplete: c.calls_incomplete.load(Ordering::Relaxed),
calls_send_error: c.calls_send_error.load(Ordering::Relaxed),
calls_channel_closed: c.calls_channel_closed.load(Ordering::Relaxed),
calls_busy: c.calls_busy.load(Ordering::Relaxed),
retransmits: c.retransmits.load(Ordering::Relaxed),
retransmit_send_errors: c.retransmit_send_errors.load(Ordering::Relaxed),
fragments_received: c.fragments_received.load(Ordering::Relaxed),
fragments_duplicate: c.fragments_duplicate.load(Ordering::Relaxed),
fragments_wrong_sender: c.fragments_wrong_sender.load(Ordering::Relaxed),
fragments_bad_symbol: c.fragments_bad_symbol.load(Ordering::Relaxed),
responses_out_of_range: c.responses_out_of_range.load(Ordering::Relaxed),
pending_now: c.pending_now.load(Ordering::Relaxed),
}
}
pub async fn reset(&self) {
let mut pending = self.pending.lock().await;
for (_, call) in pending.drain() {
let _ = call.sender.send(RpcResponseOwned {
status: 503,
body: Vec::new(),
headers: Vec::new(),
});
}
self.counters.pending_now.store(0, Ordering::Relaxed);
}
pub fn spawn_response_listener(
&self,
sock_path: impl Into<PathBuf>,
cancel: Arc<Notify>,
) -> tokio::task::JoinHandle<()> {
let proxy = self.clone();
let sock_path = sock_path.into();
tokio::spawn(async move {
run_response_listener(proxy, sock_path, cancel).await;
})
}
}
struct PendingGuard {
pending: Arc<Mutex<HashMap<u32, PendingCall>>>,
counters: Arc<RpcProxyCounters>,
id: u32,
armed: bool,
}
impl PendingGuard {
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for PendingGuard {
fn drop(&mut self) {
if !self.armed {
return;
}
let pending = Arc::clone(&self.pending);
let counters = Arc::clone(&self.counters);
let id = self.id;
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
if pending.lock().await.remove(&id).is_some() {
counters.pending_now.fetch_sub(1, Ordering::Relaxed);
}
});
}
}
}
async fn run_response_listener(proxy: AuxRpcProxy, sock_path: PathBuf, cancel: Arc<Notify>) {
let _ = std::fs::remove_file(&sock_path);
let listener = match UnixListener::bind(&sock_path) {
Ok(l) => l,
Err(e) => {
tracing::error!(error = %e, path = %sock_path.display(), "aux_rpc_response_listener_bind_failed");
return;
}
};
tracing::info!(path = %sock_path.display(), "aux_rpc_response_listener_listening");
loop {
tokio::select! {
biased;
_ = cancel.notified() => break,
accept_result = listener.accept() => {
match accept_result {
Ok((stream, _)) => {
let proxy = proxy.clone();
tokio::spawn(handle_one_connection(stream, proxy));
}
Err(e) => {
tracing::warn!(error = %e, "aux_rpc_response_listener_accept_failed");
}
}
}
}
}
proxy.reset().await;
tracing::info!("aux_rpc_response_listener_stopped");
let _ = std::fs::remove_file(&sock_path);
}
async fn handle_one_connection(mut stream: UnixStream, proxy: AuxRpcProxy) {
loop {
match read_one_frame(&mut stream).await {
Ok(payload) => match aux_rpc::decode_response(&payload) {
Ok(response) => proxy.dispatch_response(&response).await,
Err(e) => {
tracing::debug!(error = ?e, len = payload.len(), "aux_rpc_response_undecodable");
}
},
Err(e) => {
if e.kind() != std::io::ErrorKind::UnexpectedEof {
tracing::debug!(error = %e, "aux_rpc_response_reader_read_failed");
}
break;
}
}
}
}
async fn read_one_frame(stream: &mut UnixStream) -> std::io::Result<Vec<u8>> {
let mut header = [0u8; HEADER_SIZE];
stream.read_exact(&mut header).await?;
let len = u32::from_be_bytes(header) as usize;
if len > aux_mux::AUX_MAX_PAYLOAD {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"response frame exceeds the aux payload maximum",
));
}
let mut payload = vec![0u8; len];
if len > 0 {
stream.read_exact(&mut payload).await?;
}
Ok(payload)
}
#[derive(Clone)]
pub struct AuxRpcResponseIngest {
sock_path: PathBuf,
conn: Arc<Mutex<Option<UnixStream>>>,
}
impl AuxRpcResponseIngest {
pub fn new(sock_path: impl Into<PathBuf>) -> Self {
Self {
sock_path: sock_path.into(),
conn: Arc::new(Mutex::new(None)),
}
}
pub async fn send(&self, payload: &[u8]) {
let len = payload.len() as u32;
let mut frame = Vec::with_capacity(HEADER_SIZE + payload.len());
frame.extend_from_slice(&len.to_be_bytes());
frame.extend_from_slice(payload);
let mut guard = self.conn.lock().await;
match guard.as_mut() {
Some(stream) => {
let wrote = tokio::time::timeout(IPC_WRITE_TIMEOUT, stream.write_all(&frame)).await;
if !matches!(wrote, Ok(Ok(()))) {
if wrote.is_err() {
tracing::warn!(
timeout_ms = IPC_WRITE_TIMEOUT.as_millis() as u64,
"aux_rpc_ingest_write_timed_out"
);
}
*guard = None;
}
}
None => {
match tokio::time::timeout(
IPC_CONNECT_TIMEOUT,
UnixStream::connect(&self.sock_path),
)
.await
{
Ok(Ok(mut stream)) => {
let wrote =
tokio::time::timeout(IPC_WRITE_TIMEOUT, stream.write_all(&frame)).await;
if matches!(wrote, Ok(Ok(()))) {
*guard = Some(stream);
}
}
Ok(Err(_)) => {
}
Err(_) => {
tracing::warn!(
timeout_ms = IPC_CONNECT_TIMEOUT.as_millis() as u64,
"aux_rpc_ingest_connect_timed_out"
);
}
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
async fn deaf_listener(path: &std::path::Path) -> tokio::task::JoinHandle<()> {
let listener = tokio::net::UnixListener::bind(path).unwrap();
tokio::spawn(async move {
let mut held = Vec::new();
loop {
match listener.accept().await {
Ok((s, _)) => held.push(s),
Err(_) => return,
}
}
})
}
#[tokio::test]
async fn a_peer_that_stops_reading_cannot_stall_the_caller() {
let dir = tempfile::tempdir().unwrap();
let sock = dir.path().join("ingest.sock");
let _listener = deaf_listener(&sock).await;
let ingest = AuxRpcResponseIngest::new(&sock);
let big = vec![0u8; 512 * 1024];
let started = std::time::Instant::now();
for _ in 0..8 {
ingest.send(&big).await;
}
let elapsed = started.elapsed();
assert!(
elapsed < Duration::from_secs(5),
"sends took {elapsed:?}; a wedged peer is stalling the caller"
);
}
#[tokio::test]
async fn an_absent_listener_returns_promptly() {
let ingest =
AuxRpcResponseIngest::new(std::path::Path::new("/nonexistent/ados/ingest.sock"));
let started = std::time::Instant::now();
ingest.send(b"payload").await;
assert!(started.elapsed() < Duration::from_secs(1));
}
#[test]
fn the_ipc_bounds_are_far_below_the_call_bound() {
assert!(IPC_WRITE_TIMEOUT < RPC_DEFAULT_TIMEOUT / 10);
assert!(IPC_CONNECT_TIMEOUT < RPC_DEFAULT_TIMEOUT / 10);
}
use crate::aux_mux;
use tokio::net::UdpSocket;
use tokio::time::Duration;
async fn loopback_egress() -> (
AuxEgress,
tokio::sync::mpsc::Receiver<(AuxChannel, Vec<u8>)>,
) {
let listener = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let port = listener.local_addr().unwrap().port();
let sock = UdpSocket::bind("127.0.0.1:0").await.unwrap();
sock.connect(("127.0.0.1", port)).await.unwrap();
let egress = AuxEgress::connected_for_test(sock);
let (tx, rx) = tokio::sync::mpsc::channel::<(AuxChannel, Vec<u8>)>(8);
tokio::spawn(async move {
let mut buf = [0u8; 2048];
while let Ok((n, _)) = listener.recv_from(&mut buf).await {
if let Ok((ch, payload)) = aux_mux::decode(&buf[..n]) {
let _ = tx.send((ch, payload.to_vec())).await;
}
}
});
(egress, rx)
}
const DRONE: &[u8] = b"0a1b2c3d4e5f";
fn response_payloads(sender: &[u8], id: u32, status: u16, body: &[u8]) -> Vec<Vec<u8>> {
let split = aux_rpc::split_response(body).unwrap();
let total = split.symbols.len() as u16;
split
.symbols
.iter()
.enumerate()
.map(|(index, symbol)| {
aux_rpc::encode_response_fragment(
sender,
id,
status,
index as u16,
total,
split.oti,
symbol,
)
.unwrap()
})
.collect()
}
async fn deliver(proxy: &AuxRpcProxy, payload: &[u8]) {
let fragment = aux_rpc::decode_response(payload).unwrap();
proxy.dispatch_response(&fragment).await;
}
async fn deliver_body(proxy: &AuxRpcProxy, sender: &[u8], id: u32, status: u16, body: &[u8]) {
for payload in response_payloads(sender, id, status, body) {
deliver(proxy, &payload).await;
}
}
async fn pending_id(proxy: &AuxRpcProxy) -> u32 {
let pending = proxy.pending.lock().await;
*pending.keys().next().expect("one pending caller")
}
#[tokio::test]
async fn a_call_round_trips_through_egress_and_a_dispatched_response() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(500));
let proxy_for_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
let (channel, payload) = sent.recv().await.expect("no request datagram");
assert_eq!(channel, AuxChannel::Request);
let fragment = aux_rpc::decode_request_fragment(&payload).unwrap();
assert_eq!(fragment.total, 1, "a small request is one datagram");
let aux_rpc::ReassemblyOutcome::Complete(request) =
aux_rpc::RequestReassembler::new().push(&fragment, std::time::Instant::now())
else {
panic!("a one-fragment request completes on arrival");
};
assert_eq!(request.method, RpcMethod::Get);
assert_eq!(request.target, DRONE);
assert_eq!(request.path, b"/api/pairing/info");
assert!(request.body.is_empty());
let body = br#"{"device_id":"abc"}"#;
deliver_body(&proxy_for_consumer, DRONE, request.id, 200, body).await;
});
let result = proxy
.call(DRONE, RpcMethod::Get, b"/api/pairing/info", &[])
.await
.expect("call must succeed");
consumer.await.unwrap();
assert_eq!(result.status, 200);
assert_eq!(result.body, br#"{"device_id":"abc"}"#);
}
#[tokio::test]
async fn a_multi_fragment_response_completes_on_its_systematic_symbols_alone() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(500));
let body: Vec<u8> = (0..3000u32).map(|i| (i % 251) as u8).collect();
let expected = body.clone();
let proxy_call = proxy.clone();
let call = tokio::spawn(async move {
proxy_call
.call(DRONE, RpcMethod::Get, b"/api/services", &[])
.await
});
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
let fragments = response_payloads(DRONE, id, 200, &body);
assert_eq!(fragments.len(), 3 + aux_rpc::RPC_REPAIR_SYMBOLS as usize);
for fragment in &fragments[..3] {
deliver(&proxy, fragment).await;
}
let result = call.await.unwrap().expect("call must succeed");
assert_eq!(result.status, 200);
assert_eq!(result.body, expected);
}
#[tokio::test]
async fn fragments_delivered_out_of_order_reassemble_identically() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(500));
let body: Vec<u8> = (0..3000u32).map(|i| (i % 251) as u8).collect();
let expected = body.clone();
let proxy_call = proxy.clone();
let call = tokio::spawn(async move {
proxy_call
.call(DRONE, RpcMethod::Get, b"/api/services", &[])
.await
});
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
let mut fragments = response_payloads(DRONE, id, 200, &body);
fragments.reverse();
for fragment in &fragments {
deliver(&proxy, fragment).await;
}
let result = call.await.unwrap().expect("call must succeed");
assert_eq!(result.body, expected);
}
#[tokio::test]
async fn a_thirty_kilobyte_response_completes_after_losing_four_fragments() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_secs(2));
let body: Vec<u8> = (0..29_339u32).map(|i| (i % 251) as u8).collect();
let expected = body.clone();
let proxy_call = proxy.clone();
let call = tokio::spawn(async move {
proxy_call
.call(DRONE, RpcMethod::Get, b"/api/status/full", &[])
.await
});
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
let mut fragments = response_payloads(DRONE, id, 200, &body);
assert_eq!(fragments.len(), 30, "26 systematic + 4 repair");
fragments.rotate_left(7);
fragments.truncate(26);
for fragment in &fragments {
deliver(&proxy, fragment).await;
}
let result = call.await.unwrap().expect("26 of 30 must be enough");
assert_eq!(result.body, expected);
assert_eq!(proxy.stats().calls_ok, 1);
}
#[tokio::test]
async fn a_fragment_from_another_drone_is_dropped_and_counted() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(300));
let proxy_call = proxy.clone();
let call =
tokio::spawn(
async move { proxy_call.call(DRONE, RpcMethod::Get, b"/api/x", &[]).await },
);
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
let imposter = response_payloads(b"ados-elsewhere", id, 200, b"not-your-answer");
for fragment in &imposter {
deliver(&proxy, fragment).await;
}
let err = call.await.unwrap().unwrap_err();
let stats = proxy.stats();
assert_eq!(
err,
RpcError::Timeout,
"an imposter must not complete a call"
);
assert_eq!(stats.fragments_wrong_sender, imposter.len() as u64);
assert_eq!(stats.calls_ok, 0);
assert!(
stats.retransmits > 0,
"another drone's answer proves nothing about ours, so the resend \
loop must keep running"
);
}
#[tokio::test]
async fn the_addressed_drone_still_completes_the_call_after_an_imposter() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_secs(2));
let proxy_call = proxy.clone();
let call =
tokio::spawn(
async move { proxy_call.call(DRONE, RpcMethod::Get, b"/api/x", &[]).await },
);
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
for fragment in &response_payloads(b"ados-elsewhere", id, 200, b"wrong-drone") {
deliver(&proxy, fragment).await;
}
deliver_body(&proxy, DRONE, id, 200, b"right-drone").await;
let result = call.await.unwrap().expect("the addressed drone answered");
assert_eq!(result.body, b"right-drone");
}
#[tokio::test]
async fn a_broadcast_call_accepts_whichever_drone_answers() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(500));
let proxy_call = proxy.clone();
let call =
tokio::spawn(async move { proxy_call.call(&[], RpcMethod::Get, b"/api/x", &[]).await });
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
deliver_body(&proxy, b"ados-anyone", id, 200, b"answered").await;
let result = call.await.unwrap().expect("call must succeed");
assert_eq!(result.body, b"answered");
assert_eq!(proxy.stats().fragments_wrong_sender, 0);
}
#[tokio::test]
async fn a_duplicate_fragment_does_not_double_count() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(400));
let body: Vec<u8> = (0..2000u32).map(|i| (i % 251) as u8).collect();
let expected = body.clone();
let proxy_call = proxy.clone();
let call =
tokio::spawn(
async move { proxy_call.call(DRONE, RpcMethod::Get, b"/api/x", &[]).await },
);
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
let fragments = response_payloads(DRONE, id, 200, &body);
for i in [0usize, 0, 1] {
deliver(&proxy, &fragments[i]).await;
}
let result = call.await.unwrap().expect("call must succeed");
assert_eq!(result.body, expected);
assert_eq!(proxy.stats().fragments_duplicate, 1);
}
#[tokio::test]
async fn a_missing_fragment_yields_incomplete_not_timeout() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(120));
let body: Vec<u8> = (0..3000u32).map(|i| (i % 251) as u8).collect();
let proxy_for_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
let fragments = response_payloads(DRONE, id, 200, &body);
deliver(&proxy_for_consumer, &fragments[0]).await;
});
let err = proxy
.call(DRONE, RpcMethod::Get, b"/api/x", &[])
.await
.unwrap_err();
consumer.await.unwrap();
assert_eq!(
err,
RpcError::Incomplete {
received: 1,
total: 7
}
);
assert_eq!(
err.to_string(),
"relay response incomplete: 1/7 fragments",
"the operator must be able to tell a dropped fragment from silence"
);
}
#[tokio::test]
async fn a_fragment_count_past_the_ceiling_completes_with_502() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(400));
let proxy_for_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
proxy_for_consumer
.dispatch_response(&aux_rpc::RpcResponse {
id,
sender: DRONE,
status: 200,
index: 0,
total: 65,
oti: 2,
body: b"aa",
})
.await;
});
let result = proxy
.call(DRONE, RpcMethod::Get, b"/api/x", &[])
.await
.expect("the caller must be completed, not left hanging");
consumer.await.unwrap();
assert_eq!(result.status, 502);
assert_eq!(result.body, b"relay response geometry out of range");
assert_eq!(proxy.stats().responses_out_of_range, 1);
}
#[tokio::test]
async fn a_transfer_length_past_the_ceiling_completes_with_502() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(400));
let proxy_for_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
proxy_for_consumer
.dispatch_response(&aux_rpc::RpcResponse {
id,
sender: DRONE,
status: 200,
index: 0,
total: 2,
oti: u32::MAX,
body: b"aa",
})
.await;
});
let result = proxy
.call(DRONE, RpcMethod::Get, b"/api/x", &[])
.await
.expect("the caller must be completed, not left hanging");
consumer.await.unwrap();
assert_eq!(result.status, 502);
assert_eq!(proxy.stats().responses_out_of_range, 1);
}
#[tokio::test]
async fn a_symbol_of_the_wrong_size_is_dropped_and_counted() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(200));
let proxy_for_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
proxy_for_consumer
.dispatch_response(&aux_rpc::RpcResponse {
id,
sender: DRONE,
status: 200,
index: 0,
total: 5,
oti: 64,
body: b"only-nine",
})
.await;
});
let err = proxy
.call(DRONE, RpcMethod::Get, b"/api/x", &[])
.await
.unwrap_err();
consumer.await.unwrap();
assert_eq!(
err,
RpcError::Incomplete {
received: 0,
total: 5
}
);
assert_eq!(proxy.stats().fragments_bad_symbol, 1);
}
#[tokio::test]
async fn a_call_times_out_when_no_response_arrives() {
let sock = UdpSocket::bind("127.0.0.1:0").await.unwrap();
sock.connect(("127.0.0.1", 9u16)).await.unwrap();
let egress = AuxEgress::connected_for_test(sock);
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(80));
let started = std::time::Instant::now();
let err = proxy
.call(&[], RpcMethod::Get, b"/api/status", &[])
.await
.unwrap_err();
assert_eq!(err, RpcError::Timeout);
assert!(started.elapsed() < Duration::from_secs(1));
}
#[tokio::test]
async fn a_call_returns_send_error_when_egress_fails() {
let egress =
AuxEgress::with_timeout("/nonexistent/aux-cmd.sock", Duration::from_millis(40));
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(200));
let err = proxy
.call(&[], RpcMethod::Get, b"/api/status", &[])
.await
.unwrap_err();
assert!(
matches!(err, RpcError::Send(_)),
"expected Send, got {err:?}"
);
}
#[tokio::test]
async fn a_body_past_the_relay_limit_is_refused_before_sending() {
let (egress, _sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(200));
let big_body = vec![b'A'; aux_rpc::MAX_REQUEST_BODY + 1];
let err = proxy
.call(&[], RpcMethod::Post, b"/api/x", &big_body)
.await
.unwrap_err();
assert_eq!(err, RpcError::TooLarge);
}
#[tokio::test]
async fn a_body_larger_than_one_frame_goes_out_as_several_fragments() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(500));
let body = vec![b'B'; 4 * aux_mux::AUX_MAX_PAYLOAD];
let expected = body.clone();
let proxy_for_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
let mut reassembler = aux_rpc::RequestReassembler::new();
loop {
let (_, payload) = sent.recv().await.expect("request fragment");
let fragment = aux_rpc::decode_request_fragment(&payload).unwrap();
assert!(fragment.total > 1);
if let aux_rpc::ReassemblyOutcome::Complete(request) =
reassembler.push(&fragment, std::time::Instant::now())
{
assert_eq!(request.method, RpcMethod::Patch);
assert_eq!(request.body, expected);
deliver_body(&proxy_for_consumer, DRONE, request.id, 200, b"ok").await;
return;
}
}
});
let result = proxy
.call(DRONE, RpcMethod::Patch, b"/api/plugins/p/x/upload", &body)
.await
.expect("call must succeed");
consumer.await.unwrap();
assert_eq!(result.status, 200);
}
#[tokio::test]
async fn a_response_for_an_unknown_id_is_a_noop() {
let (egress, _sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(50));
for fragment in &response_payloads(DRONE, 9999, 200, b"orphan") {
deliver(&proxy, fragment).await;
}
assert!(proxy.pending.lock().await.is_empty());
}
#[tokio::test]
async fn concurrent_calls_with_distinct_ids_match_their_own_responses() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(500));
let proxy_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
for _ in 0..3 {
let (ch, payload) = sent.recv().await.expect("request datagram");
assert_eq!(ch, AuxChannel::Request);
let request = aux_rpc::decode_request_fragment(&payload).unwrap();
let body = format!("{}", request.id).into_bytes();
deliver_body(&proxy_consumer, DRONE, request.id, 200, &body).await;
}
});
let p1 = tokio::spawn({
let proxy = proxy.clone();
async move { proxy.call(&[], RpcMethod::Get, b"/api/a", &[]).await }
});
let p2 = tokio::spawn({
let proxy = proxy.clone();
async move { proxy.call(&[], RpcMethod::Get, b"/api/b", &[]).await }
});
let p3 = tokio::spawn({
let proxy = proxy.clone();
async move { proxy.call(&[], RpcMethod::Get, b"/api/c", &[]).await }
});
let r1 = p1.await.unwrap().unwrap();
let r2 = p2.await.unwrap().unwrap();
let r3 = p3.await.unwrap().unwrap();
consumer.await.unwrap();
assert_ne!(r1.body, r2.body);
assert_ne!(r2.body, r3.body);
assert_eq!(r1.status, 200);
assert_eq!(r2.status, 200);
assert_eq!(r3.status, 200);
}
#[tokio::test]
async fn reset_drains_pending_callers_with_503() {
let (egress, _sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_secs(60));
let proxy_call = proxy.clone();
let call_handle =
tokio::spawn(async move { proxy_call.call(&[], RpcMethod::Get, b"/api/x", &[]).await });
tokio::time::sleep(Duration::from_millis(20)).await;
proxy.reset().await;
let result = tokio::time::timeout(Duration::from_millis(500), call_handle)
.await
.expect("reset must unblock the caller within the bound")
.unwrap();
assert_eq!(result.unwrap().status, 503);
}
#[tokio::test]
async fn the_ingest_writes_and_the_listener_dispatches() {
let dir = tempfile::tempdir().unwrap();
let sock_path = dir.path().join("aux-rpc-responses.sock");
let sock = UdpSocket::bind("127.0.0.1:0").await.unwrap();
sock.connect(("127.0.0.1", 9u16)).await.unwrap();
let egress = AuxEgress::connected_for_test(sock);
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(500));
let cancel = Arc::new(Notify::new());
let listener_handle = proxy.spawn_response_listener(sock_path.clone(), cancel.clone());
tokio::time::sleep(Duration::from_millis(20)).await;
let proxy_call = proxy.clone();
let call_handle = tokio::spawn(async move {
proxy_call
.call(&[], RpcMethod::Get, b"/api/pairing/info", &[])
.await
});
tokio::time::sleep(Duration::from_millis(20)).await;
let id = pending_id(&proxy).await;
let ingest = AuxRpcResponseIngest::new(&sock_path);
for payload in response_payloads(DRONE, id, 200, br#"{"ok":true}"#) {
ingest.send(&payload).await;
}
let result = tokio::time::timeout(Duration::from_secs(2), call_handle)
.await
.expect("the call must complete within the bound")
.unwrap();
assert_eq!(result.unwrap().status, 200);
cancel.notify_waiters();
let _ = listener_handle.await;
}
#[tokio::test]
async fn an_undecodable_frame_does_not_poison_the_fragments_after_it() {
let dir = tempfile::tempdir().unwrap();
let sock_path = dir.path().join("aux-rpc-responses.sock");
let sock = UdpSocket::bind("127.0.0.1:0").await.unwrap();
sock.connect(("127.0.0.1", 9u16)).await.unwrap();
let egress = AuxEgress::connected_for_test(sock);
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_secs(2));
let cancel = Arc::new(Notify::new());
let listener_handle = proxy.spawn_response_listener(sock_path.clone(), cancel.clone());
tokio::time::sleep(Duration::from_millis(20)).await;
let proxy_call = proxy.clone();
let call_handle =
tokio::spawn(async move { proxy_call.call(&[], RpcMethod::Get, b"/api/x", &[]).await });
tokio::time::sleep(Duration::from_millis(20)).await;
let id = pending_id(&proxy).await;
let body: Vec<u8> = (0..2000u32).map(|i| (i % 251) as u8).collect();
let fragments = response_payloads(DRONE, id, 200, &body);
let ingest = AuxRpcResponseIngest::new(&sock_path);
ingest.send(&fragments[0]).await;
ingest.send(b"not-a-fragment").await;
ingest.send(&fragments[1]).await;
let result = tokio::time::timeout(Duration::from_secs(3), call_handle)
.await
.expect("the call must complete within the bound")
.unwrap();
assert_eq!(result.unwrap().body, body);
cancel.notify_waiters();
let _ = listener_handle.await;
}
#[tokio::test]
async fn a_lost_request_datagram_is_retransmitted_until_the_drone_answers() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_secs(3));
let proxy_for_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
let (_, first) = sent.recv().await.expect("first attempt");
let (_, second) = sent.recv().await.expect("first retransmit");
let (_, third) = sent.recv().await.expect("second retransmit");
assert_eq!(first, second, "a retransmit must be the identical datagram");
assert_eq!(second, third);
let request = aux_rpc::decode_request_fragment(&third).unwrap();
deliver_body(
&proxy_for_consumer,
DRONE,
request.id,
200,
b"late-but-answered",
)
.await;
});
let result = proxy
.call(DRONE, RpcMethod::Get, b"/api/version", &[])
.await
.expect("a retransmitted request must still complete");
consumer.await.unwrap();
assert_eq!(result.body, b"late-but-answered");
assert!(
proxy.stats().retransmits >= 2,
"one datagram on a 20-40% loss lane is a coin flip; it must be resent"
);
}
#[tokio::test]
async fn a_retransmit_that_could_not_be_sent_is_not_counted_as_one_that_was() {
let sock = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let dead = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let dead_addr = dead.local_addr().unwrap();
drop(dead);
sock.connect(dead_addr).await.unwrap();
let egress = AuxEgress::connected_for_test(sock);
let proxy = AuxRpcProxy::new(egress);
let r = tokio::time::timeout(
RPC_DEFAULT_TIMEOUT + Duration::from_secs(2),
proxy.call(&[], RpcMethod::Get, b"/api/pairing/info", &[]),
)
.await;
assert!(r.is_ok(), "the call never terminated");
let stats = proxy.stats();
assert!(
stats.retransmit_send_errors > 0,
"a broken egress must be visible as a send failure, not as silence \
from the far end: {stats:?}"
);
}
#[tokio::test]
async fn a_response_that_has_started_arriving_stops_the_retransmits() {
let (egress, mut sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(600));
let body: Vec<u8> = (0..3000u32).map(|i| (i % 251) as u8).collect();
let proxy_for_consumer = proxy.clone();
let consumer = tokio::spawn(async move {
let (_, payload) = sent.recv().await.expect("request datagram");
let id = aux_rpc::decode_request_fragment(&payload).unwrap().id;
let fragments = response_payloads(DRONE, id, 200, &body);
deliver(&proxy_for_consumer, &fragments[0]).await;
tokio::time::sleep(Duration::from_millis(450)).await;
sent.try_recv().is_ok()
});
let err = proxy
.call(DRONE, RpcMethod::Get, b"/api/x", &[])
.await
.unwrap_err();
let resent = consumer.await.unwrap();
assert_eq!(
err,
RpcError::Incomplete {
received: 1,
total: 7
}
);
assert!(!resent, "no second Request datagram may leave the ground");
assert_eq!(proxy.stats().retransmits, 0);
}
#[tokio::test]
async fn a_saturated_proxy_sheds_load_instead_of_queueing_radio_work() {
let (egress, _sent) = loopback_egress().await;
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_millis(200));
{
let mut pending = proxy.pending.lock().await;
for id in 0..MAX_PENDING_CALLS as u32 {
let (tx, _rx) = oneshot::channel::<RpcResponseOwned>();
pending.insert(
id,
PendingCall {
sender: tx,
expected_sender: None,
assembly: None,
response_started: false,
},
);
}
}
let err = proxy
.call(&[], RpcMethod::Get, b"/api/x", &[])
.await
.unwrap_err();
assert_eq!(err, RpcError::Busy);
assert_eq!(err.to_string(), "relay proxy has too many calls in flight");
assert_eq!(proxy.stats().calls_busy, 1);
}
#[tokio::test]
async fn a_cancelled_call_does_not_leak_its_pending_entry() {
let sock = UdpSocket::bind("127.0.0.1:0").await.unwrap();
sock.connect(("127.0.0.1", 9u16)).await.unwrap();
let egress = AuxEgress::connected_for_test(sock);
let proxy = AuxRpcProxy::with_timeout(egress, Duration::from_secs(60));
let proxy_call = proxy.clone();
let handle =
tokio::spawn(async move { proxy_call.call(&[], RpcMethod::Get, b"/api/x", &[]).await });
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(proxy.pending.lock().await.len(), 1);
handle.abort();
tokio::time::sleep(Duration::from_millis(80)).await;
assert!(
proxy.pending.lock().await.is_empty(),
"a dropped call future must not leak its sender and assembly forever"
);
assert_eq!(proxy.stats().pending_now, 0);
}
#[tokio::test]
async fn the_request_id_seed_is_not_a_fixed_one() {
let (egress_a, _a) = loopback_egress().await;
let (egress_b, _b) = loopback_egress().await;
let a = AuxRpcProxy::with_timeout(egress_a, Duration::from_millis(50));
let b = AuxRpcProxy::with_timeout(egress_b, Duration::from_millis(50));
assert_ne!(a.next_id.load(Ordering::Relaxed), 1);
assert_ne!(b.next_id.load(Ordering::Relaxed), 1);
}
}