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