mod common;
use common::TestCall;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use unb::{handler, Handler, HandlerError, Reply, Request};
use unb_core::ErrorCode;
use unb_server::{Endpoint, EndpointSet, Node, TransportKind};
use unb_server::{HostConfig, TcpTransport};
fn ws_endpoint(address: String) -> Endpoint {
Endpoint {
kind: TransportKind::WebSocket,
address,
cert_hash: None,
}
}
#[derive(Deserialize, Serialize, JsonSchema)]
#[serde(transparent)]
struct EchoPayload(Value);
#[derive(Deserialize, JsonSchema)]
struct Doubling {
n: i64,
}
#[derive(Deserialize, JsonSchema)]
struct Probe {}
#[handler]
async fn echo(request: Request<EchoPayload>) -> Result<Reply<Value>, HandlerError> {
Ok(Reply::new(
json!({ "served_by": request.subject(), "got": request.payload() }),
))
}
#[handler]
async fn double(request: Request<Doubling>) -> Result<Reply<Value>, HandlerError> {
let n = request.payload().n;
Ok(Reply::new(json!({ "n": n * 2 })))
}
#[handler]
async fn fail(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
Err(HandlerError::new(ErrorCode::InvalidInput, "no such city"))
}
#[handler]
async fn plan(request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
let outlook = request
.call("/weather/forecast", json!({ "city": "hue" }))
.await?;
Ok(Reply::new(
json!({ "plan": format!("{}C day", outlook["temp_c"]) }),
))
}
#[handler]
async fn forecast(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
Ok(Reply::new(json!({ "temp_c": 21 })))
}
static NO_REPLAY_INVOCATIONS: AtomicUsize = AtomicUsize::new(0);
#[handler]
async fn no_replay_stall(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
NO_REPLAY_INVOCATIONS.fetch_add(1, Ordering::SeqCst);
std::future::pending::<()>().await;
unreachable!("the original in-flight handler is cancelled with its failed session")
}
#[handler]
async fn no_replay_replacement(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
let invocation = NO_REPLAY_INVOCATIONS.fetch_add(1, Ordering::SeqCst) + 1;
Ok(Reply::new(json!({ "invocation": invocation })))
}
fn service(name: &str, subject: &'static str) -> std::sync::Arc<Node> {
Node::builder(name)
.service(echo.at_subject(subject))
.insecure_accept_declared_peer_identities()
.build()
.unwrap()
}
#[tokio::test(flavor = "multi_thread")]
async fn a_local_request_executes_the_handler_directly() {
let node = service("solo", "echo");
let reply = node.request("/solo/echo", json!({ "n": 1 })).await.unwrap();
assert_eq!(reply["served_by"], "echo");
assert_eq!(reply["got"]["n"], 1);
}
#[tokio::test(flavor = "multi_thread")]
async fn linked_nodes_request_each_other_immediately_after_link_returns() {
let a = service("node-a", "a.echo");
let b = service("node-b", "b.echo");
a.link(&b).await.unwrap();
let from_a = a
.request("/node-b/b/echo", json!({ "hello": "b" }))
.await
.unwrap();
assert_eq!(from_a["served_by"], "b.echo");
let from_b = b
.request("/node-a/a/echo", json!({ "hello": "a" }))
.await
.unwrap();
assert_eq!(from_b["served_by"], "a.echo");
}
#[tokio::test(flavor = "multi_thread")]
async fn empty_nodes_link_successfully() {
let a = Node::builder("empty-a")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let b = Node::builder("empty-b")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
a.link(&b).await.unwrap();
}
#[tokio::test(flavor = "multi_thread")]
async fn concurrent_requests_resolve_to_their_own_responses() {
let a = Node::builder("caller")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let b = Node::builder("worker")
.service(double)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
a.link(&b).await.unwrap();
let calls: Vec<_> = (0..16)
.map(|n| {
let a = a.clone();
tokio::spawn(async move { a.request("/worker/double", json!({ "n": n })).await })
})
.collect();
for (n, call) in calls.into_iter().enumerate() {
let reply = call.await.unwrap().unwrap();
assert_eq!(reply["n"], (n as i64) * 2);
}
}
#[tokio::test(flavor = "multi_thread")]
async fn an_unknown_subject_teaches_the_known_names() {
let node = service("solo", "echo");
let error = node.request("/solo/eco", json!({})).await.unwrap_err();
assert_eq!(error.code, ErrorCode::UnknownSubject);
assert!(error.message.contains("echo"), "{}", error.message);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_downstream_error_keeps_its_code_and_message() {
let a = Node::builder("caller")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let b = Node::builder("worker")
.service(fail)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
a.link(&b).await.unwrap();
let error = a.request("/worker/fail", json!({})).await.unwrap_err();
assert_eq!(error.code, ErrorCode::InvalidInput);
assert_eq!(error.message, "no such city");
}
#[tokio::test(flavor = "multi_thread")]
async fn a_handler_nested_request_reaches_a_linked_peer() {
let planner = Node::builder("planner")
.service(plan)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let weather = Node::builder("weather")
.service(forecast)
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
planner.link(&weather).await.unwrap();
let reply = planner.request("/planner/plan", json!({})).await.unwrap();
assert_eq!(reply["plan"], "21C day");
}
#[tokio::test(flavor = "multi_thread")]
async fn a_three_node_chain_routes_transit_subjects_in_both_directions() {
let a = service("chain-a", "alpha");
let b = Node::builder("chain-b")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let c = service("chain-c", "gamma");
a.link(&b).await.unwrap();
b.link(&c).await.unwrap();
common::wait_until("a learns the transit route to chain-c", || {
a.reachable_names().contains(&"chain-c".to_string())
})
.await;
common::wait_until("c learns the transit route to chain-a", || {
c.reachable_names().contains(&"chain-a".to_string())
})
.await;
let via_b = a.request("/chain-c/gamma", json!({})).await.unwrap();
assert_eq!(via_b["served_by"], "gamma");
let reverse = c.request("/chain-a/alpha", json!({})).await.unwrap();
assert_eq!(reverse["served_by"], "alpha");
}
#[tokio::test(flavor = "multi_thread")]
async fn a_request_works_immediately_after_connect_returns() {
let (server, url) = hosted("weather-1", "weather").await;
let dialer = Node::builder("caller")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
dialer
.connect(EndpointSet::from(ws_endpoint(url)))
.await
.unwrap();
let reply = dialer
.request("/weather-1/weather", json!({}))
.await
.unwrap();
assert_eq!(reply["served_by"], "weather");
drop(server);
}
#[tokio::test(flavor = "multi_thread")]
async fn websocket_candidate_stays_established_after_connect_returns() {
let (server, url) = hosted("weather-1", "weather").await;
let dialer = Node::builder("caller")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
dialer
.connect(EndpointSet::from(ws_endpoint(url)))
.await
.unwrap();
tokio::task::yield_now().await;
assert!(dialer.reachable_names().contains(&"weather-1".to_string()));
let reply = dialer
.request("/weather-1/weather", json!({}))
.await
.unwrap();
assert_eq!(reply["served_by"], "weather");
drop(server);
}
async fn hosted(
name: &str,
subject: &'static str,
) -> ((std::sync::Arc<Node>, unb_server::Hosting), String) {
let node = service(name, subject);
hosted_node(node).await
}
async fn hosted_node(
node: std::sync::Arc<Node>,
) -> ((std::sync::Arc<Node>, unb_server::Hosting), String) {
let hosting = HostConfig::tcp(([127, 0, 0, 1], 0), TcpTransport::plain())
.start(&node)
.await
.unwrap();
let url = format!("ws://{}", hosting.websocket_addr().unwrap());
((node, hosting), url)
}
async fn stalling_websocket() -> (String, Arc<AtomicUsize>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let accepted = Arc::new(AtomicUsize::new(0));
tokio::spawn({
let accepted = accepted.clone();
async move {
while let Ok((stream, _)) = listener.accept().await {
accepted.fetch_add(1, Ordering::SeqCst);
tokio::spawn(async move {
let _stream = stream;
std::future::pending::<()>().await;
});
}
}
});
(format!("ws://{address}"), accepted)
}
#[tokio::test(flavor = "multi_thread")]
async fn repeating_connect_against_a_live_connection_converges() {
let (server, url) = hosted("weather-1", "weather").await;
let dialer = Node::builder("caller")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let set = EndpointSet::from(ws_endpoint(url));
dialer.connect(set.clone()).await.unwrap();
dialer.connect(set).await.unwrap();
tokio::time::sleep(Duration::from_millis(100)).await;
let reply = dialer
.request("/weather-1/weather", json!({}))
.await
.unwrap();
assert_eq!(reply["served_by"], "weather");
drop(server);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_later_deliberate_connect_replaces_the_endpoints_after_loss() {
let dialer = Node::builder("caller")
.connect_timeout(Duration::from_secs(2))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let first = hosted("weather-1", "weather").await;
let connection = dialer
.connect(EndpointSet::from(ws_endpoint(first.1.clone())))
.await
.unwrap();
assert!(dialer.reachable_names().contains(&"weather-1".to_string()));
drop(first.0);
common::wait_until("the lost link withdraws its routes", || {
!dialer.reachable_names().contains(&"weather-1".to_string())
})
.await;
assert!(matches!(
connection.status(),
unb_server::ConnectionStatus::Connecting
));
let second = hosted("weather-1", "weather").await;
dialer
.connect(EndpointSet::from(ws_endpoint(second.1.clone())))
.await
.unwrap();
let reply = dialer
.request("/weather-1/weather", json!({}))
.await
.unwrap();
assert_eq!(reply["served_by"], "weather");
drop(second.0);
}
#[tokio::test(flavor = "multi_thread")]
async fn explicit_disconnect_withdraws_the_selected_peer_routes() {
let (server, url) = hosted("weather-1", "weather").await;
let dialer = Node::builder("caller")
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let connection = dialer
.connect(EndpointSet::from(ws_endpoint(url)))
.await
.unwrap();
assert!(dialer.reachable_names().contains(&"weather-1".to_string()));
connection.disconnect();
common::wait_until("explicit disconnect withdraws routes", || {
!dialer.reachable_names().contains(&"weather-1".to_string())
})
.await;
assert!(matches!(
connection.status(),
unb_server::ConnectionStatus::Disconnected {
reason: unb_server::DisconnectReason::ExplicitDisconnect
}
));
drop(server);
}
#[tokio::test(flavor = "multi_thread")]
async fn a_unary_request_waits_for_already_running_recovery() {
let first = hosted("weather-1", "weather").await;
let admissions = Arc::new(AtomicUsize::new(0));
let release = Arc::new(tokio::sync::Notify::new());
let replacement = Node::builder("weather-1")
.service(echo.at_subject("weather"))
.peer_layer_fn({
let admissions = admissions.clone();
let release = release.clone();
move |mut request, next| {
let admissions = admissions.clone();
let release = release.clone();
async move {
admissions.fetch_add(1, Ordering::SeqCst);
release.notified().await;
request.accept_declared();
next.admit(request).await
}
}
})
.build()
.unwrap();
let replacement_hosting = HostConfig::tcp(([127, 0, 0, 1], 0), TcpTransport::plain())
.start(&replacement)
.await
.unwrap();
let replacement_url = format!("ws://{}", replacement_hosting.websocket_addr().unwrap());
let dialer = Node::builder("caller")
.connect_timeout(Duration::from_secs(1))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let connection = dialer
.connect(EndpointSet::from([
ws_endpoint(first.1.clone()),
ws_endpoint(replacement_url),
]))
.await
.unwrap();
first.0 .0.shutdown();
drop(first);
common::wait_until("the old route is withdrawn", || {
!dialer.reachable_names().contains(&"weather-1".to_string())
})
.await;
tokio::time::timeout(Duration::from_secs(1), async {
while admissions.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
let request = dialer.request("/weather-1/weather", json!({ "after": "reconnect" }));
tokio::pin!(request);
tokio::select! {
result = &mut request => panic!("request completed before reconnect: {result:?}"),
() = tokio::task::yield_now() => {}
}
let subscription = dialer.subscribe("/weather-1/weather", json!({}));
tokio::pin!(subscription);
tokio::select! {
_ = &mut subscription => panic!("subscription completed before recovery"),
() = tokio::task::yield_now() => {}
}
release.notify_waiters();
common::wait_until("automatic recovery installs the replacement", || {
connection.status() == unb_server::ConnectionStatus::Connected
})
.await;
let reply = tokio::time::timeout(Duration::from_secs(1), &mut request)
.await
.unwrap()
.unwrap();
assert_eq!(reply["got"]["after"], "reconnect");
let _ = tokio::time::timeout(Duration::from_secs(1), &mut subscription)
.await
.unwrap();
drop(replacement_hosting);
drop(replacement);
}
#[tokio::test(flavor = "multi_thread")]
async fn an_in_flight_request_fails_on_session_loss_and_is_never_replayed() {
NO_REPLAY_INVOCATIONS.store(0, Ordering::SeqCst);
let first_owner = Node::builder("no-replay-owner")
.service(no_replay_stall.at_subject("work"))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let (first, first_url) = hosted_node(first_owner).await;
let replacement_owner = Node::builder("no-replay-owner")
.service(no_replay_replacement.at_subject("work"))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let (replacement, replacement_url) = hosted_node(replacement_owner).await;
let caller = Node::builder("no-replay-caller")
.connect_timeout(Duration::from_millis(200))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let connection = caller
.connect(EndpointSet::from([
ws_endpoint(first_url),
ws_endpoint(replacement_url),
]))
.await
.unwrap();
let in_flight = tokio::spawn({
let caller = caller.clone();
async move { caller.request("/no-replay-owner/work", json!({})).await }
});
tokio::time::timeout(Duration::from_secs(1), async {
while NO_REPLAY_INVOCATIONS.load(Ordering::SeqCst) != 1 {
tokio::task::yield_now().await;
}
})
.await
.expect("the original handler receives the request");
first.0.shutdown();
drop(first);
let error = tokio::time::timeout(Duration::from_secs(1), in_flight)
.await
.expect("the failed session terminates the in-flight operation")
.unwrap()
.unwrap_err();
assert_eq!(error.code, ErrorCode::PeerUnreachable);
common::wait_until("automatic maintenance admits the replacement", || {
connection.status() == unb_server::ConnectionStatus::Connected
})
.await;
assert_eq!(NO_REPLAY_INVOCATIONS.load(Ordering::SeqCst), 1);
let reply = caller
.request("/no-replay-owner/work", json!({}))
.await
.unwrap();
assert_eq!(reply["invocation"], 2);
assert_eq!(NO_REPLAY_INVOCATIONS.load(Ordering::SeqCst), 2);
drop(replacement);
}
#[tokio::test(flavor = "multi_thread")]
async fn reconnect_without_the_old_subject_returns_refreshed_unknown() {
let first = hosted("weather-1", "weather").await;
let admissions = Arc::new(AtomicUsize::new(0));
let release = Arc::new(tokio::sync::Notify::new());
let replacement = Node::builder("weather-1")
.peer_layer_fn({
let admissions = admissions.clone();
let release = release.clone();
move |mut request, next| {
let admissions = admissions.clone();
let release = release.clone();
async move {
admissions.fetch_add(1, Ordering::SeqCst);
release.notified().await;
request.accept_declared();
next.admit(request).await
}
}
})
.build()
.unwrap();
let (replacement_server, replacement_url) = hosted_node(replacement).await;
let dialer = Node::builder("caller")
.connect_timeout(Duration::from_secs(1))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let connection = dialer
.connect(EndpointSet::from([
ws_endpoint(first.1.clone()),
ws_endpoint(replacement_url),
]))
.await
.unwrap();
first.0 .0.shutdown();
drop(first);
common::wait_until("the old route is withdrawn", || {
!dialer.reachable_names().contains(&"weather-1".to_string())
})
.await;
tokio::time::timeout(Duration::from_secs(1), async {
while admissions.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
let request = dialer.request("/weather-1/weather", json!({}));
tokio::pin!(request);
tokio::select! {
result = &mut request => panic!("request completed before reconnect: {result:?}"),
() = tokio::task::yield_now() => {}
}
release.notify_waiters();
common::wait_until("automatic recovery installs the replacement", || {
connection.status() == unb_server::ConnectionStatus::Connected
})
.await;
let error = request.await.unwrap_err();
assert_eq!(error.code, ErrorCode::UnknownSubject);
assert_eq!(connection.status(), unb_server::ConnectionStatus::Connected);
drop(replacement_server);
}
#[tokio::test(flavor = "multi_thread")]
async fn request_joining_a_failed_reconnect_returns_peer_unreachable() {
let first = hosted("weather-1", "weather").await;
let admissions = Arc::new(AtomicUsize::new(0));
let reject = Arc::new(tokio::sync::Notify::new());
let replacement = Node::builder("weather-1")
.peer_layer_fn({
let admissions = admissions.clone();
let reject = reject.clone();
move |request, _next| {
let admissions = admissions.clone();
let reject = reject.clone();
async move {
admissions.fetch_add(1, Ordering::SeqCst);
reject.notified().await;
Err(unb_server::HandlerError::new(
ErrorCode::Unauthorized,
format!("rejected {}", request.remote().node_id),
))
}
}
})
.build()
.unwrap();
let (replacement_server, replacement_url) = hosted_node(replacement).await;
let dialer = Node::builder("caller")
.connect_timeout(Duration::from_secs(1))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let connection = dialer
.connect(EndpointSet::from([
ws_endpoint(first.1.clone()),
ws_endpoint(replacement_url),
]))
.await
.unwrap();
first.0 .0.shutdown();
drop(first);
common::wait_until("the old route is withdrawn", || {
!dialer.reachable_names().contains(&"weather-1".to_string())
})
.await;
tokio::time::timeout(Duration::from_secs(1), async {
while admissions.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
let request = dialer.request("/weather-1/weather", json!({}));
tokio::pin!(request);
tokio::select! {
result = &mut request => panic!("request completed before reconnect: {result:?}"),
() = tokio::task::yield_now() => {}
}
reject.notify_waiters();
connection.disconnect();
let error = request.await.unwrap_err();
assert_eq!(error.code, ErrorCode::PeerUnreachable);
assert!(matches!(
connection.status(),
unb_server::ConnectionStatus::Disconnected {
reason: unb_server::DisconnectReason::ExplicitDisconnect
}
));
assert_eq!(admissions.load(Ordering::SeqCst), 1);
drop(replacement_server);
}
#[tokio::test]
async fn unary_deadline_can_expire_without_cancelling_maintenance() {
let first = hosted("weather-1", "weather").await;
let (stalling_url, accepted) = stalling_websocket().await;
let dialer = Node::builder("caller")
.connect_timeout(Duration::from_secs(60))
.insecure_accept_declared_peer_identities()
.build()
.unwrap();
let connection = dialer
.connect(EndpointSet::from([
ws_endpoint(first.1.clone()),
ws_endpoint(stalling_url),
]))
.await
.unwrap();
tokio::time::pause();
first.0 .0.shutdown();
drop(first);
common::wait_until("the old route is withdrawn", || {
!dialer.reachable_names().contains(&"weather-1".to_string())
})
.await;
tokio::time::timeout(Duration::from_secs(1), async {
while accepted.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
let request = tokio::spawn({
let dialer = dialer.clone();
async move { dialer.request("/weather-1/weather", json!({})).await }
});
tokio::task::yield_now().await;
tokio::time::advance(Duration::from_secs(31)).await;
let error = request.await.unwrap().unwrap_err();
assert_eq!(error.code, ErrorCode::PeerUnreachable);
assert_eq!(
connection.status(),
unb_server::ConnectionStatus::Connecting
);
connection.disconnect();
}