plaza_session 0.8.2

//! Frames the session answers for itself.
//!
//! A [`Kind`](plaza_wire::frame::Kind) other than `Ops` is an instruction to
//! the session and never reaches the application. Probes are answered here on
//! the connection task; a `Hello` is forwarded and recorded by the deserialize
//! bridge. Both transports call in here, which is why a latency probe behaves
//! the same over a WebSocket and over TCP.
//!
//! # Two round trips
//!
//! The WebSocket transport times its own ping frame underneath everything
//! plaza does. The probe here is a `Kind::Ping`
//! frame riding the full path: encoded, queued, impaired, decoded. Comparing
//! the two says what plaza and the configured link cost this connection, and
//! on TCP, which has no ping frame of its own, this is the only round trip
//! there is.
//!
//! # Who pings whom
//!
//! On a server only the session originates probes; on a client only the
//! application does. So every `Pong` a side receives answers a `Ping` it sent,
//! and the echoed origin alone matches a `Pong` to its `Ping`.
//!
//! Several must be outstanding at once. A probe is answered a round trip after
//! it leaves and another goes out every 125ms in the fast phase, so on any link
//! slower than that the reply lands after its successor was sent. Tracking one
//! at a time would discard every such sample, so a link slower than 125ms
//! would never be measured.

use std::collections::VecDeque;
use std::time::Duration;


use plaza::agent::AgentId;
use plaza::session::ConnectionId;
use plaza_wire::frame;
use tokio::time::Instant;
use tracing::{debug, trace};

use crate::codec::WireCodec;
use crate::gate::Verdict;
use crate::manager::{ConnectionManager, Frame, Probes, SessionClock};


/// Keeps a connection's two directions from drawing the same jitter sequence.
pub const DOWN_SEED_FLIP: u64 = 0x5DEE_CE66_A5A5_1234;

pub fn earliest(a: Option<Instant>, b: Option<Instant>) -> Option<Instant> {
  match (a, b) {
    (Some(a), Some(b)) => Some(a.min(b)),
    (a, b) => a.or(b),
  }
}

/// A deadline for a `select!` arm that is switched off, so the timer it holds
/// never fires before the guard re-enables it.
pub fn far_future() -> Instant {
  Instant::now() + Duration::from_secs(86_400 * 365)
}

/// The probes this connection has in flight.
pub struct ProbeState {
  /// Oldest first. A pong is matched by its echoed origin, so no correlation
  /// id beyond that is needed, and anything older than the one it answers is
  /// known lost and dropped with it.
  outstanding: VecDeque<(u64, Instant)>,
  seq: u64,
  sent: u32,
  schedule: Probes,
}

impl ProbeState {
  /// Forgets every probe still in flight, without touching the sequence.
  ///
  /// For the link's profile changing under them: a probe launched under one
  /// profile whose pong arrives under another measured a path that existed for
  /// neither leg's whole journey. Its pong then answers no open probe and is
  /// discarded, which is the branch that already exists for the abandoned.
  pub fn forget_in_flight(&mut self) {
    self.outstanding.clear();
  }
}

impl Default for ProbeState {
  fn default() -> Self {
    Self::new(&Probes::default())
  }
}

impl ProbeState {
  pub fn new(schedule: &Probes) -> Self {
    Self {
      outstanding: VecDeque::new(),
      seq: 0,
      sent: 0,
      schedule: schedule.clone(),
    }
  }

  /// How long until the next probe, or `None` when this session does not probe.
  ///
  /// Fast at first, then sparse: a caller deciding whether a connection meets a
  /// schedule wants several samples in the first second, and after that this is
  /// upkeep.
  pub fn interval(&self) -> Option<Duration> {
    if !self.schedule.enabled {
      return None;
    }
    Some(if self.sent < self.schedule.fast_pings {
      self.schedule.fast_interval
    } else {
      self.schedule.idle_interval
    })
  }

  /// When the first probe is due, or `None` when this session does not probe.
  pub fn first_due(&self, now: Instant) -> Option<Instant> {
    self.interval().map(|gap| now + gap)
  }
}

/// What the connection task should do with an inbound frame.
pub enum Inbound {
  /// Not ours: hand it to the application.
  Forward(Frame),
  /// Ours, and it wants an answer written back to the peer.
  Reply(Frame),
  /// Ours, and finished with.
  Consumed,
  /// Theirs, and over the rate this session allows: dropped here, before the
  /// queue every connection shares. The connection stays.
  Shed,
  /// The same, on a rate whose [`Over`](crate::gate::Over) ends connections.
  Eject,
}

