datum-agent 0.10.6

use std::{
    collections::HashMap,
    future::Future,
    net::{IpAddr, Ipv4Addr, SocketAddr},
    sync::Arc,
    time::{Duration, Instant},
};

use datum::{Keep, Sink, Source};
use datum_agent::{
    AgentConfig, ClusterAgent, ClusterAgentConfig, ClusterAgentHandle, JobMat, JobSpec,
    NodeSessionConfig, NodeSessionTransport,
    dcp::{
        ClientKind, ClusterJobList, ClusterNodeList, DcpClient, DcpJobFactories,
        DcpQuicServerConfig, DcpServerConfig, DcpTcpServerConfig, Hello, PlacementSpec,
        PlacementStrategy, ResponseStatus,
    },
};
use datum_cluster::{ClusterConfig, MemberState};
use datum_net::quic::{
    crypto::rustls::{QuicClientConfig, QuicServerConfig},
    quinn,
    rustls::{
        ClientConfig as RustlsClientConfig, RootCertStore, ServerConfig as RustlsServerConfig,
        pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer},
        server::WebPkiClientVerifier,
    },
};
use rcgen::{CertifiedKey, generate_simple_self_signed};

#[derive(Clone)]
enum TestTransport {
    Tcp,
    Quic {
        server_config: Arc<quinn::ServerConfig>,
        client_config: Arc<quinn::ClientConfig>,
    },
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn node_sessions_open_on_join_and_close_on_down() {
    let mut nodes = start_cluster(3, TestTransport::Tcp).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[0], &TestTransport::Tcp).await;
    wait_sessions_connected(&client, 3).await;

    let down_id = nodes[2].cluster().node_id().to_owned();
    let down = nodes.pop().expect("down node");
    down.shutdown().await.expect("down shutdown");
    wait_until(Duration::from_secs(10), || async {
        let info = client.cluster_node_info(500).await.expect("node info");
        let node = info.nodes.iter().find(|node| node.node_id == down_id)?;
        (node.member_state == "Down" && node.session_state != "connected").then_some(())
    })
    .await;

    let info = client.cluster_node_info(500).await.expect("node info");
    let down_node = info
        .nodes
        .iter()
        .find(|node| node.node_id == down_id)
        .expect("down node still in membership view");
    assert_ne!(down_node.session_state, "connected");

    shutdown_all(nodes).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn list_cluster_jobs_aggregates_jobs_from_each_node() {
    let nodes = start_cluster(3, TestTransport::Tcp).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[0], &TestTransport::Tcp).await;
    wait_sessions_connected(&client, 3).await;

    for (index, node) in nodes.iter().enumerate() {
        let node_client = connect(node, &TestTransport::Tcp).await;
        node_client
            .start_job("ticker", format!("job-{index}"), HashMap::new())
            .await
            .expect("job starts");
    }

    let list = wait_until(Duration::from_secs(5), || {
        let client = client.clone();
        async move {
            let list = client.list_cluster_jobs(1_000).await.expect("cluster jobs");
            has_jobs(&list, &["job-0", "job-1", "job-2"]).then_some(list)
        }
    })
    .await;
    assert!(
        !list.partial,
        "unexpected partial result: {:?}",
        list_errors(&list)
    );
    assert_eq!(list.nodes.len(), 3);

    shutdown_all(nodes).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn node_kill_returns_partial_result_with_unreachable_error_and_bounded_wait() {
    let mut nodes = start_cluster(3, TestTransport::Tcp).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[0], &TestTransport::Tcp).await;
    wait_sessions_connected(&client, 3).await;

    let killed_id = nodes[2].cluster().node_id().to_owned();
    let killed = nodes.pop().expect("killed node");
    killed.shutdown().await.expect("killed shutdown");

    let started = Instant::now();
    let list = wait_until(Duration::from_secs(6), || {
        let client = client.clone();
        let killed_id = killed_id.clone();
        async move {
            let list = client.list_cluster_jobs(150).await.expect("cluster jobs");
            list.errors
                .iter()
                .any(|error| error.node_id == killed_id)
                .then_some(list)
        }
    })
    .await;
    assert!(started.elapsed() < Duration::from_secs(2));
    assert!(list.partial);
    assert!(
        list.errors.iter().any(|error| error.node_id == killed_id
            && (error.message.contains("unreachable")
                || error.message.contains("closed")
                || error.message.contains("timed out"))),
        "errors: {:?}",
        list_errors(&list)
    );

