#![forbid(unsafe_code)]
use std::{
collections::{HashMap, HashSet},
fs,
io::{self, Cursor, Write},
net::SocketAddr,
path::{Path, PathBuf},
sync::Arc,
time::Duration,
};
use clap::{Parser, Subcommand};
use datum_agent::dcp::{
ClientKind, DcpClient, DcpError, EventSubscription, Hello, MetricSubscription, ResponseStatus,
proto::{
ClusterJobList, ClusterJobNode, ClusterNodeError, ClusterNodeList, ClusterNodeStatus,
Event, JobStatus as WireJobStatus, MetricSample, PlacementSpec, PlacementStrategy,
StreamMetric,
},
};
use datum_net::quic::{
crypto::rustls::QuicClientConfig,
quinn,
rustls::{
ClientConfig as RustlsClientConfig, RootCertStore,
pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer},
},
};
use serde::Serialize;
const DEFAULT_ADDR: &str = "127.0.0.1:9555";
const EXIT_OK: i32 = 0;
const EXIT_REMOTE: i32 = 1;
const EXIT_USAGE_OR_CONNECT: i32 = 2;
const EVENT_INITIAL_IDLE: Duration = Duration::from_millis(500);
const EVENT_IDLE: Duration = Duration::from_millis(250);
const METRIC_SNAPSHOT_WAIT: Duration = Duration::from_secs(2);
const DRAIN_POLL: Duration = Duration::from_millis(20);
#[derive(Debug, Parser)]
#[command(
name = "datum",
version,
about = "Manage Datum jobs through the Datum Control Protocol"
)]
struct Cli {
#[arg(long, global = true, default_value = DEFAULT_ADDR)]
addr: SocketAddr,
#[arg(long, global = true, value_name = "PATH")]
tls_ca: Option<PathBuf>,
#[arg(long, global = true, value_name = "PATH")]
tls_cert: Option<PathBuf>,
#[arg(long, global = true, value_name = "PATH")]
tls_key: Option<PathBuf>,
#[arg(long, global = true)]
json: bool,
#[command(subcommand)]
command: Command,
}
#[derive(Debug, Subcommand)]
enum Command {
Ps {
#[arg(long)]
cluster: bool,
},
Nodes,
Status {
#[arg(long)]
cluster: bool,
job: String,
},
Start {
factory: String,
#[arg(long)]
name: Option<String>,
#[arg(long = "param", value_name = "k=v")]
params: Vec<String>,
},
Submit {
#[arg(long)]
cluster: bool,
#[arg(long)]
factory: String,
#[arg(long)]
name: Option<String>,
#[arg(long = "param", value_name = "k=v")]
params: Vec<String>,
#[arg(long)]
role: Option<String>,
#[arg(long)]
node: Option<String>,
},
Drain {
#[arg(long)]
cluster: bool,
job: String,
},
Stop {
#[arg(long)]
cluster: bool,
job: String,
},
Restart { job: String },
Events {
#[arg(long)]
follow: bool,
},
Metrics {
job: String,
#[arg(long)]
follow: bool,
},
}
#[derive(Debug)]
enum CliError {
Usage(String),
Connect { addr: SocketAddr, message: String },
Remote { message: String, fix: &'static str },
}
impl CliError {
fn exit_code(&self) -> i32 {
match self {
Self::Remote { .. } => EXIT_REMOTE,
Self::Usage(_) | Self::Connect { .. } => EXIT_USAGE_OR_CONNECT,
}
}
fn print(&self) {
match self {
Self::Usage(message) => {
eprintln!("{message}");
}
Self::Connect { addr, message } => {
eprintln!("could not connect to datum-agent at {addr}: {message}");
eprintln!(
"start datum-agent or pass --addr host:port; loopback TCP defaults to {DEFAULT_ADDR}"
);
}
Self::Remote { message, fix } => {
eprintln!("datum-agent rejected the request: {message}");
eprintln!("{fix}");
}
}
}
}
pub async fn main_entry() -> i32 {
match Cli::try_parse() {
Ok(cli) => match run(cli).await {
Ok(()) => EXIT_OK,
Err(error) => {
error.print();
error.exit_code()
}
},
Err(error) => {
let _ = error.print();
EXIT_USAGE_OR_CONNECT
}
}
}
async fn run(cli: Cli) -> Result<(), CliError> {
validate_before_connect(&cli.command)?;
let client = connect(&cli).await?;
match &cli.command {
Command::Ps { cluster } => {
if *cluster {
cluster_ps(&client, cli.addr, cli.json).await
} else {
ps(&client, cli.addr, cli.json).await
}
}
Command::Nodes => nodes(&client, cli.addr, cli.json).await,
Command::Status { cluster, job } => {
status(&client, cli.addr, job, *cluster, cli.json).await
}
Command::Start {
factory,
name,
params,
} => {
start(
&client,
cli.addr,
factory,
name.as_deref(),
params,
cli.json,
)
.await
}
Command::Submit {
cluster,
factory,
name,
params,
role,
node,
} => {
submit(
&client,
cli.addr,
SubmitArgs {
cluster: *cluster,
factory,
name: name.as_deref(),
raw_params: params,
role: role.as_deref(),
node: node.as_deref(),
},
cli.json,
)
.await
}
Command::Drain { cluster, job } => drain(&client, cli.addr, job, *cluster, cli.json).await,
Command::Stop { cluster, job } => stop(&client, cli.addr, job, *cluster, cli.json).await,
Command::Restart { job } => restart(&client, cli.addr, job, cli.json).await,
Command::Events { follow } => events(&client, cli.addr, *follow, cli.json).await,
Command::Metrics { job, follow } => {
metrics(&client, cli.addr, job, *follow, cli.json).await
}
}
}
fn validate_before_connect(command: &Command) -> Result<(), CliError> {
match command {
Command::Start { params, .. } => {
parse_params(params)?;
}
Command::Submit {
cluster,
params,
role,
node,
..
} => {
if !cluster {
return Err(CliError::Usage(
"`datum submit` currently requires --cluster; use `datum start` for local jobs."
.to_owned(),
));
}
if role.is_some() && node.is_some() {
return Err(CliError::Usage(
"pass at most one of --role or --node for cluster placement.".to_owned(),
));
}
parse_params(params)?;
}
_ => {}
}
Ok(())
}
async fn connect(cli: &Cli) -> Result<DcpClient, CliError> {
let hello = Hello::new(format!("datum-cli-{}", std::process::id()), ClientKind::Cli);
match tls_paths(cli)? {
Some(paths) => {
let config = load_quic_client_config(paths.ca, paths.cert, paths.key)?;
DcpClient::connect_quic(cli.addr, "localhost", config, hello)
.await
.map_err(|error| connect_error(cli.addr, error))
}
None => DcpClient::connect_tcp(cli.addr, hello)
.await
.map_err(|error| connect_error(cli.addr, error)),
}
}
struct TlsPaths<'a> {
ca: &'a Path,
cert: &'a Path,
key: &'a Path,
}
fn tls_paths(cli: &Cli) -> Result<Option<TlsPaths<'_>>, CliError> {
match (&cli.tls_ca, &cli.tls_cert, &cli.tls_key) {
(None, None, None) => Ok(None),
(Some(ca), Some(cert), Some(key)) => Ok(Some(TlsPaths {
ca: ca.as_path(),
cert: cert.as_path(),
key: key.as_path(),
})),
_ => Err(CliError::Usage(
"QUIC+mTLS mode needs --tls-ca, --tls-cert, and --tls-key; pass all three or omit all three for loopback TCP.".to_owned(),
)),
}
}
fn load_quic_client_config(
ca_path: &Path,
cert_path: &Path,
key_path: &Path,
) -> Result<quinn::ClientConfig, CliError> {
let mut roots = RootCertStore::empty();
for cert in read_certificates(ca_path)? {
roots.add(cert).map_err(|error| {
CliError::Usage(format!(
"could not load --tls-ca {}: {error}; pass a PEM or DER CA certificate.",
ca_path.display()
))
})?;
}
let certs = read_certificates(cert_path)?;
let key = read_private_key(key_path)?;
let rustls = RustlsClientConfig::builder()
.with_root_certificates(roots)
.with_client_auth_cert(certs, key)
.map_err(|error| {
CliError::Usage(format!(
"could not build QUIC client identity: {error}; check --tls-cert and --tls-key match."
))
})?;
let quic = QuicClientConfig::try_from(rustls).map_err(|error| {
CliError::Usage(format!(
"could not build QUIC client config: {error}; check the TLS files."
))
})?;
Ok(quinn::ClientConfig::new(Arc::new(quic)))
}
fn read_certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, CliError> {
let bytes = fs::read(path).map_err(|error| {
CliError::Usage(format!(
"could not read {}: {error}; pass a readable PEM or DER certificate file.",
path.display()
))
})?;
let mut cursor = Cursor::new(bytes.as_slice());
let certs = rustls_pemfile::certs(&mut cursor)
.collect::<Result<Vec<_>, _>>()
.map_err(|error| {
CliError::Usage(format!(
"could not parse {}: {error}; pass a PEM or DER certificate file.",
path.display()
))
})?;
if certs.is_empty() {
Ok(vec![CertificateDer::from(bytes)])
} else {
Ok(certs)
}
}
fn read_private_key(path: &Path) -> Result<PrivateKeyDer<'static>, CliError> {
let bytes = fs::read(path).map_err(|error| {
CliError::Usage(format!(
"could not read {}: {error}; pass a readable PEM or DER PKCS#8 private key.",
path.display()
))
})?;
let mut cursor = Cursor::new(bytes.as_slice());
match rustls_pemfile::private_key(&mut cursor).map_err(|error| {
CliError::Usage(format!(
"could not parse {}: {error}; pass a PEM or DER PKCS#8 private key.",
path.display()
))
})? {
Some(key) => Ok(key),
None => Ok(PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(bytes))),
}
}
fn connect_error(addr: SocketAddr, error: DcpError) -> CliError {
match error {
DcpError::Response { status, message } => remote_error(status, message),
other => CliError::Connect {
addr,
message: other.to_string(),
},
}
}
fn command_error(addr: SocketAddr, error: DcpError) -> CliError {
match error {
DcpError::Response { status, message } => remote_error(status, message),
other => CliError::Connect {
addr,
message: other.to_string(),
},
}
}
fn remote_error(status: ResponseStatus, message: String) -> CliError {
let fix = match status {
ResponseStatus::BadRequest => "Fix the command arguments and try again.",
ResponseStatus::Unauthorized => {
"Check the DCP token or mTLS client certificate configured for datum-agent."
}
ResponseStatus::NotFound => "Check the job or factory name with `datum ps`.",
ResponseStatus::Conflict => {
"Pick a different job name, or stop the existing job before starting it again."
}
ResponseStatus::ProtocolMismatch => {
"Upgrade the CLI and datum-agent to compatible versions."
}
ResponseStatus::DeadlineExceeded => {
"Retry the command or check whether datum-agent is overloaded."
}
ResponseStatus::Ok | ResponseStatus::Failed => {
"Check `datum status <job>` and the datum-agent logs, then retry."
}
};
CliError::Remote { message, fix }
}
async fn ps(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
let mut jobs = client
.list_jobs()
.await
.map_err(|error| command_error(addr, error))?;
jobs.sort_by(|left, right| left.name.cmp(&right.name));
if json {
print_json(&JobsOutput {
jobs: jobs.iter().map(JobView::from).collect(),
})?;
} else if jobs.is_empty() {
println!("no jobs");
} else {
print!("{}", render_jobs_table(&jobs));
}
Ok(())
}
async fn cluster_ps(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
let mut list = client
.list_cluster_jobs(1_000)
.await
.map_err(|error| command_error(addr, error))?;
sort_cluster_jobs(&mut list);
if json {
print_json(&ClusterJobsOutput::from(&list))?;
} else if list.nodes.iter().all(|node| node.jobs.is_empty()) && list.errors.is_empty() {
println!("no cluster jobs");
} else {
print!("{}", render_cluster_jobs_table(&list));
if !list.errors.is_empty() {
print!("{}", render_cluster_errors_table(&list.errors));
}
}
Ok(())
}
async fn nodes(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
let mut list = client
.cluster_node_info(1_000)
.await
.map_err(|error| command_error(addr, error))?;
list.nodes
.sort_by(|left, right| left.node_id.cmp(&right.node_id));
list.errors
.sort_by(|left, right| left.node_id.cmp(&right.node_id));
if json {
print_json(&ClusterNodesOutput::from(&list))?;
} else if list.nodes.is_empty() {
println!("no cluster nodes");
} else {
print!("{}", render_nodes_table(&list.nodes));
if !list.errors.is_empty() {
print!("{}", render_cluster_errors_table(&list.errors));
}
}
Ok(())
}
async fn status(
client: &DcpClient,
addr: SocketAddr,
job: &str,
cluster: bool,
json: bool,
) -> Result<(), CliError> {
let status = if cluster {
client.cluster_job_status(job, 1_000).await
} else {
client.job_status(job).await
}
.map_err(|error| command_error(addr, error))?;
if json {
print_json(&JobOutput {
job: JobView::from(&status),
})?;
} else {
print!("{}", render_status_table(&status));
}
Ok(())
}
struct SubmitArgs<'a> {
cluster: bool,
factory: &'a str,
name: Option<&'a str>,
raw_params: &'a [String],
role: Option<&'a str>,
node: Option<&'a str>,
}
async fn submit(
client: &DcpClient,
addr: SocketAddr,
args: SubmitArgs<'_>,
json: bool,
) -> Result<(), CliError> {
if !args.cluster {
return Err(CliError::Usage(
"`datum submit` currently requires --cluster; use `datum start` for local jobs."
.to_owned(),
));
}
let params = parse_params(args.raw_params)?;
let instance_name = args.name.unwrap_or(args.factory);
let placement = placement_spec(args.role, args.node);
let status = client
.submit_cluster_job(args.factory, instance_name, params, placement, 1_000)
.await
.map_err(|error| command_error(addr, error))?;
if json {
print_json(&ActionOutput {
action: "submitted",
message: format!("submitted cluster job '{}'", status.name),
job: JobView::from(&status),
})?;
} else {
println!(
"submitted cluster job '{}' from factory '{}' to node '{}' — state {}, placement generation {}",
status.name,
args.factory,
empty_dash(&status.placement_node_id),
status.state,
status.placement_generation
);
}
Ok(())
}
fn placement_spec(role: Option<&str>, node: Option<&str>) -> PlacementSpec {
PlacementSpec {
role_constraint: role.unwrap_or_default().to_owned(),
strategy: if node.is_some() {
PlacementStrategy::Pinned as i32
} else {
PlacementStrategy::LeastJobs as i32
},
pinned_node_id: node.unwrap_or_default().to_owned(),
}
}
async fn start(
client: &DcpClient,
addr: SocketAddr,
factory: &str,
name: Option<&str>,
raw_params: &[String],
json: bool,
) -> Result<(), CliError> {
let params = parse_params(raw_params)?;
let instance_name = name.unwrap_or(factory);
let status = client
.start_job(factory, instance_name, params)
.await
.map_err(|error| command_error(addr, error))?;
if json {
print_json(&ActionOutput {
action: "started",
message: format!("started '{}'", status.name),
job: JobView::from(&status),
})?;
} else {
println!(
"started '{}' from factory '{}' — state {}, generation {}",
status.name, factory, status.state, status.generation
);
}
Ok(())
}
async fn drain(
client: &DcpClient,
addr: SocketAddr,
job: &str,
cluster: bool,
json: bool,
) -> Result<(), CliError> {
let initial = if cluster {
client.drain_cluster_job(job, 1_000).await
} else {
client.drain_job(job).await
}
.map_err(|error| command_error(addr, error))?;
let status = wait_for_drain(client, job, cluster, initial)
.await
.map_err(|error| command_error(addr, error))?;
if json {
print_json(&ActionOutput {
action: "drained",
message: format!("drained '{}'", status.name),
job: JobView::from(&status),
})?;
} else if status.state == "Drained" {
println!("drained '{}' — in-flight streams completed", status.name);
} else {
println!(
"drain requested for '{}' — current state {}; run `datum status {}` to follow it",
status.name, status.state, status.name
);
}
Ok(())
}
async fn wait_for_drain(
client: &DcpClient,
job: &str,
cluster: bool,
mut status: WireJobStatus,
) -> Result<WireJobStatus, DcpError> {
let wait = status
.drain_remaining_ms
.map(|ms| Duration::from_millis(ms.saturating_add(1_000)))
.unwrap_or(Duration::from_secs(31))
.min(Duration::from_secs(60));
let started = tokio::time::Instant::now();
while !is_drain_terminal(&status) && started.elapsed() < wait {
tokio::time::sleep(DRAIN_POLL).await;
status = if cluster {
client.cluster_job_status(job, 1_000).await?
} else {
client.job_status(job).await?
};
}
Ok(status)
}
fn is_drain_terminal(status: &WireJobStatus) -> bool {
matches!(
status.state.as_str(),
"Completed" | "Drained" | "Stopped" | "Failed"
)
}
async fn stop(
client: &DcpClient,
addr: SocketAddr,
job: &str,
cluster: bool,
json: bool,
) -> Result<(), CliError> {
let status = if cluster {
client.stop_cluster_job(job, 1_000).await
} else {
client.stop_job(job).await
}
.map_err(|error| command_error(addr, error))?;
if json {
print_json(&ActionOutput {
action: "stopped",
message: format!("stopped '{}'", status.name),
job: JobView::from(&status),
})?;
} else {
println!("stopped '{}' — state {}", status.name, status.state);
}
Ok(())
}
async fn restart(
client: &DcpClient,
addr: SocketAddr,
job: &str,
json: bool,
) -> Result<(), CliError> {
let status = client
.restart_job(job)
.await
.map_err(|error| command_error(addr, error))?;
if json {
print_json(&ActionOutput {
action: "restarted",
message: format!("restarted '{}'", status.name),
job: JobView::from(&status),
})?;
} else {
println!(
"restarted '{}' — state {}, generation {}",
status.name, status.state, status.generation
);
}
Ok(())
}
async fn events(
client: &DcpClient,
addr: SocketAddr,
follow: bool,
json: bool,
) -> Result<(), CliError> {
let mut subscription = client
.subscribe_events()
.await
.map_err(|error| command_error(addr, error))?;
if follow {
while let Some(event) = subscription.recv().await {
print_event(&event, json)?;
}
if let Some(reason) = subscription.disconnect_reason() {
return Err(CliError::Connect {
addr,
message: reason,
});
}
return Ok(());
}
let events = collect_events(&mut subscription).await;
if json {
print_json(&EventsOutput {
events: events.iter().map(EventView::from).collect(),
})?;
} else if events.is_empty() {
println!("no events observed — run `datum events --follow` to keep waiting");
} else {
print!("{}", render_events_table(&events));
}
Ok(())
}
async fn collect_events(subscription: &mut EventSubscription) -> Vec<Event> {
let mut events = Vec::new();
loop {
let timeout = if events.is_empty() {
EVENT_INITIAL_IDLE
} else {
EVENT_IDLE
};
match tokio::time::timeout(timeout, subscription.recv()).await {
Ok(Some(event)) => events.push(event),
Ok(None) | Err(_) => break,
}
}
events
}
fn print_event(event: &Event, json: bool) -> Result<(), CliError> {
if json {
print_json_line(&EventLine {
event: EventView::from(event),
})?;
} else {
println!(
"{} {} '{}' generation {}",
event.sequence, event.kind, event.name, event.generation
);
io::stdout().flush().map_err(io_usage_error)?;
}
Ok(())
}
async fn metrics(
client: &DcpClient,
addr: SocketAddr,
job: &str,
follow: bool,
json: bool,
) -> Result<(), CliError> {
let mut subscription = client
.subscribe_metrics(0)
.await
.map_err(|error| command_error(addr, error))?;
if follow {
while let Some(sample) = subscription.recv().await {
let filtered = filter_metric_sample(job, &sample);
if filtered.streams.is_empty() {
continue;
}
print_metric_sample_follow(job, &filtered, json)?;
}
if let Some(reason) = subscription.disconnect_reason() {
return Err(CliError::Connect {
addr,
message: reason,
});
}
return Ok(());
}
let sample = wait_for_metric_sample(job, &mut subscription).await;
print_metric_sample(job, &sample, json)
}
async fn wait_for_metric_sample(
job: &str,
subscription: &mut MetricSubscription,
) -> MetricSampleView {
let deadline = tokio::time::Instant::now() + METRIC_SNAPSHOT_WAIT;
let mut latest = MetricSampleView {
timestamp_ms: 0,
streams: Vec::new(),
};
while tokio::time::Instant::now() < deadline {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
match tokio::time::timeout(
remaining.min(Duration::from_millis(100)),
subscription.recv(),
)
.await
{
Ok(Some(sample)) => {
latest = filter_metric_sample(job, &sample);
if !latest.streams.is_empty() {
break;
}
}
Ok(None) => break,
Err(_) => {}
}
}
latest
}
fn print_metric_sample(job: &str, sample: &MetricSampleView, json: bool) -> Result<(), CliError> {
if json {
print_json(&MetricsOutput { job, sample })?;
} else if sample.streams.is_empty() {
println!("no metrics for '{job}' yet");
} else {
print!("{}", render_metrics_table(sample));
}
Ok(())
}
fn print_metric_sample_follow(
job: &str,
sample: &MetricSampleView,
json: bool,
) -> Result<(), CliError> {
if json {
print_json_line(&MetricsOutput { job, sample })?;
} else {
print!("{}", render_metrics_table(sample));
io::stdout().flush().map_err(io_usage_error)?;
}
Ok(())
}
fn filter_metric_sample(job: &str, sample: &MetricSample) -> MetricSampleView {
let prefix = format!("{job}:");
MetricSampleView {
timestamp_ms: sample.timestamp_ms,
streams: sample
.streams
.iter()
.filter(|metric| metric.name == job || metric.name.starts_with(&prefix))
.map(StreamMetricView::from)
.collect(),
}
}
fn sort_cluster_jobs(list: &mut ClusterJobList) {
list.nodes
.sort_by(|left, right| left.node_id.cmp(&right.node_id));
for node in &mut list.nodes {
node.jobs.sort_by(|left, right| left.name.cmp(&right.name));
}
list.errors
.sort_by(|left, right| left.node_id.cmp(&right.node_id));
}
fn parse_params(raw_params: &[String]) -> Result<HashMap<String, String>, CliError> {
let mut params = HashMap::new();
let mut seen = HashSet::new();
for raw in raw_params {
let Some((key, value)) = raw.split_once('=') else {
return Err(CliError::Usage(format!(
"invalid --param '{raw}'; use --param key=value."
)));
};
if key.trim().is_empty() {
return Err(CliError::Usage(format!(
"invalid --param '{raw}'; the key before '=' cannot be empty."
)));
}
if !seen.insert(key.to_owned()) {
return Err(CliError::Usage(format!(
"duplicate --param key '{key}'; pass each key once."
)));
}
params.insert(key.to_owned(), value.to_owned());
}
Ok(params)
}
#[derive(Serialize)]
struct JobsOutput {
jobs: Vec<JobView>,
}
#[derive(Serialize)]
struct ClusterJobsOutput {
partial: bool,
nodes: Vec<ClusterJobNodeView>,
errors: Vec<ClusterNodeErrorView>,
}
impl From<&ClusterJobList> for ClusterJobsOutput {
fn from(list: &ClusterJobList) -> Self {
Self {
partial: list.partial,
nodes: list.nodes.iter().map(ClusterJobNodeView::from).collect(),
errors: list.errors.iter().map(ClusterNodeErrorView::from).collect(),
}
}
}
#[derive(Serialize)]
struct ClusterNodesOutput {
partial: bool,
nodes: Vec<ClusterNodeView>,
errors: Vec<ClusterNodeErrorView>,
}
impl From<&ClusterNodeList> for ClusterNodesOutput {
fn from(list: &ClusterNodeList) -> Self {
Self {
partial: list.partial,
nodes: list.nodes.iter().map(ClusterNodeView::from).collect(),
errors: list.errors.iter().map(ClusterNodeErrorView::from).collect(),
}
}
}
#[derive(Serialize)]
struct JobOutput {
job: JobView,
}
#[derive(Serialize)]
struct ActionOutput {
action: &'static str,
message: String,
job: JobView,
}
#[derive(Serialize)]
struct EventsOutput {
events: Vec<EventView>,
}
#[derive(Serialize)]
struct EventLine {
event: EventView,
}
#[derive(Serialize)]
struct MetricsOutput<'a> {
job: &'a str,
sample: &'a MetricSampleView,
}
#[derive(Serialize)]
struct JobView {
name: String,
job_id: u64,
state: String,
desired_state: String,
generation: u64,
starts_total: u64,
restarts_total: u64,
last_start_at_ms: Option<u64>,
last_exit_at_ms: Option<u64>,
last_exit_reason: String,
backoff_remaining_ms: Option<u64>,
drain_remaining_ms: Option<u64>,
drain_supported: bool,
active_streams: Option<u64>,
cluster_job: bool,
factory_name: String,
placement_strategy: String,
role_constraint: String,
coordinator_node_id: String,
placement_node_id: String,
placement_generation: u64,
placement_history: Vec<PlacementHistoryView>,
}
impl From<&WireJobStatus> for JobView {
fn from(status: &WireJobStatus) -> Self {
Self {
name: status.name.clone(),
job_id: status.job_id,
state: status.state.clone(),
desired_state: status.desired_state.clone(),
generation: status.generation,
starts_total: status.starts_total,
restarts_total: status.restarts_total,
last_start_at_ms: status.last_start_at_ms,
last_exit_at_ms: status.last_exit_at_ms,
last_exit_reason: status.last_exit_reason.clone(),
backoff_remaining_ms: status.backoff_remaining_ms,
drain_remaining_ms: status.drain_remaining_ms,
drain_supported: status.drain_supported,
active_streams: status.active_streams,
cluster_job: status.cluster_job,
factory_name: status.factory_name.clone(),
placement_strategy: placement_strategy_text(status),
role_constraint: status
.placement
.as_ref()
.map(|placement| placement.role_constraint.clone())
.unwrap_or_default(),
coordinator_node_id: status.coordinator_node_id.clone(),
placement_node_id: status.placement_node_id.clone(),
placement_generation: status.placement_generation,
placement_history: status
.placement_history
.iter()
.map(PlacementHistoryView::from)
.collect(),
}
}
}
#[derive(Serialize)]
struct PlacementHistoryView {
generation: u64,
from_node_id: String,
to_node_id: String,
reason: String,
timestamp_ms: u64,
}
impl From<&datum_agent::dcp::proto::ClusterPlacementHistory> for PlacementHistoryView {
fn from(history: &datum_agent::dcp::proto::ClusterPlacementHistory) -> Self {
Self {
generation: history.generation,
from_node_id: history.from_node_id.clone(),
to_node_id: history.to_node_id.clone(),
reason: history.reason.clone(),
timestamp_ms: history.timestamp_ms,
}
}
}
fn placement_strategy_text(status: &WireJobStatus) -> String {
let Some(placement) = status.placement.as_ref() else {
return String::new();
};
match PlacementStrategy::try_from(placement.strategy).unwrap_or(PlacementStrategy::LeastJobs) {
PlacementStrategy::LeastJobs => "LeastJobs".to_owned(),
PlacementStrategy::Pinned => format!("Pinned({})", placement.pinned_node_id),
}
}
#[derive(Serialize)]
struct ClusterJobNodeView {
node_id: String,
address: String,
local: bool,
jobs: Vec<JobView>,
}
impl From<&ClusterJobNode> for ClusterJobNodeView {
fn from(node: &ClusterJobNode) -> Self {
Self {
node_id: node.node_id.clone(),
address: node.address.clone(),
local: node.local,
jobs: node.jobs.iter().map(JobView::from).collect(),
}
}
}
#[derive(Serialize)]
struct ClusterNodeView {
node_id: String,
member_state: String,
address: String,
agent_addr: String,
roles: Vec<String>,
unreachable: bool,
local: bool,
session_state: String,
}
impl From<&ClusterNodeStatus> for ClusterNodeView {
fn from(node: &ClusterNodeStatus) -> Self {
Self {
node_id: node.node_id.clone(),
member_state: node.member_state.clone(),
address: node.address.clone(),
agent_addr: node.agent_addr.clone(),
roles: node.roles.clone(),
unreachable: node.unreachable,
local: node.local,
session_state: node.session_state.clone(),
}
}
}
#[derive(Serialize)]
struct ClusterNodeErrorView {
node_id: String,
message: String,
}
impl From<&ClusterNodeError> for ClusterNodeErrorView {
fn from(error: &ClusterNodeError) -> Self {
Self {
node_id: error.node_id.clone(),
message: error.message.clone(),
}
}
}
#[derive(Serialize)]
struct EventView {
sequence: u64,
timestamp_ms: u64,
name: String,
job_id: u64,
generation: u64,
kind: String,
detail: String,
}
impl From<&Event> for EventView {
fn from(event: &Event) -> Self {
Self {
sequence: event.sequence,
timestamp_ms: event.timestamp_ms,
name: event.name.clone(),
job_id: event.job_id,
generation: event.generation,
kind: event.kind.clone(),
detail: event.detail.clone(),
}
}
}
#[derive(Serialize)]
struct MetricSampleView {
timestamp_ms: u64,
streams: Vec<StreamMetricView>,
}
#[derive(Serialize)]
struct StreamMetricView {
id: u64,
name: String,
elements_through: u64,
restarts: u64,
state: String,
started_at_ms: u64,
state_changed_at_ms: u64,
finished_at_ms: Option<u64>,
uptime_ms: u64,
}
impl From<&StreamMetric> for StreamMetricView {
fn from(metric: &StreamMetric) -> Self {
Self {
id: metric.id,
name: metric.name.clone(),
elements_through: metric.elements_through,
restarts: metric.restarts,
state: metric.state.clone(),
started_at_ms: metric.started_at_ms,
state_changed_at_ms: metric.state_changed_at_ms,
finished_at_ms: metric.finished_at_ms,
uptime_ms: metric.uptime_ms,
}
}
}
fn print_json<T: Serialize>(value: &T) -> Result<(), CliError> {
let json = serde_json::to_string_pretty(value).map_err(|error| {
CliError::Usage(format!(
"could not encode JSON output: {error}; retry without --json."
))
})?;
println!("{json}");
Ok(())
}
fn print_json_line<T: Serialize>(value: &T) -> Result<(), CliError> {
let json = serde_json::to_string(value).map_err(|error| {
CliError::Usage(format!(
"could not encode JSON output: {error}; retry without --json."
))
})?;
println!("{json}");
io::stdout().flush().map_err(io_usage_error)?;
Ok(())
}
fn io_usage_error(error: io::Error) -> CliError {
CliError::Usage(format!(
"could not write CLI output: {error}; check stdout/stderr."
))
}
fn render_jobs_table(jobs: &[WireJobStatus]) -> String {
let rows = jobs
.iter()
.map(|job| {
vec![
job.name.clone(),
job.job_id.to_string(),
job.state.clone(),
job.desired_state.clone(),
job.generation.to_string(),
format_number(job.starts_total),
format_number(job.restarts_total),
format_optional_number(job.active_streams),
]
})
.collect::<Vec<_>>();
render_table(
&[
"JOB", "ID", "STATE", "DESIRED", "GEN", "STARTS", "RESTARTS", "ACTIVE",
],
&rows,
)
}
fn render_cluster_jobs_table(list: &ClusterJobList) -> String {
let rows = list
.nodes
.iter()
.flat_map(|node| {
if node.jobs.is_empty() {
vec![vec![
node.node_id.clone(),
node.address.clone(),
"-".to_owned(),
"-".to_owned(),
"-".to_owned(),
"-".to_owned(),
"-".to_owned(),
"-".to_owned(),
"-".to_owned(),
"-".to_owned(),
"-".to_owned(),
"-".to_owned(),
]]
} else {
node.jobs
.iter()
.map(|job| {
vec![
node.node_id.clone(),
node.address.clone(),
job.name.clone(),
job.job_id.to_string(),
job.state.clone(),
job.desired_state.clone(),
job.generation.to_string(),
format_number(job.starts_total),
format_number(job.restarts_total),
cluster_marker(job),
empty_dash(&job.placement_node_id).to_owned(),
placement_generation_text(job),
]
})
.collect::<Vec<_>>()
}
})
.collect::<Vec<_>>();
render_table(
&[
"NODE",
"ADDRESS",
"JOB",
"ID",
"STATE",
"DESIRED",
"GEN",
"STARTS",
"RESTARTS",
"CLUSTER",
"PLACED_ON",
"PGEN",
],
&rows,
)
}
fn render_nodes_table(nodes: &[ClusterNodeStatus]) -> String {
let rows = nodes
.iter()
.map(|node| {
vec![
node.node_id.clone(),
node.member_state.clone(),
node.unreachable.to_string(),
node.session_state.clone(),
node.agent_addr.clone(),
node.roles.join(","),
]
})
.collect::<Vec<_>>();
render_table(
&[
"NODE",
"MEMBER",
"UNREACHABLE",
"SESSION",
"AGENT_ADDR",
"ROLES",
],
&rows,
)
}
fn render_cluster_errors_table(errors: &[ClusterNodeError]) -> String {
let rows = errors
.iter()
.map(|error| vec![error.node_id.clone(), error.message.clone()])
.collect::<Vec<_>>();
render_table(&["NODE", "ERROR"], &rows)
}
fn cluster_marker(job: &WireJobStatus) -> String {
if job.cluster_job {
"yes".to_owned()
} else {
"-".to_owned()
}
}
fn placement_generation_text(job: &WireJobStatus) -> String {
if job.cluster_job {
job.placement_generation.to_string()
} else {
"-".to_owned()
}
}
fn render_status_table(status: &WireJobStatus) -> String {
let rows = vec![
vec!["job".to_owned(), status.name.clone()],
vec!["id".to_owned(), status.job_id.to_string()],
vec!["state".to_owned(), status.state.clone()],
vec!["desired".to_owned(), status.desired_state.clone()],
vec!["generation".to_owned(), status.generation.to_string()],
vec!["starts".to_owned(), format_number(status.starts_total)],
vec!["restarts".to_owned(), format_number(status.restarts_total)],
vec![
"last_start_ms".to_owned(),
format_optional_number(status.last_start_at_ms),
],
vec![
"last_exit_ms".to_owned(),
format_optional_number(status.last_exit_at_ms),
],
vec![
"last_exit_reason".to_owned(),
empty_dash(&status.last_exit_reason).to_owned(),
],
vec![
"backoff_remaining_ms".to_owned(),
format_optional_number(status.backoff_remaining_ms),
],
vec![
"drain_remaining_ms".to_owned(),
format_optional_number(status.drain_remaining_ms),
],
vec![
"drain_supported".to_owned(),
status.drain_supported.to_string(),
],
vec![
"active_streams".to_owned(),
format_optional_number(status.active_streams),
],
vec!["cluster_job".to_owned(), status.cluster_job.to_string()],
vec![
"factory".to_owned(),
empty_dash(&status.factory_name).to_owned(),
],
vec![
"placement_strategy".to_owned(),
empty_dash(&placement_strategy_text(status)).to_owned(),
],
vec![
"role_constraint".to_owned(),
status
.placement
.as_ref()
.map(|placement| empty_dash(&placement.role_constraint).to_owned())
.unwrap_or_else(|| "-".to_owned()),
],
vec![
"coordinator".to_owned(),
empty_dash(&status.coordinator_node_id).to_owned(),
],
vec![
"placed_on".to_owned(),
empty_dash(&status.placement_node_id).to_owned(),
],
vec![
"placement_generation".to_owned(),
placement_generation_text(status),
],
];
render_table(&["FIELD", "VALUE"], &rows)
}
fn render_events_table(events: &[Event]) -> String {
let rows = events
.iter()
.map(|event| {
vec![
event.sequence.to_string(),
event.timestamp_ms.to_string(),
event.name.clone(),
event.job_id.to_string(),
event.generation.to_string(),
event.kind.clone(),
empty_dash(&event.detail).to_owned(),
]
})
.collect::<Vec<_>>();
render_table(
&["SEQ", "TIME_MS", "JOB", "ID", "GEN", "KIND", "DETAIL"],
&rows,
)
}
fn render_metrics_table(sample: &MetricSampleView) -> String {
let rows = sample
.streams
.iter()
.map(|metric| {
vec![
sample.timestamp_ms.to_string(),
metric.name.clone(),
metric.id.to_string(),
metric.state.clone(),
format_number(metric.elements_through),
format_number(metric.restarts),
metric.uptime_ms.to_string(),
]
})
.collect::<Vec<_>>();
render_table(
&[
"TIME_MS",
"STREAM",
"ID",
"STATE",
"ELEMENTS",
"RESTARTS",
"UPTIME_MS",
],
&rows,
)
}
#[must_use]
pub fn render_table(headers: &[&str], rows: &[Vec<String>]) -> String {
let mut widths = headers
.iter()
.map(|header| header.len())
.collect::<Vec<_>>();
for row in rows {
for (index, cell) in row.iter().enumerate().take(widths.len()) {
widths[index] = widths[index].max(cell.len());
}
}
let mut output = String::new();
push_table_row(
&mut output,
&headers
.iter()
.map(|cell| (*cell).to_owned())
.collect::<Vec<_>>(),
&widths,
);
push_separator(&mut output, &widths);
for row in rows {
push_table_row(&mut output, row, &widths);
}
output
}
fn push_table_row(output: &mut String, row: &[String], widths: &[usize]) {
for (index, width) in widths.iter().enumerate() {
if index > 0 {
output.push_str(" ");
}
let cell = row.get(index).map(String::as_str).unwrap_or("");
output.push_str(cell);
for _ in cell.len()..*width {
output.push(' ');
}
}
output.push('\n');
}
fn push_separator(output: &mut String, widths: &[usize]) {
for (index, width) in widths.iter().enumerate() {
if index > 0 {
output.push_str(" ");
}
for _ in 0..*width {
output.push('-');
}
}
output.push('\n');
}
fn empty_dash(value: &str) -> &str {
if value.is_empty() { "-" } else { value }
}
fn format_optional_number(value: Option<u64>) -> String {
value.map(format_number).unwrap_or_else(|| "-".to_owned())
}
fn format_number(value: u64) -> String {
let raw = value.to_string();
let mut output = String::with_capacity(raw.len() + raw.len() / 3);
for (index, ch) in raw.chars().rev().enumerate() {
if index > 0 && index % 3 == 0 {
output.push(',');
}
output.push(ch);
}
output.chars().rev().collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn table_renderer_pads_columns_stably() {
let output = render_table(
&["JOB", "STATE", "STARTS"],
&[
vec!["ingest".to_owned(), "Running".to_owned(), "1".to_owned()],
vec![
"daily-rollup".to_owned(),
"Drained".to_owned(),
"12,003".to_owned(),
],
],
);
assert_eq!(
output,
"JOB STATE STARTS\n\
------------ ------- ------\n\
ingest Running 1 \n\
daily-rollup Drained 12,003\n"
);
}
#[test]
fn params_require_key_value_pairs() {
assert!(parse_params(&["threads=4".to_owned()]).is_ok());
assert!(matches!(
parse_params(&["threads".to_owned()]),
Err(CliError::Usage(_))
));
assert!(matches!(
parse_params(&["threads=4".to_owned(), "threads=8".to_owned()]),
Err(CliError::Usage(_))
));
}
#[test]
fn numbers_use_operator_grouping() {
assert_eq!(format_number(0), "0");
assert_eq!(format_number(12_003), "12,003");
assert_eq!(format_number(9_876_543_210), "9,876,543,210");
}
}