use super::buffer::{JitterBuffer as Buffer, Rejected, State};
use super::sequence::TimestampExtender;
use crate::Interceptor;
use crate::stream_info::StreamInfo;
use crate::{Packet, TaggedPacket};
use sansio::Protocol;
use shared::error::Error;
use std::collections::{HashMap, VecDeque};
use std::time::{Duration, Instant};
pub const DEFAULT_DEPTH: Duration = Duration::from_millis(120);
pub const DEFAULT_CAPACITY: usize = 512;
const DISCONTINUITY: Duration = Duration::from_secs(10);
pub struct JitterBufferBuilder {
depth: Duration,
capacity: usize,
}
impl Default for JitterBufferBuilder {
fn default() -> Self {
Self {
depth: DEFAULT_DEPTH,
capacity: DEFAULT_CAPACITY,
}
}
}
impl JitterBufferBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn with_depth(mut self, depth: Duration) -> Self {
self.depth = depth;
self
}
pub fn with_capacity(mut self, capacity: usize) -> Self {
self.capacity = capacity;
self
}
pub fn build(self) -> JitterBufferInterceptor {
JitterBufferInterceptor::new(self.depth, self.capacity)
}
}
#[derive(Debug, Clone, Copy)]
struct Anchor {
arrived: Instant,
timestamp: u64,
}
struct Stream {
buffer: Buffer,
timestamps: TimestampExtender,
clock_rate: u32,
anchor: Option<Anchor>,
deadlines: HashMap<u64, Instant>,
}
impl Stream {
fn new(clock_rate: u32, capacity: usize) -> Self {
Self {
buffer: Buffer::new(capacity),
timestamps: TimestampExtender::new(),
clock_rate: if clock_rate == 0 { 90_000 } else { clock_rate },
anchor: None,
deadlines: HashMap::new(),
}
}
fn offset_from_anchor(&self, anchor: &Anchor, timestamp: u64) -> Duration {
let ticks = timestamp.saturating_sub(anchor.timestamp);
Duration::from_secs_f64(ticks as f64 / f64::from(self.clock_rate))
}
fn restart(&mut self, arrived: Instant, timestamp: u64) {
self.buffer.reset();
self.deadlines.clear();
self.anchor = Some(Anchor { arrived, timestamp });
}
}
pub struct JitterBufferInterceptor {
depth: Duration,
capacity: usize,
streams: HashMap<u32, Stream>,
due: VecDeque<TaggedPacket>,
read_queue: VecDeque<TaggedPacket>,
write_queue: VecDeque<TaggedPacket>,
}
impl JitterBufferInterceptor {
fn new(depth: Duration, capacity: usize) -> Self {
Self {
read_queue: VecDeque::new(),
write_queue: VecDeque::new(),
due: VecDeque::new(),
depth,
capacity,
streams: HashMap::new(),
}
}
fn next_deadline(stream: &Stream) -> Option<Instant> {
let front = stream.buffer.front_sequence()?;
stream.deadlines.get(&front).copied()
}
}
impl Protocol<TaggedPacket, TaggedPacket, ()> for JitterBufferInterceptor {
type Rout = TaggedPacket;
type Wout = TaggedPacket;
type Eout = ();
type Error = Error;
type Time = Instant;
fn handle_read(&mut self, msg: TaggedPacket) -> Result<(), Self::Error> {
let Packet::Rtp(rtp) = &msg.message.packet else {
self.read_queue.push_back(msg);
return Ok(());
};
let ssrc = rtp.header.ssrc;
let timestamp = rtp.header.timestamp;
let arrived = msg.now;
let Some(stream) = self.streams.get_mut(&ssrc) else {
self.read_queue.push_back(msg);
return Ok(());
};
let extended_timestamp = stream.timestamps.extend(timestamp);
match stream.anchor {
None => {
stream.anchor = Some(Anchor {
arrived,
timestamp: extended_timestamp,
})
}
Some(anchor) => {
let ticks = extended_timestamp.abs_diff(anchor.timestamp);
let elapsed = Duration::from_secs_f64(ticks as f64 / f64::from(stream.clock_rate));
if elapsed > DISCONTINUITY {
stream.restart(arrived, extended_timestamp);
}
}
}
let anchor = stream.anchor.expect("just anchored");
let deadline =
anchor.arrived + self.depth + stream.offset_from_anchor(&anchor, extended_timestamp);
match stream.buffer.push(msg) {
Ok(extended) => {
stream.deadlines.insert(extended, deadline);
Ok(())
}
Err(
Rejected::Duplicate | Rejected::Late | Rejected::Overflow | Rejected::ForeignSsrc,
) => Ok(()),
}
}
fn poll_read(&mut self) -> Option<TaggedPacket> {
self.read_queue.pop_front().or_else(|| self.due.pop_front())
}
fn handle_write(&mut self, msg: TaggedPacket) -> Result<(), Self::Error> {
self.write_queue.push_back(msg);
Ok(())
}
fn poll_write(&mut self) -> Option<Self::Wout> {
self.write_queue.pop_front()
}
fn handle_timeout(&mut self, now: Instant) -> Result<(), Error> {
let mut due: Vec<TaggedPacket> = Vec::new();
for stream in self.streams.values_mut() {
if stream.buffer.state() == State::Buffering
&& Self::next_deadline(stream).is_some_and(|deadline| deadline <= now)
{
stream.buffer.begin_emitting();
}
while let Some(front) = stream.buffer.front_sequence() {
let Some(&deadline) = stream.deadlines.get(&front) else {
break;
};
if deadline > now {
break;
}
let Some(mut packet) = stream.buffer.pop() else {
break;
};
stream.deadlines.remove(&front);
packet.now = now;
due.push(packet);
}
if stream.buffer.is_empty() && stream.buffer.state() == State::Emitting {
stream.buffer.begin_buffering();
stream.anchor = None;
}
}
self.due.extend(due);
Ok(())
}
fn poll_timeout(&mut self) -> Option<Instant> {
self.streams.values().filter_map(Self::next_deadline).min()
}
}
impl Interceptor for JitterBufferInterceptor {
fn bind_remote_stream(&mut self, info: &StreamInfo) {
self.streams
.entry(info.ssrc)
.or_insert_with(|| Stream::new(info.clock_rate, self.capacity));
}
fn unbind_remote_stream(&mut self, info: &StreamInfo) {
self.streams.remove(&info.ssrc);
}
fn bind_local_stream(&mut self, _info: &StreamInfo) {}
fn unbind_local_stream(&mut self, _info: &StreamInfo) {}
}