    shutdown_all(nodes).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn node_session_reconnects_after_node_restarts() {
    let mut nodes = start_cluster(3, TestTransport::Tcp).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[0], &TestTransport::Tcp).await;
    wait_sessions_connected(&client, 3).await;

    let restarted_id = nodes[2].cluster().node_id().to_owned();
    let seed = nodes[0].cluster().advertise_addr();
    let old = nodes.pop().expect("old node");
    old.shutdown().await.expect("old shutdown");
    wait_until(Duration::from_secs(10), || {
        let nodes = &nodes;
        let restarted_id = restarted_id.clone();
        async move {
            (member_state(&nodes[0], &restarted_id) == Some(MemberState::Down)).then_some(())
        }
    })
    .await;

    let restarted = start_node(&restarted_id, vec![seed], TestTransport::Tcp)
        .await
        .expect("restarted node");
    nodes.push(restarted);
    wait_all_up(&nodes, 3).await;
    wait_sessions_connected(&client, 3).await;

    let restarted_client = connect(&nodes[2], &TestTransport::Tcp).await;
    restarted_client
        .start_job("ticker", "restarted-job", HashMap::new())
        .await
        .expect("restarted job starts");
    let list = wait_until(Duration::from_secs(5), || {
        let client = client.clone();
        async move {
            let list = client.list_cluster_jobs(1_000).await.expect("cluster jobs");
            has_jobs(&list, &["restarted-job"]).then_some(list)
        }
    })
    .await;
    assert!(!list.partial, "unexpected errors: {:?}", list_errors(&list));

    shutdown_all(nodes).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn node_sessions_use_quic_mtls_with_real_certs() {
    let (server_config, client_config) = mtls_quic_configs();
    let transport = TestTransport::Quic {
        server_config: Arc::new(server_config),
        client_config: Arc::new(client_config),
    };
    let nodes = start_cluster(3, transport.clone()).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[0], &transport).await;
    wait_sessions_connected(&client, 3).await;

    let info = client.cluster_node_info(1_000).await.expect("node info");
    assert!(
        info.nodes
            .iter()
            .filter(|node| !node.local)
            .all(|node| node.session_state == "connected"),
        "node info: {:?}",
        node_states(&info)
    );

    shutdown_all(nodes).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cluster_submit_places_from_any_node_by_strategy_and_role() {
    let nodes =
        start_cluster_with_roles(3, TestTransport::Tcp, &[&["red"], &["blue"], &["blue"]]).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[2], &TestTransport::Tcp).await;
    wait_sessions_connected(&client, 3).await;

    let blue_a = client
        .submit_cluster_job(
            "ticker",
            "blue-a",
            HashMap::new(),
            least_jobs(Some("blue")),
            1_000,
        )
        .await
        .expect("blue-a placed");
    let blue_b = client
        .submit_cluster_job(
            "ticker",
            "blue-b",
            HashMap::new(),
            least_jobs(Some("blue")),
            1_000,
        )
        .await
        .expect("blue-b placed");
    assert_eq!(blue_a.placement_node_id, "node-1");
    assert_eq!(blue_b.placement_node_id, "node-2");

    let pinned = client
        .submit_cluster_job(
            "ticker",
            "pinned",
            HashMap::new(),
            pinned("node-2", None),
            1_000,
        )
        .await
        .expect("pinned placed");
    assert_eq!(pinned.placement_node_id, "node-2");

    let red = client
        .submit_cluster_job(
            "ticker",
            "red",
            HashMap::new(),
            least_jobs(Some("red")),
            1_000,
        )
        .await
        .expect("red placed");
    assert_eq!(red.placement_node_id, "node-0");

    let missing = client
        .submit_cluster_job(
            "ticker",
            "missing-role",
            HashMap::new(),
            least_jobs(Some("missing")),
            1_000,
        )
        .await
        .expect_err("missing role rejected");
    assert!(matches!(
        missing,
        datum_agent::dcp::DcpError::Response {
            status: ResponseStatus::NotFound,
            ..
        }
    ));

    shutdown_all(nodes).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn node_kill_replaces_cluster_jobs_but_not_local_jobs() {
    let mut nodes = start_cluster(3, TestTransport::Tcp).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[0], &TestTransport::Tcp).await;
    wait_sessions_connected(&client, 3).await;

    client
        .submit_cluster_job(
            "ticker",
            "cluster-on-node-2",
            HashMap::new(),
            pinned("node-2", None),
            1_000,
        )
        .await
        .expect("cluster job placed");
    let node2_client = connect(&nodes[2], &TestTransport::Tcp).await;
    node2_client
        .start_job("ticker", "local-only", HashMap::new())
        .await
        .expect("local job starts");

