feat(app): add stdin-managed client target sets
This commit is contained in:
1 parent
31fd76f78d
commit
1ba4f61072
7 files changed
+704
-14
No files matched your search
@@ -20,6 +20,13 @@ pub struct ClientArgs {
|
|||||||
)]
|
)]
|
||||||
pub targets: Vec<TargetArg>,
|
pub targets: Vec<TargetArg>,
|
||||||
|
|
||||||
|
/// Read complete desired target sets from standard input in continuous mode.
|
||||||
|
#[arg(
|
||||||
|
long,
|
||||||
|
long_help = "Read one complete desired target set per stdin line in continuous mode. Commas separate targets and backslash escapes commas within a target. [] selects an empty desired set. EOF stops the client gracefully."
|
||||||
|
)]
|
||||||
|
pub targets_stdin: bool,
|
||||||
|
|
||||||
/// Managed group pacing for multi-target mode.
|
/// Managed group pacing for multi-target mode.
|
||||||
#[arg(
|
#[arg(
|
||||||
long,
|
long,
|
||||||
@@ -86,12 +93,16 @@ impl ClientArgs {
|
|||||||
/// Targets are only required here. `--list-columns` returns before this is
|
/// Targets are only required here. `--list-columns` returns before this is
|
||||||
/// called, so listing columns still needs no target.
|
/// called, so listing columns still needs no target.
|
||||||
pub fn prepare(&self) -> Result<ManagedRunSetup, String> {
|
pub fn prepare(&self) -> Result<ManagedRunSetup, String> {
|
||||||
|
if self.targets_stdin && !self.is_continuous() {
|
||||||
|
return Err("--targets-stdin requires --duration 0".to_owned());
|
||||||
|
}
|
||||||
prepare_managed_run(
|
prepare_managed_run(
|
||||||
&self.common,
|
&self.common,
|
||||||
self.duration,
|
self.duration,
|
||||||
TargetSelection {
|
TargetSelection {
|
||||||
targets: &self.targets,
|
targets: &self.targets,
|
||||||
pacing: self.pacing,
|
pacing: self.pacing,
|
||||||
|
stdin_controlled: self.targets_stdin,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -166,6 +177,7 @@ mod tests {
|
|||||||
assert_eq!(args.format, OutputFormat::Table);
|
assert_eq!(args.format, OutputFormat::Table);
|
||||||
assert_eq!(args.columns, None);
|
assert_eq!(args.columns, None);
|
||||||
assert_eq!(args.header, HeaderMode::Auto);
|
assert_eq!(args.header, HeaderMode::Auto);
|
||||||
|
assert!(!args.targets_stdin);
|
||||||
assert!(!args.is_continuous());
|
assert!(!args.is_continuous());
|
||||||
assert_eq!(args.timestamp_mode(), TimestampArg::Both);
|
assert_eq!(args.timestamp_mode(), TimestampArg::Both);
|
||||||
|
|
||||||
@@ -308,6 +320,21 @@ mod tests {
|
|||||||
assert_eq!(config.hmac_key, Some(b"secret".to_vec()));
|
assert_eq!(config.hmac_key, Some(b"secret".to_vec()));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn targets_stdin_requires_continuous_mode_and_allows_empty_initial_set() {
|
||||||
|
assert!(parse(&["--targets-stdin"]).unwrap().prepare().is_err());
|
||||||
|
let args = parse(&["--duration", "0", "--targets-stdin"]).unwrap();
|
||||||
|
let setup = args.prepare().unwrap();
|
||||||
|
assert!(setup.targets.is_empty());
|
||||||
|
assert!(setup.stdin_controlled);
|
||||||
|
let config = setup.managed_config();
|
||||||
|
assert_eq!(
|
||||||
|
config.completion,
|
||||||
|
irtt_client::managed::ManagedCompletionPolicy::ExplicitStop
|
||||||
|
);
|
||||||
|
assert_eq!(config.max_live_target_generations, 256);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn output_format_summary_policy_is_table_only() {
|
fn output_format_summary_policy_is_table_only() {
|
||||||
assert!(OutputFormat::Table.prints_summary());
|
assert!(OutputFormat::Table.prints_summary());
|
||||||
|
|||||||
@@ -1,16 +1,16 @@
|
|||||||
use std::collections::BTreeMap;
|
|
||||||
use std::{
|
use std::{
|
||||||
collections::HashSet,
|
collections::{BTreeMap, BTreeSet, HashSet, VecDeque},
|
||||||
io::{self, Write},
|
io::{self, BufRead, Write},
|
||||||
sync::atomic::AtomicBool,
|
sync::{atomic::AtomicBool, Arc, Mutex},
|
||||||
thread,
|
thread,
|
||||||
time::Duration,
|
time::Duration,
|
||||||
};
|
};
|
||||||
|
|
||||||
use irtt_client::{
|
use irtt_client::{
|
||||||
managed::{
|
managed::{
|
||||||
BlockingManagedClient, ManagedEndReason, ManagedEvent, ManagedEventSubscription,
|
BlockingManagedClient, ManagedClientHandle, ManagedCommandApplyError, ManagedEndReason,
|
||||||
ManagedEventTryRecvError, ManagedTargetEndReason, ManagedTargetOutcome, TargetInstance,
|
ManagedEvent, ManagedEventSubscription, ManagedEventTryRecvError, ManagedStatus,
|
||||||
|
ManagedTargetConfig, ManagedTargetEndReason, ManagedTargetOutcome, TargetInstance,
|
||||||
},
|
},
|
||||||
ClientEvent,
|
ClientEvent,
|
||||||
};
|
};
|
||||||
@@ -20,13 +20,13 @@ use super::{
|
|||||||
output::{EventRenderStats, OutputConfig},
|
output::{EventRenderStats, OutputConfig},
|
||||||
};
|
};
|
||||||
|
|
||||||
use crate::shared::client::expected_probe_count;
|
|
||||||
use crate::shared::client::{
|
use crate::shared::client::{
|
||||||
is_shutdown_requested,
|
expected_probe_count, is_shutdown_requested, parse_stdin_target_set,
|
||||||
session::{
|
session::{
|
||||||
drain_managed_events, peer_close_run_error, request_managed_stop_for_peer_close,
|
drain_managed_events, peer_close_run_error, request_managed_stop_for_peer_close,
|
||||||
request_managed_stop_once, should_print_final_summary, ManagedDrainState,
|
request_managed_stop_once, should_print_final_summary, ManagedDrainState,
|
||||||
},
|
},
|
||||||
|
ManagedRunSetup, STDIN_MAX_DESIRED_TARGETS, STDIN_OUTCOME_HISTORY_LIMIT,
|
||||||
};
|
};
|
||||||
|
|
||||||
use irtt_stats::{StatsCollector, StatsConfig};
|
use irtt_stats::{StatsCollector, StatsConfig};
|
||||||
@@ -41,6 +41,7 @@ const FINITE_STATS_MEMORY_STRONG_WARNING_BYTES: u64 = 512 * MIB;
|
|||||||
|
|
||||||
const FINITE_STATS_MEMORY_VERY_STRONG_WARNING_BYTES: u64 = GIB;
|
const FINITE_STATS_MEMORY_VERY_STRONG_WARNING_BYTES: u64 = GIB;
|
||||||
const MANAGED_EVENT_WAIT_SLICE: Duration = Duration::from_millis(20);
|
const MANAGED_EVENT_WAIT_SLICE: Duration = Duration::from_millis(20);
|
||||||
|
const MAX_STDIN_RECORD_BYTES: usize = 64 * 1024;
|
||||||
|
|
||||||
pub fn run_stream(
|
pub fn run_stream(
|
||||||
args: ClientArgs,
|
args: ClientArgs,
|
||||||
@@ -61,6 +62,9 @@ pub fn run_stream(
|
|||||||
args.verbose,
|
args.verbose,
|
||||||
)
|
)
|
||||||
.map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
|
.map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
|
||||||
|
if setup.stdin_controlled {
|
||||||
|
return run_stdin_stream(setup, output_config, shutdown_requested);
|
||||||
|
}
|
||||||
let continuous = args.is_continuous();
|
let continuous = args.is_continuous();
|
||||||
let target_count = setup.target_count();
|
let target_count = setup.target_count();
|
||||||
|
|
||||||
@@ -196,6 +200,561 @@ pub fn run_stream(
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct StdinTargetSet {
|
||||||
|
revision: u64,
|
||||||
|
targets: Vec<ManagedTargetConfig>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
struct StdinMailboxState {
|
||||||
|
next_revision: u64,
|
||||||
|
latest: Option<StdinTargetSet>,
|
||||||
|
eof: bool,
|
||||||
|
fatal: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
struct StdinMailbox {
|
||||||
|
state: Mutex<StdinMailboxState>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StdinMailbox {
|
||||||
|
fn publish(&self, targets: Vec<ManagedTargetConfig>) {
|
||||||
|
let mut state = self
|
||||||
|
.state
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
|
state.next_revision = state.next_revision.saturating_add(1);
|
||||||
|
state.latest = Some(StdinTargetSet {
|
||||||
|
revision: state.next_revision,
|
||||||
|
targets,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn finish(&self) {
|
||||||
|
self.state
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||||
|
.eof = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
fn fail(&self, error: String) {
|
||||||
|
let mut state = self
|
||||||
|
.state
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
|
state.fatal = Some(error);
|
||||||
|
state.latest = None;
|
||||||
|
}
|
||||||
|
|
||||||
|
fn latest(&self) -> Option<StdinTargetSet> {
|
||||||
|
self.state
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||||
|
.latest
|
||||||
|
.clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn clear_if_current(&self, revision: u64) {
|
||||||
|
let mut state = self
|
||||||
|
.state
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||||
|
if state
|
||||||
|
.latest
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|set| set.revision == revision)
|
||||||
|
{
|
||||||
|
state.latest = None;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_current(&self, revision: u64) -> bool {
|
||||||
|
self.state
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||||
|
.latest
|
||||||
|
.as_ref()
|
||||||
|
.is_some_and(|set| set.revision == revision)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn fatal(&self) -> Option<String> {
|
||||||
|
self.state
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||||
|
.fatal
|
||||||
|
.clone()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_finished(&self) -> bool {
|
||||||
|
self.state
|
||||||
|
.lock()
|
||||||
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||||
|
.eof
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_stdin_record<R: BufRead>(
|
||||||
|
reader: &mut R,
|
||||||
|
buffer: &mut Vec<u8>,
|
||||||
|
) -> io::Result<Option<String>> {
|
||||||
|
buffer.clear();
|
||||||
|
loop {
|
||||||
|
let available = reader.fill_buf()?;
|
||||||
|
if available.is_empty() {
|
||||||
|
if buffer.is_empty() {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
if buffer.len() > MAX_STDIN_RECORD_BYTES {
|
||||||
|
return Err(io::Error::new(
|
||||||
|
io::ErrorKind::InvalidData,
|
||||||
|
"stdin target record exceeds the maximum size",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
return std::str::from_utf8(buffer)
|
||||||
|
.map(str::to_owned)
|
||||||
|
.map(Some)
|
||||||
|
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "stdin is not UTF-8"));
|
||||||
|
}
|
||||||
|
|
||||||
|
if buffer.len() == MAX_STDIN_RECORD_BYTES + 1 {
|
||||||
|
if available.first() == Some(&b'\n') && buffer.last() == Some(&b'\r') {
|
||||||
|
reader.consume(1);
|
||||||
|
buffer.pop();
|
||||||
|
return std::str::from_utf8(buffer)
|
||||||
|
.map(str::to_owned)
|
||||||
|
.map(Some)
|
||||||
|
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "stdin is not UTF-8"));
|
||||||
|
}
|
||||||
|
return Err(io::Error::new(
|
||||||
|
io::ErrorKind::InvalidData,
|
||||||
|
"stdin target record exceeds the maximum size",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(newline) = available.iter().position(|byte| *byte == b'\n') {
|
||||||
|
let before_newline = &available[..newline];
|
||||||
|
if buffer.len().saturating_add(before_newline.len()) > MAX_STDIN_RECORD_BYTES + 1 {
|
||||||
|
return Err(io::Error::new(
|
||||||
|
io::ErrorKind::InvalidData,
|
||||||
|
"stdin target record exceeds the maximum size",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
buffer.extend_from_slice(before_newline);
|
||||||
|
reader.consume(newline + 1);
|
||||||
|
if buffer.last() == Some(&b'\r') {
|
||||||
|
buffer.pop();
|
||||||
|
}
|
||||||
|
if buffer.len() > MAX_STDIN_RECORD_BYTES {
|
||||||
|
return Err(io::Error::new(
|
||||||
|
io::ErrorKind::InvalidData,
|
||||||
|
"stdin target record exceeds the maximum size",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
return std::str::from_utf8(buffer)
|
||||||
|
.map(str::to_owned)
|
||||||
|
.map(Some)
|
||||||
|
.map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "stdin is not UTF-8"));
|
||||||
|
}
|
||||||
|
|
||||||
|
let remaining = MAX_STDIN_RECORD_BYTES + 1 - buffer.len();
|
||||||
|
let take = available.len().min(remaining);
|
||||||
|
buffer.extend_from_slice(&available[..take]);
|
||||||
|
reader.consume(take);
|
||||||
|
if buffer.len() > MAX_STDIN_RECORD_BYTES && buffer.last() != Some(&b'\r') {
|
||||||
|
return Err(io::Error::new(
|
||||||
|
io::ErrorKind::InvalidData,
|
||||||
|
"stdin target record exceeds the maximum size",
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_stdin_target_sets<R: BufRead>(reader: &mut R, mailbox: &StdinMailbox) {
|
||||||
|
let mut line = 0_u64;
|
||||||
|
let mut record = Vec::with_capacity(MAX_STDIN_RECORD_BYTES + 1);
|
||||||
|
loop {
|
||||||
|
match read_stdin_record(reader, &mut record) {
|
||||||
|
Ok(Some(record)) => {
|
||||||
|
line = line.saturating_add(1);
|
||||||
|
if record.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
match parse_stdin_target_set(&record, STDIN_MAX_DESIRED_TARGETS) {
|
||||||
|
Ok(targets) => {
|
||||||
|
mailbox.publish(targets.into_iter().map(|target| target.managed).collect())
|
||||||
|
}
|
||||||
|
Err(error) => {
|
||||||
|
mailbox.fail(format!("invalid --targets-stdin line {line}: {error}"));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(None) => {
|
||||||
|
mailbox.finish();
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
mailbox.fail(format!("failed to read --targets-stdin line {}", line + 1));
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn spawn_stdin_target_reader(mailbox: Arc<StdinMailbox>) -> io::Result<()> {
|
||||||
|
thread::Builder::new()
|
||||||
|
.name("irtt-targets-stdin".to_owned())
|
||||||
|
.spawn(move || {
|
||||||
|
let stdin = io::stdin();
|
||||||
|
read_stdin_target_sets(&mut stdin.lock(), &mailbox);
|
||||||
|
})
|
||||||
|
.map(|_| ())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
struct RetryAfterStatus {
|
||||||
|
revision: u64,
|
||||||
|
status: Arc<ManagedStatus>,
|
||||||
|
}
|
||||||
|
|
||||||
|
fn retry_waits_for_status_change(
|
||||||
|
retry_after: &mut Option<RetryAfterStatus>,
|
||||||
|
target_set: &StdinTargetSet,
|
||||||
|
status: &Arc<ManagedStatus>,
|
||||||
|
) -> bool {
|
||||||
|
let Some(retry) = retry_after else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
if retry.revision != target_set.revision {
|
||||||
|
*retry_after = None;
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
Arc::ptr_eq(&retry.status, status)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn retry_after_capacity_rejection(
|
||||||
|
mailbox: &StdinMailbox,
|
||||||
|
target_set: &StdinTargetSet,
|
||||||
|
status: Arc<ManagedStatus>,
|
||||||
|
) -> Option<RetryAfterStatus> {
|
||||||
|
mailbox
|
||||||
|
.is_current(target_set.revision)
|
||||||
|
.then_some(RetryAfterStatus {
|
||||||
|
revision: target_set.revision,
|
||||||
|
status,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn apply_latest_stdin_target_set(
|
||||||
|
handle: &ManagedClientHandle,
|
||||||
|
mailbox: &StdinMailbox,
|
||||||
|
retry_after: &mut Option<RetryAfterStatus>,
|
||||||
|
) -> Result<(), String> {
|
||||||
|
let Some(target_set) = mailbox.latest() else {
|
||||||
|
*retry_after = None;
|
||||||
|
return Ok(());
|
||||||
|
};
|
||||||
|
let status = handle.status();
|
||||||
|
if retry_waits_for_status_change(retry_after, &target_set, &status) {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let receipt = handle
|
||||||
|
.update_targets(target_set.targets.clone())
|
||||||
|
.map_err(|_| "failed to submit --targets-stdin update".to_owned())?;
|
||||||
|
match receipt.blocking_wait() {
|
||||||
|
Ok(_) => {
|
||||||
|
mailbox.clear_if_current(target_set.revision);
|
||||||
|
*retry_after = None;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(ManagedCommandApplyError::LiveGenerationLimitExceeded { .. }) => {
|
||||||
|
*retry_after = retry_after_capacity_rejection(mailbox, &target_set, status);
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(_) => Err("--targets-stdin update was rejected".to_owned()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
struct BoundedTargetSet {
|
||||||
|
limit: usize,
|
||||||
|
order: VecDeque<TargetInstance>,
|
||||||
|
members: HashSet<TargetInstance>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BoundedTargetSet {
|
||||||
|
fn new(limit: usize) -> Self {
|
||||||
|
Self {
|
||||||
|
limit,
|
||||||
|
order: VecDeque::with_capacity(limit),
|
||||||
|
members: HashSet::with_capacity(limit),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn insert(&mut self, target: TargetInstance) -> bool {
|
||||||
|
if !self.members.insert(target.clone()) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
if self.order.len() == self.limit {
|
||||||
|
let evicted = self
|
||||||
|
.order
|
||||||
|
.pop_front()
|
||||||
|
.expect("bounded target order is nonempty");
|
||||||
|
self.members.remove(&evicted);
|
||||||
|
}
|
||||||
|
self.order.push_back(target);
|
||||||
|
true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn process_stdin_event<W: Write>(
|
||||||
|
event: ManagedEvent,
|
||||||
|
stream_output: &mut StreamOutput<'_, W>,
|
||||||
|
stats: &mut BTreeMap<TargetInstance, StatsCollector>,
|
||||||
|
terminal_targets: &mut BoundedTargetSet,
|
||||||
|
) -> io::Result<()> {
|
||||||
|
match event {
|
||||||
|
ManagedEvent::TargetStateChanged { target, .. } => {
|
||||||
|
stats
|
||||||
|
.entry(target)
|
||||||
|
.or_insert_with(|| StatsCollector::new(stats_config(true)));
|
||||||
|
}
|
||||||
|
ManagedEvent::Client { target, event } => {
|
||||||
|
let collector = stats
|
||||||
|
.entry(target.clone())
|
||||||
|
.or_insert_with(|| StatsCollector::new(stats_config(true)));
|
||||||
|
print_events_with_stats(
|
||||||
|
stream_output,
|
||||||
|
std::slice::from_ref(&event),
|
||||||
|
Some(target.id.as_str()),
|
||||||
|
collector,
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
ManagedEvent::TargetFinished { outcome }
|
||||||
|
if terminal_targets.insert(outcome.target.clone()) =>
|
||||||
|
{
|
||||||
|
report_target_failure(&outcome);
|
||||||
|
}
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn drain_stdin_events<W: Write>(
|
||||||
|
events: &mut ManagedEventSubscription,
|
||||||
|
stream_output: &mut StreamOutput<'_, W>,
|
||||||
|
stats: &mut BTreeMap<TargetInstance, StatsCollector>,
|
||||||
|
terminal_targets: &mut BoundedTargetSet,
|
||||||
|
dropped_events: &mut u64,
|
||||||
|
) -> io::Result<ManagedDrainState> {
|
||||||
|
drain_managed_events(events, dropped_events, |event| {
|
||||||
|
process_stdin_event(event, stream_output, stats, terminal_targets)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn drain_final_stdin_events<W: Write>(
|
||||||
|
events: &mut ManagedEventSubscription,
|
||||||
|
stream_output: &mut StreamOutput<'_, W>,
|
||||||
|
stats: &mut BTreeMap<TargetInstance, StatsCollector>,
|
||||||
|
terminal_targets: &mut BoundedTargetSet,
|
||||||
|
dropped_events: &mut u64,
|
||||||
|
) -> io::Result<()> {
|
||||||
|
loop {
|
||||||
|
match events.try_recv() {
|
||||||
|
Ok(event) => process_stdin_event(event, stream_output, stats, terminal_targets)?,
|
||||||
|
Err(ManagedEventTryRecvError::Empty | ManagedEventTryRecvError::Closed) => break,
|
||||||
|
Err(ManagedEventTryRecvError::Lagged(count)) => {
|
||||||
|
*dropped_events = dropped_events.saturating_add(count);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reconcile_stdin_stats(
|
||||||
|
stats: &mut BTreeMap<TargetInstance, StatsCollector>,
|
||||||
|
status: &ManagedStatus,
|
||||||
|
final_summary_targets: Option<&BTreeSet<TargetInstance>>,
|
||||||
|
) {
|
||||||
|
let live = status
|
||||||
|
.targets
|
||||||
|
.iter()
|
||||||
|
.map(|target| target.target.clone())
|
||||||
|
.collect::<HashSet<_>>();
|
||||||
|
stats.retain(|target, _| {
|
||||||
|
live.contains(target)
|
||||||
|
|| final_summary_targets.is_some_and(|targets| targets.contains(target))
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
fn snapshot_stdin_summary_targets(
|
||||||
|
stats: &mut BTreeMap<TargetInstance, StatsCollector>,
|
||||||
|
status: &ManagedStatus,
|
||||||
|
) -> BTreeSet<TargetInstance> {
|
||||||
|
let targets = status
|
||||||
|
.targets
|
||||||
|
.iter()
|
||||||
|
.map(|target| target.target.clone())
|
||||||
|
.collect::<BTreeSet<_>>();
|
||||||
|
for target in &targets {
|
||||||
|
stats
|
||||||
|
.entry(target.clone())
|
||||||
|
.or_insert_with(|| StatsCollector::new(stats_config(true)));
|
||||||
|
}
|
||||||
|
targets
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
enum StdinStop {
|
||||||
|
Interrupted,
|
||||||
|
Eof,
|
||||||
|
Fatal(String),
|
||||||
|
}
|
||||||
|
|
||||||
|
fn stdin_stop_request(
|
||||||
|
shutdown_requested: &AtomicBool,
|
||||||
|
mailbox: &StdinMailbox,
|
||||||
|
) -> Option<StdinStop> {
|
||||||
|
if is_shutdown_requested(shutdown_requested) {
|
||||||
|
Some(StdinStop::Interrupted)
|
||||||
|
} else if let Some(error) = mailbox.fatal() {
|
||||||
|
Some(StdinStop::Fatal(error))
|
||||||
|
} else if mailbox.is_finished() {
|
||||||
|
Some(StdinStop::Eof)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn run_stdin_stream(
|
||||||
|
setup: ManagedRunSetup,
|
||||||
|
output_config: OutputConfig,
|
||||||
|
shutdown_requested: &AtomicBool,
|
||||||
|
) -> Result<(), Box<dyn std::error::Error>> {
|
||||||
|
if is_shutdown_requested(shutdown_requested) {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
let (owner, mut events) = BlockingManagedClient::start_with_subscription(
|
||||||
|
setup.managed_config(),
|
||||||
|
setup.managed_targets(),
|
||||||
|
)?;
|
||||||
|
let handle = owner.handle();
|
||||||
|
let mailbox = Arc::new(StdinMailbox::default());
|
||||||
|
spawn_stdin_target_reader(Arc::clone(&mailbox))?;
|
||||||
|
|
||||||
|
let mut stdout = io::LineWriter::new(io::stdout().lock());
|
||||||
|
let mut stream_output = StreamOutput {
|
||||||
|
config: output_config,
|
||||||
|
header_printed: false,
|
||||||
|
print_final_summary: false,
|
||||||
|
show_running_only_summary_note: false,
|
||||||
|
out: &mut stdout,
|
||||||
|
};
|
||||||
|
let mut stats = BTreeMap::new();
|
||||||
|
let mut terminal_targets = BoundedTargetSet::new(STDIN_OUTCOME_HISTORY_LIMIT);
|
||||||
|
let mut retry_after = None;
|
||||||
|
let mut dropped_events = 0_u64;
|
||||||
|
let mut stop_requested = false;
|
||||||
|
let mut stop = None;
|
||||||
|
let mut summary_targets = None;
|
||||||
|
let mut subscription_closed = false;
|
||||||
|
|
||||||
|
loop {
|
||||||
|
let status = handle.status();
|
||||||
|
if stop.is_none() {
|
||||||
|
if let Some(requested_stop) = stdin_stop_request(shutdown_requested, &mailbox) {
|
||||||
|
summary_targets = Some(snapshot_stdin_summary_targets(&mut stats, &status));
|
||||||
|
stop = Some(requested_stop);
|
||||||
|
if request_managed_stop_once(&mut stop_requested) {
|
||||||
|
drop(handle.stop());
|
||||||
|
}
|
||||||
|
} else if let Err(error) =
|
||||||
|
apply_latest_stdin_target_set(&handle, &mailbox, &mut retry_after)
|
||||||
|
{
|
||||||
|
summary_targets = Some(snapshot_stdin_summary_targets(&mut stats, &status));
|
||||||
|
stop = Some(StdinStop::Fatal(error));
|
||||||
|
if request_managed_stop_once(&mut stop_requested) {
|
||||||
|
drop(handle.stop());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let drain_state = drain_stdin_events(
|
||||||
|
&mut events,
|
||||||
|
&mut stream_output,
|
||||||
|
&mut stats,
|
||||||
|
&mut terminal_targets,
|
||||||
|
&mut dropped_events,
|
||||||
|
)?;
|
||||||
|
reconcile_stdin_stats(&mut stats, &handle.status(), summary_targets.as_ref());
|
||||||
|
match drain_state {
|
||||||
|
ManagedDrainState::Empty => thread::sleep(MANAGED_EVENT_WAIT_SLICE),
|
||||||
|
ManagedDrainState::Closed => subscription_closed = true,
|
||||||
|
ManagedDrainState::BudgetExhausted => {}
|
||||||
|
}
|
||||||
|
if handle.status().final_outcome.is_some() || subscription_closed {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let outcome = owner.join()?;
|
||||||
|
drain_final_stdin_events(
|
||||||
|
&mut events,
|
||||||
|
&mut stream_output,
|
||||||
|
&mut stats,
|
||||||
|
&mut terminal_targets,
|
||||||
|
&mut dropped_events,
|
||||||
|
)?;
|
||||||
|
if let Some(warning) = dropped_event_warning(dropped_events) {
|
||||||
|
eprintln!("{warning}");
|
||||||
|
}
|
||||||
|
for target in outcome.recent_target_outcomes.iter() {
|
||||||
|
if terminal_targets.insert(target.target.clone()) {
|
||||||
|
report_target_failure(target);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if outcome.discarded_target_outcomes != 0 {
|
||||||
|
eprintln!(
|
||||||
|
"irtt-rs: warning: {} final target outcomes were discarded",
|
||||||
|
outcome.discarded_target_outcomes
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let print_summary = matches!(stop, Some(StdinStop::Interrupted | StdinStop::Eof));
|
||||||
|
stream_output.print_final_summary = print_summary;
|
||||||
|
stream_output.show_running_only_summary_note = print_summary;
|
||||||
|
if let Some(summary_targets) = summary_targets {
|
||||||
|
let multi_target = summary_targets.len() > 1;
|
||||||
|
for target in summary_targets {
|
||||||
|
let stats = stats
|
||||||
|
.get(&target)
|
||||||
|
.expect("summary snapshot retains every selected target collector");
|
||||||
|
if multi_target && stream_output.config.prints_summary() {
|
||||||
|
writeln!(stream_output.out)?;
|
||||||
|
writeln!(
|
||||||
|
stream_output.out,
|
||||||
|
"target: {} (generation {})",
|
||||||
|
target.id, target.generation
|
||||||
|
)?;
|
||||||
|
}
|
||||||
|
stream_output.print_summary(stats)?;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
stream_output.out.flush()?;
|
||||||
|
|
||||||
|
match stop {
|
||||||
|
Some(StdinStop::Fatal(error)) => Err(error.into()),
|
||||||
|
_ => match outcome.end_reason {
|
||||||
|
ManagedEndReason::DriverFailed(failure) => {
|
||||||
|
Err(format!("managed driver failed: {failure}").into())
|
||||||
|
}
|
||||||
|
_ => Ok(()),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn process_event<W: Write>(
|
fn process_event<W: Write>(
|
||||||
event: ManagedEvent,
|
event: ManagedEvent,
|
||||||
stream_output: &mut StreamOutput<'_, W>,
|
stream_output: &mut StreamOutput<'_, W>,
|
||||||
|
|||||||
@@ -58,6 +58,7 @@ impl TuiArgs {
|
|||||||
TargetSelection {
|
TargetSelection {
|
||||||
targets: &self.targets,
|
targets: &self.targets,
|
||||||
pacing: self.pacing,
|
pacing: self.pacing,
|
||||||
|
stdin_controlled: false,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,9 +9,12 @@ pub use args::{
|
|||||||
ClockArg, CommonClientArgs, ReceivedStatsArg, TimestampArg,
|
ClockArg, CommonClientArgs, ReceivedStatsArg, TimestampArg,
|
||||||
};
|
};
|
||||||
pub use config::{expected_probe_count, DEFAULT_RECV_TIMEOUT};
|
pub use config::{expected_probe_count, DEFAULT_RECV_TIMEOUT};
|
||||||
pub use prepare::{prepare_managed_run, ManagedRunSetup, TargetSelection, MANAGED_EVENT_CAPACITY};
|
pub use prepare::{
|
||||||
|
prepare_managed_run, ManagedRunSetup, TargetSelection, MANAGED_EVENT_CAPACITY,
|
||||||
|
STDIN_MAX_DESIRED_TARGETS, STDIN_MAX_LIVE_TARGET_GENERATIONS, STDIN_OUTCOME_HISTORY_LIMIT,
|
||||||
|
};
|
||||||
pub use session::is_shutdown_requested;
|
pub use session::is_shutdown_requested;
|
||||||
pub use targets::{
|
pub use targets::{
|
||||||
parse_target, prepare_managed_targets, target_specs, GroupPacingArg, PreparedTarget, TargetArg,
|
parse_stdin_target_set, parse_target, prepare_managed_targets, target_specs, GroupPacingArg,
|
||||||
TargetSpec,
|
PreparedTarget, TargetArg, TargetSpec,
|
||||||
};
|
};
|
||||||
@@ -22,6 +22,9 @@ use super::{
|
|||||||
|
|
||||||
/// Capacity of the lossy managed presentation event channel.
|
/// Capacity of the lossy managed presentation event channel.
|
||||||
pub const MANAGED_EVENT_CAPACITY: usize = 16_384;
|
pub const MANAGED_EVENT_CAPACITY: usize = 16_384;
|
||||||
|
pub const STDIN_MAX_DESIRED_TARGETS: usize = 128;
|
||||||
|
pub const STDIN_MAX_LIVE_TARGET_GENERATIONS: usize = 256;
|
||||||
|
pub const STDIN_OUTCOME_HISTORY_LIMIT: usize = 256;
|
||||||
|
|
||||||
/// Target and pacing arguments common to the client applets.
|
/// Target and pacing arguments common to the client applets.
|
||||||
///
|
///
|
||||||
@@ -31,6 +34,7 @@ pub const MANAGED_EVENT_CAPACITY: usize = 16_384;
|
|||||||
pub struct TargetSelection<'a> {
|
pub struct TargetSelection<'a> {
|
||||||
pub targets: &'a [TargetArg],
|
pub targets: &'a [TargetArg],
|
||||||
pub pacing: GroupPacingArg,
|
pub pacing: GroupPacingArg,
|
||||||
|
pub stdin_controlled: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Everything a managed run needs, prepared and validated together.
|
/// Everything a managed run needs, prepared and validated together.
|
||||||
@@ -46,6 +50,8 @@ pub struct ManagedRunSetup {
|
|||||||
pub client: ClientConfig,
|
pub client: ClientConfig,
|
||||||
/// Send coordination across active targets.
|
/// Send coordination across active targets.
|
||||||
pub pacing: ManagedPacing,
|
pub pacing: ManagedPacing,
|
||||||
|
/// Whether this stream may replace its desired set from stdin.
|
||||||
|
pub stdin_controlled: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ManagedRunSetup {
|
impl ManagedRunSetup {
|
||||||
@@ -70,7 +76,7 @@ impl ManagedRunSetup {
|
|||||||
/// Managed driver configuration for this run.
|
/// Managed driver configuration for this run.
|
||||||
pub fn managed_config(&self) -> ManagedClientConfig {
|
pub fn managed_config(&self) -> ManagedClientConfig {
|
||||||
let target_count = self.target_count();
|
let target_count = self.target_count();
|
||||||
ManagedClientConfig {
|
let mut config = ManagedClientConfig {
|
||||||
client: self.client.clone(),
|
client: self.client.clone(),
|
||||||
pacing: self.pacing,
|
pacing: self.pacing,
|
||||||
completion: ManagedCompletionPolicy::FinishWhenQuiescent,
|
completion: ManagedCompletionPolicy::FinishWhenQuiescent,
|
||||||
@@ -78,7 +84,13 @@ impl ManagedRunSetup {
|
|||||||
outcome_history_limit: target_count,
|
outcome_history_limit: target_count,
|
||||||
max_live_target_generations: target_count,
|
max_live_target_generations: target_count,
|
||||||
..ManagedClientConfig::default()
|
..ManagedClientConfig::default()
|
||||||
|
};
|
||||||
|
if self.stdin_controlled {
|
||||||
|
config.completion = ManagedCompletionPolicy::ExplicitStop;
|
||||||
|
config.outcome_history_limit = STDIN_OUTCOME_HISTORY_LIMIT;
|
||||||
|
config.max_live_target_generations = STDIN_MAX_LIVE_TARGET_GENERATIONS;
|
||||||
}
|
}
|
||||||
|
config
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -91,12 +103,22 @@ pub fn prepare_managed_run(
|
|||||||
duration: Duration,
|
duration: Duration,
|
||||||
selection: TargetSelection<'_>,
|
selection: TargetSelection<'_>,
|
||||||
) -> Result<ManagedRunSetup, String> {
|
) -> Result<ManagedRunSetup, String> {
|
||||||
let specs = target_specs(selection.targets)?;
|
let specs = if selection.stdin_controlled {
|
||||||
|
if selection.targets.len() > STDIN_MAX_DESIRED_TARGETS {
|
||||||
|
return Err(format!(
|
||||||
|
"initial target set exceeds the {STDIN_MAX_DESIRED_TARGETS}-target limit"
|
||||||
|
));
|
||||||
|
}
|
||||||
|
super::targets::target_specs_with_empty(selection.targets, true)?
|
||||||
|
} else {
|
||||||
|
target_specs(selection.targets)?
|
||||||
|
};
|
||||||
let targets = prepare_managed_targets(specs)?;
|
let targets = prepare_managed_targets(specs)?;
|
||||||
Ok(ManagedRunSetup {
|
Ok(ManagedRunSetup {
|
||||||
targets,
|
targets,
|
||||||
client: common.to_client_config(duration),
|
client: common.to_client_config(duration),
|
||||||
pacing: selection.pacing.into(),
|
pacing: selection.pacing.into(),
|
||||||
|
stdin_controlled: selection.stdin_controlled,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -131,7 +153,15 @@ mod tests {
|
|||||||
pacing: GroupPacingArg,
|
pacing: GroupPacingArg,
|
||||||
duration: Duration,
|
duration: Duration,
|
||||||
) -> Result<ManagedRunSetup, String> {
|
) -> Result<ManagedRunSetup, String> {
|
||||||
prepare_managed_run(&common(&[]), duration, TargetSelection { targets, pacing })
|
prepare_managed_run(
|
||||||
|
&common(&[]),
|
||||||
|
duration,
|
||||||
|
TargetSelection {
|
||||||
|
targets,
|
||||||
|
pacing,
|
||||||
|
stdin_controlled: false,
|
||||||
|
},
|
||||||
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -206,6 +236,7 @@ mod tests {
|
|||||||
TargetSelection {
|
TargetSelection {
|
||||||
targets: &targets,
|
targets: &targets,
|
||||||
pacing: GroupPacingArg::Staggered,
|
pacing: GroupPacingArg::Staggered,
|
||||||
|
stdin_controlled: false,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
|||||||
@@ -1576,6 +1576,37 @@ fn authenticated_peer_close_outcome() {
|
|||||||
assert!(!has_close(&records));
|
assert!(!has_close(&records));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn peer_close_leaves_other_explicitly_managed_targets_running() {
|
||||||
|
let closing = start_server(ServerBehavior::PeerClose, None);
|
||||||
|
let continuing = start_server(ServerBehavior::Echo, None);
|
||||||
|
let mut managed = config(ManagedPacing::Staggered);
|
||||||
|
managed.completion = ManagedCompletionPolicy::ExplicitStop;
|
||||||
|
managed.client.duration = None;
|
||||||
|
let (task, handle) = ManagedClient::task(
|
||||||
|
managed,
|
||||||
|
vec![
|
||||||
|
target("closing", closing.addr),
|
||||||
|
target("continuing", continuing.addr),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
.unwrap();
|
||||||
|
let runtime = runtime();
|
||||||
|
let mut task = Box::pin(task);
|
||||||
|
|
||||||
|
drive_task_until(&runtime, &mut task, |task| {
|
||||||
|
handle.status().peer_closed_target_outcomes == 1
|
||||||
|
&& task.targets[1].counters.packets_sent != 0
|
||||||
|
});
|
||||||
|
assert_eq!(handle.status().lifecycle, ManagedLifecycle::Running);
|
||||||
|
|
||||||
|
let stop = handle.stop();
|
||||||
|
runtime.block_on(task);
|
||||||
|
runtime.block_on(stop);
|
||||||
|
closing.finish();
|
||||||
|
assert!(!probes(&continuing.finish()).is_empty());
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn peer_close_during_stop_drain_remains_authoritative() {
|
fn peer_close_during_stop_drain_remains_authoritative() {
|
||||||
let key = b"managed-drain-peer-close".to_vec();
|
let key = b"managed-drain-peer-close".to_vec();
|
||||||
|
|||||||
@@ -45,6 +45,40 @@ once per interval instead. Final per-target summaries print in the order
|
|||||||
targets were supplied on the command line, not the order in which they
|
targets were supplied on the command line, not the order in which they
|
||||||
finish or their labels sort alphabetically.
|
finish or their labels sort alphabetically.
|
||||||
|
|
||||||
|
## DYNAMIC TARGET SETS FROM STANDARD INPUT
|
||||||
|
|
||||||
|
`--targets-stdin` is available only with continuous mode (`--duration 0`).
|
||||||
|
Positional targets remain an initial desired set, and may be omitted to
|
||||||
|
start empty:
|
||||||
|
|
||||||
|
~~~sh
|
||||||
|
irtt-client --duration 0 --targets-stdin
|
||||||
|
irtt-client --duration 0 --targets-stdin ams=ams.example:2112
|
||||||
|
~~~
|
||||||
|
|
||||||
|
Each non-empty stdin record declares the complete desired target set, rather
|
||||||
|
than a delta. A later record replaces an earlier set: targets absent from it
|
||||||
|
are retired, unchanged target configurations retain their generation, and a
|
||||||
|
changed address or HMAC setting creates a fresh generation. Under transient
|
||||||
|
live-generation backpressure, only the latest unapplied desired set is kept.
|
||||||
|
|
||||||
|
Commas frame targets in stdin records; escape a literal comma within one
|
||||||
|
target as `\,`. Target parsing inside each element is otherwise the same as
|
||||||
|
for positional arguments, including `@hmac=KEY` and `@hmac=`:
|
||||||
|
|
||||||
|
~~~sh
|
||||||
|
printf '%s\n' \
|
||||||
|
'ams=ams.example:2112,sg=sg.example:2112@hmac=sg-secret' \
|
||||||
|
'sg=sg.example:2112@hmac=sg-secret,nyc=nyc.example:2112' |
|
||||||
|
irtt-client --duration 0 --targets-stdin --hmac default-secret
|
||||||
|
~~~
|
||||||
|
|
||||||
|
The exact record `[]` selects an empty desired set. Empty records are
|
||||||
|
ignored; whitespace is otherwise payload, not trimming. EOF requests a
|
||||||
|
graceful stop. The stdin-controlled mode bounds each desired set to 128
|
||||||
|
targets and retains at most 256 live target generations while replacements
|
||||||
|
drain.
|
||||||
|
|
||||||
## FINITE VERSUS CONTINUOUS OPERATION
|
## FINITE VERSUS CONTINUOUS OPERATION
|
||||||
|
|
||||||
`--duration` controls the run length:
|
`--duration` controls the run length:
|
||||||
@@ -201,6 +235,10 @@ the client exits nonzero — so a process supervisor can restart it; an
|
|||||||
interrupted run (`Ctrl-C`) never reports peer closure as an error, whichever
|
interrupted run (`Ctrl-C`) never reports peer closure as an error, whichever
|
||||||
mode it was in.
|
mode it was in.
|
||||||
|
|
||||||
|
With `--targets-stdin`, peer closure is target-local: other desired targets
|
||||||
|
and the stdin controller continue running. Including that target in a later
|
||||||
|
complete set starts a fresh generation.
|
||||||
|
|
||||||
## EXIT STATUS
|
## EXIT STATUS
|
||||||
|
|
||||||
`0` on success (including a finite run's ordinary completion, and an
|
`0` on success (including a finite run's ordinary completion, and an
|
||||||
|
|||||||
Reference in new issue
Block a user