krabka_bench_driver/

workload.rs

//! Producer and consumer workload runners.
//!
//! This module builds N producers and M consumers as members of one group, runs
//! the warmup phase and then the measurement phase, triggers an optional
//! failover mid-measurement, and merges the per-task histograms into the public
//! `LatencyPercentiles` shape.

use {
    crate::{
        failover::{FaultConfig, FaultTask},
        hist,
        ids::{MessageCount, TimeOffsetMs, WallclockMs},
        numeric::{
            event_rate, nonnegative_i64_to_u64, saturating_u64_to_i64, saturating_u128_to_u64,
            to_f64,
        },
        payload,
        prom::PromClient,
        rate::Pacer,
        scenario::{
            Acks, BrokerSample, Disturbance, FailoverSpec, LatencyPercentiles, LoadMode, ModeTag,
            Resource, RunOutput, Sample, Scenario, Stack, Throughput, Topology,
        },
    },
    anyhow::{Context, Result},
    chrono::Utc,
    futures::{FutureExt, StreamExt, stream::FuturesUnordered},
    hdrhistogram::Histogram,
    krabka_client_consumer::{AutoOffsetReset, Consumer},
    krabka_client_core::{
        ClientFrameMax, ConnectionDispatchQueueCapacity,
        security::{ClientSecurity, KeyStore, TlsConnectorConfig, TrustStore},
    },
    krabka_client_producer::{
        DEFAULT_PRODUCER_MAX_BLOCK, DEFAULT_PRODUCER_MAX_IN_FLIGHT, Producer, ProducerError,
        ProducerRecord, RecordMetadata,
    },
    krabka_macros::RefinedNewtype,
    krabka_security::ListenerProtocol,
    krabka_units::prelude::*,
    refined_type::rule::GreaterU32,
    std::{
        future::Future,
        iter, mem,
        path::PathBuf,
        pin::Pin,
        sync::{
            Arc,
            atomic::{AtomicU8, AtomicU64, Ordering},
        },
    },
    tokio::{
        task::JoinSet,
        time::{Instant, sleep, sleep_until, timeout, timeout_at},
    },
    tracing::{info, warn},
};

/// Width of one time-series sample bucket. The driver splits the measurement
/// window into slices of the configured interval. Each producer task and
/// consumer task tallies its per-slice counts and its per-slice latency
/// histogram locally, so the hot path holds no shared locks. `run()` then merges
/// them into the `samples` series.
type AckResult = Result<RecordMetadata, ProducerError>;
type AckFuture = Pin<Box<dyn Future<Output = SettledAck> + Send>>;

/// A send that settled, with the instant it was enqueued.
struct SettledAck {
    result: AckResult,
    sent_at: Instant,
}

pub const MAX_CLIENT_REQUEST_TIMEOUT: Time = millis(2_147_483_647);
pub const DEFAULT_PRODUCER_REQUEST_TIMEOUT: Time = secs(2);
pub const DEFAULT_KRABKA_CONSUMER_REQUEST_TIMEOUT: Time = secs(5);
pub const DEFAULT_KAFKA_CONSUMER_REQUEST_TIMEOUT: Time = secs(30);
pub const DEFAULT_CONSUMER_BUILD_ATTEMPTS: u32 = 6;
pub const DEFAULT_CONSUMER_BUILD_INITIAL_BACKOFF: Time = millis(100);
pub const DEFAULT_CONSUMER_BUILD_MAX_BACKOFF: Time = secs(2);
pub const DEFAULT_CONSUMER_POLL_TIMEOUT: Time = millis(50);
pub const DEFAULT_CONSUMER_POLL_ERROR_BACKOFF: Time = millis(100);
pub const DEFAULT_PRODUCER_FINAL_DRAIN_TIMEOUT: Time = secs(10);
pub const DEFAULT_SAMPLE_INTERVAL: Time = secs(2);

/// Return the validated producer request-timeout default.
///
/// # Panics
///
/// Panics if the named default is not protocol-safe.
#[must_use]
pub const fn default_producer_request_timeout() -> Time { DEFAULT_PRODUCER_REQUEST_TIMEOUT }

/// Return the validated consumer request-timeout default for `stack`.
///
/// # Panics
///
/// Panics if the selected named default is not protocol-safe.
#[must_use]
pub const fn default_consumer_request_timeout(stack: Stack) -> Time {
    match stack {
        Stack::Krabka => DEFAULT_KRABKA_CONSUMER_REQUEST_TIMEOUT,
        Stack::Kafka => DEFAULT_KAFKA_CONSUMER_REQUEST_TIMEOUT,
    }
}

/// How many times the driver tries to build a consumer before it gives up.
#[derive(Debug, Clone, Copy, PartialEq, Eq, RefinedNewtype)]
#[refined(
    rule(GreaterU32<0>),
    string_error,
    default = DEFAULT_CONSUMER_BUILD_ATTEMPTS,
    from_str,
    display
)]
pub struct ConsumerBuildAttempts(u32);

#[must_use]
pub const fn default_consumer_build_initial_backoff() -> Time {
    DEFAULT_CONSUMER_BUILD_INITIAL_BACKOFF
}

#[must_use]
pub const fn default_consumer_build_max_backoff() -> Time { DEFAULT_CONSUMER_BUILD_MAX_BACKOFF }

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ConsumerBuildRetryPolicy {
    attempts: ConsumerBuildAttempts,
    initial_backoff: Time,
    max_backoff: Time,
}

impl ConsumerBuildRetryPolicy {
    /// Validate a complete consumer-build retry policy.
    ///
    /// # Errors
    ///
    /// Returns an error when the initial backoff exceeds the maximum.
    pub fn new(
        attempts: ConsumerBuildAttempts,
        initial_backoff: Time,
        max_backoff: Time,
    ) -> Result<Self, String> {
        if initial_backoff > max_backoff {
            return Err("consumer-build initial backoff exceeds maximum".to_owned());
        }
        Ok(Self {
            attempts,
            initial_backoff,
            max_backoff,
        })
    }

    #[must_use]
    pub const fn attempts(self) -> u32 { self.attempts.into_value() }

    #[must_use]
    pub fn initial_backoff(self) -> Time { self.initial_backoff }

    #[must_use]
    pub fn max_backoff(self) -> Time { self.max_backoff }
}

impl Default for ConsumerBuildRetryPolicy {
    fn default() -> Self {
        Self {
            attempts: ConsumerBuildAttempts::default(),
            initial_backoff: default_consumer_build_initial_backoff(),
            max_backoff: default_consumer_build_max_backoff(),
        }
    }
}

#[must_use]
pub const fn default_consumer_poll_timeout() -> Time { DEFAULT_CONSUMER_POLL_TIMEOUT }

#[must_use]
pub const fn default_consumer_poll_error_backoff() -> Time { DEFAULT_CONSUMER_POLL_ERROR_BACKOFF }

/// The fixed sampling grid. Every task gets a copy, so all tasks bucket into
/// the same slices.
#[derive(Clone, Copy)]
struct Grid {
    /// When the measurement window starts, which is the end of warmup.
    meas_start: Instant,
    /// The same instant as an epoch nanosecond, and the close of the window.
    ///
    /// A consumer classifies a record by the epoch stamp the producer wrote
    /// into it, not by the monotonic clock at poll time, so it needs the window
    /// on the same scale the payload carries.
    meas_start_ns: u64,
    meas_end_ns: u64,
    interval_ms: u64,
    /// Number of slices covering the measurement window.
    n: usize,
}

impl Grid {
    /// The grid that covers `scenario`'s measurement window. That window opens
    /// after its warmup has elapsed from `started_at`. The grid is always at
    /// least one slice wide, so a run shorter than a slice still has somewhere
    /// to tally into.
    ///
    /// `started_ns` is the epoch nanosecond that `started_at` was read at. The
    /// two are one instant on two clocks.
    fn new(started_at: Instant, started_ns: u64, scenario: &Scenario, interval: Time) -> Self {
        let interval_ms = nonnegative_i64_to_u64(interval.millis_i64());
        let duration_ms = nonnegative_i64_to_u64(scenario.duration.millis_i64());
        let warmup = nonnegative_i64_to_u64(scenario.warmup.nanos_i64());
        let width = nonnegative_i64_to_u64(scenario.duration.nanos_i64());
        let meas_start_ns = started_ns.saturating_add(warmup);
        Self {
            meas_start: started_at + scenario.warmup.to_std(),
            meas_start_ns,
            meas_end_ns: meas_start_ns.saturating_add(width),
            interval_ms,
            n: usize::try_from(duration_ms.div_ceil(interval_ms).max(1)).unwrap_or(usize::MAX),
        }
    }

    /// Whether a record sent at epoch nanosecond `send_ns` belongs to the
    /// measurement window.
    ///
    /// A saturating consumer can still be draining warmup backlog when the
    /// window opens. Counting those records would report warmup throughput and
    /// warmup-aged latency as measurement, which is what the scenario
    /// discarded the warmup to avoid.
    fn sent_in_measurement(&self, send_ns: u64) -> bool {
        send_ns >= self.meas_start_ns && send_ns < self.meas_end_ns
    }

    /// Slice index for a record sent at epoch nanosecond `send_ns`, clamped
    /// into `[0, n-1]`. Callers check [`Self::sent_in_measurement`] first.
    fn idx_for_send_ns(&self, send_ns: u64) -> usize {
        let elapsed_ms = send_ns.saturating_sub(self.meas_start_ns) / 1_000_000;
        self.idx_for_elapsed_ms(elapsed_ms)
    }

    /// Slice index for an event observed at `now`, clamped into `[0, n-1]`.
    fn idx(&self, now: Instant) -> usize {
        let elapsed_ms =
            saturating_u128_to_u64(now.saturating_duration_since(self.meas_start).as_millis());
        self.idx_for_elapsed_ms(elapsed_ms)
    }

    fn idx_for_elapsed_ms(&self, elapsed_ms: u64) -> usize {
        usize::try_from(elapsed_ms / self.interval_ms)
            .unwrap_or(usize::MAX)
            .min(self.n.saturating_sub(1))
    }

    /// The width of one slice, as the extent the per-slice rates divide by.
    fn interval(&self) -> Time { Time::from_millis(saturating_u64_to_i64(self.interval_ms)) }
}

