irtt-rs/crates/irtt-integration-tests/tests/blocking_managed.rs

use std::{
    net::{SocketAddr, UdpSocket},
    sync::{Arc, Condvar, Mutex},
    thread::{self, JoinHandle},
    time::{Duration, Instant},
};

use irtt_proto::{
    decode_request, encode_echo_reply, encode_open_reply, flags, DecodedRequestKind, EchoReply,
    OpenReply, Params, ReceivedStats, StampAt, TimestampFields,
};
use tokio::{runtime::Builder, time::timeout};

use irtt_client::managed::{
    BlockingManagedClient, ManagedClientConfig, ManagedCompletionPolicy, ManagedEndReason,
    ManagedEvent, ManagedLifecycle, ManagedPacing, ManagedTargetConfig, ManagedTargetEndReason,
};
use irtt_client::ClientConfig;

const TOKEN: u64 = 0x1234_5678_90ab_cdef;

struct Signal {
    value: Mutex<bool>,
    ready: Condvar,
}

impl Signal {
    fn set(&self) {
        *self.value.lock().unwrap() = true;
        self.ready.notify_all();
    }

    fn wait(&self) {
        let value = self.value.lock().unwrap();
        let (value, timeout) = self
            .ready
            .wait_timeout_while(value, Duration::from_secs(2), |value| !*value)
            .unwrap();
        assert!(*value && !timeout.timed_out());
    }
}

struct EchoServer {
    addr: SocketAddr,
    probe: Arc<Signal>,
    close: Arc<Signal>,
    worker: JoinHandle<()>,
}

impl EchoServer {
    fn start() -> Self {
        let socket = UdpSocket::bind("127.0.0.1:0").unwrap();
        socket
            .set_read_timeout(Some(Duration::from_secs(3)))
            .unwrap();
        let addr = socket.local_addr().unwrap();
        let probe = Arc::new(Signal {
            value: Mutex::new(false),
            ready: Condvar::new(),
        });
        let close = Arc::new(Signal {
            value: Mutex::new(false),
            ready: Condvar::new(),
        });
        let worker_probe = Arc::clone(&probe);
        let worker_close = Arc::clone(&close);
        let worker = thread::spawn(move || {
            let mut buffer = [0_u8; 8192];
            let (open_len, peer) = socket.recv_from(&mut buffer).unwrap();
            let negotiated = match decode_request(&buffer[..open_len]).unwrap().kind {
                DecodedRequestKind::Open { params, .. } => Params::decode(params).unwrap(),
                other => panic!("expected an open request, got {other:?}"),
            };
            let reply = encode_open_reply(
                &OpenReply {
                    flags: flags::FLAG_OPEN | flags::FLAG_REPLY,
                    token: TOKEN,
                    params: negotiated.clone(),
                },
                None,
            )
            .unwrap();
            socket.send_to(&reply, peer).unwrap();

            loop {
                let (len, packet_peer) = socket.recv_from(&mut buffer).unwrap();
                let Ok(request) = decode_request(&buffer[..len]) else {
                    continue;
                };
                match request.kind {
                    DecodedRequestKind::Echo { sequence, .. } => {
                        worker_probe.set();
                        let reply = encode_echo_reply(
                            &EchoReply {
                                flags: flags::FLAG_REPLY,
                                token: TOKEN,
                                sequence,
                                recv_count: None,
                                recv_window: None,
                                timestamps: TimestampFields::default(),
                                payload: Vec::new(),
                            },
                            &negotiated,
                            None,
                        )
                        .unwrap();
                        socket.send_to(&reply, packet_peer).unwrap();
                    }
                    DecodedRequestKind::Close { .. } => {
                        worker_close.set();
                        return;
                    }
                    DecodedRequestKind::Open { .. } => {}
                }
            }
        });
        Self {
            addr,
            probe,
            close,
            worker,
        }
    }

    fn wait_probe(&self) {
        self.probe.wait();
    }

    fn wait_close(&self) {
        self.close.wait();
    }

    fn finish(self) {
        self.worker.join().unwrap();
    }
}

fn test_runtime() -> tokio::runtime::Runtime {
    Builder::new_current_thread()
        .enable_io()
        .enable_time()
        .build()
        .unwrap()
}

fn config(completion: ManagedCompletionPolicy) -> ManagedClientConfig {
    ManagedClientConfig {
        client: ClientConfig {
            duration: Some(Duration::from_millis(120)),
            interval: Duration::from_millis(20),
            probe_timeout: Duration::from_millis(40),
            open_timeouts: vec![Duration::from_millis(200)],
            received_stats: ReceivedStats::None,
            stamp_at: StampAt::None,
            ..ClientConfig::default()
        },
        pacing: ManagedPacing::Burst,
        completion,
        final_drain: Duration::from_millis(5),
        ..ManagedClientConfig::default()
    }
}

fn target(server: &EchoServer) -> ManagedTargetConfig {
    ManagedTargetConfig::new("target", server.addr.to_string())
}

fn assert_stopped(outcome: &irtt_client::managed::ManagedOutcome) {
    assert_eq!(outcome.end_reason, ManagedEndReason::StopRequested);
}

