unb-server 2.0.3

mod common;

use std::sync::Arc;
use std::time::Duration;

use bytes::Bytes;
use futures_util::stream;
use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::{json, Value};
use tokio::sync::Mutex;
use unb::{handler, Handler, HandlerError, Reply, Request, State, Streaming};
use unb_client::pair;
use unb_core::{Envelope, ErrorCode, Kind};
use unb_runtime::{ClientError, Wire};
use unb_server::Node;

#[derive(Deserialize, JsonSchema)]
struct Probe {}

#[derive(Deserialize, JsonSchema)]
struct FirehoseSpec {
    count: u64,
}

#[derive(Deserialize, JsonSchema)]
struct CityQuery {
    city: String,
}

fn value_stream(
    rx: tokio::sync::mpsc::Receiver<Result<Value, HandlerError>>,
) -> impl futures_util::Stream<Item = Result<Value, HandlerError>> + Send {
    stream::unfold(rx, |mut rx| async move {
        rx.recv().await.map(|item| (item, rx))
    })
}

struct TickSource {
    stream: Mutex<Option<Streaming<Value, HandlerError>>>,
    _keep: Option<tokio::sync::mpsc::Sender<Result<Value, HandlerError>>>,
}

type EventSender = tokio::sync::mpsc::Sender<Result<Value, HandlerError>>;

struct FirehoseTaps(Arc<std::sync::Mutex<Vec<EventSender>>>);

#[handler]
async fn dry_ticks(
    _request: Request<Probe>,
) -> Result<Streaming<Value, HandlerError>, HandlerError> {
    let items = vec![
        Ok(json!({ "n": 1 })),
        Ok(json!({ "n": 2 })),
        Err(HandlerError::new(ErrorCode::Internal, "tap ran dry")),
    ];
    Ok(Streaming::new(stream::iter(items)))
}

#[handler]
async fn ticks(
    source: State<TickSource>,
    _request: Request<Probe>,
) -> Result<Streaming<Value, HandlerError>, HandlerError> {
    Ok(source
        .stream
        .lock()
        .await
        .take()
        .expect("single subscription in this test"))
}

#[handler]
async fn infinite_firehose(
    _request: Request<Probe>,
) -> Result<Streaming<Value, HandlerError>, HandlerError> {
    let (tx, rx) = tokio::sync::mpsc::channel(4);
    tokio::spawn(async move {
        let mut n = 0u64;
        loop {
            n += 1;
            if tx.send(Ok(json!({ "n": n }))).await.is_err() {
                return;
            }
        }
    });
    Ok(Streaming::new(value_stream(rx)))
}

#[handler]
async fn tapped_firehose(
    taps: State<FirehoseTaps>,
    request: Request<FirehoseSpec>,
) -> Result<Streaming<Value, HandlerError>, HandlerError> {
    let count = request.payload().count;
    let (tx, rx) = tokio::sync::mpsc::channel(4);
    if count == 0 {
        taps.0.lock().unwrap().push(tx.clone());
    }
    tokio::spawn(async move {
        let mut n = 0u64;
        while count == 0 || n < count {
            n += 1;
            if tx.send(Ok(json!({ "n": n }))).await.is_err() {
                return;
            }
        }
    });
    Ok(Streaming::new(value_stream(rx)))
}

#[handler]
async fn weather(request: Request<CityQuery>) -> Result<Reply<Value>, HandlerError> {
    Ok(Reply::new(
        json!({ "temp": 21, "for": request.into_payload().city }),
    ))
}

#[handler]
async fn always_ok(_request: Request<Probe>) -> Result<Reply<Value>, HandlerError> {
    Ok(Reply::new(json!({ "ok": true })))
}

#[tokio::test(flavor = "multi_thread")]
async fn an_opening_frame_with_non_json_payload_relays_byte_identically() {
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (to_owner, at_owner) = pair();
    let owner = Arc::new(Wire::open(at_owner));
    let (connected, ()) = tokio::join!(
        relay.connect_transport("owner", to_owner),
        common::ghost_establish_with(&owner, "owner", &["blob-sink"])
    );
    connected.unwrap();

    let (client_side, relay_side) = pair();
    relay.serve_transport(relay_side).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let opaque = Bytes::from_static(&[0x00, 0xff, 0x13, 0x37, b'!', 0x80, 0x01]);
    let mut observer = owner.observe();
    let mut call = client
        .client_session()
        .start(
            "/owner/blob-sink",
            Kind::Request,
            opaque.clone(),
            None,
            Default::default(),
        )
        .await
        .unwrap();
    let corr = call.operation().as_str().to_owned();

    let request = common::deliver_observed(&mut observer).await;
    assert_eq!(request.kind, Kind::Request);
    assert_eq!(request.subject, "blob-sink");
    assert_eq!(
        request.payload, opaque,
        "the relay must forward payload bytes it cannot even parse"
    );

    let owner_corr = request.corr.unwrap();
    owner
        .respond(&owner_corr, json!({ "stored": true }))
        .await
        .unwrap();

    let response = call.next().await.unwrap().unwrap();
    assert_eq!(response.kind, Kind::Response);
    assert_eq!(response.corr.as_deref(), Some(corr.as_str()));
    assert_eq!(response.payload_json()["stored"], true);
}