/// Parameters wired in from `main` and the CLI. This struct is distinct from
/// `Scenario` because the scenario YAML describes *what* to run, and this struct
/// describes *where* it runs.
pub struct DriverConfig {
    pub bootstrap: String,
    pub client_dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
    pub client_frame_max: ClientFrameMax,
    pub topic: String,
    pub stack: Stack,
    pub namespace: String,
    pub prometheus_url: Option<String>,
    pub prometheus_request_timeout: Time,
    pub producer_request_timeout: Time,
    pub producer_final_drain_timeout: Time,
    pub sample_interval: Time,
    pub consumer_request_timeout: Time,
    pub consumer_build_retry_policy: ConsumerBuildRetryPolicy,
    /// How long one consumer `poll` waits for records.
    pub consumer_poll_timeout: Time,
    /// How long the consumer loop sleeps after a poll error before retrying.
    pub consumer_poll_error_backoff: Time,
    pub broker_count: u32,
    pub scenario_id: u64,
    /// TLS data-path config. `None` means plaintext, the default benchmark path.
    pub tls: Option<TlsParams>,
}

/// TLS knobs for the producer and consumer data path.
///
/// When this is present, both clients dial the broker's TLS listener and the
/// produce and fetch byte stream is encrypted. Krabka can serve that path with
/// kTLS sendfile, and Strimzi serves it with JVM SSL.
#[derive(Debug, Clone)]
pub struct TlsParams {
    /// Mounted CA bundle (`ca.crt`) that the client trusts to verify the broker
    /// serving cert. It comes from the per-stack cluster-CA Secret.
    pub ca_path: PathBuf,
    /// SNI server name that the client presents in the TLS `ClientHello`. The
    /// broker matches it against a SAN on its serving cert. LOAD-BEARING: the
    /// client DNS-resolves the bootstrap to a pod IP and dials by IP, so the SNI
    /// is NOT derived from the bootstrap host. You must set it to a cert-SAN
    /// name. For krabka that name is
    /// `demo-broker-headless.<ns>.svc.cluster.local`, and for Strimzi it is
    /// `demo-kafka-bootstrap`.
    pub server_name: String,
    /// Optional mTLS client identity. One-way TLS with `None` is enough for the
    /// benchmark. Set this only if the listener requires client auth.
    pub client_identity: Option<ClientIdentity>,
}

/// The PEM files of an mTLS client identity.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientIdentity {
    /// Path of the PEM certificate chain.
    pub cert_pem: PathBuf,
    /// Path of the PEM private key.
    pub key_pem: PathBuf,
}

impl TlsParams {
    /// Builds the [`ClientSecurity`] for an `Ssl` listener. The default is
    /// one-way TLS, which authenticates the server against the trust roots from
    /// the mounted CA. It becomes mutual TLS when [`Self::client_identity`] is
    /// set.
    #[must_use]
    pub fn to_security(&self) -> ClientSecurity {
        let mut tls = TlsConnectorConfig::default();
        tls.trust_store = TrustStore::PemFile(self.ca_path.clone());
        tls.server_name.clone_from(&self.server_name);
        tls.key_store = self
            .client_identity
            .clone()
            .map(|identity| KeyStore::PemFiles {
                certificate_chain: identity.cert_pem,
                private_key: identity.key_pem,
                key_password: None,
            });
        ClientSecurity {
            protocol: ListenerProtocol::Ssl,
            tls: Some(tls),
            sasl: None,
            sasl_host: None,
        }
    }
}

/// Top-level entry point that `main` calls. Returns the populated
/// `RunOutput`. The caller must serialise it to disk.
/// # Errors
/// Returns an error when input data is invalid, required I/O fails, or the destination rejects the generated report or audit event.
/// # Panics
/// Panics if synchronized state is poisoned or validated input is missing a field required to produce the output.
pub async fn run(scenario: Scenario, cfg: DriverConfig) -> Result<RunOutput> {
    // One reading of the wall clock, kept at both scales. The nanosecond form
    // is what a consumer compares a record's embedded send stamp against.
    let started_ns = nonnegative_i64_to_u64(Utc::now().timestamp_nanos_opt().unwrap_or_default());
    let wallclock_start = saturating_u64_to_i64(started_ns / 1_000_000);
    let t_start = Instant::now();

    // Time-series sampling grid: split the measurement window into fixed
    // slices. Computed up front (meas_start is t_start + warmup) so it can be
    // handed to each task at spawn time.
    let grid = Grid::new(t_start, started_ns, &scenario, cfg.sample_interval);
    // Every histogram of the run is built empty from this one, so they share
    // its bounds and merge without error.
    let histogram = hist::new().context("build the latency histogram")?;
    let n_intervals = grid.n;

    let mut notes: Vec<String> = Vec::new();
    let mut errors: Vec<String> = Vec::new();

    let TopologyCheck {
        failover_active,
        skipped,
    } = validate_topology(&scenario, &cfg, wallclock_start, &mut notes);
    if let Some(output) = skipped {
        return Ok(output);
    }

    let security = client_security(&cfg);

    let mut prod_set: JoinSet<ProducerOut> = JoinSet::new();
    let stop = Arc::new(AtomicU8::new(STATE_RUN));
    let first_ack_unix_ms = Arc::new(AtomicU64::new(0));
    // The one marker every producer reads to learn that the broker is gone.
    // It stays zero until the delete succeeds, so it is created before the
    // producers start and not beside the fault task.
    let kill_at_ms = Arc::new(AtomicU64::new(0));

    spawn_producers(
        &mut prod_set,
        &scenario,
        &cfg,
        security.as_ref(),
        grid,
        &SharedHandles {
            stop: &stop,
            first_ack: &first_ack_unix_ms,
            kill_at: &kill_at_ms,
        },
        &histogram,
    );

    let mut cons_set: JoinSet<ConsumerOut> = JoinSet::new();
    spawn_consumers(
        &mut cons_set,
        &scenario,
        &cfg,
        security.as_ref(),
        grid,
        &stop,
        &histogram,
    );

    let fault = scenario
        .failover
        .as_ref()
        .filter(|_| failover_active)
        .map(|spec| {
            start_fault(
                spec,
                &cfg,
                security.clone(),
                t_start,
                Arc::clone(&kill_at_ms),
            )
        });

    let warmup_end = t_start + scenario.warmup.to_std();
    sleep_until(warmup_end).await;
    stop.store(STATE_MEASURING, Ordering::SeqCst);

    let meas_end = warmup_end + scenario.duration.to_std();
    sleep_until(meas_end).await;
    stop.store(STATE_STOP, Ordering::SeqCst);

    let fault_at_ms = finish_fault(fault, &mut notes).await;

    // ── Join all tasks ──────────────────────────────────────────────────────
    let mut prod_hist = Histogram::new_from(&histogram);
    let mut prod_msgs = 0u64;
    let mut prod_bytes = 0u64;
    let mut prod_dropped = 0u64;
    let mut earliest_recovery_ms = 0u64;
    let mut max_spike = Time::ZERO;
    let mut prod_iv_msgs = vec![0u64; n_intervals];
    let mut prod_iv_hist: Vec<Histogram<u64>> =
        iter::repeat_with(|| Histogram::new_from(&histogram))
            .take(n_intervals)
            .collect();
    while let Some(j) = prod_set.join_next().await {
        match j {
            Ok(t) => {
                let _ = prod_hist.add(&t.latency);
                prod_msgs += t.msgs;
                prod_bytes += t.bytes;
                prod_dropped += t.dropped;
                for (total, m) in prod_iv_msgs.iter_mut().zip(&t.interval_msgs) {
                    *total += *m;
                }
                for (total, h) in prod_iv_hist.iter_mut().zip(&t.interval_hist) {
                    let _ = total.add(h);
                }
                if t.latency_spike_max > max_spike {
                    max_spike = t.latency_spike_max;
                }
                if t.recovery_unix_ms > 0
                    && (earliest_recovery_ms == 0 || t.recovery_unix_ms < earliest_recovery_ms)
                {
                    earliest_recovery_ms = t.recovery_unix_ms;
                }
                if !t.error.is_empty() {
                    errors.push(t.error);
                }
            }
            Err(e) => errors.push(format!("producer-task-panic: {e}")),
        }
    }

    let consumed = join_consumers(cons_set, n_intervals, &histogram, &mut errors).await;

    let samples = build_samples(
        grid.interval(),
        &IntervalTallies {
            messages: &prod_iv_msgs,
            histograms: &prod_iv_hist,
        },
        &IntervalTallies {
            messages: &consumed.interval_msgs,
            histograms: &consumed.interval_hist,
        },
    );

    let wallclock_end = Utc::now().timestamp_millis();
    let ResourceCapture {
        resource,
        broker_samples,
    } = capture_resources(
        &scenario,
        &cfg,
        WallclockWindow {
            start_ms: wallclock_start,
            end_ms: wallclock_end,
        },
        grid.meas_end_ns,
        prod_msgs,
        &mut notes,
    )
    .await;

    let RunTiming {
        disturbance,
        first_ack,
    } = finalize_timing(
        failover_active,
        FailoverStamps {
            kill_at_ms: fault_at_ms,
            recovery_at_ms: earliest_recovery_ms,
            dropped: prod_dropped,
            latency_spike_max: max_spike,
        },
        first_ack_unix_ms.load(Ordering::SeqCst),
        wallclock_start,
    );
    note_missing_disturbance(
        failover_active,
        disturbance.as_ref(),
        fault_at_ms,
        &mut notes,
    );

    Ok(RunOutput {
        scenario: scenario.clone(),
        stack: cfg.stack,
        topology: output_topology(&scenario, &cfg),
        wallclock_start_unix_ms: WallclockMs(wallclock_start),
        wallclock_end_unix_ms: WallclockMs(wallclock_end),
        throughput: Throughput {
            msgs_produced: MessageCount(prod_msgs),
            msgs_consumed: MessageCount(consumed.msgs),
            bytes_in: ByteSize::from_bytes(prod_bytes),
            bytes_out: ByteSize::from_bytes(consumed.bytes),
            producer_rate: event_rate(prod_msgs, scenario.duration),
            consumer_rate: event_rate(consumed.msgs, scenario.duration),
        },
        producer_latency: hist::percentiles(&prod_hist),
        consumer_e2e_latency: hist::percentiles(&consumed.latency),
        resource,
        disturbance,
        startup: None,
        first_ack,
        errors,
        notes,
        samples,
        broker_samples,
    })
}

async fn finish_fault(fault: Option<FaultTask>, notes: &mut Vec<String>) -> u64 {
    let Some(fault) = fault else {
        return 0;
    };
    let outcome = fault.stop().await;
    if let Some(note) = outcome.note() {
        notes.push(note);
    }
    outcome.killed_at_ms().unwrap_or_default()
}

