use crate::decision::AdBreakDecision;
use crate::error::{Error, Result};
use crate::splice::ConditionedSplicePoint;
use alloc::collections::BTreeMap;
use alloc::string::String;
#[derive(Debug, Clone, PartialEq)]
pub struct BreakState {
pub decision: AdBreakDecision,
pub entered_at: ConditionedSplicePoint,
pub resumed_at: Option<ConditionedSplicePoint>,
}
impl BreakState {
pub fn new(decision: AdBreakDecision, entered_at: ConditionedSplicePoint) -> Self {
BreakState {
decision,
entered_at,
resumed_at: None,
}
}
pub fn is_resumed(&self) -> bool {
self.resumed_at.is_some()
}
}
pub const DEFAULT_MAX_SESSIONS: usize = 65_536;
#[derive(Debug, Clone)]
pub struct SessionStore {
sessions: BTreeMap<String, Tracked>,
max_sessions: usize,
now_us: u64,
}
#[derive(Debug, Clone)]
struct Tracked {
state: BreakState,
touched_us: u64,
}
impl Default for SessionStore {
fn default() -> Self {
Self::new()
}
}
impl SessionStore {
pub fn new() -> Self {
SessionStore {
sessions: BTreeMap::new(),
max_sessions: DEFAULT_MAX_SESSIONS,
now_us: 0,
}
}
pub fn with_max_sessions(mut self, max_sessions: usize) -> Self {
self.max_sessions = max_sessions.max(1);
self
}
pub fn advance(&mut self, now_us: u64) {
self.now_us = self.now_us.max(now_us);
}
pub fn expire_idle(&mut self, ttl_us: u64) -> usize {
let before = self.sessions.len();
let now = self.now_us;
self.sessions
.retain(|_, t| now.saturating_sub(t.touched_us) <= ttl_us);
before - self.sessions.len()
}
pub fn begin_break(&mut self, session_id: impl Into<String>, state: BreakState) -> Result<()> {
let session_id = session_id.into();
match self.sessions.get(&session_id) {
Some(existing) if !existing.state.is_resumed() => {
return Err(Error::BreakAlreadyActive(session_id));
}
None if self.sessions.len() >= self.max_sessions => {
return Err(Error::SessionStoreFull {
max: self.max_sessions,
});
}
_ => {}
}
self.sessions.insert(
session_id,
Tracked {
state,
touched_us: self.now_us,
},
);
Ok(())
}
pub fn resume(&mut self, session_id: &str, resumed_at: ConditionedSplicePoint) -> Result<()> {
let tracked = self
.sessions
.get_mut(session_id)
.ok_or_else(|| Error::NoActiveBreak(session_id.into()))?;
if tracked.state.is_resumed() {
return Err(Error::BreakAlreadyResumed(session_id.into()));
}
tracked.state.resumed_at = Some(resumed_at);
tracked.touched_us = self.now_us;
Ok(())
}
pub fn get(&self, session_id: &str) -> Option<&BreakState> {
self.sessions.get(session_id).map(|t| &t.state)
}
pub fn end(&mut self, session_id: &str) -> Option<BreakState> {
self.sessions.remove(session_id).map(|t| t.state)
}
pub fn len(&self) -> usize {
self.sessions.len()
}
pub fn is_empty(&self) -> bool {
self.sessions.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::splice::condition_splice_point;
fn splice_point() -> ConditionedSplicePoint {
condition_splice_point(1_000, &[1_000], 0).unwrap()
}
#[test]
fn begin_get_resume_end_lifecycle() {
let mut store = SessionStore::new();
assert!(store.is_empty());
let decision = AdBreakDecision::single_asset("b1", "https://x/a.m3u8");
store
.begin_break("s1", BreakState::new(decision.clone(), splice_point()))
.unwrap();
assert_eq!(store.len(), 1);
assert!(!store.get("s1").unwrap().is_resumed());
store.resume("s1", splice_point()).unwrap();
assert!(store.get("s1").unwrap().is_resumed());
let evicted = store.end("s1").unwrap();
assert_eq!(evicted.decision, decision);
assert!(store.is_empty());
}
#[test]
fn cannot_begin_a_break_while_one_is_already_active() {
let mut store = SessionStore::new();
let decision = AdBreakDecision::single_asset("b1", "https://x/a.m3u8");
store
.begin_break("s1", BreakState::new(decision.clone(), splice_point()))
.unwrap();
let err = store
.begin_break("s1", BreakState::new(decision, splice_point()))
.unwrap_err();
assert!(matches!(err, Error::BreakAlreadyActive(ref s) if s == "s1"));
}
#[test]
fn begin_break_is_allowed_again_after_resume() {
let mut store = SessionStore::new();
let decision = AdBreakDecision::single_asset("b1", "https://x/a.m3u8");
store
.begin_break("s1", BreakState::new(decision.clone(), splice_point()))
.unwrap();
store.resume("s1", splice_point()).unwrap();
store
.begin_break("s1", BreakState::new(decision, splice_point()))
.unwrap();
assert_eq!(store.len(), 1);
}
#[test]
fn resume_without_an_active_break_errors() {
let mut store = SessionStore::new();
let err = store.resume("nope", splice_point()).unwrap_err();
assert!(matches!(err, Error::NoActiveBreak(ref s) if s == "nope"));
}
#[test]
fn resume_on_an_already_resumed_break_keeps_the_first_point() {
let mut store = SessionStore::new();
let decision = AdBreakDecision::single_asset("b1", "https://x/a.m3u8");
store
.begin_break("s1", BreakState::new(decision, splice_point()))
.unwrap();
let first = condition_splice_point(2_000, &[2_000], 0).unwrap();
store.resume("s1", first).unwrap();
let err = store.resume("s1", splice_point()).unwrap_err();
assert!(matches!(err, Error::BreakAlreadyResumed(ref s) if s == "s1"));
assert_eq!(store.get("s1").unwrap().resumed_at, Some(first));
}
#[test]
fn the_store_is_bounded_and_existing_sessions_still_begin_breaks() {
let mut store = SessionStore::new().with_max_sessions(2);
let decision = AdBreakDecision::single_asset("b1", "https://x/a.m3u8");
for id in ["a", "b"] {
store
.begin_break(id, BreakState::new(decision.clone(), splice_point()))
.unwrap();
}
let err = store
.begin_break("c", BreakState::new(decision.clone(), splice_point()))
.unwrap_err();
assert!(matches!(err, Error::SessionStoreFull { max: 2 }));
store.resume("a", splice_point()).unwrap();
store
.begin_break("a", BreakState::new(decision, splice_point()))
.unwrap();
assert_eq!(store.len(), 2);
}
#[test]
fn expire_idle_drops_only_sessions_untouched_for_longer_than_the_ttl() {
let mut store = SessionStore::new();
let decision = AdBreakDecision::single_asset("b1", "https://x/a.m3u8");
store.advance(1_000);
store
.begin_break("old", BreakState::new(decision.clone(), splice_point()))
.unwrap();
store.advance(5_000);
store
.begin_break("new", BreakState::new(decision, splice_point()))
.unwrap();
store.advance(6_000);
assert_eq!(store.expire_idle(2_000), 1);
assert!(store.get("old").is_none());
assert!(store.get("new").is_some());
store.advance(1);
assert_eq!(
store.expire_idle(1_000),
0,
"the clock never goes backwards"
);
}
}