melin-server-runtime 0.19.0

//! DPDK response stage — encodes responses and queues them for the DPDK
//! poll thread instead of writing to kernel sockets.
//!
//! The response stage still runs on its own pinned thread for cursor
//! gating and response encoding. Instead of calling `write_all` on kernel
//! sockets, it pushes `(connection_id, encoded_bytes)` into a shared
//! lock-free queue. The DPDK poll thread drains this queue into smoltcp
//! TCP sockets during each poll iteration.
//!
//! This decoupling is necessary because smoltcp is single-threaded — only
//! the DPDK poll thread can call `socket.send_slice()`.

use rustc_hash::FxHashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::mpsc;
use std::time::{Duration, Instant};

use melin_pipeline::ring;
use melin_pipeline::spsc;
use melin_pipeline::wait::WaitStrategy;

use crate::durability_gate::{DurabilityGate, GateInputs, GateOutcome};
use crate::halt::RefusalQueue;
use melin_app::Application;
use melin_app::amortized_timer::AmortizedTimer;
use melin_transport_core::DurableWireSeqCursor;
use melin_transport_core::pipeline::{OutputSlot, StageUtilization};

use melin_wire_protocol::control::TransportResponse;
use melin_wire_protocol::control_codec;

use crate::response_frame::{EncodeBuf, MAX_APP_FRAME, encode_slot_payload, frame_app_response};

#[cfg(feature = "latency-trace")]
use melin_transport_core::trace;

/// Maximum number of output slots consumed per batch.
const MAX_BATCH: usize = 1024;

/// Upper bound on the encoded `BatchEnd` terminator, which is a fixed byte
/// string (a length prefix and a discriminant).
const MAX_BATCH_END_FRAME: usize = 8;

/// Maximum wire frame size: the largest framed application response, plus
/// room for the `BatchEnd` terminator carried in the same frame when the
/// slot closes a request.
const MAX_TX_FRAME: usize = MAX_APP_FRAME + MAX_BATCH_END_FRAME;

/// An encoded frame destined for a specific connection.
/// Sent from the response stage to the DPDK poll thread via lock-free SPSC.
///
/// Fixed-size and `Copy` to fit the SPSC queue's requirements (no heap
/// allocation per frame). The slot is sized for the largest framed
/// response an application's encoder may produce (`MAX_APP_FRAME`), plus
/// the `BatchEnd` terminator.
///
/// A frame carries the response *and* the `BatchEnd` terminator when the slot
/// closes a request. The struct is copied by value out of the SPSC on the poll
/// thread, which is also the ingress thread, so one frame per request rather
/// than two halves both that copy and the per-frame `queue_send` work.
#[derive(Clone, Copy)]
pub struct TxFrame {
    pub connection_id: u64,
    /// Number of valid bytes in `data`.
    pub len: u16,
    /// Wire frame: `[u32 length prefix][payload]`. Only `data[..len]` is valid.
    pub data: [u8; MAX_TX_FRAME],
}

impl Default for TxFrame {
    fn default() -> Self {
        TxFrame {
            connection_id: 0,
            len: 0,
            data: [0u8; MAX_TX_FRAME],
        }
    }
}

impl TxFrame {
    /// The valid wire frame bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.data[..self.len as usize]
    }
}

/// Control plane events for connection registration (DPDK variant).
///
/// Unlike the TCP variant, this doesn't carry a socket writer —
/// the DPDK poll thread owns all socket state.
pub enum ControlEvent {
    /// A new connection was accepted by the DPDK poll thread.
    Connected { connection_id: u64 },
    /// A connection was closed.
    Disconnected { connection_id: u64 },
}