/// Builds this connection's next probe.
///
/// The origin is a sequence number rather than a clock reading: the session
/// times the round trip with an `Instant` it keeps, so its own probes need no
/// unit and no clock.
pub fn make_probe<C: WireCodec>(codec: &C, probe: &mut ProbeState, now: Instant) -> Frame {
  probe.seq = probe.seq.wrapping_add(1);
  probe.sent = probe.sent.saturating_add(1);
  if probe.outstanding.len() >= probe.schedule.slots {
    probe.outstanding.pop_front();
  }
  probe.outstanding.push_back((probe.seq, now));

  let mut buf = Vec::with_capacity(frame::PROBE_FRAME_HINT);
  frame::begin(frame::Kind::Ping, &mut buf);
  codec
    .encode_into(&frame::Ping { origin: probe.seq }, &mut buf)
    .expect("a u64 always encodes");
  Frame::from(buf)
}

/// Handles one inbound frame, answering it if it is the session's business.
pub fn handle_inbound<ID: AgentId, C: WireCodec>(
  frame_bytes: Frame,
  codec: &C,
  clock: Option<&SessionClock>,
  probe: &mut ProbeState,
  conn_id: ConnectionId,
  manager: &ConnectionManager<ID>,
) -> Inbound {
  // An empty frame is malformed rather than unknown, and the bridge already
  // reports it, so forwarding it avoids a second report.
  let Some((tag, body)) = frame::split(&frame_bytes) else {
    return gated(frame_bytes, conn_id, manager);
  };

  match frame::Kind::from_byte(tag) {
    Some(frame::Kind::Ping) => match frame::answer_ping(codec, body, clock.map(|c| c())) {
      Some(reply) => Inbound::Reply(Frame::from(reply)),
      None => {
        trace!(conn_id, "Discarding a malformed Ping.");
        Inbound::Consumed
      }
    },
    Some(frame::Kind::Pong) => {
      if let Ok(pong) = codec.decode::<frame::Pong>(body) {
        match probe.outstanding.iter().position(|(origin, _)| *origin == pong.origin) {
          Some(index) => {
            // Everything before it went out earlier and is still unanswered, so
            // it is lost rather than merely late: dropped with the one that
            // arrived, which keeps the queue from filling with the abandoned.
            let (_, sent) = probe.outstanding.drain(..=index).next_back().expect("index is in range");
            manager.record_link_rtt(conn_id, sent.elapsed());
          }
          None => trace!(conn_id, "Discarding a Pong that answers no open probe."),
        }
      }
      Inbound::Consumed
    }
    _ => gated(frame_bytes, conn_id, manager),
  }
}

/// Counts a data frame and returns what the session's rate says to do with it.
///
/// A probe never arrives here, which is what keeps a busy link's own upkeep
/// from spending a client's budget.
fn gated<ID: AgentId>(frame_bytes: Frame, conn_id: ConnectionId, manager: &ConnectionManager<ID>) -> Inbound {
  match manager.record_inbound_activity(conn_id, frame_bytes.len()) {
    Verdict::Admit => Inbound::Forward(frame_bytes),
    Verdict::Shed => Inbound::Shed,
    Verdict::Close => Inbound::Eject,
  }
}

/// What a routed frame leaves the connection task to do.
#[derive(Debug, Default, PartialEq, Eq)]
pub enum Routed {
  /// Delivered, answered internally, or shed. Carry on reading.
  #[default]
  Nothing,
  /// Write this back to the peer.
  Reply(Frame),
  /// Close this connection: it exceeded a rate that ends connections.
  Eject,
}

