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 return Ok(());
732 }
733
734 let events = collect_events(&mut subscription).await;
735 if json {
736 print_json(&EventsOutput {
737 events: events.iter().map(EventView::from).collect(),
738 })?;
739 } else if events.is_empty() {
740 println!("no events observed — run `datum events --follow` to keep waiting");
741 } else {
742 print!("{}", render_events_table(&events));
743 }
744 Ok(())
745}
746
747async fn collect_events(subscription: &mut EventSubscription) -> Vec<Event> {
748 let mut events = Vec::new();
749 loop {
750 let timeout = if events.is_empty() {
751 EVENT_INITIAL_IDLE
752 } else {
753 EVENT_IDLE
754 };
755 match tokio::time::timeout(timeout, subscription.recv()).await {
756 Ok(Some(event)) => events.push(event),
757 Ok(None) | Err(_) => break,
758 }
759 }
760 events
761}
762
763fn print_event(event: &Event, json: bool) -> Result<(), CliError> {
764 if json {
765 print_json_line(&EventLine {
766 event: EventView::from(event),
767 })?;
768 } else {
769 println!(
770 "{} {} '{}' generation {}",
771 event.sequence, event.kind, event.name, event.generation
772 );
773 io::stdout().flush().map_err(io_usage_error)?;
774 }
775 Ok(())
776}
777
778async fn metrics(
779 client: &DcpClient,
780 addr: SocketAddr,
781 job: &str,
782 follow: bool,
783 json: bool,
784) -> Result<(), CliError> {
785 let mut subscription = client
786 .subscribe_metrics(0)
787 .await
788 .map_err(|error| command_error(addr, error))?;
789 if follow {
790 while let Some(sample) = subscription.recv().await {
791 let filtered = filter_metric_sample(job, &sample);
792 if filtered.streams.is_empty() {
793 continue;
794 }
795 print_metric_sample_follow(job, &filtered, json)?;
796 }
797 return Ok(());
798 }
799
800 let sample = wait_for_metric_sample(job, &mut subscription).await;
801 print_metric_sample(job, &sample, json)
802}
803
804async fn wait_for_metric_sample(
805 job: &str,
806 subscription: &mut MetricSubscription,
807) -> MetricSampleView {
808 let deadline = tokio::time::Instant::now() + METRIC_SNAPSHOT_WAIT;
809 let mut latest = MetricSampleView {
810 timestamp_ms: 0,
811 streams: Vec::new(),
812 };
813 while tokio::time::Instant::now() < deadline {
814 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
815 match tokio::time::timeout(
816 remaining.min(Duration::from_millis(100)),
817 subscription.recv(),
818 )
819 .await
820 {
821 Ok(Some(sample)) => {
822 latest = filter_metric_sample(job, &sample);
823 if !latest.streams.is_empty() {
824 break;
825 }
826 }
827 Ok(None) => break,
828 Err(_) => {}
829 }
830 }
831 latest
832}
833
834fn print_metric_sample(job: &str, sample: &MetricSampleView, json: bool) -> Result<(), CliError> {
835 if json {
836 print_json(&MetricsOutput { job, sample })?;
837 } else if sample.streams.is_empty() {
838 println!("no metrics for '{job}' yet");
839 } else {
840 print!("{}", render_metrics_table(sample));
841 }
842 Ok(())
843}
844
845fn print_metric_sample_follow(
846 job: &str,
847 sample: &MetricSampleView,
848 json: bool,
849) -> Result<(), CliError> {
850 if json {
851 print_json_line(&MetricsOutput { job, sample })?;
852 } else {
853 print!("{}", render_metrics_table(sample));
854 io::stdout().flush().map_err(io_usage_error)?;
855 }
856 Ok(())
857}
858
859fn filter_metric_sample(job: &str, sample: &MetricSample) -> MetricSampleView {
860 let prefix = format!("{job}:");
861 MetricSampleView {
862 timestamp_ms: sample.timestamp_ms,
863 streams: sample
864 .streams
865 .iter()
866 .filter(|metric| metric.name == job || metric.name.starts_with(&prefix))
867 .map(StreamMetricView::from)
868 .collect(),
869 }
870}
871
872fn sort_cluster_jobs(list: &mut ClusterJobList) {
873 list.nodes
874 .sort_by(|left, right| left.node_id.cmp(&right.node_id));
875 for node in &mut list.nodes {
876 node.jobs.sort_by(|left, right| left.name.cmp(&right.name));
877 }
878 list.errors
879 .sort_by(|left, right| left.node_id.cmp(&right.node_id));
880}
881
882fn parse_params(raw_params: &[String]) -> Result<HashMap<String, String>, CliError> {
883 let mut params = HashMap::new();
884 let mut seen = HashSet::new();
885 for raw in raw_params {
886 let Some((key, value)) = raw.split_once('=') else {
887 return Err(CliError::Usage(format!(
888 "invalid --param '{raw}'; use --param key=value."
889 )));
890 };
891 if key.trim().is_empty() {
892 return Err(CliError::Usage(format!(
893 "invalid --param '{raw}'; the key before '=' cannot be empty."
894 )));
895 }
896 if !seen.insert(key.to_owned()) {
897 return Err(CliError::Usage(format!(
898 "duplicate --param key '{key}'; pass each key once."
899 )));
900 }
901 params.insert(key.to_owned(), value.to_owned());
902 }
903 Ok(params)
904}
905
906#[derive(Serialize)]
907struct JobsOutput {
908 jobs: Vec<JobView>,
909}
910
911#[derive(Serialize)]
912struct ClusterJobsOutput {
913 partial: bool,
914 nodes: Vec<ClusterJobNodeView>,
915 errors: Vec<ClusterNodeErrorView>,
916}
917
918impl From<&ClusterJobList> for ClusterJobsOutput {
919 fn from(list: &ClusterJobList) -> Self {
920 Self {
921 partial: list.partial,
922 nodes: list.nodes.iter().map(ClusterJobNodeView::from).collect(),
923 errors: list.errors.iter().map(ClusterNodeErrorView::from).collect(),
924 }
925 }
926}
927
928#[derive(Serialize)]
929struct ClusterNodesOutput {
930 partial: bool,
931 nodes: Vec<ClusterNodeView>,
932 errors: Vec<ClusterNodeErrorView>,
933}
934
935impl From<&ClusterNodeList> for ClusterNodesOutput {
936 fn from(list: &ClusterNodeList) -> Self {
937 Self {
938 partial: list.partial,
939 nodes: list.nodes.iter().map(ClusterNodeView::from).collect(),
940 errors: list.errors.iter().map(ClusterNodeErrorView::from).collect(),
941 }
942 }
943}
944
945#[derive(Serialize)]
946struct JobOutput {
947 job: JobView,
948}
949
950#[derive(Serialize)]
951struct ActionOutput {
952 action: &'static str,
953 message: String,
954 job: JobView,
955}
956
957#[derive(Serialize)]
958struct EventsOutput {
959 events: Vec<EventView>,
960}
961
962#[derive(Serialize)]
963struct EventLine {
964 event: EventView,
965}
966
967#[derive(Serialize)]
968struct MetricsOutput<'a> {
969 job: &'a str,
970 sample: &'a MetricSampleView,
971}
972
973#[derive(Serialize)]
974struct JobView {
975 name: String,
976 job_id: u64,
977 state: String,
978 desired_state: String,
979 generation: u64,
980 starts_total: u64,
981 restarts_total: u64,
982 last_start_at_ms: Option<u64>,
983 last_exit_at_ms: Option<u64>,
984 last_exit_reason: String,
985 backoff_remaining_ms: Option<u64>,
986 drain_remaining_ms: Option<u64>,
987 drain_supported: bool,
988 active_streams: Option<u64>,
989 cluster_job: bool,
990 factory_name: String,
991 placement_strategy: String,
992 role_constraint: String,
993 coordinator_node_id: String,
994 placement_node_id: String,
995 placement_generation: u64,
996 placement_history: Vec<PlacementHistoryView>,
997}
998
999impl From<&WireJobStatus> for JobView {
1000 fn from(status: &WireJobStatus) -> Self {
1001 Self {
1002 name: status.name.clone(),
1003 job_id: status.job_id,
1004 state: status.state.clone(),
1005 desired_state: status.desired_state.clone(),
1006 generation: status.generation,
1007 starts_total: status.starts_total,
1008 restarts_total: status.restarts_total,
1009 last_start_at_ms: status.last_start_at_ms,
1010 last_exit_at_ms: status.last_exit_at_ms,
1011 last_exit_reason: status.last_exit_reason.clone(),
1012 backoff_remaining_ms: status.backoff_remaining_ms,
1013 drain_remaining_ms: status.drain_remaining_ms,
1014 drain_supported: status.drain_supported,
1015 active_streams: status.active_streams,
1016 cluster_job: status.cluster_job,
1017 factory_name: status.factory_name.clone(),
1018 placement_strategy: placement_strategy_text(status),
1019 role_constraint: status
1020 .placement
1021 .as_ref()
1022 .map(|placement| placement.role_constraint.clone())
1023 .unwrap_or_default(),
1024 coordinator_node_id: status.coordinator_node_id.clone(),
1025 placement_node_id: status.placement_node_id.clone(),
1026 placement_generation: status.placement_generation,
1027 placement_history: status
1028 .placement_history
1029 .iter()
1030 .map(PlacementHistoryView::from)
1031 .collect(),
1032 }
1033 }
1034}
1035
1036#[derive(Serialize)]
1037struct PlacementHistoryView {
1038 generation: u64,
1039 from_node_id: String,
1040 to_node_id: String,
1041 reason: String,
1042 timestamp_ms: u64,
1043}
1044
1045impl From<&datum_agent::dcp::proto::ClusterPlacementHistory> for PlacementHistoryView {
1046 fn from(history: &datum_agent::dcp::proto::ClusterPlacementHistory) -> Self {
1047 Self {
1048 generation: history.generation,
1049 from_node_id: history.from_node_id.clone(),
1050 to_node_id: history.to_node_id.clone(),
1051 reason: history.reason.clone(),
1052 timestamp_ms: history.timestamp_ms,
1053 }
1054 }
1055}
1056
1057fn placement_strategy_text(status: &WireJobStatus) -> String {
1058 let Some(placement) = status.placement.as_ref() else {
1059 return String::new();
1060 };
1061 match PlacementStrategy::try_from(placement.strategy).unwrap_or(PlacementStrategy::LeastJobs) {
1062 PlacementStrategy::LeastJobs => "LeastJobs".to_owned(),
1063 PlacementStrategy::Pinned => format!("Pinned({})", placement.pinned_node_id),
1064 }
1065}
1066
1067#[derive(Serialize)]
1068struct ClusterJobNodeView {
1069 node_id: String,
1070 address: String,
1071 local: bool,
1072 jobs: Vec<JobView>,
1073}
1074
1075impl From<&ClusterJobNode> for ClusterJobNodeView {
1076 fn from(node: &ClusterJobNode) -> Self {
1077 Self {
1078 node_id: node.node_id.clone(),
1079 address: node.address.clone(),
1080 local: node.local,
1081 jobs: node.jobs.iter().map(JobView::from).collect(),
1082 }
1083 }
1084}
1085
1086#[derive(Serialize)]
1087struct ClusterNodeView {
1088 node_id: String,
1089 member_state: String,
1090 address: String,
1091 agent_addr: String,
1092 roles: Vec<String>,
1093 unreachable: bool,
1094 local: bool,
1095 session_state: String,
1096}
1097
1098impl From<&ClusterNodeStatus> for ClusterNodeView {
1099 fn from(node: &ClusterNodeStatus) -> Self {
1100 Self {
1101 node_id: node.node_id.clone(),
1102 member_state: node.member_state.clone(),
1103 address: node.address.clone(),
1104 agent_addr: node.agent_addr.clone(),
1105 roles: node.roles.clone(),
1106 unreachable: node.unreachable,
1107 local: node.local,
1108 session_state: node.session_state.clone(),
1109 }
1110 }
1111}
1112
1113#[derive(Serialize)]
1114struct ClusterNodeErrorView {
1115 node_id: String,
1116 message: String,
1117}
1118
1119impl From<&ClusterNodeError> for ClusterNodeErrorView {
1120 fn from(error: &ClusterNodeError) -> Self {
1121 Self {
1122 node_id: error.node_id.clone(),
1123 message: error.message.clone(),
1124 }
1125 }
1126}
1127
1128#[derive(Serialize)]
1129struct EventView {
1130 sequence: u64,
1131 timestamp_ms: u64,
1132 name: String,
1133 job_id: u64,
1134 generation: u64,
1135 kind: String,
1136 detail: String,
1137}
1138
1139impl From<&Event> for EventView {
1140 fn from(event: &Event) -> Self {
1141 Self {
1142 sequence: event.sequence,
1143 timestamp_ms: event.timestamp_ms,
1144 name: event.name.clone(),
1145 job_id: event.job_id,
1146 generation: event.generation,
1147 kind: event.kind.clone(),
1148 detail: event.detail.clone(),
1149 }
1150 }
1151}
1152
1153#[derive(Serialize)]
1154struct MetricSampleView {
1155 timestamp_ms: u64,
1156 streams: Vec<StreamMetricView>,
1157}
1158
1159#[derive(Serialize)]
1160struct StreamMetricView {
1161 id: u64,
1162 name: String,
1163 elements_through: u64,
1164 restarts: u64,
1165 state: String,
1166 started_at_ms: u64,
1167 state_changed_at_ms: u64,
1168 finished_at_ms: Option<u64>,
1169 uptime_ms: u64,
1170}
1171
1172impl From<&StreamMetric> for StreamMetricView {
1173 fn from(metric: &StreamMetric) -> Self {
1174 Self {
1175 id: metric.id,
1176 name: metric.name.clone(),
1177 elements_through: metric.elements_through,
1178 restarts: metric.restarts,
1179 state: metric.state.clone(),
1180 started_at_ms: metric.started_at_ms,
1181 state_changed_at_ms: metric.state_changed_at_ms,
1182 finished_at_ms: metric.finished_at_ms,
1183 uptime_ms: metric.uptime_ms,
1184 }
1185 }
1186}
1187
1188fn print_json<T: Serialize>(value: &T) -> Result<(), CliError> {
1189 let json = serde_json::to_string_pretty(value).map_err(|error| {
1190 CliError::Usage(format!(
1191 "could not encode JSON output: {error}; retry without --json."
1192 ))
1193 })?;
1194 println!("{json}");
1195 Ok(())
1196}
1197
1198fn print_json_line<T: Serialize>(value: &T) -> Result<(), CliError> {
1199 let json = serde_json::to_string(value).map_err(|error| {
1200 CliError::Usage(format!(
1201 "could not encode JSON output: {error}; retry without --json."
1202 ))
1203 })?;
1204 println!("{json}");
1205 io::stdout().flush().map_err(io_usage_error)?;
1206 Ok(())
1207}
1208
1209fn io_usage_error(error: io::Error) -> CliError {
1210 CliError::Usage(format!(
1211 "could not write CLI output: {error}; check stdout/stderr."
1212 ))
1213}
1214
1215fn render_jobs_table(jobs: &[WireJobStatus]) -> String {
1216 let rows = jobs
1217 .iter()
1218 .map(|job| {
1219 vec![
1220 job.name.clone(),
1221 job.job_id.to_string(),
1222 job.state.clone(),
1223 job.desired_state.clone(),
1224 job.generation.to_string(),
1225 format_number(job.starts_total),
1226 format_number(job.restarts_total),
1227 format_optional_number(job.active_streams),
1228 ]
1229 })
1230 .collect::<Vec<_>>();
1231 render_table(
1232 &[
1233 "JOB", "ID", "STATE", "DESIRED", "GEN", "STARTS", "RESTARTS", "ACTIVE",
1234 ],
1235 &rows,
1236 )
1237}
1238
1239fn render_cluster_jobs_table(list: &ClusterJobList) -> String {
1240 let rows = list
1241 .nodes
1242 .iter()
1243 .flat_map(|node| {
1244 if node.jobs.is_empty() {
1245 vec![vec![
1246 node.node_id.clone(),
1247 node.address.clone(),
1248 "-".to_owned(),
1249 "-".to_owned(),
1250 "-".to_owned(),
1251 "-".to_owned(),
1252 "-".to_owned(),
1253 "-".to_owned(),
1254 "-".to_owned(),
1255 "-".to_owned(),
1256 "-".to_owned(),
1257 "-".to_owned(),
1258 ]]
1259 } else {
1260 node.jobs
1261 .iter()
1262 .map(|job| {
1263 vec![
1264 node.node_id.clone(),
1265 node.address.clone(),
1266 job.name.clone(),
1267 job.job_id.to_string(),
1268 job.state.clone(),
1269 job.desired_state.clone(),
1270 job.generation.to_string(),
1271 format_number(job.starts_total),
1272 format_number(job.restarts_total),
1273 cluster_marker(job),
1274 empty_dash(&job.placement_node_id).to_owned(),
1275 placement_generation_text(job),
1276 ]
1277 })
1278 .collect::<Vec<_>>()
1279 }
1280 })
1281 .collect::<Vec<_>>();
1282 render_table(
1283 &[
1284 "NODE",
1285 "ADDRESS",
1286 "JOB",
1287 "ID",
1288 "STATE",
1289 "DESIRED",
1290 "GEN",
1291 "STARTS",
1292 "RESTARTS",
1293 "CLUSTER",
1294 "PLACED_ON",
1295 "PGEN",
1296 ],
1297 &rows,
1298 )
1299}
1300
1301fn render_nodes_table(nodes: &[ClusterNodeStatus]) -> String {
1302 let rows = nodes
1303 .iter()
1304 .map(|node| {
1305 vec![
1306 node.node_id.clone(),
1307 node.member_state.clone(),
1308 node.unreachable.to_string(),
1309 node.session_state.clone(),
1310 node.agent_addr.clone(),
1311 node.roles.join(","),
1312 ]
1313 })
1314 .collect::<Vec<_>>();
1315 render_table(
1316 &[
1317 "NODE",
1318 "MEMBER",
1319 "UNREACHABLE",
1320 "SESSION",
1321 "AGENT_ADDR",
1322 "ROLES",
1323 ],
1324 &rows,
1325 )
1326}
1327
1328fn render_cluster_errors_table(errors: &[ClusterNodeError]) -> String {
1329 let rows = errors
1330 .iter()
1331 .map(|error| vec![error.node_id.clone(), error.message.clone()])
1332 .collect::<Vec<_>>();
1333 render_table(&["NODE", "ERROR"], &rows)
1334}
1335
1336fn cluster_marker(job: &WireJobStatus) -> String {
1337 if job.cluster_job {
1338 "yes".to_owned()
1339 } else {
1340 "-".to_owned()
1341 }
1342}
1343
1344fn placement_generation_text(job: &WireJobStatus) -> String {
1345 if job.cluster_job {
1346 job.placement_generation.to_string()
1347 } else {
1348 "-".to_owned()
1349 }
1350}
1351
1352fn render_status_table(status: &WireJobStatus) -> String {
1353 let rows = vec![
1354 vec!["job".to_owned(), status.name.clone()],
1355 vec!["id".to_owned(), status.job_id.to_string()],
1356 vec!["state".to_owned(), status.state.clone()],
1357 vec!["desired".to_owned(), status.desired_state.clone()],
1358 vec!["generation".to_owned(), status.generation.to_string()],
1359 vec!["starts".to_owned(), format_number(status.starts_total)],
1360 vec!["restarts".to_owned(), format_number(status.restarts_total)],
1361 vec![
1362 "last_start_ms".to_owned(),
1363 format_optional_number(status.last_start_at_ms),
1364 ],
1365 vec![
1366 "last_exit_ms".to_owned(),
1367 format_optional_number(status.last_exit_at_ms),
1368 ],
1369 vec![
1370 "last_exit_reason".to_owned(),
1371 empty_dash(&status.last_exit_reason).to_owned(),
1372 ],
1373 vec![
1374 "backoff_remaining_ms".to_owned(),
1375 format_optional_number(status.backoff_remaining_ms),
1376 ],
1377 vec![
1378 "drain_remaining_ms".to_owned(),
1379 format_optional_number(status.drain_remaining_ms),
1380 ],
1381 vec![
1382 "drain_supported".to_owned(),
1383 status.drain_supported.to_string(),
1384 ],
1385 vec![
1386 "active_streams".to_owned(),
1387 format_optional_number(status.active_streams),
1388 ],
1389 vec!["cluster_job".to_owned(), status.cluster_job.to_string()],
1390 vec![
1391 "factory".to_owned(),
1392 empty_dash(&status.factory_name).to_owned(),
1393 ],
1394 vec![
1395 "placement_strategy".to_owned(),
1396 empty_dash(&placement_strategy_text(status)).to_owned(),
1397 ],
1398 vec![
1399 "role_constraint".to_owned(),
1400 status
1401 .placement
1402 .as_ref()
1403 .map(|placement| empty_dash(&placement.role_constraint).to_owned())
1404 .unwrap_or_else(|| "-".to_owned()),
1405 ],
1406 vec![
1407 "coordinator".to_owned(),
1408 empty_dash(&status.coordinator_node_id).to_owned(),
1409 ],
1410 vec![
1411 "placed_on".to_owned(),
1412 empty_dash(&status.placement_node_id).to_owned(),
1413 ],
1414 vec![
1415 "placement_generation".to_owned(),
1416 placement_generation_text(status),
1417 ],
1418 ];
1419 render_table(&["FIELD", "VALUE"], &rows)
1420}
1421
1422fn render_events_table(events: &[Event]) -> String {
1423 let rows = events
1424 .iter()
1425 .map(|event| {
1426 vec![
1427 event.sequence.to_string(),
1428 event.timestamp_ms.to_string(),
1429 event.name.clone(),
1430 event.job_id.to_string(),
1431 event.generation.to_string(),
1432 event.kind.clone(),
1433 empty_dash(&event.detail).to_owned(),
1434 ]
1435 })
1436 .collect::<Vec<_>>();
1437 render_table(
1438 &["SEQ", "TIME_MS", "JOB", "ID", "GEN", "KIND", "DETAIL"],
1439 &rows,
1440 )
1441}
1442
1443fn render_metrics_table(sample: &MetricSampleView) -> String {
1444 let rows = sample
1445 .streams
1446 .iter()
1447 .map(|metric| {
1448 vec![
1449 sample.timestamp_ms.to_string(),
1450 metric.name.clone(),
1451 metric.id.to_string(),
1452 metric.state.clone(),
1453 format_number(metric.elements_through),
1454 format_number(metric.restarts),
1455 metric.uptime_ms.to_string(),
1456 ]
1457 })
1458 .collect::<Vec<_>>();
1459 render_table(
1460 &[
1461 "TIME_MS",
1462 "STREAM",
1463 "ID",
1464 "STATE",
1465 "ELEMENTS",
1466 "RESTARTS",
1467 "UPTIME_MS",
1468 ],
1469 &rows,
1470 )
1471}
1472
1473#[must_use]
1475pub fn render_table(headers: &[&str], rows: &[Vec<String>]) -> String {
1476 let mut widths = headers
1477 .iter()
1478 .map(|header| header.len())
1479 .collect::<Vec<_>>();
1480 for row in rows {
1481 for (index, cell) in row.iter().enumerate().take(widths.len()) {
1482 widths[index] = widths[index].max(cell.len());
1483 }
1484 }
1485
1486 let mut output = String::new();
1487 push_table_row(
1488 &mut output,
1489 &headers
1490 .iter()
1491 .map(|cell| (*cell).to_owned())
1492 .collect::<Vec<_>>(),
1493 &widths,
1494 );
1495 push_separator(&mut output, &widths);
1496 for row in rows {
1497 push_table_row(&mut output, row, &widths);
1498 }
1499 output
1500}
1501
1502fn push_table_row(output: &mut String, row: &[String], widths: &[usize]) {
1503 for (index, width) in widths.iter().enumerate() {
1504 if index > 0 {
1505 output.push_str(" ");
1506 }
1507 let cell = row.get(index).map(String::as_str).unwrap_or("");
1508 output.push_str(cell);
1509 for _ in cell.len()..*width {
1510 output.push(' ');
1511 }
1512 }
1513 output.push('\n');
1514}
1515
1516fn push_separator(output: &mut String, widths: &[usize]) {
1517 for (index, width) in widths.iter().enumerate() {
1518 if index > 0 {
1519 output.push_str(" ");
1520 }
1521 for _ in 0..*width {
1522 output.push('-');
1523 }
1524 }
1525 output.push('\n');
1526}
1527
1528fn empty_dash(value: &str) -> &str {
1529 if value.is_empty() { "-" } else { value }
1530}
1531
1532fn format_optional_number(value: Option<u64>) -> String {
1533 value.map(format_number).unwrap_or_else(|| "-".to_owned())
1534}
1535
1536fn format_number(value: u64) -> String {
1537 let raw = value.to_string();
1538 let mut output = String::with_capacity(raw.len() + raw.len() / 3);
1539 for (index, ch) in raw.chars().rev().enumerate() {
1540 if index > 0 && index % 3 == 0 {
1541 output.push(',');
1542 }
1543 output.push(ch);
1544 }
1545 output.chars().rev().collect()
1546}
1547
1548#[cfg(test)]
1549mod tests {
1550 use super::*;
1551
1552 #[test]
1553 fn table_renderer_pads_columns_stably() {
1554 let output = render_table(
1555 &["JOB", "STATE", "STARTS"],
1556 &[
1557 vec!["ingest".to_owned(), "Running".to_owned(), "1".to_owned()],
1558 vec![
1559 "daily-rollup".to_owned(),
1560 "Drained".to_owned(),
1561 "12,003".to_owned(),
1562 ],
1563 ],
1564 );
1565
1566 assert_eq!(
1567 output,
1568 "JOB STATE STARTS\n\
1569 ------------ ------- ------\n\
1570 ingest Running 1 \n\
1571 daily-rollup Drained 12,003\n"
1572 );
1573 }
1574
1575 #[test]
1576 fn params_require_key_value_pairs() {
1577 assert!(parse_params(&["threads=4".to_owned()]).is_ok());
1578 assert!(matches!(
1579 parse_params(&["threads".to_owned()]),
1580 Err(CliError::Usage(_))
1581 ));
1582 assert!(matches!(
1583 parse_params(&["threads=4".to_owned(), "threads=8".to_owned()]),
1584 Err(CliError::Usage(_))
1585 ));
1586 }
1587
1588 #[test]
1589 fn numbers_use_operator_grouping() {
1590 assert_eq!(format_number(0), "0");
1591 assert_eq!(format_number(12_003), "12,003");
1592 assert_eq!(format_number(9_876_543_210), "9,876,543,210");
1593 }
1594}