#[test]
fn starts_without_caller_runtime_and_finishes_naturally() {
    let server = EchoServer::start();
    let owner = BlockingManagedClient::start(
        config(ManagedCompletionPolicy::FinishWhenQuiescent),
        vec![target(&server)],
    )
    .unwrap();

    server.wait_probe();
    let outcome = owner.join().unwrap();
    assert_eq!(outcome.end_reason, ManagedEndReason::TargetsComplete);
    assert_eq!(outcome.recent_target_outcomes.len(), 1);
    assert_eq!(
        outcome.recent_target_outcomes[0].end_reason,
        ManagedTargetEndReason::TestComplete
    );
    server.wait_close();
    server.finish();
}

#[test]
fn subscribed_start_captures_events_from_a_short_run() {
    let server = EchoServer::start();
    let (owner, mut events) = BlockingManagedClient::start_with_subscription(
        config(ManagedCompletionPolicy::FinishWhenQuiescent),
        vec![target(&server)],
    )
    .unwrap();

    server.wait_probe();
    let outcome = owner.join().unwrap();
    assert_eq!(outcome.end_reason, ManagedEndReason::TargetsComplete);

    let mut saw_started = false;
    let mut saw_finished = false;
    let mut saw_completed = false;
    while let Ok(event) = events.try_recv() {
        match event {
            ManagedEvent::Started => saw_started = true,
            ManagedEvent::TargetFinished { .. } => saw_finished = true,
            ManagedEvent::Completed { .. } => saw_completed = true,
            _ => {}
        }
    }
    assert!(saw_started && saw_finished && saw_completed);
    server.wait_close();
    server.finish();
}

#[test]
fn external_handle_stop_ends_with_graceful_close() {
    let server = EchoServer::start();
    let owner = BlockingManagedClient::start(
        config(ManagedCompletionPolicy::ExplicitStop),
        vec![target(&server)],
    )
    .unwrap();
    let handle = owner.handle();

    server.wait_probe();
    test_runtime().block_on(async {
        timeout(Duration::from_secs(2), handle.stop())
            .await
            .unwrap();
    });
    let outcome = owner.join().unwrap();
    assert_stopped(&outcome);
    server.wait_close();
    server.finish();
}

#[test]
fn dynamic_update_reaches_background_worker() {
    let server = EchoServer::start();
    let owner = BlockingManagedClient::start(config(ManagedCompletionPolicy::ExplicitStop), vec![])
        .unwrap();
    let handle = owner.handle();

    let acknowledgement = handle
        .update_targets(vec![target(&server)])
        .unwrap()
        .blocking_wait()
        .unwrap();
    assert_eq!(acknowledgement.sequence, 1);
    server.wait_probe();
    test_runtime().block_on(async {
        timeout(Duration::from_secs(2), owner.handle().stop())
            .await
            .unwrap();
    });
    assert_stopped(&owner.join().unwrap());
    server.wait_close();
    server.finish();
}

#[test]
fn dropping_external_handles_does_not_stop_owner() {
    let server = EchoServer::start();
    let owner = BlockingManagedClient::start(
        config(ManagedCompletionPolicy::FinishWhenQuiescent),
        vec![target(&server)],
    )
    .unwrap();
    drop(owner.handle());
    drop(owner.handle());

    server.wait_probe();
    assert_eq!(
        owner.join().unwrap().end_reason,
        ManagedEndReason::TargetsComplete
    );
    server.wait_close();
    server.finish();
}

#[test]
fn dropping_owner_stops_and_joins_worker() {
    let server = EchoServer::start();
    let owner = BlockingManagedClient::start(
        config(ManagedCompletionPolicy::ExplicitStop),
        vec![target(&server)],
    )
    .unwrap();
    let handle = owner.handle();

    server.wait_probe();
    drop(owner);
    let status = handle.status();
    assert_eq!(status.lifecycle, ManagedLifecycle::Completed);
    assert_eq!(
        status.final_outcome.as_ref().unwrap().end_reason,
        ManagedEndReason::StopRequested
    );
    server.wait_close();
    server.finish();
}

#[test]
fn drop_after_natural_completion_preserves_outcome() {
    let server = EchoServer::start();
    let owner = BlockingManagedClient::start(
        config(ManagedCompletionPolicy::FinishWhenQuiescent),
        vec![target(&server)],
    )
    .unwrap();
    let handle = owner.handle();
    server.wait_probe();
    server.wait_close();
    let deadline = Instant::now() + Duration::from_secs(2);
    let status = loop {
        let status = handle.status();
        if status.lifecycle == ManagedLifecycle::Completed && status.final_outcome.is_some() {
            break status;
        }
        assert!(Instant::now() < deadline, "managed worker did not complete");
        thread::sleep(Duration::from_millis(5));
    };
    assert_eq!(
        status.final_outcome.as_ref().unwrap().end_reason,
        ManagedEndReason::TargetsComplete
    );
    drop(owner);
    assert_eq!(
        handle.status().final_outcome.as_ref().unwrap().end_reason,
        ManagedEndReason::TargetsComplete
    );
    server.finish();
}