fn start_fault(
    spec: &FailoverSpec,
    cfg: &DriverConfig,
    security: Option<ClientSecurity>,
    started_at: Instant,
    kill_at: Arc<AtomicU64>,
) -> FaultTask {
    FaultTask::spawn(
        spec,
        FaultConfig {
            stack: cfg.stack,
            namespace: cfg.namespace.clone(),
            bootstrap: cfg.bootstrap.clone(),
            topic: cfg.topic.clone(),
            security,
            dispatch_queue_capacity: cfg.client_dispatch_queue_capacity,
            frame_max: cfg.client_frame_max,
        },
        started_at,
        kill_at,
    )
}

fn client_security(cfg: &DriverConfig) -> Option<ClientSecurity> {
    let security = cfg.tls.as_ref().map(TlsParams::to_security);
    if let Some(value) = &security {
        info!(
            server_name = value
                .tls
                .as_ref()
                .map_or("", |tls| tls.server_name.as_str()),
            "TLS data path enabled (protocol=Ssl)"
        );
    }
    security
}

/// The per-interval message counts and latency histograms that one side of the
/// run (producers or consumers) tallied.
struct IntervalTallies<'a> {
    messages: &'a [u64],
    histograms: &'a [Histogram<u64>],
}

fn build_samples(
    interval: Time,
    producer: &IntervalTallies<'_>,
    consumer: &IntervalTallies<'_>,
) -> Vec<Sample> {
    let interval_ms = nonnegative_i64_to_u64(interval.millis_i64());
    // The histogram counts microseconds; an empty slice measured nothing.
    let percentile = |histogram: &Histogram<u64>, quantile: f64| {
        if histogram.is_empty() {
            Time::ZERO
        } else {
            Time::from_micros(saturating_u64_to_i64(histogram.value_at_quantile(quantile)))
        }
    };
    producer
        .messages
        .iter()
        .zip(producer.histograms)
        .zip(consumer.messages.iter().zip(consumer.histograms))
        .enumerate()
        .map(
            |(index, ((messages, producer_histogram), (consumer_messages, consumer_histogram)))| {
                Sample {
                    t_offset_ms: TimeOffsetMs(
                        u64::try_from(index)
                            .unwrap_or(u64::MAX)
                            .saturating_mul(interval_ms),
                    ),
                    producer_rate: event_rate(*messages, interval),
                    consumer_rate: event_rate(*consumer_messages, interval),
                    producer_p50: percentile(producer_histogram, 0.50),
                    producer_p99: percentile(producer_histogram, 0.99),
                    consumer_e2e_p99: percentile(consumer_histogram, 0.99),
                }
            },
        )
        .collect()
}

/// The run-phase signal, the first-ack stamp, and the kill marker that every
/// producer task shares.
#[derive(Clone, Copy)]
struct SharedHandles<'a> {
    stop: &'a Arc<AtomicU8>,
    first_ack: &'a Arc<AtomicU64>,
    kill_at: &'a Arc<AtomicU64>,
}

/// Spawn one producer task per `scenario.producers`.
///
/// The three shared handles travel together because they are one thing: the
/// run-phase signal, the first-ack stamp, and the kill marker every task reads.
fn spawn_producers(
    tasks: &mut JoinSet<ProducerOut>,
    scenario: &Scenario,
    cfg: &DriverConfig,
    security: Option<&ClientSecurity>,
    grid: Grid,
    shared: &SharedHandles<'_>,
    histogram: &Histogram<u64>,
) {
    let SharedHandles {
        stop,
        first_ack,
        kill_at,
    } = *shared;
    for idx in 0..scenario.producers {
        tasks.spawn(run_producer(ProducerTask {
            setup: TaskSetup::new(
                idx,
                scenario,
                cfg,
                security,
                grid,
                stop,
                cfg.producer_request_timeout,
            ),
            first_ack: Arc::clone(first_ack),
            kill_at: Arc::clone(kill_at),
            final_drain_timeout: cfg.producer_final_drain_timeout,
            histogram: Histogram::new_from(histogram),
        }));
    }
}

/// Wait for every consumer task, merging what each measured. A task that
/// failed or panicked leaves its error in `errors`.
async fn join_consumers(
    mut tasks: JoinSet<ConsumerOut>,
    intervals: usize,
    histogram: &Histogram<u64>,
    errors: &mut Vec<String>,
) -> TaskMeasurements {
    let mut merged = TaskMeasurements::new(intervals, histogram);
    while let Some(joined) = tasks.join_next().await {
        match joined {
            Ok(task) => {
                let _ = merged.latency.add(&task.latency);
                merged.msgs += task.msgs;
                merged.bytes += task.bytes;
                for (total, m) in merged.interval_msgs.iter_mut().zip(&task.interval_msgs) {
                    *total += *m;
                }
                for (total, h) in merged.interval_hist.iter_mut().zip(&task.interval_hist) {
                    let _ = total.add(h);
                }
                if !task.error.is_empty() {
                    errors.push(task.error);
                }
            }
            Err(e) => errors.push(format!("consumer-task-panic: {e}")),
        }
    }
    merged
}

/// Spawn one consumer task per `scenario.consumers`.
fn spawn_consumers(
    tasks: &mut JoinSet<ConsumerOut>,
    scenario: &Scenario,
    cfg: &DriverConfig,
    security: Option<&ClientSecurity>,
    grid: Grid,
    stop: &Arc<AtomicU8>,
    histogram: &Histogram<u64>,
) {
    for idx in 0..scenario.consumers {
        tasks.spawn(run_consumer(ConsumerTask {
            setup: TaskSetup::new(
                idx,
                scenario,
                cfg,
                security,
                grid,
                stop,
                cfg.consumer_request_timeout,
            ),
            build_retry_policy: cfg.consumer_build_retry_policy,
            poll_timeout: cfg.consumer_poll_timeout,
            poll_error_backoff: cfg.consumer_poll_error_backoff,
            histogram: Histogram::new_from(histogram),
        }));
    }
}

fn note_missing_disturbance(
    failover_active: bool,
    disturbance: Option<&Disturbance>,
    kill_at_ms: u64,
    notes: &mut Vec<String>,
) {
    if !failover_active || disturbance.is_some() {
        return;
    }
    notes.push(if kill_at_ms == 0 {
        "failover:no-broker-deleted".to_owned()
    } else {
        "failover:no-ack-after-kill".to_owned()
    });
}

/// What the producers measured around a failover kill.
#[derive(Debug, Clone, Copy)]
struct FailoverStamps {
    /// Epoch milliseconds of the pod delete, or zero when it never landed.
    kill_at_ms: u64,
    /// Epoch milliseconds of the first ack at or after the kill, or zero.
    recovery_at_ms: u64,
    /// Sends the producers counted as dropped.
    dropped: u64,
    /// The largest ack latency seen after the kill.
    latency_spike_max: Time,
}

/// The disturbance a run measured and the extent from run start to the first
/// ack.
struct RunTiming {
    disturbance: Option<Disturbance>,
    first_ack: Time,
}

/// The disturbance a failover run measured, and the run's first-ack extent.
///
/// A disturbance is written only when the delete actually landed and a record
/// was acknowledged after it. `failover.kill_at_ms` is the shared kill marker,
/// which stays zero when the pod delete failed; `failover.recovery_at_ms` is
/// the first ack at or after that marker. Either one at zero means the run has
/// no failover evidence, and an empty disturbance is the honest answer. The
/// report's gate reads that as a violation rather than as an instant recovery.
fn finalize_timing(
    failover_active: bool,
    failover: FailoverStamps,
    first_ack: u64,
    wallclock_start: i64,
) -> RunTiming {
    let measured_kill = failover.kill_at_ms > 0 && failover.recovery_at_ms >= failover.kill_at_ms;
    let disturbance = (failover_active && measured_kill).then_some(Disturbance {
        kill_at_ms: TimeOffsetMs(failover.kill_at_ms),
        recovery_at_ms: TimeOffsetMs(failover.recovery_at_ms),
        dropped: MessageCount(failover.dropped),
        latency_spike_max: failover.latency_spike_max,
    });
    // Both stamps are epoch milliseconds — coordinates — so their difference is
    // the extent from run start to the first ack. No ack at all means no extent.
    let first_ack = if first_ack == 0 {
        Time::ZERO
    } else {
        Time::from_millis(
            i64::try_from(first_ack)
                .unwrap_or(i64::MAX)
                .saturating_sub(wallclock_start)
                .max(0),
        )
    };
    RunTiming {
        disturbance,
        first_ack,
    }
}

/// Whether failover is active for the run, and the early output of a run that
/// the topology cannot host.
struct TopologyCheck {
    failover_active: bool,
    skipped: Option<RunOutput>,
}

fn validate_topology(
    scenario: &Scenario,
    cfg: &DriverConfig,
    wallclock_start: i64,
    notes: &mut Vec<String>,
) -> TopologyCheck {
    let failover_active =
        scenario.failover.is_some() && scenario.replication_factor >= 3 && cfg.broker_count >= 3;
    if scenario.failover.is_some() && !failover_active {
        notes.push("skipped:failover-needs-rf3".into());
    }
    let topology_mismatch = matches!(scenario.mode_tag, ModeTag::Cluster)
        && scenario.replication_factor >= 3
        && cfg.broker_count < 3;
    if !topology_mismatch {
        return TopologyCheck {
            failover_active,
            skipped: None,
        };
    }
    notes.push(format!(
        "skipped:topology-mismatch (rf={} brokers={})",
        scenario.replication_factor, cfg.broker_count
    ));
    let output = empty_output(scenario, cfg, wallclock_start, mem::take(notes), Vec::new());
    TopologyCheck {
        failover_active,
        skipped: Some(output),
    }
}

/// Capture broker CPU and memory for the window the run measured.
///
/// `measured_until_ns` is the epoch nanosecond at which the measurement window
/// closed, which is not the same instant as this call: joining the tasks and
/// draining the producers happens in between, and a failover run can wait out
/// its final-drain timeout there. The queries are evaluated at the close of the
/// window, so the lookback covers what was measured.
/// The wall-clock window of a run, in epoch milliseconds.
#[derive(Debug, Clone, Copy)]
struct WallclockWindow {
    start_ms: i64,
    end_ms: i64,
}

/// The resource summary of a run and its broker resource time series.
struct ResourceCapture {
    resource: Resource,
    broker_samples: Vec<BrokerSample>,
}