    let killed = nodes.pop().expect("killed node");
    killed.shutdown().await.expect("killed shutdown");

    let list = wait_until(Duration::from_secs(10), || {
        let client = client.clone();
        async move {
            let list = client.list_cluster_jobs(1_000).await.expect("cluster jobs");
            let cluster_replaced = find_job(&list, "cluster-on-node-2").is_some_and(|job| {
                job.cluster_job
                    && job.placement_node_id != "node-2"
                    && job.placement_generation >= 2
            });
            let local_replaced = find_job(&list, "local-only").is_some();
            (cluster_replaced && !local_replaced).then_some(list)
        }
    })
    .await;
    let replaced = find_job(&list, "cluster-on-node-2").expect("replaced job");
    assert_ne!(replaced.placement_node_id, "node-2");
    assert_eq!(replaced.placement_history.len(), 2);

    shutdown_all(nodes).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn coordinator_kill_promotes_next_oldest_and_replaces_its_jobs() {
    let mut nodes = start_cluster(3, TestTransport::Tcp).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[1], &TestTransport::Tcp).await;
    wait_sessions_connected(&client, 3).await;

    client
        .submit_cluster_job(
            "ticker",
            "coordinator-job",
            HashMap::new(),
            pinned("node-0", None),
            1_000,
        )
        .await
        .expect("coordinator job placed");
    tokio::time::sleep(Duration::from_millis(200)).await;

    let old = nodes.remove(0);
    old.shutdown().await.expect("old coordinator shutdown");

    let list = wait_until(Duration::from_secs(10), || {
        let client = client.clone();
        async move {
            let list = client.list_cluster_jobs(1_000).await.ok()?;
            let replaced = find_job(&list, "coordinator-job").is_some_and(|job| {
                job.placement_node_id != "node-0" && job.placement_generation >= 2
            });
            replaced.then_some(list)
        }
    })
    .await;
    let replaced = find_job(&list, "coordinator-job").expect("replaced job");
    assert_ne!(replaced.placement_node_id, "node-0");
    assert!(
        replaced
            .placement_history
            .iter()
            .any(|event| event.reason.contains("node_down:node-0"))
    );

