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::{Event, JobStatus as WireJobStatus, MetricSample, StreamMetric},
17};
18use datum_net::quic::{
19 crypto::rustls::QuicClientConfig,
20 quinn,
21 rustls::{
22 ClientConfig as RustlsClientConfig, RootCertStore,
23 pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer},
24 },
25};
26use serde::Serialize;
27
28const DEFAULT_ADDR: &str = "127.0.0.1:9555";
29const EXIT_OK: i32 = 0;
30const EXIT_REMOTE: i32 = 1;
31const EXIT_USAGE_OR_CONNECT: i32 = 2;
32const EVENT_INITIAL_IDLE: Duration = Duration::from_millis(500);
33const EVENT_IDLE: Duration = Duration::from_millis(250);
34const METRIC_SNAPSHOT_WAIT: Duration = Duration::from_secs(2);
35const DRAIN_POLL: Duration = Duration::from_millis(20);
36
37#[derive(Debug, Parser)]
38#[command(
39 name = "datum",
40 version,
41 about = "Manage Datum jobs through the Datum Control Protocol"
42)]
43struct Cli {
44 #[arg(long, global = true, default_value = DEFAULT_ADDR)]
46 addr: SocketAddr,
47
48 #[arg(long, global = true, value_name = "PATH")]
50 tls_ca: Option<PathBuf>,
51
52 #[arg(long, global = true, value_name = "PATH")]
54 tls_cert: Option<PathBuf>,
55
56 #[arg(long, global = true, value_name = "PATH")]
58 tls_key: Option<PathBuf>,
59
60 #[arg(long, global = true)]
62 json: bool,
63
64 #[command(subcommand)]
65 command: Command,
66}
67
68#[derive(Debug, Subcommand)]
69enum Command {
70 Ps,
72 Status { job: String },
74 Start {
76 factory: String,
77 #[arg(long)]
78 name: Option<String>,
79 #[arg(long = "param", value_name = "k=v")]
80 params: Vec<String>,
81 },
82 Drain { job: String },
84 Stop { job: String },
86 Restart { job: String },
88 Events {
90 #[arg(long)]
91 follow: bool,
92 },
93 Metrics {
95 job: String,
96 #[arg(long)]
97 follow: bool,
98 },
99}
100
101#[derive(Debug)]
102enum CliError {
103 Usage(String),
104 Connect { addr: SocketAddr, message: String },
105 Remote { message: String, fix: &'static str },
106}
107
108impl CliError {
109 fn exit_code(&self) -> i32 {
110 match self {
111 Self::Remote { .. } => EXIT_REMOTE,
112 Self::Usage(_) | Self::Connect { .. } => EXIT_USAGE_OR_CONNECT,
113 }
114 }
115
116 fn print(&self) {
117 match self {
118 Self::Usage(message) => {
119 eprintln!("{message}");
120 }
121 Self::Connect { addr, message } => {
122 eprintln!("could not connect to datum-agent at {addr}: {message}");
123 eprintln!(
124 "start datum-agent or pass --addr host:port; loopback TCP defaults to {DEFAULT_ADDR}"
125 );
126 }
127 Self::Remote { message, fix } => {
128 eprintln!("datum-agent rejected the request: {message}");
129 eprintln!("{fix}");
130 }
131 }
132 }
133}
134
135pub async fn main_entry() -> i32 {
137 match Cli::try_parse() {
138 Ok(cli) => match run(cli).await {
139 Ok(()) => EXIT_OK,
140 Err(error) => {
141 error.print();
142 error.exit_code()
143 }
144 },
145 Err(error) => {
146 let _ = error.print();
147 EXIT_USAGE_OR_CONNECT
148 }
149 }
150}
151
152async fn run(cli: Cli) -> Result<(), CliError> {
153 validate_before_connect(&cli.command)?;
154 let client = connect(&cli).await?;
155 match &cli.command {
156 Command::Ps => ps(&client, cli.addr, cli.json).await,
157 Command::Status { job } => status(&client, cli.addr, job, cli.json).await,
158 Command::Start {
159 factory,
160 name,
161 params,
162 } => {
163 start(
164 &client,
165 cli.addr,
166 factory,
167 name.as_deref(),
168 params,
169 cli.json,
170 )
171 .await
172 }
173 Command::Drain { job } => drain(&client, cli.addr, job, cli.json).await,
174 Command::Stop { job } => stop(&client, cli.addr, job, cli.json).await,
175 Command::Restart { job } => restart(&client, cli.addr, job, cli.json).await,
176 Command::Events { follow } => events(&client, cli.addr, *follow, cli.json).await,
177 Command::Metrics { job, follow } => {
178 metrics(&client, cli.addr, job, *follow, cli.json).await
179 }
180 }
181}
182
183fn validate_before_connect(command: &Command) -> Result<(), CliError> {
184 if let Command::Start { params, .. } = command {
185 parse_params(params)?;
186 }
187 Ok(())
188}
189
190async fn connect(cli: &Cli) -> Result<DcpClient, CliError> {
191 let hello = Hello::new(format!("datum-cli-{}", std::process::id()), ClientKind::Cli);
192 match tls_paths(cli)? {
193 Some(paths) => {
194 let config = load_quic_client_config(paths.ca, paths.cert, paths.key)?;
195 DcpClient::connect_quic(cli.addr, "localhost", config, hello)
196 .await
197 .map_err(|error| connect_error(cli.addr, error))
198 }
199 None => DcpClient::connect_tcp(cli.addr, hello)
200 .await
201 .map_err(|error| connect_error(cli.addr, error)),
202 }
203}
204
205struct TlsPaths<'a> {
206 ca: &'a Path,
207 cert: &'a Path,
208 key: &'a Path,
209}
210
211fn tls_paths(cli: &Cli) -> Result<Option<TlsPaths<'_>>, CliError> {
212 match (&cli.tls_ca, &cli.tls_cert, &cli.tls_key) {
213 (None, None, None) => Ok(None),
214 (Some(ca), Some(cert), Some(key)) => Ok(Some(TlsPaths {
215 ca: ca.as_path(),
216 cert: cert.as_path(),
217 key: key.as_path(),
218 })),
219 _ => Err(CliError::Usage(
220 "QUIC+mTLS mode needs --tls-ca, --tls-cert, and --tls-key; pass all three or omit all three for loopback TCP.".to_owned(),
221 )),
222 }
223}
224
225fn load_quic_client_config(
226 ca_path: &Path,
227 cert_path: &Path,
228 key_path: &Path,
229) -> Result<quinn::ClientConfig, CliError> {
230 let mut roots = RootCertStore::empty();
231 for cert in read_certificates(ca_path)? {
232 roots.add(cert).map_err(|error| {
233 CliError::Usage(format!(
234 "could not load --tls-ca {}: {error}; pass a PEM or DER CA certificate.",
235 ca_path.display()
236 ))
237 })?;
238 }
239 let certs = read_certificates(cert_path)?;
240 let key = read_private_key(key_path)?;
241 let rustls = RustlsClientConfig::builder()
242 .with_root_certificates(roots)
243 .with_client_auth_cert(certs, key)
244 .map_err(|error| {
245 CliError::Usage(format!(
246 "could not build QUIC client identity: {error}; check --tls-cert and --tls-key match."
247 ))
248 })?;
249 let quic = QuicClientConfig::try_from(rustls).map_err(|error| {
250 CliError::Usage(format!(
251 "could not build QUIC client config: {error}; check the TLS files."
252 ))
253 })?;
254 Ok(quinn::ClientConfig::new(Arc::new(quic)))
255}
256
257fn read_certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, CliError> {
258 let bytes = fs::read(path).map_err(|error| {
259 CliError::Usage(format!(
260 "could not read {}: {error}; pass a readable PEM or DER certificate file.",
261 path.display()
262 ))
263 })?;
264 let mut cursor = Cursor::new(bytes.as_slice());
265 let certs = rustls_pemfile::certs(&mut cursor)
266 .collect::<Result<Vec<_>, _>>()
267 .map_err(|error| {
268 CliError::Usage(format!(
269 "could not parse {}: {error}; pass a PEM or DER certificate file.",
270 path.display()
271 ))
272 })?;
273 if certs.is_empty() {
274 Ok(vec![CertificateDer::from(bytes)])
275 } else {
276 Ok(certs)
277 }
278}
279
280fn read_private_key(path: &Path) -> Result<PrivateKeyDer<'static>, CliError> {
281 let bytes = fs::read(path).map_err(|error| {
282 CliError::Usage(format!(
283 "could not read {}: {error}; pass a readable PEM or DER PKCS#8 private key.",
284 path.display()
285 ))
286 })?;
287 let mut cursor = Cursor::new(bytes.as_slice());
288 match rustls_pemfile::private_key(&mut cursor).map_err(|error| {
289 CliError::Usage(format!(
290 "could not parse {}: {error}; pass a PEM or DER PKCS#8 private key.",
291 path.display()
292 ))
293 })? {
294 Some(key) => Ok(key),
295 None => Ok(PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(bytes))),
296 }
297}
298
299fn connect_error(addr: SocketAddr, error: DcpError) -> CliError {
300 match error {
301 DcpError::Response { status, message } => remote_error(status, message),
302 other => CliError::Connect {
303 addr,
304 message: other.to_string(),
305 },
306 }
307}
308
309fn command_error(addr: SocketAddr, error: DcpError) -> CliError {
310 match error {
311 DcpError::Response { status, message } => remote_error(status, message),
312 other => CliError::Connect {
313 addr,
314 message: other.to_string(),
315 },
316 }
317}
318
319fn remote_error(status: ResponseStatus, message: String) -> CliError {
320 let fix = match status {
321 ResponseStatus::BadRequest => "Fix the command arguments and try again.",
322 ResponseStatus::Unauthorized => {
323 "Check the DCP token or mTLS client certificate configured for datum-agent."
324 }
325 ResponseStatus::NotFound => "Check the job or factory name with `datum ps`.",
326 ResponseStatus::Conflict => {
327 "Pick a different job name, or stop the existing job before starting it again."
328 }
329 ResponseStatus::ProtocolMismatch => {
330 "Upgrade the CLI and datum-agent to compatible versions."
331 }
332 ResponseStatus::DeadlineExceeded => {
333 "Retry the command or check whether datum-agent is overloaded."
334 }
335 ResponseStatus::Ok | ResponseStatus::Failed => {
336 "Check `datum status <job>` and the datum-agent logs, then retry."
337 }
338 };
339 CliError::Remote { message, fix }
340}
341
342async fn ps(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
343 let mut jobs = client
344 .list_jobs()
345 .await
346 .map_err(|error| command_error(addr, error))?;
347 jobs.sort_by(|left, right| left.name.cmp(&right.name));
348 if json {
349 print_json(&JobsOutput {
350 jobs: jobs.iter().map(JobView::from).collect(),
351 })?;
352 } else if jobs.is_empty() {
353 println!("no jobs");
354 } else {
355 print!("{}", render_jobs_table(&jobs));
356 }
357 Ok(())
358}
359
360async fn status(
361 client: &DcpClient,
362 addr: SocketAddr,
363 job: &str,
364 json: bool,
365) -> Result<(), CliError> {
366 let status = client
367 .job_status(job)
368 .await
369 .map_err(|error| command_error(addr, error))?;
370 if json {
371 print_json(&JobOutput {
372 job: JobView::from(&status),
373 })?;
374 } else {
375 print!("{}", render_status_table(&status));
376 }
377 Ok(())
378}
379
380async fn start(
381 client: &DcpClient,
382 addr: SocketAddr,
383 factory: &str,
384 name: Option<&str>,
385 raw_params: &[String],
386 json: bool,
387) -> Result<(), CliError> {
388 let params = parse_params(raw_params)?;
389 let instance_name = name.unwrap_or(factory);
390 let status = client
391 .start_job(factory, instance_name, params)
392 .await
393 .map_err(|error| command_error(addr, error))?;
394 if json {
395 print_json(&ActionOutput {
396 action: "started",
397 message: format!("started '{}'", status.name),
398 job: JobView::from(&status),
399 })?;
400 } else {
401 println!(
402 "started '{}' from factory '{}' — state {}, generation {}",
403 status.name, factory, status.state, status.generation
404 );
405 }
406 Ok(())
407}
408
409async fn drain(
410 client: &DcpClient,
411 addr: SocketAddr,
412 job: &str,
413 json: bool,
414) -> Result<(), CliError> {
415 let initial = client
416 .drain_job(job)
417 .await
418 .map_err(|error| command_error(addr, error))?;
419 let status = wait_for_drain(client, job, initial)
420 .await
421 .map_err(|error| command_error(addr, error))?;
422 if json {
423 print_json(&ActionOutput {
424 action: "drained",
425 message: format!("drained '{}'", status.name),
426 job: JobView::from(&status),
427 })?;
428 } else if status.state == "Drained" {
429 println!("drained '{}' — in-flight streams completed", status.name);
430 } else {
431 println!(
432 "drain requested for '{}' — current state {}; run `datum status {}` to follow it",
433 status.name, status.state, status.name
434 );
435 }
436 Ok(())
437}
438
439async fn wait_for_drain(
440 client: &DcpClient,
441 job: &str,
442 mut status: WireJobStatus,
443) -> Result<WireJobStatus, DcpError> {
444 let wait = status
445 .drain_remaining_ms
446 .map(|ms| Duration::from_millis(ms.saturating_add(1_000)))
447 .unwrap_or(Duration::from_secs(31))
448 .min(Duration::from_secs(60));
449 let started = tokio::time::Instant::now();
450 while !is_drain_terminal(&status) && started.elapsed() < wait {
451 tokio::time::sleep(DRAIN_POLL).await;
452 status = client.job_status(job).await?;
453 }
454 Ok(status)
455}
456
457fn is_drain_terminal(status: &WireJobStatus) -> bool {
458 matches!(
459 status.state.as_str(),
460 "Completed" | "Drained" | "Stopped" | "Failed"
461 )
462}
463
464async fn stop(client: &DcpClient, addr: SocketAddr, job: &str, json: bool) -> Result<(), CliError> {
465 let status = client
466 .stop_job(job)
467 .await
468 .map_err(|error| command_error(addr, error))?;
469 if json {
470 print_json(&ActionOutput {
471 action: "stopped",
472 message: format!("stopped '{}'", status.name),
473 job: JobView::from(&status),
474 })?;
475 } else {
476 println!("stopped '{}' — state {}", status.name, status.state);
477 }
478 Ok(())
479}
480
481async fn restart(
482 client: &DcpClient,
483 addr: SocketAddr,
484 job: &str,
485 json: bool,
486) -> Result<(), CliError> {
487 let status = client
488 .restart_job(job)
489 .await
490 .map_err(|error| command_error(addr, error))?;
491 if json {
492 print_json(&ActionOutput {
493 action: "restarted",
494 message: format!("restarted '{}'", status.name),
495 job: JobView::from(&status),
496 })?;
497 } else {
498 println!(
499 "restarted '{}' — state {}, generation {}",
500 status.name, status.state, status.generation
501 );
502 }
503 Ok(())
504}
505
506async fn events(
507 client: &DcpClient,
508 addr: SocketAddr,
509 follow: bool,
510 json: bool,
511) -> Result<(), CliError> {
512 let mut subscription = client
513 .subscribe_events()
514 .await
515 .map_err(|error| command_error(addr, error))?;
516 if follow {
517 while let Some(event) = subscription.recv().await {
518 print_event(&event, json)?;
519 }
520 return Ok(());
521 }
522
523 let events = collect_events(&mut subscription).await;
524 if json {
525 print_json(&EventsOutput {
526 events: events.iter().map(EventView::from).collect(),
527 })?;
528 } else if events.is_empty() {
529 println!("no events observed — run `datum events --follow` to keep waiting");
530 } else {
531 print!("{}", render_events_table(&events));
532 }
533 Ok(())
534}
535
536async fn collect_events(subscription: &mut EventSubscription) -> Vec<Event> {
537 let mut events = Vec::new();
538 loop {
539 let timeout = if events.is_empty() {
540 EVENT_INITIAL_IDLE
541 } else {
542 EVENT_IDLE
543 };
544 match tokio::time::timeout(timeout, subscription.recv()).await {
545 Ok(Some(event)) => events.push(event),
546 Ok(None) | Err(_) => break,
547 }
548 }
549 events
550}
551
552fn print_event(event: &Event, json: bool) -> Result<(), CliError> {
553 if json {
554 print_json_line(&EventLine {
555 event: EventView::from(event),
556 })?;
557 } else {
558 println!(
559 "{} {} '{}' generation {}",
560 event.sequence, event.kind, event.name, event.generation
561 );
562 io::stdout().flush().map_err(io_usage_error)?;
563 }
564 Ok(())
565}
566
567async fn metrics(
568 client: &DcpClient,
569 addr: SocketAddr,
570 job: &str,
571 follow: bool,
572 json: bool,
573) -> Result<(), CliError> {
574 let mut subscription = client
575 .subscribe_metrics(0)
576 .await
577 .map_err(|error| command_error(addr, error))?;
578 if follow {
579 while let Some(sample) = subscription.recv().await {
580 let filtered = filter_metric_sample(job, &sample);
581 if filtered.streams.is_empty() {
582 continue;
583 }
584 print_metric_sample_follow(job, &filtered, json)?;
585 }
586 return Ok(());
587 }
588
589 let sample = wait_for_metric_sample(job, &mut subscription).await;
590 print_metric_sample(job, &sample, json)
591}
592
593async fn wait_for_metric_sample(
594 job: &str,
595 subscription: &mut MetricSubscription,
596) -> MetricSampleView {
597 let deadline = tokio::time::Instant::now() + METRIC_SNAPSHOT_WAIT;
598 let mut latest = MetricSampleView {
599 timestamp_ms: 0,
600 streams: Vec::new(),
601 };
602 while tokio::time::Instant::now() < deadline {
603 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
604 match tokio::time::timeout(
605 remaining.min(Duration::from_millis(100)),
606 subscription.recv(),
607 )
608 .await
609 {
610 Ok(Some(sample)) => {
611 latest = filter_metric_sample(job, &sample);
612 if !latest.streams.is_empty() {
613 break;
614 }
615 }
616 Ok(None) => break,
617 Err(_) => {}
618 }
619 }
620 latest
621}
622
623fn print_metric_sample(job: &str, sample: &MetricSampleView, json: bool) -> Result<(), CliError> {
624 if json {
625 print_json(&MetricsOutput { job, sample })?;
626 } else if sample.streams.is_empty() {
627 println!("no metrics for '{job}' yet");
628 } else {
629 print!("{}", render_metrics_table(sample));
630 }
631 Ok(())
632}
633
634fn print_metric_sample_follow(
635 job: &str,
636 sample: &MetricSampleView,
637 json: bool,
638) -> Result<(), CliError> {
639 if json {
640 print_json_line(&MetricsOutput { job, sample })?;
641 } else {
642 print!("{}", render_metrics_table(sample));
643 io::stdout().flush().map_err(io_usage_error)?;
644 }
645 Ok(())
646}
647
648fn filter_metric_sample(job: &str, sample: &MetricSample) -> MetricSampleView {
649 let prefix = format!("{job}:");
650 MetricSampleView {
651 timestamp_ms: sample.timestamp_ms,
652 streams: sample
653 .streams
654 .iter()
655 .filter(|metric| metric.name == job || metric.name.starts_with(&prefix))
656 .map(StreamMetricView::from)
657 .collect(),
658 }
659}
660
661fn parse_params(raw_params: &[String]) -> Result<HashMap<String, String>, CliError> {
662 let mut params = HashMap::new();
663 let mut seen = HashSet::new();
664 for raw in raw_params {
665 let Some((key, value)) = raw.split_once('=') else {
666 return Err(CliError::Usage(format!(
667 "invalid --param '{raw}'; use --param key=value."
668 )));
669 };
670 if key.trim().is_empty() {
671 return Err(CliError::Usage(format!(
672 "invalid --param '{raw}'; the key before '=' cannot be empty."
673 )));
674 }
675 if !seen.insert(key.to_owned()) {
676 return Err(CliError::Usage(format!(
677 "duplicate --param key '{key}'; pass each key once."
678 )));
679 }
680 params.insert(key.to_owned(), value.to_owned());
681 }
682 Ok(params)
683}
684
685#[derive(Serialize)]
686struct JobsOutput {
687 jobs: Vec<JobView>,
688}
689
690#[derive(Serialize)]
691struct JobOutput {
692 job: JobView,
693}
694
695#[derive(Serialize)]
696struct ActionOutput {
697 action: &'static str,
698 message: String,
699 job: JobView,
700}
701
702#[derive(Serialize)]
703struct EventsOutput {
704 events: Vec<EventView>,
705}
706
707#[derive(Serialize)]
708struct EventLine {
709 event: EventView,
710}
711
712#[derive(Serialize)]
713struct MetricsOutput<'a> {
714 job: &'a str,
715 sample: &'a MetricSampleView,
716}
717
718#[derive(Serialize)]
719struct JobView {
720 name: String,
721 job_id: u64,
722 state: String,
723 desired_state: String,
724 generation: u64,
725 starts_total: u64,
726 restarts_total: u64,
727 last_start_at_ms: Option<u64>,
728 last_exit_at_ms: Option<u64>,
729 last_exit_reason: String,
730 backoff_remaining_ms: Option<u64>,
731 drain_remaining_ms: Option<u64>,
732 drain_supported: bool,
733 active_streams: Option<u64>,
734}
735
736impl From<&WireJobStatus> for JobView {
737 fn from(status: &WireJobStatus) -> Self {
738 Self {
739 name: status.name.clone(),
740 job_id: status.job_id,
741 state: status.state.clone(),
742 desired_state: status.desired_state.clone(),
743 generation: status.generation,
744 starts_total: status.starts_total,
745 restarts_total: status.restarts_total,
746 last_start_at_ms: status.last_start_at_ms,
747 last_exit_at_ms: status.last_exit_at_ms,
748 last_exit_reason: status.last_exit_reason.clone(),
749 backoff_remaining_ms: status.backoff_remaining_ms,
750 drain_remaining_ms: status.drain_remaining_ms,
751 drain_supported: status.drain_supported,
752 active_streams: status.active_streams,
753 }
754 }
755}
756
757#[derive(Serialize)]
758struct EventView {
759 sequence: u64,
760 timestamp_ms: u64,
761 name: String,
762 job_id: u64,
763 generation: u64,
764 kind: String,
765 detail: String,
766}
767
768impl From<&Event> for EventView {
769 fn from(event: &Event) -> Self {
770 Self {
771 sequence: event.sequence,
772 timestamp_ms: event.timestamp_ms,
773 name: event.name.clone(),
774 job_id: event.job_id,
775 generation: event.generation,
776 kind: event.kind.clone(),
777 detail: event.detail.clone(),
778 }
779 }
780}
781
782#[derive(Serialize)]
783struct MetricSampleView {
784 timestamp_ms: u64,
785 streams: Vec<StreamMetricView>,
786}
787
788#[derive(Serialize)]
789struct StreamMetricView {
790 id: u64,
791 name: String,
792 elements_through: u64,
793 restarts: u64,
794 state: String,
795 started_at_ms: u64,
796 state_changed_at_ms: u64,
797 finished_at_ms: Option<u64>,
798 uptime_ms: u64,
799}
800
801impl From<&StreamMetric> for StreamMetricView {
802 fn from(metric: &StreamMetric) -> Self {
803 Self {
804 id: metric.id,
805 name: metric.name.clone(),
806 elements_through: metric.elements_through,
807 restarts: metric.restarts,
808 state: metric.state.clone(),
809 started_at_ms: metric.started_at_ms,
810 state_changed_at_ms: metric.state_changed_at_ms,
811 finished_at_ms: metric.finished_at_ms,
812 uptime_ms: metric.uptime_ms,
813 }
814 }
815}
816
817fn print_json<T: Serialize>(value: &T) -> Result<(), CliError> {
818 let json = serde_json::to_string_pretty(value).map_err(|error| {
819 CliError::Usage(format!(
820 "could not encode JSON output: {error}; retry without --json."
821 ))
822 })?;
823 println!("{json}");
824 Ok(())
825}
826
827fn print_json_line<T: Serialize>(value: &T) -> Result<(), CliError> {
828 let json = serde_json::to_string(value).map_err(|error| {
829 CliError::Usage(format!(
830 "could not encode JSON output: {error}; retry without --json."
831 ))
832 })?;
833 println!("{json}");
834 io::stdout().flush().map_err(io_usage_error)?;
835 Ok(())
836}
837
838fn io_usage_error(error: io::Error) -> CliError {
839 CliError::Usage(format!(
840 "could not write CLI output: {error}; check stdout/stderr."
841 ))
842}
843
844fn render_jobs_table(jobs: &[WireJobStatus]) -> String {
845 let rows = jobs
846 .iter()
847 .map(|job| {
848 vec![
849 job.name.clone(),
850 job.job_id.to_string(),
851 job.state.clone(),
852 job.desired_state.clone(),
853 job.generation.to_string(),
854 format_number(job.starts_total),
855 format_number(job.restarts_total),
856 format_optional_number(job.active_streams),
857 ]
858 })
859 .collect::<Vec<_>>();
860 render_table(
861 &[
862 "JOB", "ID", "STATE", "DESIRED", "GEN", "STARTS", "RESTARTS", "ACTIVE",
863 ],
864 &rows,
865 )
866}
867
868fn render_status_table(status: &WireJobStatus) -> String {
869 let rows = vec![
870 vec!["job".to_owned(), status.name.clone()],
871 vec!["id".to_owned(), status.job_id.to_string()],
872 vec!["state".to_owned(), status.state.clone()],
873 vec!["desired".to_owned(), status.desired_state.clone()],
874 vec!["generation".to_owned(), status.generation.to_string()],
875 vec!["starts".to_owned(), format_number(status.starts_total)],
876 vec!["restarts".to_owned(), format_number(status.restarts_total)],
877 vec![
878 "last_start_ms".to_owned(),
879 format_optional_number(status.last_start_at_ms),
880 ],
881 vec![
882 "last_exit_ms".to_owned(),
883 format_optional_number(status.last_exit_at_ms),
884 ],
885 vec![
886 "last_exit_reason".to_owned(),
887 empty_dash(&status.last_exit_reason).to_owned(),
888 ],
889 vec![
890 "backoff_remaining_ms".to_owned(),
891 format_optional_number(status.backoff_remaining_ms),
892 ],
893 vec![
894 "drain_remaining_ms".to_owned(),
895 format_optional_number(status.drain_remaining_ms),
896 ],
897 vec![
898 "drain_supported".to_owned(),
899 status.drain_supported.to_string(),
900 ],
901 vec![
902 "active_streams".to_owned(),
903 format_optional_number(status.active_streams),
904 ],
905 ];
906 render_table(&["FIELD", "VALUE"], &rows)
907}
908
909fn render_events_table(events: &[Event]) -> String {
910 let rows = events
911 .iter()
912 .map(|event| {
913 vec![
914 event.sequence.to_string(),
915 event.timestamp_ms.to_string(),
916 event.name.clone(),
917 event.job_id.to_string(),
918 event.generation.to_string(),
919 event.kind.clone(),
920 empty_dash(&event.detail).to_owned(),
921 ]
922 })
923 .collect::<Vec<_>>();
924 render_table(
925 &["SEQ", "TIME_MS", "JOB", "ID", "GEN", "KIND", "DETAIL"],
926 &rows,
927 )
928}
929
930fn render_metrics_table(sample: &MetricSampleView) -> String {
931 let rows = sample
932 .streams
933 .iter()
934 .map(|metric| {
935 vec![
936 sample.timestamp_ms.to_string(),
937 metric.name.clone(),
938 metric.id.to_string(),
939 metric.state.clone(),
940 format_number(metric.elements_through),
941 format_number(metric.restarts),
942 metric.uptime_ms.to_string(),
943 ]
944 })
945 .collect::<Vec<_>>();
946 render_table(
947 &[
948 "TIME_MS",
949 "STREAM",
950 "ID",
951 "STATE",
952 "ELEMENTS",
953 "RESTARTS",
954 "UPTIME_MS",
955 ],
956 &rows,
957 )
958}
959
960#[must_use]
962pub fn render_table(headers: &[&str], rows: &[Vec<String>]) -> String {
963 let mut widths = headers
964 .iter()
965 .map(|header| header.len())
966 .collect::<Vec<_>>();
967 for row in rows {
968 for (index, cell) in row.iter().enumerate().take(widths.len()) {
969 widths[index] = widths[index].max(cell.len());
970 }
971 }
972
973 let mut output = String::new();
974 push_table_row(
975 &mut output,
976 &headers
977 .iter()
978 .map(|cell| (*cell).to_owned())
979 .collect::<Vec<_>>(),
980 &widths,
981 );
982 push_separator(&mut output, &widths);
983 for row in rows {
984 push_table_row(&mut output, row, &widths);
985 }
986 output
987}
988
989fn push_table_row(output: &mut String, row: &[String], widths: &[usize]) {
990 for (index, width) in widths.iter().enumerate() {
991 if index > 0 {
992 output.push_str(" ");
993 }
994 let cell = row.get(index).map(String::as_str).unwrap_or("");
995 output.push_str(cell);
996 for _ in cell.len()..*width {
997 output.push(' ');
998 }
999 }
1000 output.push('\n');
1001}
1002
1003fn push_separator(output: &mut String, widths: &[usize]) {
1004 for (index, width) in widths.iter().enumerate() {
1005 if index > 0 {
1006 output.push_str(" ");
1007 }
1008 for _ in 0..*width {
1009 output.push('-');
1010 }
1011 }
1012 output.push('\n');
1013}
1014
1015fn empty_dash(value: &str) -> &str {
1016 if value.is_empty() { "-" } else { value }
1017}
1018
1019fn format_optional_number(value: Option<u64>) -> String {
1020 value.map(format_number).unwrap_or_else(|| "-".to_owned())
1021}
1022
1023fn format_number(value: u64) -> String {
1024 let raw = value.to_string();
1025 let mut output = String::with_capacity(raw.len() + raw.len() / 3);
1026 for (index, ch) in raw.chars().rev().enumerate() {
1027 if index > 0 && index % 3 == 0 {
1028 output.push(',');
1029 }
1030 output.push(ch);
1031 }
1032 output.chars().rev().collect()
1033}
1034
1035#[cfg(test)]
1036mod tests {
1037 use super::*;
1038
1039 #[test]
1040 fn table_renderer_pads_columns_stably() {
1041 let output = render_table(
1042 &["JOB", "STATE", "STARTS"],
1043 &[
1044 vec!["ingest".to_owned(), "Running".to_owned(), "1".to_owned()],
1045 vec![
1046 "daily-rollup".to_owned(),
1047 "Drained".to_owned(),
1048 "12,003".to_owned(),
1049 ],
1050 ],
1051 );
1052
1053 assert_eq!(
1054 output,
1055 "JOB STATE STARTS\n\
1056 ------------ ------- ------\n\
1057 ingest Running 1 \n\
1058 daily-rollup Drained 12,003\n"
1059 );
1060 }
1061
1062 #[test]
1063 fn params_require_key_value_pairs() {
1064 assert!(parse_params(&["threads=4".to_owned()]).is_ok());
1065 assert!(matches!(
1066 parse_params(&["threads".to_owned()]),
1067 Err(CliError::Usage(_))
1068 ));
1069 assert!(matches!(
1070 parse_params(&["threads=4".to_owned(), "threads=8".to_owned()]),
1071 Err(CliError::Usage(_))
1072 ));
1073 }
1074
1075 #[test]
1076 fn numbers_use_operator_grouping() {
1077 assert_eq!(format_number(0), "0");
1078 assert_eq!(format_number(12_003), "12,003");
1079 assert_eq!(format_number(9_876_543_210), "9,876,543,210");
1080 }
1081}