#[tokio::test(flavor = "multi_thread")]
async fn relayed_events_carry_payload_bytes_verbatim_and_errors_collect_the_path() {
    let owner = Node::builder("owner")
        .service(dry_ticks.at_subject("ticks"))
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (to_owner, at_owner) = pair();
    let (connected, _) = tokio::join!(
        relay.connect_transport("owner", to_owner),
        owner.serve_transport(at_owner)
    );
    connected.unwrap();

    let (client_side, relay_side) = pair();
    relay.serve_transport(relay_side).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let mut call = common::stream(&client, "/owner/ticks", Kind::Subscribe, json!({})).await;
    let corr = call.operation().as_str().to_owned();

    let first = call.next().await.unwrap().unwrap();
    assert_eq!(first.kind, Kind::Event);
    assert_eq!(first.corr.as_deref(), Some(corr.as_str()));
    assert_eq!(
        first.payload,
        Envelope::encode_payload(&json!({ "n": 1 })),
        "event payload bytes are the owner's encoding, untouched by the relay"
    );

    let second = call.next().await.unwrap().unwrap();
    assert_eq!(second.payload_json()["n"], 2);
    assert_eq!(
        (first.seq, second.seq),
        (Some(1), Some(2)),
        "the owner's seq numbering survives the splice"
    );

    assert!(matches!(
        call.next().await,
        Err(ClientError::Protocol {
            code: unb_core::ErrorCode::Internal,
            ..
        })
    ));
}

fn infinite_firehose_owner() -> Arc<Node> {
    Node::builder("owner")
        .service(infinite_firehose.at_subject("firehose"))
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap()
}

#[tokio::test(flavor = "multi_thread")]
async fn a_stalled_then_resuming_consumer_keeps_its_stream() {
    let owner = infinite_firehose_owner();
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (to_owner, at_owner) = pair();
    let (connected, _) = tokio::join!(
        relay.connect_transport("owner", to_owner),
        owner.serve_transport(at_owner)
    );
    connected.unwrap();

    let (client_side, relay_side) = pair();
    relay.serve_transport(relay_side).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let mut call = common::stream(&client, "/owner/firehose", Kind::Subscribe, json!({})).await;
    let corr = call.operation().as_str().to_owned();

    for _ in 0..20 {
        let envelope = call.next().await.unwrap().unwrap();
        assert_eq!(envelope.kind, Kind::Event);
    }

    tokio::time::sleep(Duration::from_millis(700)).await;

    for _ in 0..400 {
        let envelope = tokio::time::timeout(Duration::from_secs(5), call.next())
            .await
            .expect("a resuming consumer must keep receiving its stream")
            .unwrap()
            .unwrap();
        assert_eq!(
            envelope.kind,
            Kind::Event,
            "a consumer that stalled briefly then resumed must NOT be torn down: {}",
            envelope.payload_json()
        );
        assert_eq!(envelope.corr.as_deref(), Some(corr.as_str()));
    }
}

