use crate::{Error, Result, Timescale, Timestamp, coding};
use crate::{broadcast, cache, frame, group, stats};
use super::{Datagram, Requests};
pub use super::subscription::Subscription;
use std::{
collections::{HashSet, VecDeque},
sync::Arc,
task::{Poll, ready},
time::Duration,
};
pub const DEFAULT_LATENCY_MAX: Duration = Duration::from_secs(5);
const MAX_DATAGRAM_AGE: Duration = Duration::from_millis(50);
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct Info {
pub timescale: Timescale,
pub latency_max: Duration,
pub priority: u8,
pub ordered: bool,
pub(crate) broadcast: Arc<broadcast::Info>,
}
fn default_broadcast() -> Arc<broadcast::Info> {
static DEFAULT: std::sync::LazyLock<Arc<broadcast::Info>> =
std::sync::LazyLock::new(|| Arc::new(broadcast::Info::default()));
DEFAULT.clone()
}
impl Default for Info {
fn default() -> Self {
Self {
timescale: Timescale::default(),
latency_max: DEFAULT_LATENCY_MAX,
priority: 0,
ordered: false,
broadcast: default_broadcast(),
}
}
}
impl Info {
pub fn with_timescale(mut self, timescale: Timescale) -> Self {
self.timescale = timescale;
self
}
pub fn with_latency_max(mut self, latency_max: Duration) -> Self {
self.latency_max = latency_max;
self
}
pub fn with_priority(mut self, priority: u8) -> Self {
self.priority = priority;
self
}
pub fn with_ordered(mut self, ordered: bool) -> Self {
self.ordered = ordered;
self
}
}
#[derive(Default)]
struct TrackState {
info: Option<Info>,
broadcast: Arc<broadcast::Info>,
latest_entry: Option<Arc<cache::Entry>>,
groups: VecDeque<Option<(group::Producer, web_async::time::Instant)>>,
datagrams: VecDeque<(Datagram, web_async::time::Instant)>,
datagram_offset: usize,
duplicates: HashSet<u64>,
offset: usize,
max_sequence: Option<u64>,
final_sequence: Option<u64>,
abort: Option<Error>,
subscriptions: kio::Shared<Subscriptions>,
fetch: kio::Shared<FetchState>,
}
type Subscriptions = Vec<kio::Consumer<Subscription>>;
type FetchState = Requests<u64, PendingFetch>;
struct PendingFetch {
priority: u8,
result: kio::Producer<FetchOutcome>,
}
#[derive(Default)]
struct FetchOutcome {
rejected: Option<Error>,
}
impl TrackState {
fn poll_info(&self) -> Poll<Result<Info>> {
if let Some(info) = &self.info {
Poll::Ready(Ok(info.clone()))
} else {
Poll::Pending
}
}
fn poll_recv_group(&self, index: usize, min_sequence: u64) -> Poll<Result<Option<(group::Consumer, usize)>>> {
let start = index.saturating_sub(self.offset);
for (i, slot) in self.groups.iter().enumerate().skip(start) {
if let Some((group, _)) = slot
&& group.sequence >= min_sequence
&& !group.is_aborted()
{
return Poll::Ready(Ok(Some((group.consume(), self.offset + i))));
}
}
if self.is_complete() {
Poll::Ready(Ok(None))
} else if let Some(err) = &self.abort {
Poll::Ready(Err(err.clone()))
} else {
Poll::Pending
}
}
fn poll_recv_datagram(&self, index: usize) -> Poll<Result<Option<(Datagram, usize)>>> {
let start = index.saturating_sub(self.datagram_offset);
if let Some((datagram, _)) = self.datagrams.get(start) {
return Poll::Ready(Ok(Some((datagram.clone(), self.datagram_offset + start))));
}
if self.is_complete() {
Poll::Ready(Ok(None))
} else if let Some(err) = &self.abort {
Poll::Ready(Err(err.clone()))
} else {
Poll::Pending
}
}
fn push_datagram(&mut self, datagram: Datagram) {
let now = web_async::time::Instant::now();
self.datagrams.push_back((datagram, now));
while let Some((_, at)) = self.datagrams.front() {
if now.duration_since(*at) <= MAX_DATAGRAM_AGE {
break;
}
self.datagrams.pop_front();
self.datagram_offset += 1;
}
}
fn poll_read_frame(
&self,
index: usize,
next_sequence: u64,
waiter: &kio::Waiter,
) -> Poll<Result<Option<(frame::Frame, usize, u64)>>> {
let start = index.saturating_sub(self.offset);
let mut pending_seen = false;
for (i, slot) in self.groups.iter().enumerate().skip(start) {
let Some((group, _)) = slot else { continue };
if group.sequence < next_sequence {
continue;
}
let mut consumer = group.consume();
match consumer.poll_read_frame(waiter) {
Poll::Ready(Ok(Some(frame))) => {
return Poll::Ready(Ok(Some((frame, self.offset + i, group.sequence))));
}
Poll::Ready(Ok(None)) => continue,
Poll::Ready(Err(_)) => continue,
Poll::Pending => {
pending_seen = true;
continue;
}
}
}
if pending_seen {
Poll::Pending
} else if self.is_complete() {
Poll::Ready(Ok(None))
} else if let Some(err) = &self.abort {
Poll::Ready(Err(err.clone()))
} else {
Poll::Pending
}
}
fn poll_next_in_range(
&self,
next_sequence: u64,
end_sequence: Option<u64>,
) -> Poll<Result<Option<group::Consumer>>> {
if let Some(end) = end_sequence
&& end < next_sequence
{
if let Some(err) = &self.abort {
return Poll::Ready(Err(err.clone()));
}
return Poll::Pending;
}
let mut best: Option<&group::Producer> = None;
for (group, _) in self.groups.iter().flatten() {
if group.sequence < next_sequence {
continue;
}
if let Some(end) = end_sequence
&& group.sequence > end
{
continue;
}
if group.is_aborted() {
continue;
}
if best.is_none_or(|b| group.sequence < b.sequence) {
best = Some(group);
}
}
if let Some(group) = best {
return Poll::Ready(Ok(Some(group.consume())));
}
if let Some(err) = &self.abort {
return Poll::Ready(Err(err.clone()));
}
if let Some(fin) = self.final_sequence
&& next_sequence >= fin
{
return Poll::Ready(Ok(None));
}
Poll::Pending
}
fn cached_group(&self, sequence: u64) -> Option<group::Consumer> {
self.groups
.iter()
.flatten()
.find(|(group, _)| group.sequence == sequence && !group.is_aborted())
.map(|(group, _)| group.consume())
}
fn latency_bound(&self) -> Option<Duration> {
self.info.as_ref().map(|info| info.latency_max)
}
fn poll_fetch_cached(&self, sequence: u64) -> Poll<Result<group::Consumer>> {
if let Some(group) = self.cached_group(sequence) {
return Poll::Ready(Ok(group));
}
if let Some(err) = &self.abort {
return Poll::Ready(Err(err.clone()));
}
if self.final_sequence.is_some_and(|fin| sequence >= fin) {
return Poll::Ready(Err(Error::NotFound));
}
Poll::Pending
}
fn evict_expired(&mut self, now: web_async::time::Instant, max_age: Duration) {
for slot in self.groups.iter_mut() {
let Some((group, created_at)) = slot else { continue };
if group.is_aborted() {
self.duplicates.remove(&group.sequence);
*slot = None;
continue;
}
if Some(group.sequence) == self.max_sequence {
continue;
}
if now.duration_since(*created_at) <= max_age {
break;
}
self.duplicates.remove(&group.sequence);
if let Some((group, _)) = slot.take() {
let _ = group.abort(Error::Old);
}
}
while let Some(None) = self.groups.front() {
self.groups.pop_front();
self.offset += 1;
}
}
fn pin_latest(&mut self, group: &group::Producer) {
if Some(group.sequence) != self.max_sequence {
return;
}
if let Some(prev) = self.latest_entry.take() {
prev.set_pinned(false);
}
if let Some(entry) = group.cache_entry() {
entry.set_pinned(true);
self.latest_entry = Some(entry);
}
}
fn set_final(&mut self, final_sequence: u64) -> Result<()> {
if self.final_sequence.is_some() {
return Err(Error::Closed);
}
if let Some(max) = self.max_sequence
&& final_sequence <= max
{
return Err(Error::ProtocolViolation);
}
self.final_sequence = Some(final_sequence);
Ok(())
}
fn is_complete(&self) -> bool {
self.final_sequence
.is_some_and(|fin| self.max_sequence.map_or(0, |max| max.saturating_add(1)) >= fin)
}
fn poll_finished(&self) -> Poll<Result<u64>> {
if let Some(fin) = self.final_sequence {
Poll::Ready(Ok(fin))
} else if let Some(err) = &self.abort {
Poll::Ready(Err(err.clone()))
} else {
Poll::Pending
}
}
fn modify(producer: &kio::Producer<Self>) -> Result<kio::Mut<'_, Self>> {
producer.write().map_err(|r| r.abort.clone().unwrap_or(Error::Dropped))
}
fn replace_evicted(
&mut self,
sequence: u64,
track: Info,
now: web_async::time::Instant,
) -> Option<Result<group::Producer>> {
let slot = self
.groups
.iter_mut()
.find(|slot| matches!(slot, Some((group, _)) if group.sequence == sequence))?;
let (existing, _) = slot.as_ref().unwrap();
if !existing.is_aborted() {
return Some(Err(Error::Duplicate));
}
let group = group::Producer::new(group::Info { sequence }, track);
*slot = Some((group.clone(), now));
self.pin_latest(&group);
Some(Ok(group))
}
fn insert_group_request(&mut self, sequence: u64, info: Option<Info>) -> Result<group::Producer> {
if let Some(err) = &self.abort {
return Err(err.clone());
}
if let Some(fin) = self.final_sequence
&& sequence >= fin
{
return Err(Error::Closed);
}
let now = web_async::time::Instant::now();
let broadcast = self.broadcast.clone();
let info = self
.info
.get_or_insert_with(|| {
let mut info = info.unwrap_or_default();
info.broadcast = broadcast;
info
})
.clone();
if !self.duplicates.insert(sequence) {
return self
.replace_evicted(sequence, info, now)
.unwrap_or(Err(Error::Duplicate));
}
let latency_max = info.latency_max;
let group = group::Producer::new(group::Info { sequence }, info);
self.max_sequence = Some(self.max_sequence.unwrap_or(0).max(sequence));
self.groups.push_back(Some((group.clone(), now)));
self.pin_latest(&group);
self.evict_expired(now, latency_max);
Ok(group)
}
}
#[derive(Clone)]
pub struct Producer {
name: Arc<str>,
broadcast: Arc<broadcast::Info>,
state: kio::Producer<TrackState>,
prev_subscription: Option<Subscription>,
stats: stats::Scope,
}
impl Producer {
pub(crate) fn new(
broadcast: Arc<broadcast::Info>,
name: impl Into<Arc<str>>,
info: impl Into<Option<Info>>,
) -> Self {
let mut info = info.into().unwrap_or_default();
info.broadcast = broadcast.clone();
Self {
name: name.into(),
state: kio::Producer::new(TrackState {
info: Some(info),
broadcast: broadcast.clone(),
..Default::default()
}),
broadcast,
prev_subscription: None,
stats: stats::Scope::default(),
}
}
pub(crate) fn with_stats(mut self, scope: stats::Scope) -> Self {
scope.open_subscription();
self.stats = scope;
self
}
pub fn name(&self) -> &str {
&self.name
}
pub fn broadcast(&self) -> &broadcast::Info {
&self.broadcast
}
pub fn create_group(&mut self, group: group::Info) -> Result<group::Producer> {
let mut state = self.modify()?;
if let Some(fin) = state.final_sequence
&& group.sequence >= fin
{
return Err(Error::Closed);
}
let info = state.info.as_ref().unwrap();
let track = info.clone();
let latency_max = info.latency_max;
let now = web_async::time::Instant::now();
if !state.duplicates.insert(group.sequence) {
return state
.replace_evicted(group.sequence, track, now)
.unwrap_or(Err(Error::Duplicate));
}
let group = group::Producer::new(group, track).with_meter(self.stats.meter());
state.max_sequence = Some(state.max_sequence.unwrap_or(0).max(group.sequence));
state.groups.push_back(Some((group.clone(), now)));
state.pin_latest(&group);
state.evict_expired(now, latency_max);
Ok(group)
}
pub fn append_group(&mut self) -> Result<group::Producer> {
let mut state = self.modify()?;
let sequence = match state.max_sequence {
Some(s) => s.checked_add(1).ok_or(coding::BoundsExceeded)?,
None => 0,
};
if let Some(fin) = state.final_sequence
&& sequence >= fin
{
return Err(Error::Closed);
}
let info = state.info.as_ref().unwrap();
let track = info.clone();
let latency_max = info.latency_max;
let group = group::Producer::new(group::Info { sequence }, track).with_meter(self.stats.meter());
let now = web_async::time::Instant::now();
state.duplicates.insert(sequence);
state.max_sequence = Some(sequence);
state.groups.push_back(Some((group.clone(), now)));
state.pin_latest(&group);
state.evict_expired(now, latency_max);
Ok(group)
}
pub fn append_datagram<B: crate::IntoBytes>(&mut self, timestamp: Timestamp, payload: B) -> Result<u64> {
let payload = payload.into_bytes();
if payload.len() > super::datagram::MAX_DATAGRAM_PAYLOAD {
return Err(Error::FrameTooLarge);
}
let meter = self.stats.meter();
let mut state = self.modify()?;
let timescale = state.info.as_ref().unwrap().timescale;
let timestamp = timestamp.convert(timescale).map_err(|_| Error::TimestampMismatch)?;
let sequence = match state.max_sequence {
Some(s) => s.checked_add(1).ok_or(coding::BoundsExceeded)?,
None => 0,
};
if let Some(fin) = state.final_sequence
&& sequence >= fin
{
return Err(Error::Closed);
}
state.max_sequence = Some(sequence);
meter.datagram(payload.len() as u64);
state.push_datagram(Datagram {
sequence,
timestamp,
payload,
});
Ok(sequence)
}
pub fn write_datagram(&mut self, mut datagram: Datagram) -> Result<()> {
if datagram.payload.len() > super::datagram::MAX_DATAGRAM_PAYLOAD {
return Err(Error::FrameTooLarge);
}
let meter = self.stats.meter();
let mut state = self.modify()?;
let timescale = state.info.as_ref().unwrap().timescale;
datagram.timestamp = datagram
.timestamp
.convert(timescale)
.map_err(|_| Error::TimestampMismatch)?;
if let Some(fin) = state.final_sequence
&& datagram.sequence >= fin
{
return Err(Error::Closed);
}
state.max_sequence = Some(state.max_sequence.unwrap_or(0).max(datagram.sequence));
meter.datagram(datagram.payload.len() as u64);
state.push_datagram(datagram);
Ok(())
}
pub fn write_frame<B: crate::IntoBytes>(&mut self, timestamp: Timestamp, frame: B) -> Result<()> {
let mut group = self.append_group()?;
group.write_frame(timestamp, frame)?;
group.finish()?;
Ok(())
}
pub fn finish(&mut self) -> Result<()> {
let mut state = self.modify()?;
let final_sequence = match state.max_sequence {
Some(max) => max.checked_add(1).ok_or(coding::BoundsExceeded)?,
None => 0,
};
state.set_final(final_sequence)
}
pub fn finish_at(&mut self, final_sequence: u64) -> Result<()> {
self.modify()?.set_final(final_sequence)
}
pub fn final_sequence(&self) -> Option<u64> {
self.state.read().final_sequence
}
pub fn abort(self, err: Error) -> Result<()> {
let mut guard = self.modify()?;
guard.abort = Some(err);
guard.groups.clear();
guard.datagrams.clear();
guard.duplicates.clear();
guard.latest_entry = None;
guard.close();
Ok(())
}
pub async fn unused(&self) -> Result<()> {
self.state.unused().await.map_err(|_| self.abort_reason())
}
pub async fn used(&self) -> Result<()> {
self.state.used().await.map_err(|_| self.abort_reason())
}
pub async fn closed(&self) -> Error {
kio::wait(|waiter| self.poll_closed(waiter)).await
}
pub fn poll_closed(&self, waiter: &kio::Waiter) -> Poll<Error> {
self.state.poll_closed(waiter).map(|()| self.abort_reason())
}
fn abort_reason(&self) -> Error {
self.state.read().abort.clone().unwrap_or(Error::Dropped)
}
pub fn is_closed(&self) -> bool {
self.state.read().is_closed()
}
pub fn latest(&self) -> Option<u64> {
self.state.read().max_sequence
}
pub fn is_clone(&self, other: &Self) -> bool {
self.state.same_channel(&other.state)
}
pub(crate) fn weak(&self) -> TrackWeak {
TrackWeak {
name: self.name.clone(),
state: self.state.weak(),
}
}
pub fn demand(&self) -> Demand {
Demand {
name: self.name.clone(),
state: self.state.weak(),
}
}
pub fn consume(&self) -> Consumer {
Consumer::plain(self.name.clone(), self.state.consume())
}
pub fn subscribe(&self, subscription: impl Into<Option<Subscription>>) -> Subscriber {
let preferences = subscription.into().unwrap_or_default();
let info = self
.state
.read()
.info
.as_ref()
.expect("producer always has info")
.clone();
let subscription = kio::Producer::new(preferences);
register_subscription(self.state.read(), &subscription);
Subscriber {
name: self.name.clone(),
info,
inner: SubscriberKind::Plain(PlainSubscriber {
state: self.state.consume(),
subscription,
index: 0,
datagram_index: 0,
min_sequence: 0,
next_sequence: 0,
end_sequence: None,
}),
stats: stats::Scope::default(),
_stats_sub: stats::Subscription::default(),
}
}
pub async fn subscription_changed(&mut self) -> Result<Option<Subscription>> {
kio::wait(|waiter| self.poll_subscription_changed(waiter)).await
}
pub fn subscription(&self) -> Option<Subscription> {
let state = self.state.read();
let (subs, bound) = (state.subscriptions.clone(), state.latency_bound());
drop(state);
snapshot_subscription(&subs, bound)
}
pub fn poll_subscription_changed(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Subscription>>> {
if self.state.poll_closed(waiter).is_ready() {
let abort = self.state.read().abort.clone();
return Poll::Ready(Err(abort.unwrap_or(Error::Dropped)));
}
let state = self.state.read();
let (subs, bound) = (state.subscriptions.clone(), state.latency_bound());
drop(state);
let prev = &self.prev_subscription;
let mut combined = None;
let mut guard = match subs.poll(waiter, |subs| {
let next = combined_subscription(subs, bound, waiter);
if &next == prev {
Poll::Pending
} else {
combined = next;
Poll::Ready(())
}
}) {
Poll::Ready(guard) => guard,
Poll::Pending => return Poll::Pending,
};
guard.retain(|sub| !sub.is_closed());
drop(guard);
self.prev_subscription = combined.clone();
Poll::Ready(Ok(combined))
}
pub fn poll_unused(&self, waiter: &kio::Waiter) -> Poll<()> {
self.state.poll_unused(waiter).map(|_| ())
}
pub fn dynamic(&self) -> Dynamic {
Dynamic::new(self.name.clone(), self.state.clone())
}
fn modify(&self) -> Result<kio::Mut<'_, TrackState>> {
TrackState::modify(&self.state)
}
}
fn poll_requested_group(
state: &kio::Producer<TrackState>,
fetch: &kio::Shared<FetchState>,
waiter: &kio::Waiter,
) -> Poll<Result<GroupRequest>> {
if let Poll::Ready(mut guard) = fetch.poll(waiter, |fetch| {
if fetch.has_queued() {
Poll::Ready(())
} else {
Poll::Pending
}
}) {
let sequence = guard.pop().expect("predicate guaranteed a request");
let pending = guard.get(&sequence).expect("popped key must be pending");
let priority = pending.priority;
let result = pending.result.clone();
drop(guard);
return Poll::Ready(Ok(GroupRequest {
state: state.clone(),
fetch: fetch.clone(),
sequence,
priority,
result,
done: false,
}));
}
match state.poll_ref(waiter, |state| match &state.abort {
Some(err) => Poll::Ready(err.clone()),
None => Poll::Pending,
}) {
Poll::Ready(Ok(err)) => Poll::Ready(Err(err)),
Poll::Ready(Err(closed)) => Poll::Ready(Err(closed.abort.clone().unwrap_or(Error::Dropped))),
Poll::Pending => Poll::Pending,
}
}
pub struct Dynamic {
name: Arc<str>,
state: kio::Producer<TrackState>,
fetch: kio::Shared<FetchState>,
}
impl Dynamic {
fn new(name: Arc<str>, state: kio::Producer<TrackState>) -> Self {
let fetch = state.read().fetch.clone();
fetch.lock().add_handler();
Self { name, state, fetch }
}
pub fn name(&self) -> &str {
&self.name
}
pub async fn requested_group(&self) -> Result<GroupRequest> {
kio::wait(|waiter| self.poll_requested_group(waiter)).await
}
pub fn poll_requested_group(&self, waiter: &kio::Waiter) -> Poll<Result<GroupRequest>> {
poll_requested_group(&self.state, &self.fetch, waiter)
}
pub fn poll_unused(&self, waiter: &kio::Waiter) -> Poll<()> {
self.state.poll_unused(waiter).map(|_| ())
}
}
impl Clone for Dynamic {
fn clone(&self) -> Self {
self.fetch.lock().add_handler();
Self {
name: self.name.clone(),
state: self.state.clone(),
fetch: self.fetch.clone(),
}
}
}
impl Drop for Dynamic {
fn drop(&mut self) {
let mut fetch = self.fetch.lock();
if fetch.remove_handler() {
fetch.drain_queued();
}
}
}
impl Drop for Producer {
fn drop(&mut self) {
if !self.state.is_last() {
return;
}
self.stats.close_subscription();
if let Ok(mut state) = self.state.write()
&& state.final_sequence.is_none()
{
tracing::warn!(
track = %self.name(),
"track::Producer dropped without finish() or abort()"
);
state.groups.clear();
state.datagrams.clear();
state.duplicates.clear();
state.latest_entry = None;
}
}
}
fn combined_subscription(subs: &Subscriptions, bound: Option<Duration>, waiter: &kio::Waiter) -> Option<Subscription> {
let mut combined = None;
for sub in subs.iter() {
if sub.is_closed() {
continue;
}
if let Poll::Ready(Ok(sub)) = sub.poll(waiter, |sub| sub.poll_combined(&combined)) {
combined = Some(sub);
}
}
clamp_combined(combined, bound)
}
fn snapshot_subscription(subs: &kio::Shared<Subscriptions>, bound: Option<Duration>) -> Option<Subscription> {
let mut combined: Option<Subscription> = None;
for sub in subs.read().iter() {
if sub.is_closed() {
continue;
}
if let Poll::Ready(merged) = sub.read().poll_combined(&combined) {
combined = Some(merged);
}
}
clamp_combined(combined, bound)
}
fn clamp_combined(combined: Option<Subscription>, bound: Option<Duration>) -> Option<Subscription> {
let mut combined = combined?;
if let Some(bound) = bound {
combined.latency_max = combined.latency_max.min(bound);
}
Some(combined)
}
fn register_subscription(state: kio::Ref<'_, TrackState>, subscription: &kio::Producer<Subscription>) {
if state.is_closed() {
return;
}
let subs = state.subscriptions.clone();
drop(state);
subs.lock().push(subscription.consume());
}
#[derive(Clone)]
pub(crate) struct TrackWeak {
name: Arc<str>,
state: kio::ProducerWeak<TrackState>,
}
impl TrackWeak {
pub fn consume(&self) -> Consumer {
Consumer::plain(self.name.clone(), self.state.consume())
}
pub(crate) fn name(&self) -> &Arc<str> {
&self.name
}
pub(crate) fn is_used(&self) -> bool {
!self.state.is_closed() && self.state.is_used()
}
pub(crate) fn poll_used(&self, waiter: &kio::Waiter) {
let _ = self.state.poll_used(waiter);
}
pub(crate) fn poll_unused(&self, waiter: &kio::Waiter) {
let _ = self.state.poll_unused(waiter);
}
}
impl super::WeakEntry for TrackWeak {
fn is_closed(&self) -> bool {
self.state.is_closed()
}
fn same_channel(&self, other: &Self) -> bool {
self.state.same_channel(&other.state)
}
}
#[derive(Clone)]
pub struct Demand {
name: Arc<str>,
state: kio::ProducerWeak<TrackState>,
}
impl Demand {
pub fn name(&self) -> &str {
&self.name
}
pub async fn used(&self) -> Result<()> {
self.state.used().await.map_err(|_| self.abort_reason())
}
pub async fn unused(&self) -> Result<()> {
self.state.unused().await.map_err(|_| self.abort_reason())
}
pub async fn closed(&self) -> Error {
self.state.closed().await;
self.abort_reason()
}
fn abort_reason(&self) -> Error {
self.state.read().abort.clone().unwrap_or(Error::Dropped)
}
}
#[derive(Clone)]
pub struct Consumer {
name: Arc<str>,
inner: ConsumerKind,
stats: stats::Scope,
}
#[derive(Clone)]
enum ConsumerKind {
Plain(kio::Consumer<TrackState>),
Spliced(super::resume::Consumer),
}
impl Consumer {
fn plain(name: Arc<str>, state: kio::Consumer<TrackState>) -> Self {
Self {
name,
inner: ConsumerKind::Plain(state),
stats: stats::Scope::default(),
}
}
pub(crate) fn spliced(name: Arc<str>, resume: super::resume::Consumer) -> Self {
Self {
name,
inner: ConsumerKind::Spliced(resume),
stats: stats::Scope::default(),
}
}
pub(crate) fn with_stats(mut self, scope: stats::Scope) -> Self {
self.stats = scope;
self
}
pub fn name(&self) -> &str {
&self.name
}
pub fn subscribe(&self, subscription: impl Into<Option<Subscription>>) -> kio::Pending<Subscribing> {
let subscription = kio::Producer::new(subscription.into().unwrap_or_default());
let inner = match &self.inner {
ConsumerKind::Plain(state) => {
register_subscription(state.read(), &subscription);
SubscribingKind::Plain(state.clone())
}
ConsumerKind::Spliced(resume) => SubscribingKind::Spliced(resume.clone()),
};
kio::Pending::new(Subscribing {
name: self.name.clone(),
inner,
subscription,
stats: self.stats.clone(),
})
}
#[cfg(test)]
pub(crate) fn peek_group(&self, sequence: u64) -> Option<group::Consumer> {
match &self.inner {
ConsumerKind::Plain(state) => state.read().cached_group(sequence),
ConsumerKind::Spliced(_) => None,
}
}
pub fn fetch_group(&self, sequence: u64, options: impl Into<Option<group::Fetch>>) -> kio::Pending<Fetching> {
let options = options.into().unwrap_or_default();
self.stats.fetch();
let state = match &self.inner {
ConsumerKind::Plain(state) => state,
ConsumerKind::Spliced(resume) => {
return kio::Pending::new(Fetching {
inner: FetchingKind::Spliced(resume.fetch_group(sequence, options)),
stats: self.stats.clone(),
});
}
};
let mut result = None;
let (fetch, unresolved) = {
let state = state.read();
(state.fetch.clone(), state.poll_fetch_cached(sequence).is_pending())
};
if unresolved {
let mut fetch = fetch.lock();
if let Some(pending) = fetch.join(&sequence) {
pending.priority = pending.priority.max(options.priority);
result = Some(pending.result.consume());
} else {
let producer = kio::Producer::<FetchOutcome>::default();
let consumer = producer.consume();
let attempt = PendingFetch {
priority: options.priority,
result: producer,
};
if fetch.insert(sequence, attempt).is_ok() {
result = Some(consumer);
}
}
}
kio::Pending::new(Fetching {
inner: FetchingKind::Plain {
state: state.clone(),
fetch,
sequence,
result,
},
stats: self.stats.clone(),
})
}
pub fn info(&self) -> kio::Pending<Querying> {
kio::Pending::new(Querying {
inner: match &self.inner {
ConsumerKind::Plain(state) => QueryingKind::Plain(state.clone()),
ConsumerKind::Spliced(resume) => QueryingKind::Spliced(resume.clone()),
},
})
}
pub fn latest(&self) -> Option<u64> {
match &self.inner {
ConsumerKind::Plain(state) => state.read().max_sequence,
ConsumerKind::Spliced(resume) => resume.latest(),
}
}
pub(crate) fn poll_complete(&self, waiter: &kio::Waiter) -> Poll<Result<()>> {
let ConsumerKind::Plain(state) = &self.inner else {
return Poll::Pending;
};
match ready!(state.poll(waiter, |state| {
if state.is_complete() {
Poll::Ready(())
} else {
Poll::Pending
}
})) {
Ok(_) => Poll::Ready(Ok(())),
Err(closed) => Poll::Ready(Err(closed.abort.clone().unwrap_or(Error::Dropped))),
}
}
}
pub struct Subscribing {
name: Arc<str>,
inner: SubscribingKind,
subscription: kio::Producer<Subscription>,
stats: stats::Scope,
}
enum SubscribingKind {
Plain(kio::Consumer<TrackState>),
Spliced(super::resume::Consumer),
}
impl Subscribing {
pub fn poll_ok(&self, waiter: &kio::Waiter) -> Poll<Result<Subscriber>> {
match &self.inner {
SubscribingKind::Plain(state) => {
let info = ready!(state.poll(waiter, |state| state.poll_info()))
.map_err(|e| e.abort.clone().unwrap_or(Error::Dropped))??;
Poll::Ready(Ok(Subscriber {
name: self.name.clone(),
info,
inner: SubscriberKind::Plain(PlainSubscriber {
state: state.clone(),
subscription: self.subscription.clone(),
index: 0,
datagram_index: 0,
min_sequence: 0,
next_sequence: 0,
end_sequence: None,
}),
stats: self.stats.clone(),
_stats_sub: self.stats.subscribe(),
}))
}
SubscribingKind::Spliced(resume) => {
let info = ready!(resume.poll_info(waiter))?;
Poll::Ready(Ok(Subscriber {
name: self.name.clone(),
info,
inner: SubscriberKind::Spliced(Box::new(resume.subscribe_shared(self.subscription.clone()))),
stats: self.stats.clone(),
_stats_sub: self.stats.subscribe(),
}))
}
}
}
pub fn update(&mut self, subscription: Subscription) -> Result<()> {
let mut state = self.subscription.write().map_err(|_| Error::Closed)?;
*state = subscription;
Ok(())
}
}
impl kio::Pollable for Subscribing {
type Output = Result<Subscriber>;
fn poll(&self, waiter: &kio::Waiter) -> Poll<Self::Output> {
self.poll_ok(waiter)
}
}
pub struct Querying {
inner: QueryingKind,
}
enum QueryingKind {
Plain(kio::Consumer<TrackState>),
Spliced(super::resume::Consumer),
}
impl Querying {
pub fn poll_ok(&self, waiter: &kio::Waiter) -> Poll<Result<Info>> {
match &self.inner {
QueryingKind::Plain(state) => {
let info = ready!(state.poll(waiter, |state| state.poll_info()))
.map_err(|e| e.abort.clone().unwrap_or(Error::Dropped))??;
Poll::Ready(Ok(info))
}
QueryingKind::Spliced(resume) => resume.poll_info(waiter),
}
}
}
impl kio::Pollable for Querying {
type Output = Result<Info>;
fn poll(&self, waiter: &kio::Waiter) -> Poll<Self::Output> {
self.poll_ok(waiter)
}
}
pub struct GroupRequest {
state: kio::Producer<TrackState>,
fetch: kio::Shared<FetchState>,
sequence: u64,
priority: u8,
result: kio::Producer<FetchOutcome>,
done: bool,
}
impl GroupRequest {
pub fn sequence(&self) -> u64 {
self.sequence
}
pub fn priority(&self) -> u8 {
self.priority
}
pub fn accept(mut self, info: impl Into<Option<Info>>) -> Result<group::Producer> {
self.done = true;
let res = TrackState::modify(&self.state)
.and_then(|mut state| state.insert_group_request(self.sequence, info.into()));
self.remove();
res
}
pub fn reject(mut self, err: Error) {
self.done = true;
self.remove();
if let Ok(mut outcome) = self.result.write() {
outcome.rejected = Some(err);
}
}
fn remove(&self) {
self.fetch
.lock()
.remove_if(&self.sequence, |pending| pending.result.same_channel(&self.result));
}
}
impl Drop for GroupRequest {
fn drop(&mut self) {
if self.done {
return;
}
self.remove();
if let Ok(mut outcome) = self.result.write() {
outcome.rejected = Some(Error::Dropped);
}
}
}
pub struct Fetching {
inner: FetchingKind,
stats: stats::Scope,
}
enum FetchingKind {
Plain {
state: kio::Consumer<TrackState>,
fetch: kio::Shared<FetchState>,
sequence: u64,
result: Option<kio::Consumer<FetchOutcome>>,
},
Spliced(kio::Pending<super::resume::Fetching>),
}
impl kio::Pollable for Fetching {
type Output = Result<group::Consumer>;
fn poll(&self, waiter: &kio::Waiter) -> Poll<Self::Output> {
let (state, fetch, sequence, result) = match &self.inner {
FetchingKind::Plain {
state,
fetch,
sequence,
result,
} => (state, fetch, *sequence, result.as_ref()),
FetchingKind::Spliced(spliced) => {
return kio::Pollable::poll(&**spliced, waiter)
.map(|res| res.map(|group| group.with_meter(self.stats.meter())));
}
};
match state.poll(waiter, |state| state.poll_fetch_cached(sequence)) {
Poll::Ready(Ok(res)) => return Poll::Ready(res.map(|group| group.with_meter(self.stats.meter()))),
Poll::Ready(Err(closed)) => {
return Poll::Ready(Err(closed.abort.clone().unwrap_or(Error::Dropped)));
}
Poll::Pending => {}
}
let Some(result) = result else {
return match fetch.poll(waiter, |fetch| match fetch.has_handlers() {
false => Poll::Ready(()),
true => Poll::Pending,
}) {
Poll::Ready(_guard) => Poll::Ready(Err(Error::NotFound)),
Poll::Pending => Poll::Pending,
};
};
match result.poll(waiter, |outcome| match &outcome.rejected {
Some(err) => Poll::Ready(err.clone()),
None => Poll::Pending,
}) {
Poll::Ready(Ok(err)) => Poll::Ready(Err(err)),
Poll::Ready(Err(_closed)) => Poll::Ready(Err(Error::NotFound)),
Poll::Pending => Poll::Pending,
}
}
}
pub struct Subscriber {
name: Arc<str>,
info: Info,
inner: SubscriberKind,
stats: stats::Scope,
_stats_sub: stats::Subscription,
}
enum SubscriberKind {
Plain(PlainSubscriber),
Spliced(Box<super::resume::Subscriber>),
}
struct PlainSubscriber {
state: kio::Consumer<TrackState>,
subscription: kio::Producer<Subscription>,
index: usize,
datagram_index: usize,
min_sequence: u64,
next_sequence: u64,
end_sequence: Option<u64>,
}
impl PlainSubscriber {
fn poll<F, R>(&self, waiter: &kio::Waiter, f: F) -> Poll<Result<R>>
where
F: Fn(&kio::Ref<'_, TrackState>) -> Poll<Result<R>>,
{
Poll::Ready(match ready!(self.state.poll(waiter, f)) {
Ok(res) => res,
Err(state) => Err(state.abort.clone().unwrap_or(Error::Dropped)),
})
}
fn poll_recv_group(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<group::Consumer>>> {
let Some((consumer, found_index)) =
ready!(self.poll(waiter, |state| state.poll_recv_group(self.index, self.min_sequence))?)
else {
return Poll::Ready(Ok(None));
};
self.index = found_index + 1;
Poll::Ready(Ok(Some(consumer)))
}
fn poll_recv_datagram(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Datagram>>> {
let Some((datagram, found_index)) =
ready!(self.poll(waiter, |state| state.poll_recv_datagram(self.datagram_index))?)
else {
return Poll::Ready(Ok(None));
};
self.datagram_index = found_index + 1;
self.next_sequence = self.next_sequence.max(datagram.sequence.saturating_add(1));
Poll::Ready(Ok(Some(datagram)))
}
fn poll_next_group(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<group::Consumer>>> {
let floor = self.next_sequence.max(self.min_sequence);
let Some(group) = ready!(self.poll(waiter, |state| state.poll_next_in_range(floor, self.end_sequence))?) else {
return Poll::Ready(Ok(None));
};
self.next_sequence = group.sequence.saturating_add(1);
Poll::Ready(Ok(Some(group)))
}
fn poll_read_frame(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<frame::Frame>>> {
let lower = self.min_sequence.max(self.next_sequence);
let Some((frame, found_index, sequence)) =
ready!(self.poll(waiter, |state| { state.poll_read_frame(self.index, lower, waiter) })?)
else {
return Poll::Ready(Ok(None));
};
self.index = found_index + 1;
self.next_sequence = sequence.saturating_add(1);
Poll::Ready(Ok(Some(frame)))
}
}
#[derive(Clone)]
pub struct SubscriberControl {
subscription: kio::Producer<Subscription>,
}
impl SubscriberControl {
pub fn subscription(&self) -> Subscription {
self.subscription.read().clone()
}
pub fn update(&self, subscription: Subscription) -> Result<()> {
let mut state = self.subscription.write().map_err(|_| Error::Closed)?;
*state = subscription;
Ok(())
}
}
impl Subscriber {
pub fn info(&self) -> &Info {
&self.info
}
pub fn name(&self) -> &str {
&self.name
}
pub fn control(&self) -> SubscriberControl {
SubscriberControl {
subscription: match &self.inner {
SubscriberKind::Plain(plain) => plain.subscription.clone(),
SubscriberKind::Spliced(spliced) => spliced.prefs(),
},
}
}
pub fn poll_recv_group(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<group::Consumer>>> {
let meter = self.stats.meter();
let res = match &mut self.inner {
SubscriberKind::Plain(plain) => plain.poll_recv_group(waiter),
SubscriberKind::Spliced(spliced) => spliced.poll_recv_group(waiter),
};
res.map(|res| res.map(|group| group.map(|group| group.with_meter(meter))))
}
pub async fn recv_group(&mut self) -> Result<Option<group::Consumer>> {
kio::wait(|waiter| self.poll_recv_group(waiter)).await
}
pub fn poll_recv_datagram(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Datagram>>> {
let meter = self.stats.meter();
let res = match &mut self.inner {
SubscriberKind::Plain(plain) => plain.poll_recv_datagram(waiter),
SubscriberKind::Spliced(spliced) => spliced.poll_recv_datagram(waiter),
};
if let Poll::Ready(Ok(Some(datagram))) = &res {
meter.datagram(datagram.payload.len() as u64);
}
res
}
pub async fn recv_datagram(&mut self) -> Result<Option<Datagram>> {
kio::wait(|waiter| self.poll_recv_datagram(waiter)).await
}
pub fn poll_next_group(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<group::Consumer>>> {
let meter = self.stats.meter();
let res = match &mut self.inner {
SubscriberKind::Plain(plain) => plain.poll_next_group(waiter),
SubscriberKind::Spliced(spliced) => spliced.poll_next_group(waiter),
};
res.map(|res| res.map(|group| group.map(|group| group.with_meter(meter))))
}
pub async fn next_group(&mut self) -> Result<Option<group::Consumer>> {
kio::wait(|waiter| self.poll_next_group(waiter)).await
}
pub fn poll_read_frame(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<frame::Frame>>> {
let meter = self.stats.meter();
let res = match &mut self.inner {
SubscriberKind::Plain(plain) => plain.poll_read_frame(waiter),
SubscriberKind::Spliced(spliced) => spliced.poll_read_frame(waiter),
};
if let Poll::Ready(Ok(Some(frame))) = &res {
meter.group();
meter.frames(1);
meter.bytes(frame.payload.len() as u64);
}
res
}
pub async fn read_frame(&mut self) -> Result<Option<frame::Frame>> {
kio::wait(|waiter| self.poll_read_frame(waiter)).await
}
pub fn is_clone(&self, other: &Self) -> bool {
match (&self.inner, &other.inner) {
(SubscriberKind::Plain(a), SubscriberKind::Plain(b)) => a.state.same_channel(&b.state),
(SubscriberKind::Spliced(a), SubscriberKind::Spliced(b)) => a.is_clone(b),
_ => false,
}
}
pub fn poll_finished(&mut self, waiter: &kio::Waiter) -> Poll<Result<u64>> {
match &mut self.inner {
SubscriberKind::Plain(plain) => plain.poll(waiter, |state| state.poll_finished()),
SubscriberKind::Spliced(spliced) => spliced.poll_finished(waiter),
}
}
pub async fn finished(&mut self) -> Result<u64> {
kio::wait(|waiter| self.poll_finished(waiter)).await
}
pub fn start_at(&mut self, sequence: u64) {
match &mut self.inner {
SubscriberKind::Plain(plain) => plain.min_sequence = sequence,
SubscriberKind::Spliced(spliced) => spliced.start_at(sequence),
}
}
pub fn end_at(&mut self, sequence: impl Into<Option<u64>>) {
match &mut self.inner {
SubscriberKind::Plain(plain) => plain.end_sequence = sequence.into(),
SubscriberKind::Spliced(spliced) => spliced.end_at(sequence),
}
}
pub fn subscription(&self) -> Subscription {
self.control().subscription()
}
pub fn update(&mut self, subscription: Subscription) -> Result<()> {
match &mut self.inner {
SubscriberKind::Plain(plain) => {
let mut state = plain.subscription.write().map_err(|_| Error::Closed)?;
*state = subscription;
}
SubscriberKind::Spliced(spliced) => spliced.update(subscription),
}
Ok(())
}
pub fn latest(&self) -> Option<u64> {
match &self.inner {
SubscriberKind::Plain(plain) => plain.state.read().max_sequence,
SubscriberKind::Spliced(spliced) => spliced.latest(),
}
}
}
pub struct Request {
name: Arc<str>,
broadcast: Arc<broadcast::Info>,
state: kio::Producer<TrackState>,
prev_subscription: Option<Subscription>,
_dynamic: Dynamic,
stats: stats::Scope,
}
impl Request {
pub(crate) fn new(broadcast: Arc<broadcast::Info>, name: impl Into<Arc<str>>) -> Self {
let name = name.into();
let state = kio::Producer::new(TrackState {
broadcast: broadcast.clone(),
..Default::default()
});
let dynamic = Dynamic::new(name.clone(), state.clone());
Self {
name,
broadcast,
state,
prev_subscription: None,
_dynamic: dynamic,
stats: stats::Scope::default(),
}
}
pub(crate) fn with_stats(mut self, scope: stats::Scope) -> Self {
self.stats = scope;
self
}
pub fn name(&self) -> &str {
&self.name
}
pub fn consume(&self) -> Consumer {
Consumer::plain(self.name.clone(), self.state.consume())
}
pub fn dynamic(&self) -> Dynamic {
Dynamic::new(self.name.clone(), self.state.clone())
}
pub fn poll_unused(&self, waiter: &kio::Waiter) -> Poll<()> {
self.state.poll_unused(waiter).map(|_| ())
}
pub fn accept(self, info: impl Into<Option<Info>>) -> Producer {
let mut info = info.into().unwrap_or_default();
info.broadcast = self.broadcast.clone();
if let Ok(mut state) = self.state.write() {
state.info = Some(info);
}
self.stats.open_subscription();
Producer {
name: self.name,
broadcast: self.broadcast,
state: self.state,
prev_subscription: None,
stats: self.stats,
}
}
pub fn reject(self, err: Error) {
if let Ok(mut state) = self.state.write() {
state.abort = Some(err);
}
}
pub fn subscription(&self) -> Option<Subscription> {
let state = self.state.read();
let (subs, bound) = (state.subscriptions.clone(), state.latency_bound());
drop(state);
snapshot_subscription(&subs, bound)
}
pub async fn subscription_changed(&mut self) -> Option<Subscription> {
kio::wait(|waiter| self.poll_subscription_changed(waiter)).await
}
pub fn poll_subscription_changed(&mut self, waiter: &kio::Waiter) -> Poll<Option<Subscription>> {
let state = self.state.read();
let (subs, bound) = (state.subscriptions.clone(), state.latency_bound());
drop(state);
let prev = &self.prev_subscription;
let mut combined = None;
let mut guard = ready!(subs.poll(waiter, |subs| {
let next = combined_subscription(subs, bound, waiter);
if &next == prev {
Poll::Pending
} else {
combined = next;
Poll::Ready(())
}
}));
guard.retain(|sub| !sub.is_closed());
drop(guard);
self.prev_subscription = combined.clone();
Poll::Ready(combined)
}
pub(super) fn weak(&self) -> TrackWeak {
TrackWeak {
name: self.name.clone(),
state: self.state.weak(),
}
}
}
#[cfg(test)]
use futures::FutureExt;
#[cfg(test)]
#[allow(missing_docs)] impl Subscriber {
pub fn assert_group(&mut self) -> group::Consumer {
self.recv_group()
.now_or_never()
.expect("group would have blocked")
.expect("would have errored")
.expect("track was closed")
}
pub fn assert_no_group(&mut self) {
assert!(
self.recv_group().now_or_never().is_none(),
"recv_group would not have blocked"
);
}
pub fn assert_not_closed(&mut self) {
assert!(self.finished().now_or_never().is_none(), "should not be closed");
}
pub fn assert_closed(&mut self) {
assert!(self.finished().now_or_never().is_some(), "should be closed");
}
pub fn assert_error(&mut self) {
assert!(
self.finished().now_or_never().expect("should not block").is_err(),
"should be error"
);
}
pub fn assert_is_clone(&self, other: &Self) {
assert!(self.is_clone(other), "should be clone");
}
pub fn assert_not_clone(&self, other: &Self) {
assert!(!self.is_clone(other), "should not be clone");
}
}
#[cfg(test)]
mod test {
use super::*;
fn track_producer(name: impl Into<Arc<str>>, info: impl Into<Option<Info>>) -> Producer {
Producer::new(Arc::new(broadcast::Info::default()), name, info)
}
fn live_groups(state: &TrackState) -> usize {
state.groups.iter().flatten().count()
}
fn first_live_sequence(state: &TrackState) -> u64 {
state.groups.iter().flatten().next().unwrap().0.sequence
}
fn recv_datagram(dg: &mut Subscriber) -> Datagram {
dg.recv_datagram()
.now_or_never()
.expect("datagram would have blocked")
.expect("would have errored")
.expect("track was closed")
}
#[tokio::test]
async fn append_datagram_shares_group_sequence() {
let mut producer = track_producer("test", None);
let ts = Timestamp::from_millis(10).unwrap();
assert_eq!(producer.append_group().unwrap().sequence, 0);
assert_eq!(producer.append_datagram(ts, &b"a"[..]).unwrap(), 1);
assert_eq!(producer.append_group().unwrap().sequence, 2);
assert_eq!(producer.append_datagram(ts, &b"b"[..]).unwrap(), 3);
assert_eq!(producer.latest(), Some(3));
}
#[tokio::test]
async fn append_datagram_roundtrip() {
let mut producer = track_producer("test", None);
let mut dg = producer.subscribe(None);
let ts = Timestamp::from_millis(42).unwrap();
let seq = producer.append_datagram(ts, &b"hello"[..]).unwrap();
let got = recv_datagram(&mut dg);
assert_eq!(got.sequence, seq);
assert_eq!(got.timestamp, ts);
assert_eq!(&got.payload[..], b"hello");
}
#[tokio::test]
async fn write_datagram_preserves_sequence() {
let mut producer = track_producer("test", None);
let mut dg = producer.subscribe(None);
let ts = Timestamp::from_millis(5).unwrap();
producer
.write_datagram(Datagram {
sequence: 100,
timestamp: ts,
payload: bytes::Bytes::from_static(b"x"),
})
.unwrap();
assert_eq!(recv_datagram(&mut dg).sequence, 100);
assert_eq!(producer.append_group().unwrap().sequence, 101);
}
#[tokio::test]
async fn recv_datagram_advances_ordered_group_cursor() {
let mut producer = track_producer("test", None);
let mut subscriber = producer.subscribe(None);
let ts = Timestamp::from_millis(5).unwrap();
producer
.write_datagram(Datagram {
sequence: 5,
timestamp: ts,
payload: bytes::Bytes::from_static(b"x"),
})
.unwrap();
assert_eq!(recv_datagram(&mut subscriber).sequence, 5);
producer.create_group(group::Info { sequence: 3 }).unwrap();
producer.create_group(group::Info { sequence: 6 }).unwrap();
let group = subscriber
.next_group()
.now_or_never()
.expect("group would have blocked")
.expect("would have errored")
.expect("track was closed");
assert_eq!(group.sequence, 6);
}
#[tokio::test]
async fn datagram_normalized_to_track_timescale() {
let info = Info::default().with_timescale(Timescale::MICRO);
let mut producer = track_producer("test", info);
let mut dg = producer.subscribe(None);
producer
.append_datagram(Timestamp::from_millis(2).unwrap(), &b"z"[..])
.unwrap();
let got = recv_datagram(&mut dg);
assert_eq!(got.timestamp.scale(), Timescale::MICRO);
assert_eq!(got.timestamp.value(), 2_000);
}
#[tokio::test]
async fn datagram_rejects_oversized() {
let mut producer = track_producer("test", None);
let big = bytes::Bytes::from(vec![0u8; crate::model::datagram::MAX_DATAGRAM_PAYLOAD + 1]);
let ts = Timestamp::from_millis(0).unwrap();
assert!(matches!(
producer.append_datagram(ts, big.clone()),
Err(Error::FrameTooLarge)
));
assert!(matches!(
producer.write_datagram(Datagram {
sequence: 0,
timestamp: ts,
payload: big,
}),
Err(Error::FrameTooLarge)
));
}
#[tokio::test]
async fn datagram_fanout_to_subscribers() {
let mut producer = track_producer("test", None);
let mut a = producer.subscribe(None);
let mut b = producer.subscribe(None);
let ts = Timestamp::from_millis(1).unwrap();
producer.append_datagram(ts, &b"first"[..]).unwrap();
producer.append_datagram(ts, &b"second"[..]).unwrap();
assert_eq!(&recv_datagram(&mut a).payload[..], b"first");
assert_eq!(&recv_datagram(&mut a).payload[..], b"second");
assert_eq!(&recv_datagram(&mut b).payload[..], b"first");
assert_eq!(&recv_datagram(&mut b).payload[..], b"second");
}
#[tokio::test]
async fn datagram_evicts_stale() {
tokio::time::pause();
let mut producer = track_producer("test", None);
let mut dg = producer.subscribe(None);
let ts = Timestamp::from_millis(0).unwrap();
producer.append_datagram(ts, &b"old"[..]).unwrap();
tokio::time::advance(MAX_DATAGRAM_AGE + Duration::from_millis(10)).await;
producer.append_datagram(ts, &b"new"[..]).unwrap();
let got = recv_datagram(&mut dg);
assert_eq!(got.sequence, 1);
assert_eq!(&got.payload[..], b"new");
}
#[tokio::test]
async fn datagram_recv_pends_until_written() {
let mut producer = track_producer("test", None);
let mut dg = producer.subscribe(None);
assert!(
dg.recv_datagram().now_or_never().is_none(),
"should block with no datagrams"
);
producer
.append_datagram(Timestamp::from_millis(0).unwrap(), &b"go"[..])
.unwrap();
assert_eq!(&recv_datagram(&mut dg).payload[..], b"go");
}
#[tokio::test]
async fn datagram_wire_roundtrip_between_tracks() {
use crate::coding::{Decode, Encode};
use crate::lite;
let version = lite::Version::Lite05;
let mut origin = track_producer("test", None);
let mut origin_dg = origin.subscribe(None);
let ts = Timestamp::from_millis(7).unwrap();
let seq = origin.append_datagram(ts, &b"payload"[..]).unwrap();
let d = recv_datagram(&mut origin_dg);
let body = lite::Datagram {
subscribe: 5,
sequence: d.sequence,
timestamp: d.timestamp.value(),
payload: d.payload.clone(),
}
.encode_bytes(version)
.unwrap();
let mut slice = &body[..];
let wire = lite::Datagram::decode(&mut slice, version).unwrap();
let mut downstream = track_producer("test", None);
let mut downstream_dg = downstream.subscribe(None);
downstream
.write_datagram(Datagram {
sequence: wire.sequence,
timestamp: Timestamp::new(wire.timestamp, Timescale::MILLI).unwrap(),
payload: wire.payload,
})
.unwrap();
let got = recv_datagram(&mut downstream_dg);
assert_eq!(got.sequence, seq);
assert_eq!(got.timestamp, ts);
assert_eq!(&got.payload[..], b"payload");
}
#[tokio::test]
async fn evict_expired_groups() {
tokio::time::pause();
let mut producer = track_producer("test", None);
producer.append_group().unwrap(); producer.append_group().unwrap(); producer.append_group().unwrap();
{
let state = producer.state.read();
assert_eq!(live_groups(&state), 3);
assert_eq!(state.offset, 0);
}
tokio::time::advance(DEFAULT_LATENCY_MAX + Duration::from_secs(1)).await;
producer.append_group().unwrap();
{
let state = producer.state.read();
assert_eq!(live_groups(&state), 1);
assert_eq!(first_live_sequence(&state), 3);
assert_eq!(state.offset, 3);
assert!(!state.duplicates.contains(&0));
assert!(!state.duplicates.contains(&1));
assert!(!state.duplicates.contains(&2));
assert!(state.duplicates.contains(&3));
}
}
#[tokio::test]
async fn evict_keeps_max_sequence() {
tokio::time::pause();
let mut producer = track_producer("test", None);
producer.append_group().unwrap();
tokio::time::advance(DEFAULT_LATENCY_MAX + Duration::from_secs(1)).await;
producer.append_group().unwrap();
{
let state = producer.state.read();
assert_eq!(live_groups(&state), 1);
assert_eq!(first_live_sequence(&state), 1);
assert_eq!(state.offset, 1);
}
}
#[tokio::test]
async fn no_eviction_when_fresh() {
tokio::time::pause();
let mut producer = track_producer("test", None);
producer.append_group().unwrap(); producer.append_group().unwrap(); producer.append_group().unwrap();
{
let state = producer.state.read();
assert_eq!(live_groups(&state), 3);
assert_eq!(state.offset, 0);
}
}
#[tokio::test]
async fn consumer_skips_evicted_groups() {
tokio::time::pause();
let mut producer = track_producer("test", None);
producer.append_group().unwrap();
let mut consumer = producer.subscribe(None);
tokio::time::advance(DEFAULT_LATENCY_MAX + Duration::from_secs(1)).await;
producer.append_group().unwrap();
let group = consumer.assert_group();
assert_eq!(group.sequence, 1);
}
#[tokio::test]
async fn cache_age_controls_eviction() {
tokio::time::pause();
let mut producer = track_producer("test", Info::default().with_latency_max(Duration::from_secs(1)));
producer.append_group().unwrap();
tokio::time::advance(Duration::from_secs(2)).await;
producer.append_group().unwrap();
let state = producer.state.read();
assert_eq!(live_groups(&state), 1);
assert_eq!(first_live_sequence(&state), 1);
}
#[test]
fn latency_max_clamped_to_cache() {
let producer = track_producer("test", Info::default().with_latency_max(Duration::from_secs(2)));
let mut subscriber = producer.subscribe(Subscription::default().with_latency_max(Duration::from_secs(10)));
assert_eq!(subscriber.subscription().latency_max, Duration::from_secs(10));
assert_eq!(producer.subscription().unwrap().latency_max, Duration::from_secs(2));
subscriber
.update(Subscription::default().with_latency_max(Duration::from_millis(500)))
.unwrap();
assert_eq!(producer.subscription().unwrap().latency_max, Duration::from_millis(500));
subscriber
.update(Subscription::default().with_latency_max(Duration::ZERO))
.unwrap();
assert_eq!(producer.subscription().unwrap().latency_max, Duration::ZERO);
}
#[test]
fn latency_max_clamped_via_every_update_path() {
let producer = track_producer("test", Info::default().with_latency_max(Duration::from_secs(2)));
let over = Subscription::default().with_latency_max(Duration::from_secs(10));
let mut subscriber = producer.subscribe(over.clone());
assert_eq!(producer.subscription().unwrap().latency_max, Duration::from_secs(2));
subscriber.control().update(over.clone()).unwrap();
assert_eq!(producer.subscription().unwrap().latency_max, Duration::from_secs(2));
subscriber.update(over).unwrap();
assert_eq!(producer.subscription().unwrap().latency_max, Duration::from_secs(2));
}
#[test]
fn latency_max_aggregate_clamps_the_max_across_subscribers() {
let producer = track_producer("test", Info::default().with_latency_max(Duration::from_secs(2)));
let _a = producer.subscribe(Subscription::default().with_latency_max(Duration::from_millis(500)));
let _b = producer.subscribe(Subscription::default().with_latency_max(Duration::from_secs(10)));
assert_eq!(producer.subscription().unwrap().latency_max, Duration::from_secs(2));
}
#[test]
fn subscriber_control_updates_while_read_future_is_pending() {
let producer = track_producer("test", None);
let mut subscriber = producer.subscribe(None);
let control = subscriber.control();
let mut recv = Box::pin(subscriber.recv_group());
assert!(recv.as_mut().now_or_never().is_none());
control
.update(Subscription::default().with_priority(7).with_ordered(false))
.unwrap();
let aggregate = producer.subscription().expect("expected an active subscription");
assert_eq!(aggregate.priority, 7);
assert!(!aggregate.ordered);
}
#[test]
fn dropped_subscriber_leaves_no_ghost_in_aggregate() {
let mut producer = track_producer("test", None);
let a = producer.subscribe(Subscription::default().with_priority(5));
let waiter = kio::Waiter::noop();
assert!(
matches!(producer.poll_subscription_changed(&waiter), Poll::Ready(Ok(Some(_)))),
"one live subscriber should aggregate to Some",
);
drop(a);
assert!(
matches!(producer.poll_subscription_changed(&waiter), Poll::Ready(Ok(None))),
"a dropped subscriber must not linger in the aggregate",
);
assert!(
producer.subscription().is_none(),
"snapshot must exclude a dropped subscriber",
);
}
#[tokio::test]
async fn out_of_order_max_sequence_at_front() {
tokio::time::pause();
let mut producer = track_producer("test", None);
producer.create_group(group::Info { sequence: 5 }).unwrap();
producer.create_group(group::Info { sequence: 3 }).unwrap();
producer.create_group(group::Info { sequence: 4 }).unwrap();
{
let state = producer.state.read();
assert_eq!(state.max_sequence, Some(5));
}
tokio::time::advance(DEFAULT_LATENCY_MAX + Duration::from_secs(1)).await;
producer.append_group().unwrap();
{
let state = producer.state.read();
assert_eq!(live_groups(&state), 1);
assert_eq!(first_live_sequence(&state), 6);
assert!(!state.duplicates.contains(&3));
assert!(!state.duplicates.contains(&4));
assert!(!state.duplicates.contains(&5));
assert!(state.duplicates.contains(&6));
}
}
#[tokio::test]
async fn max_sequence_at_front_blocks_trim() {
tokio::time::pause();
let mut producer = track_producer("test", None);
producer.create_group(group::Info { sequence: 5 }).unwrap();
tokio::time::advance(DEFAULT_LATENCY_MAX + Duration::from_secs(1)).await;
producer.create_group(group::Info { sequence: 3 }).unwrap();
{
let state = producer.state.read();
assert_eq!(live_groups(&state), 2);
assert_eq!(state.offset, 0);
}
tokio::time::advance(DEFAULT_LATENCY_MAX + Duration::from_secs(1)).await;
producer.create_group(group::Info { sequence: 2 }).unwrap();
{
let state = producer.state.read();
assert_eq!(live_groups(&state), 2);
assert_eq!(state.offset, 0);
assert!(state.duplicates.contains(&5));
assert!(!state.duplicates.contains(&3));
assert!(state.duplicates.contains(&2));
}
let mut consumer = producer.subscribe(None);
let group = consumer.assert_group();
assert_eq!(group.sequence, 5);
}
#[tokio::test]
async fn abort_clears_cached_groups() {
let mut producer = track_producer("test", None);
producer.append_group().unwrap();
producer.append_group().unwrap();
let mut consumer = producer.subscribe(None);
assert_eq!(live_groups(&producer.state.read()), 2);
producer.clone().abort(Error::Cancel).unwrap();
{
let state = producer.state.read();
assert!(state.groups.is_empty(), "cached groups should be dropped on abort");
assert!(state.duplicates.is_empty());
}
let result = consumer.recv_group().now_or_never().expect("should not block");
assert!(matches!(result, Err(Error::Cancel)));
}
#[tokio::test]
async fn drop_unfinished_clears_cached_groups() {
let producer = track_producer("test", None);
let mut writer = producer.clone();
writer.append_group().unwrap();
let mut consumer = producer.subscribe(None);
assert_eq!(live_groups(&producer.state.read()), 1);
drop(writer);
drop(producer);
let result = consumer.recv_group().now_or_never().expect("should not block");
assert!(matches!(result, Err(Error::Dropped)));
}
#[tokio::test]
async fn drop_finished_keeps_cached_groups() {
let mut producer = track_producer("test", None);
producer.append_group().unwrap();
producer.finish().unwrap();
let mut consumer = producer.subscribe(None);
drop(producer);
assert_eq!(consumer.assert_group().sequence, 0);
let done = consumer.recv_group().now_or_never().expect("should not block").unwrap();
assert!(done.is_none(), "consumer should drain then see clean finish");
}
#[test]
fn append_finish_cannot_be_rewritten() {
let mut producer = track_producer("test", None);
assert!(producer.finish().is_ok());
assert!(producer.finish().is_err());
assert!(producer.append_group().is_err());
}
#[test]
fn finish_after_groups() {
let mut producer = track_producer("test", None);
producer.append_group().unwrap();
assert!(producer.finish().is_ok());
assert!(producer.finish().is_err());
assert!(producer.append_group().is_err());
}
#[test]
fn finish_at_rejects_a_boundary_at_or_below_the_live_edge() {
let mut producer = track_producer("test", None);
producer.create_group(group::Info { sequence: 5 }).unwrap();
assert!(producer.finish_at(4).is_err());
assert!(producer.finish_at(5).is_err());
assert!(producer.finish_at(6).is_ok());
{
let state = producer.state.read();
assert_eq!(state.final_sequence, Some(6));
}
assert!(producer.finish_at(6).is_err());
assert!(producer.create_group(group::Info { sequence: 4 }).is_ok());
assert!(producer.create_group(group::Info { sequence: 6 }).is_err());
}
#[test]
fn final_sequence_reports_the_declared_boundary() {
let mut producer = track_producer("test", None);
assert_eq!(producer.final_sequence(), None);
producer.create_group(group::Info { sequence: 5 }).unwrap();
assert_eq!(producer.final_sequence(), None, "a group does not declare a boundary");
producer.finish_at(9).unwrap();
assert_eq!(producer.final_sequence(), Some(9));
assert!(producer.finish().is_err());
}
#[test]
fn final_sequence_reports_the_live_edge_after_finish() {
let mut producer = track_producer("test", None);
producer.create_group(group::Info { sequence: 5 }).unwrap();
producer.finish().unwrap();
assert_eq!(producer.final_sequence(), Some(6));
}
#[tokio::test]
async fn finish_at_declares_a_future_boundary() {
let mut producer = track_producer("test", None);
producer.create_group(group::Info { sequence: 5 }).unwrap();
producer.finish_at(7).unwrap();
let mut consumer = producer.subscribe(None);
assert_eq!(consumer.assert_group().sequence, 5);
let boundary = consumer
.finished()
.now_or_never()
.expect("boundary is known immediately")
.expect("would have errored");
assert_eq!(boundary, 7);
assert!(
consumer.recv_group().now_or_never().is_none(),
"should wait for the outstanding group"
);
producer.create_group(group::Info { sequence: 6 }).unwrap();
assert_eq!(consumer.assert_group().sequence, 6);
let done = consumer
.recv_group()
.now_or_never()
.expect("should not block")
.expect("would have errored");
assert!(done.is_none(), "track completes once the boundary is reached");
}
#[tokio::test]
async fn recv_group_finishes_without_waiting_for_gaps() {
let mut producer = track_producer("test", None);
producer.create_group(group::Info { sequence: 1 }).unwrap();
producer.finish().unwrap();
let mut consumer = producer.subscribe(None);
assert_eq!(consumer.assert_group().sequence, 1);
let done = consumer
.recv_group()
.now_or_never()
.expect("should not block")
.expect("would have errored");
assert!(done.is_none(), "track should finish without waiting for gaps");
}
#[tokio::test]
async fn next_group_skips_late_arrivals() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
producer.create_group(group::Info { sequence: 5 }).unwrap();
let group = consumer
.next_group()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(group.sequence, 5);
producer.create_group(group::Info { sequence: 3 }).unwrap();
producer.create_group(group::Info { sequence: 4 }).unwrap();
producer.create_group(group::Info { sequence: 7 }).unwrap();
let group = consumer
.next_group()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(group.sequence, 7);
assert!(
consumer.next_group().now_or_never().is_none(),
"should block waiting for a higher sequence"
);
}
#[tokio::test]
async fn next_group_returns_arrivals_in_order() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
producer.create_group(group::Info { sequence: 3 }).unwrap();
producer.create_group(group::Info { sequence: 5 }).unwrap();
let group = consumer
.next_group()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(group.sequence, 3);
let group = consumer
.next_group()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(group.sequence, 5);
}
#[tokio::test]
async fn next_group_and_recv_group_use_independent_cursors() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
producer.create_group(group::Info { sequence: 5 }).unwrap();
producer.create_group(group::Info { sequence: 3 }).unwrap();
let group = consumer
.next_group()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(group.sequence, 3);
assert_eq!(consumer.assert_group().sequence, 5);
}
#[tokio::test]
async fn end_at_caps_next_group() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
for s in 0..6 {
producer.create_group(group::Info { sequence: s }).unwrap();
}
consumer.end_at(2);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
0
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
1
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
2
);
assert!(
consumer.next_group().now_or_never().is_none(),
"capped consumer must block instead of returning out-of-range groups"
);
}
#[tokio::test]
async fn end_at_release_drains_cached_groups() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
for s in 0..6 {
producer.create_group(group::Info { sequence: s }).unwrap();
}
consumer.end_at(1);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
0
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
1
);
assert!(consumer.next_group().now_or_never().is_none(), "capped at 1");
consumer.end_at(4);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
2
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
3
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
4
);
assert!(consumer.next_group().now_or_never().is_none(), "capped at 4");
consumer.end_at(None);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
5
);
assert!(consumer.next_group().now_or_never().is_none(), "no more groups");
}
#[tokio::test]
async fn end_at_lower_than_cursor_parks_consumer() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
for s in 0..3 {
producer.create_group(group::Info { sequence: s }).unwrap();
}
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
0
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
1
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
2
);
consumer.end_at(1);
producer.create_group(group::Info { sequence: 3 }).unwrap();
producer.create_group(group::Info { sequence: 4 }).unwrap();
assert!(
consumer.next_group().now_or_never().is_none(),
"cap is below cursor; nothing returnable until cap rises"
);
consumer.end_at(None);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
3
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
4
);
}
#[tokio::test]
async fn end_at_toggling_around_late_arrivals() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
consumer.end_at(5);
producer.create_group(group::Info { sequence: 2 }).unwrap();
producer.create_group(group::Info { sequence: 5 }).unwrap();
producer.create_group(group::Info { sequence: 3 }).unwrap();
producer.create_group(group::Info { sequence: 8 }).unwrap();
producer.create_group(group::Info { sequence: 4 }).unwrap();
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
2
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
3
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
4
);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
5
);
assert!(consumer.next_group().now_or_never().is_none());
consumer.end_at(10);
assert_eq!(
consumer.next_group().now_or_never().unwrap().unwrap().unwrap().sequence,
8
);
}
#[tokio::test]
async fn read_frame_returns_single_frame_per_group() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
producer.write_frame(Timestamp::ZERO, b"hello".as_slice()).unwrap();
producer.write_frame(Timestamp::ZERO, b"world".as_slice()).unwrap();
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(&frame.payload[..], b"hello");
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(&frame.payload[..], b"world");
}
#[tokio::test]
async fn read_frame_preserves_timestamp() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
producer
.write_frame(Timestamp::from_micros(20_000).unwrap(), b"hello".as_slice())
.unwrap();
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(frame.timestamp.as_micros(), 20_000);
assert_eq!(&frame.payload[..], b"hello");
}
#[tokio::test]
async fn read_frame_skips_stalled_group_for_newer_ready_frame() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
let _stalled = producer.create_group(group::Info { sequence: 3 }).unwrap();
let mut g5 = producer.create_group(group::Info { sequence: 5 }).unwrap();
g5.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"later"))
.unwrap();
g5.finish().unwrap();
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block on stalled earlier group")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(&frame.payload[..], b"later");
}
#[tokio::test]
async fn read_frame_discards_rest_of_multi_frame_group() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
let mut g0 = producer.create_group(group::Info { sequence: 0 }).unwrap();
g0.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"one"))
.unwrap();
g0.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"two"))
.unwrap();
g0.finish().unwrap();
producer.write_frame(Timestamp::ZERO, b"next".as_slice()).unwrap();
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(&frame.payload[..], b"one");
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(&frame.payload[..], b"next");
}
#[tokio::test]
async fn read_frame_waits_for_pending_group_after_finish() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
let mut g0 = producer.create_group(group::Info { sequence: 0 }).unwrap();
producer.finish().unwrap();
assert!(
consumer.read_frame().now_or_never().is_none(),
"read_frame must block on a pending group even after finish()"
);
g0.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"late"))
.unwrap();
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block once a frame is written")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(&frame.payload[..], b"late");
}
#[tokio::test]
async fn read_frame_respects_start_at() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
consumer.start_at(5);
let mut g3 = producer.create_group(group::Info { sequence: 3 }).unwrap();
g3.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"skip-me"))
.unwrap();
g3.finish().unwrap();
let mut g5 = producer.create_group(group::Info { sequence: 5 }).unwrap();
g5.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"keep"))
.unwrap();
g5.finish().unwrap();
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(&frame.payload[..], b"keep");
}
#[tokio::test]
async fn read_frame_returns_none_when_finished() {
let mut producer = track_producer("test", None);
let mut consumer = producer.subscribe(None);
producer.write_frame(Timestamp::ZERO, b"only".as_slice()).unwrap();
producer.finish().unwrap();
let frame = consumer
.read_frame()
.now_or_never()
.expect("should not block")
.expect("would have errored")
.expect("track should not be closed");
assert_eq!(&frame.payload[..], b"only");
let done = consumer
.read_frame()
.now_or_never()
.expect("should not block")
.expect("would have errored");
assert!(done.is_none());
}
#[test]
fn append_group_returns_bounds_exceeded_on_sequence_overflow() {
let mut producer = track_producer("test", None);
{
let mut state = producer.state.write().ok().unwrap();
state.max_sequence = Some(u64::MAX);
}
assert!(matches!(producer.append_group(), Err(Error::BoundsExceeded(_))));
}
#[tokio::test]
async fn fetch_cache_hit() {
let mut producer = track_producer("test", None);
let mut group = producer.append_group().unwrap(); group
.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"hello"))
.unwrap();
group.finish().unwrap();
let dynamic = producer.dynamic();
let consumer = producer.consume();
assert!(consumer.peek_group(0).is_some());
let mut g = consumer.fetch_group(0, None).await.unwrap();
assert_eq!(g.sequence, 0);
assert_eq!(&g.read_frame().await.unwrap().unwrap().payload[..], b"hello");
assert!(dynamic.poll_requested_group(&kio::Waiter::noop()).is_pending());
}
#[tokio::test]
async fn fetch_miss_signals_dynamic() {
let producer = track_producer("test", None);
let dynamic = producer.dynamic();
let consumer = producer.consume();
assert!(consumer.peek_group(5).is_none());
let pending = consumer.fetch_group(5, group::Fetch::default().with_priority(7));
assert!(kio::Pollable::poll(&*pending, &kio::Waiter::noop()).is_pending());
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
assert_eq!(req.sequence(), 5);
assert_eq!(req.priority(), 7);
let mut group = req.accept(None).unwrap();
group
.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"hi"))
.unwrap();
group.finish().unwrap();
let mut g = pending.await.unwrap();
assert_eq!(g.sequence, 5);
assert_eq!(&g.read_frame().await.unwrap().unwrap().payload[..], b"hi");
}
#[tokio::test]
async fn fetch_miss_rejects() {
let producer = track_producer("test", None);
let dynamic = producer.dynamic();
let consumer = producer.consume();
let pending = consumer.fetch_group(5, None);
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
req.reject(Error::Cancel);
assert!(matches!(pending.await, Err(Error::Cancel)));
let fetch = producer.state.read().fetch.clone();
assert!(fetch.read().is_empty());
}
#[tokio::test]
async fn fetch_miss_drop_rejects() {
let producer = track_producer("test", None);
let dynamic = producer.dynamic();
let consumer = producer.consume();
let pending = consumer.fetch_group(5, None);
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
drop(req);
assert!(matches!(pending.await, Err(Error::Dropped)));
}
#[tokio::test]
async fn fetch_reject_does_not_poison_retry() {
let producer = track_producer("test", None);
let dynamic = producer.dynamic();
let consumer = producer.consume();
let pending = consumer.fetch_group(5, None);
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
req.reject(Error::Cancel);
assert!(matches!(pending.await, Err(Error::Cancel)));
let retry = consumer.fetch_group(5, None);
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
let mut group = req.accept(None).unwrap();
group
.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"retry"))
.unwrap();
group.finish().unwrap();
let mut group = retry.await.unwrap();
assert_eq!(&group.read_frame().await.unwrap().unwrap().payload[..], b"retry");
}
#[tokio::test]
async fn fetch_coalesces_concurrent() {
let producer = track_producer("test", None);
let dynamic = producer.dynamic();
let consumer = producer.consume();
let first = consumer.fetch_group(5, group::Fetch::default().with_priority(1));
let second = consumer.fetch_group(5, group::Fetch::default().with_priority(7));
assert!(kio::Pollable::poll(&*first, &kio::Waiter::noop()).is_pending());
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
assert_eq!(req.sequence(), 5);
assert_eq!(req.priority(), 7);
assert!(
dynamic.poll_requested_group(&kio::Waiter::noop()).is_pending(),
"the second fetch queued a duplicate request"
);
let third = consumer.fetch_group(5, None);
let mut group = req.accept(None).unwrap();
group
.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"hi"))
.unwrap();
group.finish().unwrap();
assert_eq!(first.await.unwrap().sequence, 5);
assert_eq!(second.await.unwrap().sequence, 5);
assert_eq!(third.await.unwrap().sequence, 5);
}
#[tokio::test]
async fn fetch_coalesced_reject_fails_all() {
let producer = track_producer("test", None);
let dynamic = producer.dynamic();
let consumer = producer.consume();
let first = consumer.fetch_group(5, None);
let second = consumer.fetch_group(5, None);
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
req.reject(Error::Cancel);
assert!(matches!(first.await, Err(Error::Cancel)));
assert!(matches!(second.await, Err(Error::Cancel)));
let retry = consumer.fetch_group(5, None);
assert!(kio::Pollable::poll(&*retry, &kio::Waiter::noop()).is_pending());
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
assert_eq!(req.sequence(), 5);
}
#[tokio::test]
async fn fetch_queued_fails_when_handlers_leave() {
let producer = track_producer("test", None);
let dynamic = producer.dynamic();
let consumer = producer.consume();
let pending = consumer.fetch_group(5, None);
assert!(kio::Pollable::poll(&*pending, &kio::Waiter::noop()).is_pending());
drop(dynamic);
assert!(matches!(pending.await, Err(Error::NotFound)));
let fetch = producer.state.read().fetch.clone();
assert!(fetch.read().is_empty());
}
#[tokio::test]
async fn fetch_miss_no_dynamic_not_found() {
let mut producer = track_producer("test", None);
producer.append_group().unwrap(); let consumer = producer.consume();
assert!(matches!(consumer.fetch_group(5, None).await, Err(Error::NotFound)));
}
#[tokio::test]
async fn fetch_past_final_not_found() {
let mut producer = track_producer("test", None);
producer.append_group().unwrap(); producer.finish().unwrap();
let dynamic = producer.dynamic();
let consumer = producer.consume();
assert!(matches!(consumer.fetch_group(5, None).await, Err(Error::NotFound)));
assert!(dynamic.poll_requested_group(&kio::Waiter::noop()).is_pending());
}
fn pooled_producer(capacity: u64) -> (Producer, cache::Pool) {
let pool = cache::Pool::new(capacity);
let broadcast = broadcast::Info {
origin: crate::origin::Info::default().with_pool(pool.clone()),
..Default::default()
};
let producer = Producer::new(Arc::new(broadcast), "test", None);
(producer, pool)
}
fn finished_group(producer: &mut Producer, size: usize) -> u64 {
let mut group = producer.append_group().unwrap();
group
.write_frame(Timestamp::ZERO, bytes::Bytes::from(vec![0u8; size]))
.unwrap();
group.finish().unwrap();
group.sequence
}
#[tokio::test]
async fn pool_evicts_oldest_group() {
tokio::time::pause();
let (mut producer, pool) = pooled_producer(3000);
finished_group(&mut producer, 1000); tokio::time::advance(Duration::from_millis(10)).await;
finished_group(&mut producer, 1000); tokio::time::advance(Duration::from_millis(10)).await;
finished_group(&mut producer, 1000);
assert!(pool.used() <= 3000, "pool should be back under budget");
let consumer = producer.consume();
assert!(consumer.peek_group(0).is_none(), "evicted group is a cache miss");
assert!(consumer.peek_group(1).is_some());
assert!(consumer.peek_group(2).is_some());
let mut subscriber = producer.subscribe(None);
assert_eq!(subscriber.assert_group().sequence, 1);
assert_eq!(subscriber.assert_group().sequence, 2);
}
#[tokio::test]
async fn pool_never_evicts_latest() {
tokio::time::pause();
let (mut producer, pool) = pooled_producer(100);
finished_group(&mut producer, 1000);
assert!(pool.used() > 100, "pinned latest may exceed the budget");
let mut subscriber = producer.subscribe(None);
let mut group = subscriber.assert_group();
assert_eq!(group.read_frame().await.unwrap().unwrap().payload.len(), 1000);
}
#[tokio::test]
async fn pool_reads_bump_recency() {
tokio::time::pause();
let (mut producer, pool) = pooled_producer(3000);
let mut subscriber = producer.subscribe(None);
finished_group(&mut producer, 1000); tokio::time::advance(Duration::from_millis(10)).await;
finished_group(&mut producer, 1000); tokio::time::advance(Duration::from_millis(10)).await;
let mut group = subscriber.assert_group();
assert_eq!(group.sequence, 0);
group.read_frame().await.unwrap().unwrap();
tokio::time::advance(Duration::from_millis(10)).await;
finished_group(&mut producer, 1000);
let consumer = producer.consume();
assert!(
consumer.peek_group(0).is_some(),
"recently read group survives: {:?}",
pool.debug_entries()
);
assert!(
consumer.peek_group(1).is_none(),
"stale group is evicted: {:?}",
pool.debug_entries()
);
}
#[tokio::test]
async fn pool_eviction_aborts_readers() {
tokio::time::pause();
let (mut producer, pool) = pooled_producer(3000);
let mut subscriber = producer.subscribe(None);
finished_group(&mut producer, 1000); let group0 = subscriber.assert_group();
tokio::time::advance(Duration::from_millis(10)).await;
finished_group(&mut producer, 1000); tokio::time::advance(Duration::from_millis(10)).await;
finished_group(&mut producer, 1000);
let mut group0 = group0;
let read = group0.read_frame().await;
assert!(
matches!(read, Err(Error::Evicted)),
"expected Evicted, got {read:?}: {:?}",
pool.debug_entries()
);
}
#[tokio::test]
async fn pool_growth_on_old_group_charges() {
tokio::time::pause();
let (mut producer, pool) = pooled_producer(3000);
let mut group0 = producer.append_group().unwrap();
tokio::time::advance(Duration::from_millis(10)).await;
let _group1 = producer.append_group().unwrap();
tokio::time::advance(Duration::from_millis(10)).await;
group0
.write_frame(Timestamp::ZERO, bytes::Bytes::from(vec![0u8; 4000]))
.unwrap();
assert!(pool.used() <= 3000, "growth on an old group triggers eviction");
assert!(matches!(group0.abort(Error::Cancel), Err(Error::Evicted)));
}
#[tokio::test]
async fn refetched_latest_group_is_repinned() {
tokio::time::pause();
let (mut producer, pool) = pooled_producer(3000);
let dynamic = producer.dynamic();
let mut straggler = producer.append_group().unwrap();
straggler
.write_frame(Timestamp::ZERO, bytes::Bytes::from(vec![0u8; 1000]))
.unwrap();
tokio::time::advance(Duration::from_millis(10)).await;
let latest = producer.append_group().unwrap(); latest.abort(Error::Cancel).unwrap();
tokio::time::advance(Duration::from_millis(10)).await;
let consumer = producer.consume();
let pending = consumer.fetch_group(1, None);
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
let mut group = req.accept(None).unwrap();
group
.write_frame(Timestamp::ZERO, bytes::Bytes::from(vec![0u8; 1000]))
.unwrap();
group.finish().unwrap();
pending.await.unwrap();
tokio::time::advance(Duration::from_millis(10)).await;
straggler
.write_frame(Timestamp::ZERO, bytes::Bytes::from(vec![0u8; 4000]))
.unwrap();
assert!(pool.used() <= 3000);
let mut group = consumer.peek_group(1).expect("refetched latest must stay pinned");
assert_eq!(group.read_frame().await.unwrap().unwrap().payload.len(), 1000);
}
#[tokio::test]
async fn pool_eviction_allows_refetch() {
tokio::time::pause();
let (mut producer, _pool) = pooled_producer(3000);
let dynamic = producer.dynamic();
finished_group(&mut producer, 1000); tokio::time::advance(Duration::from_millis(10)).await;
finished_group(&mut producer, 1000); tokio::time::advance(Duration::from_millis(10)).await;
finished_group(&mut producer, 1000);
let consumer = producer.consume();
assert!(consumer.peek_group(0).is_none());
let pending = consumer.fetch_group(0, None);
let req = dynamic
.requested_group()
.now_or_never()
.expect("should not block")
.unwrap();
assert_eq!(req.sequence(), 0);
let mut group = req.accept(None).unwrap();
group
.write_frame(Timestamp::ZERO, bytes::Bytes::from_static(b"refetched"))
.unwrap();
group.finish().unwrap();
let mut group = pending.await.unwrap();
assert_eq!(&group.read_frame().await.unwrap().unwrap().payload[..], b"refetched");
}
#[tokio::test]
async fn fetch_aborts_with_track() {
let producer = track_producer("test", None);
let dynamic = producer.dynamic();
let consumer = producer.consume();
let pending = consumer.fetch_group(3, None);
assert!(kio::Pollable::poll(&*pending, &kio::Waiter::noop()).is_pending());
producer.abort(Error::Cancel).unwrap();
assert!(pending.await.is_err());
drop(dynamic);
}
}