    shutdown_all(nodes).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cluster_drain_routes_to_current_placement_node() {
    let nodes = start_cluster(3, TestTransport::Tcp).await;
    wait_all_up(&nodes, 3).await;
    let client = connect(&nodes[0], &TestTransport::Tcp).await;
    wait_sessions_connected(&client, 3).await;

    let status = client
        .submit_cluster_job(
            "ticker",
            "drain-me",
            HashMap::new(),
            pinned("node-2", None),
            1_000,
        )
        .await
        .expect("cluster job placed");
    assert_eq!(status.placement_node_id, "node-2");

    let drained = client
        .drain_cluster_job("drain-me", 1_000)
        .await
        .expect("cluster drain routes");
    assert_eq!(drained.name, "drain-me");
    let node2_client = connect(&nodes[2], &TestTransport::Tcp).await;
    wait_until(Duration::from_secs(5), || {
        let node2_client = node2_client.clone();
        async move {
            let status = node2_client.job_status("drain-me").await.expect("status");
            matches!(status.state.as_str(), "Draining" | "Drained" | "Stopped").then_some(())
        }
    })
    .await;

    shutdown_all(nodes).await;
}

async fn start_cluster(count: usize, transport: TestTransport) -> Vec<ClusterAgentHandle> {
    start_cluster_with_roles(count, transport, &[]).await
}

async fn start_cluster_with_roles(
    count: usize,
    transport: TestTransport,
    roles: &[&[&str]],
) -> Vec<ClusterAgentHandle> {
    let mut nodes = Vec::new();
    for index in 0..count {
        let seeds = nodes
            .first()
            .map(|node: &ClusterAgentHandle| vec![node.cluster().advertise_addr()])
            .unwrap_or_default();
        let node = start_node_with_roles(
            &format!("node-{index}"),
            seeds,
            transport.clone(),
            roles.get(index).copied().unwrap_or(&[]),
        )
        .await
        .expect("cluster node starts");
        nodes.push(node);
    }
    nodes
}

async fn start_node(
    node_id: &str,
    seed_nodes: Vec<SocketAddr>,
    transport: TestTransport,
) -> Result<ClusterAgentHandle, datum_agent::ClusterAgentError> {
    start_node_with_roles(node_id, seed_nodes, transport, &[]).await
}

async fn start_node_with_roles(
    node_id: &str,
    seed_nodes: Vec<SocketAddr>,
    transport: TestTransport,
    extra_roles: &[&str],
) -> Result<ClusterAgentHandle, datum_agent::ClusterAgentError> {
    ClusterAgent::start(
        cluster_config_with_roles(node_id, seed_nodes, &transport, extra_roles),
        factories(),
    )
    .await
}

fn cluster_config_with_roles(
    node_id: &str,
    seed_nodes: Vec<SocketAddr>,
    transport: &TestTransport,
    extra_roles: &[&str],
) -> ClusterAgentConfig {
    let dcp = match transport {
        TestTransport::Tcp => DcpServerConfig {
            node_id: node_id.to_owned(),
            tcp: Some(DcpTcpServerConfig {
                addr: loopback_addr(),
            }),
            quic: None,
            auth_token: None,
            metrics_interval: Duration::from_millis(20),
            frame_buffer: 128,
        },
        TestTransport::Quic { server_config, .. } => DcpServerConfig {
            node_id: node_id.to_owned(),
            tcp: None,
            quic: Some(DcpQuicServerConfig {
                addr: loopback_addr(),
                server_config: server_config.as_ref().clone(),
            }),
            auth_token: None,
            metrics_interval: Duration::from_millis(20),
            frame_buffer: 128,
        },
    };
    let session_transport = match transport {
        TestTransport::Tcp => NodeSessionTransport::TcpLoopback,
        TestTransport::Quic { client_config, .. } => NodeSessionTransport::QuicMtls {
            server_name: "localhost".to_owned(),
            client_config: client_config.as_ref().clone(),
        },
    };
    ClusterAgentConfig {
        agent: AgentConfig {
            poll_interval: Duration::from_millis(5),
            event_buffer: 256,
            ..AgentConfig::default()
        },
        cluster: ClusterConfig {
            node_id: node_id.to_owned(),
            roles: {
                let mut roles = vec!["agent".to_owned()];
                roles.extend(extra_roles.iter().map(|role| (*role).to_owned()));
                roles
            },
            bind_addr: loopback_addr(),
            advertise_addr: loopback_addr(),
            seed_nodes,
            gossip_interval: Duration::from_millis(60),
            probe_timeout: Duration::from_millis(15),
            suspect_timeout: Duration::from_secs(2),
            downing_timeout: Duration::from_millis(150),
            remove_down_after: Duration::from_secs(20),
            event_buffer: 512,
            ..ClusterConfig::default()
        },
        dcp,
        sessions: NodeSessionConfig {
            transport: session_transport,
            reconnect_min_backoff: Duration::from_millis(20),
            reconnect_max_backoff: Duration::from_millis(100),
            request_timeout: Duration::from_millis(300),
            command_buffer: 16,
            ..NodeSessionConfig::default()
        },
    }
}

fn factories() -> DcpJobFactories {
    let factories = DcpJobFactories::new();
    factories
        .register("ticker", |instance_name, _params| {
            Ok(JobSpec::new(instance_name, |context| {
                let control = context.control();
                Ok(
                    Source::tick(Duration::ZERO, Duration::from_millis(20), 1_u64)
                        .via_mat(context.drain_flow(), Keep::right)
                        .to_mat(Sink::ignore(), move |_switch, completion| {
                            JobMat::new(completion, control.clone())
                        }),
                )
            }))
        })
        .expect("ticker factory registers");
    factories
}

async fn connect(node: &ClusterAgentHandle, transport: &TestTransport) -> DcpClient {
    match transport {
        TestTransport::Tcp => DcpClient::connect_tcp(
            node.tcp_addr().expect("TCP addr"),
            Hello::new("test-cli", ClientKind::Cli),
        )
        .await
        .expect("TCP client connects"),
        TestTransport::Quic { client_config, .. } => DcpClient::connect_quic(
            node.quic_addr().expect("QUIC addr"),
            "localhost",
            client_config.as_ref().clone(),
            Hello::new("test-cli", ClientKind::Cli),
        )
        .await
        .expect("QUIC client connects"),
    }
}

async fn wait_all_up(nodes: &[ClusterAgentHandle], expected: usize) {
    wait_until(Duration::from_secs(8), || async {
        nodes
            .iter()
            .all(|node| {
                let state = node.cluster().current_state();
                state.members.len() == expected
                    && state
                        .members
                        .values()
                        .all(|member| member.state == MemberState::Up && !member.unreachable)
            })
            .then_some(())
    })
    .await;
}

async fn wait_sessions_connected(client: &DcpClient, expected_nodes: usize) {
    wait_until(Duration::from_secs(8), || async {
        let info = client.cluster_node_info(1_000).await.expect("node info");
        let connected = info.nodes.len() == expected_nodes
            && info
                .nodes
                .iter()
                .filter(|node| !node.local)
                .all(|node| node.session_state == "connected");
        connected.then_some(info)
    })
    .await;
}

async fn wait_until<T, Fut>(timeout: Duration, mut condition: impl FnMut() -> Fut) -> T
where
    Fut: Future<Output = Option<T>>,
{
    let deadline = Instant::now() + timeout;
    loop {
        if let Some(value) = condition().await {
            return value;
        }
        assert!(Instant::now() < deadline, "condition timed out");
        tokio::time::sleep(Duration::from_millis(10)).await;
    }
}

fn has_jobs(list: &ClusterJobList, jobs: &[&str]) -> bool {
    jobs.iter().all(|job| {
        list.nodes
            .iter()
            .any(|node| node.jobs.iter().any(|status| status.name == *job))
    })
}

fn find_job<'a>(
    list: &'a ClusterJobList,
    name: &str,
) -> Option<&'a datum_agent::dcp::proto::JobStatus> {
    list.nodes
        .iter()
        .flat_map(|node| node.jobs.iter())
        .find(|status| status.name == name)
}

