Small improvements and comments.

I actually kind of like the simplicity of wake() over changed()
This commit is contained in:
Luke Curley 2023-09-17 11:01:03 -07:00
parent ddf22012e0
commit 73151fbcfe
2 changed files with 13 additions and 9 deletions

View File

@ -10,7 +10,7 @@
//! The segment is closed with [Error::Closed] when all publishers or subscribers are dropped.
use core::fmt;
use std::{
future::{poll_fn, Future},
future::poll_fn,
io,
ops::Deref,
pin::Pin,
@ -21,7 +21,6 @@ use std::{
use crate::{Error, VarInt};
use bytes::{Bytes, BytesMut};
use tokio::pin;
use super::Watch;
@ -171,29 +170,29 @@ impl Subscriber {
/// Check if there is a chunk available.
pub fn poll_chunk(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Bytes>, Error>> {
// If there's already buffered data, return it.
if !self.buffer.is_empty() {
let chunk = self.buffer.split().freeze();
return Poll::Ready(Ok(Some(chunk)));
}
// Grab the lock and check if there's a new chunk available.
let state = self.state.lock();
if self.index < state.data.len() {
// Yep, clone and return it.
let chunk = state.data[self.index].clone();
self.index += 1;
return Poll::Ready(Ok(Some(chunk)));
}
let notify = match state.closed {
// Otherwise we wait until the state changes and try again.
match state.closed {
Err(Error::Closed) => return Poll::Ready(Ok(None)),
Err(err) => return Poll::Ready(Err(err)),
Ok(()) => state.changed(), // Wake up when the state changes
Ok(()) => state.waker(cx), // Wake us up when the state changes.
};
// Register our context with the notify waker.
pin!(notify);
let _ = notify.poll(cx);
Poll::Pending
}
@ -227,7 +226,7 @@ impl tokio::io::AsyncRead for Subscriber {
// No more data.
Ok(None) => return Poll::Ready(Ok(())),
// TODO cast to io::Error
// Crudely cast to io::Error
Err(err) => return Poll::Ready(Err(err.as_io())),
};

View File

@ -108,6 +108,11 @@ impl<'a, T> WatchRef<'a, T> {
}
}
// Release the lock and provide a context to wake when next updated.
pub fn waker(mut self, cx: &mut task::Context<'_>) {
self.lock.register(cx.waker());
}
// Upgrade to a mutable references that automatically calls notify on drop.
pub fn into_mut(self) -> WatchMut<'a, T> {
WatchMut { lock: self.lock }