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. //! The segment is closed with [Error::Closed] when all publishers or subscribers are dropped.
use core::fmt; use core::fmt;
use std::{ use std::{
future::{poll_fn, Future}, future::poll_fn,
io, io,
ops::Deref, ops::Deref,
pin::Pin, pin::Pin,
@ -21,7 +21,6 @@ use std::{
use crate::{Error, VarInt}; use crate::{Error, VarInt};
use bytes::{Bytes, BytesMut}; use bytes::{Bytes, BytesMut};
use tokio::pin;
use super::Watch; use super::Watch;
@ -171,29 +170,29 @@ impl Subscriber {
/// Check if there is a chunk available. /// Check if there is a chunk available.
pub fn poll_chunk(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<Bytes>, Error>> { 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() { if !self.buffer.is_empty() {
let chunk = self.buffer.split().freeze(); let chunk = self.buffer.split().freeze();
return Poll::Ready(Ok(Some(chunk))); return Poll::Ready(Ok(Some(chunk)));
} }
// Grab the lock and check if there's a new chunk available.
let state = self.state.lock(); let state = self.state.lock();
if self.index < state.data.len() { if self.index < state.data.len() {
// Yep, clone and return it.
let chunk = state.data[self.index].clone(); let chunk = state.data[self.index].clone();
self.index += 1; self.index += 1;
return Poll::Ready(Ok(Some(chunk))); 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(Error::Closed) => return Poll::Ready(Ok(None)),
Err(err) => return Poll::Ready(Err(err)), 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 Poll::Pending
} }
@ -227,7 +226,7 @@ impl tokio::io::AsyncRead for Subscriber {
// No more data. // No more data.
Ok(None) => return Poll::Ready(Ok(())), 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())), 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. // Upgrade to a mutable references that automatically calls notify on drop.
pub fn into_mut(self) -> WatchMut<'a, T> { pub fn into_mut(self) -> WatchMut<'a, T> {
WatchMut { lock: self.lock } WatchMut { lock: self.lock }