70 lines
1.3 KiB
Rust
70 lines
1.3 KiB
Rust
use std::{
|
|
collections::{BTreeMap, btree_map::Entry},
|
|
future::Future,
|
|
ops::RangeToInclusive,
|
|
sync::Mutex,
|
|
};
|
|
|
|
use futures::pin_mut;
|
|
use tokio::sync::watch::{Sender, channel};
|
|
use tuwunel_core::implement;
|
|
|
|
use crate::keyval::KeyBuf;
|
|
|
|
type Watchers = Mutex<BTreeMap<KeyBuf, Sender<()>>>;
|
|
|
|
#[derive(Default)]
|
|
pub(super) struct Watch {
|
|
watchers: Watchers,
|
|
}
|
|
|
|
#[implement(super::Map)]
|
|
pub fn watch_raw_prefix<K>(&self, prefix: &K) -> impl Future<Output = ()> + Send + use<K>
|
|
where
|
|
K: AsRef<[u8]> + ?Sized,
|
|
{
|
|
let rx = match self
|
|
.watch
|
|
.watchers
|
|
.lock()
|
|
.expect("locked")
|
|
.entry(prefix.as_ref().into())
|
|
{
|
|
| Entry::Occupied(node) => node.get().subscribe(),
|
|
| Entry::Vacant(node) => {
|
|
let (tx, rx) = channel(());
|
|
node.insert(tx);
|
|
rx
|
|
},
|
|
};
|
|
|
|
async move {
|
|
pin_mut!(rx);
|
|
rx.changed()
|
|
.await
|
|
.expect("watcher sender dropped");
|
|
}
|
|
}
|
|
|
|
#[implement(super::Map)]
|
|
pub(crate) fn notify<K>(&self, key: &K)
|
|
where
|
|
K: AsRef<[u8]> + Ord + ?Sized,
|
|
{
|
|
let range = RangeToInclusive::<KeyBuf> { end: key.as_ref().into() };
|
|
|
|
let mut watchers = self.watch.watchers.lock().expect("locked");
|
|
|
|
watchers
|
|
.range(range)
|
|
.rev()
|
|
.take_while(|(k, _)| key.as_ref().starts_with(k))
|
|
.filter_map(|(k, tx)| tx.send(()).is_err().then_some(k))
|
|
.cloned()
|
|
.collect::<Vec<_>>()
|
|
.into_iter()
|
|
.for_each(|k| {
|
|
watchers.remove(&k);
|
|
});
|
|
}
|