use std::fmt;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, PoisonError};
use arc_swap::ArcSwapOption;
use async_channel::{Receiver, Sender, unbounded};
pub struct Snapshot<T> {
slot: Arc<SnapshotSlot<T>>,
}
struct SnapshotSlot<T> {
value: ArcSwapOption<T>,
readers: AtomicUsize,
}
impl<T> Snapshot<T> {
pub fn new() -> Self {
Self {
slot: Arc::new(SnapshotSlot {
value: ArcSwapOption::new(None),
readers: AtomicUsize::new(0),
}),
}
}
pub fn publish(&self, value: T) -> usize {
self.slot.value.store(Some(Arc::new(value)));
self.slot.readers.load(Ordering::Relaxed)
}
pub fn reader(&self) -> SnapshotReader<T> {
self.slot.readers.fetch_add(1, Ordering::Relaxed);
SnapshotReader {
slot: Arc::clone(&self.slot),
}
}
pub fn has_readers(&self) -> bool {
self.slot.readers.load(Ordering::Relaxed) > 0
}
}
impl<T> Clone for Snapshot<T> {
fn clone(&self) -> Self {
Self {
slot: Arc::clone(&self.slot),
}
}
}
impl<T> Default for Snapshot<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> fmt::Debug for Snapshot<T> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("Snapshot")
.field("readers", &self.slot.readers.load(Ordering::Relaxed))
.finish_non_exhaustive()
}
}
pub struct SnapshotReader<T> {
slot: Arc<SnapshotSlot<T>>,
}
impl<T> SnapshotReader<T> {
#[must_use]
pub fn latest(&self) -> Option<Arc<T>> {
self.slot.value.load_full()
}
}
impl<T> Clone for SnapshotReader<T> {
fn clone(&self) -> Self {
self.slot.readers.fetch_add(1, Ordering::Relaxed);
Self {
slot: Arc::clone(&self.slot),
}
}
}
impl<T> Drop for SnapshotReader<T> {
fn drop(&mut self) {
self.slot.readers.fetch_sub(1, Ordering::Relaxed);
}
}
impl<T> fmt::Debug for SnapshotReader<T> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("SnapshotReader")
.finish_non_exhaustive()
}
}
pub struct Stream<T> {
inner: Arc<StreamInner<T>>,
}
struct StreamInner<T> {
senders: Mutex<Vec<Sender<T>>>,
}
impl<T: Clone> Stream<T> {
pub fn new() -> Self {
Self {
inner: Arc::new(StreamInner {
senders: Mutex::new(Vec::new()),
}),
}
}
pub fn subscribe(&self) -> StreamReader<T> {
let (sender, receiver) = unbounded();
self.inner
.senders
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push(sender);
StreamReader {
inner: Arc::new(ReaderInner {
queue: receiver,
batch: Mutex::new(Vec::new()),
}),
}
}
pub fn publish(&self, value: T) -> usize {
self.publish_all([value])
}
pub fn publish_all(&self, values: impl IntoIterator<Item = T>) -> usize {
let mut senders = self
.inner
.senders
.lock()
.unwrap_or_else(PoisonError::into_inner);
senders.retain(|sender| !sender.is_closed());
for value in values {
for sender in senders.iter() {
let _ = sender.try_send(value.clone());
}
}
senders.retain(|sender| !sender.is_closed());
senders.len()
}
pub fn has_readers(&self) -> bool {
let mut senders = self
.inner
.senders
.lock()
.unwrap_or_else(PoisonError::into_inner);
senders.retain(|sender| !sender.is_closed());
!senders.is_empty()
}
}
impl<T: Clone> Clone for Stream<T> {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
impl<T: Clone> Default for Stream<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> fmt::Debug for Stream<T> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.debug_struct("Stream").finish_non_exhaustive()
}
}
pub struct StreamReader<T> {
inner: Arc<ReaderInner<T>>,
}
struct ReaderInner<T> {
queue: Receiver<T>,
batch: Mutex<Vec<T>>,
}
impl<T> StreamReader<T> {
pub fn pump(&self) {
let mut batch = self
.inner
.batch
.lock()
.unwrap_or_else(PoisonError::into_inner);
batch.clear();
while let Ok(value) = self.inner.queue.try_recv() {
batch.push(value);
}
}
pub fn with_batch<R>(&self, f: impl FnOnce(&[T]) -> R) -> R {
let batch = self
.inner
.batch
.lock()
.unwrap_or_else(PoisonError::into_inner);
f(&batch)
}
pub fn copy_into(&self, out: &mut Vec<T>)
where
T: Clone,
{
self.with_batch(|batch| out.extend_from_slice(batch));
}
}
impl<T> Clone for StreamReader<T> {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
impl<T> fmt::Debug for StreamReader<T> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("StreamReader")
.finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_snapshot_keeps_the_newest_value() {
let snapshot = Snapshot::new();
let reader = snapshot.reader();
assert!(reader.latest().is_none());
assert_eq!(snapshot.publish(1), 1);
assert_eq!(snapshot.publish(2), 1);
assert_eq!(*reader.latest().unwrap(), 2);
}
#[test]
fn a_late_reader_finds_the_value_published_before_it() {
let snapshot = Snapshot::new();
assert_eq!(snapshot.publish(1), 0);
let reader = snapshot.reader();
assert_eq!(*reader.latest().expect("the publish was kept"), 1);
assert_eq!(snapshot.publish(2), 1);
assert_eq!(*reader.latest().expect("the newest publish"), 2);
}
#[test]
fn dropping_the_last_reader_unsubscribes() {
let snapshot = Snapshot::new();
let reader = snapshot.reader();
assert_eq!(snapshot.publish(1), 1);
drop(reader);
assert_eq!(snapshot.publish(2), 0);
assert_eq!(*snapshot.reader().latest().expect("the last publish"), 2);
}
#[test]
fn a_stream_gives_every_reader_the_same_values() {
let stream = Stream::new();
let one = stream.subscribe();
let other = stream.subscribe();
assert_eq!(stream.publish(1), 2);
assert_eq!(stream.publish(2), 2);
one.pump();
other.pump();
assert_eq!(one.with_batch(|batch| batch.to_vec()), vec![1, 2]);
assert_eq!(other.with_batch(|batch| batch.to_vec()), vec![1, 2]);
}
#[test]
fn a_pump_replaces_the_previous_batch() {
let stream = Stream::new();
let reader = stream.subscribe();
stream.publish(1);
reader.pump();
stream.publish(2);
reader.pump();
assert_eq!(reader.with_batch(|batch| batch.to_vec()), vec![2]);
}
#[test]
fn a_stream_prunes_a_reader_that_went_away() {
let stream = Stream::new();
let reader = stream.subscribe();
assert!(stream.has_readers());
drop(reader);
assert!(!stream.has_readers());
assert_eq!(stream.publish(1), 0);
}
#[test]
fn a_reader_clone_shares_the_queue() {
let stream = Stream::new();
let reader = stream.subscribe();
let clone = reader.clone();
stream.publish(1);
reader.pump();
assert_eq!(clone.with_batch(|batch| batch.to_vec()), vec![1]);
drop(reader);
assert!(stream.has_readers());
}
}