|
| 1 | +// Copyright ⓒ 2024-2025 Peter Morgan <peter.james.morgan@gmail.com> |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +use clap::{Args, Subcommand}; |
| 16 | +use human_units::iec::iec_unit; |
| 17 | +use tansu_perf::Perf; |
| 18 | +use tansu_sans_io::ErrorCode; |
| 19 | +use url::Url; |
| 20 | + |
| 21 | +use crate::{EnvVarExp, Result, cli::DEFAULT_BROKER}; |
| 22 | + |
| 23 | +#[derive(Clone, Debug, Subcommand)] |
| 24 | +pub(super) enum Command { |
| 25 | + /// Produce messages |
| 26 | + Produce { |
| 27 | + /// The partition to produce messages into |
| 28 | + #[arg(long, default_value = "0")] |
| 29 | + partition: i32, |
| 30 | + |
| 31 | + /// Message batch size used by every producer |
| 32 | + #[arg(long, default_value = "1")] |
| 33 | + batch_size: u32, |
| 34 | + |
| 35 | + /// Message batch size used by every producer |
| 36 | + #[arg(long, default_value = "1k", value_parser=clap::value_parser!(human_units::Size))] |
| 37 | + record_size: human_units::Size, |
| 38 | + |
| 39 | + /// The maximum number of messages per second |
| 40 | + #[clap(long, group = "output")] |
| 41 | + per_second: Option<u32>, |
| 42 | + |
| 43 | + /// Message throughput |
| 44 | + #[clap(long, group = "output")] |
| 45 | + throughput: Option<Throughput>, |
| 46 | + |
| 47 | + /// The number of producers generating messages |
| 48 | + #[arg(long, default_value = "1")] |
| 49 | + producers: u32, |
| 50 | + |
| 51 | + /// Stop sending messages after this time |
| 52 | + #[arg(long, default_value = "1m", value_parser=clap::value_parser!(human_units::Duration))] |
| 53 | + duration: Option<human_units::Duration>, |
| 54 | + }, |
| 55 | + |
| 56 | + Consume, |
| 57 | +} |
| 58 | + |
| 59 | +#[iec_unit(symbol = "B/s")] |
| 60 | +#[derive(Copy, Clone, Debug, Default, Eq, Hash, Ord, PartialEq, PartialOrd)] |
| 61 | +pub(super) struct Throughput(pub u32); |
| 62 | + |
| 63 | +#[derive(Args, Clone, Debug)] |
| 64 | +pub(super) struct Arg { |
| 65 | + #[command(subcommand)] |
| 66 | + command: Command, |
| 67 | + |
| 68 | + /// The URL of the broker |
| 69 | + #[arg(long, default_value = DEFAULT_BROKER, env = "ADVERTISED_LISTENER_URL")] |
| 70 | + broker: EnvVarExp<Url>, |
| 71 | + |
| 72 | + /// The topic to generate messages into |
| 73 | + #[clap(value_parser)] |
| 74 | + topic: String, |
| 75 | +} |
| 76 | + |
| 77 | +impl Arg { |
| 78 | + pub(super) async fn main(self) -> Result<ErrorCode> { |
| 79 | + match self.command { |
| 80 | + Command::Produce { |
| 81 | + partition, |
| 82 | + batch_size, |
| 83 | + record_size, |
| 84 | + per_second, |
| 85 | + throughput, |
| 86 | + producers, |
| 87 | + duration, |
| 88 | + } => Perf::builder() |
| 89 | + .broker(self.broker.into_inner()) |
| 90 | + .topic(self.topic) |
| 91 | + .partition(partition) |
| 92 | + .batch_size(batch_size) |
| 93 | + .record_size(record_size.0 as usize) |
| 94 | + .per_second(per_second) |
| 95 | + .throughput(throughput.map(|throughput| throughput.0)) |
| 96 | + .producers(producers) |
| 97 | + .duration(duration.map(|duration| duration.0)) |
| 98 | + .build() |
| 99 | + .main() |
| 100 | + .await |
| 101 | + .map_err(Into::into), |
| 102 | + |
| 103 | + Command::Consume => todo!(), |
| 104 | + } |
| 105 | + } |
| 106 | +} |
0 commit comments