use std::{collections::HashSet, hash::Hash, sync::Arc, time::Duration};
use parking_lot::Mutex;
use tokio::{select, sync::Notify};
use crate::raw::Raw;
mod raw;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Update<V> {
Add(V),
Delete,
}
impl<V> Update<V> {
pub fn map<W>(self, f: impl FnOnce(V) -> W) -> Update<W> {
match self {
Update::Add(v) => Update::Add(f(v)),
Update::Delete => Update::Delete,
}
}
pub fn into_option(self) -> Option<V> {
match self {
Update::Add(v) => Some(v),
Update::Delete => None,
}
}
}
pub fn channel<K, V>(cooldown: Duration) -> (Sender<K, V>, Receiver<K, V>) {
channel_with_starting_keys(HashSet::new(), cooldown)
}
pub fn channel_with_starting_keys<K, V>(
starting_keys: HashSet<K>,
cooldown: Duration,
) -> (Sender<K, V>, Receiver<K, V>) {
let channel = Arc::new(Channel {
inner: Mutex::new(Inner {
raw: Raw::new(starting_keys, cooldown),
dropped: false,
}),
changed: Notify::new(),
});
(Sender(Arc::clone(&channel)), Receiver(channel))
}
struct Channel<K, V> {
inner: Mutex<Inner<K, V>>,
changed: Notify,
}
struct Inner<K, V> {
raw: Raw<K, V>,
dropped: bool,
}
pub struct Sender<K, V>(Arc<Channel<K, V>>);
impl<K: Clone + Eq + Hash, V> Sender<K, V> {
pub fn send(&self, key: K, update: Update<V>) -> Result<(), SendError<V>> {
let Channel { inner, changed } = self.0.as_ref();
let replaced = {
let mut inner = inner.lock();
let Inner { raw, dropped } = &mut *inner;
if *dropped {
return Err(SendError(update));
}
raw.send(key, update)
};
changed.notify_waiters();
drop(replaced);
Ok(())
}
pub fn clear(&self) -> Result<(), SendError<V>> {
let Channel { inner, changed } = self.0.as_ref();
{
let mut inner = inner.lock();
let Inner { raw, dropped } = &mut *inner;
if *dropped {
return Err(SendError(Update::Delete));
}
raw.clear();
}
changed.notify_waiters();
Ok(())
}
}
#[derive(Debug)]
pub struct SendError<V>(pub Update<V>);
pub struct Receiver<K, V>(Arc<Channel<K, V>>);
impl<K: Clone + Eq + Hash, V> Receiver<K, V> {
pub async fn recv(&mut self) -> Option<(K, Update<V>)> {
let Channel { inner, changed } = self.0.as_ref();
loop {
let notified_fut = changed.notified();
let next_cooldown = {
let mut inner = inner.lock();
let Inner { raw, dropped } = &mut *inner;
if let Some(received) = raw.try_recv() {
return Some(received);
}
if *dropped {
return None;
}
raw.wait_next_cooldown()
};
select! {
() = notified_fut => {},
() = next_cooldown => {},
}
}
}
pub fn try_recv(&mut self) -> Result<(K, Update<V>), TryRecvError> {
let mut inner = self.0.inner.lock();
let Inner { raw, dropped } = &mut *inner;
match raw.try_recv() {
Some(output) => Ok(output),
None => Err(TryRecvError::from_dropped(*dropped)),
}
}
pub fn try_recv_now(&mut self) -> Result<(K, Update<V>), TryRecvError> {
let mut inner = self.0.inner.lock();
let Inner { raw, dropped } = &mut *inner;
match raw.try_recv_now() {
Some(output) => Ok(output),
None => Err(TryRecvError::from_dropped(*dropped)),
}
}
}
#[derive(Debug, PartialEq, Eq)]
pub enum TryRecvError {
Disconnected,
Empty,
}
impl TryRecvError {
fn from_dropped(dropped: bool) -> Self {
if dropped {
Self::Disconnected
} else {
Self::Empty
}
}
}
impl<K, V> Drop for Sender<K, V> {
fn drop(&mut self) {
let Channel { inner, changed } = self.0.as_ref();
inner.lock().dropped = true;
changed.notify_waiters();
}
}
impl<K, V> Drop for Receiver<K, V> {
fn drop(&mut self) {
let Channel { inner, changed } = self.0.as_ref();
inner.lock().dropped = true;
changed.notify_waiters();
}
}