Skip to main content

datum_cli/
lib.rs

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    /// DCP address. Defaults to the loopback TCP development listener.
49    #[arg(long, global = true, default_value = DEFAULT_ADDR)]
50    addr: SocketAddr,
51
52    /// CA certificate for QUIC+mTLS mode.
53    #[arg(long, global = true, value_name = "PATH")]
54    tls_ca: Option<PathBuf>,
55
56    /// Client certificate for QUIC+mTLS mode.
57    #[arg(long, global = true, value_name = "PATH")]
58    tls_cert: Option<PathBuf>,
59
60    /// Client private key for QUIC+mTLS mode.
61    #[arg(long, global = true, value_name = "PATH")]
62    tls_key: Option<PathBuf>,
63
64    /// Emit machine-readable JSON instead of human tables.
65    #[arg(long, global = true)]
66    json: bool,
67
68    #[command(subcommand)]
69    command: Command,
70}
71
72#[derive(Debug, Subcommand)]
73enum Command {
74    /// List jobs.
75    Ps {
76        /// Aggregate jobs across cluster nodes.
77        #[arg(long)]
78        cluster: bool,
79    },
80    /// List cluster nodes.
81    Nodes,
82    /// Show one job.
83    Status {
84        #[arg(long)]
85        cluster: bool,
86        job: String,
87    },
88    /// Start a registered factory.
89    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 a job through cluster placement.
97    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    /// Gracefully drain a job.
112    Drain {
113        #[arg(long)]
114        cluster: bool,
115        job: String,
116    },
117    /// Stop a job immediately.
118    Stop {
119        #[arg(long)]
120        cluster: bool,
121        job: String,
122    },
123    /// Restart a job.
124    Restart { job: String },
125    /// Tail lifecycle events.
126    Events {
127        #[arg(long)]
128        follow: bool,
129    },
130    /// Show stream metrics for one job.
131    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
172/// Run the CLI using process arguments and return the requested process exit code.
173pub 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/// Render a simple left-aligned ASCII table with stable column widths.
1474#[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}