async fn capture_resources(
    scenario: &Scenario,
    cfg: &DriverConfig,
    wallclock: WallclockWindow,
    measured_until_ns: u64,
    produced: u64,
    notes: &mut Vec<String>,
) -> ResourceCapture {
    let Some(url) = &cfg.prometheus_url else {
        notes.push("prometheus-url-not-set".into());
        return ResourceCapture {
            resource: Resource::default(),
            broker_samples: Vec::new(),
        };
    };
    let client = match PromClient::new(url, cfg.prometheus_request_timeout) {
        Ok(client) => client,
        Err(error) => {
            notes.push(format!("prometheus-client-failed: {error}"));
            return ResourceCapture {
                resource: Resource::default(),
                broker_samples: Vec::new(),
            };
        }
    };
    let resource = match client
        .capture_resource(
            cfg.stack,
            &cfg.namespace,
            scenario.duration,
            to_f64(measured_until_ns) / 1.0e9,
            MessageCount(produced),
        )
        .await
    {
        Ok(resource) => resource,
        Err(error) => {
            warn!(%error, "prometheus capture failed");
            notes.push(format!("prometheus-capture-failed: {error}"));
            Resource::default()
        }
    };
    let broker_samples = match client
        .capture_resource_series(
            cfg.stack,
            &cfg.namespace,
            to_f64(wallclock.start_ms) / 1000.0,
            to_f64(wallclock.end_ms) / 1000.0,
            15,
        )
        .await
    {
        Ok(samples) => samples,
        Err(error) => {
            warn!(%error, "prometheus resource time-series capture failed");
            notes.push(format!("prometheus-resource-series-failed: {error}"));
            Vec::new()
        }
    };
    ResourceCapture {
        resource,
        broker_samples,
    }
}

const STATE_RUN: u8 = 0; // warmup phase, record-but-discard
const STATE_MEASURING: u8 = 1;
const STATE_STOP: u8 = 2;

struct ProducerOut {
    latency: Histogram<u64>,
    msgs: u64,
    bytes: u64,
    dropped: u64,
    recovery_unix_ms: u64,
    latency_spike_max: Time,
    error: String,
    /// Measurement-window tallies, one per slice of the [`Grid`]. This is empty
    /// on a build failure. `run()` merges by index with bounds checks.
    interval_msgs: Vec<u64>,
    interval_hist: Vec<Histogram<u64>>,
}

struct ConsumerOut {
    latency: Histogram<u64>,
    msgs: u64,
    bytes: u64,
    error: String,
    interval_msgs: Vec<u64>,
    interval_hist: Vec<Histogram<u64>>,
}

fn empty_output(
    scenario: &Scenario,
    cfg: &DriverConfig,
    start: i64,
    notes: Vec<String>,
    errors: Vec<String>,
) -> RunOutput {
    RunOutput {
        scenario: scenario.clone(),
        stack: cfg.stack,
        topology: output_topology(scenario, cfg),
        wallclock_start_unix_ms: WallclockMs(start),
        wallclock_end_unix_ms: WallclockMs(start),
        throughput: Throughput::default(),
        producer_latency: LatencyPercentiles::default(),
        consumer_e2e_latency: LatencyPercentiles::default(),
        resource: Resource::default(),
        disturbance: None,
        startup: None,
        first_ack: Time::ZERO,
        errors,
        notes,
        samples: Vec::new(),
        broker_samples: Vec::new(),
    }
}

fn output_topology(scenario: &Scenario, cfg: &DriverConfig) -> Topology {
    Topology {
        partitions: scenario.partitions,
        replication_factor: scenario.replication_factor,
        broker_count: cfg.broker_count,
    }
}

// ── Producer task ───────────────────────────────────────────────────────────

struct TaskSetup {
    idx: usize,
    scenario: Scenario,
    bootstrap: String,
    topic: String,
    scenario_id: u64,
    stop: Arc<AtomicU8>,
    security: Option<ClientSecurity>,
    grid: Grid,
    request_timeout: Time,
    dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
    frame_max: ClientFrameMax,
}

impl TaskSetup {
    fn new(
        idx: usize,
        scenario: &Scenario,
        cfg: &DriverConfig,
        security: Option<&ClientSecurity>,
        grid: Grid,
        stop: &Arc<AtomicU8>,
        request_timeout: Time,
    ) -> Self {
        Self {
            idx,
            scenario: scenario.clone(),
            bootstrap: cfg.bootstrap.clone(),
            topic: cfg.topic.clone(),
            scenario_id: cfg.scenario_id,
            stop: Arc::clone(stop),
            security: security.cloned(),
            grid,
            request_timeout,
            dispatch_queue_capacity: cfg.client_dispatch_queue_capacity,
            frame_max: cfg.client_frame_max,
        }
    }
}

struct ProducerTask {
    setup: TaskSetup,
    first_ack: Arc<AtomicU64>,
    /// Epoch millisecond at which the failover task deleted the broker pod, or
    /// zero while no delete has succeeded. This is the only thing that says a
    /// kill happened. A producer error says the send failed and nothing more.
    kill_at: Arc<AtomicU64>,
    final_drain_timeout: Time,
    /// The empty histogram every latency histogram of the task is built from.
    histogram: Histogram<u64>,
}

struct TaskMeasurements {
    latency: Histogram<u64>,
    interval_msgs: Vec<u64>,
    interval_hist: Vec<Histogram<u64>>,
    msgs: u64,
    bytes: u64,
}

impl TaskMeasurements {
    /// Empty measurements over `intervals` slices, each histogram built empty
    /// from `histogram`.
    fn new(intervals: usize, histogram: &Histogram<u64>) -> Self {
        Self {
            latency: Histogram::new_from(histogram),
            interval_msgs: vec![0; intervals],
            interval_hist: iter::repeat_with(|| Histogram::new_from(histogram))
                .take(intervals)
                .collect(),
            msgs: 0,
            bytes: 0,
        }
    }
}

/// One producer's acknowledgement tallies and the run controls they observe.
struct ProducerObservations<'a> {
    idx: usize,
    msg_size: u64,
    grid: Grid,
    shared: SharedHandles<'a>,
    output: ProducerOut,
}

impl<'a> ProducerObservations<'a> {
    fn new(
        idx: usize,
        msg_size: u64,
        grid: Grid,
        shared: SharedHandles<'a>,
        histogram: &Histogram<u64>,
    ) -> Self {
        let measurements = TaskMeasurements::new(grid.n, histogram);
        Self {
            idx,
            msg_size,
            grid,
            shared,
            output: ProducerOut {
                latency: measurements.latency,
                msgs: measurements.msgs,
                bytes: measurements.bytes,
                dropped: 0,
                recovery_unix_ms: 0,
                latency_spike_max: Time::ZERO,
                error: String::new(),
                interval_msgs: measurements.interval_msgs,
                interval_hist: measurements.interval_hist,
            },
        }
    }

    /// Record a send that `enqueue` held past the producer's own blocking
    /// limit. The record counts as dropped, and the first error stays.
    fn record_enqueue_timeout(&mut self) {
        self.output.dropped += 1;
        if self.output.error.is_empty() {
            self.output.error = format!("producer-{}-enqueue-timeout", self.idx);
        }
    }

    /// Record a shutdown step (`flush` or `close`) that outlived its limit.
    /// The first error a producer hit stays the one it reports.
    fn record_teardown_timeout(&mut self, step: &str) {
        if self.output.error.is_empty() {
            self.output.error = format!("producer-{}-{step}-timeout", self.idx);
        }
    }

    /// Account for every send still outstanding at the final drain deadline.
    fn record_drain_timeout(&mut self, unresolved: usize) {
        self.output.dropped += unresolved as u64;
        if self.output.error.is_empty() {
            self.output.error = format!("producer-{}-final-drain-timeout:{unresolved}", self.idx);
        }
    }

    /// Record one settled send in the phase observed at receipt.
    fn record_ack(&mut self, result: AckResult, sent_at: Instant) {
        match result {
            Ok(_meta) => {
                let latency = sent_at.elapsed().as_time();
                if self.shared.stop.load(Ordering::Relaxed) == STATE_MEASURING {
                    hist::record(&mut self.output.latency, latency);
                    self.output.msgs += 1;
                    self.output.bytes += self.msg_size;
                    let iv = self.grid.idx(Instant::now());
                    if let Some(count) = self.output.interval_msgs.get_mut(iv) {
                        *count += 1;
                    }
                    if let Some(histogram) = self.output.interval_hist.get_mut(iv) {
                        hist::record(histogram, latency);
                    }
                    // Recovery and the spike are measured against the
                    // shared kill marker, not against this task's own
                    // errors. The marker is published only after the pod
                    // delete succeeds, so an ack before it belongs to the
                    // healthy cluster and is not a recovery.
                    let killed_at = self.shared.kill_at.load(Ordering::Relaxed);
                    if killed_at > 0 {
                        let now_ms = nonnegative_i64_to_u64(Utc::now().timestamp_millis());
                        if self.output.recovery_unix_ms == 0 && now_ms >= killed_at {
                            self.output.recovery_unix_ms = now_ms;
                        }
                        if latency > self.output.latency_spike_max {
                            self.output.latency_spike_max = latency;
                        }
                    }
                }
                if self.shared.first_ack.load(Ordering::Relaxed) == 0 {
                    let now_ms = nonnegative_i64_to_u64(Utc::now().timestamp_millis());
                    let _ = self.shared.first_ack.compare_exchange(
                        0,
                        now_ms,
                        Ordering::SeqCst,
                        Ordering::Relaxed,
                    );
                }
            }
            Err(e) => {
                if self.shared.stop.load(Ordering::Relaxed) == STATE_MEASURING {
                    self.output.dropped += 1;
                }
                if self.output.dropped == 1 && self.output.error.is_empty() {
                    self.output.error = format!("producer-{}-err: {e}", self.idx);
                }
            }
        }
    }
}

