use std::{
fmt,
hint::unreachable_unchecked,
mem,
pin::Pin,
task::{ready, Context, Poll},
vec,
};
use crate::reusable_box::ReusableBoxFuture;
use futures_core::Stream;
use imbl::Vector;
use tokio::sync::broadcast::{
self,
error::{RecvError, TryRecvError},
Receiver,
};
#[cfg(feature = "tracing")]
use tracing::info;
use super::{BroadcastMessage, OneOrManyDiffs, VectorDiff};
/// A subscriber for updates of a [`Vector`].
#[derive(Debug)]
pub struct VectorSubscriber<T> {
values: Vector<T>,
rx: Receiver<BroadcastMessage<T>>,
}
impl<T: Clone + 'static> VectorSubscriber<T> {
pub(super) fn new(items: Vector<T>, rx: Receiver<BroadcastMessage<T>>) -> Self {
Self { values: items, rx }
}
/// Get the items the [`ObservableVector`][super::ObservableVector]
/// contained when this subscriber was created.
pub fn values(&self) -> Vector<T> {
self.values.clone()
}
/// Turn this `VectorSubcriber` into a stream of `VectorDiff`s.
pub fn into_stream(self) -> VectorSubscriberStream<T> {
VectorSubscriberStream::new(ReusableBoxRecvFuture::new(self.rx))
}
/// Turn this `VectorSubcriber` into a stream of `Vec<VectorDiff>`s.
pub fn into_batched_stream(self) -> VectorSubscriberBatchedStream<T> {
VectorSubscriberBatchedStream::new(ReusableBoxRecvFuture::new(self.rx))
}
/// Destructure this `VectorSubscriber` into the initial values and a stream
/// of `VectorDiff`s.
///
/// Semantically equivalent to calling `.values()` and `.into_stream()`
/// separately, but guarantees that the values are not unnecessarily cloned.
pub fn into_values_and_stream(self) -> (Vector<T>, VectorSubscriberStream<T>) {
let Self { values, rx } = self;
(values, VectorSubscriberStream::new(ReusableBoxRecvFuture::new(rx)))
}
/// Destructure this `VectorSubscriber` into the initial values and a stream
/// of `Vec<VectorDiff>`s.
///
/// Semantically equivalent to calling `.values()` and
/// `.into_batched_stream()` separately, but guarantees that the values
/// are not unnecessarily cloned.
pub fn into_values_and_batched_stream(self) -> (Vector<T>, VectorSubscriberBatchedStream<T>) {
let Self { values, rx } = self;
(values, VectorSubscriberBatchedStream::new(ReusableBoxRecvFuture::new(rx)))
}
}
/// A stream of `VectorDiff`s created from a [`VectorSubscriber`].
///
/// Use its [`Stream`] implementation to interact with it (futures-util and
/// other futures-related crates have extension traits with convenience
/// methods).
#[derive(Debug)]
pub struct VectorSubscriberStream<T> {
inner: ReusableBoxRecvFuture<T>,
state: VectorSubscriberStreamState<T>,
}
impl<T> VectorSubscriberStream<T> {
fn new(inner: ReusableBoxRecvFuture<T>) -> Self {
Self { inner, state: VectorSubscriberStreamState::Recv }
}
}
#[derive(Debug)]
enum VectorSubscriberStreamState<T> {
// Stream is waiting on a new message from the inner broadcast receiver.
Recv,
// Stream is yielding remaining items from a previous message with multiple
// diffs.
YieldBatch { iter: vec::IntoIter<VectorDiff<T>>, rx: Receiver<BroadcastMessage<T>> },
}
// Not clear why this explicit impl is needed, but it's not unsafe so it is fine
impl<T> Unpin for VectorSubscriberStreamState<T> {}
impl<T: Clone + 'static> Stream for VectorSubscriberStream<T> {
type Item = VectorDiff<T>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match &mut self.state {
VectorSubscriberStreamState::Recv => {
let (result, mut rx) = ready!(self.inner.poll(cx));
let poll = match result {
Ok(msg) => match msg.diffs {
OneOrManyDiffs::One(diff) => Poll::Ready(Some(diff)),
OneOrManyDiffs::Many(diffs) if diffs.is_empty() => {
unreachable!("ObservableVectorTransaction never sends empty diffs")
}
OneOrManyDiffs::Many(mut diffs) if diffs.len() == 1 => {
Poll::Ready(Some(diffs.pop().unwrap()))
}
OneOrManyDiffs::Many(diffs) => {
let mut iter = diffs.into_iter();
let fst = iter.next().unwrap();
self.state = VectorSubscriberStreamState::YieldBatch { iter, rx };
return Poll::Ready(Some(fst));
}
},
Err(RecvError::Closed) => Poll::Ready(None),
Err(RecvError::Lagged(_)) => {
Poll::Ready(handle_lag(&mut rx).map(|values| VectorDiff::Reset { values }))
}
};
self.inner.set(rx);
poll
}
VectorSubscriberStreamState::YieldBatch { iter, .. } => {
let diff =
iter.next().expect("YieldBatch is never left empty when exiting poll_next");
if iter.len() == 0 {
let old_state =
mem::replace(&mut self.state, VectorSubscriberStreamState::Recv);
let rx = match old_state {
VectorSubscriberStreamState::YieldBatch { rx, .. } => rx,
// Safety: We would not be in the outer branch otherwise
_ => unsafe { unreachable_unchecked() },
};
self.inner.set(rx);
}
Poll::Ready(Some(diff))
}
}
}
}
/// A batched stream of `VectorDiff`s created from a [`VectorSubscriber`].
///
/// Use its [`Stream`] implementation to interact with it (futures-util and
/// other futures-related crates have extension traits with convenience
/// methods).
#[derive(Debug)]
pub struct VectorSubscriberBatchedStream<T> {
inner: ReusableBoxRecvFuture<T>,
}
impl<T> VectorSubscriberBatchedStream<T> {
fn new(inner: ReusableBoxRecvFuture<T>) -> Self {
Self { inner }
}
}
impl<T: Clone + 'static> Stream for VectorSubscriberBatchedStream<T> {
type Item = Vec<VectorDiff<T>>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
fn append<T>(target: &mut Vec<VectorDiff<T>>, source: OneOrManyDiffs<T>) {
match source {
OneOrManyDiffs::One(diff) => target.push(diff),
OneOrManyDiffs::Many(mut diffs) => target.append(&mut diffs),
}
}
let (result, mut rx) = ready!(self.inner.poll(cx));
let poll = match result {
Ok(msg) => {
let mut batch = msg.diffs.into_vec();
loop {
match rx.try_recv() {
Ok(msg) => append(&mut batch, msg.diffs),
Err(TryRecvError::Empty | TryRecvError::Closed) => {
break Poll::Ready(Some(batch));
}
Err(TryRecvError::Lagged(_)) => {
break Poll::Ready(
handle_lag(&mut rx)
.map(|values| vec![VectorDiff::Reset { values }]),
);
}
}
}
}
Err(RecvError::Closed) => Poll::Ready(None),
Err(RecvError::Lagged(_)) => {
Poll::Ready(handle_lag(&mut rx).map(|values| vec![VectorDiff::Reset { values }]))
}
};
self.inner.set(rx);
poll
}
}
fn handle_lag<T: Clone + 'static>(rx: &mut Receiver<BroadcastMessage<T>>) -> Option<Vector<T>> {
let mut msg = None;
loop {
match rx.try_recv() {
// There's a newer message in the receiver's buffer, use that for reset.
Ok(m) => {
msg = Some(m);
}
// Ideally we'd return a `VecDiff::Reset` with the last state before the
// channel was closed here, but we have no way of obtaining the last state.
Err(TryRecvError::Closed) => {
#[cfg(feature = "tracing")]
info!("Channel closed after lag, can't return last state");
return None;
}
// Lagged twice in a row, is this possible? If it is, it's fine to just
// loop again and look at the next try_recv result.
Err(TryRecvError::Lagged(_)) => {}
Err(TryRecvError::Empty) => match msg {
// We exhausted the internal buffer using try_recv, msg contains the
// last message from it, which we use for the reset.
Some(msg) => return Some(msg.state),
// We exhausted the internal buffer using try_recv but there was no
// message in it, even though we got TryRecvError::Lagged(_) before.
None => unreachable!("got no new message via try_recv after lag"),
},
}
}
}
type SubscriberFutureReturn<T> = (Result<T, RecvError>, Receiver<T>);
struct ReusableBoxRecvFuture<T> {
inner: ReusableBoxFuture<'static, SubscriberFutureReturn<BroadcastMessage<T>>>,
}
async fn make_recv_future<T: Clone>(mut rx: Receiver<T>) -> SubscriberFutureReturn<T> {
let result = rx.recv().await;
(result, rx)
}
impl<T> ReusableBoxRecvFuture<T>
where
T: Clone + 'static,
{
fn set(&mut self, rx: Receiver<BroadcastMessage<T>>) {
self.inner.set(make_recv_future(rx));
}
fn poll(&mut self, cx: &mut Context<'_>) -> Poll<SubscriberFutureReturn<BroadcastMessage<T>>> {
self.inner.poll(cx)
}
}
impl<T> ReusableBoxRecvFuture<T>
where
T: Clone + 'static,
{
fn new(rx: Receiver<BroadcastMessage<T>>) -> Self {
Self { inner: ReusableBoxFuture::new(make_recv_future(rx)) }
}
}
fn assert_send<T: Send>(_val: T) {}
#[allow(unused)]
fn assert_make_future_send() {
#[derive(Clone)]
struct IsSend(*mut ());
unsafe impl Send for IsSend {}
let (_sender, receiver): (_, Receiver<IsSend>) = broadcast::channel(1);
assert_send(make_recv_future(receiver));
}
// SAFETY: make_future is Send if T is, as proven by assert_make_future_send.
unsafe impl<T: Send> Send for ReusableBoxRecvFuture<T> {}
impl<T> fmt::Debug for ReusableBoxRecvFuture<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ReusableBoxRecvFuture").finish()
}
}