/// Run the DPDK response stage loop. Blocks the calling thread until shutdown.
///
/// Identical to the TCP response stage except:
/// - No socket writers — encoded frames are sent via `tx_out` channel
/// - No flush syscalls — the DPDK poll thread handles transmission
/// - Heartbeats are sent via the same `tx_out` channel
///
/// Top-level thread entry point — the wide arg list mirrors stage state
/// owned elsewhere; bundling into a config struct adds indirection
/// without simplifying.
#[allow(clippy::too_many_arguments)]
// `journal_persisted_wire_seq`: highest wire seq durably persisted on
// this node's journal — same semantics as
// `response::Response::journal_persisted_wire_seq`. See that field's docs
// for why the gate must use wire-seq space rather than the
// journal-consumer cursor.
pub fn run<A: Application>(
    mut consumer: ring::Consumer<OutputSlot<A::Report, A::QueryResponse>>,
    control_rx: mpsc::Receiver<ControlEvent>,
    journal_persisted_wire_seq: DurableWireSeqCursor,
    ack_policy: Arc<std::sync::atomic::AtomicU8>,
    replication_metrics: Option<Arc<crate::replication::ReplicationMetrics>>,
    replica_active: Option<[Arc<AtomicBool>; 2]>,
    shutdown: &AtomicBool,
    heartbeat_interval: Option<Duration>,
    active_connections: Arc<AtomicU64>,
    mut tx_producers: Vec<spsc::Producer<TxFrame>>,
    utilization: Arc<StageUtilization>,
    wait: WaitStrategy,
    encoder: crate::response::ResponseEncoderArc<A>,
    // Writes the poll thread refused while halted, each sent once
    // everything published before it is answered — see `crate::halt`.
    mut refusals: RefusalQueue<A::Report>,
) {
    // The ack policy in force and the durability gate — the same state
    // machine the kernel response stage runs. See `DurabilityGate`.
    let mut gate = DurabilityGate::new(
        GateInputs {
            journal_persisted_wire_seq,
            ack_policy,
            replication_metrics,
            replica_active,
            utilization: Arc::clone(&utilization),
            wait,
        },
        "dpdk response",
    );
    // Track known connections (for heartbeat scheduling).
    // FxHash over SipHash: keys are server-generated connection ids, never
    // attacker-chosen, so HashDoS resistance buys nothing and the default
    // hasher costs more per lookup than the lookup itself.
    let mut connections: FxHashMap<u64, ConnectionHeartbeat> =
        FxHashMap::with_capacity_and_hasher(256, Default::default());

    let mut encode_buf: EncodeBuf = [0u8; MAX_APP_FRAME];

    // The BatchEnd terminator is the same bytes on every request, so it is
    // encoded once here instead of per slot.
    let batch_end_wire_frame = {
        let mut buf = [0u8; MAX_BATCH_END_FRAME];
        let written =
            control_codec::encode_transport_response(&TransportResponse::BatchEnd, &mut buf)
                .expect("BatchEnd encodes");
        buf[..written].to_vec()
    };

    // Pre-encode heartbeat frame (fixed-size, no heap allocation).
    let mut heartbeat_frame = [0u8; 8];
    let heartbeat_len = control_codec::encode_transport_response(
        &TransportResponse::Heartbeat,
        &mut heartbeat_frame,
    )
    .expect("heartbeat encodes");

    // Gate the heartbeat scan's clock read so the count==0 spin doesn't
    // spend the response thread's CPU on `__vdso_clock_gettime`. Reads
    // the clock every ~1 M idle iterations while spinning; heartbeat
    // interval is seconds, so this is plenty.
    let mut heartbeat_timer = AmortizedTimer::new();
    let mut waiter = wait.waiter();
    let mut busy_count: u64 = 0;
    let mut idle_count: u64 = 0;
    // Input sequence of the last output slot handled — refusals are
    // released only where it changes. See `response::run`.
    let mut last_input_seq = u64::MAX;

    // Stage histograms — mirror the TCP response stage but without
    // an `egress` histogram. DPDK egress lives in the poll thread
    // (`dpdk_transport.rs`), which is where the actual TX happens;
    // sampling here would only capture SPSC-publish time and
    // mislead the bench's tick-to-trade breakdown.
    #[cfg(feature = "latency-trace")]
    let mut spsc_rec =
        trace::register_stage("response: SPSC wakeup (matching publish → response consume)");
    #[cfg(feature = "latency-trace")]
    let mut dispatch_rec = trace::register_stage("response: dispatch (consume → SPSC publish)");
    #[cfg(feature = "latency-trace")]
    let mut server_e2e_rec =
        trace::register_stage("server e2e (reader recv → response SPSC publish)");
    #[cfg(feature = "tick-to-trade")]
    let mut journal_wait_rec =
        trace::register_stage("response: journal-wait (match_complete → journal cursor crossed)");
    #[cfg(feature = "tick-to-trade")]
    let mut replica_wait_rec = trace::register_stage(
        "response: replica-wait (match_complete → replication cursor crossed)",
    );
    #[cfg(feature = "tick-to-trade")]
    let mut encode_rec = trace::register_stage("response: encode (per-kind wire encoding)");
    // Paces the idle-path recorder flush — see `response::run`.
    #[cfg(feature = "latency-trace")]
    let mut last_stats_flush = Instant::now();

    'run: loop {
        // Observe runtime policy swaps from the admin `ACK-POLICY`
        // command — see `DurabilityGate::observe_policy_swap`.
        gate.observe_policy_swap();

        // Fencing note: a superseded ex-primary exits through this same
        // shutdown branch (`FenceState::fence_if_superseded` co-sets
        // `shutdown`). Unlike the kernel response stage there is no
        // best-effort flush here to suppress, so the fence needs no
        // dedicated check — exiting stops releasing response frames onto
        // the TX rings, which is the whole ack gate.
        if shutdown.load(Ordering::Relaxed) {
            utilization.busy.store(busy_count, Ordering::Relaxed);
            utilization.idle.store(idle_count, Ordering::Relaxed);
            return;
        }

        // Borrow output slots from the matching stage in place — see
        // `response::run` for why this is a borrow and not a copy, and
        // for the consequence of publishing progress after the batch
        // instead of before it. Refusals follow the order `RefusalQueue`
        // documents: the idle bound before the read, the sync between the
        // read and the control drain.
        let idle_refusal_bound = refusals.idle_bound();
        let slots = consumer.read_contiguous(MAX_BATCH);
        let refusals_waiting = refusals.sync();

        // Poll control channel for connect/disconnect — after the ring
        // is read, never before: the poll thread queues `Connected` on
        // authentication and parses that connection's requests only
        // afterwards, so every slot read above has its `Connected`
        // already queued, and this drain sees it. See `response::run`
        // for the reply this order protects.
        // Counter accounting: the response stage is the sole owner of
        // active_connections decrements. The poll thread increments on
        // auth success and sends ControlEvent::Disconnected on close.
        process_control_events(
            &control_rx,
            &mut connections,
            &active_connections,
            Instant::now,
        );

        if slots.is_empty() {
            if let Some(bound) = idle_refusal_bound {
                let now = Instant::now();
                refusals.release(bound, |refusal| {
                    push_refusal(
                        refusal,
                        &*encoder,
                        &mut connections,
                        &mut tx_producers,
                        &mut encode_buf,
                        &batch_end_wire_frame,
                        now,
                    );
                });
            }

            idle_count += 1;
            if idle_count.is_multiple_of(1024) {
                utilization.busy.store(busy_count, Ordering::Relaxed);
                utilization.idle.store(idle_count, Ordering::Relaxed);
            }

            // Heartbeat scan, gated by AmortizedTimer to keep the busy-
            // spin path off `clock_gettime`. Without this, perf-annotate
            // showed ~22 % of the response thread's CPU on the vDSO.
            if let Some(interval) = heartbeat_interval
                && heartbeat_timer
                    .tick(Duration::from_secs(1), waiter.spinning())
                    .is_some()
            {
                let now = Instant::now();
                let mut failed: Vec<u64> = Vec::new();
                for (&conn_id, state) in connections.iter_mut() {
                    if now.duration_since(state.last_send) >= interval {
                        let mut frame = TxFrame {
                            connection_id: conn_id,
                            len: heartbeat_len as u16,
                            ..Default::default()
                        };
                        frame.data[..heartbeat_len]
                            .copy_from_slice(&heartbeat_frame[..heartbeat_len]);
                        let tid = (conn_id >> 56) as usize % tx_producers.len();
                        if tx_producers[tid].try_publish(frame).is_err() {
                            // SPSC full — DPDK poll thread fell behind.
                            failed.push(conn_id);
                            continue;
                        }
                        state.last_send = now;
                    }
                }
                for conn_id in failed {
                    connections.remove(&conn_id);
                    active_connections.fetch_sub(1, Ordering::Relaxed);
                }
            }

            // Re-evaluate the policy on a slow timer so the
            // `policy_degraded` flag and warn-log track the cluster
            // state even when no batches are flowing. See response.rs
            // for the rationale.
            {
                let now_ts = Instant::now();
                gate.idle_recheck(now_ts);

                // Hand buffered latency samples to the stats registry
                // while idle. Reuses the policy check's clock read —
                // see `response::run` for why this must happen at all.
                #[cfg(feature = "latency-trace")]
                if now_ts.duration_since(last_stats_flush) >= trace::IDLE_FLUSH_INTERVAL {
                    last_stats_flush = now_ts;
                    spsc_rec.flush();
                    dispatch_rec.flush();
                    server_e2e_rec.flush();
                    #[cfg(feature = "tick-to-trade")]
                    {
                        journal_wait_rec.flush();
                        replica_wait_rec.flush();
                        encode_rec.flush();
                    }
                }
            }

            waiter.idle();
            continue;
        }
        waiter.reset();
        busy_count += 1;

        #[cfg(feature = "latency-trace")]
        let consume_ts = trace::mono_trace_ns();

        // One Instant::now() per batch for heartbeat tracking instead of
        // per response — heartbeat interval is 10s, sub-ms precision is plenty.
        let batch_now = Instant::now();

        // Encode and queue responses. Each slot expands to at most two
        // wire frames: the payload (Report / QueryResponse / EngineError)
        // and an optional trailing `BatchEnd` when `is_last_in_request`
        // is set. `OutputPayload::BatchEnd` carries no payload of its
        // own — the wire BatchEnd is emitted purely from the flag.
        //
        // The durability gate is evaluated per slot, inline — see
        // `response::run` for why batch-max gating was a head-of-line
        // block and why the extra waits are self-limiting.
        for slot in slots {
            #[cfg(feature = "latency-trace")]
            spsc_rec.record_elapsed(slot.match_complete_ts, consume_ts);

            // First slot of an event: refusals stamped up to its sequence
            // go first, ahead of the gate. See `response::run`.
            let first_of_event = slot.input_seq != last_input_seq;
            last_input_seq = slot.input_seq;
            if refusals_waiting && first_of_event {
                refusals.release(slot.input_seq, |refusal| {
                    push_refusal(
                        refusal,
                        &*encoder,
                        &mut connections,
                        &mut tx_producers,
                        &mut encode_buf,
                        &batch_end_wire_frame,
                        batch_now,
                    );
                });
            }

            // Per-slot journal-wait / replica-wait tracker. Same shape as
            // the TCP response — see `crate::response::GateCrossTracker`
            // for the rationale and edge cases.
            #[cfg(feature = "tick-to-trade")]
            let mut gate_tracker = crate::response::GateCrossTracker::new(slot.wire_seq);

            // Gate on `wire_seq` (matches `response::run`) — see that
            // module's notes on the input-seq vs wire-seq space mismatch
            // that motivated the field. With the wire-seq stamp, the
            // slot's own `wire_seq` is the exact value the gate must
            // observe for that slot's event to count as durable; no `+1`
            // (the old `+1` compensated for the input-seq
            // off-by-(starting-1), which is gone now).
            //
            // A gate that cannot open does not hold the shutdown sequence;
            // the reply the policy never confirmed is not queued.
            if gate.needs_wait(slot)
                && let GateOutcome::Shutdown = gate.wait_durable(
                    slot.wire_seq,
                    shutdown,
                    #[cfg(feature = "tick-to-trade")]
                    &mut gate_tracker,
                )
            {
                continue 'run;
            }

            #[cfg(feature = "tick-to-trade")]
            if let Some(ts) = gate_tracker.journal_crossed() {
                journal_wait_rec.record_elapsed(slot.match_complete_ts, ts);
            }
            #[cfg(feature = "tick-to-trade")]
            if let Some(ts) = gate_tracker.replica_crossed() {
                replica_wait_rec.record_elapsed(slot.match_complete_ts, ts);
            }

            // One lookup, not two: the entry decides whether the
            // connection is still known, and the same borrow stamps its
            // heartbeat clock once the frame is queued.
            let Some(conn_state) = connections.get_mut(&slot.connection_id) else {
                // A reply for a connection this stage does not hold: it
                // is gone. The drain runs after the ring is read, so a
                // connection that is merely new is never a miss. Logged
                // so a lost reply leaves a trace. Connection 0 is a
                // server-originated event — a seed, a tick — whose
                // reports have no client to go to, and is skipped
                // without a word.
                if slot.connection_id != 0 {
                    tracing::debug!(
                        connection_id = slot.connection_id,
                        wire_seq = slot.wire_seq,
                        "reply dropped: connection not registered"
                    );
                }
                continue;
            };

            // All frames for this slot share the same destination tid
            // (single connection_id), so we can compute it once and
            // bundle the slot's 1–2 frames under one Release at the end
            // of the slot. Per-slot flush — rather than once per outer
            // batch — keeps individual requests' RTT short under
            // saturation (each request's response ships as soon as
            // encoded) while still letting the DPDK poll thread drain
            // Report + BatchEnd in a single consume cycle for low-rate
            // workloads (consumer-side win for single-request p99).
            let tid = (slot.connection_id >> 56) as usize % tx_producers.len();
            let conn_id = slot.connection_id;

            // Frame 1: application payload (Report / Query via encoder;
            // EngineError via codec). BatchEnd payloads carry no body —
            // the terminator below handles them via is_last_in_request.
            let payload_result = encode_slot_payload(&slot.payload, &*encoder, &mut encode_buf);

            // The BatchEnd terminator rides in the same frame. It is
            // transport-shaped and byte-identical every time, so it is
            // encoded once at startup rather than per slot.
            let trailer: &[u8] = if slot.is_last_in_request {
                &batch_end_wire_frame
            } else {
                &[]
            };

            #[cfg(feature = "tick-to-trade")]
            let encode_start = trace::mono_trace_ns();
            push_frame(
                payload_result,
                conn_id,
                &mut tx_producers[tid],
                &encode_buf,
                trailer,
            );
            #[cfg(feature = "tick-to-trade")]
            encode_rec.record_elapsed(encode_start, trace::mono_trace_ns());

            // Release this slot's frames as a unit. `flush` is a no-op
            // when nothing was pushed (e.g. every kind hit an encode
            // error), so the call is safe regardless.
            tx_producers[tid].flush();
            #[cfg(feature = "latency-trace")]
            if slot.is_last_in_request {
                // Record server-e2e relative to post-flush wall clock —
                // i.e. the moment the DPDK poll thread can see the
                // response.
                server_e2e_rec.record_elapsed(slot.recv_ts, trace::mono_trace_ns());
            }

            conn_state.last_send = batch_now;
        }

        // Borrowed slots are dead here — hand the batch back to the
        // producer now that every one of them has cleared its gate.
        consumer.commit();

        // Log degradation transitions / re-emit the reminder — after
        // dispatch, not before it. See `DurabilityGate::after_batch`.
        gate.after_batch();

        #[cfg(feature = "latency-trace")]
        dispatch_rec.record_elapsed(consume_ts, trace::mono_trace_ns());
    }
}