async fn run_producer(task: ProducerTask) -> ProducerOut {
    let ProducerTask {
        setup:
            TaskSetup {
                idx,
                scenario,
                bootstrap,
                topic,
                scenario_id,
                stop,
                security,
                grid,
                request_timeout,
                dispatch_queue_capacity,
                frame_max,
            },
        first_ack,
        kill_at,
        final_drain_timeout,
        histogram,
    } = task;
    // Idempotence forces acks=All; turn it off whenever the scenario
    // requested something weaker.
    let enable_idempotence = matches!(scenario.acks, Acks::All);
    match Producer::builder()
        .bootstrap(bootstrap.clone())
        .dispatch_queue_capacity(dispatch_queue_capacity.get())
        .frame_max(frame_max.size())
        .client_id(format!("bench-producer-{idx}"))
        .acks(scenario.acks.into_producer())
        .compression(scenario.compression.into_producer())
        .enable_idempotence(enable_idempotence)
        .linger(scenario.linger.to_std())
        .batch_size(scenario.batch_size.bytes_usize())
        .max_in_flight_per_connection(producer_max_in_flight(enable_idempotence))
        // A produce to a healthy broker completes in low single-digit ms, so a
        // 2s ceiling should not trip under normal load — but it bounds how long a
        // send to a *killed* leader blocks before the client re-routes, so the
        // measured failover recovery reflects the cluster's leader re-election
        // rather than the 30s default request timeout.
        .request_timeout(request_timeout.to_std())
        // `None` → plaintext (default). `Some` → all produce traffic for this
        // task goes over the broker's TLS listener.
        .maybe_security(security)
        .build()
        .await
        .context("build producer")
    {
        Err(e) => ProducerOut {
            latency: histogram,
            msgs: 0,
            bytes: 0,
            dropped: 0,
            recovery_unix_ms: 0,
            latency_spike_max: Time::ZERO,
            error: format!("producer-{idx}-build: {e:#}"),
            interval_msgs: Vec::new(),
            interval_hist: Vec::new(),
        },
        Ok(producer) => {
            let mut tmpl = payload::template(scenario.msg_size);
            // Counter behind each record key, so keys vary within this task. The task
            // index is folded in, so two tasks never mint the same key.
            let mut sent_records = (idx as u64) << 40;
            let mut observations = ProducerObservations::new(
                idx,
                scenario.msg_size.bytes_u64(),
                grid,
                SharedHandles {
                    stop: &stop,
                    first_ack: &first_ack,
                    kill_at: &kill_at,
                },
                &histogram,
            );

            let mut pacer = match scenario.mode {
                LoadMode::Saturate => None,
                // The scenario rate is for the whole run, so each task paces at its
                // share of it. `Pacer` floors the result at one message a second.
                LoadMode::FixedRate { rate } => {
                    Some(Pacer::new(rate / to_f64(scenario.producers.max(1))))
                }
            };

            // Pipeline depth: how many records may be in flight (awaiting their ack)
            // before the loop applies back-pressure. Without pipelining the loop sent
            // one record and awaited its ack before the next, capping a task's
            // throughput at 1/RTT regardless of cluster capacity — so the numbers
            // measured the driver, not the cluster. A bounded window lets throughput
            // track the cluster while keeping memory in hand.
            let max_inflight: usize = 512;
            let mut inflight: FuturesUnordered<AckFuture> = FuturesUnordered::new();
            let enqueue_limit = enqueue_limit(request_timeout);

            loop {
                let state = stop.load(Ordering::Relaxed);
                if state == STATE_STOP {
                    break;
                }
                // Drain every ack that has already settled — promptly, so the recorded
                // latency is the real send→ack time, not how long a record waited in
                // the in-flight set behind an older stuck send.
                while let Some(SettledAck {
                    result: res,
                    sent_at: t0,
                }) = inflight.next().now_or_never().flatten()
                {
                    observations.record_ack(res, t0);
                }
                // Back-pressure: at capacity, block until any send completes. Waiting
                // for the oldest can make one dead-leader send hide live-partition acks.
                while inflight.len() >= max_inflight {
                    if let Some(SettledAck {
                        result: res,
                        sent_at: t0,
                    }) = inflight.next().await
                    {
                        observations.record_ack(res, t0);
                    }
                }
                if let Some(p) = pacer.as_mut() {
                    p.await_token().await;
                }
                let value = payload::stamp_into(&mut tmpl, scenario_id);
                // A scenario that sets `key_size` gets keyed records. The key changes
                // per record, so the partitioner spreads them the way a real keyed
                // workload spreads, and the key's bytes are on the wire the way the
                // scenario asked for. `key_size: 0` stays keyless.
                let key = payload::key(scenario.key_size, sent_records);
                sent_records += 1;
                let rec = ProducerRecord {
                    topic: topic.clone(),
                    key,
                    value: Some(value),
                    ..Default::default()
                };
                let t0 = Instant::now();
                match within(enqueue_limit, producer.enqueue(rec)).await {
                    Some(Ok(delivery)) => inflight.push(Box::pin(async move {
                        SettledAck {
                            result: delivery.await,
                            sent_at: t0,
                        }
                    })),
                    Some(Err(error)) => observations.record_ack(Err(error), t0),
                    // `enqueue` outlived the producer's own `max_block`: the producer
                    // is stuck, and another send would only wait again. Stop sending;
                    // the bounded drain and teardown below still let the run report.
                    None => {
                        observations.record_enqueue_timeout();
                        break;
                    }
                }
            }

            // Drain anything still outstanding when the measurement window closed, but
            // do not let a stuck producer retry slot keep the whole benchmark job alive
            // forever. Unsettled sends are counted as drops so failover reports capture
            // the stall instead of timing out without a JSON result.
            let drain_until = Instant::now() + final_drain_timeout.to_std();
            while !inflight.is_empty() {
                let now = Instant::now();
                if now >= drain_until {
                    observations.record_drain_timeout(inflight.len());
                    break;
                }
                match timeout_at(drain_until, inflight.next()).await {
                    Ok(Some(SettledAck {
                        result: res,
                        sent_at: t0,
                    })) => observations.record_ack(res, t0),
                    Ok(None) => break,
                    Err(_) => {
                        observations.record_drain_timeout(inflight.len());
                        break;
                    }
                }
            }

            // `flush` and `close` wait on the same sends the drain gave up on, so
            // they get the same bound: a producer that stops issuing requests must
            // not hold the run open past its report.
            if !finished_within(final_drain_timeout, producer.flush()).await {
                observations.record_teardown_timeout("flush");
            }
            if !finished_within(final_drain_timeout, Box::pin(producer.close())).await {
                observations.record_teardown_timeout("close");
            }

            observations.output
        }
    }
}

// ── Consumer task ───────────────────────────────────────────────────────────

struct ConsumerTask {
    setup: TaskSetup,
    build_retry_policy: ConsumerBuildRetryPolicy,
    poll_timeout: Time,
    poll_error_backoff: Time,
    /// The empty histogram every latency histogram of the task is built from.
    histogram: Histogram<u64>,
}

async fn run_consumer(task: ConsumerTask) -> ConsumerOut {
    let ConsumerTask {
        setup:
            TaskSetup {
                idx,
                scenario,
                bootstrap,
                topic,
                scenario_id,
                stop,
                security,
                grid,
                request_timeout,
                dispatch_queue_capacity,
                frame_max,
            },
        build_retry_policy,
        poll_timeout,
        poll_error_backoff,
        histogram,
    } = task;
    let group_id = format!("krabka-bench-{}", scenario.name);
    let mut consumer = match build_consumer_with_retry(
        idx,
        &bootstrap,
        group_id,
        &topic,
        security,
        ConsumerBuildPolicy {
            request_timeout,
            retry: build_retry_policy,
            dispatch_queue_capacity,
            frame_max,
        },
    )
    .await
    {
        Ok(c) => c,
        Err(e) => {
            return ConsumerOut {
                latency: histogram,
                msgs: 0,
                bytes: 0,
                error: format!("consumer-{idx}-build: {e:#}"),
                interval_msgs: Vec::new(),
                interval_hist: Vec::new(),
            };
        }
    };

    let mut measurements = TaskMeasurements::new(grid.n, &histogram);
    let mut error = String::new();

    loop {
        if stop.load(Ordering::Relaxed) == STATE_STOP {
            break;
        }
        match consumer.poll(poll_timeout).await {
            Ok(records) => {
                let now_ns =
                    nonnegative_i64_to_u64(Utc::now().timestamp_nanos_opt().unwrap_or_default());
                let phase = stop.load(Ordering::Relaxed);
                for r in records {
                    if let Some(val) = &r.value {
                        let bytes = val.len() as u64;
                        if let Some(send_nanos) = payload::read_send_nanos(val, scenario_id) {
                            // Both stamps are epoch nanoseconds; their
                            // difference is the end-to-end extent.
                            let latency = Time::from_nanos(saturating_u64_to_i64(
                                now_ns.saturating_sub(send_nanos),
                            ));
                            // Classify by when the producer SENT the record,
                            // not by the phase this poll returned in. A lagging
                            // consumer drains warmup backlog after the window
                            // opens, and counting it reports warmup throughput
                            // and warmup-aged latency as measurement.
                            if grid.sent_in_measurement(send_nanos) {
                                let slice = grid.idx_for_send_ns(send_nanos);
                                hist::record(&mut measurements.latency, latency);
                                measurements.msgs += 1;
                                measurements.bytes += bytes;
                                if let Some(count) = measurements.interval_msgs.get_mut(slice) {
                                    *count += 1;
                                }
                                if let Some(histogram) = measurements.interval_hist.get_mut(slice) {
                                    hist::record(histogram, latency);
                                }
                            }
                        } else if phase == STATE_MEASURING {
                            // Non-bench record, left over from an earlier run.
                            // It carries no send stamp, so the poll phase is
                            // the only thing that can place it. Count bytes and
                            // not end-to-end latency.
                            measurements.bytes += bytes;
                        }
                    }
                }
            }
            Err(e) => {
                if error.is_empty() {
                    error = format!("consumer-{idx}-poll: {e}");
                }
                sleep(poll_error_backoff.to_std()).await;
            }
        }
    }
    // A close that outlives one request timeout would only hold the run open.
    if !finished_within(request_timeout, Box::pin(consumer.close())).await && error.is_empty() {
        error = format!("consumer-{idx}-close-timeout");
    }
    ConsumerOut {
        latency: measurements.latency,
        msgs: measurements.msgs,
        bytes: measurements.bytes,
        error,
        interval_msgs: measurements.interval_msgs,
        interval_hist: measurements.interval_hist,
    }
}

#[derive(Clone, Copy)]
struct ConsumerBuildPolicy {
    request_timeout: Time,
    retry: ConsumerBuildRetryPolicy,
    dispatch_queue_capacity: ConnectionDispatchQueueCapacity,
    frame_max: ClientFrameMax,
}

