unb-server 1.0.0

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("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 })))
}

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("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("b.echo", json!({ "hello": "b" })).await.unwrap();
    assert_eq!(from_a["served_by"], "b.echo");
    let from_b = b.request("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("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("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("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("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 gamma", || {
        a.reachable_names().contains(&"gamma".to_string())
    })
    .await;
    common::wait_until("c learns the transit route to alpha", || {
        c.reachable_names().contains(&"alpha".to_string())
    })
    .await;
    let via_b = a.request("gamma", json!({})).await.unwrap();
    assert_eq!(via_b["served_by"], "gamma");
    let reverse = c.request("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", 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".to_string()));
    let reply = dialer.request("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", json!({})).await.unwrap();
    assert_eq!(reply["served_by"], "weather");
    drop(server);
}

#[tokio::test(flavor = "multi_thread")]
async fn reconnecting_after_loss_relearns_the_routes() {
    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".to_string()));

    drop(first.0);
    common::wait_until("the lost link withdraws its routes", || {
        !dialer.reachable_names().contains(&"weather".to_string())
    })
    .await;
    assert!(dialer.request("weather", json!({})).await.is_err());
    assert!(matches!(
        connection.status(),
        unb_server::ConnectionStatus::Disconnected { .. }
    ));

    let second = hosted("weather-1", "weather").await;
    dialer
        .connect(EndpointSet::from(ws_endpoint(second.1.clone())))
        .await
        .unwrap();
    let reply = dialer.request("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".to_string()));

    connection.disconnect();

    common::wait_until("explicit disconnect withdraws routes", || {
        !dialer.reachable_names().contains(&"weather".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_joins_an_already_running_reconnect() {
    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".to_string())
    })
    .await;

    let reconnect = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().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", json!({ "after": "reconnect" }));
    tokio::pin!(request);
    tokio::select! {
        result = &mut request => panic!("request completed before reconnect: {result:?}"),
        () = tokio::task::yield_now() => {}
    }
    let subscription = tokio::time::timeout(
        Duration::from_millis(100),
        dialer.subscribe("weather", json!({})),
    )
    .await
    .expect("subscription waited for reconnect");
    assert!(subscription.is_err());

    release.notify_waiters();

    reconnect.await.unwrap().unwrap();
    let reply = tokio::time::timeout(Duration::from_secs(1), &mut request)
        .await
        .unwrap()
        .unwrap();
    assert_eq!(reply["got"]["after"], "reconnect");
    drop(replacement_hosting);
    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".to_string())
    })
    .await;
    let reconnect = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().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", json!({}));
    tokio::pin!(request);
    tokio::select! {
        result = &mut request => panic!("request completed before reconnect: {result:?}"),
        () = tokio::task::yield_now() => {}
    }

    release.notify_waiters();

    reconnect.await.unwrap().unwrap();
    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".to_string())
    })
    .await;
    let reconnect = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().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", json!({}));
    tokio::pin!(request);
    tokio::select! {
        result = &mut request => panic!("request completed before reconnect: {result:?}"),
        () = tokio::task::yield_now() => {}
    }

    reject.notify_waiters();

    assert!(reconnect.await.unwrap().is_err());
    let error = request.await.unwrap_err();
    assert_eq!(error.code, ErrorCode::PeerUnreachable);
    assert!(matches!(
        connection.status(),
        unb_server::ConnectionStatus::Disconnected {
            reason: unb_server::DisconnectReason::ReconnectFailed { .. }
        }
    ));
    assert_eq!(admissions.load(Ordering::SeqCst), 1);
    drop(replacement_server);
}

#[tokio::test]
async fn unary_deadline_can_expire_without_cancelling_the_reconnect() {
    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".to_string())
    })
    .await;
    let reconnect = tokio::spawn({
        let connection = connection.clone();
        async move { connection.reconnect().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", 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!(!reconnect.is_finished());
    connection.disconnect();
    assert_eq!(
        reconnect.await.unwrap(),
        Err(unb_server::ConnectError::Cancelled)
    );
}