use crate::cache;
use crate::frame::{self, Frame, FrameBuf};
use crate::{Cap, Timescale, stats, track};
use std::collections::VecDeque;
use std::mem::MaybeUninit;
use std::ops::{Bound, RangeBounds};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Poll, ready};
use crate::{Error, IntoBytes, Result, Timestamp};
pub const MAX_CACHE_BYTES: u64 = 32 * 1024 * 1024;
pub const MAX_GROUP_FRAMES: usize = 8192;
const FRAME_SLOTS: usize = 4;
pub(crate) const CACHE_OVERHEAD: u64 = (kio::Producer::<GroupState>::HEAP
+ 2 * size_of::<usize>()
+ size_of::<Alive>()
+ FRAME_SLOTS * size_of::<Frame>()) as u64;
#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
pub struct Info {
pub sequence: u64,
}
impl Info {
#[cfg(test)]
pub(crate) fn produce(self) -> Producer {
Producer::new(self, track::Info::default(), Default::default())
}
}
impl From<usize> for Info {
fn from(sequence: usize) -> Self {
Self {
sequence: sequence as u64,
}
}
}
impl From<u64> for Info {
fn from(sequence: u64) -> Self {
Self { sequence }
}
}
impl From<u32> for Info {
fn from(sequence: u32) -> Self {
Self {
sequence: sequence as u64,
}
}
}
impl From<u16> for Info {
fn from(sequence: u16) -> Self {
Self {
sequence: sequence as u64,
}
}
}
pub(crate) struct Partial {
timestamp: Timestamp,
buf: FrameBuf,
}
#[derive(Default)]
pub(crate) struct GroupState {
pub(crate) frames: VecDeque<Frame>,
pub(crate) partial: Option<Partial>,
pub(crate) offset: usize,
next_index: usize,
committed: usize,
pub(crate) cache: u64,
charge: cache::Charge,
timestamp: Option<Timestamp>,
latest: Option<Timestamp>,
pub(crate) fin: Option<usize>,
pub(crate) abort: Option<Error>,
}
impl GroupState {
fn content(&self) -> stats::Content {
stats::Content {
bytes: self.cache,
frames: self.next_index.saturating_sub(self.offset) as u64,
groups: 1,
datagrams: 0,
}
}
pub(crate) fn content_range(&self, start: usize, end: usize) -> stats::Content {
let start = start.max(self.offset);
let end = end.min(self.next_index);
if start >= end {
return stats::Content::default();
}
let local_start = start.saturating_sub(self.offset).min(self.frames.len());
let local_end = end.saturating_sub(self.offset).min(self.frames.len());
let mut bytes = self
.frames
.range(local_start..local_end)
.map(|frame| frame.payload.len() as u64)
.sum();
if start <= self.committed
&& self.committed < end
&& let Some(partial) = &self.partial
{
bytes += partial.buf.capacity() as u64;
}
stats::Content {
bytes,
frames: (end - start) as u64,
groups: 0,
datagrams: 0,
}
}
fn poll_frame_source(&self, index: usize) -> Poll<Result<Option<(frame::Info, frame::Source)>>> {
if index < self.offset {
return Poll::Ready(Err(Error::Lagged));
}
let local = index - self.offset;
if let Some(f) = self.frames.get(local) {
self.charge.refresh();
let info = frame::Info {
size: f.payload.len() as u64,
timestamp: f.timestamp,
};
return Poll::Ready(Ok(Some((info, frame::Source::Complete(f.payload.clone())))));
}
if local == self.frames.len()
&& let Some(p) = &self.partial
{
self.charge.refresh();
let info = frame::Info {
size: p.buf.capacity() as u64,
timestamp: p.timestamp,
};
return Poll::Ready(Ok(Some((info, frame::Source::Partial(p.buf.clone())))));
}
ready!(self.poll_terminal(index))?;
Poll::Ready(Ok(None))
}
fn poll_terminal(&self, index: usize) -> Poll<Result<()>> {
match (self.fin, &self.abort) {
(Some(total), Some(err)) if index < total => Poll::Ready(Err(err.clone())),
(Some(_), _) => Poll::Ready(Ok(())),
(None, Some(err)) => Poll::Ready(Err(err.clone())),
(None, None) => Poll::Pending,
}
}
fn poll_end(&self, index: usize) -> Poll<Result<()>> {
if index < self.offset {
return Poll::Ready(Err(Error::Lagged));
}
self.poll_terminal(index)
}
fn stamp(&mut self, timestamp: Timestamp) {
self.timestamp.get_or_insert(timestamp);
self.latest = Some(timestamp);
}
fn would_overflow(&self, extra_frames: usize, extra_bytes: u64) -> bool {
self.next_index.saturating_sub(self.offset).saturating_add(extra_frames) > MAX_GROUP_FRAMES
|| self.cache.saturating_add(extra_bytes) > MAX_CACHE_BYTES
}
fn release(&mut self) {
self.frames.clear();
self.partial = None;
self.cache = 0;
self.charge.clear();
}
}
fn modify(state: &kio::Producer<GroupState>) -> Result<kio::Mut<'_, GroupState>> {
state.write().map_err(|r| r.abort.clone().unwrap_or(Error::Dropped))
}
pub struct Producer {
state: kio::Producer<GroupState>,
info: Info,
track: track::Info,
cache: Arc<cache::Track>,
stats: stats::Meter,
alive: Arc<Alive>,
}
struct Alive {
info: Info,
state: kio::Producer<GroupState>,
aborted: AtomicBool,
access: Arc<cache::Access>,
}
impl Drop for Alive {
fn drop(&mut self) {
match self.state.write() {
Ok(mut state) => {
if state.fin.is_some() || state.abort.is_some() {
return;
}
tracing::warn!(
sequence = self.info.sequence,
"group::Producer dropped without finish() or abort()"
);
state.release();
}
Err(state) => {
if state.fin.is_some() || state.abort.is_some() {
return;
}
tracing::warn!(
sequence = self.info.sequence,
"group::Producer dropped without finish() or abort()"
);
}
}
}
}
impl std::ops::Deref for Producer {
type Target = Info;
fn deref(&self) -> &Self::Target {
&self.info
}
}
impl Producer {
pub(crate) fn new(info: Info, track: track::Info, cache: Arc<cache::Track>) -> Self {
let state = kio::Producer::<GroupState>::default();
let charge = cache.charge();
let access = charge.access();
state.write().ok().expect("a new group is open").charge = charge;
let alive = Arc::new(Alive {
info,
state: state.clone(),
aborted: AtomicBool::new(false),
access,
});
Self {
info,
state,
track,
cache,
stats: stats::Meter::default(),
alive,
}
}
pub(crate) fn with_meter(mut self, meter: stats::Meter) -> Self {
meter.group();
self.stats = meter;
self
}
pub(crate) fn info(&self) -> Info {
self.info
}
pub fn timescale(&self) -> Timescale {
self.track.timescale
}
pub fn start_at(&mut self, index: u64) -> Result<()> {
let index = usize::try_from(index).map_err(|_| Error::BoundsExceeded(crate::coding::BoundsExceeded))?;
if index == usize::MAX {
return Err(Error::BoundsExceeded(crate::coding::BoundsExceeded));
}
let mut state = modify(&self.state)?;
if state.fin.is_some() || state.next_index != state.offset {
return Err(Error::Closed);
}
state.offset = index;
state.next_index = index;
state.committed = index;
Ok(())
}
pub fn write_frame<B: IntoBytes>(&mut self, timestamp: Timestamp, data: B) -> Result<()> {
let timestamp = timestamp
.convert(self.track.timescale)
.map_err(|_| Error::TimestampMismatch)?;
let payload = data.into_bytes();
if payload.len() as u64 > MAX_CACHE_BYTES {
return Err(Error::FrameTooLarge);
}
let mut state = modify(&self.state)?;
if state.fin.is_some() {
return Err(Error::Closed);
}
if state.partial.is_some() {
return Err(Error::FrameOpen);
}
let next_index = state
.next_index
.checked_add(1)
.ok_or(Error::BoundsExceeded(crate::coding::BoundsExceeded))?;
debug_assert!(state.partial.is_none(), "a frame is already open");
let size = payload.len() as u64;
if state.would_overflow(1, size) {
return Err(self.abort_too_large(state));
}
state.cache += size;
let now = state.charge.add(size);
state.frames.push_back(Frame { timestamp, payload });
state.next_index = next_index;
state.committed = state.next_index;
state.stamp(timestamp);
drop(state);
self.cache.settle(now);
self.stats.frames(1);
self.stats.bytes(size);
Ok(())
}
pub fn write_frames<const N: usize>(&mut self, frames: &mut frame::Buffer<N>) -> Result<()> {
for frame in frames.filled() {
frame
.timestamp
.convert(self.track.timescale)
.map_err(|_| Error::TimestampMismatch)?;
if frame.payload.len() as u64 > MAX_CACHE_BYTES {
return Err(Error::FrameTooLarge);
}
}
let count = frames.len();
let bytes: u64 = frames.filled().iter().map(|frame| frame.payload.len() as u64).sum();
let mut state = modify(&self.state)?;
if state.fin.is_some() {
return Err(Error::Closed);
}
if state.partial.is_some() {
return Err(Error::FrameOpen);
}
let next_index = state
.next_index
.checked_add(count)
.ok_or(Error::BoundsExceeded(crate::coding::BoundsExceeded))?;
if state.would_overflow(count, bytes) {
return Err(self.abort_too_large(state));
}
let mut now = None;
for mut frame in frames.drain() {
frame.timestamp = frame
.timestamp
.convert(self.track.timescale)
.expect("timestamp scale checked above");
let size = frame.payload.len() as u64;
state.cache += size;
now = state.charge.add(size);
state.stamp(frame.timestamp);
state.frames.push_back(frame);
}
state.next_index = next_index;
state.committed = next_index;
drop(state);
self.cache.settle(now);
self.stats.frames(count as u64);
self.stats.bytes(bytes);
Ok(())
}
pub fn create_frame(&mut self, frame: frame::Info) -> Result<frame::Producer<'_>> {
let timestamp = frame
.timestamp
.convert(self.track.timescale)
.map_err(|_| Error::TimestampMismatch)?;
if frame.size > MAX_CACHE_BYTES {
return Err(Error::FrameTooLarge);
}
let buf = FrameBuf::new(frame.size as usize);
let mut state = modify(&self.state)?;
if state.fin.is_some() {
return Err(Error::Closed);
}
if state.partial.is_some() {
return Err(Error::FrameOpen);
}
let next_index = state
.next_index
.checked_add(1)
.ok_or(Error::BoundsExceeded(crate::coding::BoundsExceeded))?;
if state.would_overflow(1, frame.size) {
return Err(self.abort_too_large(state));
}
state.cache += frame.size;
let now = state.charge.add(frame.size);
state.partial = Some(Partial {
timestamp,
buf: buf.clone(),
});
state.next_index = next_index;
state.stamp(timestamp);
drop(state);
self.cache.settle(now);
self.stats.frames(1);
let meter = self.stats.clone();
let info = frame::Info {
size: frame.size,
timestamp,
};
Ok(frame::Producer::new(self, buf, info).with_meter(meter))
}
pub(crate) fn create_frame_owned(&mut self, frame: frame::Info) -> Result<frame::ProducerOwned> {
let timestamp = frame
.timestamp
.convert(self.track.timescale)
.map_err(|_| Error::TimestampMismatch)?;
if frame.size > MAX_CACHE_BYTES {
return Err(Error::FrameTooLarge);
}
let buf = FrameBuf::new(frame.size as usize);
let mut state = modify(&self.state)?;
if state.fin.is_some() {
return Err(Error::Closed);
}
if state.partial.is_some() {
return Err(Error::FrameOpen);
}
let next_index = state
.next_index
.checked_add(1)
.ok_or(Error::BoundsExceeded(crate::coding::BoundsExceeded))?;
if state.would_overflow(1, frame.size) {
return Err(self.abort_too_large(state));
}
state.cache += frame.size;
let now = state.charge.add(frame.size);
state.partial = Some(Partial {
timestamp,
buf: buf.clone(),
});
state.next_index = next_index;
state.stamp(timestamp);
drop(state);
self.cache.settle(now);
self.stats.frames(1);
let meter = self.stats.clone();
let info = frame::Info {
size: frame.size,
timestamp,
};
Ok(frame::ProducerOwned::new(self.clone(), buf, info).with_meter(meter))
}
pub(crate) fn frame_notify(&self) {
let now = self
.state
.write()
.ok()
.and_then(|mut state| state.charge.record_write());
self.cache.settle(now);
}
pub(crate) fn frame_commit(&mut self, frame: Frame) -> Result<()> {
let mut state = modify(&self.state)?;
state.partial = None;
state.frames.push_back(frame);
state.committed = state.next_index;
let now = state.charge.record_write();
drop(state);
self.cache.settle(now);
Ok(())
}
pub(crate) fn frame_abort(&mut self, err: Error) {
let _ = self.clone().abort(err);
}
pub fn frame_count(&self) -> usize {
self.state.read().next_index
}
pub fn finish(&self) -> Result<()> {
let mut state = modify(&self.state)?;
if state.partial.is_some() {
return Err(Error::FrameOpen);
}
state.fin = Some(state.next_index);
Ok(())
}
pub fn abort(self, err: Error) -> Result<()> {
let mut guard = modify(&self.state)?;
guard.abort = Some(err);
self.alive.aborted.store(true, Ordering::Release);
guard.release();
guard.close();
Ok(())
}
fn abort_too_large(&self, mut state: kio::Mut<'_, GroupState>) -> Error {
let err = Error::GroupTooLarge;
state.abort = Some(err.clone());
self.alive.aborted.store(true, Ordering::Release);
state.release();
state.close();
err
}
pub(crate) fn is_aborted(&self) -> bool {
self.alive.aborted.load(Ordering::Acquire)
}
pub(crate) fn is_finished(&self) -> bool {
self.state.read().fin.is_some()
}
pub(crate) fn live_first_frame(&self) -> Option<usize> {
let state = self.state.read();
state.abort.is_none().then_some(state.offset)
}
pub(crate) fn resume_frame(&self) -> Option<usize> {
let state = self.state.read();
if state.fin.is_some() {
return None;
}
(state.committed > state.offset).then_some(state.committed)
}
pub(crate) fn timestamp(&self) -> Option<Timestamp> {
self.state.read().timestamp
}
pub(crate) fn latest(&self) -> Option<Timestamp> {
self.state.read().latest
}
pub(crate) fn cache_size(&self) -> u64 {
self.state.read().charge.size()
}
pub(crate) fn cache_accessed(&self) -> u64 {
self.alive.access.get()
}
pub(crate) fn cache_accessed_tick(&self, now: Option<u64>) -> Option<u64> {
self.alive.access.tick(now)
}
pub(crate) fn cache_demote(&self) {
if let Ok(mut state) = self.state.write() {
state.charge.demote();
}
}
pub(crate) fn cache_refresh(&self) {
self.state.read().charge.refresh();
}
pub fn consume(&self) -> Consumer {
Consumer {
info: self.info,
track: self.track.clone(),
inner: ConsumerKind::Plain(Plain {
state: self.state.consume(),
index: 0,
end: None,
prefetch: Prefetch::default(),
cache: self.cache.clone(),
access: self.alive.access.clone(),
refreshed: self.cache.pool().now(),
}),
stats: stats::Meter::default(),
stale_stats: stats::Meter::default(),
expiry: None,
expired: false,
ended: false,
stale_counted: Arc::default(),
}
}
pub(crate) fn poll_timestamp(&self, waiter: &kio::Waiter) -> Poll<()> {
match self.state.poll(waiter, |state| {
if state.timestamp.is_some() || state.fin.is_some() || state.abort.is_some() {
Poll::Ready(())
} else {
Poll::Pending
}
}) {
Poll::Ready(_) => Poll::Ready(()),
Poll::Pending => Poll::Pending,
}
}
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())
}
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())
}
fn abort_reason(&self) -> Error {
self.state.read().abort.clone().unwrap_or(Error::Dropped)
}
}
impl Clone for Producer {
fn clone(&self) -> Self {
Self {
info: self.info,
state: self.state.clone(),
track: self.track.clone(),
cache: self.cache.clone(),
stats: self.stats.clone(),
alive: self.alive.clone(),
}
}
}
struct Prefetch {
frames: [MaybeUninit<Frame>; Self::CAP],
pos: usize,
len: usize,
}
impl Prefetch {
const CAP: usize = 8;
fn pop(&mut self) -> Option<Frame> {
if self.pos == self.len {
return None;
}
let frame = unsafe { self.frames[self.pos].assume_init_read() };
self.pos += 1;
Some(frame)
}
fn fill(&mut self, frames: impl Iterator<Item = Frame>) {
debug_assert_eq!(self.pos, self.len, "fill on a non-empty batch would leak frames");
self.pos = 0;
self.len = 0;
for frame in frames.take(Self::CAP) {
self.frames[self.len].write(frame);
self.len += 1;
}
}
fn buffered(&self) -> (u64, u64) {
let mut bytes = 0u64;
for slot in &self.frames[self.pos..self.len] {
bytes += unsafe { slot.assume_init_ref() }.payload.len() as u64;
}
((self.len - self.pos) as u64, bytes)
}
}
impl Default for Prefetch {
fn default() -> Self {
Self {
frames: [const { MaybeUninit::uninit() }; Self::CAP],
pos: 0,
len: 0,
}
}
}
impl Drop for Prefetch {
fn drop(&mut self) {
for slot in &mut self.frames[self.pos..self.len] {
unsafe { slot.assume_init_drop() };
}
}
}
pub struct Consumer {
inner: ConsumerKind,
info: Info,
track: track::Info,
stats: stats::Meter,
stale_stats: stats::Meter,
expiry: Option<Arc<dyn Expiry>>,
expired: bool,
ended: bool,
stale_counted: Arc<AtomicBool>,
}
pub(crate) trait Expiry: Send + Sync {
fn is_expired(&self, waiter: &kio::Waiter) -> bool;
}
#[expect(clippy::large_enum_variant)]
enum ConsumerKind {
Plain(Plain),
Spliced(Box<super::resume::Group>),
}
struct Plain {
state: kio::Consumer<GroupState>,
index: usize,
end: Option<usize>,
prefetch: Prefetch,
cache: Arc<cache::Track>,
access: Arc<cache::Access>,
refreshed: u64,
}
impl Clone for Plain {
fn clone(&self) -> Self {
Self {
state: self.state.clone(),
index: self.index,
end: self.end,
prefetch: Prefetch::default(),
cache: self.cache.clone(),
access: self.access.clone(),
refreshed: self.refreshed,
}
}
}
impl Clone for Consumer {
fn clone(&self) -> Self {
Self {
inner: match &self.inner {
ConsumerKind::Plain(plain) => ConsumerKind::Plain(plain.clone()),
ConsumerKind::Spliced(spliced) => ConsumerKind::Spliced(Box::new((**spliced).clone())),
},
info: self.info,
track: self.track.clone(),
stats: self.stats.clone(),
stale_stats: self.stale_stats.clone(),
expiry: self.expiry.clone(),
expired: self.expired,
ended: self.ended,
stale_counted: self.stale_counted.clone(),
}
}
}
impl std::ops::Deref for Consumer {
type Target = Info;
fn deref(&self) -> &Self::Target {
&self.info
}
}
impl Consumer {
pub(crate) fn content(&self) -> stats::Content {
match &self.inner {
ConsumerKind::Plain(plain) => plain.state.read().content(),
ConsumerKind::Spliced(_) => stats::Content {
groups: 1,
..Default::default()
},
}
}
fn unread_content(&self) -> stats::Content {
match &self.inner {
ConsumerKind::Plain(plain) => plain.unread_content(),
ConsumerKind::Spliced(_) => stats::Content::default(),
}
}
pub(crate) fn into_spliced(self, mut spliced: super::resume::Group) -> Self {
spliced.set_stale_meter(self.stale_stats.clone());
Self {
inner: ConsumerKind::Spliced(Box::new(spliced)),
info: self.info,
track: self.track,
stats: self.stats,
stale_stats: self.stale_stats,
expiry: None,
expired: false,
ended: false,
stale_counted: self.stale_counted,
}
}
pub(crate) fn with_meter(mut self, meter: stats::Meter) -> Self {
meter.group();
self.stats = meter.clone();
self.set_stale_meter(meter);
self
}
pub(crate) fn set_stale_meter(&mut self, meter: stats::Meter) {
if let ConsumerKind::Spliced(spliced) = &mut self.inner {
spliced.set_stale_meter(meter.clone());
}
self.stale_stats = meter;
}
pub(crate) fn with_expiry(mut self, expiry: Arc<dyn Expiry>) -> Self {
self.expiry = Some(expiry);
self
}
pub(crate) fn poll_expired(&mut self, waiter: &kio::Waiter) -> bool {
self.poll_expired_while_pending(waiter, false)
}
fn poll_expired_if_blocked(&mut self, waiter: &kio::Waiter) -> Option<bool> {
if self.ended {
return Some(false);
}
if !self.poll_expired(waiter) {
return None;
}
let truncates = self.expired_truncates();
self.ended = !truncates;
Some(truncates)
}
fn expired_truncates(&self) -> bool {
let unread = self.unread_content();
unread.frames > 0 || unread.bytes > 0
}
pub(crate) fn poll_expired_while_pending(&mut self, waiter: &kio::Waiter, pending: bool) -> bool {
if !self.expired
&& (pending || self.expiry_pending())
&& self.expiry.as_ref().is_some_and(|expiry| expiry.is_expired(waiter))
{
self.expired = true;
if !self.stale_counted.swap(true, Ordering::Relaxed) {
self.stale_stats.stale(self.unread_content());
}
}
self.expired
}
fn expiry_pending(&self) -> bool {
match &self.inner {
ConsumerKind::Plain(plain) => plain.expiry_pending(),
ConsumerKind::Spliced(_) => false,
}
}
pub(crate) fn latency_expired(&self) -> bool {
self.expired
}
pub(crate) fn is_aborted(&self) -> bool {
match &self.inner {
ConsumerKind::Plain(plain) => plain.state.read().abort.is_some(),
ConsumerKind::Spliced(_) => false,
}
}
pub fn keep_alive(&self) {
if let ConsumerKind::Plain(plain) = &self.inner {
plain.state.read().charge.refresh();
}
}
pub(crate) fn cache_refresh(&self) {
self.keep_alive();
}
pub(crate) fn poll_closed(&self, waiter: &kio::Waiter) -> Poll<()> {
match &self.inner {
ConsumerKind::Plain(plain) => plain.state.poll_closed(waiter),
ConsumerKind::Spliced(_) => Poll::Ready(()),
}
}
pub fn timescale(&self) -> Timescale {
self.track.timescale
}
pub fn index(&self) -> u64 {
match &self.inner {
ConsumerKind::Plain(plain) => plain.index as u64,
ConsumerKind::Spliced(spliced) => spliced.index(),
}
}
pub fn set_frames(&mut self, frames: impl RangeBounds<u64>) {
let (start, end) = super::subscription::sequence_bounds(frames);
self.start_at(start);
self.end_at(end.map_or(Bound::Unbounded, Bound::Excluded));
}
pub(crate) fn start_at(&mut self, index: u64) {
match &mut self.inner {
ConsumerKind::Plain(plain) => plain.start_at(index),
ConsumerKind::Spliced(spliced) => spliced.start_at(index),
}
}
pub fn skip_to(&mut self, index: u64) {
match &mut self.inner {
ConsumerKind::Plain(plain) => plain.skip_to(index),
ConsumerKind::Spliced(spliced) => spliced.start_at(index),
}
}
pub(crate) fn end_at(&mut self, end: impl Into<Cap>) {
let end = end.into().exclusive();
match &mut self.inner {
ConsumerKind::Plain(plain) => {
plain.end = end.map(|end| usize::try_from(end).unwrap_or(usize::MAX));
}
ConsumerKind::Spliced(spliced) => spliced.end_at(end),
}
}
pub fn frame_count(&self) -> usize {
match &self.inner {
ConsumerKind::Plain(plain) => {
let state = plain.state.read();
state.fin.unwrap_or(state.next_index)
}
ConsumerKind::Spliced(spliced) => spliced.frame_count(),
}
}
pub async fn next_frame(&mut self) -> Result<Option<frame::Consumer>> {
kio::wait(|waiter| self.poll_next_frame(waiter)).await
}
pub fn poll_next_frame(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<frame::Consumer>>> {
if self.ended {
return Poll::Ready(Ok(None));
}
if self.expired {
return Poll::Ready(Err(Error::Old));
}
let stats = self.stats.clone();
let expiry = self
.expiry
.as_ref()
.map(|policy| frame::Expiry::new(policy.clone(), self.stale_stats.clone(), self.stale_counted.clone()));
let res = match &mut self.inner {
ConsumerKind::Plain(plain) => plain.poll_next_frame(waiter, &stats, expiry),
ConsumerKind::Spliced(spliced) => {
let res = ready!(spliced.poll_next_frame(waiter))?;
if res.is_some() {
stats.frames(1);
}
Poll::Ready(Ok(res.map(|frame| frame.with_meter(stats))))
}
};
match res.is_pending().then(|| self.poll_expired_if_blocked(waiter)).flatten() {
Some(true) => Poll::Ready(Err(Error::Old)),
Some(false) => Poll::Ready(Ok(None)),
None => res,
}
}
pub fn poll_read_frame(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<frame::Frame>>> {
if self.ended {
return Poll::Ready(Ok(None));
}
if self.expired {
return Poll::Ready(Err(Error::Old));
}
let stats = self.stats.clone();
let res = match &mut self.inner {
ConsumerKind::Plain(plain) => plain.poll_read_frame(waiter, &stats),
ConsumerKind::Spliced(spliced) => {
let res = ready!(spliced.poll_read_frame(waiter))?;
if let Some(frame) = &res {
stats.frames(1);
stats.bytes(frame.payload.len() as u64);
}
Poll::Ready(Ok(res))
}
};
match res.is_pending().then(|| self.poll_expired_if_blocked(waiter)).flatten() {
Some(true) => Poll::Ready(Err(Error::Old)),
Some(false) => Poll::Ready(Ok(None)),
None => res,
}
}
pub async fn read_frame(&mut self) -> Result<Option<frame::Frame>> {
if !self.expired
&& let ConsumerKind::Plain(plain) = &mut self.inner
{
if !plain.capped()
&& let Some(frame) = plain.prefetch.pop()
{
plain.refresh_if_stale();
plain.index += 1;
return Ok(Some(frame));
}
}
kio::wait(|waiter| self.poll_read_frame(waiter)).await
}
pub fn poll_read_frames<const N: usize>(
&mut self,
waiter: &kio::Waiter,
out: &mut frame::Buffer<N>,
) -> Poll<Result<usize>> {
out.clear();
if out.capacity() == 0 {
return Poll::Ready(Ok(0));
}
while !out.is_full() {
match self.poll_read_frame(waiter) {
Poll::Ready(Ok(Some(frame))) => out.push(frame).expect("buffer capacity checked"),
Poll::Ready(Ok(None)) => break,
Poll::Ready(Err(err)) => {
if out.is_empty() {
return Poll::Ready(Err(err));
}
break;
}
Poll::Pending if !out.is_empty() => break,
Poll::Pending => return Poll::Pending,
}
}
Poll::Ready(Ok(out.len()))
}
pub async fn read_frames<'a, const N: usize>(
&mut self,
out: &'a mut frame::Buffer<N>,
) -> Result<&'a mut [frame::Frame]> {
kio::wait(|waiter| self.poll_read_frames(waiter, out)).await?;
Ok(out.filled_mut())
}
pub fn poll_finished(&mut self, waiter: &kio::Waiter) -> Poll<Result<u64>> {
if self.ended {
return Poll::Ready(Ok(self.index()));
}
if self.expired {
return Poll::Ready(Err(Error::Old));
}
let res = match &mut self.inner {
ConsumerKind::Plain(plain) => {
let index = plain.index;
plain
.poll(waiter, |state| state.poll_end(index))
.map(|res| res.map(|()| index as u64))
}
ConsumerKind::Spliced(spliced) => spliced.poll_finished(waiter),
};
match res.is_pending().then(|| self.poll_expired_if_blocked(waiter)).flatten() {
Some(true) => Poll::Ready(Err(Error::Old)),
Some(false) => Poll::Ready(Ok(self.index())),
None => res,
}
}
pub async fn finished(&mut self) -> Result<u64> {
kio::wait(|waiter| self.poll_finished(waiter)).await
}
}
impl Plain {
fn expiry_pending(&self) -> bool {
if self.capped() {
return false;
}
let state = self.state.read();
state.abort.is_none() && state.fin.is_none_or(|fin| self.index < fin)
}
fn unread_content(&self) -> stats::Content {
let prefetched = self.prefetch.buffered().0 as usize;
let start = self.index.saturating_add(prefetched);
let end = self.end.unwrap_or(usize::MAX);
self.state.read().content_range(start, end)
}
fn refresh_if_stale(&mut self) {
self.access.touch();
let tick = self.cache.pool().now();
if tick != self.refreshed {
self.state.read().charge.refresh();
self.refreshed = tick;
}
}
fn poll<F, R>(&self, waiter: &kio::Waiter, f: F) -> Poll<Result<R>>
where
F: Fn(&kio::Ref<'_, GroupState>) -> 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 capped(&self) -> bool {
self.end.is_some_and(|end| self.index >= end)
}
fn start_at(&mut self, index: u64) {
let index = usize::try_from(index).unwrap_or(usize::MAX);
let index = index.max(self.state.read().offset);
if index <= self.index {
return;
}
self.index = index;
self.prefetch = Prefetch::default();
}
fn skip_to(&mut self, index: u64) {
let index = usize::try_from(index).unwrap_or(usize::MAX);
if index <= self.index {
return;
}
self.index = index;
self.prefetch = Prefetch::default();
}
fn poll_next_frame(
&mut self,
waiter: &kio::Waiter,
stats: &stats::Meter,
expiry: Option<frame::Expiry>,
) -> Poll<Result<Option<frame::Consumer>>> {
if self.capped() {
return Poll::Ready(Ok(None));
}
let end = self.end.unwrap_or(usize::MAX);
if let Some(frame) = self.prefetch.pop() {
self.refresh_if_stale();
self.index += 1;
let tail = self.index.saturating_add(self.prefetch.buffered().0 as usize)..end;
let info = frame::Info {
size: frame.payload.len() as u64,
timestamp: frame.timestamp,
};
let source = frame::Source::Complete(frame.payload);
let frame = frame::Consumer::new(self.state.clone(), info, source);
return Poll::Ready(Ok(Some(match expiry {
Some(expiry) => frame.with_expiry(expiry.for_frame(tail, false)),
None => frame,
})));
}
let index = self.index;
let Some((info, source)) = ready!(self.poll(waiter, |state| state.poll_frame_source(index))?) else {
return Poll::Ready(Ok(None));
};
self.index += 1;
stats.frames(1);
let frame = frame::Consumer::new(self.state.clone(), info, source).with_meter(stats.clone());
Poll::Ready(Ok(Some(match expiry {
Some(expiry) => frame.with_expiry(expiry.for_frame(self.index..end, true)),
None => frame,
})))
}
fn poll_read_frame(&mut self, waiter: &kio::Waiter, stats: &stats::Meter) -> Poll<Result<Option<frame::Frame>>> {
if self.capped() {
return Poll::Ready(Ok(None));
}
if let Some(frame) = self.prefetch.pop() {
self.refresh_if_stale();
self.index += 1;
return Poll::Ready(Ok(Some(frame)));
}
let index = self.index;
let budget = self.end.map_or(usize::MAX, |end| end.saturating_sub(index));
let prefetch = &mut self.prefetch;
let res = self.state.poll(waiter, |state| {
if index < state.offset {
return Poll::Ready(Err(Error::Lagged));
}
let local = (index - state.offset).min(state.frames.len());
prefetch.fill(state.frames.range(local..).take(budget).cloned());
if prefetch.len > 0 {
state.charge.refresh();
return Poll::Ready(Ok(()));
}
state.poll_terminal(index)
});
match ready!(res) {
Ok(Ok(())) => {}
Ok(Err(err)) => return Poll::Ready(Err(err)),
Err(state) => return Poll::Ready(Err(state.abort.clone().unwrap_or(Error::Dropped))),
}
self.refreshed = self.cache.pool().now();
let (frames, bytes) = self.prefetch.buffered();
stats.frames(frames);
stats.bytes(bytes);
Poll::Ready(Ok(self.prefetch.pop().inspect(|_| {
self.index += 1;
})))
}
}
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct Fetch {
pub priority: u8,
pub frame_start: u64,
}
impl Fetch {
pub fn with_priority(mut self, priority: u8) -> Self {
self.priority = priority;
self
}
pub fn with_frame_start(mut self, frame_start: u64) -> Self {
self.frame_start = frame_start;
self
}
}
pub struct Request {
pub(crate) state: kio::Producer<track::TrackState>,
pub(crate) fetch: kio::Shared<track::FetchState>,
pub(crate) sequence: u64,
pub(crate) priority: u8,
pub(crate) frame_start: u64,
pub(crate) result: kio::Producer<track::FetchOutcome>,
pub(crate) done: bool,
}
#[cfg(test)]
mod test {
use super::*;
use crate::model::test_tracing::count_drop_warnings;
use bytes::Bytes;
use futures::FutureExt;
#[test]
fn one_frame_fits_the_charged_slots() {
let mut frames: VecDeque<Frame> = VecDeque::new();
frames.push_back(Frame {
timestamp: Timestamp::ZERO,
payload: Bytes::new(),
});
let capacity = frames.capacity();
assert!(
capacity <= FRAME_SLOTS,
"a one-frame deque now allocates {capacity} slots"
);
}
#[test]
fn basic_frame_reading() {
let mut producer = Info { sequence: 0 }.produce();
producer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"frame0"))
.unwrap();
producer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"frame1"))
.unwrap();
producer.finish().unwrap();
let mut consumer = producer.consume();
let f0 = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(f0.size, 6);
let f1 = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(f1.size, 6);
let end = consumer.next_frame().now_or_never().unwrap().unwrap();
assert!(end.is_none());
}
#[test]
fn read_frame_all_at_once() {
let mut producer = Info { sequence: 0 }.produce();
producer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"hello"))
.unwrap();
producer.finish().unwrap();
let mut consumer = producer.consume();
let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(frame.payload, Bytes::from_static(b"hello"));
}
#[test]
fn read_frame_preserves_timestamp() {
let mut producer = Info { sequence: 0 }.produce();
let timestamp = Timestamp::from_micros(20_000).unwrap();
producer.write_frame(timestamp, Bytes::from_static(b"hello")).unwrap();
producer.finish().unwrap();
let mut consumer = producer.consume();
let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(frame.timestamp.as_micros(), 20_000);
assert_eq!(frame.payload, Bytes::from_static(b"hello"));
}
#[test]
fn chunked_frame_reads_whole() {
let mut producer = Info { sequence: 0 }.produce();
{
let mut frame = producer
.create_frame(frame::Info {
size: 10,
timestamp: Timestamp::ZERO,
})
.unwrap();
frame.write(Bytes::from_static(b"hello")).unwrap();
frame.write(Bytes::from_static(b"world")).unwrap();
frame.finish().unwrap();
}
producer.finish().unwrap();
let mut consumer = producer.consume();
let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(frame.payload, Bytes::from_static(b"helloworld"));
}
#[test]
fn chunked_frame_streams_partial() {
let mut producer = Info { sequence: 0 }.produce();
let mut consumer = producer.consume();
let mut frame = producer
.create_frame(frame::Info {
size: 6,
timestamp: Timestamp::ZERO,
})
.unwrap();
frame.write(Bytes::from_static(b"foo")).unwrap();
let mut f = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
let c1 = f.read_chunk().now_or_never().unwrap().unwrap();
assert_eq!(c1, Some(Bytes::from_static(b"foo")));
assert!(f.read_chunk().now_or_never().is_none());
frame.write(Bytes::from_static(b"bar")).unwrap();
frame.finish().unwrap();
let c2 = f.read_chunk().now_or_never().unwrap().unwrap();
assert_eq!(c2, Some(Bytes::from_static(b"bar")));
let c3 = f.read_chunk().now_or_never().unwrap().unwrap();
assert_eq!(c3, None);
}
#[test]
fn group_finish_returns_none() {
let producer = Info { sequence: 0 }.produce();
producer.finish().unwrap();
let mut consumer = producer.consume();
let end = consumer.next_frame().now_or_never().unwrap().unwrap();
assert!(end.is_none());
}
#[test]
fn abort_propagates() {
let producer = Info { sequence: 0 }.produce();
let mut consumer = producer.consume();
producer.abort(crate::Error::Cancel).unwrap();
let result = consumer.next_frame().now_or_never().unwrap();
assert!(matches!(result, Err(crate::Error::Cancel)));
}
#[test]
fn abort_clears_cached_frames() {
let mut producer = Info { sequence: 0 }.produce();
producer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
.unwrap();
let _consumer = producer.consume();
assert_eq!(producer.state.read().frames.len(), 1);
producer.clone().abort(crate::Error::Cancel).unwrap();
let state = producer.state.read();
assert!(state.frames.is_empty(), "cached frames should be dropped on abort");
assert_eq!(state.cache, 0);
}
#[test]
fn drop_unfinished_clears_cached_frames() {
let producer = Info { sequence: 0 }.produce();
let mut writer = producer.clone();
writer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
.unwrap();
let mut consumer = producer.consume();
assert_eq!(producer.state.read().frames.len(), 1);
drop(writer);
drop(producer);
let result = consumer.next_frame().now_or_never().unwrap();
assert!(matches!(result, Err(crate::Error::Dropped)));
}
#[test]
fn drop_after_abort_does_not_warn() {
let warns = count_drop_warnings("group::Producer dropped without finish", || {
let producer = Info { sequence: 0 }.produce();
let keep = producer.clone();
let mut writer = producer.clone();
writer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
.unwrap();
let _consumer = producer.consume();
writer.abort(crate::Error::Cancel).unwrap();
drop(keep);
});
assert_eq!(warns, 0, "abort-then-drop must not emit unfinished-producer WARN");
}
#[test]
fn drop_unfinished_warns() {
let warns = count_drop_warnings("group::Producer dropped without finish", || {
let producer = Info { sequence: 0 }.produce();
let mut writer = producer.clone();
writer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
.unwrap();
let _consumer = producer.consume();
drop(writer);
drop(producer);
});
assert!(warns >= 1, "unfinished drop must emit unfinished-producer WARN");
}
#[test]
fn drop_finished_keeps_cached_frames() {
let mut producer = Info { sequence: 0 }.produce();
producer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
.unwrap();
producer.finish().unwrap();
let mut consumer = producer.consume();
drop(producer);
let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(frame.payload, Bytes::from_static(b"data"));
}
#[tokio::test]
async fn pending_then_ready() {
let mut producer = Info { sequence: 0 }.produce();
let mut consumer = producer.consume();
assert!(consumer.next_frame().now_or_never().is_none());
producer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"data"))
.unwrap();
producer.finish().unwrap();
let frame = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(frame.size, 4);
}
#[test]
fn overflow_aborts_the_group() {
let mut producer = Info { sequence: 0 }.produce();
let mut consumer = producer.consume();
let big = Bytes::from(vec![0u8; MAX_CACHE_BYTES as usize]);
producer.write_frame(Timestamp::ZERO, big.clone()).unwrap();
assert!(matches!(
producer.write_frame(Timestamp::ZERO, big),
Err(Error::GroupTooLarge)
));
{
let state = producer.state.read();
assert!(matches!(state.abort, Some(Error::GroupTooLarge)));
assert!(state.frames.is_empty());
assert_eq!(state.offset, 0);
}
let result = consumer.next_frame().now_or_never().unwrap();
assert!(matches!(result, Err(Error::GroupTooLarge)));
}
#[test]
fn no_overflow_under_budget() {
let mut producer = Info { sequence: 0 }.produce();
for _ in 0..MAX_GROUP_FRAMES {
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"x")).unwrap();
}
producer.finish().unwrap();
let state = producer.state.read();
assert_eq!(state.offset, 0);
assert_eq!(state.frames.len(), MAX_GROUP_FRAMES);
assert!(state.abort.is_none());
}
#[test]
fn writer_sees_group_too_large_on_the_8193rd_frame() {
let mut producer = Info { sequence: 0 }.produce();
for _ in 0..MAX_GROUP_FRAMES {
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"x")).unwrap();
}
assert!(matches!(
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"x")),
Err(Error::GroupTooLarge)
));
assert!(matches!(producer.state.read().abort, Some(Error::GroupTooLarge)));
}
#[test]
fn clone_consumer_independent() {
let mut producer = Info { sequence: 0 }.produce();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"a")).unwrap();
let mut c1 = producer.consume();
let _ = c1.next_frame().now_or_never().unwrap().unwrap().unwrap();
let mut c2 = c1.clone();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"b")).unwrap();
producer.finish().unwrap();
let f = c2.next_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(f.size, 1);
let end = c2.next_frame().now_or_never().unwrap().unwrap();
assert!(end.is_none());
}
fn prefetched_consumer(pool: &cache::Pool, max_age: std::time::Duration) -> (Producer, Consumer) {
let cache = cache::Track::new(pool.clone(), kio::Weak::new());
let track = track::Info::default().with_max_age(max_age);
let mut producer = Producer::new(Info { sequence: 0 }, track, cache);
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"a")).unwrap();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"b")).unwrap();
producer.finish().unwrap();
let mut consumer = producer.consume();
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
(producer, consumer)
}
#[test]
fn prefetch_refresh_honors_pool_expiry() {
let config = cache::Config::default().with_expiry(std::time::Duration::from_secs(1));
let pool = cache::Pool::new(config);
let (producer, mut consumer) = prefetched_consumer(&pool, std::time::Duration::MAX);
let before = producer.cache_accessed();
crate::model::clock::advance(std::time::Duration::from_millis(600));
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert!(producer.cache_accessed() > before, "the pool cadence is used");
}
#[test]
fn prefetch_refresh_honors_track_max_age() {
let config = cache::Config::default().with_expiry(std::time::Duration::from_secs(30));
let pool = cache::Pool::new(config);
let (producer, mut consumer) = prefetched_consumer(&pool, std::time::Duration::from_secs(1));
let before = producer.cache_accessed();
crate::model::clock::advance(std::time::Duration::from_millis(600));
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert!(producer.cache_accessed() > before, "the track cadence remains in force");
}
#[test]
fn read_frame_crosses_prefetch_batches() {
let n = Prefetch::CAP * 3 + 5;
let mut producer = Info { sequence: 0 }.produce();
for i in 0..n {
producer
.write_frame(Timestamp::ZERO, Bytes::from(vec![i as u8; 4]))
.unwrap();
}
producer.finish().unwrap();
let mut consumer = producer.consume();
for i in 0..n {
let frame = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(frame.payload, Bytes::from(vec![i as u8; 4]));
}
assert!(consumer.read_frame().now_or_never().unwrap().unwrap().is_none());
}
#[test]
fn abort_after_finish_keeps_the_clean_end_for_a_drained_reader() {
let mut producer = Info { sequence: 0 }.produce();
producer
.write_frame(Timestamp::ZERO, Bytes::from_static(b"hello"))
.unwrap();
producer.finish().unwrap();
let mut drained = producer.consume();
let mut behind = producer.consume();
let frame = drained.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(frame.payload, Bytes::from_static(b"hello"));
producer.abort(Error::Old).unwrap();
assert!(drained.read_frame().now_or_never().unwrap().unwrap().is_none());
assert!(drained.next_frame().now_or_never().unwrap().unwrap().is_none());
assert!(matches!(behind.read_frame().now_or_never().unwrap(), Err(Error::Old)));
}
#[test]
fn finished_answers_for_the_cursor() {
let mut producer = Info { sequence: 0 }.produce();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"a")).unwrap();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"b")).unwrap();
producer.finish().unwrap();
let mut drained = producer.consume();
let mut behind = producer.consume();
while drained.read_frame().now_or_never().unwrap().unwrap().is_some() {}
behind.read_frame().now_or_never().unwrap().unwrap().unwrap();
producer.abort(Error::Old).unwrap();
assert_eq!(drained.finished().now_or_never().unwrap().unwrap(), 2);
assert!(matches!(behind.finished().now_or_never().unwrap(), Err(Error::Old)));
assert_eq!(behind.frame_count(), 2);
}
#[test]
fn finished_reports_a_group_too_large() {
let mut producer = Info { sequence: 0 }.produce();
let mut consumer = producer.consume();
let big = Bytes::from(vec![0u8; MAX_CACHE_BYTES as usize]);
producer.write_frame(Timestamp::ZERO, big.clone()).unwrap();
assert!(matches!(
producer.write_frame(Timestamp::ZERO, big),
Err(Error::GroupTooLarge)
));
assert!(matches!(
consumer.finished().now_or_never().unwrap(),
Err(Error::GroupTooLarge)
));
}
#[test]
fn interleave_read_and_next_frame() {
let mut producer = Info { sequence: 0 }.produce();
for i in 0..5u8 {
producer.write_frame(Timestamp::ZERO, Bytes::from(vec![i; 1])).unwrap();
}
producer.finish().unwrap();
let mut consumer = producer.consume();
let f0 = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
assert_eq!(f0.payload, Bytes::from(vec![0u8; 1]));
for i in 1..5u8 {
let mut f = consumer.next_frame().now_or_never().unwrap().unwrap().unwrap();
let data = f.read_all().now_or_never().unwrap().unwrap();
assert_eq!(data, Bytes::from(vec![i; 1]));
}
assert!(consumer.next_frame().now_or_never().unwrap().unwrap().is_none());
}
#[test]
fn read_frame_past_cleared_frames_does_not_panic() {
let mut producer = Info { sequence: 0 }.produce();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"a")).unwrap();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"b")).unwrap();
let mut consumer = producer.consume();
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
producer.abort(Error::Cancel).unwrap();
let result = consumer.read_frame().now_or_never().unwrap();
assert!(matches!(result, Err(Error::Cancel)), "expected Cancel, got {result:?}");
}
#[test]
fn drop_with_partial_batch() {
let mut producer = Info { sequence: 0 }.produce();
for _ in 0..Prefetch::CAP {
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"x")).unwrap();
}
producer.finish().unwrap();
let mut consumer = producer.consume();
let _ = consumer.read_frame().now_or_never().unwrap().unwrap().unwrap();
drop(consumer);
}
#[tokio::test]
async fn chunk_write_wakes_parked_reader() {
let mut producer = Info { sequence: 0 }.produce();
let mut consumer = producer.consume();
let mut frame = producer
.create_frame(frame::Info {
size: 6,
timestamp: Timestamp::ZERO,
})
.unwrap();
let mut f = consumer.next_frame().await.unwrap().unwrap();
let handle = tokio::spawn(async move { f.read_chunk().await });
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
frame.write(Bytes::from_static(b"foo")).unwrap();
let chunk = tokio::time::timeout(std::time::Duration::from_secs(2), handle)
.await
.expect("parked chunk reader was never woken by the chunk write")
.unwrap()
.unwrap();
assert_eq!(chunk, Some(Bytes::from_static(b"foo")));
}
#[test]
fn create_frame_converts_mismatched_scale() {
use crate::{Timescale, Timestamp};
let mut producer = Producer::new(
Info { sequence: 0 },
track::Info::default().with_timescale(Timescale::MICRO),
Default::default(),
);
let frame = frame::Info {
size: 3,
timestamp: Timestamp::from_millis(1).unwrap(), };
let writer = producer.create_frame(frame).unwrap();
assert_eq!(writer.timestamp.scale(), Timescale::MICRO);
assert_eq!(writer.timestamp.value(), 1000);
}
#[tokio::test]
async fn create_frame_converts_current_timestamp() {
use crate::Timescale;
let mut producer = Producer::new(
Info { sequence: 0 },
track::Info::default().with_timescale(Timescale::MICRO),
Default::default(),
);
let writer = producer
.create_frame(frame::Info {
size: 3,
timestamp: Timestamp::now(),
})
.unwrap();
assert_eq!(writer.timestamp.scale(), Timescale::MICRO);
assert!(!writer.timestamp.is_zero(), "local clock should be non-zero");
}
#[test]
fn start_at_starts_the_group_later() {
let mut producer = Info { sequence: 0 }.produce();
producer.start_at(3).unwrap();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"d")).unwrap();
producer.finish().unwrap();
assert_eq!(producer.frame_count(), 4);
let mut consumer = producer.consume();
assert_eq!(consumer.frame_count(), 4);
assert!(matches!(
consumer.finished().now_or_never().unwrap(),
Err(Error::Lagged)
));
assert!(matches!(
consumer.read_frame().now_or_never().unwrap(),
Err(Error::Lagged)
));
}
#[test]
fn start_at_clamps_up_to_the_first_frame() {
let mut producer = Info { sequence: 0 }.produce();
producer.start_at(3).unwrap();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"d")).unwrap();
producer.finish().unwrap();
let mut consumer = producer.consume();
consumer.start_at(1);
assert_eq!(consumer.index(), 3, "clamped up to the first frame that exists");
assert_eq!(
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap().payload,
Bytes::from_static(b"d")
);
}
#[test]
fn end_at_caps_and_reopens() {
let mut producer = Info { sequence: 0 }.produce();
for i in 0..4u8 {
producer.write_frame(Timestamp::ZERO, Bytes::from(vec![i])).unwrap();
}
producer.finish().unwrap();
let mut consumer = producer.consume();
consumer.set_frames(..2);
assert_eq!(
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap().payload[0],
0
);
assert_eq!(
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap().payload[0],
1
);
assert!(
consumer.read_frame().now_or_never().unwrap().unwrap().is_none(),
"capped reads end cleanly"
);
consumer.set_frames(..);
assert_eq!(
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap().payload[0],
2
);
}
#[test]
fn frame_ranges_preserve_progress_and_make_inclusion_explicit() {
let mut producer = Info { sequence: 0 }.produce();
for i in 0..4u8 {
producer.write_frame(Timestamp::ZERO, Bytes::from(vec![i])).unwrap();
}
producer.finish().unwrap();
let mut consumer = producer.consume();
consumer.set_frames(1..=1);
assert_eq!(
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap().payload[0],
1
);
assert!(consumer.read_frame().now_or_never().unwrap().unwrap().is_none());
consumer.set_frames(..3);
assert_eq!(
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap().payload[0],
2
);
assert!(consumer.read_frame().now_or_never().unwrap().unwrap().is_none());
consumer.set_frames(0..=3);
assert_eq!(
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap().payload[0],
3
);
}
#[test]
fn end_at_zero_is_empty() {
let mut producer = Info { sequence: 0 }.produce();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"x")).unwrap();
producer.finish().unwrap();
let mut consumer = producer.consume();
consumer.set_frames(..0);
assert!(
consumer.read_frame().now_or_never().unwrap().unwrap().is_none(),
"empty cap delivers nothing"
);
consumer.set_frames(..1);
assert_eq!(
consumer.read_frame().now_or_never().unwrap().unwrap().unwrap().payload,
Bytes::from_static(b"x")
);
}
#[test]
fn start_at_rejected_after_a_frame() {
let mut producer = Info { sequence: 0 }.produce();
producer.start_at(2).unwrap();
producer.start_at(3).unwrap();
producer.write_frame(Timestamp::ZERO, Bytes::from_static(b"a")).unwrap();
assert!(matches!(producer.start_at(4), Err(Error::Closed)));
assert_eq!(producer.frame_count(), 4, "the frame landed at index 3");
let mut producer = Info { sequence: 1 }.produce();
producer.finish().unwrap();
assert!(matches!(producer.start_at(1), Err(Error::Closed)));
}
#[test]
fn start_at_rejects_the_largest_index() {
let mut producer = Info { sequence: 0 }.produce();
assert!(matches!(
producer.start_at(usize::MAX as u64),
Err(Error::BoundsExceeded(_))
));
}
#[test]
fn create_frame_rejects_oversized() {
let mut producer = Info { sequence: 0 }.produce();
let result = producer.create_frame(frame::Info {
size: MAX_CACHE_BYTES + 1,
timestamp: Timestamp::ZERO,
});
assert!(matches!(result, Err(Error::FrameTooLarge)));
}
}