1use std::ops::Range;
9use std::sync::Arc;
10use std::sync::OnceLock;
11use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
12use std::task::{Poll, ready};
13
14use arrayvec::ArrayVec;
15use bytes::Bytes;
16
17use crate::group::{self, GroupState};
18use crate::{Error, IntoBytes, Result, Timestamp, stats};
19
20#[derive(Clone, Copy, Debug)]
25pub struct Info {
26 pub size: u64,
28 pub timestamp: Timestamp,
34}
35
36#[derive(Clone, Debug)]
41pub struct Frame {
42 pub timestamp: Timestamp,
44 pub payload: Bytes,
46}
47
48#[derive(Debug, Default)]
70pub struct Buffer<const N: usize = 8>(ArrayVec<Frame, N>);
71
72impl<const N: usize> Buffer<N> {
73 pub fn new() -> Self {
75 Self(ArrayVec::new())
76 }
77
78 pub const fn capacity(&self) -> usize {
80 N
81 }
82
83 pub fn len(&self) -> usize {
85 self.0.len()
86 }
87
88 pub fn is_empty(&self) -> bool {
90 self.0.is_empty()
91 }
92
93 pub fn is_full(&self) -> bool {
95 self.0.is_full()
96 }
97
98 pub fn filled(&self) -> &[Frame] {
100 &self.0
101 }
102
103 pub fn filled_mut(&mut self) -> &mut [Frame] {
105 &mut self.0
106 }
107
108 pub fn push(&mut self, frame: Frame) -> std::result::Result<(), Frame> {
113 self.0.try_push(frame).map_err(|err| err.element())
114 }
115
116 pub fn drain(&mut self) -> impl ExactSizeIterator<Item = Frame> + '_ {
121 self.0.drain(..)
122 }
123
124 pub fn clear(&mut self) {
126 self.0.clear();
127 }
128}
129
130#[derive(Clone)]
139pub(crate) struct FrameBuf(Arc<FrameBufInner>);
140
141struct FrameBufInner {
142 capacity: usize,
143 written: AtomicUsize,
144 storage: OnceLock<FrameStorage>,
145}
146
147enum FrameStorage {
148 Shared(Bytes),
149 Mutable(MutableFrameBuf),
150}
151
152struct MutableFrameBuf {
153 data: *mut u8,
155 capacity: usize,
156}
157
158unsafe impl Send for MutableFrameBuf {}
161unsafe impl Sync for MutableFrameBuf {}
162
163impl Drop for MutableFrameBuf {
164 fn drop(&mut self) {
165 unsafe {
168 let slice = std::ptr::slice_from_raw_parts_mut(self.data, self.capacity);
169 drop(Box::from_raw(slice));
170 }
171 }
172}
173
174impl MutableFrameBuf {
175 fn new(size: usize) -> Self {
176 let boxed: Box<[u8]> = vec![0u8; size].into_boxed_slice();
177 let capacity = boxed.len();
178 let data = Box::into_raw(boxed) as *mut u8;
179 Self { data, capacity }
180 }
181}
182
183impl FrameBuf {
184 pub(crate) fn new(size: usize) -> Self {
189 Self(Arc::new(FrameBufInner {
190 capacity: size,
191 written: AtomicUsize::new(0),
192 storage: OnceLock::new(),
193 }))
194 }
195
196 pub(crate) fn capacity(&self) -> usize {
197 self.0.capacity
198 }
199
200 pub(crate) fn written(&self, ord: Ordering) -> usize {
201 self.0.written.load(ord)
202 }
203
204 fn try_set_bytes(&self, bytes: Bytes) -> std::result::Result<(), Bytes> {
205 if bytes.len() != self.capacity() || self.written(Ordering::Acquire) != 0 {
206 return Err(bytes);
207 }
208 self.0
209 .storage
210 .set(FrameStorage::Shared(bytes))
211 .map_err(|storage| match storage {
212 FrameStorage::Shared(bytes) => bytes,
213 FrameStorage::Mutable(_) => unreachable!("try_set_bytes only installs shared storage"),
214 })
215 }
216
217 fn mutable(&self) -> Option<&MutableFrameBuf> {
221 match self
222 .0
223 .storage
224 .get_or_init(|| FrameStorage::Mutable(MutableFrameBuf::new(self.capacity())))
225 {
226 FrameStorage::Shared(_) => None,
227 FrameStorage::Mutable(buf) => Some(buf),
228 }
229 }
230
231 unsafe fn store_written(&self, new_written: usize) {
233 self.0.written.store(new_written, Ordering::Release);
235 }
236
237 fn append(&self, src: &[u8]) {
243 if src.is_empty() {
244 return;
245 }
246 let prev = self.written(Ordering::Relaxed);
247 let Some(buf) = self.mutable() else {
248 return;
251 };
252 unsafe {
255 std::ptr::copy_nonoverlapping(src.as_ptr(), buf.data.add(prev), src.len());
256 self.store_written(prev + src.len());
257 }
258 }
259
260 fn freeze(&self, size: usize) -> Bytes {
265 match self.0.storage.get() {
266 Some(FrameStorage::Shared(bytes)) => bytes.clone(),
267 _ => self.slice(0, size),
268 }
269 }
270
271 fn slice(&self, start: usize, end: usize) -> Bytes {
273 Bytes::from_owner(self.clone()).slice(start..end)
274 }
275}
276
277impl AsRef<[u8]> for FrameBuf {
278 fn as_ref(&self) -> &[u8] {
279 let written = self.0.written.load(Ordering::Acquire);
282 match self.0.storage.get() {
283 Some(FrameStorage::Shared(bytes)) => &bytes[..written],
284 Some(FrameStorage::Mutable(buf)) => {
285 unsafe { std::slice::from_raw_parts(buf.data, written) }
289 }
290 None => &[],
291 }
292 }
293}
294
295struct Raw<G: std::borrow::BorrowMut<group::Producer>> {
298 group: G,
299 buf: FrameBuf,
300 info: Info,
301 done: bool,
303 stats: stats::Meter,
306}
307
308impl<G: std::borrow::BorrowMut<group::Producer>> Raw<G> {
309 fn remaining(&self) -> usize {
310 self.buf.capacity() - self.buf.written(Ordering::Acquire)
311 }
312
313 fn write<B: IntoBytes>(&mut self, chunk: B) -> Result<()> {
314 let len = chunk.as_ref().len();
315 if len > self.remaining() {
316 return Err(Error::WrongSize);
317 }
318 self.stats.bytes(len as u64);
320 if len == self.buf.capacity() && self.buf.written(Ordering::Acquire) == 0 {
322 match self.buf.try_set_bytes(chunk.into_bytes()) {
323 Ok(()) => {
324 let cap = self.buf.capacity();
325 unsafe { self.buf.store_written(cap) };
328 }
329 Err(chunk) => self.buf.append(&chunk),
330 }
331 } else {
332 self.buf.append(chunk.as_ref());
333 }
334 Ok(())
335 }
336
337 fn finish(&mut self) -> Result<()> {
338 if self.buf.written(Ordering::Acquire) != self.buf.capacity() {
339 return Err(Error::WrongSize);
340 }
341 let payload = self.buf.freeze(self.buf.capacity());
342 self.group.borrow_mut().frame_commit(Frame {
343 timestamp: self.info.timestamp,
344 payload,
345 })?;
346 self.done = true;
347 Ok(())
348 }
349
350 fn abort(&mut self, err: Error) -> Result<()> {
351 self.group.borrow_mut().frame_abort(err);
352 self.done = true;
353 Ok(())
354 }
355}
356
357impl<G: std::borrow::BorrowMut<group::Producer>> Drop for Raw<G> {
358 fn drop(&mut self) {
359 if !self.done {
360 tracing::warn!(
363 group = self.group.borrow_mut().info().sequence,
364 "frame::Producer dropped before writing all bytes"
365 );
366 self.group.borrow_mut().frame_abort(Error::Dropped);
367 }
368 }
369}
370
371pub struct Producer<'a>(Raw<&'a mut group::Producer>);
382
383impl std::ops::Deref for Producer<'_> {
384 type Target = Info;
385
386 fn deref(&self) -> &Self::Target {
387 &self.0.info
388 }
389}
390
391impl<'a> Producer<'a> {
392 pub(crate) fn new(group: &'a mut group::Producer, buf: FrameBuf, info: Info) -> Self {
393 Self(Raw {
394 group,
395 buf,
396 info,
397 done: false,
398 stats: stats::Meter::default(),
399 })
400 }
401
402 pub(crate) fn with_meter(mut self, meter: stats::Meter) -> Self {
404 self.0.stats = meter;
405 self
406 }
407
408 pub fn group(&self) -> group::Info {
410 self.0.group.info()
411 }
412
413 pub fn remaining(&self) -> usize {
415 self.0.remaining()
416 }
417
418 pub fn write<B: IntoBytes>(&mut self, chunk: B) -> Result<()> {
422 self.0.write(chunk)?;
423 self.0.group.frame_notify();
424 Ok(())
425 }
426
427 pub fn finish(mut self) -> Result<()> {
431 self.0.finish()
432 }
433
434 pub fn abort(mut self, err: Error) -> Result<()> {
436 self.0.abort(err)
437 }
438}
439
440pub(crate) struct ProducerOwned(Raw<group::Producer>);
447
448impl std::ops::Deref for ProducerOwned {
449 type Target = Info;
450
451 fn deref(&self) -> &Self::Target {
452 &self.0.info
453 }
454}
455
456impl ProducerOwned {
457 pub(crate) fn new(group: group::Producer, buf: FrameBuf, info: Info) -> Self {
458 Self(Raw {
459 group,
460 buf,
461 info,
462 done: false,
463 stats: stats::Meter::default(),
464 })
465 }
466
467 pub(crate) fn with_meter(mut self, meter: stats::Meter) -> Self {
469 self.0.stats = meter;
470 self
471 }
472
473 pub fn remaining(&self) -> usize {
475 self.0.remaining()
476 }
477
478 pub(crate) fn write<B: IntoBytes>(&mut self, chunk: B) -> Result<()> {
487 self.0.write(chunk)
488 }
489
490 pub(crate) fn notify(&self) {
492 self.0.group.frame_notify();
493 }
494
495 pub fn finish(mut self) -> Result<()> {
497 self.0.finish()
498 }
499
500 pub fn abort(mut self, err: Error) -> Result<()> {
502 self.0.abort(err)
503 }
504}
505
506#[derive(Clone)]
509pub(crate) enum Source {
510 Complete(Bytes),
511 Partial(FrameBuf),
512}
513
514#[derive(Clone)]
516pub(crate) struct Expiry {
517 policy: Arc<dyn group::Expiry>,
518 stale_stats: stats::Meter,
519 stale_counted: Arc<AtomicBool>,
520 tail: Range<usize>,
521 count_payload: bool,
522}
523
524impl Expiry {
525 pub(crate) fn new(
526 policy: Arc<dyn group::Expiry>,
527 stale_stats: stats::Meter,
528 stale_counted: Arc<AtomicBool>,
529 ) -> Self {
530 Self {
531 policy,
532 stale_stats,
533 stale_counted,
534 tail: 0..0,
535 count_payload: false,
536 }
537 }
538
539 pub(crate) fn for_frame(mut self, tail: Range<usize>, count_payload: bool) -> Self {
540 self.tail = tail;
541 self.count_payload = count_payload;
542 self
543 }
544}
545
546#[derive(Clone)]
552pub struct Consumer {
553 state: kio::Consumer<GroupState>,
555 info: Info,
556 source: Source,
557 read_idx: usize,
559 stats: stats::Meter,
562 expiry: Option<Expiry>,
564 expired: bool,
565}
566
567impl std::ops::Deref for Consumer {
568 type Target = Info;
569
570 fn deref(&self) -> &Self::Target {
571 &self.info
572 }
573}
574
575impl Consumer {
576 pub(crate) fn new(state: kio::Consumer<GroupState>, info: Info, source: Source) -> Self {
577 Self {
578 state,
579 info,
580 source,
581 read_idx: 0,
582 stats: stats::Meter::default(),
583 expiry: None,
584 expired: false,
585 }
586 }
587
588 pub(crate) fn with_meter(mut self, meter: stats::Meter) -> Self {
591 self.stats = meter;
592 self
593 }
594
595 pub(crate) fn with_expiry(mut self, expiry: Expiry) -> Self {
596 self.expiry = Some(expiry);
597 self
598 }
599
600 fn size(&self) -> usize {
601 match &self.source {
602 Source::Complete(bytes) => bytes.len(),
603 Source::Partial(_) => self.info.size as usize,
604 }
605 }
606
607 fn poll_expired(&mut self, waiter: &kio::Waiter) -> bool {
613 if self.expired || self.read_idx >= self.size() {
614 return self.expired;
615 }
616 let Some(expiry) = &self.expiry else {
617 return false;
618 };
619 if !expiry.policy.is_expired(waiter) {
620 return false;
621 }
622
623 self.expired = true;
624 if !expiry.stale_counted.swap(true, Ordering::Relaxed) {
625 let mut stale = self.state.read().content_range(expiry.tail.start, expiry.tail.end);
626 if expiry.count_payload {
627 stale.bytes += self.size().saturating_sub(self.read_idx) as u64;
628 }
629 expiry.stale_stats.stale(stale);
630 }
631 true
632 }
633
634 pub fn poll_read_chunk(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Bytes>>> {
638 if self.expired {
639 return Poll::Ready(Err(Error::Old));
640 }
641 let buf = match &self.source {
642 Source::Complete(bytes) => {
643 if self.read_idx >= bytes.len() {
644 return Poll::Ready(Ok(None));
645 }
646 let out = bytes.slice(self.read_idx..);
647 self.read_idx = bytes.len();
648 self.stats.bytes(out.len() as u64);
649 return Poll::Ready(Ok(Some(out)));
650 }
651 Source::Partial(buf) => buf.clone(),
652 };
653
654 let size = self.info.size as usize;
655 loop {
656 let written = buf.written(Ordering::Acquire);
657 if written > self.read_idx {
658 let out = buf.slice(self.read_idx, written);
659 self.read_idx = written;
660 self.stats.bytes(out.len() as u64);
661 return Poll::Ready(Ok(Some(out)));
662 }
663 if written >= size {
664 return Poll::Ready(Ok(None));
665 }
666 if self.poll_expired(waiter) {
669 return Poll::Ready(Err(Error::Old));
670 }
671 let read_idx = self.read_idx;
672 ready!(poll_state(&self.state, waiter, |state| {
675 if let Some(err) = &state.abort {
676 return Poll::Ready(Err(err.clone()));
677 }
678 let w = buf.written(Ordering::Acquire);
679 if w > read_idx || w >= size {
680 Poll::Ready(Ok(()))
681 } else {
682 Poll::Pending
683 }
684 })?);
685 }
686 }
687
688 pub async fn read_chunk(&mut self) -> Result<Option<Bytes>> {
690 kio::wait(|waiter| self.poll_read_chunk(waiter)).await
691 }
692
693 pub fn poll_read_all(&mut self, waiter: &kio::Waiter) -> Poll<Result<Bytes>> {
695 if self.expired {
696 return Poll::Ready(Err(Error::Old));
697 }
698 let buf = match &self.source {
699 Source::Complete(bytes) => {
700 let out = bytes.slice(self.read_idx..);
701 self.read_idx = bytes.len();
702 self.stats.bytes(out.len() as u64);
703 return Poll::Ready(Ok(out));
704 }
705 Source::Partial(buf) => buf.clone(),
706 };
707
708 let size = self.info.size as usize;
709 let read_idx = self.read_idx;
710 if buf.written(Ordering::Acquire) < size && self.poll_expired(waiter) {
712 return Poll::Ready(Err(Error::Old));
713 }
714 ready!(poll_state(&self.state, waiter, |state| {
715 if let Some(err) = &state.abort {
716 return Poll::Ready(Err(err.clone()));
717 }
718 if buf.written(Ordering::Acquire) >= size {
719 Poll::Ready(Ok(()))
720 } else {
721 Poll::Pending
722 }
723 })?);
724 let out = buf.slice(read_idx, size);
725 self.read_idx = size;
726 self.stats.bytes(out.len() as u64);
727 Poll::Ready(Ok(out))
728 }
729
730 pub async fn read_all(&mut self) -> Result<Bytes> {
732 kio::wait(|waiter| self.poll_read_all(waiter)).await
733 }
734}
735
736fn poll_state<F, R>(state: &kio::Consumer<GroupState>, waiter: &kio::Waiter, f: F) -> Poll<Result<R>>
739where
740 F: Fn(&kio::Ref<'_, GroupState>) -> Poll<Result<R>>,
741{
742 Poll::Ready(match ready!(state.poll(waiter, f)) {
743 Ok(res) => res,
744 Err(state) => Err(state.abort.clone().unwrap_or(Error::Dropped)),
745 })
746}