use alloc::boxed::Box;
use alloc::string::{String, ToString};
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::sync::atomic::{AtomicUsize, Ordering};
use async_channel::{Receiver, Sender};
use futures_channel::oneshot;
use futures_util::{select_biased, FutureExt, StreamExt};
use hashbrown::HashMap;
use super::{
BoxFut, BoxStream, Connection, Dialer, EnvelopeCodec, Inbound, Outbound, Payload, RpcError,
SubUpdate, TransportError,
};
use crate::router::RouterBuilder;
use crate::AimDb;
const SUBSCRIBE_CHANNEL_CAP: usize = 256;
const MAX_COMMAND_QUEUE_CAP: usize = 8192;
#[derive(Debug, Clone)]
pub struct ClientConfig {
pub reconnect: bool,
pub reconnect_delay: u64,
pub max_reconnect_delay: u64,
pub max_reconnect_attempts: usize,
pub keepalive_interval: Option<u64>,
pub max_offline_queue: usize,
pub sends_hello: bool,
}
impl Default for ClientConfig {
fn default() -> Self {
Self {
reconnect: true,
reconnect_delay: 200,
max_reconnect_delay: 200,
max_reconnect_attempts: 0,
keepalive_interval: None,
max_offline_queue: 256,
sends_hello: false,
}
}
}
fn backoff_delay(config: &ClientConfig, attempt: usize) -> u64 {
let base = config.reconnect_delay;
let cap = config.max_reconnect_delay.max(base);
let shift = attempt.saturating_sub(1).min(16) as u32;
base.saturating_mul(1u64 << shift).min(cap)
}
#[derive(Clone)]
pub struct ClientHandle {
cmd_tx: Sender<ClientCmd>,
prune_tx: Sender<()>,
next_id: Arc<AtomicUsize>,
}
enum ClientCmd {
Call {
id: u64,
method: String,
params: Payload,
reply: oneshot::Sender<CallReply>,
},
CancelCall {
id: u64,
},
Subscribe {
id: u64,
topic: String,
events: Sender<Result<SubUpdate, RpcError>>,
},
Write {
topic: String,
payload: Payload,
},
}
struct CancelOnDrop<'a> {
handle: &'a ClientHandle,
id: u64,
armed: bool,
}
impl Drop for CancelOnDrop<'_> {
fn drop(&mut self) {
if self.armed {
if self
.handle
.cmd_tx
.try_send(ClientCmd::CancelCall { id: self.id })
.is_err()
{
self.handle.request_prune();
}
}
}
}
impl ClientHandle {
fn enqueue(&self, cmd: ClientCmd) -> Result<(), RpcError> {
match self.cmd_tx.force_send(cmd) {
Ok(Some(ClientCmd::CancelCall { .. })) => {
self.request_prune();
Ok(())
}
Ok(_) => Ok(()),
Err(_) => Err(RpcError::Internal),
}
}
fn request_prune(&self) {
let _ = self.prune_tx.try_send(());
}
fn next_id(&self) -> Result<u64, RpcError> {
self.next_id
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
.map(|id| id as u64)
.map_err(|_| RpcError::Internal)
}
pub async fn call(
&self,
method: impl Into<String>,
params: Payload,
) -> Result<Payload, RpcError> {
let id = self.next_id()?;
let (reply, rx) = oneshot::channel();
self.enqueue(ClientCmd::Call {
id,
method: method.into(),
params,
reply,
})?;
let mut cancel = CancelOnDrop {
handle: self,
id,
armed: true,
};
let reply = rx.await;
cancel.armed = false;
reply.map_err(|_| RpcError::Internal)?
}
pub fn subscribe(
&self,
topic: impl Into<String>,
) -> Result<BoxStream<'static, Result<SubUpdate, RpcError>>, RpcError> {
let (events, rx) =
async_channel::bounded::<Result<SubUpdate, RpcError>>(SUBSCRIBE_CHANNEL_CAP);
self.enqueue(ClientCmd::Subscribe {
id: self.next_id()?,
topic: topic.into(),
events,
})?;
Ok(Box::pin(rx))
}
pub fn write(&self, topic: impl Into<String>, payload: Payload) -> Result<(), RpcError> {
self.enqueue(ClientCmd::Write {
topic: topic.into(),
payload,
})
}
}
pub fn run_client<D, C>(
dialer: D,
codec: C,
config: ClientConfig,
clock: Arc<dyn crate::executor::RuntimeOps>,
) -> (ClientHandle, BoxFut<'static, ()>)
where
D: Dialer + 'static,
C: EnvelopeCodec + 'static,
{
let (cmd_tx, cmd_rx) =
async_channel::bounded(config.max_offline_queue.clamp(1, MAX_COMMAND_QUEUE_CAP));
let (prune_tx, prune_rx) = async_channel::bounded(1);
let handle = ClientHandle {
cmd_tx,
prune_tx,
next_id: Arc::new(AtomicUsize::new(1)),
};
let fut = Box::pin(client_loop(dialer, codec, config, cmd_rx, prune_rx, clock));
(handle, fut)
}
enum Ended {
Disconnected,
HandlesDropped,
}
struct DrainOnExit<'a>(&'a Receiver<ClientCmd>);
impl Drop for DrainOnExit<'_> {
fn drop(&mut self) {
self.0.close();
while self.0.try_recv().is_ok() {}
}
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum Delivery {
Event,
BurstBody,
BurstEnd,
}
fn deliver(
subs: &mut HashMap<String, Sender<Result<SubUpdate, RpcError>>>,
last_seq: &mut HashMap<String, u64>,
sub: &str,
seq: u64,
topic: Option<Arc<str>>,
data: Payload,
mode: Delivery,
) {
let prev = last_seq.get(sub).copied().unwrap_or(0);
let skipped = seq.saturating_sub(prev.saturating_add(1));
if let Some(tx) = subs.get(sub) {
if mode == Delivery::BurstBody && tx.capacity().is_some_and(|cap| tx.len() + 1 >= cap) {
return; }
let update = SubUpdate {
topic,
data,
skipped,
snapshot_end: mode == Delivery::BurstEnd,
};
match tx.try_send(Ok(update)) {
Ok(()) => {
last_seq.insert(sub.to_string(), seq.max(prev));
}
Err(e) if e.is_full() => {}
Err(_) => {
subs.remove(sub);
last_seq.remove(sub);
}
}
}
}
enum ClientStep {
Inbound(super::TransportResult<Option<Vec<u8>>>),
Keepalive,
Prune,
Cmd(Option<ClientCmd>),
}
async fn client_loop<D, C>(
dialer: D,
codec: C,
config: ClientConfig,
cmd_rx: Receiver<ClientCmd>,
prune_rx: Receiver<()>,
clock: Arc<dyn crate::executor::RuntimeOps>,
) where
D: Dialer,
C: EnvelopeCodec,
{
let _drain = DrainOnExit(&cmd_rx);
let mut attempt: usize = 0;
loop {
let conn = match dialer.connect().await {
Ok(conn) => {
attempt = 0;
conn
}
Err(e) => {
log_warn!("client dial failed: {:?}", e);
if e == TransportError::Closed {
return;
}
match reconnect_after(&mut attempt, &config, &*clock).await {
true => continue,
false => return,
}
}
};
match drive_connection(conn, &codec, &cmd_rx, &prune_rx, &config, &*clock).await {
Ended::HandlesDropped => return,
Ended::Disconnected => match reconnect_after(&mut attempt, &config, &*clock).await {
true => continue,
false => return,
},
}
}
}
async fn reconnect_after(
attempt: &mut usize,
config: &ClientConfig,
clock: &dyn crate::executor::RuntimeOps,
) -> bool {
if !config.reconnect {
return false;
}
*attempt += 1;
if config.max_reconnect_attempts != 0 && *attempt >= config.max_reconnect_attempts {
log_warn!(
"client giving up after {} reconnect attempts",
config.max_reconnect_attempts
);
return false;
}
clock
.sleep(core::time::Duration::from_millis(backoff_delay(
config, *attempt,
)))
.await;
true
}
type CallReply = Result<Payload, RpcError>;
type PendingCalls = HashMap<u64, oneshot::Sender<CallReply>>;
fn sweep_abandoned(pending: &mut PendingCalls) {
pending.retain(|_, reply| !reply.is_canceled());
}
async fn drive_connection<C>(
mut conn: Box<dyn Connection>,
codec: &C,
cmd_rx: &Receiver<ClientCmd>,
prune_rx: &Receiver<()>,
config: &ClientConfig,
clock: &dyn crate::executor::RuntimeOps,
) -> Ended
where
C: EnvelopeCodec + ?Sized,
{
let mut pending: PendingCalls = HashMap::new();
let mut subs: HashMap<String, Sender<Result<SubUpdate, RpcError>>> = HashMap::new();
let mut last_seq: HashMap<String, u64> = HashMap::new();
let mut out = Vec::new();
let keepalive_ms = config.keepalive_interval;
let mut last_activity = clock.now_nanos();
let mut keepalive_timer =
keepalive_ms.map(|ms| clock.sleep(core::time::Duration::from_millis(ms)).fuse());
if config.sends_hello {
out.clear();
if codec.encode_inbound(Inbound::Ping, &mut out).is_err() || conn.send(&out).await.is_err()
{
return Ended::Disconnected;
}
match conn.recv().await {
Ok(Some(frame)) => match codec.decode_outbound(&frame) {
Ok(Outbound::Pong) => {}
_ => return Ended::Disconnected,
},
_ => return Ended::Disconnected,
}
}
loop {
let step = {
let mut recv = conn.recv().fuse();
let mut keepalive = match keepalive_timer.as_mut() {
Some(timer) => futures_util::future::Either::Left(timer),
None => futures_util::future::Either::Right(futures_util::future::pending::<()>()),
};
let mut cmd = core::pin::pin!(cmd_rx.recv().fuse());
let mut prune = core::pin::pin!(async {
if prune_rx.recv().await.is_err() {
futures_util::future::pending::<()>().await;
}
}
.fuse());
select_biased! {
r = recv => ClientStep::Inbound(r),
_ = keepalive => ClientStep::Keepalive,
_ = prune => ClientStep::Prune,
c = cmd => ClientStep::Cmd(c.ok()),
}
};
if !matches!(step, ClientStep::Keepalive | ClientStep::Prune) {
last_activity = clock.now_nanos();
}
match step {
ClientStep::Inbound(recv) => {
let frame = match recv {
Ok(Some(frame)) => frame,
Ok(None) | Err(_) => return Ended::Disconnected,
};
match codec.decode_outbound(&frame) {
Ok(Outbound::Reply { id, result }) => {
if let Some(tx) = pending.remove(&id) {
let _ = tx.send(result);
} else if let Err(err) = result {
if let Some(tx) = subs.remove(&id.to_string()) {
let _ = tx.try_send(Err(err));
last_seq.remove(&id.to_string());
}
}
}
Ok(Outbound::Event {
sub,
seq,
topic,
data,
}) => deliver(
&mut subs,
&mut last_seq,
sub,
seq,
topic.map(Arc::from),
data,
Delivery::Event,
),
Ok(Outbound::Snapshot {
sub,
seq,
last,
topic,
data,
}) => deliver(
&mut subs,
&mut last_seq,
sub,
seq,
Some(Arc::from(topic)),
data,
if last {
Delivery::BurstEnd
} else {
Delivery::BurstBody
},
),
Ok(Outbound::Pong) => {}
Ok(Outbound::Subscribed { .. }) => {}
Err(_e) => continue, }
}
ClientStep::Prune => sweep_abandoned(&mut pending),
ClientStep::Keepalive => {
let interval_ms = keepalive_ms.unwrap_or(0);
let idle_ms = clock.now_nanos().saturating_sub(last_activity) / 1_000_000;
if idle_ms >= interval_ms {
out.clear();
if codec.encode_inbound(Inbound::Ping, &mut out).is_ok()
&& conn.send(&out).await.is_err()
{
return Ended::Disconnected;
}
last_activity = clock.now_nanos();
keepalive_timer = Some(
clock
.sleep(core::time::Duration::from_millis(interval_ms))
.fuse(),
);
} else {
keepalive_timer = Some(
clock
.sleep(core::time::Duration::from_millis(interval_ms - idle_ms))
.fuse(),
);
}
}
ClientStep::Cmd(cmd) => {
let cmd = match cmd {
Some(cmd) => cmd,
None => return Ended::HandlesDropped, };
match cmd {
ClientCmd::Call {
id,
method,
params,
reply,
} => {
if reply.is_canceled() {
continue;
}
pending.insert(id, reply);
out.clear();
let sent = codec
.encode_inbound(Inbound::Request { id, method, params }, &mut out)
.is_ok()
&& conn.send(&out).await.is_ok();
if !sent {
if let Some(tx) = pending.remove(&id) {
let _ = tx.send(Err(RpcError::Internal));
}
return Ended::Disconnected;
}
}
ClientCmd::CancelCall { id } => {
pending.remove(&id);
}
ClientCmd::Subscribe { id, topic, events } => {
subs.insert(id.to_string(), events);
out.clear();
let sent = codec
.encode_inbound(Inbound::Subscribe { id, topic }, &mut out)
.is_ok()
&& conn.send(&out).await.is_ok();
if !sent {
return Ended::Disconnected;
}
}
ClientCmd::Write { topic, payload } => {
out.clear();
let sent = codec
.encode_inbound(Inbound::Write { topic, payload }, &mut out)
.is_ok()
&& conn.send(&out).await.is_ok();
if !sent {
return Ended::Disconnected;
}
}
}
}
}
}
}
pub fn pump_client(db: &AimDb, scheme: &str, handle: &ClientHandle) -> Vec<BoxFut<'static, ()>> {
let ctx = db.runtime_ctx();
let mut pumps: Vec<BoxFut<'static, ()>> = Vec::new();
for crate::OutboundRoute {
topic: destination,
source,
..
} in db.collect_outbound_routes(scheme)
{
let handle = handle.clone();
let ctx = ctx.clone();
pumps.push(Box::pin(async move {
let mut reader = source.subscribe();
loop {
let msg = match reader.recv(&ctx).await {
Ok(m) => m,
Err(crate::DbError::BufferLagged { .. }) => continue,
Err(_) => break,
};
let dest = msg.dest.unwrap_or_else(|| destination.clone());
if handle
.write(dest, Payload::from(msg.payload.as_slice()))
.is_err()
{
break; }
}
}));
}
let router = Arc::new(RouterBuilder::from_routes(db.collect_inbound_routes(scheme)).build());
for id in router.resource_ids() {
pumps.push(Box::pin(inbound_pump(
handle.clone(),
router.clone(),
id,
ctx.clone(),
)));
}
pumps
}
async fn inbound_pump(
handle: ClientHandle,
router: Arc<crate::router::Router>,
id: Arc<str>,
ctx: crate::RuntimeContext,
) {
let mut stream = match handle.subscribe(id.as_ref()) {
Ok(s) => s,
Err(_e) => return,
};
while let Some(update) = stream.next().await {
match update {
Ok(update) => {
if update.skipped > 0 {
log_warn!(
"mirror gap on '{}': {} update(s) lost (latest-state, no resync)",
id.as_ref(),
update.skipped
);
#[cfg(feature = "defmt")]
defmt::warn!(
"mirror gap on '{}': {} update(s) lost (latest-state, no resync)",
id.as_ref(),
update.skipped
);
}
let _ = router.route(id.as_ref(), &update.data, &ctx);
}
Err(_e) => break,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::session::{CodecError, PeerInfo, TransportResult};
use core::future::pending;
#[derive(Default)]
struct SilentPeer {
peer: PeerInfo,
}
impl Connection for SilentPeer {
fn recv(&mut self) -> BoxFut<'_, TransportResult<Option<Vec<u8>>>> {
Box::pin(pending())
}
fn send<'a>(&'a mut self, _frame: &'a [u8]) -> BoxFut<'a, TransportResult<()>> {
Box::pin(async { Ok(()) })
}
fn peer(&self) -> &PeerInfo {
&self.peer
}
}
struct NullCodec;
impl EnvelopeCodec for NullCodec {
fn decode(&self, _frame: &[u8]) -> Result<Inbound, CodecError> {
Err(CodecError::Malformed)
}
fn encode(&self, _msg: Outbound<'_>, _out: &mut Vec<u8>) -> Result<(), CodecError> {
Ok(())
}
fn encode_inbound(&self, _msg: Inbound, _out: &mut Vec<u8>) -> Result<(), CodecError> {
Ok(())
}
fn decode_outbound<'a>(&self, _frame: &'a [u8]) -> Result<Outbound<'a>, CodecError> {
Err(CodecError::Malformed)
}
}
fn test_handle_capped(cap: usize) -> (ClientHandle, Receiver<ClientCmd>, Receiver<()>) {
let (cmd_tx, cmd_rx) = async_channel::bounded(cap.max(1));
let (prune_tx, prune_rx) = async_channel::bounded(1);
(
ClientHandle {
cmd_tx,
prune_tx,
next_id: Arc::new(AtomicUsize::new(1)),
},
cmd_rx,
prune_rx,
)
}
fn test_handle() -> (ClientHandle, Receiver<ClientCmd>, Receiver<()>) {
test_handle_capped(256)
}
async fn abandon_call(handle: &ClientHandle, method: &'static str) {
let call = handle.call(method, Payload::from(&b"x"[..]));
assert!(
tokio::time::timeout(core::time::Duration::ZERO, call)
.await
.is_err(),
"the call must not resolve"
);
}
#[tokio::test]
async fn abandoning_a_call_queues_its_cancellation() {
let (handle, cmd_rx, _prune_rx) = test_handle();
abandon_call(&handle, "one").await;
match cmd_rx.try_recv() {
Ok(ClientCmd::Call { id: 1, .. }) => {}
_ => panic!("the call must be queued first"),
}
match cmd_rx.try_recv() {
Ok(ClientCmd::CancelCall { id: 1 }) => {}
_ => panic!("abandoning the call must queue its cancellation"),
}
assert!(cmd_rx.is_empty(), "no other command is emitted");
}
#[tokio::test]
async fn a_resolved_call_queues_no_cancellation() {
let (handle, cmd_rx, _prune_rx) = test_handle();
let call = handle.call("one", Payload::from(&b"x"[..]));
let answer = async {
match cmd_rx.recv().await {
Ok(ClientCmd::Call { reply, .. }) => {
let _ = reply.send(Ok(Payload::from(&b"ok"[..])));
}
_ => panic!("the call must be queued"),
}
};
let (reply, ()) = futures_util::future::join(call, answer).await;
assert_eq!(&*reply.expect("the call resolves"), b"ok");
assert!(
cmd_rx.is_empty(),
"a resolved call must not queue a cancellation"
);
}
#[tokio::test]
async fn a_concurrent_batch_of_cancellations_frees_every_entry() {
const BATCH: u64 = 1_000;
let (cmd_tx, cmd_rx) = async_channel::unbounded::<ClientCmd>();
let config = ClientConfig {
reconnect: false,
keepalive_interval: None,
..Default::default()
};
let clock = crate::executor::test_support::NoopRuntimeOps;
let (_prune_tx, prune_rx) = async_channel::bounded::<()>(1);
let engine = drive_connection(
Box::new(SilentPeer::default()),
&NullCodec,
&cmd_rx,
&prune_rx,
&config,
&clock,
);
let exercise = async {
let mut replies = Vec::new();
for id in 1..=BATCH {
let (reply, rx) = oneshot::channel();
cmd_tx
.try_send(ClientCmd::Call {
id,
method: String::from("hang"),
params: Payload::from(&b"x"[..]),
reply,
})
.expect("the channel is unbounded");
replies.push(rx);
}
for id in 1..=BATCH {
cmd_tx
.try_send(ClientCmd::CancelCall { id })
.expect("the channel is unbounded");
}
for (i, rx) in replies.into_iter().enumerate() {
assert!(
rx.await.is_err(),
"call {} was still held by the engine",
i + 1
);
}
};
futures_util::pin_mut!(engine);
futures_util::pin_mut!(exercise);
match futures_util::future::select(engine, exercise).await {
futures_util::future::Either::Left(_) => panic!("the engine ended early"),
futures_util::future::Either::Right(((), _)) => {}
}
}
#[tokio::test]
async fn exhausted_ids_are_refused_not_wrapped() {
let (cmd_tx, cmd_rx) = async_channel::bounded(256);
let (prune_tx, _prune_rx) = async_channel::bounded(1);
let handle = ClientHandle {
cmd_tx,
prune_tx,
next_id: Arc::new(AtomicUsize::new(usize::MAX)),
};
let refused = tokio::time::timeout(
core::time::Duration::ZERO,
handle.call("one", Payload::from(&b"x"[..])),
)
.await
.expect("a refused call must not reach the engine and await a reply");
assert!(
matches!(refused, Err(RpcError::Internal)),
"a call at the id ceiling must be refused"
);
assert!(
handle.subscribe("tele").is_err(),
"a subscription at the id ceiling must be refused"
);
assert!(
cmd_rx.is_empty(),
"a refused allocation must not reach the engine"
);
}
#[tokio::test]
async fn an_evicted_cancellation_raises_the_prune_signal() {
let (handle, cmd_rx, prune_rx) = test_handle_capped(1);
handle
.enqueue(ClientCmd::CancelCall { id: 1 })
.expect("the engine is still running");
assert!(prune_rx.is_empty(), "nothing has been evicted yet");
handle
.write("tele", Payload::from(&b"x"[..]))
.expect("the engine is still running");
assert_eq!(cmd_rx.len(), 1, "the cap still holds");
assert_eq!(
prune_rx.len(),
1,
"an evicted cancellation must fall back to a sweep"
);
}
#[tokio::test]
async fn an_evicted_write_raises_no_prune_signal() {
let (handle, _cmd_rx, prune_rx) = test_handle_capped(1);
for topic in ["one", "two"] {
handle
.write(topic, Payload::from(&b"x"[..]))
.expect("the engine is still running");
}
assert!(
prune_rx.is_empty(),
"an evicted write is expected loss, not a leak"
);
}
#[tokio::test]
async fn a_cancellation_never_evicts_a_queued_write() {
let (handle, cmd_rx, prune_rx) = test_handle_capped(2);
handle
.write("tele", Payload::from(&b"x"[..]))
.expect("the engine is still running");
abandon_call(&handle, "hang").await;
assert_eq!(cmd_rx.len(), 2, "the cap still holds");
match cmd_rx.try_recv() {
Ok(ClientCmd::Write { topic, .. }) => assert_eq!(topic, "tele"),
_ => panic!("the caller's write must survive the cancellation"),
}
match cmd_rx.try_recv() {
Ok(ClientCmd::Call { id: 1, .. }) => {}
_ => panic!("the call itself must still be queued"),
}
assert_eq!(
prune_rx.len(),
1,
"the cancellation it could not queue must fall back to a sweep"
);
}
#[test]
fn the_sweep_frees_abandoned_entries_and_keeps_live_ones() {
let mut pending: PendingCalls = HashMap::new();
let (abandoned, rx) = oneshot::channel();
let (live, _live_rx) = oneshot::channel();
pending.insert(1, abandoned);
pending.insert(2, live);
drop(rx);
sweep_abandoned(&mut pending);
assert!(
!pending.contains_key(&1),
"the abandoned call must be freed"
);
assert!(
pending.contains_key(&2),
"a caller still awaiting its reply must not be dropped"
);
}
#[tokio::test]
async fn an_evicted_cancellation_is_still_reclaimed() {
use core::future::Future;
let (handle, cmd_rx, prune_rx) = test_handle_capped(1);
let config = ClientConfig {
reconnect: false,
keepalive_interval: None,
..Default::default()
};
let clock = crate::executor::test_support::NoopRuntimeOps;
let engine = drive_connection(
Box::new(SilentPeer::default()),
&NullCodec,
&cmd_rx,
&prune_rx,
&config,
&clock,
);
let exercise = async {
let mut call = Box::pin(handle.call("hang", Payload::from(&b"x"[..])));
core::future::poll_fn(|cx| {
assert!(
call.as_mut().poll(cx).is_pending(),
"the call must not resolve"
);
core::task::Poll::Ready(())
})
.await;
tokio::task::yield_now().await;
assert!(cmd_rx.is_empty(), "the engine must have filed the call");
drop(call);
assert_eq!(cmd_rx.len(), 1, "the cancellation must be queued");
assert!(prune_rx.is_empty(), "nothing has fallen back to a sweep");
handle
.write("tele", Payload::from(&b"x"[..]))
.expect("the engine is still running");
assert_eq!(prune_rx.len(), 1, "the eviction must have been noticed");
tokio::time::timeout(core::time::Duration::from_secs(5), async {
while !prune_rx.is_empty() {
tokio::task::yield_now().await;
}
})
.await
.expect("the signal must wake the engine and be consumed");
};
futures_util::pin_mut!(engine);
futures_util::pin_mut!(exercise);
match futures_util::future::select(engine, exercise).await {
futures_util::future::Either::Left(_) => panic!("the engine ended early"),
futures_util::future::Either::Right(((), _)) => {}
}
}
#[tokio::test]
async fn the_command_queue_is_bounded() {
let (handle, cmd_rx, _prune_rx) = test_handle_capped(2);
for topic in ["one", "two", "three", "four"] {
handle
.write(topic, Payload::from(&b"x"[..]))
.expect("the engine is still running");
}
assert_eq!(cmd_rx.len(), 2, "the backlog must not exceed the cap");
for expected in ["three", "four"] {
match cmd_rx.try_recv() {
Ok(ClientCmd::Write { topic, .. }) => assert_eq!(topic, expected),
_ => panic!("the newest writes must survive eviction"),
}
}
}
#[tokio::test]
async fn an_unbounded_cap_degrades_to_the_ceiling() {
struct UndialableRemote;
impl Dialer for UndialableRemote {
fn connect(&self) -> BoxFut<'_, TransportResult<Box<dyn Connection>>> {
Box::pin(async { Err(TransportError::Closed) })
}
}
let config = ClientConfig {
max_offline_queue: usize::MAX,
..Default::default()
};
let (handle, _engine) = run_client(
UndialableRemote,
NullCodec,
config,
Arc::new(crate::executor::test_support::NoopRuntimeOps),
);
handle
.write("t", Payload::from(&b"x"[..]))
.expect("the engine is still running");
}
#[tokio::test]
async fn a_zero_cap_is_clamped_to_one() {
let (handle, cmd_rx, _prune_rx) = test_handle_capped(0);
let (reply, rx) = oneshot::channel();
handle
.enqueue(ClientCmd::Call {
id: 1,
method: String::from("one"),
params: Payload::from(&b"x"[..]),
reply,
})
.expect("the engine is still running");
handle
.write("later", Payload::from(&b"x"[..]))
.expect("the engine is still running");
assert_eq!(cmd_rx.len(), 1, "a clamped cap still holds exactly one");
assert!(rx.await.is_err(), "and the evicted caller learns it");
}
#[tokio::test]
async fn cloned_handles_cannot_exceed_the_cap() {
let (handle, cmd_rx, _prune_rx) = test_handle_capped(1);
let other = handle.clone();
handle
.write("first", Payload::from(&b"x"[..]))
.expect("the engine is still running");
other
.write("second", Payload::from(&b"x"[..]))
.expect("the engine is still running");
assert_eq!(cmd_rx.len(), 1, "clones share one bounded channel");
match cmd_rx.try_recv() {
Ok(ClientCmd::Write { topic, .. }) => assert_eq!(topic, "second"),
_ => panic!("oldest-first eviction must hold across clones"),
}
}
#[tokio::test]
async fn a_trimmed_call_fails_its_caller() {
let (handle, _cmd_rx, _prune_rx) = test_handle_capped(1);
let (reply, rx) = oneshot::channel();
handle
.enqueue(ClientCmd::Call {
id: 1,
method: String::from("one"),
params: Payload::from(&b"x"[..]),
reply,
})
.expect("the engine is still running");
handle
.write("later", Payload::from(&b"x"[..]))
.expect("the engine is still running");
assert!(
rx.await.is_err(),
"a trimmed call must not leave its caller waiting"
);
}
type SubSinks = HashMap<String, Sender<Result<SubUpdate, RpcError>>>;
type SubCursors = HashMap<String, u64>;
fn test_sink() -> (SubSinks, SubCursors, Receiver<Result<SubUpdate, RpcError>>) {
let (tx, rx) = async_channel::bounded(SUBSCRIBE_CHANNEL_CAP);
let mut subs = HashMap::new();
subs.insert(String::from("1"), tx);
(subs, HashMap::new(), rx)
}
fn deliver_event(subs: &mut SubSinks, last_seq: &mut SubCursors, seq: u64) {
deliver(
subs,
last_seq,
"1",
seq,
None,
Payload::from(&b"x"[..]),
Delivery::Event,
);
}
#[tokio::test]
async fn a_saturated_sequence_does_not_overflow() {
let (mut subs, mut last_seq, rx) = test_sink();
deliver_event(&mut subs, &mut last_seq, u64::MAX);
deliver_event(&mut subs, &mut last_seq, 5);
let skips: Vec<u64> = core::iter::from_fn(|| rx.try_recv().ok())
.map(|u| u.unwrap().skipped)
.collect();
assert_eq!(skips.len(), 2, "both frames must be delivered");
assert_eq!(skips[1], 0, "a saturated cursor cannot manufacture a gap");
}
#[tokio::test]
async fn a_rewound_sequence_does_not_invent_a_gap() {
let (mut subs, mut last_seq, rx) = test_sink();
for seq in [1, 3, 2, 4] {
deliver_event(&mut subs, &mut last_seq, seq);
}
let skips: Vec<u64> = core::iter::from_fn(|| rx.try_recv().ok())
.map(|u| u.unwrap().skipped)
.collect();
assert_eq!(
skips,
alloc::vec![0, 1, 0, 0],
"only the real gap (2 missing between 1 and 3) may be reported"
);
assert_eq!(last_seq.get("1"), Some(&4));
}
#[tokio::test]
async fn a_mirror_gap_still_routes_and_keeps_mirroring() {
let seen: Arc<spin::Mutex<Vec<Vec<u8>>>> = Arc::new(spin::Mutex::new(Vec::new()));
let recorder = seen.clone();
let ingest: crate::connector::IngestFn =
Arc::new(move |_ctx: &crate::RuntimeContext, bytes: &[u8]| {
recorder.lock().push(bytes.to_vec());
Ok(())
});
let routes = alloc::vec![(String::from("tele"), ingest)];
let router = Arc::new(RouterBuilder::from_routes(routes).build());
let ctx =
crate::RuntimeContext::new(Arc::new(crate::executor::test_support::NoopRuntimeOps));
let (handle, cmd_rx, _prune_rx) = test_handle();
let pump = inbound_pump(handle, router, Arc::from("tele"), ctx);
let remote = async {
let events = match cmd_rx.recv().await {
Ok(ClientCmd::Subscribe { events, .. }) => events,
_ => panic!("the pump must subscribe"),
};
for (data, skipped) in [(&b"gapped"[..], 7u64), (&b"after"[..], 0)] {
events
.send(Ok(SubUpdate {
topic: None,
data: Payload::from(data),
skipped,
snapshot_end: false,
}))
.await
.expect("the pump holds the receiver");
}
drop(events);
};
futures_util::future::join(pump, remote).await;
let seen = seen.lock();
assert_eq!(seen.len(), 2, "the gap must not drop or end the mirror");
assert_eq!(
seen[0].as_slice(),
b"gapped",
"the update reporting the gap must still be produced"
);
assert_eq!(seen[1].as_slice(), b"after", "mirroring must continue");
}
}