/// Handles an inbound frame and delivers it wherever it belongs, returning what
/// the peer is owed.
///
/// The caller decides how a reply reaches the socket and how a connection ends,
/// because those are the two parts the transports do differently.
pub async fn route_inbound<ID: AgentId, C: WireCodec>(
  frame_bytes: Frame,
  codec: &C,
  clock: Option<&SessionClock>,
  probe: &mut ProbeState,
  conn_id: ConnectionId,
  manager: &ConnectionManager<ID>,
  agent: &plaza::agent::Agent<ID>,
) -> Routed {
  match handle_inbound(frame_bytes, codec, clock, probe, conn_id, manager) {
    Inbound::Forward(frame_bytes) => {
      manager.forward_incoming(agent.clone(), frame_bytes).await;
      Routed::Nothing
    }
    Inbound::Reply(reply) => Routed::Reply(reply),
    Inbound::Consumed | Inbound::Shed => Routed::Nothing,
    Inbound::Eject => {
      debug!(conn_id, "Closing a connection that exceeded its inbound rate.");
      Routed::Eject
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;
  use crate::codec::JsonCodec;
  use crate::manager::DEFAULT_PROBE_SLOTS;
  use plaza::agent::Agent;
  use std::sync::Arc;

  fn manager() -> Arc<ConnectionManager<u32>> {
    Arc::new(ConnectionManager::new("test", 8))
  }

  fn ping_frame(origin: u64) -> Frame {
    let mut buf = Vec::new();
    frame::begin(frame::Kind::Ping, &mut buf);
    JsonCodec.encode_into(&frame::Ping { origin }, &mut buf).unwrap();
    Frame::from(buf)
  }

  fn pong_frame(origin: u64) -> Frame {
    let mut buf = Vec::new();
    frame::begin(frame::Kind::Pong, &mut buf);
    JsonCodec
      .encode_into(&frame::Pong { origin, responder: None }, &mut buf)
      .unwrap();
    Frame::from(buf)
  }

  fn decode_pong(frame_bytes: &Frame) -> frame::Pong {
    let (kind, body) = frame::split(frame_bytes).unwrap();
    assert_eq!(frame::Kind::from_byte(kind), Some(frame::Kind::Pong));
    JsonCodec.decode(body).unwrap()
  }

  #[tokio::test]
  async fn a_pong_that_straddles_a_profile_change_is_discarded_not_recorded() {
    // A probe launched under one profile whose answer arrives under another
    // went out on the old link and came back on the new one, so it measured
    // neither and would leave a wrong minimum. Forgetting the in-flight probes
    // on the change sends its pong down the answers-no-open-probe branch.
    let manager = manager();
    let agent = Agent::new_human(7u32);
    let (tx, _rx) = plaza::session::session_channel(4);
    let conn_id = manager.register(agent, tx).await;

    let mut probe = ProbeState::default();
    let sent = make_probe(&JsonCodec, &mut probe, Instant::now());
    let (_, body) = frame::split(&sent).unwrap();
    let origin = JsonCodec.decode::<frame::Ping>(body).unwrap().origin;

    probe.forget_in_flight();
    let out = handle_inbound(pong_frame(origin), &JsonCodec, None, &mut probe, conn_id, &manager);
    assert!(matches!(out, Inbound::Consumed));
    assert!(
      manager.agent_link_rtt(&7).is_none(),
      "a straddling sample records nothing"
    );

    // And the next full-journey probe measures normally.
    let sent = make_probe(&JsonCodec, &mut probe, Instant::now());
    let (_, body) = frame::split(&sent).unwrap();
    let origin = JsonCodec.decode::<frame::Ping>(body).unwrap().origin;
    handle_inbound(pong_frame(origin), &JsonCodec, None, &mut probe, conn_id, &manager);
    assert!(manager.agent_link_rtt(&7).is_some(), "and measuring resumes at once");
  }

  #[tokio::test]
  async fn a_ping_is_answered_without_the_application_seeing_it() {
    let manager = manager();
    let mut probe = ProbeState::default();
    let clock: SessionClock = Arc::new(|| 4242);

    let out = handle_inbound(ping_frame(9), &JsonCodec, Some(&clock), &mut probe, 1, &manager);
    let Inbound::Reply(reply) = out else {
      panic!("a ping is answered")
    };
    let pong = decode_pong(&reply);
    assert_eq!(pong.origin, 9);
    assert_eq!(pong.responder, Some(4242));
  }

  #[tokio::test]
  async fn without_a_clock_the_answer_still_carries_the_echo() {
    let manager = manager();
    let mut probe = ProbeState::default();

    let out = handle_inbound(ping_frame(9), &JsonCodec, None, &mut probe, 1, &manager);
    let Inbound::Reply(reply) = out else {
      panic!("a ping is answered")
    };
    let pong = decode_pong(&reply);
    assert_eq!(pong.origin, 9, "a round trip is measurable without a clock");
    assert_eq!(pong.responder, None, "only the offset needs one");
  }

  #[tokio::test]
  async fn a_matching_pong_records_a_sample() {
    let manager = manager();
    let (tx, _rx) = plaza::session::session_channel(4);
    let conn_id = manager.register(Agent::new_human(7u32), tx).await;

    let mut probe = ProbeState::default();
    let sent = make_probe(&JsonCodec, &mut probe, Instant::now());
    let (_, body) = frame::split(&sent).unwrap();
    let origin = JsonCodec.decode::<frame::Ping>(body).unwrap().origin;

    let mut reply = Vec::new();
    frame::begin(frame::Kind::Pong, &mut reply);
    JsonCodec
      .encode_into(
        &frame::Pong {
          origin,
          responder: None,
        },
        &mut reply,
      )
      .unwrap();

    let out = handle_inbound(
      Frame::from(reply),
      &JsonCodec,
      None,
      &mut probe,
      conn_id,
      &manager,
    );
    assert!(matches!(out, Inbound::Consumed));
    assert_eq!(manager.link_rtt_samples(conn_id), 1);
    assert!(manager.agent_link_rtt(&7).is_some());
  }

  #[tokio::test]
  async fn a_pong_answering_nothing_records_nothing() {
    let manager = manager();
    let (tx, _rx) = plaza::session::session_channel(4);
    let conn_id = manager.register(Agent::new_human(7u32), tx).await;
    let mut probe = ProbeState::default();

    let mut stray = Vec::new();
    frame::begin(frame::Kind::Pong, &mut stray);
    JsonCodec
      .encode_into(
        &frame::Pong {
          origin: 12345,
          responder: None,
        },
        &mut stray,
      )
      .unwrap();

    handle_inbound(
      Frame::from(stray),
      &JsonCodec,
      None,
      &mut probe,
      conn_id,
      &manager,
    );
    assert_eq!(manager.link_rtt_samples(conn_id), 0);
  }

  #[tokio::test]
  async fn ops_and_hello_are_not_this_module_s_business() {
    let manager = manager();
    let mut probe = ProbeState::default();

    for kind in [frame::Kind::Ops, frame::Kind::Hello] {
      let mut buf = Vec::new();
      frame::begin(kind, &mut buf);
      buf.extend_from_slice(b"[]");
      let out = handle_inbound(Frame::from(buf), &JsonCodec, None, &mut probe, 1, &manager);
      assert!(matches!(out, Inbound::Forward(_)), "{kind:?} belongs to the bridge");
    }
  }

  #[tokio::test]
  async fn a_pong_slower_than_the_probe_interval_is_still_matched() {
    // With a 200ms link and a 125ms schedule every reply arrives after its
    // successor went out, so a single probe slot would record nothing.
    let manager = manager();
    let (tx, _rx) = plaza::session::session_channel(4);
    let conn_id = manager.register(Agent::new_human(7u32), tx).await;
    let mut probe = ProbeState::default();

    let mut origins = Vec::new();
    for _ in 0..4 {
      let frame_bytes = make_probe(&JsonCodec, &mut probe, Instant::now());
      let (_, body) = frame::split(&frame_bytes).unwrap();
      origins.push(JsonCodec.decode::<frame::Ping>(body).unwrap().origin);
    }

    // The first probe answered only after three more went out.
    handle_inbound(pong_frame(origins[0]), &JsonCodec, None, &mut probe, conn_id, &manager);
    assert_eq!(manager.link_rtt_samples(conn_id), 1, "a late reply still counts");

    handle_inbound(pong_frame(origins[3]), &JsonCodec, None, &mut probe, conn_id, &manager);
    assert_eq!(manager.link_rtt_samples(conn_id), 2);

    // Ones it skipped over are gone rather than left to accumulate.
    handle_inbound(pong_frame(origins[1]), &JsonCodec, None, &mut probe, conn_id, &manager);
    assert_eq!(manager.link_rtt_samples(conn_id), 2, "answered out of order and already dropped");
  }

  #[tokio::test]
  async fn an_unanswered_link_does_not_grow_without_bound() {
    let mut probe = ProbeState::default();
    for _ in 0..(DEFAULT_PROBE_SLOTS * 4) {
      make_probe(&JsonCodec, &mut probe, Instant::now());
    }
    assert_eq!(probe.outstanding.len(), DEFAULT_PROBE_SLOTS, "the oldest are abandoned");
  }

  #[test]
  fn a_session_that_does_not_probe_never_schedules_one() {
    let probe = ProbeState::new(&Probes::off());
    assert_eq!(probe.interval(), None);
    assert_eq!(probe.first_due(Instant::now()), None, "nothing for the timer arm to wait on");
  }

  #[tokio::test]
  async fn probes_start_fast_and_settle() {
    let mut probe = ProbeState::default();
    let schedule = Probes::default();
    for _ in 0..schedule.fast_pings {
      assert_eq!(probe.interval(), Some(schedule.fast_interval));
      make_probe(&JsonCodec, &mut probe, Instant::now());
    }
    assert_eq!(probe.interval(), Some(schedule.idle_interval));
  }
}