krabka_bench_driver/soak/drill/
cli.rs1use {
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#[derive(Debug, Parser)]
17#[command(name = "krabka-soak-drill", version, about)]
18pub struct Cli {
19 #[command(subcommand)]
20 command: Command,
21}
22
23#[derive(Debug, Subcommand)]
25pub enum Command {
26 Objectives {
29 #[arg(long)]
30 contract: PathBuf,
31 },
32 Load {
34 #[arg(long)]
35 contract: PathBuf,
36 #[command(flatten)]
37 login: Login,
38 #[arg(long, default_value = "")]
40 broker_image: String,
41 #[arg(long)]
42 out: PathBuf,
43 },
44 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 Readback {
55 #[arg(long)]
56 contract: PathBuf,
57 #[arg(long)]
59 load: PathBuf,
60 #[command(flatten)]
61 login: Login,
62 #[arg(long)]
64 capture_started_ms: i64,
65 #[arg(long)]
67 disaster_ms: i64,
68 #[arg(long)]
71 serving_ms: Option<i64>,
72 #[arg(long)]
73 out: PathBuf,
74 },
75}
76
77#[derive(Debug, Args)]
79pub struct Login {
80 #[arg(long)]
81 bootstrap: String,
82 #[arg(long)]
84 tls_ca: PathBuf,
85 #[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 #[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
107pub 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}