Skip to main content

krabka_bench_driver/soak/drill/
cli.rs

1//! `krabka-soak-drill`'s command line, and [`run`], which carries a parsed one
2//! out. The binary only parses its arguments and listens for `SIGTERM`.
3
4use {
5    super::{load_contract, objectives, probe, record_load, record_readback},
6    crate::soak::{
7        client::{Endpoint, ScramLogin},
8        recovery::Boundary,
9    },
10    anyhow::Result,
11    clap::{Args, Parser, Subcommand},
12    std::{future::Future, path::PathBuf},
13};
14
15/// `krabka-soak-drill`'s command line.
16#[derive(Debug, Parser)]
17#[command(name = "krabka-soak-drill", version, about)]
18pub struct Cli {
19    #[command(subcommand)]
20    command: Command,
21}
22
23/// What `krabka-soak-drill` was asked to do.
24#[derive(Debug, Subcommand)]
25pub enum Command {
26    /// Validate a disaster-recovery contract and print its objectives and
27    /// paths as JSON, for the drill that creates the topics.
28    Objectives {
29        #[arg(long)]
30        contract: PathBuf,
31    },
32    /// Run the contract's load against the source cluster until `SIGTERM`.
33    Load {
34        #[arg(long)]
35        contract: PathBuf,
36        #[command(flatten)]
37        login: Login,
38        /// The broker image the source cluster runs, recorded in the evidence.
39        #[arg(long, default_value = "")]
40        broker_image: String,
41        #[arg(long)]
42        out: PathBuf,
43    },
44    /// Write one record as the authenticated principal and read it back.
45    Probe {
46        #[command(flatten)]
47        login: Login,
48        #[arg(long)]
49        topic: String,
50        #[arg(long, default_value_t = 0)]
51        partition: i32,
52    },
53    /// Read the restored cluster back and write the evidence the judge grades.
54    Readback {
55        #[arg(long)]
56        contract: PathBuf,
57        /// The `load` output.
58        #[arg(long)]
59        load: PathBuf,
60        #[command(flatten)]
61        login: Login,
62        /// When the capture started, in Unix milliseconds: the capture boundary.
63        #[arg(long)]
64        capture_started_ms: i64,
65        /// When the source cluster was lost, in Unix milliseconds.
66        #[arg(long)]
67        disaster_ms: i64,
68        /// When the restored cluster first served authenticated produce and
69        /// consume, in Unix milliseconds.
70        #[arg(long)]
71        serving_ms: Option<i64>,
72        #[arg(long)]
73        out: PathBuf,
74    },
75}
76
77/// A SCRAM-SHA-512 login over TLS.
78#[derive(Debug, Args)]
79pub struct Login {
80    #[arg(long)]
81    bootstrap: String,
82    /// The PEM CA bundle the brokers' certificates chain to.
83    #[arg(long)]
84    tls_ca: PathBuf,
85    /// A name the brokers' certificate carries.
86    #[arg(long)]
87    tls_server_name: String,
88    #[arg(long)]
89    username: String,
90    #[arg(long, env = "KRABKA_DRILL_PASSWORD", hide_env_values = true)]
91    password: String,
92}
93
94impl Login {
95    /// The endpoint this login reaches.
96    #[must_use]
97    pub fn endpoint(&self) -> Endpoint {
98        Endpoint::sasl_ssl(self.bootstrap.clone(), ScramLogin {
99            ca: self.tls_ca.clone(),
100            server_name: self.tls_server_name.clone(),
101            username: self.username.clone(),
102            password: self.password.clone(),
103        })
104    }
105}
106
107/// Carry out `cli`, and return what it prints on standard output.
108///
109/// `lost` is the loss of the source cluster: it ends a `load`.
110///
111/// # Errors
112///
113/// The contract, a load record or the cluster refuses the command, or a file
114/// cannot be written.
115pub async fn run(cli: Cli, lost: impl Future<Output = ()>) -> Result<String> {
116    match cli.command {
117        Command::Objectives { contract } => Ok(serde_json::to_string_pretty(&objectives(
118            &load_contract(&contract)?,
119        )?)?),
120        Command::Load {
121            contract,
122            login,
123            broker_image,
124            out,
125        } => {
126            let contract = load_contract(&contract)?;
127            record_load(&contract, login.endpoint(), broker_image, lost, &out).await?;
128            Ok(String::new())
129        }
130        Command::Probe {
131            login,
132            topic,
133            partition,
134        } => probe(&login.endpoint(), &topic, partition).await,
135        Command::Readback {
136            contract,
137            load,
138            login,
139            capture_started_ms,
140            disaster_ms,
141            serving_ms,
142            out,
143        } => {
144            let boundary = Boundary {
145                capture_started_unix_ms: capture_started_ms,
146                disaster_unix_ms: disaster_ms,
147            };
148            record_readback(
149                &load_contract(&contract)?,
150                &load,
151                &login.endpoint(),
152                boundary,
153                serving_ms,
154                &out,
155            )
156            .await?;
157            Ok(String::new())
158        }
159    }
160}