#[tokio::test(flavor = "multi_thread")]
async fn stalled_session_preserves_cross_session_fairness_then_fails_with_busy() {
    let producers: Arc<std::sync::Mutex<Vec<tokio::sync::mpsc::Sender<_>>>> =
        Arc::new(std::sync::Mutex::new(Vec::new()));
    let owner = Node::builder("owner")
        .state(FirehoseTaps(producers.clone()))
        .service(tapped_firehose.at_subject("firehose"))
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (to_owner, at_owner) = pair();
    let (connected, _) = tokio::join!(
        relay.connect_transport("owner", to_owner),
        owner.serve_transport(at_owner)
    );
    connected.unwrap();

    let (slow_side, relay_slow) = pair();
    relay.serve_transport(relay_slow).await;
    let slow = Wire::open(slow_side);
    common::ready_client(&slow).await;

    let (live_side, relay_live) = pair();
    relay.serve_transport(relay_live).await;
    let live = Wire::open(live_side);
    common::ready_client(&live).await;

    let mut slow_call = common::stream(
        &slow,
        "/owner/firehose",
        Kind::Subscribe,
        json!({ "count": 0 }),
    )
    .await;
    tokio::time::sleep(Duration::from_millis(300)).await;
    let slow_producer = producers.lock().unwrap().remove(0);

    let mut live_call = common::stream(
        &live,
        "/owner/firehose",
        Kind::Subscribe,
        json!({ "count": 25 }),
    )
    .await;
    let live_corr = live_call.operation().as_str().to_owned();
    let live_events = tokio::time::timeout(Duration::from_secs(1), async {
        let mut live_events = 0;
        loop {
            let envelope = live_call.next().await.unwrap().unwrap();
            match envelope.kind {
                Kind::Event => live_events += 1,
                Kind::Response => return live_events,
                other => panic!("unexpected frame on the live stream: {other:?}"),
            }
            assert_eq!(envelope.corr.as_deref(), Some(live_corr.as_str()));
        }
    })
    .await
    .expect("a stalled session must not consume another session's relay opportunity");
    assert_eq!(
        live_events, 25,
        "a sibling stream still receives every event (delivery preserved while a peer stalls; \
         it may be delayed by the stalled peer's backpressure window on a shared upstream)"
    );

    tokio::time::sleep(Duration::from_millis(5_100)).await;

    let mut saw_busy = false;
    for _ in 0..20_000 {
        match tokio::time::timeout(Duration::from_secs(10), slow_call.next())
            .await
            .expect("the slow client must eventually see its BUSY teardown")
        {
            Err(ClientError::Protocol {
                code: unb_core::ErrorCode::Busy,
                ..
            }) => {
                saw_busy = true;
                break;
            }
            Ok(Some(envelope)) => assert_eq!(envelope.kind, Kind::Event),
            other => panic!("unexpected slow stream result: {other:?}"),
        }
    }
    assert!(saw_busy, "the overflowing stream ends in a BUSY error");

    tokio::time::timeout(Duration::from_secs(5), slow_producer.closed())
        .await
        .expect("the BUSY teardown must cancel the owner-side stream (receiver dropped)");
}

#[tokio::test(flavor = "multi_thread")]
async fn a_forwarded_request_returns_the_owner_result_through_the_relay() {
    let owner = Node::builder("owner-1")
        .service(weather)
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (relay_side, owner_side) = pair();
    let (connected, _) = tokio::join!(
        relay.connect_transport("owner-1", relay_side),
        owner.serve_transport(owner_side)
    );
    connected.unwrap();

    let (client_side, relay_server) = pair();
    relay.serve_transport(relay_server).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let mut call = common::stream(
        &client,
        "/owner-1/weather",
        Kind::Request,
        json!({ "city": "hue" }),
    )
    .await;
    let corr = call.operation().as_str().to_owned();
    let response = call.next().await.unwrap().unwrap();
    assert_eq!(response.kind, Kind::Response);
    assert_eq!(response.corr.as_deref(), Some(corr.as_str()));
    assert_eq!(response.payload_json()["temp"], 21);
    assert_eq!(
        response.payload_json()["for"],
        "hue",
        "the owner's result crosses the splice untouched"
    );
    assert!(call.next().await.unwrap().is_none());
}

#[tokio::test(flavor = "multi_thread")]
async fn a_relayed_subscription_streams_and_cancel_tears_down_the_owner_stream() {
    let (tx, rx) = tokio::sync::mpsc::channel(4);
    tx.send(Ok(json!({ "n": 1 }))).await.unwrap();
    let owner = Node::builder("owner-1")
        .state(TickSource {
            stream: Mutex::new(Some(Streaming::new(value_stream(rx)))),
            _keep: Some(tx.clone()),
        })
        .service(ticks)
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (relay_side, owner_side) = pair();
    let (connected, _) = tokio::join!(
        relay.connect_transport("owner-1", relay_side),
        owner.serve_transport(owner_side)
    );
    connected.unwrap();
    let (client_side, relay_server) = pair();
    relay.serve_transport(relay_server).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let mut call = common::stream(&client, "/owner-1/ticks", Kind::Subscribe, json!({})).await;
    let first = call.next().await.unwrap().unwrap();
    assert_eq!(first.kind, Kind::Event);
    assert_eq!(
        first.payload_json()["n"],
        1,
        "events flow through the splice"
    );

    drop(call);
    tokio::time::timeout(Duration::from_secs(5), tx.closed())
        .await
        .expect("cancel must tear down the owner-side stream through the relay");
}