/// Push one SPSC frame carrying an encoded payload and/or a pre-encoded
/// trailer onto the DPDK tx ring.
///
/// `payload` is `None` for slot kinds with no body (`BatchEnd`), and `trailer`
/// is empty for slots that do not close a request; a call with neither pushes
/// nothing. Encode failures are logged at error level and drop the payload,
/// but still ship the trailer, so a request that fails to encode one response
/// is still terminated rather than leaving the client waiting.
#[inline]
fn push_frame(
    payload: Option<Result<usize, &'static str>>,
    conn_id: u64,
    tx: &mut spsc::Producer<TxFrame>,
    encode_buf: &[u8],
    trailer: &[u8],
) {
    let written = match payload {
        Some(Ok(n)) => n,
        Some(Err(reason)) => {
            tracing::error!(connection_id = conn_id, reason, "encode error");
            0
        }
        None => 0,
    };
    if written == 0 && trailer.is_empty() {
        return;
    }
    let total = written + trailer.len();
    debug_assert!(total <= MAX_TX_FRAME);
    let len = total as u16;
    tx.push_with(|frame| {
        frame.connection_id = conn_id;
        frame.len = len;
        frame.data[..written].copy_from_slice(&encode_buf[..written]);
        frame.data[written..total].copy_from_slice(trailer);
    });
}