fn least_jobs(role: Option<&str>) -> PlacementSpec {
    PlacementSpec {
        role_constraint: role.unwrap_or_default().to_owned(),
        strategy: PlacementStrategy::LeastJobs as i32,
        pinned_node_id: String::new(),
    }
}

fn pinned(node_id: &str, role: Option<&str>) -> PlacementSpec {
    PlacementSpec {
        role_constraint: role.unwrap_or_default().to_owned(),
        strategy: PlacementStrategy::Pinned as i32,
        pinned_node_id: node_id.to_owned(),
    }
}

fn list_errors(list: &ClusterJobList) -> Vec<(String, String)> {
    list.errors
        .iter()
        .map(|error| (error.node_id.clone(), error.message.clone()))
        .collect()
}

fn node_states(info: &ClusterNodeList) -> Vec<(String, String, String)> {
    info.nodes
        .iter()
        .map(|node| {
            (
                node.node_id.clone(),
                node.member_state.clone(),
                node.session_state.clone(),
            )
        })
        .collect()
}

fn member_state(node: &ClusterAgentHandle, node_id: &str) -> Option<MemberState> {
    node.cluster()
        .current_state()
        .member(node_id)
        .map(|member| member.state)
}

async fn shutdown_all(nodes: Vec<ClusterAgentHandle>) {
    for node in nodes {
        let _ = node.shutdown().await;
    }
}

fn loopback_addr() -> SocketAddr {
    SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)
}

fn mtls_quic_configs() -> (quinn::ServerConfig, quinn::ClientConfig) {
    let CertifiedKey {
        cert: server_cert,
        key_pair: server_key,
    } = generate_simple_self_signed(["localhost".to_owned()]).expect("server cert");
    let CertifiedKey {
        cert: client_cert,
        key_pair: client_key,
    } = generate_simple_self_signed(["datum-client".to_owned()]).expect("client cert");

    let mut client_roots = RootCertStore::empty();
    client_roots
        .add(client_cert.der().clone())
        .expect("trust client cert");
    let verifier = WebPkiClientVerifier::builder(Arc::new(client_roots))
        .build()
        .expect("client verifier");
    let server_crypto = RustlsServerConfig::builder()
        .with_client_cert_verifier(verifier)
        .with_single_cert(
            vec![server_cert.der().clone()],
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(server_key.serialize_der())),
        )
        .expect("server crypto");

    let mut server_roots = RootCertStore::empty();
    server_roots
        .add(server_cert.der().clone())
        .expect("trust server cert");
    let client_crypto = RustlsClientConfig::builder()
        .with_root_certificates(server_roots)
        .with_client_auth_cert(
            vec![client_cert.der().clone()],
            PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(client_key.serialize_der())),
        )
        .expect("client crypto");

    (
        quinn::ServerConfig::with_crypto(Arc::new(
            QuicServerConfig::try_from(server_crypto).expect("QUIC server config"),
        )),
        quinn::ClientConfig::new(Arc::new(
            QuicClientConfig::try_from(client_crypto).expect("QUIC client config"),
        )),
    )
}