async fn build_consumer_with_retry(
    idx: usize,
    bootstrap: &str,
    group_id: String,
    topic: &str,
    security: Option<ClientSecurity>,
    policy: ConsumerBuildPolicy,
) -> Result<Consumer> {
    let backoff = exponential_backoff::Backoff::new(
        policy.retry.attempts(),
        policy.retry.initial_backoff().to_std(),
        Some(policy.retry.max_backoff().to_std()),
    );
    for (attempt_idx, delay) in backoff.into_iter().enumerate() {
        let attempt = attempt_idx + 1;
        let result = Consumer::builder()
            .bootstrap(bootstrap.to_owned())
            .dispatch_queue_capacity(policy.dispatch_queue_capacity.get())
            .frame_max(policy.frame_max.size())
            .client_id(format!("bench-consumer-{idx}"))
            .group_id(group_id.clone())
            .subscribe(vec![topic.to_owned()])
            .auto_offset_reset(AutoOffsetReset::Earliest)
            .request_timeout(policy.request_timeout)
            // `None` → plaintext (default). `Some` → all fetch traffic for this
            // task goes over the broker's TLS listener (kTLS sendfile on krabka).
            .maybe_security(security.clone())
            .build()
            .await
            .context("build consumer");

        match result {
            Ok(consumer) => return Ok(consumer),
            Err(e) => match delay {
                Some(delay) => {
                    warn!(
                        attempt,
                        retry_after_ms = delay.as_millis(),
                        error = %e,
                        "consumer build failed; retrying"
                    );
                    sleep(delay).await;
                }
                None => {
                    return Err(
                        e.context(format!("build consumer failed after {attempt} attempts"))
                    );
                }
            },
        }
    }
    unreachable!("exponential_backoff::Backoff yields a terminal attempt");
}

/// Await `step` for at most `limit`. `None` when the limit ran out first.
async fn within<T>(limit: Time, step: impl Future<Output = T>) -> Option<T> {
    timeout(limit.to_std(), step).await.ok()
}

/// Await `step` for at most `limit`. `false` when the limit ran out first.
async fn finished_within(limit: Time, step: impl Future) -> bool {
    within(limit, step).await.is_some()
}

/// How long the driver lets one `enqueue` take. The producer itself fails a
/// send that waits past `max_block`, which the driver leaves at its default,
/// so a send still pending one request timeout after that is stuck.
fn enqueue_limit(request_timeout: Time) -> Time {
    Time::from_std(DEFAULT_PRODUCER_MAX_BLOCK) + request_timeout
}

/// The producer's `max_in_flight_per_connection` for a scenario.
///
/// The producer pipelines one in-flight request per partition and uses this
/// value as the cross-partition fan-out cap: how many distinct partitions it
/// services concurrently per drain cycle. The default (5) throttles topics with
/// more partitions than that, and the scenarios run 6 to 100 partitions, so a
/// non-idempotent producer gets a cap well past the largest partition count.
///
/// An idempotent producer refuses more than 5, as Kafka's
/// `ProducerConfig.postProcessAndValidateIdempotenceConfigs` does, so an
/// `acks: all` scenario, which turns idempotence on, takes Kafka's limit.
fn producer_max_in_flight(enable_idempotence: bool) -> usize {
    if enable_idempotence {
        DEFAULT_PRODUCER_MAX_IN_FLIGHT
    } else {
        128
    }
}

#[cfg(test)]
mod tests {
    use {
        super::*,
        crate::{
            prom::DEFAULT_PROMETHEUS_REQUEST_TIMEOUT,
            scenario::{FailoverSpec, ModeTag, fixture::short_scenario},
        },
        assert2::{assert, check},
        std::{future, sync::atomic::AtomicU16},
        tokio::{
            sync::{oneshot, oneshot::Receiver},
            time,
        },
    };

    #[test]
    fn an_idempotent_producer_stays_within_kafka_s_in_flight_limit() {
        // (enable_idempotence, expected cap)
        let cases = [(true, 5), (false, 128)];
        for (enable_idempotence, expected) in cases {
            check!(
                producer_max_in_flight(enable_idempotence) == expected,
                "{enable_idempotence}"
            );
        }
    }

    fn cfg(broker_count: u32) -> DriverConfig {
        DriverConfig {
            bootstrap: "broker:9092".into(),
            client_dispatch_queue_capacity: ConnectionDispatchQueueCapacity::default(),
            client_frame_max: ClientFrameMax::default(),
            topic: "t".into(),
            stack: Stack::Krabka,
            namespace: "default".into(),
            prometheus_url: None,
            prometheus_request_timeout: DEFAULT_PROMETHEUS_REQUEST_TIMEOUT,
            producer_request_timeout: default_producer_request_timeout(),
            producer_final_drain_timeout: DEFAULT_PRODUCER_FINAL_DRAIN_TIMEOUT,
            sample_interval: DEFAULT_SAMPLE_INTERVAL,
            consumer_request_timeout: default_consumer_request_timeout(Stack::Krabka),
            consumer_build_retry_policy: ConsumerBuildRetryPolicy::default(),
            consumer_poll_timeout: default_consumer_poll_timeout(),
            consumer_poll_error_backoff: default_consumer_poll_error_backoff(),
            broker_count,
            scenario_id: 0,
            tls: None,
        }
    }

    fn scenario(rf: i16) -> Scenario { short_scenario(rf) }

    #[derive(Debug, PartialEq, Eq)]
    struct ProducerCounters {
        messages: u64,
        bytes: u64,
        dropped: u64,
    }

    fn producer_counters(output: &ProducerOut) -> ProducerCounters {
        ProducerCounters {
            messages: output.msgs,
            bytes: output.bytes,
            dropped: output.dropped,
        }
    }

    #[test]
    fn an_enqueue_may_block_for_max_block_and_one_request_timeout() {
        assert!(enqueue_limit(secs(2)) == secs(62));
    }

    #[tokio::test(start_paused = true)]
    async fn a_shutdown_step_is_bounded_by_its_limit() {
        let outcomes = (
            finished_within(secs(1), future::ready(())).await,
            finished_within(secs(1), future::pending::<()>()).await,
        );
        assert!(outcomes == (true, false));
    }

    /// The phase signal and failover markers a producer's observations
    /// read, for a one-producer scenario sampled every 10 ms.
    struct ObservationHandles {
        grid: Grid,
        stop: Arc<AtomicU8>,
        first_ack: Arc<AtomicU64>,
        kill_at: Arc<AtomicU64>,
    }

    impl ObservationHandles {
        fn new(phase: u8) -> Self {
            Self {
                grid: Grid::new(Instant::now(), 0, &scenario(1), millis(10)),
                stop: Arc::new(AtomicU8::new(phase)),
                first_ack: Arc::new(AtomicU64::new(0)),
                kill_at: Arc::new(AtomicU64::new(0)),
            }
        }

