1#![forbid(unsafe_code)]
2
3use std::{
4 collections::{HashMap, HashSet},
5 fs,
6 io::{self, Cursor, Write},
7 net::SocketAddr,
8 path::{Path, PathBuf},
9 sync::Arc,
10 time::Duration,
11};
12
13use clap::{Parser, Subcommand, error::ErrorKind};
14use datum_agent::dcp::{
15 ClientKind, DcpClient, DcpError, EventSubscription, Hello, MetricSubscription, ResponseStatus,
16 proto::{
17 ClusterJobList, ClusterJobNode, ClusterNodeError, ClusterNodeList, ClusterNodeStatus,
18 Event, JobStatus as WireJobStatus, MetricSample, PlacementSpec, PlacementStrategy,
19 StreamMetric,
20 },
21};
22use datum_net::quic::{
23 crypto::rustls::QuicClientConfig,
24 quinn,
25 rustls::{
26 ClientConfig as RustlsClientConfig, RootCertStore,
27 pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer},
28 },
29};
30use serde::Serialize;
31
32const DEFAULT_ADDR: &str = "127.0.0.1:9555";
33const EXIT_OK: i32 = 0;
34const EXIT_REMOTE: i32 = 1;
35const EXIT_USAGE_OR_CONNECT: i32 = 2;
36const EVENT_INITIAL_IDLE: Duration = Duration::from_millis(500);
37const EVENT_IDLE: Duration = Duration::from_millis(250);
38const METRIC_SNAPSHOT_WAIT: Duration = Duration::from_secs(2);
39const DRAIN_POLL: Duration = Duration::from_millis(20);
40
41#[derive(Debug, Parser)]
42#[command(
43 name = "datum",
44 version,
45 about = "Manage Datum jobs through the Datum Control Protocol"
46)]
47struct Cli {
48 #[arg(long, global = true, default_value = DEFAULT_ADDR)]
50 addr: SocketAddr,
51
52 #[arg(long, global = true, value_name = "PATH")]
54 tls_ca: Option<PathBuf>,
55
56 #[arg(long, global = true, value_name = "PATH")]
58 tls_cert: Option<PathBuf>,
59
60 #[arg(long, global = true, value_name = "PATH")]
62 tls_key: Option<PathBuf>,
63
64 #[arg(long, global = true)]
66 json: bool,
67
68 #[command(subcommand)]
69 command: Command,
70}
71
72#[derive(Debug, Subcommand)]
73enum Command {
74 Ps {
76 #[arg(long)]
78 cluster: bool,
79 },
80 Nodes,
82 Status {
84 #[arg(long)]
85 cluster: bool,
86 job: String,
87 },
88 Start {
90 factory: String,
91 #[arg(long)]
92 name: Option<String>,
93 #[arg(long = "param", value_name = "k=v")]
94 params: Vec<String>,
95 },
96 Submit {
98 #[arg(long)]
99 cluster: bool,
100 #[arg(long)]
101 factory: String,
102 #[arg(long)]
103 name: Option<String>,
104 #[arg(long = "param", value_name = "k=v")]
105 params: Vec<String>,
106 #[arg(long)]
107 role: Option<String>,
108 #[arg(long)]
109 node: Option<String>,
110 },
111 Drain {
113 #[arg(long)]
114 cluster: bool,
115 job: String,
116 },
117 Stop {
119 #[arg(long)]
120 cluster: bool,
121 job: String,
122 },
123 Restart { job: String },
125 Events {
127 #[arg(long)]
128 follow: bool,
129 },
130 Metrics {
132 job: String,
133 #[arg(long)]
134 follow: bool,
135 },
136}
137
138#[derive(Debug)]
139enum CliError {
140 Usage(String),
141 Connect { addr: SocketAddr, message: String },
142 Remote { message: String, fix: &'static str },
143}
144
145impl CliError {
146 fn exit_code(&self) -> i32 {
147 match self {
148 Self::Remote { .. } => EXIT_REMOTE,
149 Self::Usage(_) | Self::Connect { .. } => EXIT_USAGE_OR_CONNECT,
150 }
151 }
152
153 fn print(&self) {
154 match self {
155 Self::Usage(message) => {
156 eprintln!("{message}");
157 }
158 Self::Connect { addr, message } => {
159 eprintln!("could not connect to datum-agent at {addr}: {message}");
160 eprintln!(
161 "start datum-agent or pass --addr host:port; loopback TCP defaults to {DEFAULT_ADDR}"
162 );
163 }
164 Self::Remote { message, fix } => {
165 eprintln!("datum-agent rejected the request: {message}");
166 eprintln!("{fix}");
167 }
168 }
169 }
170}
171
172pub async fn main_entry() -> i32 {
174 match Cli::try_parse() {
175 Ok(cli) => match run(cli).await {
176 Ok(()) => EXIT_OK,
177 Err(error) => {
178 error.print();
179 error.exit_code()
180 }
181 },
182 Err(error) => {
183 let _ = error.print();
184 clap_error_exit_code(&error)
185 }
186 }
187}
188
189fn clap_error_exit_code(error: &clap::Error) -> i32 {
192 match error.kind() {
193 ErrorKind::DisplayHelp
194 | ErrorKind::DisplayVersion
195 | ErrorKind::DisplayHelpOnMissingArgumentOrSubcommand => EXIT_OK,
196 _ => EXIT_USAGE_OR_CONNECT,
197 }
198}
199
200async fn run(cli: Cli) -> Result<(), CliError> {
201 validate_before_connect(&cli.command)?;
202 let client = connect(&cli).await?;
203 match &cli.command {
204 Command::Ps { cluster } => {
205 if *cluster {
206 cluster_ps(&client, cli.addr, cli.json).await
207 } else {
208 ps(&client, cli.addr, cli.json).await
209 }
210 }
211 Command::Nodes => nodes(&client, cli.addr, cli.json).await,
212 Command::Status { cluster, job } => {
213 status(&client, cli.addr, job, *cluster, cli.json).await
214 }
215 Command::Start {
216 factory,
217 name,
218 params,
219 } => {
220 start(
221 &client,
222 cli.addr,
223 factory,
224 name.as_deref(),
225 params,
226 cli.json,
227 )
228 .await
229 }
230 Command::Submit {
231 cluster,
232 factory,
233 name,
234 params,
235 role,
236 node,
237 } => {
238 submit(
239 &client,
240 cli.addr,
241 SubmitArgs {
242 cluster: *cluster,
243 factory,
244 name: name.as_deref(),
245 raw_params: params,
246 role: role.as_deref(),
247 node: node.as_deref(),
248 },
249 cli.json,
250 )
251 .await
252 }
253 Command::Drain { cluster, job } => drain(&client, cli.addr, job, *cluster, cli.json).await,
254 Command::Stop { cluster, job } => stop(&client, cli.addr, job, *cluster, cli.json).await,
255 Command::Restart { job } => restart(&client, cli.addr, job, cli.json).await,
256 Command::Events { follow } => events(&client, cli.addr, *follow, cli.json).await,
257 Command::Metrics { job, follow } => {
258 metrics(&client, cli.addr, job, *follow, cli.json).await
259 }
260 }
261}
262
263fn validate_before_connect(command: &Command) -> Result<(), CliError> {
264 match command {
265 Command::Start { params, .. } => {
266 parse_params(params)?;
267 }
268 Command::Submit {
269 cluster,
270 params,
271 role,
272 node,
273 ..
274 } => {
275 if !cluster {
276 return Err(CliError::Usage(
277 "`datum submit` currently requires --cluster; use `datum start` for local jobs."
278 .to_owned(),
279 ));
280 }
281 if role.is_some() && node.is_some() {
282 return Err(CliError::Usage(
283 "pass at most one of --role or --node for cluster placement.".to_owned(),
284 ));
285 }
286 parse_params(params)?;
287 }
288 _ => {}
289 }
290 Ok(())
291}
292
293async fn connect(cli: &Cli) -> Result<DcpClient, CliError> {
294 let hello = Hello::new(format!("datum-cli-{}", std::process::id()), ClientKind::Cli);
295 match tls_paths(cli)? {
296 Some(paths) => {
297 let config = load_quic_client_config(paths.ca, paths.cert, paths.key)?;
298 DcpClient::connect_quic(cli.addr, "localhost", config, hello)
299 .await
300 .map_err(|error| connect_error(cli.addr, error))
301 }
302 None => DcpClient::connect_tcp(cli.addr, hello)
303 .await
304 .map_err(|error| connect_error(cli.addr, error)),
305 }
306}
307
308struct TlsPaths<'a> {
309 ca: &'a Path,
310 cert: &'a Path,
311 key: &'a Path,
312}
313
314fn tls_paths(cli: &Cli) -> Result<Option<TlsPaths<'_>>, CliError> {
315 match (&cli.tls_ca, &cli.tls_cert, &cli.tls_key) {
316 (None, None, None) => Ok(None),
317 (Some(ca), Some(cert), Some(key)) => Ok(Some(TlsPaths {
318 ca: ca.as_path(),
319 cert: cert.as_path(),
320 key: key.as_path(),
321 })),
322 _ => Err(CliError::Usage(
323 "QUIC+mTLS mode needs --tls-ca, --tls-cert, and --tls-key; pass all three or omit all three for loopback TCP.".to_owned(),
324 )),
325 }
326}
327
328fn load_quic_client_config(
329 ca_path: &Path,
330 cert_path: &Path,
331 key_path: &Path,
332) -> Result<quinn::ClientConfig, CliError> {
333 let mut roots = RootCertStore::empty();
334 for cert in read_certificates(ca_path)? {
335 roots.add(cert).map_err(|error| {
336 CliError::Usage(format!(
337 "could not load --tls-ca {}: {error}; pass a PEM or DER CA certificate.",
338 ca_path.display()
339 ))
340 })?;
341 }
342 let certs = read_certificates(cert_path)?;
343 let key = read_private_key(key_path)?;
344 let rustls = RustlsClientConfig::builder()
345 .with_root_certificates(roots)
346 .with_client_auth_cert(certs, key)
347 .map_err(|error| {
348 CliError::Usage(format!(
349 "could not build QUIC client identity: {error}; check --tls-cert and --tls-key match."
350 ))
351 })?;
352 let quic = QuicClientConfig::try_from(rustls).map_err(|error| {
353 CliError::Usage(format!(
354 "could not build QUIC client config: {error}; check the TLS files."
355 ))
356 })?;
357 Ok(quinn::ClientConfig::new(Arc::new(quic)))
358}
359
360fn read_certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, CliError> {
361 let bytes = fs::read(path).map_err(|error| {
362 CliError::Usage(format!(
363 "could not read {}: {error}; pass a readable PEM or DER certificate file.",
364 path.display()
365 ))
366 })?;
367 let mut cursor = Cursor::new(bytes.as_slice());
368 let certs = rustls_pemfile::certs(&mut cursor)
369 .collect::<Result<Vec<_>, _>>()
370 .map_err(|error| {
371 CliError::Usage(format!(
372 "could not parse {}: {error}; pass a PEM or DER certificate file.",
373 path.display()
374 ))
375 })?;
376 if certs.is_empty() {
377 Ok(vec![CertificateDer::from(bytes)])
378 } else {
379 Ok(certs)
380 }
381}
382
383fn read_private_key(path: &Path) -> Result<PrivateKeyDer<'static>, CliError> {
384 let bytes = fs::read(path).map_err(|error| {
385 CliError::Usage(format!(
386 "could not read {}: {error}; pass a readable PEM or DER PKCS#8 private key.",
387 path.display()
388 ))
389 })?;
390 let mut cursor = Cursor::new(bytes.as_slice());
391 match rustls_pemfile::private_key(&mut cursor).map_err(|error| {
392 CliError::Usage(format!(
393 "could not parse {}: {error}; pass a PEM or DER PKCS#8 private key.",
394 path.display()
395 ))
396 })? {
397 Some(key) => Ok(key),
398 None => Ok(PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(bytes))),
399 }
400}
401
402fn connect_error(addr: SocketAddr, error: DcpError) -> CliError {
403 match error {
404 DcpError::Response { status, message } => remote_error(status, message),
405 other => CliError::Connect {
406 addr,
407 message: other.to_string(),
408 },
409 }
410}
411
412fn command_error(addr: SocketAddr, error: DcpError) -> CliError {
413 match error {
414 DcpError::Response { status, message } => remote_error(status, message),
415 other => CliError::Connect {
416 addr,
417 message: other.to_string(),
418 },
419 }
420}
421
422fn remote_error(status: ResponseStatus, message: String) -> CliError {
423 let fix = match status {
424 ResponseStatus::BadRequest => "Fix the command arguments and try again.",
425 ResponseStatus::Unauthorized => {
426 "Check the DCP token or mTLS client certificate configured for datum-agent."
427 }
428 ResponseStatus::NotFound => "Check the job or factory name with `datum ps`.",
429 ResponseStatus::Conflict => {
430 "Pick a different job name, or stop the existing job before starting it again."
431 }
432 ResponseStatus::ProtocolMismatch => {
433 "Upgrade the CLI and datum-agent to compatible versions."
434 }
435 ResponseStatus::DeadlineExceeded => {
436 "Retry the command or check whether datum-agent is overloaded."
437 }
438 ResponseStatus::Ok | ResponseStatus::Failed => {
439 "Check `datum status <job>` and the datum-agent logs, then retry."
440 }
441 };
442 CliError::Remote { message, fix }
443}
444
445async fn ps(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
446 let mut jobs = client
447 .list_jobs()
448 .await
449 .map_err(|error| command_error(addr, error))?;
450 jobs.sort_by(|left, right| left.name.cmp(&right.name));
451 if json {
452 print_json(&JobsOutput {
453 jobs: jobs.iter().map(JobView::from).collect(),
454 })?;
455 } else if jobs.is_empty() {
456 println!("no jobs");
457 } else {
458 print!("{}", render_jobs_table(&jobs));
459 }
460 Ok(())
461}
462
463async fn cluster_ps(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
464 let mut list = client
465 .list_cluster_jobs(1_000)
466 .await
467 .map_err(|error| command_error(addr, error))?;
468 sort_cluster_jobs(&mut list);
469 if json {
470 print_json(&ClusterJobsOutput::from(&list))?;
471 } else if list.nodes.iter().all(|node| node.jobs.is_empty()) && list.errors.is_empty() {
472 println!("no cluster jobs");
473 } else {
474 print!("{}", render_cluster_jobs_table(&list));
475 if !list.errors.is_empty() {
476 print!("{}", render_cluster_errors_table(&list.errors));
477 }
478 }
479 Ok(())
480}
481
482async fn nodes(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
483 let mut list = client
484 .cluster_node_info(1_000)
485 .await
486 .map_err(|error| command_error(addr, error))?;
487 list.nodes
488 .sort_by(|left, right| left.node_id.cmp(&right.node_id));
489 list.errors
490 .sort_by(|left, right| left.node_id.cmp(&right.node_id));
491 let coordinator_node_id = cluster_coordinator_node_id(client, &list).await;
492 if json {
493 print_json(&ClusterNodesOutput::new(&list, coordinator_node_id.clone()))?;
494 } else if list.nodes.is_empty() {
495 println!("no cluster nodes");
496 } else {
497 print!(
498 "{}",
499 render_nodes_table(&list.nodes, coordinator_node_id.as_deref())
500 );
501 if !list.errors.is_empty() {
502 print!("{}", render_cluster_errors_table(&list.errors));
503 }
504 }
505 Ok(())
506}
507
508async fn cluster_coordinator_node_id(
511 client: &DcpClient,
512 node_info: &ClusterNodeList,
513) -> Option<String> {
514 if let Some(coordinator_node_id) = cluster_coordinator_node_id_from(node_info, None) {
515 return Some(coordinator_node_id);
516 }
517 let jobs = client.list_cluster_jobs(1_000).await.ok()?;
518 cluster_coordinator_node_id_from(node_info, Some(&jobs))
519}
520
521fn cluster_coordinator_node_id_from(
522 node_info: &ClusterNodeList,
523 jobs: Option<&ClusterJobList>,
524) -> Option<String> {
525 if !node_info.coordinator_node_id.is_empty() {
526 return Some(node_info.coordinator_node_id.clone());
527 }
528 jobs.and_then(cluster_coordinator_node_id_from_jobs)
529}
530
531fn cluster_coordinator_node_id_from_jobs(jobs: &ClusterJobList) -> Option<String> {
532 jobs.nodes
533 .iter()
534 .flat_map(|node| node.jobs.iter())
535 .map(|job| job.coordinator_node_id.as_str())
536 .find(|id| !id.is_empty())
537 .map(str::to_owned)
538}
539
540async fn status(
541 client: &DcpClient,
542 addr: SocketAddr,
543 job: &str,
544 cluster: bool,
545 json: bool,
546) -> Result<(), CliError> {
547 let status = if cluster {
548 client.cluster_job_status(job, 1_000).await
549 } else {
550 client.job_status(job).await
551 }
552 .map_err(|error| command_error(addr, error))?;
553 if json {
554 print_json(&JobOutput {
555 job: JobView::from(&status),
556 })?;
557 } else {
558 print!("{}", render_status_table(&status));
559 }
560 Ok(())
561}
562
563struct SubmitArgs<'a> {
564 cluster: bool,
565 factory: &'a str,
566 name: Option<&'a str>,
567 raw_params: &'a [String],
568 role: Option<&'a str>,
569 node: Option<&'a str>,
570}
571
572async fn submit(
573 client: &DcpClient,
574 addr: SocketAddr,
575 args: SubmitArgs<'_>,
576 json: bool,
577) -> Result<(), CliError> {
578 if !args.cluster {
579 return Err(CliError::Usage(
580 "`datum submit` currently requires --cluster; use `datum start` for local jobs."
581 .to_owned(),
582 ));
583 }
584 let params = parse_params(args.raw_params)?;
585 let instance_name = args.name.unwrap_or(args.factory);
586 let placement = placement_spec(args.role, args.node);
587 let status = client
588 .submit_cluster_job(args.factory, instance_name, params, placement, 1_000)
589 .await
590 .map_err(|error| command_error(addr, error))?;
591 if json {
592 print_json(&ActionOutput {
593 action: "submitted",
594 message: format!("submitted cluster job '{}'", status.name),
595 job: JobView::from(&status),
596 })?;
597 } else {
598 println!(
599 "submitted cluster job '{}' from factory '{}' to node '{}' — state {}, placement generation {}",
600 status.name,
601 args.factory,
602 empty_dash(&status.placement_node_id),
603 status.state,
604 status.placement_generation
605 );
606 }
607 Ok(())
608}
609
610fn placement_spec(role: Option<&str>, node: Option<&str>) -> PlacementSpec {
611 PlacementSpec {
612 role_constraint: role.unwrap_or_default().to_owned(),
613 strategy: if node.is_some() {
614 PlacementStrategy::Pinned as i32
615 } else {
616 PlacementStrategy::LeastJobs as i32
617 },
618 pinned_node_id: node.unwrap_or_default().to_owned(),
619 }
620}
621
622async fn start(
623 client: &DcpClient,
624 addr: SocketAddr,
625 factory: &str,
626 name: Option<&str>,
627 raw_params: &[String],
628 json: bool,
629) -> Result<(), CliError> {
630 let params = parse_params(raw_params)?;
631 let instance_name = name.unwrap_or(factory);
632 let status = client
633 .start_job(factory, instance_name, params)
634 .await
635 .map_err(|error| command_error(addr, error))?;
636 if json {
637 print_json(&ActionOutput {
638 action: "started",
639 message: format!("started '{}'", status.name),
640 job: JobView::from(&status),
641 })?;
642 } else {
643 println!(
644 "started '{}' from factory '{}' — state {}, generation {}",
645 status.name, factory, status.state, status.generation
646 );
647 }
648 Ok(())
649}
650
651async fn drain(
652 client: &DcpClient,
653 addr: SocketAddr,
654 job: &str,
655 cluster: bool,
656 json: bool,
657) -> Result<(), CliError> {
658 let initial = if cluster {
659 client.drain_cluster_job(job, 1_000).await
660 } else {
661 client.drain_job(job).await
662 }
663 .map_err(|error| command_error(addr, error))?;
664 let status = wait_for_drain(client, job, cluster, initial)
665 .await
666 .map_err(|error| command_error(addr, error))?;
667 if json {
668 print_json(&ActionOutput {
669 action: "drained",
670 message: format!("drained '{}'", status.name),
671 job: JobView::from(&status),
672 })?;
673 } else if status.state == "Drained" {
674 println!("drained '{}' — in-flight streams completed", status.name);
675 } else {
676 println!(
677 "drain requested for '{}' — current state {}; run `datum status {}` to follow it",
678 status.name, status.state, status.name
679 );
680 }
681 Ok(())
682}
683
684async fn wait_for_drain(
685 client: &DcpClient,
686 job: &str,
687 cluster: bool,
688 mut status: WireJobStatus,
689) -> Result<WireJobStatus, DcpError> {
690 let wait = status
691 .drain_remaining_ms
692 .map(|ms| Duration::from_millis(ms.saturating_add(1_000)))
693 .unwrap_or(Duration::from_secs(31))
694 .min(Duration::from_secs(60));
695 let started = tokio::time::Instant::now();
696 while !is_drain_terminal(&status) && started.elapsed() < wait {
697 tokio::time::sleep(DRAIN_POLL).await;
698 status = if cluster {
699 client.cluster_job_status(job, 1_000).await?
700 } else {
701 client.job_status(job).await?
702 };
703 }
704 Ok(status)
705}
706
707fn is_drain_terminal(status: &WireJobStatus) -> bool {
708 matches!(
709 status.state.as_str(),
710 "Completed" | "Drained" | "Stopped" | "Failed"
711 )
712}
713
714async fn stop(
715 client: &DcpClient,
716 addr: SocketAddr,
717 job: &str,
718 cluster: bool,
719 json: bool,
720) -> Result<(), CliError> {
721 let status = if cluster {
722 client.stop_cluster_job(job, 1_000).await
723 } else {
724 client.stop_job(job).await
725 }
726 .map_err(|error| command_error(addr, error))?;
727 if json {
728 print_json(&ActionOutput {
729 action: "stopped",
730 message: format!("stopped '{}'", status.name),
731 job: JobView::from(&status),
732 })?;
733 } else {
734 println!("stopped '{}' — state {}", status.name, status.state);
735 }
736 Ok(())
737}
738
739async fn restart(
740 client: &DcpClient,
741 addr: SocketAddr,
742 job: &str,
743 json: bool,
744) -> Result<(), CliError> {
745 let status = client
746 .restart_job(job)
747 .await
748 .map_err(|error| command_error(addr, error))?;
749 if json {
750 print_json(&ActionOutput {
751 action: "restarted",
752 message: format!("restarted '{}'", status.name),
753 job: JobView::from(&status),
754 })?;
755 } else {
756 println!(
757 "restarted '{}' — state {}, generation {}",
758 status.name, status.state, status.generation
759 );
760 }
761 Ok(())
762}
763
764async fn events(
765 client: &DcpClient,
766 addr: SocketAddr,
767 follow: bool,
768 json: bool,
769) -> Result<(), CliError> {
770 let mut subscription = client
771 .subscribe_events()
772 .await
773 .map_err(|error| command_error(addr, error))?;
774 if follow {
775 while let Some(event) = subscription.recv().await {
776 print_event(&event, json)?;
777 }
778 if let Some(reason) = subscription.disconnect_reason() {
779 return Err(CliError::Connect {
780 addr,
781 message: reason,
782 });
783 }
784 return Ok(());
785 }
786
787 let events = collect_events(&mut subscription).await;
788 if json {
789 print_json(&EventsOutput {
790 events: events.iter().map(EventView::from).collect(),
791 })?;
792 } else if events.is_empty() {
793 println!("no events observed — run `datum events --follow` to keep waiting");
794 } else {
795 print!("{}", render_events_table(&events));
796 }
797 Ok(())
798}
799
800async fn collect_events(subscription: &mut EventSubscription) -> Vec<Event> {
801 let mut events = Vec::new();
802 loop {
803 let timeout = if events.is_empty() {
804 EVENT_INITIAL_IDLE
805 } else {
806 EVENT_IDLE
807 };
808 match tokio::time::timeout(timeout, subscription.recv()).await {
809 Ok(Some(event)) => events.push(event),
810 Ok(None) | Err(_) => break,
811 }
812 }
813 events
814}
815
816fn print_event(event: &Event, json: bool) -> Result<(), CliError> {
817 if json {
818 print_json_line(&EventLine {
819 event: EventView::from(event),
820 })?;
821 } else {
822 println!(
823 "{} {} '{}' generation {}",
824 event.sequence, event.kind, event.name, event.generation
825 );
826 io::stdout().flush().map_err(io_usage_error)?;
827 }
828 Ok(())
829}
830
831async fn metrics(
832 client: &DcpClient,
833 addr: SocketAddr,
834 job: &str,
835 follow: bool,
836 json: bool,
837) -> Result<(), CliError> {
838 let mut subscription = client
839 .subscribe_metrics(0)
840 .await
841 .map_err(|error| command_error(addr, error))?;
842 if follow {
843 while let Some(sample) = subscription.recv().await {
844 let filtered = filter_metric_sample(job, &sample);
845 if filtered.streams.is_empty() {
846 continue;
847 }
848 print_metric_sample_follow(job, &filtered, json)?;
849 }
850 if let Some(reason) = subscription.disconnect_reason() {
851 return Err(CliError::Connect {
852 addr,
853 message: reason,
854 });
855 }
856 return Ok(());
857 }
858
859 let sample = wait_for_metric_sample(job, &mut subscription).await;
860 print_metric_sample(job, &sample, json)
861}
862
863async fn wait_for_metric_sample(
864 job: &str,
865 subscription: &mut MetricSubscription,
866) -> MetricSampleView {
867 let deadline = tokio::time::Instant::now() + METRIC_SNAPSHOT_WAIT;
868 let mut latest = MetricSampleView {
869 timestamp_ms: 0,
870 streams: Vec::new(),
871 };
872 while tokio::time::Instant::now() < deadline {
873 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
874 match tokio::time::timeout(
875 remaining.min(Duration::from_millis(100)),
876 subscription.recv(),
877 )
878 .await
879 {
880 Ok(Some(sample)) => {
881 latest = filter_metric_sample(job, &sample);
882 if !latest.streams.is_empty() {
883 break;
884 }
885 }
886 Ok(None) => break,
887 Err(_) => {}
888 }
889 }
890 latest
891}
892
893fn print_metric_sample(job: &str, sample: &MetricSampleView, json: bool) -> Result<(), CliError> {
894 if json {
895 print_json(&MetricsOutput { job, sample })?;
896 } else if sample.streams.is_empty() {
897 println!("no metrics for '{job}' yet");
898 } else {
899 print!("{}", render_metrics_table(sample));
900 }
901 Ok(())
902}
903
904fn print_metric_sample_follow(
905 job: &str,
906 sample: &MetricSampleView,
907 json: bool,
908) -> Result<(), CliError> {
909 if json {
910 print_json_line(&MetricsOutput { job, sample })?;
911 } else {
912 print!("{}", render_metrics_table(sample));
913 io::stdout().flush().map_err(io_usage_error)?;
914 }
915 Ok(())
916}
917
918fn filter_metric_sample(job: &str, sample: &MetricSample) -> MetricSampleView {
919 let prefix = format!("{job}:");
920 MetricSampleView {
921 timestamp_ms: sample.timestamp_ms,
922 streams: sample
923 .streams
924 .iter()
925 .filter(|metric| metric.name == job || metric.name.starts_with(&prefix))
926 .map(StreamMetricView::from)
927 .collect(),
928 }
929}
930
931fn sort_cluster_jobs(list: &mut ClusterJobList) {
932 list.nodes
933 .sort_by(|left, right| left.node_id.cmp(&right.node_id));
934 for node in &mut list.nodes {
935 node.jobs.sort_by(|left, right| left.name.cmp(&right.name));
936 }
937 list.errors
938 .sort_by(|left, right| left.node_id.cmp(&right.node_id));
939}
940
941fn parse_params(raw_params: &[String]) -> Result<HashMap<String, String>, CliError> {
942 let mut params = HashMap::new();
943 let mut seen = HashSet::new();
944 for raw in raw_params {
945 let Some((key, value)) = raw.split_once('=') else {
946 return Err(CliError::Usage(format!(
947 "invalid --param '{raw}'; use --param key=value."
948 )));
949 };
950 if key.trim().is_empty() {
951 return Err(CliError::Usage(format!(
952 "invalid --param '{raw}'; the key before '=' cannot be empty."
953 )));
954 }
955 if !seen.insert(key.to_owned()) {
956 return Err(CliError::Usage(format!(
957 "duplicate --param key '{key}'; pass each key once."
958 )));
959 }
960 params.insert(key.to_owned(), value.to_owned());
961 }
962 Ok(params)
963}
964
965#[derive(Serialize)]
966struct JobsOutput {
967 jobs: Vec<JobView>,
968}
969
970#[derive(Serialize)]
971struct ClusterJobsOutput {
972 partial: bool,
973 nodes: Vec<ClusterJobNodeView>,
974 errors: Vec<ClusterNodeErrorView>,
975}
976
977impl From<&ClusterJobList> for ClusterJobsOutput {
978 fn from(list: &ClusterJobList) -> Self {
979 Self {
980 partial: list.partial,
981 nodes: list.nodes.iter().map(ClusterJobNodeView::from).collect(),
982 errors: list.errors.iter().map(ClusterNodeErrorView::from).collect(),
983 }
984 }
985}
986
987#[derive(Serialize)]
988struct ClusterNodesOutput {
989 partial: bool,
990 coordinator_node_id: Option<String>,
991 nodes: Vec<ClusterNodeView>,
992 errors: Vec<ClusterNodeErrorView>,
993}
994
995impl ClusterNodesOutput {
996 fn new(list: &ClusterNodeList, coordinator_node_id: Option<String>) -> Self {
997 Self {
998 partial: list.partial,
999 coordinator_node_id,
1000 nodes: list.nodes.iter().map(ClusterNodeView::from).collect(),
1001 errors: list.errors.iter().map(ClusterNodeErrorView::from).collect(),
1002 }
1003 }
1004}
1005
1006#[derive(Serialize)]
1007struct JobOutput {
1008 job: JobView,
1009}
1010
1011#[derive(Serialize)]
1012struct ActionOutput {
1013 action: &'static str,
1014 message: String,
1015 job: JobView,
1016}
1017
1018#[derive(Serialize)]
1019struct EventsOutput {
1020 events: Vec<EventView>,
1021}
1022
1023#[derive(Serialize)]
1024struct EventLine {
1025 event: EventView,
1026}
1027
1028#[derive(Serialize)]
1029struct MetricsOutput<'a> {
1030 job: &'a str,
1031 sample: &'a MetricSampleView,
1032}
1033
1034#[derive(Serialize)]
1035struct JobView {
1036 name: String,
1037 job_id: u64,
1038 state: String,
1039 desired_state: String,
1040 generation: u64,
1041 starts_total: u64,
1042 restarts_total: u64,
1043 last_start_at_ms: Option<u64>,
1044 last_exit_at_ms: Option<u64>,
1045 last_exit_reason: String,
1046 backoff_remaining_ms: Option<u64>,
1047 drain_remaining_ms: Option<u64>,
1048 drain_supported: bool,
1049 active_streams: Option<u64>,
1050 cluster_job: bool,
1051 factory_name: String,
1052 placement_strategy: String,
1053 role_constraint: String,
1054 coordinator_node_id: String,
1055 placement_node_id: String,
1056 placement_generation: u64,
1057 placement_history: Vec<PlacementHistoryView>,
1058}
1059
1060impl From<&WireJobStatus> for JobView {
1061 fn from(status: &WireJobStatus) -> Self {
1062 Self {
1063 name: status.name.clone(),
1064 job_id: status.job_id,
1065 state: status.state.clone(),
1066 desired_state: status.desired_state.clone(),
1067 generation: status.generation,
1068 starts_total: status.starts_total,
1069 restarts_total: status.restarts_total,
1070 last_start_at_ms: status.last_start_at_ms,
1071 last_exit_at_ms: status.last_exit_at_ms,
1072 last_exit_reason: status.last_exit_reason.clone(),
1073 backoff_remaining_ms: status.backoff_remaining_ms,
1074 drain_remaining_ms: status.drain_remaining_ms,
1075 drain_supported: status.drain_supported,
1076 active_streams: status.active_streams,
1077 cluster_job: status.cluster_job,
1078 factory_name: status.factory_name.clone(),
1079 placement_strategy: placement_strategy_text(status),
1080 role_constraint: status
1081 .placement
1082 .as_ref()
1083 .map(|placement| placement.role_constraint.clone())
1084 .unwrap_or_default(),
1085 coordinator_node_id: status.coordinator_node_id.clone(),
1086 placement_node_id: status.placement_node_id.clone(),
1087 placement_generation: status.placement_generation,
1088 placement_history: status
1089 .placement_history
1090 .iter()
1091 .map(PlacementHistoryView::from)
1092 .collect(),
1093 }
1094 }
1095}
1096
1097#[derive(Serialize)]
1098struct PlacementHistoryView {
1099 generation: u64,
1100 from_node_id: String,
1101 to_node_id: String,
1102 reason: String,
1103 timestamp_ms: u64,
1104}
1105
1106impl From<&datum_agent::dcp::proto::ClusterPlacementHistory> for PlacementHistoryView {
1107 fn from(history: &datum_agent::dcp::proto::ClusterPlacementHistory) -> Self {
1108 Self {
1109 generation: history.generation,
1110 from_node_id: history.from_node_id.clone(),
1111 to_node_id: history.to_node_id.clone(),
1112 reason: history.reason.clone(),
1113 timestamp_ms: history.timestamp_ms,
1114 }
1115 }
1116}
1117
1118fn placement_strategy_text(status: &WireJobStatus) -> String {
1119 let Some(placement) = status.placement.as_ref() else {
1120 return String::new();
1121 };
1122 match PlacementStrategy::try_from(placement.strategy).unwrap_or(PlacementStrategy::LeastJobs) {
1123 PlacementStrategy::LeastJobs => "LeastJobs".to_owned(),
1124 PlacementStrategy::Pinned => format!("Pinned({})", placement.pinned_node_id),
1125 }
1126}
1127
1128#[derive(Serialize)]
1129struct ClusterJobNodeView {
1130 node_id: String,
1131 address: String,
1132 local: bool,
1133 jobs: Vec<JobView>,
1134}
1135
1136impl From<&ClusterJobNode> for ClusterJobNodeView {
1137 fn from(node: &ClusterJobNode) -> Self {
1138 Self {
1139 node_id: node.node_id.clone(),
1140 address: node.address.clone(),
1141 local: node.local,
1142 jobs: node.jobs.iter().map(JobView::from).collect(),
1143 }
1144 }
1145}
1146
1147#[derive(Serialize)]
1148struct ClusterNodeView {
1149 node_id: String,
1150 member_state: String,
1151 address: String,
1152 agent_addr: String,
1153 roles: Vec<String>,
1154 unreachable: bool,
1155 local: bool,
1156 session_state: String,
1157}
1158
1159impl From<&ClusterNodeStatus> for ClusterNodeView {
1160 fn from(node: &ClusterNodeStatus) -> Self {
1161 Self {
1162 node_id: node.node_id.clone(),
1163 member_state: node.member_state.clone(),
1164 address: node.address.clone(),
1165 agent_addr: node.agent_addr.clone(),
1166 roles: node.roles.clone(),
1167 unreachable: node.unreachable,
1168 local: node.local,
1169 session_state: node.session_state.clone(),
1170 }
1171 }
1172}
1173
1174#[derive(Serialize)]
1175struct ClusterNodeErrorView {
1176 node_id: String,
1177 message: String,
1178}
1179
1180impl From<&ClusterNodeError> for ClusterNodeErrorView {
1181 fn from(error: &ClusterNodeError) -> Self {
1182 Self {
1183 node_id: error.node_id.clone(),
1184 message: error.message.clone(),
1185 }
1186 }
1187}
1188
1189#[derive(Serialize)]
1190struct EventView {
1191 sequence: u64,
1192 timestamp_ms: u64,
1193 name: String,
1194 job_id: u64,
1195 generation: u64,
1196 kind: String,
1197 detail: String,
1198}
1199
1200impl From<&Event> for EventView {
1201 fn from(event: &Event) -> Self {
1202 Self {
1203 sequence: event.sequence,
1204 timestamp_ms: event.timestamp_ms,
1205 name: event.name.clone(),
1206 job_id: event.job_id,
1207 generation: event.generation,
1208 kind: event.kind.clone(),
1209 detail: event.detail.clone(),
1210 }
1211 }
1212}
1213
1214#[derive(Serialize)]
1215struct MetricSampleView {
1216 timestamp_ms: u64,
1217 streams: Vec<StreamMetricView>,
1218}
1219
1220#[derive(Serialize)]
1221struct StreamMetricView {
1222 id: u64,
1223 name: String,
1224 elements_through: u64,
1225 restarts: u64,
1226 state: String,
1227 started_at_ms: u64,
1228 state_changed_at_ms: u64,
1229 finished_at_ms: Option<u64>,
1230 uptime_ms: u64,
1231}
1232
1233impl From<&StreamMetric> for StreamMetricView {
1234 fn from(metric: &StreamMetric) -> Self {
1235 Self {
1236 id: metric.id,
1237 name: metric.name.clone(),
1238 elements_through: metric.elements_through,
1239 restarts: metric.restarts,
1240 state: metric.state.clone(),
1241 started_at_ms: metric.started_at_ms,
1242 state_changed_at_ms: metric.state_changed_at_ms,
1243 finished_at_ms: metric.finished_at_ms,
1244 uptime_ms: metric.uptime_ms,
1245 }
1246 }
1247}
1248
1249fn print_json<T: Serialize>(value: &T) -> Result<(), CliError> {
1250 let json = serde_json::to_string_pretty(value).map_err(|error| {
1251 CliError::Usage(format!(
1252 "could not encode JSON output: {error}; retry without --json."
1253 ))
1254 })?;
1255 println!("{json}");
1256 Ok(())
1257}
1258
1259fn print_json_line<T: Serialize>(value: &T) -> Result<(), CliError> {
1260 let json = serde_json::to_string(value).map_err(|error| {
1261 CliError::Usage(format!(
1262 "could not encode JSON output: {error}; retry without --json."
1263 ))
1264 })?;
1265 println!("{json}");
1266 io::stdout().flush().map_err(io_usage_error)?;
1267 Ok(())
1268}
1269
1270fn io_usage_error(error: io::Error) -> CliError {
1271 CliError::Usage(format!(
1272 "could not write CLI output: {error}; check stdout/stderr."
1273 ))
1274}
1275
1276fn render_jobs_table(jobs: &[WireJobStatus]) -> String {
1277 let rows = jobs
1278 .iter()
1279 .map(|job| {
1280 vec![
1281 job.name.clone(),
1282 job.job_id.to_string(),
1283 job.state.clone(),
1284 job.desired_state.clone(),
1285 job.generation.to_string(),
1286 format_number(job.starts_total),
1287 format_number(job.restarts_total),
1288 format_optional_number(job.active_streams),
1289 ]
1290 })
1291 .collect::<Vec<_>>();
1292 render_table(
1293 &[
1294 "JOB", "ID", "STATE", "DESIRED", "GEN", "STARTS", "RESTARTS", "ACTIVE",
1295 ],
1296 &rows,
1297 )
1298}
1299
1300fn render_cluster_jobs_table(list: &ClusterJobList) -> String {
1301 let rows = list
1302 .nodes
1303 .iter()
1304 .flat_map(|node| {
1305 if node.jobs.is_empty() {
1306 vec![vec![
1307 node.node_id.clone(),
1308 node.address.clone(),
1309 "-".to_owned(),
1310 "-".to_owned(),
1311 "-".to_owned(),
1312 "-".to_owned(),
1313 "-".to_owned(),
1314 "-".to_owned(),
1315 "-".to_owned(),
1316 "-".to_owned(),
1317 "-".to_owned(),
1318 "-".to_owned(),
1319 ]]
1320 } else {
1321 node.jobs
1322 .iter()
1323 .map(|job| {
1324 vec![
1325 node.node_id.clone(),
1326 node.address.clone(),
1327 job.name.clone(),
1328 job.job_id.to_string(),
1329 job.state.clone(),
1330 job.desired_state.clone(),
1331 job.generation.to_string(),
1332 format_number(job.starts_total),
1333 format_number(job.restarts_total),
1334 cluster_marker(job),
1335 empty_dash(&job.placement_node_id).to_owned(),
1336 placement_generation_text(job),
1337 ]
1338 })
1339 .collect::<Vec<_>>()
1340 }
1341 })
1342 .collect::<Vec<_>>();
1343 render_table(
1344 &[
1345 "NODE",
1346 "ADDRESS",
1347 "JOB",
1348 "ID",
1349 "STATE",
1350 "DESIRED",
1351 "GEN",
1352 "STARTS",
1353 "RESTARTS",
1354 "CLUSTER",
1355 "PLACED_ON",
1356 "PGEN",
1357 ],
1358 &rows,
1359 )
1360}
1361
1362fn render_nodes_table(nodes: &[ClusterNodeStatus], coordinator_node_id: Option<&str>) -> String {
1363 let rows = nodes
1364 .iter()
1365 .map(|node| {
1366 vec![
1367 node.node_id.clone(),
1368 node.member_state.clone(),
1369 coordinator_marker(&node.node_id, coordinator_node_id).to_owned(),
1370 node.unreachable.to_string(),
1371 node.session_state.clone(),
1372 node.agent_addr.clone(),
1373 node.roles.join(","),
1374 ]
1375 })
1376 .collect::<Vec<_>>();
1377 render_table(
1378 &[
1379 "NODE",
1380 "MEMBER",
1381 "COORD",
1382 "UNREACHABLE",
1383 "SESSION",
1384 "AGENT_ADDR",
1385 "ROLES",
1386 ],
1387 &rows,
1388 )
1389}
1390
1391fn coordinator_marker(node_id: &str, coordinator_node_id: Option<&str>) -> &'static str {
1392 if coordinator_node_id.is_some_and(|coordinator| coordinator == node_id) {
1393 "*"
1394 } else {
1395 "-"
1396 }
1397}
1398
1399fn render_cluster_errors_table(errors: &[ClusterNodeError]) -> String {
1400 let rows = errors
1401 .iter()
1402 .map(|error| vec![error.node_id.clone(), error.message.clone()])
1403 .collect::<Vec<_>>();
1404 render_table(&["NODE", "ERROR"], &rows)
1405}
1406
1407fn cluster_marker(job: &WireJobStatus) -> String {
1408 if job.cluster_job {
1409 "yes".to_owned()
1410 } else {
1411 "-".to_owned()
1412 }
1413}
1414
1415fn placement_generation_text(job: &WireJobStatus) -> String {
1416 if job.cluster_job {
1417 job.placement_generation.to_string()
1418 } else {
1419 "-".to_owned()
1420 }
1421}
1422
1423fn render_status_table(status: &WireJobStatus) -> String {
1424 let rows = vec![
1425 vec!["job".to_owned(), status.name.clone()],
1426 vec!["id".to_owned(), status.job_id.to_string()],
1427 vec!["state".to_owned(), status.state.clone()],
1428 vec!["desired".to_owned(), status.desired_state.clone()],
1429 vec!["generation".to_owned(), status.generation.to_string()],
1430 vec!["starts".to_owned(), format_number(status.starts_total)],
1431 vec!["restarts".to_owned(), format_number(status.restarts_total)],
1432 vec![
1433 "last_start_ms".to_owned(),
1434 format_optional_number(status.last_start_at_ms),
1435 ],
1436 vec![
1437 "last_exit_ms".to_owned(),
1438 format_optional_number(status.last_exit_at_ms),
1439 ],
1440 vec![
1441 "last_exit_reason".to_owned(),
1442 empty_dash(&status.last_exit_reason).to_owned(),
1443 ],
1444 vec![
1445 "backoff_remaining_ms".to_owned(),
1446 format_optional_number(status.backoff_remaining_ms),
1447 ],
1448 vec![
1449 "drain_remaining_ms".to_owned(),
1450 format_optional_number(status.drain_remaining_ms),
1451 ],
1452 vec![
1453 "drain_supported".to_owned(),
1454 status.drain_supported.to_string(),
1455 ],
1456 vec![
1457 "active_streams".to_owned(),
1458 format_optional_number(status.active_streams),
1459 ],
1460 vec!["cluster_job".to_owned(), status.cluster_job.to_string()],
1461 vec![
1462 "factory".to_owned(),
1463 empty_dash(&status.factory_name).to_owned(),
1464 ],
1465 vec![
1466 "placement_strategy".to_owned(),
1467 empty_dash(&placement_strategy_text(status)).to_owned(),
1468 ],
1469 vec![
1470 "role_constraint".to_owned(),
1471 status
1472 .placement
1473 .as_ref()
1474 .map(|placement| empty_dash(&placement.role_constraint).to_owned())
1475 .unwrap_or_else(|| "-".to_owned()),
1476 ],
1477 vec![
1478 "coordinator".to_owned(),
1479 empty_dash(&status.coordinator_node_id).to_owned(),
1480 ],
1481 vec![
1482 "placed_on".to_owned(),
1483 empty_dash(&status.placement_node_id).to_owned(),
1484 ],
1485 vec![
1486 "placement_generation".to_owned(),
1487 placement_generation_text(status),
1488 ],
1489 ];
1490 render_table(&["FIELD", "VALUE"], &rows)
1491}
1492
1493fn render_events_table(events: &[Event]) -> String {
1494 let rows = events
1495 .iter()
1496 .map(|event| {
1497 vec![
1498 event.sequence.to_string(),
1499 event.timestamp_ms.to_string(),
1500 event.name.clone(),
1501 event.job_id.to_string(),
1502 event.generation.to_string(),
1503 event.kind.clone(),
1504 empty_dash(&event.detail).to_owned(),
1505 ]
1506 })
1507 .collect::<Vec<_>>();
1508 render_table(
1509 &["SEQ", "TIME_MS", "JOB", "ID", "GEN", "KIND", "DETAIL"],
1510 &rows,
1511 )
1512}
1513
1514fn render_metrics_table(sample: &MetricSampleView) -> String {
1515 let rows = sample
1516 .streams
1517 .iter()
1518 .map(|metric| {
1519 vec![
1520 sample.timestamp_ms.to_string(),
1521 metric.name.clone(),
1522 metric.id.to_string(),
1523 metric.state.clone(),
1524 format_number(metric.elements_through),
1525 format_number(metric.restarts),
1526 metric.uptime_ms.to_string(),
1527 ]
1528 })
1529 .collect::<Vec<_>>();
1530 render_table(
1531 &[
1532 "TIME_MS",
1533 "STREAM",
1534 "ID",
1535 "STATE",
1536 "ELEMENTS",
1537 "RESTARTS",
1538 "UPTIME_MS",
1539 ],
1540 &rows,
1541 )
1542}
1543
1544#[must_use]
1546pub fn render_table(headers: &[&str], rows: &[Vec<String>]) -> String {
1547 let mut widths = headers
1548 .iter()
1549 .map(|header| header.len())
1550 .collect::<Vec<_>>();
1551 for row in rows {
1552 for (index, cell) in row.iter().enumerate().take(widths.len()) {
1553 widths[index] = widths[index].max(cell.len());
1554 }
1555 }
1556
1557 let mut output = String::new();
1558 push_table_row(
1559 &mut output,
1560 &headers
1561 .iter()
1562 .map(|cell| (*cell).to_owned())
1563 .collect::<Vec<_>>(),
1564 &widths,
1565 );
1566 push_separator(&mut output, &widths);
1567 for row in rows {
1568 push_table_row(&mut output, row, &widths);
1569 }
1570 output
1571}
1572
1573fn push_table_row(output: &mut String, row: &[String], widths: &[usize]) {
1574 for (index, width) in widths.iter().enumerate() {
1575 if index > 0 {
1576 output.push_str(" ");
1577 }
1578 let cell = row.get(index).map(String::as_str).unwrap_or("");
1579 output.push_str(cell);
1580 for _ in cell.len()..*width {
1581 output.push(' ');
1582 }
1583 }
1584 output.push('\n');
1585}
1586
1587fn push_separator(output: &mut String, widths: &[usize]) {
1588 for (index, width) in widths.iter().enumerate() {
1589 if index > 0 {
1590 output.push_str(" ");
1591 }
1592 for _ in 0..*width {
1593 output.push('-');
1594 }
1595 }
1596 output.push('\n');
1597}
1598
1599fn empty_dash(value: &str) -> &str {
1600 if value.is_empty() { "-" } else { value }
1601}
1602
1603fn format_optional_number(value: Option<u64>) -> String {
1604 value.map(format_number).unwrap_or_else(|| "-".to_owned())
1605}
1606
1607fn format_number(value: u64) -> String {
1608 let raw = value.to_string();
1609 let mut output = String::with_capacity(raw.len() + raw.len() / 3);
1610 for (index, ch) in raw.chars().rev().enumerate() {
1611 if index > 0 && index % 3 == 0 {
1612 output.push(',');
1613 }
1614 output.push(ch);
1615 }
1616 output.chars().rev().collect()
1617}
1618
1619#[cfg(test)]
1620mod tests {
1621 use super::*;
1622
1623 #[test]
1624 fn table_renderer_pads_columns_stably() {
1625 let output = render_table(
1626 &["JOB", "STATE", "STARTS"],
1627 &[
1628 vec!["ingest".to_owned(), "Running".to_owned(), "1".to_owned()],
1629 vec![
1630 "daily-rollup".to_owned(),
1631 "Drained".to_owned(),
1632 "12,003".to_owned(),
1633 ],
1634 ],
1635 );
1636
1637 assert_eq!(
1638 output,
1639 "JOB STATE STARTS\n\
1640 ------------ ------- ------\n\
1641 ingest Running 1 \n\
1642 daily-rollup Drained 12,003\n"
1643 );
1644 }
1645
1646 #[test]
1647 fn params_require_key_value_pairs() {
1648 assert!(parse_params(&["threads=4".to_owned()]).is_ok());
1649 assert!(matches!(
1650 parse_params(&["threads".to_owned()]),
1651 Err(CliError::Usage(_))
1652 ));
1653 assert!(matches!(
1654 parse_params(&["threads=4".to_owned(), "threads=8".to_owned()]),
1655 Err(CliError::Usage(_))
1656 ));
1657 }
1658
1659 #[test]
1660 fn numbers_use_operator_grouping() {
1661 assert_eq!(format_number(0), "0");
1662 assert_eq!(format_number(12_003), "12,003");
1663 assert_eq!(format_number(9_876_543_210), "9,876,543,210");
1664 }
1665
1666 #[test]
1667 fn help_and_version_exit_zero() {
1668 let help_err = Cli::try_parse_from(["datum", "--help"]).unwrap_err();
1669 assert_eq!(clap_error_exit_code(&help_err), EXIT_OK);
1670
1671 let subcommand_help_err = Cli::try_parse_from(["datum", "nodes", "--help"]).unwrap_err();
1672 assert_eq!(clap_error_exit_code(&subcommand_help_err), EXIT_OK);
1673
1674 let version_err = Cli::try_parse_from(["datum", "--version"]).unwrap_err();
1675 assert_eq!(clap_error_exit_code(&version_err), EXIT_OK);
1676 }
1677
1678 #[test]
1679 fn unknown_flag_exits_nonzero() {
1680 let error = Cli::try_parse_from(["datum", "--not-a-real-flag"]).unwrap_err();
1681 assert_ne!(clap_error_exit_code(&error), EXIT_OK);
1682
1683 let missing_job = Cli::try_parse_from(["datum", "status"]).unwrap_err();
1684 assert_ne!(clap_error_exit_code(&missing_job), EXIT_OK);
1685 }
1686
1687 fn cluster_node_status(node_id: &str) -> ClusterNodeStatus {
1688 ClusterNodeStatus {
1689 node_id: node_id.to_owned(),
1690 member_state: "Up".to_owned(),
1691 address: "127.0.0.1:9555".to_owned(),
1692 agent_addr: "127.0.0.1:9556".to_owned(),
1693 roles: vec!["worker".to_owned()],
1694 unreachable: false,
1695 local: false,
1696 session_state: "connected".to_owned(),
1697 }
1698 }
1699
1700 #[test]
1701 fn coordinator_marker_flags_only_the_matching_node() {
1702 assert_eq!(coordinator_marker("node-a", Some("node-a")), "*");
1703 assert_eq!(coordinator_marker("node-a", Some("node-b")), "-");
1704 assert_eq!(coordinator_marker("node-a", None), "-");
1705 }
1706
1707 #[test]
1708 fn nodes_table_marks_the_coordinator_row() {
1709 let nodes = vec![cluster_node_status("node-a"), cluster_node_status("node-b")];
1710 let output = render_nodes_table(&nodes, Some("node-b"));
1711
1712 assert!(output.contains("COORD"));
1713 let lines: Vec<&str> = output.lines().collect();
1714 let node_a_row = lines[2];
1715 let node_b_row = lines[3];
1716 assert!(!node_a_row.split_whitespace().any(|token| token == "*"));
1717 assert!(node_b_row.split_whitespace().any(|token| token == "*"));
1718 }
1719
1720 #[test]
1721 fn cluster_nodes_output_json_carries_coordinator_node_id() {
1722 let list = ClusterNodeList {
1723 nodes: vec![cluster_node_status("node-a")],
1724 errors: Vec::new(),
1725 partial: false,
1726 coordinator_node_id: "node-a".to_owned(),
1727 };
1728 let coordinator_node_id = cluster_coordinator_node_id_from(&list, None);
1729 let output = ClusterNodesOutput::new(&list, coordinator_node_id.clone());
1730
1731 let value = serde_json::to_value(&output).expect("serializes");
1732 assert_eq!(value["coordinator_node_id"], "node-a");
1733 assert_eq!(value["nodes"][0]["node_id"], "node-a");
1734
1735 let table = render_nodes_table(&list.nodes, coordinator_node_id.as_deref());
1736 let node_row = table.lines().nth(2).expect("node row");
1737 assert!(node_row.split_whitespace().any(|token| token == "*"));
1738 }
1739
1740 #[test]
1741 fn cluster_coordinator_node_id_falls_back_to_job_metadata() {
1742 let list = ClusterNodeList {
1743 nodes: vec![cluster_node_status("node-a")],
1744 errors: Vec::new(),
1745 partial: false,
1746 coordinator_node_id: String::new(),
1747 };
1748 let jobs = ClusterJobList {
1749 nodes: vec![ClusterJobNode {
1750 node_id: "node-a".to_owned(),
1751 address: "127.0.0.1:9555".to_owned(),
1752 local: true,
1753 jobs: vec![WireJobStatus {
1754 coordinator_node_id: "node-a".to_owned(),
1755 ..WireJobStatus::default()
1756 }],
1757 }],
1758 errors: Vec::new(),
1759 partial: false,
1760 };
1761
1762 assert!(list.coordinator_node_id.is_empty());
1763 assert_eq!(
1764 cluster_coordinator_node_id_from(&list, Some(&jobs)).as_deref(),
1765 Some("node-a")
1766 );
1767 }
1768}