use std::collections::HashMap;
use std::num::NonZeroUsize;
use std::sync::Arc;
use broadcast_common::{Stage, Timestamp};
#[cfg(test)]
use bytes::Bytes;
use transmux::{Sample, TrackSpec};
use crate::trunk::{RetentionClass, Trunk, TrunkConfig, TrunkWriter};
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProgramId(pub u32);
pub const DEFAULT_MAX_PROGRAMS: NonZeroUsize = match NonZeroUsize::new(64) {
Some(n) => n,
None => panic!("64 is non-zero"),
};
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum SessionEvent {
Established,
NewProgram {
program: ProgramId,
tracks: Vec<TrackSpec>,
},
Sample {
program: ProgramId,
track_id: u32,
retention: RetentionClass,
sample: Sample,
},
TracksChanged {
program: ProgramId,
tracks: Vec<TrackSpec>,
},
}
pub trait IngestSession: for<'a> Stage<Out = SessionEvent> + Send {
type Request: Send;
fn poll_transmit(&mut self) -> Option<Self::Request> {
None
}
}
pub trait Dialer: Send {
type Session: IngestSession;
type Error;
fn dial(&mut self) -> Result<Self::Session, Self::Error>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SessionId(pub u64);
pub trait Listener: Send {
type Session: IngestSession;
type Error;
fn max_sessions(&self) -> usize;
fn poll_accept(&mut self) -> Result<Option<Self::Session>, Self::Error>;
}
#[derive(Debug)]
#[non_exhaustive]
pub enum HealthState<E> {
Establishing,
Live,
Ended,
Failed(E),
HandshakeTimedOut {
deadline: Timestamp,
},
}
impl<E: PartialEq> PartialEq for HealthState<E> {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(HealthState::Establishing, HealthState::Establishing) => true,
(HealthState::Live, HealthState::Live) => true,
(HealthState::Ended, HealthState::Ended) => true,
(HealthState::Failed(a), HealthState::Failed(b)) => a == b,
(
HealthState::HandshakeTimedOut { deadline: a },
HealthState::HandshakeTimedOut { deadline: b },
) => a == b,
_ => false,
}
}
}
impl<E> HealthState<E> {
pub fn is_running(&self) -> bool {
matches!(self, HealthState::Establishing | HealthState::Live)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct HandshakePolicy {
pub establish_by: Timestamp,
}
impl HandshakePolicy {
pub fn establish_by(establish_by: Timestamp) -> Self {
HandshakePolicy { establish_by }
}
}
pub struct IngestDriver<S: IngestSession> {
session: S,
trunk_config: TrunkConfig,
handshake: HandshakePolicy,
max_programs: NonZeroUsize,
programs: HashMap<ProgramId, Arc<Trunk>>,
writers: HashMap<ProgramId, TrunkWriter>,
refused_programs: u64,
health: HealthState<S::Error>,
}
impl<S: IngestSession> IngestDriver<S> {
pub fn new(
session: S,
trunk_config: TrunkConfig,
handshake: HandshakePolicy,
max_programs: NonZeroUsize,
) -> Self {
IngestDriver {
session,
trunk_config,
handshake,
max_programs,
programs: HashMap::new(),
writers: HashMap::new(),
refused_programs: 0,
health: HealthState::Establishing,
}
}
pub fn max_programs(&self) -> NonZeroUsize {
self.max_programs
}
pub fn program_count(&self) -> usize {
self.programs.len()
}
pub fn refused_program_count(&self) -> u64 {
self.refused_programs
}
pub fn feed(&mut self, input: S::In<'_>, now: Timestamp) {
if !self.health.is_running() {
return;
}
match self.session.feed(input, now) {
Ok(()) => {
self.drain();
self.check_handshake_deadline(now);
}
Err(e) => self.health = HealthState::Failed(e),
}
}
pub fn on_deadline(&mut self, now: Timestamp) {
if !self.health.is_running() {
return;
}
self.session.on_deadline(now);
self.drain();
self.check_handshake_deadline(now);
}
pub fn finish(&mut self) {
if !self.health.is_running() {
return;
}
match self.session.finish() {
Ok(()) => {
self.drain();
self.health = HealthState::Ended;
}
Err(e) => self.health = HealthState::Failed(e),
}
}
pub fn poll_transmit(&mut self) -> Option<S::Request> {
self.session.poll_transmit()
}
pub fn next_deadline(&self) -> Option<Timestamp> {
let session = self.session.next_deadline();
let handshake =
matches!(self.health, HealthState::Establishing).then_some(self.handshake.establish_by);
match (session, handshake) {
(Some(a), Some(b)) => Some(a.min(b)),
(a, b) => a.or(b),
}
}
pub fn health(&self) -> &HealthState<S::Error> {
&self.health
}
pub fn into_health(self) -> HealthState<S::Error> {
self.health
}
fn check_handshake_deadline(&mut self, now: Timestamp) {
if matches!(self.health, HealthState::Establishing) && now >= self.handshake.establish_by {
self.health = HealthState::HandshakeTimedOut {
deadline: self.handshake.establish_by,
};
}
}
pub fn trunk(&self, program: ProgramId) -> Option<&Arc<Trunk>> {
self.programs.get(&program)
}
pub fn programs(&self) -> impl Iterator<Item = ProgramId> + '_ {
self.programs.keys().copied()
}
pub fn session(&self) -> &S {
&self.session
}
fn drain(&mut self) {
while let Some(event) = self.session.poll() {
match event {
SessionEvent::Established => {
if matches!(self.health, HealthState::Establishing) {
self.health = HealthState::Live;
}
}
SessionEvent::NewProgram { program, tracks } => {
if !self.programs.contains_key(&program)
&& self.programs.len() >= self.max_programs.get()
{
self.refused_programs += 1;
continue;
}
if let Some(writer) = self.writers.get(&program) {
writer.set_tracks(tracks);
continue;
}
let trunk = Trunk::new(self.trunk_config);
let writer = trunk
.writer()
.expect("a freshly constructed Trunk always has an unclaimed writer");
writer.set_tracks(tracks);
self.programs.insert(program, trunk);
self.writers.insert(program, writer);
}
SessionEvent::Sample {
program,
track_id,
retention,
sample,
} => {
if let Some(writer) = self.writers.get(&program) {
writer.publish(track_id, retention, sample);
}
}
SessionEvent::TracksChanged { program, tracks } => {
if let Some(writer) = self.writers.get(&program) {
writer.set_tracks(tracks);
}
}
}
}
}
}
pub fn run_dial<D: Dialer>(
dialer: &mut D,
trunk_config: TrunkConfig,
handshake: HandshakePolicy,
max_programs: NonZeroUsize,
) -> Result<IngestDriver<D::Session>, D::Error> {
let session = dialer.dial()?;
Ok(IngestDriver::new(
session,
trunk_config,
handshake,
max_programs,
))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct ReconnectPolicy {
pub max_attempts: u32,
}
impl ReconnectPolicy {
pub fn new(max_attempts: u32) -> Self {
assert!(max_attempts > 0, "ReconnectPolicy max_attempts must be > 0");
ReconnectPolicy { max_attempts }
}
}
#[non_exhaustive]
pub enum DialAttempt<S: IngestSession, E> {
Connected(IngestDriver<S>),
Retry(E),
GaveUp(E),
Exhausted,
}
pub struct DialSupervisor<D: Dialer> {
dialer: D,
policy: ReconnectPolicy,
attempts: u32,
exhausted: bool,
}
impl<D: Dialer> DialSupervisor<D> {
pub fn new(dialer: D, policy: ReconnectPolicy) -> Self {
DialSupervisor {
dialer,
policy,
attempts: 0,
exhausted: false,
}
}
pub fn attempts(&self) -> u32 {
self.attempts
}
pub fn is_exhausted(&self) -> bool {
self.exhausted
}
pub fn try_dial(
&mut self,
trunk_config: TrunkConfig,
handshake: HandshakePolicy,
max_programs: NonZeroUsize,
) -> DialAttempt<D::Session, D::Error> {
if self.exhausted {
return DialAttempt::Exhausted;
}
self.attempts += 1;
match self.dialer.dial() {
Ok(session) => {
self.attempts = 0;
DialAttempt::Connected(IngestDriver::new(
session,
trunk_config,
handshake,
max_programs,
))
}
Err(e) => {
if self.attempts >= self.policy.max_attempts {
self.exhausted = true;
DialAttempt::GaveUp(e)
} else {
DialAttempt::Retry(e)
}
}
}
}
}
#[derive(Debug)]
#[non_exhaustive]
pub enum AcceptOutcome<E> {
Admitted(SessionId),
Idle,
Refused,
Error(E),
}
pub struct ListenDriver<L: Listener> {
listener: L,
trunk_config: TrunkConfig,
handshake: HandshakePolicy,
max_programs: NonZeroUsize,
sessions: HashMap<SessionId, IngestDriver<L::Session>>,
next_id: u64,
}
impl<L: Listener> ListenDriver<L> {
pub fn new(
listener: L,
trunk_config: TrunkConfig,
handshake: HandshakePolicy,
max_programs: NonZeroUsize,
) -> Self {
ListenDriver {
listener,
trunk_config,
handshake,
max_programs,
sessions: HashMap::new(),
next_id: 0,
}
}
pub fn max_programs(&self) -> NonZeroUsize {
self.max_programs
}
pub fn refused_program_count(&self, id: SessionId) -> Option<u64> {
self.sessions
.get(&id)
.map(IngestDriver::refused_program_count)
}
pub fn session_count(&self) -> usize {
self.sessions.len()
}
pub fn max_sessions(&self) -> usize {
self.listener.max_sessions()
}
pub fn poll_accept(&mut self) -> AcceptOutcome<L::Error> {
match self.listener.poll_accept() {
Ok(None) => AcceptOutcome::Idle,
Ok(Some(session)) => {
if self.sessions.len() >= self.listener.max_sessions() {
drop(session);
AcceptOutcome::Refused
} else {
let id = SessionId(self.next_id);
self.next_id += 1;
self.sessions.insert(
id,
IngestDriver::new(
session,
self.trunk_config,
self.handshake,
self.max_programs,
),
);
AcceptOutcome::Admitted(id)
}
}
Err(e) => AcceptOutcome::Error(e),
}
}
pub fn feed(
&mut self,
id: SessionId,
input: &[u8],
now: Timestamp,
) -> Option<HealthState<<L::Session as Stage>::Error>>
where
L::Session: for<'a> Stage<In<'a> = &'a [u8]>,
{
self.drive(id, |d| d.feed(input, now))
}
pub fn on_deadline(
&mut self,
id: SessionId,
now: Timestamp,
) -> Option<HealthState<<L::Session as Stage>::Error>> {
self.drive(id, |d| d.on_deadline(now))
}
pub fn finish(&mut self, id: SessionId) -> Option<HealthState<<L::Session as Stage>::Error>> {
self.drive(id, IngestDriver::finish)
}
pub fn health(&self, id: SessionId) -> Option<&HealthState<<L::Session as Stage>::Error>> {
self.sessions.get(&id).map(IngestDriver::health)
}
pub fn trunk(&self, id: SessionId, program: ProgramId) -> Option<&Arc<Trunk>> {
self.sessions.get(&id).and_then(|d| d.trunk(program))
}
pub fn driver(&self, id: SessionId) -> Option<&IngestDriver<L::Session>> {
self.sessions.get(&id)
}
pub fn driver_mut(&mut self, id: SessionId) -> Option<&mut IngestDriver<L::Session>> {
self.sessions.get_mut(&id)
}
pub fn reap_if_terminal(
&mut self,
id: SessionId,
) -> Option<HealthState<<L::Session as Stage>::Error>> {
let driver = self.sessions.get(&id)?;
if driver.health().is_running() {
None
} else {
self.sessions.remove(&id).map(|d| d.health)
}
}
fn drive(
&mut self,
id: SessionId,
op: impl FnOnce(&mut IngestDriver<L::Session>),
) -> Option<HealthState<<L::Session as Stage>::Error>> {
let driver = self.sessions.get_mut(&id)?;
op(driver);
if driver.health().is_running() {
None
} else {
self.sessions.remove(&id).map(|d| d.health)
}
}
}
pub fn run_listen<L: Listener>(
listener: L,
trunk_config: TrunkConfig,
handshake: HandshakePolicy,
max_programs: NonZeroUsize,
) -> ListenDriver<L> {
ListenDriver::new(listener, trunk_config, handshake, max_programs)
}
#[cfg(test)]
mod tests {
use super::*;
use broadcast_common::Demand;
use std::collections::VecDeque;
use transmux::pipeline::{CodecConfig, DataCarriage};
fn nz(n: usize) -> std::num::NonZeroUsize {
std::num::NonZeroUsize::new(n).expect("test capacity must be non-zero")
}
fn trunk_config() -> TrunkConfig {
TrunkConfig::new(nz(64), nz(16), nz(8), nz(8), nz(8))
}
fn handshake() -> HandshakePolicy {
HandshakePolicy::establish_by(Timestamp::from_nanos(u64::MAX))
}
fn max_programs() -> std::num::NonZeroUsize {
nz(1024)
}
fn sample(byte: u8) -> Sample {
Sample::new(Bytes::from(vec![byte; 4]), Some(0), Some(0), Some(1), true)
}
fn opaque_track(track_id: u32) -> TrackSpec {
TrackSpec::new(
track_id,
90_000,
CodecConfig::Data {
stream_type: 0x06,
descriptors: Vec::new(),
carriage: DataCarriage::Pes,
},
)
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct FakeError(&'static str);
enum FeedOutcome {
Events(Vec<SessionEvent>),
Err(FakeError),
}
struct ScriptedSession {
script: VecDeque<FeedOutcome>,
pending: VecDeque<SessionEvent>,
finish_outcome: Result<(), FakeError>,
}
impl ScriptedSession {
fn new(script: Vec<FeedOutcome>) -> Self {
ScriptedSession {
script: script.into(),
pending: VecDeque::from(vec![SessionEvent::Established]),
finish_outcome: Ok(()),
}
}
fn failing_finish(mut self, err: FakeError) -> Self {
self.finish_outcome = Err(err);
self
}
}
impl Stage for ScriptedSession {
type In<'a> = &'a [u8];
type Out = SessionEvent;
type Error = FakeError;
fn feed(&mut self, _input: &[u8], _now: Timestamp) -> Result<(), FakeError> {
match self.script.pop_front() {
Some(FeedOutcome::Events(evs)) => {
self.pending.extend(evs);
Ok(())
}
Some(FeedOutcome::Err(e)) => Err(e),
None => Ok(()),
}
}
fn poll(&mut self) -> Option<SessionEvent> {
self.pending.pop_front()
}
fn finish(&mut self) -> Result<(), FakeError> {
self.finish_outcome.clone()
}
fn next_deadline(&self) -> Option<Timestamp> {
None
}
fn on_deadline(&mut self, _now: Timestamp) {}
fn demand(&self) -> Demand {
Demand::new(4096)
}
}
impl IngestSession for ScriptedSession {
type Request = Bytes;
}
struct ScriptedDialer {
sessions: VecDeque<ScriptedSession>,
fail_with: FakeError,
}
impl Dialer for ScriptedDialer {
type Session = ScriptedSession;
type Error = FakeError;
fn dial(&mut self) -> Result<ScriptedSession, FakeError> {
self.sessions
.pop_front()
.ok_or_else(|| self.fail_with.clone())
}
}
#[test]
fn run_dial_drives_fake_session_end_to_end_samples_land_in_trunk() {
let session = ScriptedSession::new(vec![
FeedOutcome::Events(vec![SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(7)],
}]),
FeedOutcome::Events(vec![SessionEvent::Sample {
program: ProgramId(1),
track_id: 7,
retention: RetentionClass::Timed,
sample: sample(0xAB),
}]),
]);
let mut dialer = ScriptedDialer {
sessions: VecDeque::from(vec![session]),
fail_with: FakeError("unused"),
};
let mut driver = run_dial(&mut dialer, trunk_config(), handshake(), max_programs())
.expect("fake dial succeeds");
let trunk_before = driver.trunk(ProgramId(1)).cloned();
assert!(
trunk_before.is_none(),
"no Trunk before NewProgram is announced"
);
driver.feed(b"pat", Timestamp::ZERO);
let trunk = driver
.trunk(ProgramId(1))
.cloned()
.expect("NewProgram announced a Trunk for program 1");
let mut cursor = trunk.subscribe();
driver.feed(b"pes", Timestamp::from_nanos(1));
let item = cursor.poll().expect("the published sample is on the ring");
match item {
crate::SampleCursorItem::Timed { track_id, sample } => {
assert_eq!(track_id, 7);
assert_eq!(sample.data.as_ref(), &[0xAB; 4]);
}
other => panic!("expected Timed, got {other:?}"),
}
}
#[test]
fn new_program_seeds_the_trunk_with_exactly_the_announced_tracks() {
let session =
ScriptedSession::new(vec![FeedOutcome::Events(vec![SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(3), opaque_track(9)],
}])]);
let mut dialer = ScriptedDialer {
sessions: VecDeque::from(vec![session]),
fail_with: FakeError("unused"),
};
let mut driver = run_dial(&mut dialer, trunk_config(), handshake(), max_programs())
.expect("fake dial succeeds");
driver.feed(b"pat", Timestamp::ZERO);
let trunk = driver
.trunk(ProgramId(1))
.cloned()
.expect("NewProgram announced a Trunk for program 1");
let track_ids: Vec<u32> = trunk.tracks().iter().map(|t| t.track_id).collect();
assert_eq!(
track_ids,
vec![3, 9],
"the Trunk must expose exactly the tracks NewProgram carried"
);
}
#[test]
fn tracks_changed_replaces_the_set_and_bumps_generation() {
let session = ScriptedSession::new(vec![
FeedOutcome::Events(vec![SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(1)],
}]),
FeedOutcome::Events(vec![SessionEvent::TracksChanged {
program: ProgramId(1),
tracks: vec![opaque_track(1), opaque_track(2)],
}]),
]);
let mut dialer = ScriptedDialer {
sessions: VecDeque::from(vec![session]),
fail_with: FakeError("unused"),
};
let mut driver = run_dial(&mut dialer, trunk_config(), handshake(), max_programs())
.expect("fake dial succeeds");
driver.feed(b"pat", Timestamp::from_nanos(0));
let trunk = driver.trunk(ProgramId(1)).cloned().unwrap();
assert_eq!(
trunk.track_generation(),
1,
"NewProgram's seed counts as the first set_tracks call"
);
driver.feed(b"pmt-bump", Timestamp::from_nanos(1));
let track_ids: Vec<u32> = trunk.tracks().iter().map(|t| t.track_id).collect();
assert_eq!(
track_ids,
vec![1, 2],
"TracksChanged must replace the set with the new complete snapshot"
);
assert_eq!(
trunk.track_generation(),
2,
"TracksChanged must bump the generation exactly once"
);
}
#[test]
fn tracks_changed_for_an_unannounced_program_is_dropped_not_panicking_and_mints_nothing() {
let session = ScriptedSession::new(vec![FeedOutcome::Events(vec![
SessionEvent::TracksChanged {
program: ProgramId(1),
tracks: vec![opaque_track(1)],
},
])]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), max_programs());
driver.feed(b"stray", Timestamp::ZERO);
assert!(
driver.trunk(ProgramId(1)).is_none(),
"TracksChanged must never mint a Trunk on its own"
);
assert_eq!(
driver.program_count(),
0,
"an unannounced program must not be admitted by TracksChanged"
);
}
#[test]
fn track_generation_is_stable_when_nothing_changes() {
let session = ScriptedSession::new(vec![
FeedOutcome::Events(vec![SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(1)],
}]),
FeedOutcome::Events(vec![SessionEvent::Sample {
program: ProgramId(1),
track_id: 1,
retention: RetentionClass::Timed,
sample: sample(0xAB),
}]),
FeedOutcome::Events(vec![]),
]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), max_programs());
driver.feed(b"1", Timestamp::from_nanos(0));
let trunk = driver.trunk(ProgramId(1)).cloned().unwrap();
assert_eq!(trunk.track_generation(), 1);
driver.feed(b"2", Timestamp::from_nanos(1));
driver.feed(b"3", Timestamp::from_nanos(2));
assert_eq!(
trunk.track_generation(),
1,
"samples and no-op feeds must never bump track_generation"
);
}
#[test]
fn repeat_new_program_updates_in_place_and_does_not_strand_subscribers() {
let session = ScriptedSession::new(vec![
FeedOutcome::Events(vec![SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(7)],
}]),
FeedOutcome::Events(vec![SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(7), opaque_track(8)],
}]),
FeedOutcome::Events(vec![SessionEvent::Sample {
program: ProgramId(1),
track_id: 7,
retention: RetentionClass::Timed,
sample: sample(0xCD),
}]),
]);
let mut dialer = ScriptedDialer {
sessions: VecDeque::from(vec![session]),
fail_with: FakeError("unused"),
};
let mut driver = run_dial(&mut dialer, trunk_config(), handshake(), max_programs())
.expect("fake dial succeeds");
driver.feed(b"pat", Timestamp::ZERO);
let trunk = driver
.trunk(ProgramId(1))
.cloned()
.expect("NewProgram announced a Trunk for program 1");
let mut cursor = trunk.subscribe();
driver.feed(b"pat-again", Timestamp::from_nanos(1));
assert_eq!(
driver.program_count(),
1,
"a repeat announcement must not add a program"
);
let track_ids: Vec<u32> = trunk.tracks().iter().map(|t| t.track_id).collect();
assert_eq!(
track_ids,
vec![7, 8],
"the re-announcement's track set must land on the SAME Trunk the \
subscriber already holds"
);
driver.feed(b"pes", Timestamp::from_nanos(2));
let item = cursor
.poll()
.expect("a cursor subscribed before the re-announcement must still receive samples");
match item {
crate::SampleCursorItem::Timed { track_id, sample } => {
assert_eq!(track_id, 7);
assert_eq!(sample.data.as_ref(), &[0xCD; 4]);
}
other => panic!("expected Timed, got {other:?}"),
}
}
#[test]
fn clean_finish_yields_ended_not_failed() {
let session = ScriptedSession::new(vec![]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), max_programs());
assert!(
matches!(driver.health(), HealthState::Establishing),
"a freshly dialled session has not established yet"
);
driver.finish();
assert!(
matches!(driver.health(), HealthState::Ended),
"a session that finished cleanly must be Ended, not Failed: {:?}",
driver.health()
);
}
#[test]
fn erroring_feed_yields_failed_not_ended() {
let session = ScriptedSession::new(vec![FeedOutcome::Err(FakeError("bad continuity"))]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), max_programs());
driver.feed(b"garbage", Timestamp::ZERO);
match driver.health() {
HealthState::Failed(FakeError(reason)) => assert_eq!(*reason, "bad continuity"),
other => panic!("expected Failed(\"bad continuity\"), got {other:?}"),
}
}
#[test]
fn erroring_finish_yields_failed_not_ended() {
let session = ScriptedSession::new(vec![]).failing_finish(FakeError("truncated tail"));
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), max_programs());
driver.finish();
match driver.health() {
HealthState::Failed(FakeError(reason)) => assert_eq!(*reason, "truncated tail"),
other => panic!("expected Failed(\"truncated tail\"), got {other:?}"),
}
}
#[test]
fn terminated_driver_ignores_further_feed_and_finish() {
let session = ScriptedSession::new(vec![FeedOutcome::Err(FakeError("boom"))]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), max_programs());
driver.feed(b"x", Timestamp::ZERO);
assert!(matches!(driver.health(), HealthState::Failed(_)));
driver.finish();
assert!(matches!(driver.health(), HealthState::Failed(_)));
}
#[test]
fn one_connection_two_programs_yields_two_trunks_including_one_announced_late() {
let session = ScriptedSession::new(vec![
FeedOutcome::Events(vec![
SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(1)],
},
SessionEvent::Sample {
program: ProgramId(1),
track_id: 1,
retention: RetentionClass::Timed,
sample: sample(0x01),
},
]),
FeedOutcome::Events(vec![]),
FeedOutcome::Events(vec![
SessionEvent::NewProgram {
program: ProgramId(2),
tracks: vec![opaque_track(9)],
},
SessionEvent::Sample {
program: ProgramId(2),
track_id: 9,
retention: RetentionClass::Timed,
sample: sample(0x02),
},
]),
]);
let mut dialer = ScriptedDialer {
sessions: VecDeque::from(vec![session]),
fail_with: FakeError("unused"),
};
let mut driver =
run_dial(&mut dialer, trunk_config(), handshake(), max_programs()).unwrap();
driver.feed(b"1", Timestamp::from_nanos(0));
assert!(driver.trunk(ProgramId(1)).is_some());
assert!(
driver.trunk(ProgramId(2)).is_none(),
"program 2 must not exist before it is announced"
);
driver.feed(b"2", Timestamp::from_nanos(1));
assert!(
driver.trunk(ProgramId(2)).is_none(),
"a no-op feed must not fabricate a program"
);
driver.feed(b"3", Timestamp::from_nanos(2));
let trunk1 = driver.trunk(ProgramId(1)).cloned().unwrap();
let trunk2 = driver
.trunk(ProgramId(2))
.cloned()
.expect("program 2 announced mid-session must get its own Trunk");
assert!(
!Arc::ptr_eq(&trunk1, &trunk2),
"each program must get a genuinely distinct Trunk"
);
let mut programs: Vec<_> = driver.programs().collect();
programs.sort();
assert_eq!(programs, vec![ProgramId(1), ProgramId(2)]);
assert_eq!(trunk1.timed_len(), 1);
assert_eq!(trunk2.timed_len(), 1);
}
#[test]
fn programs_up_to_max_get_a_trunk_each_the_next_one_is_refused_and_reported() {
let cap = 2;
let session = ScriptedSession::new(vec![FeedOutcome::Events(vec![
SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(1)],
},
SessionEvent::NewProgram {
program: ProgramId(2),
tracks: vec![opaque_track(2)],
},
SessionEvent::NewProgram {
program: ProgramId(3),
tracks: vec![opaque_track(3)],
},
])]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), nz(cap));
driver.feed(b"pat", Timestamp::ZERO);
assert!(driver.trunk(ProgramId(1)).is_some(), "program 1 admitted");
assert!(driver.trunk(ProgramId(2)).is_some(), "program 2 admitted");
assert!(
driver.trunk(ProgramId(3)).is_none(),
"the (cap+1)th program must be refused a Trunk"
);
assert_eq!(
driver.program_count(),
cap,
"admitted program count must sit exactly at the cap, not above it"
);
assert_eq!(
driver.refused_program_count(),
1,
"the refusal must be reported via a queryable counter, never silent"
);
}
#[test]
fn repeat_announcement_of_an_already_admitted_program_is_never_refused() {
let cap = 1;
let session = ScriptedSession::new(vec![FeedOutcome::Events(vec![
SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(1)],
},
SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(1)],
},
])]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), nz(cap));
driver.feed(b"pat", Timestamp::ZERO);
assert_eq!(driver.program_count(), 1);
assert_eq!(
driver.refused_program_count(),
0,
"re-announcing an already-admitted program must not be refused"
);
}
#[test]
fn refusal_does_not_disturb_already_admitted_programs_their_samples_keep_flowing() {
let cap = 1;
let session = ScriptedSession::new(vec![
FeedOutcome::Events(vec![SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(1)],
}]),
FeedOutcome::Events(vec![
SessionEvent::NewProgram {
program: ProgramId(2),
tracks: vec![opaque_track(2)],
},
SessionEvent::Sample {
program: ProgramId(1),
track_id: 1,
retention: RetentionClass::Timed,
sample: sample(0x01),
},
]),
]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), nz(cap));
driver.feed(b"1", Timestamp::from_nanos(0));
let trunk1 = driver.trunk(ProgramId(1)).cloned().unwrap();
let mut cursor = trunk1.subscribe();
driver.feed(b"2", Timestamp::from_nanos(1));
assert!(
driver.trunk(ProgramId(2)).is_none(),
"program 2 must be refused, not given a Trunk"
);
assert_eq!(driver.refused_program_count(), 1);
match cursor
.poll()
.expect("program 1's sample must still land despite program 2 being refused")
{
crate::SampleCursorItem::Timed { track_id, sample } => {
assert_eq!(track_id, 1);
assert_eq!(sample.data.as_ref(), &[0x01; 4]);
}
other => panic!("expected Timed, got {other:?}"),
}
assert!(
matches!(driver.health(), HealthState::Live),
"refusing an extra program must not fail the session: {:?}",
driver.health()
);
}
#[test]
fn newprogram_flood_is_bounded_admits_exactly_max_programs_and_allocates_no_more_trunks() {
let cap = 3;
let mut events = Vec::with_capacity(10_000);
for i in 0..10_000u32 {
events.push(SessionEvent::NewProgram {
program: ProgramId(i),
tracks: vec![opaque_track(i)],
});
}
let session = ScriptedSession::new(vec![FeedOutcome::Events(events)]);
let mut driver = IngestDriver::new(session, trunk_config(), handshake(), nz(cap));
driver.feed(b"flood", Timestamp::ZERO);
assert_eq!(
driver.program_count(),
cap,
"a 10,000-program flood must admit exactly max_programs Trunks, never more"
);
assert_eq!(
driver.refused_program_count(),
10_000 - cap as u64,
"every program past the cap must be counted as refused"
);
}
struct FloodingListener {
max_sessions: usize,
}
impl Listener for FloodingListener {
type Session = ScriptedSession;
type Error = FakeError;
fn max_sessions(&self) -> usize {
self.max_sessions
}
fn poll_accept(&mut self) -> Result<Option<ScriptedSession>, FakeError> {
Ok(Some(ScriptedSession::new(vec![])))
}
}
#[test]
fn run_listen_admits_up_to_max_sessions_then_refuses_and_stays_bounded() {
let max_sessions = 3;
let mut driver = run_listen(
FloodingListener { max_sessions },
trunk_config(),
handshake(),
max_programs(),
);
for _ in 0..max_sessions {
assert!(matches!(driver.poll_accept(), AcceptOutcome::Admitted(_)));
}
assert_eq!(driver.session_count(), max_sessions);
for _ in 0..10_000 {
assert!(matches!(driver.poll_accept(), AcceptOutcome::Refused));
assert!(
driver.session_count() <= max_sessions,
"session_count grew past max_sessions under flood"
);
}
assert_eq!(driver.session_count(), max_sessions);
}
#[test]
fn ended_session_is_reaped_freeing_a_slot() {
let mut driver = run_listen(
FloodingListener { max_sessions: 1 },
trunk_config(),
handshake(),
max_programs(),
);
let AcceptOutcome::Admitted(id) = driver.poll_accept() else {
panic!("expected admission");
};
assert_eq!(driver.session_count(), 1);
assert!(matches!(driver.poll_accept(), AcceptOutcome::Refused));
let health = driver.finish(id).expect("finish terminates the session");
assert!(matches!(health, HealthState::Ended));
assert_eq!(
driver.session_count(),
0,
"a terminated session must be reaped, freeing its slot"
);
assert!(
driver.health(id).is_none(),
"a reaped session is no longer queryable by id"
);
assert!(matches!(driver.poll_accept(), AcceptOutcome::Admitted(_)));
}
#[test]
fn driver_mut_and_reap_if_terminal_mirror_feed_semantics() {
let mut driver = run_listen(
FloodingListener { max_sessions: 1 },
trunk_config(),
handshake(),
max_programs(),
);
let AcceptOutcome::Admitted(id) = driver.poll_accept() else {
panic!("expected admission");
};
assert!(
matches!(
driver.driver(id).map(IngestDriver::health),
Some(HealthState::Establishing)
),
"a freshly admitted session must be Establishing, observable via driver()"
);
driver
.driver_mut(id)
.expect("session just admitted")
.feed(b"reply", Timestamp::from_nanos(1));
assert!(
matches!(
driver.driver(id).map(IngestDriver::health),
Some(HealthState::Live)
),
"driver_mut's feed must reach the session exactly like Self::feed would"
);
assert_eq!(driver.session_count(), 1, "not reaped: still Live");
driver.driver_mut(id).expect("still admitted").finish();
assert!(
matches!(
driver.driver(id).map(IngestDriver::health),
Some(HealthState::Ended)
),
"finish() through driver_mut must be observable via driver() before reaping"
);
assert_eq!(driver.session_count(), 1, "not yet reaped");
let health = driver
.reap_if_terminal(id)
.expect("a terminal session must be reaped");
assert!(matches!(health, HealthState::Ended));
assert_eq!(
driver.session_count(),
0,
"reap_if_terminal must free the slot"
);
assert!(driver.driver(id).is_none());
assert!(driver.reap_if_terminal(id).is_none());
}
#[test]
fn permanently_failing_dial_is_bounded_and_does_not_spin_or_grow() {
let dialer = ScriptedDialer {
sessions: VecDeque::new(),
fail_with: FakeError("connection refused"),
};
let mut supervisor = DialSupervisor::new(dialer, ReconnectPolicy::new(3));
assert!(matches!(
supervisor.try_dial(trunk_config(), handshake(), max_programs()),
DialAttempt::Retry(_)
));
assert_eq!(supervisor.attempts(), 1);
assert!(matches!(
supervisor.try_dial(trunk_config(), handshake(), max_programs()),
DialAttempt::Retry(_)
));
assert_eq!(supervisor.attempts(), 2);
assert!(matches!(
supervisor.try_dial(trunk_config(), handshake(), max_programs()),
DialAttempt::GaveUp(_)
));
assert_eq!(supervisor.attempts(), 3);
assert!(supervisor.is_exhausted());
for _ in 0..10_000 {
assert!(matches!(
supervisor.try_dial(trunk_config(), handshake(), max_programs()),
DialAttempt::Exhausted
));
assert_eq!(
supervisor.attempts(),
3,
"attempts must not grow past max_attempts under flood"
);
}
}
#[test]
fn dial_supervisor_succeeds_within_the_bound_and_resets_attempts() {
let good_session = ScriptedSession::new(vec![]);
let dialer = ScriptedDialer {
sessions: VecDeque::from(vec![good_session]),
fail_with: FakeError("refused"),
};
let mut supervisor = DialSupervisor::new(dialer, ReconnectPolicy::new(2));
match supervisor.try_dial(trunk_config(), handshake(), max_programs()) {
DialAttempt::Connected(_) => {}
DialAttempt::Retry(_) => panic!("expected Connected, got Retry"),
DialAttempt::GaveUp(_) => panic!("expected Connected, got GaveUp"),
DialAttempt::Exhausted => panic!("expected Connected, got Exhausted"),
}
assert_eq!(supervisor.attempts(), 0);
assert!(!supervisor.is_exhausted());
}
struct HandshakeSession {
step: usize,
outbound: VecDeque<Bytes>,
pending: VecDeque<SessionEvent>,
}
const HANDSHAKE_REQUESTS: [&[u8]; 3] = [b"DESCRIBE", b"SETUP", b"PLAY"];
impl HandshakeSession {
fn new() -> Self {
HandshakeSession {
step: 0,
outbound: VecDeque::from(vec![Bytes::from_static(HANDSHAKE_REQUESTS[0])]),
pending: VecDeque::new(),
}
}
}
impl Stage for HandshakeSession {
type In<'a> = &'a [u8];
type Out = SessionEvent;
type Error = FakeError;
fn feed(&mut self, input: &[u8], _now: Timestamp) -> Result<(), FakeError> {
if self.step >= HANDSHAKE_REQUESTS.len() {
self.pending.push_back(SessionEvent::Sample {
program: ProgramId(1),
track_id: 7,
retention: RetentionClass::Timed,
sample: sample(0xAB),
});
return Ok(());
}
let expected = format!(
"200 {}",
String::from_utf8_lossy(HANDSHAKE_REQUESTS[self.step])
);
if input != expected.as_bytes() {
return Err(FakeError("handshake reply out of sequence"));
}
self.step += 1;
match HANDSHAKE_REQUESTS.get(self.step) {
Some(next) => self.outbound.push_back(Bytes::from_static(next)),
None => {
self.pending.push_back(SessionEvent::Established);
self.pending.push_back(SessionEvent::NewProgram {
program: ProgramId(1),
tracks: vec![opaque_track(7)],
});
}
}
Ok(())
}
fn poll(&mut self) -> Option<SessionEvent> {
self.pending.pop_front()
}
fn finish(&mut self) -> Result<(), FakeError> {
Ok(())
}
fn next_deadline(&self) -> Option<Timestamp> {
None
}
fn on_deadline(&mut self, _now: Timestamp) {}
fn demand(&self) -> Demand {
Demand::new(4096)
}
}
impl IngestSession for HandshakeSession {
type Request = Bytes;
fn poll_transmit(&mut self) -> Option<Bytes> {
self.outbound.pop_front()
}
}
struct HandshakeDialer;
impl Dialer for HandshakeDialer {
type Session = HandshakeSession;
type Error = FakeError;
fn dial(&mut self) -> Result<HandshakeSession, FakeError> {
Ok(HandshakeSession::new())
}
}
#[test]
fn multi_round_trip_handshake_completes_through_feed_and_poll_transmit_only() {
let mut dialer = HandshakeDialer;
let mut driver = run_dial(&mut dialer, trunk_config(), handshake(), max_programs())
.expect("dial constructs a session");
assert!(
matches!(driver.health(), HealthState::Establishing),
"dial() must not establish the session: {:?}",
driver.health()
);
let mut wire: Vec<Bytes> = Vec::new();
let mut now = 0u64;
for _ in 0..HANDSHAKE_REQUESTS.len() {
let req = driver
.poll_transmit()
.expect("the session has a handshake request to send");
assert!(
driver.poll_transmit().is_none(),
"one request in flight at a time"
);
wire.push(req.clone());
let reply = format!("200 {}", String::from_utf8_lossy(&req));
now += 1;
driver.feed(reply.as_bytes(), Timestamp::from_nanos(now));
}
let sent: Vec<&[u8]> = wire.iter().map(|b| b.as_ref()).collect();
assert_eq!(sent, HANDSHAKE_REQUESTS, "handshake request sequence");
assert!(
matches!(driver.health(), HealthState::Live),
"after the final handshake reply the session must be Live: {:?}",
driver.health()
);
let trunk = driver
.trunk(ProgramId(1))
.cloned()
.expect("the handshake announced program 1");
let mut cursor = trunk.subscribe();
driver.feed(b"media", Timestamp::from_nanos(now + 1));
match cursor.poll().expect("post-handshake sample on the ring") {
crate::SampleCursorItem::Timed { track_id, .. } => assert_eq!(track_id, 7),
other => panic!("expected Timed, got {other:?}"),
}
}
struct StallingSession {
outbound: VecDeque<Bytes>,
}
impl Stage for StallingSession {
type In<'a> = &'a [u8];
type Out = SessionEvent;
type Error = FakeError;
fn feed(&mut self, _input: &[u8], _now: Timestamp) -> Result<(), FakeError> {
Ok(()) }
fn poll(&mut self) -> Option<SessionEvent> {
None
}
fn finish(&mut self) -> Result<(), FakeError> {
Ok(())
}
fn next_deadline(&self) -> Option<Timestamp> {
None
}
fn on_deadline(&mut self, _now: Timestamp) {}
fn demand(&self) -> Demand {
Demand::new(4096)
}
}
impl IngestSession for StallingSession {
type Request = Bytes;
fn poll_transmit(&mut self) -> Option<Bytes> {
self.outbound.pop_front()
}
}
struct StallingListener {
max_sessions: usize,
}
impl Listener for StallingListener {
type Session = StallingSession;
type Error = FakeError;
fn max_sessions(&self) -> usize {
self.max_sessions
}
fn poll_accept(&mut self) -> Result<Option<StallingSession>, FakeError> {
Ok(Some(StallingSession {
outbound: VecDeque::from(vec![Bytes::from_static(HANDSHAKE_REQUESTS[0])]),
}))
}
}
#[test]
fn never_completing_handshake_is_bounded_and_reported_not_leaked() {
const DEADLINE: Timestamp = Timestamp::from_nanos(1_000);
let mut driver = run_listen(
StallingListener { max_sessions: 1 },
trunk_config(),
HandshakePolicy::establish_by(DEADLINE),
max_programs(),
);
let AcceptOutcome::Admitted(id) = driver.poll_accept() else {
panic!("expected admission");
};
assert!(matches!(driver.health(id), Some(HealthState::Establishing)));
assert!(matches!(driver.poll_accept(), AcceptOutcome::Refused));
assert!(
driver
.feed(id, b"...", Timestamp::from_nanos(DEADLINE.as_nanos() - 1))
.is_none(),
"must not time out before the deadline"
);
assert!(matches!(driver.health(id), Some(HealthState::Establishing)));
assert_eq!(driver.session_count(), 1);
let health = driver
.on_deadline(id, DEADLINE)
.expect("the blown deadline must terminate the session");
assert_eq!(
health,
HealthState::HandshakeTimedOut { deadline: DEADLINE },
"a never-completing handshake must be reported as HandshakeTimedOut"
);
assert_eq!(
driver.session_count(),
0,
"a timed-out session must be reaped, not left pinning its slot"
);
assert!(driver.health(id).is_none());
assert!(matches!(driver.poll_accept(), AcceptOutcome::Admitted(_)));
}
#[test]
fn handshake_completing_exactly_at_the_deadline_still_establishes() {
const DEADLINE: Timestamp = Timestamp::from_nanos(500);
let session = ScriptedSession::new(vec![]);
let mut driver = IngestDriver::new(
session,
trunk_config(),
HandshakePolicy::establish_by(DEADLINE),
max_programs(),
);
driver.feed(b"reply", DEADLINE);
assert!(
matches!(driver.health(), HealthState::Live),
"a handshake completing exactly at the deadline must establish, \
not be rejected by a nanosecond: {:?}",
driver.health()
);
}
#[test]
fn next_deadline_surfaces_the_handshake_bound_while_establishing() {
const DEADLINE: Timestamp = Timestamp::from_nanos(9_000);
let mut dialer = HandshakeDialer;
let mut driver = run_dial(
&mut dialer,
trunk_config(),
HandshakePolicy::establish_by(DEADLINE),
max_programs(),
)
.unwrap();
assert_eq!(
driver.next_deadline(),
Some(DEADLINE),
"while Establishing, next_deadline must surface the handshake bound"
);
for _ in 0..HANDSHAKE_REQUESTS.len() {
let req = driver.poll_transmit().expect("handshake request");
let reply = format!("200 {}", String::from_utf8_lossy(&req));
driver.feed(reply.as_bytes(), Timestamp::ZERO);
}
assert!(matches!(driver.health(), HealthState::Live));
assert_eq!(
driver.next_deadline(),
None,
"the handshake bound must not linger after establishment"
);
}
}