use crate::RUNTIME;
#[derive(Clone, Debug)]
struct PoolStatus {
generation: u32,
}
impl Default for PoolStatus {
fn default() -> Self {
PoolStatus { generation: 0 }
}
}
pub(super) fn channel() -> (PoolGenerationPublisher, PoolGenerationSubscriber) {
let (sender, mut receiver) = tokio::sync::watch::channel(Default::default());
RUNTIME.block_in_place(receiver.recv());
(
PoolGenerationPublisher { sender },
PoolGenerationSubscriber { receiver },
)
}
#[derive(Debug)]
pub(super) struct PoolGenerationPublisher {
sender: tokio::sync::watch::Sender<PoolStatus>,
}
impl PoolGenerationPublisher {
pub(super) fn publish(&self, new_generation: u32) {
let new_status = PoolStatus {
generation: new_generation,
};
let _: std::result::Result<_, _> = self.sender.broadcast(new_status);
}
}
#[derive(Clone, Debug)]
pub(crate) struct PoolGenerationSubscriber {
receiver: tokio::sync::watch::Receiver<PoolStatus>,
}
impl PoolGenerationSubscriber {
pub(crate) fn generation(&self) -> u32 {
self.receiver.borrow().generation
}
#[cfg(test)]
pub(crate) async fn wait_for_generation_change(
&mut self,
timeout: std::time::Duration,
) -> Option<u32> {
crate::RUNTIME
.timeout(timeout, self.receiver.recv())
.await
.ok()
.flatten()
.map(|status| status.generation)
}
}