use std::{
future::Future,
net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener as StdTcpListener},
path::PathBuf,
process::{Command, Output, Stdio},
time::{Duration, Instant},
};
use datum::{Keep, Sink, Source};
use datum_agent::{
Agent, AgentConfig, ClusterAgent, ClusterAgentConfig, ClusterAgentHandle, JobMat, JobSpec,
NodeSessionConfig, NodeSessionTransport,
dcp::{
ClientKind, DcpClient, DcpJobFactories, DcpServer, DcpServerConfig, DcpServerHandle,
DcpTcpServerConfig, Hello, proto::DcpFrame,
},
};
use datum_cluster::{ClusterConfig, MemberState};
use prost::Message as ProstMessage;
use serde_json::Value;
use tokio::{
io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader},
net::TcpListener,
process::Command as TokioCommand,
};
struct TestDaemon {
agent: datum_agent::AgentHandle,
handle: DcpServerHandle,
addr: SocketAddr,
}
impl TestDaemon {
async fn start() -> Self {
let agent = Agent::start_with_config(AgentConfig {
poll_interval: Duration::from_millis(5),
event_buffer: 128,
..AgentConfig::default()
})
.expect("agent starts");
let server = DcpServer::from_agent(
&agent,
factories(),
DcpServerConfig {
node_id: "cli-test-agent".to_owned(),
tcp: Some(DcpTcpServerConfig {
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
}),
quic: None,
auth_token: None,
metrics_interval: Duration::from_millis(10),
frame_buffer: 128,
},
);
let handle = server.start().await.expect("DCP server starts");
let addr = handle.tcp_addr().expect("TCP listener address");
Self {
agent,
handle,
addr,
}
}
async fn shutdown(self) {
self.handle.shutdown().await;
self.agent
.registry()
.shutdown()
.expect("registry shuts down");
}
}
struct TestCluster {
nodes: Vec<ClusterAgentHandle>,
}
impl TestCluster {
async fn start(count: usize) -> Self {
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 config = cluster_agent_config(&format!("node-{index}"), seeds);
nodes.push(
ClusterAgent::start(config, factories())
.await
.expect("cluster agent starts"),
);
}
wait_cluster_up(&nodes, count).await;
let client = DcpClient::connect_tcp(
nodes[0].tcp_addr().expect("node 0 TCP"),
Hello::new("cli-test", ClientKind::Cli),
)
.await
.expect("cluster client connects");
wait_cluster_sessions(&client, count).await;
Self { nodes }
}
async fn shutdown(self) {
for node in self.nodes {
let _ = node.shutdown().await;
}
}
}
fn cluster_agent_config(node_id: &str, seed_nodes: Vec<SocketAddr>) -> ClusterAgentConfig {
ClusterAgentConfig {
agent: AgentConfig {
poll_interval: Duration::from_millis(5),
event_buffer: 128,
..AgentConfig::default()
},
cluster: ClusterConfig {
node_id: node_id.to_owned(),
roles: vec!["agent".to_owned()],
bind_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
advertise_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
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: DcpServerConfig {
node_id: node_id.to_owned(),
tcp: Some(DcpTcpServerConfig {
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
}),
quic: None,
auth_token: None,
metrics_interval: Duration::from_millis(10),
frame_buffer: 128,
},
sessions: NodeSessionConfig {
transport: NodeSessionTransport::TcpLoopback,
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(10), 1_u64)
.instrumented(
format!("{}:{}", context.name(), context.generation()),
context.instrumentation_registry(),
)
.via_mat(context.drain_flow(), Keep::right)
.to_mat(Sink::ignore(), move |_switch, completion| {
JobMat::new(completion, control.clone())
}),
)
}))
})
.expect("ticker factory registers");
factories
.register("empty", |instance_name, _params| {
Ok(JobSpec::new(instance_name, |context| {
let control = context.control();
Ok(Source::<u64>::empty()
.instrumented(
format!("{}:{}", context.name(), context.generation()),
context.instrumentation_registry(),
)
.via_mat(context.drain_flow(), Keep::right)
.to_mat(Sink::ignore(), move |_switch, completion| {
JobMat::new(completion, control.clone())
}))
}))
})
.expect("empty factory registers");
factories
}
async fn wait_cluster_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)
})
})
.await;
}
async fn wait_cluster_sessions(client: &DcpClient, expected: usize) {
wait_until(Duration::from_secs(8), || async {
client
.cluster_node_info(1_000)
.await
.ok()
.is_some_and(|info| {
info.nodes.len() == expected
&& info
.nodes
.iter()
.filter(|node| !node.local)
.all(|node| node.session_state == "connected")
})
})
.await;
}
async fn wait_until<Fut>(timeout: Duration, mut condition: impl FnMut() -> Fut)
where
Fut: Future<Output = bool>,
{
let deadline = Instant::now() + timeout;
while Instant::now() < deadline {
if condition().await {
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
assert!(condition().await, "condition timed out");
}
fn datum_bin() -> PathBuf {
PathBuf::from(env!("CARGO_BIN_EXE_datum"))
}
fn run_datum(addr: SocketAddr, args: &[&str]) -> Output {
Command::new(datum_bin())
.arg("--addr")
.arg(addr.to_string())
.args(args)
.output()
.expect("datum process runs")
}
fn assert_success(output: &Output) {
assert!(
output.status.success(),
"expected success\nstdout:\n{}\nstderr:\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
}
fn stdout(output: &Output) -> String {
String::from_utf8(output.stdout.clone()).expect("stdout is UTF-8")
}
fn stderr(output: &Output) -> String {
String::from_utf8(output.stderr.clone()).expect("stderr is UTF-8")
}
async fn stop_child(child: &mut tokio::process::Child) {
let _ = child.start_kill();
let _ = tokio::time::timeout(Duration::from_secs(1), child.wait()).await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn lifecycle_commands_and_json_snapshots() {
let daemon = TestDaemon::start().await;
let output = run_datum(daemon.addr, &["--json", "ps"]);
assert_success(&output);
assert_eq!(stdout(&output), "{\n \"jobs\": []\n}\n");
let output = run_datum(
daemon.addr,
&["start", "ticker", "--name", "ingest", "--param", "shard=0"],
);
assert_success(&output);
assert!(stdout(&output).contains("started 'ingest' from factory 'ticker'"));
let output = run_datum(daemon.addr, &["ps"]);
assert_success(&output);
let jobs_table = stdout(&output);
assert!(jobs_table.contains("JOB"));
assert!(jobs_table.contains("ingest"));
assert!(jobs_table.contains("Running"));
let output = run_datum(daemon.addr, &["--json", "status", "ingest"]);
assert_success(&output);
let status: Value = serde_json::from_slice(&output.stdout).expect("status JSON parses");
assert_eq!(status["job"]["name"], "ingest");
assert_eq!(status["job"]["state"], "Running");
assert_eq!(status["job"]["generation"], 1);
let output = run_datum(daemon.addr, &["--json", "restart", "ingest"]);
assert_success(&output);
let restarted: Value = serde_json::from_slice(&output.stdout).expect("restart JSON parses");
assert_eq!(restarted["action"], "restarted");
assert_eq!(restarted["job"]["name"], "ingest");
assert_eq!(restarted["job"]["generation"], 2);
let output = run_datum(daemon.addr, &["--json", "metrics", "ingest"]);
assert_success(&output);
let metrics: Value = serde_json::from_slice(&output.stdout).expect("metrics JSON parses");
assert_eq!(metrics["job"], "ingest");
assert!(
metrics["sample"]["streams"]
.as_array()
.expect("metric streams array")
.iter()
.any(|stream| stream["name"]
.as_str()
.is_some_and(|name| name.starts_with("ingest:")))
);
let output = run_datum(daemon.addr, &["drain", "ingest"]);
assert_success(&output);
assert!(stdout(&output).contains("drained 'ingest'"));
let output = run_datum(daemon.addr, &["start", "ticker", "--name", "stop-me"]);
assert_success(&output);
let output = run_datum(daemon.addr, &["--json", "stop", "stop-me"]);
assert_success(&output);
let stopped: Value = serde_json::from_slice(&output.stdout).expect("stop JSON parses");
assert_eq!(stopped["action"], "stopped");
assert_eq!(stopped["job"]["state"], "Stopped");
daemon.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cluster_ps_and_nodes_use_cluster_dcp_requests() {
let cluster = TestCluster::start(2).await;
let node0 = cluster.nodes[0].tcp_addr().expect("node 0 TCP");
let node1 = cluster.nodes[1].tcp_addr().expect("node 1 TCP");
let output = run_datum(node1, &["start", "ticker", "--name", "remote-job"]);
assert_success(&output);
let output = run_datum(
node0,
&[
"submit",
"--cluster",
"--factory",
"ticker",
"--name",
"placed-job",
"--node",
"node-1",
],
);
assert_success(&output);
assert!(stdout(&output).contains("submitted cluster job 'placed-job'"));
let output = run_datum(node0, &["nodes"]);
assert_success(&output);
let nodes_table = stdout(&output);
assert!(nodes_table.contains("NODE"));
assert!(nodes_table.contains("node-0"));
assert!(nodes_table.contains("node-1"));
assert!(nodes_table.contains("connected"));
let output = run_datum(node0, &["ps", "--cluster"]);
assert_success(&output);
let jobs_table = stdout(&output);
assert!(jobs_table.contains("NODE"));
assert!(jobs_table.contains("node-1"));
assert!(jobs_table.contains("remote-job"));
assert!(jobs_table.contains("placed-job"));
assert!(jobs_table.contains("PLACED_ON"));
cluster.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn events_command_tails_subscription() {
let daemon = TestDaemon::start().await;
let child = TokioCommand::new(datum_bin())
.arg("--addr")
.arg(daemon.addr.to_string())
.arg("events")
.stdout(Stdio::piped())
.spawn()
.expect("events command starts");
tokio::time::sleep(Duration::from_millis(100)).await;
let output = run_datum(daemon.addr, &["start", "ticker", "--name", "event-job"]);
assert_success(&output);
let output = tokio::time::timeout(Duration::from_secs(4), child.wait_with_output())
.await
.expect("events command exits after idle")
.expect("events output");
assert_success(&output);
let output = stdout(&output);
assert!(output.contains("event-job"));
assert!(output.contains("Submitted") || output.contains("Started"));
daemon.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn follow_streams_emit_json_lines() {
let daemon = TestDaemon::start().await;
let mut child = TokioCommand::new(datum_bin())
.arg("--addr")
.arg(daemon.addr.to_string())
.arg("--json")
.arg("events")
.arg("--follow")
.stdout(Stdio::piped())
.spawn()
.expect("events --follow command starts");
tokio::time::sleep(Duration::from_millis(100)).await;
let output = run_datum(daemon.addr, &["start", "ticker", "--name", "follow-events"]);
assert_success(&output);
let stdout = child.stdout.take().expect("event stdout is piped");
let mut reader = BufReader::new(stdout);
let mut line = String::new();
tokio::time::timeout(Duration::from_secs(4), reader.read_line(&mut line))
.await
.expect("event line arrives")
.expect("event line reads");
drop(reader);
stop_child(&mut child).await;
let event: Value = serde_json::from_str(&line).expect("event JSON line parses");
assert_eq!(event["event"]["name"], "follow-events");
let output = run_datum(
daemon.addr,
&["start", "ticker", "--name", "follow-metrics"],
);
assert_success(&output);
let mut child = TokioCommand::new(datum_bin())
.arg("--addr")
.arg(daemon.addr.to_string())
.arg("--json")
.arg("metrics")
.arg("follow-metrics")
.arg("--follow")
.stdout(Stdio::piped())
.spawn()
.expect("metrics --follow command starts");
let stdout = child.stdout.take().expect("metric stdout is piped");
let mut reader = BufReader::new(stdout);
let mut line = String::new();
tokio::time::timeout(Duration::from_secs(4), reader.read_line(&mut line))
.await
.expect("metric line arrives")
.expect("metric line reads");
drop(reader);
stop_child(&mut child).await;
let metric: Value = serde_json::from_str(&line).expect("metric JSON line parses");
assert_eq!(metric["job"], "follow-metrics");
assert!(
metric["sample"]["streams"]
.as_array()
.expect("metric streams array")
.iter()
.any(|stream| stream["name"]
.as_str()
.is_some_and(|name| name.starts_with("follow-metrics:")))
);
daemon.shutdown().await;
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn connect_failure_uses_exit_code_2_and_fix_text() {
let listener = StdTcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0))
.expect("ephemeral port binds");
let addr = listener.local_addr().expect("listener address");
drop(listener);
let output = run_datum(addr, &["ps"]);
assert_eq!(output.status.code(), Some(2));
let stderr = stderr(&output);
assert!(stderr.contains("could not connect to datum-agent"));
assert!(stderr.contains("start datum-agent or pass --addr"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_errors_use_exit_code_1_and_fix_text() {
let daemon = TestDaemon::start().await;
let output = run_datum(daemon.addr, &["status", "missing"]);
assert_eq!(output.status.code(), Some(1));
let stderr = stderr(&output);
assert!(stderr.contains("datum-agent rejected the request"));
assert!(stderr.contains("Check the job or factory name with `datum ps`."));
daemon.shutdown().await;
}
#[test]
fn usage_errors_use_exit_code_2() {
let output = Command::new(datum_bin())
.arg("start")
.arg("ticker")
.arg("--param")
.arg("bad-param")
.output()
.expect("datum process runs");
assert_eq!(output.status.code(), Some(2));
assert!(stderr(&output).contains("invalid --param 'bad-param'"));
}
async fn recv_raw_frame<R: AsyncRead + Unpin>(reader: &mut R) -> std::io::Result<Option<DcpFrame>> {
let mut header = [0_u8; 4];
match reader.read_exact(&mut header).await {
Ok(_) => {}
Err(error) if error.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
Err(error) => return Err(error),
}
let len = u32::from_be_bytes(header) as usize;
if len == 0 || len > 16 * 1024 * 1024 {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"bad frame length",
));
}
let mut payload = vec![0_u8; len];
reader.read_exact(&mut payload).await?;
Ok(Some(DcpFrame::decode(payload.as_slice()).map_err(|e| {
std::io::Error::new(std::io::ErrorKind::InvalidData, e)
})?))
}
async fn send_raw_frame<W: AsyncWrite + Unpin>(
writer: &mut W,
frame: &DcpFrame,
) -> std::io::Result<()> {
let payload = frame.encode_to_vec();
let len = u32::try_from(payload.len())
.map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "frame too large"))?;
writer.write_all(&len.to_be_bytes()).await?;
writer.write_all(&payload).await?;
writer.flush().await
}
async fn accept_and_respond(listener: &TcpListener) {
let (mut stream, _) = listener.accept().await.expect("accept");
let hello = recv_raw_frame(&mut stream)
.await
.expect("read hello")
.expect("hello frame");
assert!(matches!(
hello.frame,
Some(datum_agent::dcp::proto::dcp_frame::Frame::Hello(_))
));
let hello_resp = datum_agent::dcp::proto::Response::ok(0, Vec::new());
send_raw_frame(&mut stream, &DcpFrame::response(hello_resp))
.await
.expect("write hello response");
let subscribe = recv_raw_frame(&mut stream)
.await
.expect("read subscribe")
.expect("subscribe frame");
assert!(matches!(
subscribe.frame,
Some(datum_agent::dcp::proto::dcp_frame::Frame::Request(_))
));
let subscribe_resp = datum_agent::dcp::proto::Response::ok(1, Vec::new());
send_raw_frame(&mut stream, &DcpFrame::response(subscribe_resp))
.await
.expect("write subscribe response");
drop(stream);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn events_follow_exits_nonzero_on_daemon_disconnect() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind listener");
let addr = listener.local_addr().expect("listener addr");
let server_task = tokio::spawn(async move {
accept_and_respond(&listener).await;
});
let child = TokioCommand::new(datum_bin())
.arg("--addr")
.arg(addr.to_string())
.arg("events")
.arg("--follow")
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.expect("events --follow command starts");
server_task.await.expect("server completes handshake");
let output = tokio::time::timeout(Duration::from_secs(8), child.wait_with_output())
.await
.expect("events --follow exits after disconnect")
.expect("events output");
assert!(
!output.status.success(),
"events --follow should exit non-zero on disconnect"
);
let err = String::from_utf8_lossy(&output.stderr);
assert!(
!err.is_empty(),
"events --follow should print an error on disconnect, got empty stderr"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn metrics_follow_exits_nonzero_on_daemon_disconnect() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind listener");
let addr = listener.local_addr().expect("listener addr");
let server_task = tokio::spawn(async move {
accept_and_respond(&listener).await;
});
let child = TokioCommand::new(datum_bin())
.arg("--addr")
.arg(addr.to_string())
.arg("metrics")
.arg("any-job")
.arg("--follow")
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
.expect("metrics --follow command starts");
server_task.await.expect("server completes handshake");
let output = tokio::time::timeout(Duration::from_secs(8), child.wait_with_output())
.await
.expect("metrics --follow exits after disconnect")
.expect("metrics output");
assert!(
!output.status.success(),
"metrics --follow should exit non-zero on disconnect"
);
let err = String::from_utf8_lossy(&output.stderr);
assert!(
!err.is_empty(),
"metrics --follow should print an error on disconnect, got empty stderr"
);
}