/// Queue the reply to a write refused at ingress — the application's
/// rejection and the request terminator, in one frame — for the poll
/// thread. No durability gate: nothing was published. A connection that
/// has gone away is skipped.
fn push_refusal<R: Copy, Q: Copy>(
    refusal: crate::halt::Refusal<R>,
    encoder: &dyn melin_app::encoder::ResponseEncoder<Report = R, Query = Q>,
    connections: &mut FxHashMap<u64, ConnectionHeartbeat>,
    tx_producers: &mut [spsc::Producer<TxFrame>],
    encode_buf: &mut EncodeBuf,
    batch_end_wire_frame: &[u8],
    now: Instant,
) {
    let conn_id = refusal.connection_id;
    let Some(conn_state) = connections.get_mut(&conn_id) else {
        tracing::debug!(
            connection_id = conn_id,
            "refusal dropped: connection not registered"
        );
        return;
    };
    let tid = (conn_id >> 56) as usize % tx_producers.len();
    let payload = frame_app_response(encode_buf, |body| {
        encoder.encode_report(&refusal.report, body)
    });
    push_frame(
        Some(payload),
        conn_id,
        &mut tx_producers[tid],
        encode_buf,
        batch_end_wire_frame,
    );
    tx_producers[tid].flush();
    conn_state.last_send = now;
}