        fn shared(&self) -> SharedHandles<'_> {
            SharedHandles {
                stop: &self.stop,
                first_ack: &self.first_ack,
                kill_at: &self.kill_at,
            }
        }
    }

    #[tokio::test(start_paused = true)]
    async fn a_teardown_timeout_names_its_step_and_keeps_an_earlier_error() {
        let handles = ObservationHandles::new(STATE_STOP);
        let (grid, shared) = (handles.grid, handles.shared());

        let mut stalled = ProducerObservations::new(4, 100, grid, shared, &hist::new().unwrap());
        stalled.record_teardown_timeout("flush");
        stalled.record_teardown_timeout("close");

        let mut drained = ProducerObservations::new(4, 100, grid, shared, &hist::new().unwrap());
        drained.record_drain_timeout(3);
        drained.record_teardown_timeout("flush");

        let mut blocked = ProducerObservations::new(4, 100, grid, shared, &hist::new().unwrap());
        blocked.record_enqueue_timeout();
        blocked.record_teardown_timeout("flush");

        assert!(
            [
                (stalled.output.error.as_str(), stalled.output.dropped),
                (drained.output.error.as_str(), drained.output.dropped),
                (blocked.output.error.as_str(), blocked.output.dropped),
            ] == [
                ("producer-4-flush-timeout", 0),
                ("producer-4-final-drain-timeout:3", 3),
                ("producer-4-enqueue-timeout", 1),
            ]
        );
    }

    #[tokio::test(start_paused = true)]
    async fn settled_acknowledgements_observe_phase_and_shared_failover_markers() {
        let handles = ObservationHandles::new(STATE_RUN);
        let ObservationHandles {
            grid,
            stop,
            first_ack,
            kill_at,
        } = &handles;
        let (grid, shared) = (*grid, handles.shared());
        let metadata = RecordMetadata {
            topic: "t".into(),
            partition: 0,
            offset: 7,
            timestamp_ms: 0,
            serialized_key_size: -1,
            serialized_value_size: 600,
        };

        // Warmup and stopped receipts publish the first acknowledgement, but
        // only measuring receipts contribute throughput, latency, or drops.
        for phase in [STATE_RUN, STATE_MEASURING, STATE_STOP] {
            stop.store(phase, Ordering::SeqCst);
            first_ack.store(0, Ordering::SeqCst);
            let mut observations =
                ProducerObservations::new(4, 100, grid, shared, &hist::new().unwrap());
            let sent_at = Instant::now();
            observations.record_ack(Err(ProducerError::Server(29)), sent_at);
            assert!(first_ack.load(Ordering::SeqCst) == 0);
            time::advance(millis(5).to_std()).await;
            observations.record_ack(Ok(metadata.clone()), sent_at);
            let first_receipt = first_ack.load(Ordering::SeqCst);
            assert!(first_receipt > 0);
            let measured = u64::from(phase == STATE_MEASURING);
            assert!(
                producer_counters(&observations.output)
                    == ProducerCounters {
                        messages: measured,
                        bytes: measured * 100,
                        dropped: measured,
                    }
            );
            assert!(observations.output.latency.len() == measured);
            assert!(observations.output.interval_msgs.iter().sum::<u64>() == measured);
            assert!(
                observations
                    .output
                    .interval_hist
                    .iter()
                    .map(Histogram::len)
                    .sum::<u64>()
                    == measured
            );
            observations.record_ack(Err(ProducerError::Server(31)), sent_at);
            assert!(observations.output.dropped == 2 * measured);
            assert!(
                observations.output.error
                    == if measured == 1 {
                        "producer-4-err: broker error_code 29"
                    } else {
                        ""
                    }
            );
            assert!(
                observations.output.recovery_unix_ms == 0
                    && observations.output.latency_spike_max == Time::ZERO
            );

            stop.store(STATE_STOP, Ordering::SeqCst);
            observations.record_ack(Ok(metadata.clone()), sent_at);
            observations.record_ack(Err(ProducerError::Server(32)), sent_at);
            assert!(
                observations.output.msgs == measured && observations.output.dropped == 2 * measured
            );
            assert!(first_ack.load(Ordering::SeqCst) == first_receipt);
            let first_error = observations.output.error.clone();
            observations.record_drain_timeout(3);
            assert!(observations.output.dropped == 2 * measured + 3);
            assert!(
                observations.output.error
                    == if first_error.is_empty() {
                        "producer-4-final-drain-timeout:3"
                    } else {
                        first_error.as_str()
                    }
            );
        }

        stop.store(STATE_MEASURING, Ordering::SeqCst);
        first_ack.store(777, Ordering::SeqCst);
        let mut observations =
            ProducerObservations::new(4, 100, grid, shared, &hist::new().unwrap());
        let sent_at = Instant::now();
        time::advance(millis(20).to_std()).await;
        observations.record_ack(Ok(metadata.clone()), sent_at);
        assert!(observations.output.latency_spike_max == Time::ZERO);

        // A marker beyond this receipt does not establish recovery. The
        // original spike ledger still records latency once a marker exists.
        kill_at.store(u64::MAX, Ordering::SeqCst);
        let sent_at = Instant::now();
        time::advance(millis(50).to_std()).await;
        observations.record_ack(Ok(metadata.clone()), sent_at);
        assert!(observations.output.recovery_unix_ms == 0);
        assert!(observations.output.latency_spike_max == millis(50));

        kill_at.store(1, Ordering::SeqCst);
        let sent_at = Instant::now();
        time::advance(millis(20).to_std()).await;
        observations.record_ack(Ok(metadata.clone()), sent_at);
        let recovered = observations.output.recovery_unix_ms;
        assert!(recovered >= 1);
        assert!(observations.output.latency_spike_max == millis(50));
        let sent_at = Instant::now();
        time::advance(millis(70).to_std()).await;
        observations.record_ack(Ok(metadata), sent_at);
        assert!(observations.output.recovery_unix_ms == recovered);
        assert!(observations.output.latency_spike_max == millis(70));
        assert!(
            producer_counters(&observations.output)
                == ProducerCounters {
                    messages: 4,
                    bytes: 400,
                    dropped: 0,
                }
        );
        assert!(first_ack.load(Ordering::SeqCst) == 777);
    }

    /// A mock broker for the producer pipeline test, with the signal that a
    /// metadata request named a topic and the largest batch it saw.
    struct ProducerBrokerFixture {
        broker: krabka_client_core::MockBroker,
        started: Receiver<()>,
        largest_batch: Arc<AtomicU64>,
    }

    async fn producer_broker(metadata_error: i16, delivery_error: i16) -> ProducerBrokerFixture {
        use {
            bytes::BytesMut,
            krabka_protocol::{
                Decode as _, Encode as _,
                owned::{
                    api_versions_request,
                    api_versions_response::{ApiVersion, ApiVersionsResponse},
                    metadata_request::{self, MetadataRequest},
                    metadata_response::{
                        MetadataResponse, MetadataResponseBroker, MetadataResponsePartition,
                        MetadataResponseTopic,
                    },
                    produce_request::{self, ProduceRequest},
                    produce_response::{
                        PartitionProduceResponse, ProduceResponse, TopicProduceResponse,
                    },
                },
                records::RecordsPayload,
            },
        };

        let port = Arc::new(AtomicU16::new(0));
        let handler_port = Arc::clone(&port);
        let largest_batch = Arc::new(AtomicU64::new(0));
        let observed_batch = Arc::clone(&largest_batch);
        let (started, receiving) = oneshot::channel();
        let mut started = Some(started);
        let broker = krabka_client_core::MockBroker::start(move |api, version, _, body| {
            let mut encoded = BytesMut::new();
            match api {
                api_versions_request::API_KEY => ApiVersionsResponse {
                    api_keys: [
                        (metadata_request::API_KEY, 0),
                        (produce_request::API_KEY, 3),
                    ]
                    .into_iter()
                    .map(|(api_key, max_version)| ApiVersion {
                        api_key,
                        min_version: max_version,
                        max_version,
                        ..Default::default()
                    })
                    .collect(),
                    ..Default::default()
                }
                .encode(&mut encoded, 0)
                .unwrap(),
                metadata_request::API_KEY => {
                    let mut body = &body[2 + "bench-producer-0".len()..];
                    let request = MetadataRequest::decode(&mut body, version).unwrap();
                    let named = request.topics.is_some_and(|topics| !topics.is_empty());
                    if named && let Some(started) = started.take() {
                        let _ = started.send(());
                    }
                    MetadataResponse {
                        brokers: vec![MetadataResponseBroker {
                            node_id: 1,
                            host: "127.0.0.1".into(),
                            port: i32::from(handler_port.load(Ordering::SeqCst)),
                            ..Default::default()
                        }],
                        topics: named
                            .then(|| MetadataResponseTopic {
                                name: Some("t".into()),
                                error_code: metadata_error,
                                partitions: vec![MetadataResponsePartition {
                                    partition_index: 0,
                                    leader_id: 1,
                                    replica_nodes: vec![1],
                                    isr_nodes: vec![1],
                                    ..Default::default()
                                }],
                                ..Default::default()
                            })
                            .into_iter()
                            .collect(),
                        ..Default::default()
                    }
                    .encode(&mut encoded, version)
                    .unwrap();
                }
                produce_request::API_KEY => {
                    let mut body = &body[2 + "bench-producer-0".len()..];
                    let request = ProduceRequest::decode(&mut body, version).unwrap();
                    ProduceResponse {
                        responses: request
                            .topic_data
                            .into_iter()
                            .map(|topic| TopicProduceResponse {
                                name: topic.name,
                                partition_responses: topic
                                    .partition_data
                                    .into_iter()
                                    .map(|partition| {
                                        let Some(RecordsPayload::V2(batches)) = partition.records
                                        else {
                                            panic!("benchmark sends v2 record batches");
                                        };
                                        for batch in batches {
                                            observed_batch.fetch_max(
                                                batch.records.len() as u64,
                                                Ordering::SeqCst,
                                            );
                                        }
                                        PartitionProduceResponse {
                                            index: partition.index,
                                            error_code: delivery_error,
                                            base_offset: 0,
                                            ..Default::default()
                                        }
                                    })
                                    .collect(),
                                ..Default::default()
                            })
                            .collect(),
                        ..Default::default()
                    }
                    .encode(&mut encoded, version)
                    .unwrap();
                }
                _ => panic!("unexpected benchmark API {api}"),
            }
            Some(encoded.to_vec())
        })
        .await;
        port.store(broker.addr.port(), Ordering::SeqCst);
        ProducerBrokerFixture {
            broker,
            started: receiving,
            largest_batch,
        }
    }

    #[yare::parameterized(
        success = { 0, 0 },
        delivery_rejected = { 0, 29 },
        metadata_rejected = { 29, 0 },
    )]
    #[test_macro(tokio::test(flavor = "multi_thread", worker_threads = 2))]
    async fn producer_pipeline_counts_delivery_and_admission_outcomes(
        metadata_error: i16,
        delivery_error: i16,
    ) {
        let ProducerBrokerFixture {
            broker,
            started,
            largest_batch,
        } = producer_broker(metadata_error, delivery_error).await;
        let mut scenario = scenario(1);
        scenario.linger = millis(20);
        let stop = Arc::new(AtomicU8::new(STATE_MEASURING));
        let first_ack = Arc::new(AtomicU64::new(0));
        let kill_at = Arc::new(AtomicU64::new(0));
        let cfg = DriverConfig {
            bootstrap: broker.addr.to_string(),
            producer_request_timeout: secs(1),
            producer_final_drain_timeout: secs(1),
            ..cfg(1)
        };
        let grid = Grid::new(Instant::now(), 0, &scenario, cfg.sample_interval);
        let mut producers = JoinSet::new();
        spawn_producers(
            &mut producers,
            &scenario,
            &cfg,
            None,
            grid,
            &SharedHandles {
                stop: &stop,
                first_ack: &first_ack,
                kill_at: &kill_at,
            },
            &hist::new().unwrap(),
        );
        assert2::assert!(producers.len() == scenario.producers);
        time::timeout(secs(10).to_std(), started)
            .await
            .unwrap()
            .expect("producer reaches topic admission");
        // A short measurement window exercises the bounded pipeline and its
        // final drain, using the same phase signal as a benchmark run. When
        // delivery succeeds the window opens at the first ack, so a slow
        // runner still measures something.
        if metadata_error == 0 && delivery_error == 0 {
            time::timeout(secs(10).to_std(), async {
                while first_ack.load(Ordering::SeqCst) == 0 {
                    time::sleep(millis(10).to_std()).await;
                }
            })
            .await
            .expect("the first record is acknowledged");
        }
        time::sleep(millis(200).to_std()).await;
        stop.store(STATE_STOP, Ordering::SeqCst);
        let output = time::timeout(secs(10).to_std(), producers.join_next())
            .await
            .unwrap()
            .expect("one producer task")
            .unwrap();
        assert2::assert!(producers.is_empty());
        broker.stop();

        assert2::assert!(output.interval_msgs.iter().sum::<u64>() == output.msgs);
        assert2::assert!(output.latency.len() == output.msgs);
        assert2::assert!(output.bytes == output.msgs * 100);
        if metadata_error == 0 && delivery_error == 0 {
            assert2::assert!(output.msgs > 0);
            assert2::assert!((output.dropped, output.error.as_str()) == (0, ""));
            assert2::assert!(first_ack.load(Ordering::SeqCst) > 0);
        } else {
            assert2::assert!(
                (output.msgs, output.bytes, first_ack.load(Ordering::SeqCst)) == (0, 0, 0)
            );
            assert2::assert!(output.dropped > 0);
            assert2::assert!(output.error == "producer-0-err: broker error_code 29");
        }
        if metadata_error == 0 {
            assert2::assert!(
                largest_batch.load(Ordering::SeqCst) > 1,
                "records are queued before awaiting delivery"
            );
        } else {
            assert2::assert!(
                largest_batch.load(Ordering::SeqCst) == 0,
                "admission failures never reach Produce"
            );
        }
    }

    #[tokio::test(start_paused = true)]
    async fn spawned_consumers_report_startup_failure_without_measurements() {
        let mut scenario = scenario(1);
        scenario.consumers = 2;
        let cfg = DriverConfig {
            bootstrap: String::new(),
            consumer_build_retry_policy: ConsumerBuildRetryPolicy::new(
                ConsumerBuildAttempts::new(1).unwrap(),
                millis(1),
                millis(1),
            )
            .unwrap(),
            ..cfg(1)
        };
        let stop = Arc::new(AtomicU8::new(STATE_MEASURING));
        let grid = Grid::new(Instant::now(), 0, &scenario, cfg.sample_interval);
        let mut consumers = JoinSet::new();
        spawn_consumers(
            &mut consumers,
            &scenario,
            &cfg,
            None,
            grid,
            &stop,
            &hist::new().unwrap(),
        );
        assert!(consumers.len() == scenario.consumers);
        let mut errors = Vec::new();
        // Empty bootstrap does no DNS or TCP work; paused time advances the
        // consumer client's internal startup retries to their deadline.
        while let Some(joined) = time::timeout(secs(600).to_std(), consumers.join_next())
            .await
            .expect("consumer startup retries end")
        {
            let output = joined.expect("consumer reports its build error");
            assert!((output.msgs, output.bytes, output.latency.len()) == (0, 0, 0));
            assert!(output.interval_msgs.is_empty() && output.interval_hist.is_empty());
            errors.push(output.error);
        }
        errors.sort();
        for (idx, error) in errors.iter().enumerate() {
            assert!(error.starts_with(&format!(
                "consumer-{idx}-build: build consumer failed after 1 attempts:"
            )));
        }
        assert!(errors.len() == scenario.consumers);
    }

    #[test]
    fn client_request_timeout_defaults_preserve_policy() {
        check!(default_producer_request_timeout() == secs(2));
        check!(default_consumer_request_timeout(Stack::Krabka) == secs(5));
        check!(default_consumer_request_timeout(Stack::Kafka) == secs(30));
    }

    #[test]
    fn consumer_build_retry_defaults_preserve_policy() {
        let policy = ConsumerBuildRetryPolicy::default();

        check!(
            (
                policy.attempts(),
                policy.initial_backoff(),
                policy.max_backoff()
            ) == (6, millis(100), secs(2))
        );
    }

    #[test]
    fn consumer_build_retry_accepts_positive_minimum_and_equal_backoffs() {
        let attempts = ConsumerBuildAttempts::new(1).expect("one attempt is valid");
        let policy = ConsumerBuildRetryPolicy::new(attempts, millis(1), millis(1))
            .expect("equal bounds are valid");

        check!(policy.attempts() == 1);
        check!(policy.initial_backoff() == millis(1));
        check!(policy.max_backoff() == millis(1));
    }

    #[yare::parameterized(
        zero = { "0" },
        text = { "not-a-number" },
        negative = { "-1" },
        overflow = { "4294967296" },
    )]
    fn consumer_build_retry_rejects_invalid_primitive_values(invalid: &str) {
        check!(
            invalid.parse::<ConsumerBuildAttempts>().is_err(),
            "attempts {invalid:?} must be rejected"
        );
    }

    #[test]
    fn consumer_build_retry_rejects_inverted_backoff_range() {
        let attempts = ConsumerBuildAttempts::new(1).expect("valid attempts");
        assert!(ConsumerBuildRetryPolicy::new(attempts, millis(2), millis(1)).is_err());
    }

    #[test]
    fn consumer_poll_timing_defaults_preserve_behavior() {
        check!(default_consumer_poll_timeout() == millis(50));
        check!(default_consumer_poll_error_backoff() == millis(100));
    }

    #[test]
    fn producer_final_drain_timeout_default_preserves_behavior() {
        check!(DEFAULT_PRODUCER_FINAL_DRAIN_TIMEOUT == secs(10));
    }

    #[test]
    fn sample_interval_default_preserves_behavior() {
        assert!(DEFAULT_SAMPLE_INTERVAL == secs(2));
    }

    // TLS enabled: a CA path + server_name must build a `ClientSecurity` whose
    // protocol is `Ssl`, whose TLS config carries the mounted CA as trust roots
    // and the explicit SNI, and (one-way TLS) no client identity. This is the
    // exact shape passed to `.maybe_security(...)` on both clients.
    #[test]
    fn tls_params_build_ssl_security_with_server_name() {
        let params = TlsParams {
            ca_path: "/etc/bench-ca/ca.crt".into(),
            server_name: "demo-broker-headless.default.svc.cluster.local".into(),
            client_identity: None,
        };
        let sec = params.to_security();
        assert2::assert!(sec.protocol == ListenerProtocol::Ssl);
        assert2::assert!(sec.protocol.requires_tls());
        assert2::assert!(sec.sasl.is_none());
        let tls = sec.tls.expect("Ssl security carries a TLS config");
        assert2::assert!(
            tls.server_name.as_str() == "demo-broker-headless.default.svc.cluster.local"
        );
        assert2::assert!(tls.trust_store == TrustStore::PemFile("/etc/bench-ca/ca.crt".into()));
        assert2::assert!(tls.key_store == None);
    }

    // mTLS variant: a client identity threads through to the TLS config.
    #[test]
    fn tls_params_carry_client_identity_for_mtls() {
        let params = TlsParams {
            ca_path: "/etc/bench-ca/ca.crt".into(),
            server_name: "demo-kafka-bootstrap".into(),
            client_identity: Some(ClientIdentity {
                cert_pem: "/c/tls.crt".into(),
                key_pem: "/c/tls.key".into(),
            }),
        };
        let sec = params.to_security();
        let tls = sec.tls.expect("TLS config present");
        assert2::assert!(
            tls.key_store
                == Some(KeyStore::PemFiles {
                    certificate_chain: "/c/tls.crt".into(),
                    private_key: "/c/tls.key".into(),
                    key_password: None,
                })
        );
    }

    // Disabled: when `DriverConfig.tls` is `None`, the derived security is
    // `None` (plaintext) — the unchanged default benchmark path.
    #[test]
    fn no_tls_params_means_plaintext_security_is_none() {
        let c = cfg(1);
        assert2::assert!(c.tls.is_none());
        let security: Option<ClientSecurity> = c.tls.as_ref().map(TlsParams::to_security);
        assert2::assert!(security.is_none());
    }

    /// A record produced during warmup but polled after the window opens
    /// belongs to the warmup. The grid says so from the record's own send
    /// stamp, whatever the consumer's lag was.
    #[test]
    fn grid_places_a_record_by_its_send_stamp() {
        let mut s = scenario(1);
        s.warmup = secs(10);
        s.duration = secs(60);
        let started_ns = 1_000_000_000_000u64;
        let grid = Grid::new(Instant::now(), started_ns, &s, secs(2));

        let warmup_send = started_ns + nonnegative_i64_to_u64(secs(5).nanos_i64());
        let first_slice = started_ns + nonnegative_i64_to_u64(secs(11).nanos_i64());
        let last_slice = started_ns + nonnegative_i64_to_u64(secs(69).nanos_i64());
        let after_window = started_ns + nonnegative_i64_to_u64(secs(71).nanos_i64());

        check!(!grid.sent_in_measurement(warmup_send));
        check!(grid.sent_in_measurement(first_slice));
        check!(grid.sent_in_measurement(last_slice));
        check!(!grid.sent_in_measurement(after_window));

        // 1s into a 2s-slice grid is slice 0; 59s in is the last slice.
        check!(grid.idx_for_send_ns(first_slice) == 0);
        check!(grid.idx_for_send_ns(last_slice) == grid.n - 1);
    }

    #[yare::parameterized(
        no_kill_marker = {
            FailoverStamps {
                kill_at_ms: 0,
                recovery_at_ms: 5_000,
                dropped: 3,
                latency_spike_max: millis(9),
            },
        },
        no_ack_after_kill = {
            FailoverStamps {
                kill_at_ms: 5_000,
                recovery_at_ms: 0,
                dropped: 3,
                latency_spike_max: millis(9),
            },
        },
    )]
    fn finalize_timing_records_no_disturbance(failover: FailoverStamps) {
        let RunTiming { disturbance, .. } = finalize_timing(true, failover, 1_000, 0);
        check!(disturbance.is_none());
    }

    #[test]
    fn finalize_timing_records_a_measured_kill() {
        let RunTiming { disturbance, .. } = finalize_timing(
            true,
            FailoverStamps {
                kill_at_ms: 5_000,
                recovery_at_ms: 7_500,
                dropped: 3,
                latency_spike_max: millis(9),
            },
            1_000,
            0,
        );
        let disturbance = disturbance.expect("a measured kill gives a disturbance");
        check!(disturbance.kill_at_ms == TimeOffsetMs(5_000));
        check!(disturbance.recovery_at_ms == TimeOffsetMs(7_500));
        check!(disturbance.dropped == MessageCount(3));
    }

    #[test]
    fn empty_output_preserves_inputs() {
        let s = scenario(1);
        let c = cfg(1);
        let out = empty_output(&s, &c, 42, vec!["a-note".into()], vec!["an-error".into()]);
        check!(out.wallclock_start_unix_ms == WallclockMs(42));
        check!(out.wallclock_end_unix_ms == WallclockMs(42));
        check!(out.topology.broker_count == 1);
        check!(out.notes == vec!["a-note".to_owned()]);
        check!(out.errors == vec!["an-error".to_owned()]);
        check!(out.first_ack == Time::ZERO);
        check!(out.disturbance.is_none());
    }

    // The state byte is shared across producer/consumer tasks; verify the
    // three values are pairwise distinct so a flat AtomicU8 can encode them
    // without ambiguity.
    #[test]
    fn state_constants_are_distinct() {
        check!(STATE_RUN != STATE_MEASURING);
        check!(STATE_MEASURING != STATE_STOP);
        check!(STATE_RUN != STATE_STOP);
    }

    #[tokio::test(start_paused = true)]
    async fn cluster_mode_rf3_with_one_broker_is_skipped() {
        let mut s = scenario(3);
        s.mode_tag = ModeTag::Cluster;
        let out = run(s, cfg(1)).await.expect("run returned");
        assert2::assert!(out.throughput.msgs_produced == 0);
        assert2::assert!(out.notes.iter().any(|n| n.contains("topology-mismatch")));
    }

    #[tokio::test(start_paused = true)]
    async fn failover_request_without_rf3_records_skip_note() {
        // Scenario asks for failover, but RF=1 + 1 broker → driver must
        // record a skip note. The scenario has no producers or consumers, so
        // the run builds no client and needs no broker. A consumer pointed at
        // the unresolvable "broker:9092" retries its startup for five minutes
        // and looks the name up again on every attempt: the paused clock makes
        // the backoff free, but each lookup takes real time, a few
        // milliseconds on one machine and far longer on another.
        let mut s = scenario(1);
        s.failover = Some(FailoverSpec {
            kill_after: secs(1),
            target: "partition0_leader".into(),
        });
        s.producers = 0;
        s.consumers = 0;
        s.warmup = Time::ZERO;
        s.duration = Time::ZERO;
        let out = run(s, cfg(1)).await.expect("run returned");
        check!(
            out.notes
                .iter()
                .any(|n| n.contains("skipped:failover-needs-rf3"))
        );
    }
}