use std::collections::HashMap;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::time::Duration;
use std::sync::Arc;
use async_trait::async_trait;
use bytes::Bytes;
use parking_lot::RwLock;
use plaza::agent::{Agent, AgentId};
use plaza::error::PlazaError;
use plaza::session::{
session_channel, ConnectionId, MessageTarget, PresenceEvent, Session, SessionMessage, SessionReceiver,
SessionSender,
};
use serde::de::DeserializeOwned;
use serde::Serialize;
use tracing::{debug, trace, warn};
use crate::codec::WireCodec;
use plaza_wire::frame::{self, ProtocolVersion};
use crate::conditioner::LinkProfile;
use crate::error::SessionLayerError;
use crate::stats::TransportStats;
pub const DEFAULT_BROADCAST_CAPACITY: usize = 256;
pub const DEFAULT_CLIENT_QUEUE_CAPACITY: usize = 64;
pub const DEFAULT_CONDITIONER_CAPACITY: usize = 1024;
pub const DEFAULT_PROBE_FAST_PINGS: u32 = 8;
pub const DEFAULT_PROBE_FAST_INTERVAL: Duration = Duration::from_millis(125);
pub const DEFAULT_PROBE_IDLE_INTERVAL: Duration = Duration::from_secs(5);
pub const DEFAULT_PROBE_SLOTS: usize = 16;
pub const DEFAULT_MAX_FRAME_BYTES: usize = 8 * 1024 * 1024;
pub const DEFAULT_MAX_MESSAGE_BYTES: usize = 1024 * 1024;
#[derive(Debug, Clone)]
pub struct Queues {
pub inbound: usize,
pub decoded: usize,
pub presence: usize,
pub outbound: usize,
pub conditioner: usize,
}
impl Default for Queues {
fn default() -> Self {
Self {
inbound: DEFAULT_BROADCAST_CAPACITY,
decoded: DEFAULT_BROADCAST_CAPACITY,
presence: DEFAULT_BROADCAST_CAPACITY,
outbound: DEFAULT_CLIENT_QUEUE_CAPACITY,
conditioner: DEFAULT_CONDITIONER_CAPACITY,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum OutboundOverflow {
#[default]
Drop,
Disconnect,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum InboundOverflow {
#[default]
Drop,
Backpressure,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PresenceOverflow {
#[default]
Drop,
Backpressure,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct Overflow {
pub outbound: OutboundOverflow,
pub inbound: InboundOverflow,
pub presence: PresenceOverflow,
}
impl Overflow {
pub fn drop_everywhere() -> Self {
Self::default()
}
pub fn block_where_possible() -> Self {
Self {
outbound: OutboundOverflow::Drop,
inbound: InboundOverflow::Backpressure,
presence: PresenceOverflow::Backpressure,
}
}
}
#[derive(Debug, Clone)]
pub struct Limits {
pub max_frame_bytes: usize,
pub max_message_bytes: usize,
}
impl Default for Limits {
fn default() -> Self {
Self {
max_frame_bytes: DEFAULT_MAX_FRAME_BYTES,
max_message_bytes: DEFAULT_MAX_MESSAGE_BYTES,
}
}
}
#[derive(Debug, Clone)]
pub struct Probes {
pub enabled: bool,
pub slots: usize,
pub fast_pings: u32,
pub fast_interval: Duration,
pub idle_interval: Duration,
}
impl Default for Probes {
fn default() -> Self {
Self {
enabled: true,
slots: DEFAULT_PROBE_SLOTS,
fast_pings: DEFAULT_PROBE_FAST_PINGS,
fast_interval: DEFAULT_PROBE_FAST_INTERVAL,
idle_interval: DEFAULT_PROBE_IDLE_INTERVAL,
}
}
}
impl Probes {
pub fn off() -> Self {
Self {
enabled: false,
..Self::default()
}
}
}
pub type SessionClock = Arc<dyn Fn() -> u64 + Send + Sync>;
#[derive(Clone)]
pub struct SessionOptions {
pub protocol: ProtocolVersion,
pub clock: Option<SessionClock>,
pub queues: Queues,
pub limits: Limits,
pub overflow: Overflow,
pub probes: Probes,
}
impl Default for SessionOptions {
fn default() -> Self {
Self {
protocol: ProtocolVersion::UNKNOWN,
clock: None,
queues: Queues::default(),
limits: Limits::default(),
overflow: Overflow::default(),
probes: Probes::default(),
}
}
}
impl SessionOptions {
pub fn with_protocol(protocol: ProtocolVersion) -> Self {
Self {
protocol,
..Self::default()
}
}
pub fn clock(mut self, clock: impl Fn() -> u64 + Send + Sync + 'static) -> Self {
self.clock = Some(Arc::new(clock));
self
}
pub fn workload(mut self, workload: &crate::workload::Workload) -> Self {
self.queues = Queues::for_workload(workload);
self.limits = Limits::for_workload(workload);
self.overflow = Overflow::for_workload(workload);
self
}
pub fn queues(mut self, queues: Queues) -> Self {
self.queues = queues;
self
}
pub fn limits(mut self, limits: Limits) -> Self {
self.limits = limits;
self
}
pub fn overflow(mut self, overflow: Overflow) -> Self {
self.overflow = overflow;
self
}
pub fn disconnect_slow_clients(mut self) -> Self {
self.overflow.outbound = OutboundOverflow::Disconnect;
self
}
pub fn backpressure_inbound(mut self) -> Self {
self.overflow.inbound = InboundOverflow::Backpressure;
self
}
pub fn backpressure_presence(mut self) -> Self {
self.overflow.presence = PresenceOverflow::Backpressure;
self
}
pub fn inbound_capacity(mut self, depth: usize) -> Self {
self.queues.inbound = depth;
self
}
pub fn decoded_capacity(mut self, depth: usize) -> Self {
self.queues.decoded = depth;
self
}
pub fn presence_capacity(mut self, depth: usize) -> Self {
self.queues.presence = depth;
self
}
pub fn outbound_capacity(mut self, depth: usize) -> Self {
self.queues.outbound = depth;
self
}
pub fn conditioner_capacity(mut self, depth: usize) -> Self {
self.queues.conditioner = depth;
self
}
pub fn max_frame_bytes(mut self, bytes: usize) -> Self {
self.limits.max_frame_bytes = bytes;
self
}
pub fn max_message_bytes(mut self, bytes: usize) -> Self {
self.limits.max_message_bytes = bytes;
self
}
pub fn probe_slots(mut self, slots: usize) -> Self {
self.probes.slots = slots;
self
}
pub fn probes(mut self, probes: Probes) -> Self {
self.probes = probes;
self
}
pub fn without_probes(mut self) -> Self {
self.probes.enabled = false;
self
}
pub fn probe_schedule(mut self, fast_pings: u32, fast: Duration, idle: Duration) -> Self {
self.probes.fast_pings = fast_pings;
self.probes.fast_interval = fast;
self.probes.idle_interval = idle;
self
}
}
impl Debug for SessionOptions {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SessionOptions")
.field("protocol", &self.protocol)
.field("clock", &self.clock.is_some())
.field("queues", &self.queues)
.field("limits", &self.limits)
.field("overflow", &self.overflow)
.field("probes", &self.probes)
.finish()
}
}
pub struct IncomingFrame<ID: AgentId> {
pub from: Agent<ID>,
pub frame: Frame,
}
#[derive(Clone, PartialEq, Eq)]
pub struct Frame(Bytes);
impl Frame {
pub fn len(&self) -> usize {
self.0.len()
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn into_bytes(self) -> Bytes {
self.0
}
}
impl From<Bytes> for Frame {
fn from(bytes: Bytes) -> Self {
Self(bytes)
}
}
impl From<Vec<u8>> for Frame {
fn from(bytes: Vec<u8>) -> Self {
Self(Bytes::from(bytes))
}
}
impl AsRef<[u8]> for Frame {
fn as_ref(&self) -> &[u8] {
&self.0
}
}
impl std::ops::Deref for Frame {
type Target = [u8];
fn deref(&self) -> &[u8] {
&self.0
}
}
impl Debug for Frame {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Frame({} bytes)", self.0.len())
}
}
pub type OutboundFrame = Frame;
#[derive(Debug)]
pub enum ConnectionOrder {
Close { farewell: Option<OutboundFrame> },
Deadline {
after: Option<Duration>,
farewell: Option<OutboundFrame>,
},
}
const ORDER_QUEUE_DEPTH: usize = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct InboundVolume {
pub frames: u64,
pub bytes: u64,
}
struct ClientHandle<ID: AgentId> {
agent: Agent<ID>,
to_client_tx: SessionSender<OutboundFrame>,
orders_tx: SessionSender<ConnectionOrder>,
orders_rx: RwLock<Option<SessionReceiver<ConnectionOrder>>>,
last_activity_us: AtomicU64,
inbound_frames: AtomicU64,
inbound_bytes: AtomicU64,
rtt_us: AtomicU64,
min_rtt_us: AtomicU64,
samples: AtomicU64,
protocol: AtomicU64,
link_rtt_us: AtomicU64,
min_link_rtt_us: AtomicU64,
link_samples: AtomicU64,
link_dropped: AtomicU64,
link: Arc<LinkHandle>,
}
pub struct LinkHandle {
impaired: AtomicBool,
profile: RwLock<LinkProfile>,
}
impl LinkHandle {
fn new() -> Self {
Self {
impaired: AtomicBool::new(false),
profile: RwLock::new(LinkProfile::default()),
}
}
pub fn impaired(&self) -> bool {
self.impaired.load(Ordering::Acquire)
}
pub fn read(&self) -> LinkProfile {
*self.profile.read()
}
fn set(&self, profile: LinkProfile) {
*self.profile.write() = profile;
self.impaired.store(!profile.is_passthrough(), Ordering::Release);
}
}
impl<ID: AgentId> ClientHandle<ID> {
fn new(agent: Agent<ID>, to_client_tx: SessionSender<OutboundFrame>, now_us: u64) -> Self {
let (orders_tx, orders_rx) = session_channel::<ConnectionOrder>(ORDER_QUEUE_DEPTH);
Self {
agent,
to_client_tx,
orders_tx,
orders_rx: RwLock::new(Some(orders_rx)),
last_activity_us: AtomicU64::new(now_us),
inbound_frames: AtomicU64::new(0),
inbound_bytes: AtomicU64::new(0),
rtt_us: AtomicU64::new(0),
min_rtt_us: AtomicU64::new(0),
samples: AtomicU64::new(0),
protocol: AtomicU64::new(0),
link_rtt_us: AtomicU64::new(0),
min_link_rtt_us: AtomicU64::new(0),
link_samples: AtomicU64::new(0),
link_dropped: AtomicU64::new(0),
link: Arc::new(LinkHandle::new()),
}
}
}
fn record_sample(samples: &AtomicU64, smoothed_us: &AtomicU64, min_us: &AtomicU64, us: u64) {
samples.fetch_add(1, Ordering::Relaxed);
let previous = smoothed_us.load(Ordering::Relaxed);
let smoothed = if previous == 0 { us } else { previous - previous / 8 + us / 8 };
smoothed_us.store(smoothed, Ordering::Relaxed);
let min = min_us.load(Ordering::Relaxed);
if min == 0 || us < min {
min_us.store(us, Ordering::Relaxed);
}
}
fn take_stream<T: Send + 'static>(slot: &RwLock<Option<T>>, name: &str) -> T {
slot
.write()
.take()
.unwrap_or_else(|| panic!("the {name} stream was already taken; it has a single consumer"))
}
pub fn target_matches<ID: AgentId>(target: &MessageTarget<ID>, agent: &Agent<ID>) -> bool {
let agent_id = match agent.id() {
Some(id) => id,
None => return false,
};
match target {
MessageTarget::All => true,
MessageTarget::Agent(id) => id == agent_id,
MessageTarget::Agents(ids) => ids.contains(agent_id),
MessageTarget::AllExcept(id) => id != agent_id,
MessageTarget::AllExceptThese(ids) => !ids.contains(agent_id),
}
}
struct Registry<ID: AgentId> {
by_conn: HashMap<ConnectionId, ClientHandle<ID>>,
by_agent: HashMap<ID, Vec<ConnectionId>>,
}
impl<ID: AgentId> Registry<ID> {
fn new() -> Self {
Self {
by_conn: HashMap::new(),
by_agent: HashMap::new(),
}
}
fn insert(&mut self, conn_id: ConnectionId, handle: ClientHandle<ID>) {
if let Some(id) = handle.agent.id_cloned() {
self.by_agent.entry(id).or_default().push(conn_id);
}
self.by_conn.insert(conn_id, handle);
}
fn remove(&mut self, conn_id: ConnectionId) -> Option<ClientHandle<ID>> {
let handle = self.by_conn.remove(&conn_id)?;
if let Some(id) = handle.agent.id() {
if let Some(conns) = self.by_agent.get_mut(id) {
conns.retain(|held| *held != conn_id);
if conns.is_empty() {
self.by_agent.remove(id);
}
}
}
Some(handle)
}
fn get(&self, conn_id: ConnectionId) -> Option<&ClientHandle<ID>> {
self.by_conn.get(&conn_id)
}
fn len(&self) -> usize {
self.by_conn.len()
}
fn handles(&self) -> impl Iterator<Item = &ClientHandle<ID>> {
self.by_conn.values()
}
fn handles_with_ids(&self) -> impl Iterator<Item = (ConnectionId, &ClientHandle<ID>)> {
self.by_conn.iter().map(|(conn_id, handle)| (*conn_id, handle))
}
fn for_agent(&self, id: &ID) -> impl Iterator<Item = (ConnectionId, &ClientHandle<ID>)> {
self
.by_agent
.get(id)
.map_or(&[][..], |conns| conns.as_slice())
.iter()
.filter_map(|conn_id| self.by_conn.get(conn_id).map(|handle| (*conn_id, handle)))
}
fn for_target(&self, target: &MessageTarget<ID>, mut visit: impl FnMut(ConnectionId, &ClientHandle<ID>)) {
match target {
MessageTarget::All => {
for (conn_id, handle) in &self.by_conn {
if handle.agent.id().is_some() {
visit(*conn_id, handle);
}
}
}
MessageTarget::Agent(id) => {
for (conn_id, handle) in self.for_agent(id) {
visit(conn_id, handle);
}
}
MessageTarget::Agents(ids) => {
for (position, id) in ids.iter().enumerate() {
if ids[..position].contains(id) {
continue;
}
for (conn_id, handle) in self.for_agent(id) {
visit(conn_id, handle);
}
}
}
MessageTarget::AllExcept(id) => {
for (conn_id, handle) in &self.by_conn {
if handle.agent.id().is_some_and(|agent_id| agent_id != id) {
visit(*conn_id, handle);
}
}
}
MessageTarget::AllExceptThese(ids) => {
for (conn_id, handle) in &self.by_conn {
if handle.agent.id().is_some_and(|agent_id| !ids.contains(agent_id)) {
visit(*conn_id, handle);
}
}
}
}
}
}
pub struct ConnectionManager<ID: AgentId> {
transport: &'static str,
epoch: std::time::Instant,
next_conn_id: AtomicU64,
connections: RwLock<Registry<ID>>,
raw_incoming_tx: SessionSender<IncomingFrame<ID>>,
raw_incoming_rx: RwLock<Option<SessionReceiver<IncomingFrame<ID>>>>,
presence_tx: SessionSender<PresenceEvent<ID>>,
presence_rx: RwLock<Option<SessionReceiver<PresenceEvent<ID>>>>,
stats: Arc<TransportStats>,
hello: Option<OutboundFrame>,
clock: Option<SessionClock>,
queues: Queues,
limits: Limits,
overflow: Overflow,
probes: Probes,
}
impl<ID: AgentId> Debug for ConnectionManager<ID> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ConnectionManager")
.field("transport", &self.transport)
.field("connections", &self.connections.read().len())
.field("declares_protocol", &self.hello.is_some())
.field("has_clock", &self.clock.is_some())
.finish()
}
}
impl<ID: AgentId> Debug for ClientHandle<ID> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ClientHandle").field("agent", &self.agent).finish()
}
}
impl<ID: AgentId> ConnectionManager<ID> {
pub fn stats(&self) -> Arc<TransportStats> {
Arc::clone(&self.stats)
}
pub fn new(transport: &'static str, capacity: usize) -> Self {
Self::with_hello(transport, capacity, None)
}
pub fn with_hello(transport: &'static str, capacity: usize, hello: Option<OutboundFrame>) -> Self {
Self::with_hello_and_clock(transport, capacity, hello, None)
}
pub fn with_hello_and_clock(
transport: &'static str,
capacity: usize,
hello: Option<OutboundFrame>,
clock: Option<SessionClock>,
) -> Self {
let options = SessionOptions {
clock,
queues: Queues {
inbound: capacity,
decoded: capacity,
presence: capacity,
..Queues::default()
},
..SessionOptions::default()
};
Self::with_options(transport, hello, &options)
}
pub fn with_options(transport: &'static str, hello: Option<OutboundFrame>, options: &SessionOptions) -> Self {
let (raw_incoming_tx, raw_incoming_rx) = session_channel(options.queues.inbound);
let (presence_tx, presence_rx) = session_channel(options.queues.presence);
Self {
transport,
epoch: std::time::Instant::now(),
next_conn_id: AtomicU64::new(1),
connections: RwLock::new(Registry::new()),
raw_incoming_tx,
raw_incoming_rx: RwLock::new(Some(raw_incoming_rx)),
presence_tx,
presence_rx: RwLock::new(Some(presence_rx)),
stats: TransportStats::new(),
hello,
clock: options.clock.clone(),
queues: options.queues.clone(),
limits: options.limits.clone(),
overflow: options.overflow,
probes: options.probes.clone(),
}
}
pub fn probes(&self) -> &Probes {
&self.probes
}
pub fn overflow(&self) -> Overflow {
self.overflow
}
pub fn queues(&self) -> &Queues {
&self.queues
}
pub fn limits(&self) -> &Limits {
&self.limits
}
pub fn clock(&self) -> Option<&SessionClock> {
self.clock.as_ref()
}
pub async fn register(&self, agent: Agent<ID>, to_client_tx: SessionSender<OutboundFrame>) -> ConnectionId {
let conn_id = self.next_conn_id.fetch_add(1, Ordering::Relaxed);
if let Some(hello) = &self.hello {
let _ = to_client_tx.try_send(hello.clone());
}
self
.connections
.write()
.insert(conn_id, ClientHandle::new(agent.clone(), to_client_tx, self.now_us()));
debug!(transport = self.transport, conn_id, agent = %agent, "Connection registered.");
if !self.announce(PresenceEvent::Joined { agent, conn_id }).await {
warn!(
transport = self.transport,
conn_id,
depth = self.queues.presence,
"Join notification dropped: no controller listening, or its queue is full. \
Raise `presence_capacity` or set `backpressure_presence`."
);
}
conn_id
}
async fn announce(&self, event: PresenceEvent<ID>) -> bool {
let taken = match self.overflow.presence {
PresenceOverflow::Drop => self.presence_tx.try_send(event).is_ok(),
PresenceOverflow::Backpressure => self.presence_tx.send(event).await.is_ok(),
};
if !taken {
self.stats.record_presence_dropped();
}
taken
}
pub async fn deregister(&self, conn_id: ConnectionId) {
self.remove(conn_id, true).await
}
fn now_us(&self) -> u64 {
self.epoch.elapsed().as_micros() as u64
}
pub fn record_inbound_activity(&self, conn_id: ConnectionId, bytes: usize) {
let now_us = self.now_us();
let connections = self.connections.read();
if let Some(handle) = connections.get(conn_id) {
handle.last_activity_us.store(now_us, Ordering::Relaxed);
handle.inbound_frames.fetch_add(1, Ordering::Relaxed);
handle.inbound_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
}
}
pub fn idle_for(&self, conn_id: ConnectionId) -> Option<Duration> {
let connections = self.connections.read();
let handle = connections.get(conn_id)?;
let last = handle.last_activity_us.load(Ordering::Relaxed);
Some(Duration::from_micros(self.now_us().saturating_sub(last)))
}
pub fn agent_idle_for(&self, id: &ID) -> Option<Duration> {
let connections = self.connections.read();
let now_us = self.now_us();
connections
.for_agent(id)
.map(|(_, handle)| now_us.saturating_sub(handle.last_activity_us.load(Ordering::Relaxed)))
.min()
.map(Duration::from_micros)
}
pub fn connection_inbound(&self, conn_id: ConnectionId) -> Option<InboundVolume> {
let connections = self.connections.read();
let handle = connections.get(conn_id)?;
Some(InboundVolume {
frames: handle.inbound_frames.load(Ordering::Relaxed),
bytes: handle.inbound_bytes.load(Ordering::Relaxed),
})
}
pub fn agent_inbound(&self, id: &ID) -> InboundVolume {
let connections = self.connections.read();
connections.for_agent(id).fold(InboundVolume::default(), |sum, (_, handle)| InboundVolume {
frames: sum.frames + handle.inbound_frames.load(Ordering::Relaxed),
bytes: sum.bytes + handle.inbound_bytes.load(Ordering::Relaxed),
})
}
pub fn connections_of(&self, id: &ID) -> Vec<ConnectionId> {
self
.connections
.read()
.for_agent(id)
.map(|(conn_id, _)| conn_id)
.collect()
}
pub fn close_connection(&self, conn_id: ConnectionId, farewell: Option<OutboundFrame>) -> bool {
let connections = self.connections.read();
match connections.get(conn_id) {
Some(handle) => handle.orders_tx.try_send(ConnectionOrder::Close { farewell }).is_ok(),
None => false,
}
}
pub fn set_deadline(&self, conn_id: ConnectionId, after: Option<Duration>, farewell: Option<OutboundFrame>) -> bool {
let connections = self.connections.read();
match connections.get(conn_id) {
Some(handle) => handle
.orders_tx
.try_send(ConnectionOrder::Deadline { after, farewell })
.is_ok(),
None => false,
}
}
pub fn deregister_agent(&self, id: &ID, farewell: Option<OutboundFrame>) -> usize {
self
.connections_of(id)
.into_iter()
.filter(|conn_id| self.close_connection(*conn_id, farewell.clone()))
.count()
}
pub fn disconnect_all(&self, farewell: Option<OutboundFrame>) -> usize {
let connections = self.connections.read();
connections
.handles_with_ids()
.filter(|(_, handle)| {
handle
.orders_tx
.try_send(ConnectionOrder::Close {
farewell: farewell.clone(),
})
.is_ok()
})
.count()
}
pub fn take_orders(&self, conn_id: ConnectionId) -> Option<SessionReceiver<ConnectionOrder>> {
self
.connections
.read()
.get(conn_id)
.map(|handle| take_stream(&handle.orders_rx, "connection orders"))
}
async fn remove(&self, conn_id: ConnectionId, may_wait: bool) {
let handle = self.connections.write().remove(conn_id);
match handle {
Some(handle) => {
debug!(transport = self.transport, conn_id, agent = %handle.agent, "Connection deregistered.");
if let Some(id) = handle.agent.id_cloned() {
let event = PresenceEvent::Left { agent_id: id, conn_id };
let announced = if may_wait {
self.announce(event).await
} else {
let taken = self.presence_tx.try_send(event).is_ok();
if !taken {
self.stats.record_presence_dropped();
}
taken
};
if !announced {
warn!(transport = self.transport, conn_id, "Leave notification dropped.");
}
}
}
None => {
trace!(
transport = self.transport,
conn_id,
"Deregister for unknown connection; already gone."
);
}
}
}
pub async fn forward_incoming(&self, from: Agent<ID>, frame: impl Into<Frame>) {
let incoming = IncomingFrame { from, frame: frame.into() };
let taken = match self.overflow.inbound {
InboundOverflow::Backpressure => self.raw_incoming_tx.send(incoming).await.is_ok(),
InboundOverflow::Drop => self.raw_incoming_tx.try_send(incoming).is_ok(),
};
self.stats.record_inbound(!taken);
if !taken {
warn!(
transport = self.transport,
depth = self.queues.inbound,
"Inbound ops dropped: controller queue full or closed. \
Raise `inbound_capacity` or set `backpressure_inbound`."
);
}
}
pub fn broadcast(
&self,
target: &MessageTarget<ID>,
frame: OutboundFrame,
) -> Result<Vec<ConnectionId>, SessionLayerError> {
let disconnecting = self.overflow.outbound == OutboundOverflow::Disconnect;
let connections = self.connections.read();
let mut full: Option<ConnectionId> = None;
let mut overflowed = Vec::new();
let (mut sent, mut dropped) = (0u64, 0u64);
connections.for_target(target, |conn_id, handle| {
if handle.to_client_tx.try_send(frame.clone()).is_err() {
dropped += 1;
full.get_or_insert(conn_id);
if disconnecting {
overflowed.push(conn_id);
}
} else {
sent += 1;
}
});
drop(connections);
self.stats.record_outbound(sent, dropped);
match full {
Some(conn_id) if !disconnecting => Err(SessionLayerError::ClientSendFailed {
transport: self.transport,
conn_id,
reason: "client queue full or closed",
}),
_ => Ok(overflowed),
}
}
pub async fn disconnect_overflowed(&self, overflowed: Vec<ConnectionId>) {
for conn_id in overflowed {
warn!(
transport = self.transport,
conn_id,
depth = self.queues.outbound,
"Disconnecting a client whose outbound queue was full."
);
self.remove(conn_id, false).await;
}
}
pub fn take_raw_incoming(&self) -> SessionReceiver<IncomingFrame<ID>> {
take_stream(&self.raw_incoming_rx, "raw incoming")
}
pub fn take_presence(&self) -> SessionReceiver<PresenceEvent<ID>> {
take_stream(&self.presence_rx, "presence")
}
pub fn record_protocol(&self, agent: &Agent<ID>, version: ProtocolVersion) {
let Some(id) = agent.id() else { return };
let connections = self.connections.read();
for (_, handle) in connections.for_agent(id) {
handle.protocol.store(version.0 as u64, Ordering::Relaxed);
}
}
pub fn protocol(&self, id: &ID) -> Option<ProtocolVersion> {
let connections = self.connections.read();
connections
.for_agent(id)
.next()
.map(|(_, h)| ProtocolVersion(h.protocol.load(Ordering::Relaxed) as u32))
.filter(|v| *v != ProtocolVersion::UNKNOWN)
}
pub fn record_rtt(&self, conn_id: ConnectionId, rtt: Duration) {
let connections = self.connections.read();
let Some(handle) = connections.get(conn_id) else {
return;
};
record_sample(
&handle.samples,
&handle.rtt_us,
&handle.min_rtt_us,
rtt.as_micros() as u64,
);
}
pub fn record_link_rtt(&self, conn_id: ConnectionId, rtt: Duration) {
let connections = self.connections.read();
let Some(handle) = connections.get(conn_id) else {
return;
};
record_sample(
&handle.link_samples,
&handle.link_rtt_us,
&handle.min_link_rtt_us,
rtt.as_micros() as u64,
);
}
pub fn link_rtt(&self, conn_id: ConnectionId) -> Option<Duration> {
self.sample(conn_id, |h| h.link_rtt_us.load(Ordering::Relaxed))
}
pub fn min_link_rtt(&self, conn_id: ConnectionId) -> Option<Duration> {
self.sample(conn_id, |h| h.min_link_rtt_us.load(Ordering::Relaxed))
}
pub fn link_rtt_samples(&self, conn_id: ConnectionId) -> u64 {
self
.connections
.read()
.get(conn_id)
.map(|h| h.link_samples.load(Ordering::Relaxed))
.unwrap_or(0)
}
pub fn agent_link_rtt(&self, id: &ID) -> Option<(Duration, u64)> {
let connections = self.connections.read();
let (_, handle) = connections.for_agent(id).next()?;
let us = handle.min_link_rtt_us.load(Ordering::Relaxed);
(us > 0).then(|| (Duration::from_micros(us), handle.link_samples.load(Ordering::Relaxed)))
}
pub fn record_link_drop(&self, conn_id: ConnectionId) {
let connections = self.connections.read();
if let Some(handle) = connections.get(conn_id) {
handle.link_dropped.fetch_add(1, Ordering::Relaxed);
}
}
pub fn link_dropped(&self, conn_id: ConnectionId) -> u64 {
self
.connections
.read()
.get(conn_id)
.map(|h| h.link_dropped.load(Ordering::Relaxed))
.unwrap_or(0)
}
pub fn agent_link_dropped(&self, id: &ID) -> u64 {
self
.connections
.read()
.for_agent(id)
.map(|(_, h)| h.link_dropped.load(Ordering::Relaxed))
.sum()
}
pub fn total_link_dropped(&self) -> u64 {
self
.connections
.read()
.handles()
.map(|h| h.link_dropped.load(Ordering::Relaxed))
.sum()
}
pub fn set_link_profile(&self, conn_id: ConnectionId, profile: LinkProfile) {
let connections = self.connections.read();
if let Some(handle) = connections.get(conn_id) {
handle.link.set(profile);
}
}
pub fn set_agent_link_profile(&self, id: &ID, profile: LinkProfile) {
let connections = self.connections.read();
for (_, handle) in connections.for_agent(id) {
handle.link.set(profile);
}
}
pub fn set_all_link_profiles(&self, profile: LinkProfile) {
let connections = self.connections.read();
for handle in connections.handles() {
handle.link.set(profile);
}
}
pub fn link_profile(&self, conn_id: ConnectionId) -> Option<LinkProfile> {
self.connections.read().get(conn_id).map(|h| h.link.read())
}
pub fn link_handle(&self, conn_id: ConnectionId) -> Option<Arc<LinkHandle>> {
self.connections.read().get(conn_id).map(|h| Arc::clone(&h.link))
}
pub fn rtt(&self, conn_id: ConnectionId) -> Option<Duration> {
self.sample(conn_id, |h| h.rtt_us.load(Ordering::Relaxed))
}
pub fn min_rtt(&self, conn_id: ConnectionId) -> Option<Duration> {
self.sample(conn_id, |h| h.min_rtt_us.load(Ordering::Relaxed))
}
pub fn rtt_samples(&self, conn_id: ConnectionId) -> u64 {
self
.connections
.read()
.get(conn_id)
.map(|h| h.samples.load(Ordering::Relaxed))
.unwrap_or(0)
}
pub fn agent_rtt(&self, id: &ID) -> Option<(Duration, u64)> {
let connections = self.connections.read();
let (_, handle) = connections.for_agent(id).next()?;
let us = handle.min_rtt_us.load(Ordering::Relaxed);
(us > 0).then(|| (Duration::from_micros(us), handle.samples.load(Ordering::Relaxed)))
}
fn sample(&self, conn_id: ConnectionId, read: impl Fn(&ClientHandle<ID>) -> u64) -> Option<Duration> {
let connections = self.connections.read();
let us = read(connections.get(conn_id)?);
(us > 0).then(|| Duration::from_micros(us))
}
pub fn connection_count(&self) -> usize {
self.connections.read().len()
}
}
pub struct TransportSession<Op: Send + 'static, ID: AgentId, C: WireCodec> {
transport: &'static str,
codec: C,
manager: Arc<ConnectionManager<ID>>,
deserialized_rx: RwLock<Option<SessionReceiver<SessionMessage<Op, ID>>>>,
encode_hint: AtomicUsize,
_phantom: PhantomData<fn() -> Op>,
}
impl<Op, ID, C> Debug for TransportSession<Op, ID, C>
where
Op: Send + 'static,
ID: AgentId,
C: WireCodec,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TransportSession")
.field("transport", &self.transport)
.field("codec", &self.codec.name())
.field("connections", &self.manager.connection_count())
.finish()
}
}
impl<Op, ID, C> TransportSession<Op, ID, C>
where
Op: Serialize + DeserializeOwned + Clone + Debug + Send + Sync + 'static,
ID: AgentId,
C: WireCodec,
{
pub fn with_protocol(transport: &'static str, codec: C, capacity: usize, protocol: ProtocolVersion) -> Arc<Self> {
let options = SessionOptions {
protocol,
queues: Queues {
inbound: capacity,
decoded: capacity,
presence: capacity,
..Queues::default()
},
..SessionOptions::default()
};
Self::with_options(transport, codec, options)
}
pub fn with_options(transport: &'static str, codec: C, options: SessionOptions) -> Arc<Self> {
let protocol = options.protocol;
let hello = (protocol != ProtocolVersion::UNKNOWN).then(|| {
let mut buf = Vec::new();
frame::begin(frame::Kind::Hello, &mut buf);
codec.encode_into(&protocol, &mut buf).expect("a u32 always encodes");
Frame::from(buf)
});
let manager = Arc::new(ConnectionManager::with_options(transport, hello, &options));
let (deserialized_tx, deserialized_rx) = session_channel(options.queues.decoded);
let session = Arc::new(Self {
transport,
codec: codec.clone(),
manager: manager.clone(),
deserialized_rx: RwLock::new(Some(deserialized_rx)),
encode_hint: AtomicUsize::new(0),
_phantom: PhantomData,
});
tokio::spawn(deserialize_bridge::<Op, ID, C>(
transport,
codec,
manager.clone(),
protocol,
manager.take_raw_incoming(),
deserialized_tx,
));
session
}
pub fn new(transport: &'static str, codec: C, capacity: usize) -> Arc<Self> {
Self::with_protocol(transport, codec, capacity, ProtocolVersion::UNKNOWN)
}
pub fn manager(&self) -> &Arc<ConnectionManager<ID>> {
&self.manager
}
pub fn codec(&self) -> &C {
&self.codec
}
pub fn encode_message(&self, msg: SessionMessage<Op, ID>) -> Result<OutboundFrame, SessionLayerError> {
let mut buf = Vec::with_capacity(self.encode_hint.load(Ordering::Relaxed));
frame::begin(frame::Kind::Ops, &mut buf);
self
.codec
.encode_into(&msg.ops, &mut buf)
.map_err(|source| SessionLayerError::Serialization {
transport: self.transport,
context: "session message",
source,
})?;
let hint = self.encode_hint.load(Ordering::Relaxed);
self.encode_hint.store(buf.len().max(hint / 2), Ordering::Relaxed);
Ok(Frame::from(buf))
}
}
async fn deserialize_bridge<Op, ID, C>(
transport: &'static str,
codec: C,
manager: Arc<ConnectionManager<ID>>,
expected: ProtocolVersion,
raw_rx: SessionReceiver<IncomingFrame<ID>>,
typed_tx: SessionSender<SessionMessage<Op, ID>>,
) where
Op: DeserializeOwned + Clone + Debug + Send + 'static,
ID: AgentId,
C: WireCodec,
{
let mut probe_warned: std::collections::HashSet<ID> = std::collections::HashSet::new();
loop {
let Ok(raw) = raw_rx.recv().await else {
debug!(transport, "Raw incoming channel closed; deserialize bridge stopping.");
return;
};
let (from, bytes) = (raw.from, raw.frame);
let Some((tag, body)) = frame::split(&bytes) else {
warn!(transport, agent = %from, "Discarding an empty frame.");
continue;
};
match frame::Kind::from_byte(tag) {
Some(frame::Kind::Ops) => {}
Some(frame::Kind::Hello) => {
match codec.decode::<ProtocolVersion>(body) {
Ok(theirs) => {
manager.record_protocol(&from, theirs);
if !theirs.agrees_with(expected) {
debug!(
transport,
agent = %from,
theirs = theirs.0,
ours = expected.0,
"Client declared a different protocol version; see ConnectionManager::protocol."
);
}
}
Err(source) => warn!(transport, error = %source, agent = %from, "Discarding a malformed Hello."),
}
continue;
}
Some(frame::Kind::Ping) | Some(frame::Kind::Pong) => {
if from.id().is_none_or(|id| probe_warned.insert(id.clone())) {
warn!(
transport,
agent = %from,
"A probe frame reached the deserialize bridge, so this transport is not answering probes. \
Call `control::route_inbound` (or at least `frame::answer_ping`) on the connection task \
before forwarding: a client measuring its round trip will never hear back."
);
}
continue;
}
_ => {
trace!(transport, kind = tag, agent = %from, "Skipping a frame of unknown kind.");
continue;
}
}
let typed = match codec.decode::<Vec<Op>>(body) {
Ok(ops) => SessionMessage::new(from, ops),
Err(source) => {
warn!(
transport,
error = %source,
agent = %from,
"Discarding client ops that failed to decode."
);
continue;
}
};
if typed_tx.send(typed).await.is_err() {
debug!(transport, "Controller stopped consuming; deserialize bridge stopping.");
return;
}
}
}
#[async_trait]
impl<Op, ID, C> Session<Op, ID> for TransportSession<Op, ID, C>
where
Op: Serialize + DeserializeOwned + Clone + Debug + Send + Sync + 'static,
ID: AgentId,
C: WireCodec,
{
async fn send_message(
&self,
target: MessageTarget<ID>,
msg: SessionMessage<Op, ID>,
) -> Result<(), PlazaError<ID>> {
let encoded = self.encode_message(msg)?;
let overflowed = self.manager.broadcast(&target, encoded)?;
self.manager.disconnect_overflowed(overflowed).await;
Ok(())
}
fn subscribe_to_incoming_messages(&self) -> SessionReceiver<SessionMessage<Op, ID>> {
take_stream(&self.deserialized_rx, "incoming messages")
}
fn on_presence_change(&self) -> SessionReceiver<PresenceEvent<ID>> {
self.manager.take_presence()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::conditioner::DirectionProfile;
fn handle(agent: Agent<u32>) -> ClientHandle<u32> {
ClientHandle::new(agent, session_channel(4).0, 0)
}
#[test]
fn routing_through_the_index_answers_what_a_scan_answers() {
let members = [
(1u64, Agent::Human(10u32)),
(2, Agent::Human(20)),
(3, Agent::Human(10)),
(4, Agent::Bot(30)),
(5, Agent::System),
];
let mut registry = Registry::new();
for (conn_id, agent) in &members {
registry.insert(*conn_id, handle(agent.clone()));
}
let targets = [
MessageTarget::All,
MessageTarget::Agent(10),
MessageTarget::Agent(99),
MessageTarget::Agents(vec![10, 30, 10]),
MessageTarget::Agents(vec![]),
MessageTarget::AllExcept(10),
MessageTarget::AllExcept(99),
MessageTarget::AllExceptThese(vec![10, 30]),
MessageTarget::AllExceptThese(vec![]),
MessageTarget::Agents((0..40).chain([10, 10]).collect()),
MessageTarget::AllExceptThese((0..40).collect()),
];
for target in &targets {
let mut addressed = Vec::new();
registry.for_target(target, |conn_id, _| addressed.push(conn_id));
addressed.sort_unstable();
let mut scanned: Vec<ConnectionId> = members
.iter()
.filter(|(_, agent)| target_matches(target, agent))
.map(|(conn_id, _)| *conn_id)
.collect();
scanned.sort_unstable();
assert_eq!(addressed, scanned, "{target:?}");
}
}
#[test]
fn a_removed_connection_leaves_no_index_entry() {
let mut registry = Registry::new();
registry.insert(1, handle(Agent::Human(10u32)));
registry.insert(2, handle(Agent::Human(10)));
registry.remove(1);
let mut addressed = Vec::new();
registry.for_target(&MessageTarget::Agent(10), |conn_id, _| addressed.push(conn_id));
assert_eq!(addressed, vec![2]);
registry.remove(2);
assert!(registry.by_agent.is_empty());
}
#[test]
fn fanning_a_frame_out_shares_one_buffer() {
let frame: OutboundFrame = Frame::from(vec![7u8; 4096]);
let queued: Vec<OutboundFrame> = (0..32).map(|_| frame.clone()).collect();
for copy in &queued {
assert_eq!(copy.as_ptr(), frame.as_ptr(), "a recipient got its own copy of the frame");
assert_eq!(copy.len(), frame.len());
}
}
fn manager_with(overflow: Overflow) -> ConnectionManager<u32> {
let options = SessionOptions {
queues: Queues {
outbound: 1,
presence: 1,
..Queues::default()
},
overflow,
..SessionOptions::default()
};
ConnectionManager::with_options("test", None, &options)
}
#[tokio::test]
async fn a_full_client_queue_is_named_under_disconnect_and_only_counted_under_drop() {
for (overflow, expected) in [
(Overflow::default(), 0),
(
Overflow {
outbound: OutboundOverflow::Disconnect,
..Overflow::default()
},
1,
),
] {
let manager = manager_with(overflow);
let (tx, _rx) = session_channel(1);
let conn_id = manager.register(Agent::new_human(1u32), tx).await;
let frame: OutboundFrame = Frame::from(b"x".to_vec());
let _ = manager.broadcast(&MessageTarget::All, frame.clone());
let outcome = manager.broadcast(&MessageTarget::All, frame);
let named = outcome.map(|full| full.len()).unwrap_or(0);
assert_eq!(named, expected, "{overflow:?}");
if expected == 1 {
manager.disconnect_overflowed(vec![conn_id]).await;
assert_eq!(manager.connection_count(), 0, "the wedged client is gone");
} else {
assert_eq!(manager.connection_count(), 1, "dropping keeps the connection");
}
}
}
#[tokio::test]
async fn dropping_a_presence_event_is_counted_rather_than_waited_on() {
let manager = manager_with(Overflow::default());
manager.register(Agent::new_human(1u32), session_channel(4).0).await;
manager.register(Agent::new_human(2u32), session_channel(4).0).await;
assert_eq!(manager.stats().presence_dropped(), 1);
assert_eq!(manager.connection_count(), 2, "a lost announcement is not a lost connection");
}
#[tokio::test]
async fn every_way_of_setting_a_profile_moves_the_flag_the_frame_path_reads() {
let impaired = LinkProfile::symmetric(DirectionProfile::delayed(Duration::from_millis(50)));
let manager = manager_with(Overflow::default());
let conn_id = manager.register(Agent::new_human(1u32), session_channel(4).0).await;
let link = manager.link_handle(conn_id).expect("just registered");
assert!(!link.impaired(), "a fresh link is passthrough");
for set in [
&|m: &ConnectionManager<u32>, p| m.set_link_profile(1, p) as _,
&|m: &ConnectionManager<u32>, p| m.set_agent_link_profile(&1u32, p) as _,
&|m: &ConnectionManager<u32>, p| m.set_all_link_profiles(p) as _,
] as [&dyn Fn(&ConnectionManager<u32>, LinkProfile); 3]
{
set(&manager, impaired);
assert!(link.impaired(), "an impaired link says so");
assert_eq!(manager.link_profile(conn_id), Some(impaired));
set(&manager, LinkProfile::default());
assert!(!link.impaired(), "clearing it says so too");
}
}
#[tokio::test]
async fn one_broadcast_costs_one_buffer_and_a_pointer_per_recipient() {
let manager = manager_with(Overflow::default());
let mut inboxes = Vec::new();
for seat in 1..=4u32 {
let (tx, rx) = session_channel(4);
manager.register(Agent::new_human(seat), tx).await;
inboxes.push(rx);
}
let shared: OutboundFrame = Frame::from(vec![9u8; 512]);
manager.broadcast(&MessageTarget::All, shared.clone()).expect("fan out");
for inbox in &inboxes {
let queued = inbox.try_recv().expect("every recipient was sent the frame");
assert_eq!(queued.as_ptr(), shared.as_ptr(), "a recipient got its own copy");
}
let mut addressed = Vec::new();
for seat in 1..=4u32 {
let own: OutboundFrame = Frame::from(vec![seat as u8; 512]);
addressed.push(own.as_ptr());
manager.broadcast(&MessageTarget::Agent(seat), own).expect("one recipient");
}
for (index, inbox) in inboxes.iter().enumerate() {
let queued = inbox.try_recv().expect("the agent's own frame");
assert_eq!(queued.as_ptr(), addressed[index], "an addressed frame is that agent's");
}
}
#[tokio::test]
async fn a_registration_that_waits_still_counts_the_connection() {
let manager = Arc::new(manager_with(Overflow {
presence: PresenceOverflow::Backpressure,
..Overflow::default()
}));
manager.register(Agent::new_human(1u32), session_channel(4).0).await;
let waiting = {
let manager = manager.clone();
tokio::spawn(async move { manager.register(Agent::new_human(2u32), session_channel(4).0).await })
};
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(!waiting.is_finished(), "the announcement is still waiting");
assert_eq!(
manager.connection_count(),
2,
"a connection held at its announcement is still a connection"
);
waiting.abort();
}
#[tokio::test]
async fn disconnecting_a_client_never_waits_on_the_controller_that_caused_it() {
let manager = manager_with(Overflow {
outbound: OutboundOverflow::Disconnect,
presence: PresenceOverflow::Backpressure,
..Overflow::default()
});
let (tx, _rx) = session_channel(1);
let conn_id = manager.register(Agent::new_human(1u32), tx).await;
let frame: OutboundFrame = Frame::from(b"x".to_vec());
let _ = manager.broadcast(&MessageTarget::All, frame.clone());
let overflowed = manager
.broadcast(&MessageTarget::All, frame)
.expect("disconnecting reports rather than erroring");
assert_eq!(overflowed, vec![conn_id]);
tokio::time::timeout(Duration::from_secs(2), manager.disconnect_overflowed(overflowed))
.await
.expect("the departure is announced without waiting for a full presence queue");
assert_eq!(manager.connection_count(), 0);
assert_eq!(manager.stats().presence_dropped(), 1, "the lost Left is counted");
}
#[tokio::test]
async fn backpressure_waits_where_dropping_would_not() {
let manager = Arc::new(manager_with(Overflow {
presence: PresenceOverflow::Backpressure,
..Overflow::default()
}));
let presence = manager.take_presence();
manager.register(Agent::new_human(1u32), session_channel(4).0).await;
let waiting = {
let manager = manager.clone();
tokio::spawn(async move { manager.register(Agent::new_human(2u32), session_channel(4).0).await })
};
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(!waiting.is_finished(), "the second join is waiting on a full queue");
presence.recv().await.expect("the first join");
tokio::time::timeout(Duration::from_secs(1), waiting)
.await
.expect("draining one event releases the waiter")
.expect("the registration task");
assert_eq!(manager.stats().presence_dropped(), 0, "nothing was lost");
}
}