#[tokio::test(flavor = "multi_thread")]
async fn peer_death_mid_stream_terminates_spliced_clients_with_peer_unreachable() {
    let (tx, rx) = tokio::sync::mpsc::channel(4);
    tx.send(Ok(json!({ "n": 1 }))).await.unwrap();
    let owner = Node::builder("owner-1")
        .state(TickSource {
            stream: Mutex::new(Some(Streaming::new(value_stream(rx)))),
            _keep: None,
        })
        .service(ticks)
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (relay_side, owner_side) = pair();
    let (connected, owner_wire) = tokio::join!(
        relay.connect_transport("owner-1", relay_side),
        owner.serve_transport(owner_side)
    );
    connected.unwrap();
    let (client_side, relay_server) = pair();
    relay.serve_transport(relay_server).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let mut call = common::stream(&client, "/owner-1/ticks", Kind::Subscribe, json!({})).await;
    let first = call.next().await.unwrap().unwrap();
    assert_eq!(first.payload_json()["n"], 1);

    owner_wire.shutdown();
    let error = tokio::time::timeout(Duration::from_secs(5), call.next())
        .await
        .expect("peer death must reach the client within 5s")
        .unwrap_err();
    assert!(matches!(
        error,
        ClientError::Protocol {
            code: unb_core::ErrorCode::PeerUnreachable,
            ..
        }
    ));
}

#[tokio::test(flavor = "multi_thread")]
async fn client_disconnect_tears_down_the_owner_leg() {
    let (tx, rx) = tokio::sync::mpsc::channel(4);
    tx.send(Ok(json!({ "n": 1 }))).await.unwrap();
    let owner = Node::builder("owner-1")
        .state(TickSource {
            stream: Mutex::new(Some(Streaming::new(value_stream(rx)))),
            _keep: Some(tx.clone()),
        })
        .service(ticks)
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (relay_side, owner_side) = pair();
    let (connected, _) = tokio::join!(
        relay.connect_transport("owner-1", relay_side),
        owner.serve_transport(owner_side)
    );
    connected.unwrap();
    let (client_side, relay_server) = pair();
    relay.serve_transport(relay_server).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let mut call = common::stream(&client, "/owner-1/ticks", Kind::Subscribe, json!({})).await;
    let first = call.next().await.unwrap().unwrap();
    assert_eq!(first.payload_json()["n"], 1);

    client.shutdown();
    tokio::time::timeout(Duration::from_secs(5), tx.closed())
        .await
        .expect("client disconnect must tear down the owner-side stream through the relay");
}

#[tokio::test(flavor = "multi_thread")]
async fn hop_exhaustion_returns_the_walked_path() {
    let owner = Node::builder("owner-1")
        .service(always_ok.at_subject("far.subject"))
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let relay = Node::builder("relay-a")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (relay_side, owner_side) = pair();
    let (connected, _) = tokio::join!(
        relay.connect_transport("owner-1", relay_side),
        owner.serve_transport(owner_side)
    );
    connected.unwrap();
    let (client_side, a_server) = pair();
    relay.serve_transport(a_server).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let mut call = client
        .client_session()
        .start(
            "/owner-1/far/subject",
            Kind::Request,
            Envelope::encode_payload(&json!({})),
            Some(0),
            Default::default(),
        )
        .await
        .unwrap();
    assert!(matches!(
        call.next().await,
        Err(ClientError::Protocol {
            code: unb_core::ErrorCode::HopLimitExceeded,
            ..
        })
    ));
}

#[tokio::test(flavor = "multi_thread")]
async fn a_dynamically_joined_subject_is_reachable_through_the_relay() {
    let owner = Node::builder("owner-1")
        .service(weather)
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let relay = Node::builder("relay-1")
        .insecure_accept_declared_peer_identities()
        .build()
        .unwrap();
    let (owner_side, relay_side) = pair();
    let (connected, _) = tokio::join!(
        owner.connect_transport("relay-1", owner_side),
        relay.serve_transport(relay_side)
    );
    connected.unwrap();
    let (client_side, relay_server) = pair();
    relay.serve_transport(relay_server).await;
    let client = Wire::open(client_side);
    common::ready_client(&client).await;

    let response = tokio::time::timeout(Duration::from_secs(5), async {
        loop {
            let mut call = common::stream(
                &client,
                "/owner-1/weather",
                Kind::Request,
                json!({ "city": "hue" }),
            )
            .await;
            if let Ok(Some(env)) = call.next().await {
                return env;
            }
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await
    .expect("the announced subject becomes reachable through the relay");
    assert_eq!(response.payload_json()["temp"], 21);
    assert_eq!(
        response.payload_json()["for"],
        "hue",
        "the owner's result crosses the dynamically-established splice"
    );
}