/// Per-connection heartbeat state. No socket writer — the DPDK poll
/// thread owns socket state.
struct ConnectionHeartbeat {
    last_send: Instant,
}

/// Process a batch of control events, updating the connection map and
/// active_connections counter.
///
/// Extracted from the `run()` loop so the counter accounting invariant
/// can be unit-tested: the response stage is the **sole owner** of
/// `active_connections` decrements. The poll thread increments on auth
/// success and sends `Disconnected`; this function handles the decrement.
///
/// A new connection's heartbeat clock starts at `clock()`, read at most once
/// per call and only when a `Connected` is drained: this runs every loop
/// iteration, and almost every drain is empty.
fn process_control_events(
    control_rx: &mpsc::Receiver<ControlEvent>,
    connections: &mut FxHashMap<u64, ConnectionHeartbeat>,
    active_connections: &AtomicU64,
    clock: impl Fn() -> Instant,
) {
    let mut now = None;
    while let Ok(event) = control_rx.try_recv() {
        match event {
            ControlEvent::Connected { connection_id } => {
                let last_send = *now.get_or_insert_with(&clock);
                connections.insert(connection_id, ConnectionHeartbeat { last_send });
            }
            ControlEvent::Disconnected { connection_id } => {
                if connections.remove(&connection_id).is_some() {
                    active_connections.fetch_sub(1, Ordering::Relaxed);
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicU64;
    use std::sync::mpsc;
    use std::time::Instant;

    /// Simulate the poll thread's side: increment counter on auth, send
    /// Disconnected on close. The response stage (process_control_events)
    /// owns the decrement.
    #[test]
    fn active_connections_single_lifecycle() {
        let counter = AtomicU64::new(0);
        let (tx, rx) = mpsc::channel();
        let mut connections = FxHashMap::default();
        let now = Instant::now();

        // Poll thread: auth succeeds → increment.
        counter.fetch_add(1, Ordering::Relaxed);
        tx.send(ControlEvent::Connected { connection_id: 1 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, || now);
        assert_eq!(counter.load(Ordering::Relaxed), 1);
        assert_eq!(connections.len(), 1);

        // Poll thread: connection closes → send Disconnected (no decrement).
        tx.send(ControlEvent::Disconnected { connection_id: 1 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, || now);
        assert_eq!(counter.load(Ordering::Relaxed), 0);
        assert_eq!(connections.len(), 0);
    }

    /// Disconnected for an unknown connection (e.g., pre-auth drop or
    /// duplicate event) must not decrement the counter.
    #[test]
    fn disconnect_unknown_connection_no_decrement() {
        let counter = AtomicU64::new(0);
        let (tx, rx) = mpsc::channel();
        let mut connections = FxHashMap::default();
        let now = Instant::now();

        tx.send(ControlEvent::Disconnected { connection_id: 999 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, || now);
        // Counter must stay at 0 — not wrap to u64::MAX.
        assert_eq!(counter.load(Ordering::Relaxed), 0);
    }

    /// Multiple connections with interleaved connect/disconnect.
    #[test]
    fn active_connections_multiple_lifecycle() {
        let counter = AtomicU64::new(0);
        let (tx, rx) = mpsc::channel();
        let mut connections = FxHashMap::default();
        let now = Instant::now();

        // Three connections authenticate.
        for id in 1..=3 {
            counter.fetch_add(1, Ordering::Relaxed);
            tx.send(ControlEvent::Connected { connection_id: id })
                .unwrap();
        }
        process_control_events(&rx, &mut connections, &counter, || now);
        assert_eq!(counter.load(Ordering::Relaxed), 3);
        assert_eq!(connections.len(), 3);

        // Connection 2 disconnects.
        tx.send(ControlEvent::Disconnected { connection_id: 2 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, || now);
        assert_eq!(counter.load(Ordering::Relaxed), 2);
        assert_eq!(connections.len(), 2);

        // Remaining two disconnect.
        tx.send(ControlEvent::Disconnected { connection_id: 1 })
            .unwrap();
        tx.send(ControlEvent::Disconnected { connection_id: 3 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, || now);
        assert_eq!(counter.load(Ordering::Relaxed), 0);
        assert_eq!(connections.len(), 0);
    }

    /// Duplicate Disconnected for the same connection must only decrement
    /// once (the second remove returns None).
    #[test]
    fn duplicate_disconnect_single_decrement() {
        let counter = AtomicU64::new(0);
        let (tx, rx) = mpsc::channel();
        let mut connections = FxHashMap::default();
        let now = Instant::now();

        counter.fetch_add(1, Ordering::Relaxed);
        tx.send(ControlEvent::Connected { connection_id: 1 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, || now);

        // Two Disconnected events for the same connection.
        tx.send(ControlEvent::Disconnected { connection_id: 1 })
            .unwrap();
        tx.send(ControlEvent::Disconnected { connection_id: 1 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, || now);
        assert_eq!(counter.load(Ordering::Relaxed), 0);
    }

    /// A new connection's heartbeat clock starts when its `Connected` is
    /// drained, not at some earlier stored instant; and a drain with no
    /// `Connected` reads no clock at all, since it runs every iteration.
    #[test]
    fn connected_stamps_the_clock_read_lazily_at_drain() {
        let counter = AtomicU64::new(0);
        let (tx, rx) = mpsc::channel();
        let mut connections = FxHashMap::default();
        let reads = std::cell::Cell::new(0u32);
        let at_drain = Instant::now() + std::time::Duration::from_secs(60);
        let clock = || {
            reads.set(reads.get() + 1);
            at_drain
        };

        // Empty drain, and a drain of a Disconnected only: no clock read.
        process_control_events(&rx, &mut connections, &counter, clock);
        tx.send(ControlEvent::Disconnected { connection_id: 9 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, clock);
        assert_eq!(reads.get(), 0);

        // Two Connected in one drain: one read, both stamped with it.
        tx.send(ControlEvent::Connected { connection_id: 1 })
            .unwrap();
        tx.send(ControlEvent::Connected { connection_id: 2 })
            .unwrap();
        process_control_events(&rx, &mut connections, &counter, clock);
        assert_eq!(reads.get(), 1);
        assert_eq!(connections[&1].last_send, at_drain);
        assert_eq!(connections[&2].last_send, at_drain);
    }

    mod push_frame {
        use super::*;

        const TRAILER: &[u8] = b"END";

        /// Drive `push_frame` and return every frame the poll thread would
        /// see, as `(connection_id, bytes)`.
        fn pushed(
            payload: Option<Result<usize, &'static str>>,
            trailer: &[u8],
        ) -> Vec<(u64, Vec<u8>)> {
            let (mut tx, mut rx) = spsc::channel::<TxFrame>(8, WaitStrategy::SpinThenYield);
            let encode_buf = b"response-body";
            super::super::push_frame(payload, 42, &mut tx, encode_buf, trailer);
            tx.flush();

            let mut out = Vec::new();
            while let Some((_, frame)) = rx.try_consume() {
                out.push((frame.connection_id, frame.as_bytes().to_vec()));
            }
            out
        }

        #[test]
        fn payload_and_trailer_ship_as_one_frame() {
            let frames = pushed(Some(Ok(8)), TRAILER);

            assert_eq!(frames.len(), 1, "two SPSC slots where one suffices");
            assert_eq!(frames[0].0, 42);
            assert_eq!(frames[0].1, b"responseEND".to_vec());
        }

        #[test]
        fn payload_without_trailer_ships_alone() {
            let frames = pushed(Some(Ok(8)), &[]);

            assert_eq!(frames.len(), 1);
            assert_eq!(frames[0].1, b"response".to_vec());
        }

        #[test]
        fn trailer_without_payload_ships_alone() {
            // A `BatchEnd` slot: no body of its own, but it closes the
            // request.
            let frames = pushed(None, TRAILER);

            assert_eq!(frames.len(), 1);
            assert_eq!(frames[0].1, TRAILER.to_vec());
        }

        #[test]
        fn nothing_to_send_pushes_no_frame() {
            assert!(pushed(None, &[]).is_empty());
        }

        #[test]
        fn encode_failure_still_terminates_the_request() {
            // The response is lost, but dropping the terminator too would
            // leave the client waiting on a batch that never ends.
            let frames = pushed(Some(Err("encode error")), TRAILER);

            assert_eq!(frames.len(), 1);
            assert_eq!(frames[0].1, TRAILER.to_vec());
        }

        #[test]
        fn encode_failure_without_trailer_pushes_no_frame() {
            assert!(pushed(Some(Err("encode error")), &[]).is_empty());
        }

        #[test]
        fn the_encoded_terminator_fits_its_reserved_room() {
            let batch_end = {
                let mut buf = [0u8; MAX_BATCH_END_FRAME];
                let written = control_codec::encode_transport_response(
                    &TransportResponse::BatchEnd,
                    &mut buf,
                )
                .expect("BatchEnd encodes");
                buf[..written].to_vec()
            };
            assert!(batch_end.len() <= MAX_BATCH_END_FRAME);
        }
    }
}