mirror of https://github.com/smol-rs/smol
Only consume budget on successful I/O operations
This commit is contained in:
parent
e48a8787f4
commit
19d8d5a16d
|
@ -273,6 +273,11 @@ pub fn iter<T: Send + 'static>(
|
||||||
*self = State::Idle(Some(iter));
|
*self = State::Idle(Some(iter));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if opt.is_some() {
|
||||||
|
// Consume a unit of I/O budget.
|
||||||
|
throttle::bump();
|
||||||
|
}
|
||||||
|
|
||||||
Poll::Ready(opt)
|
Poll::Ready(opt)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -379,6 +384,11 @@ pub fn reader(reader: impl Read + Send + 'static) -> impl AsyncRead + Send + Unp
|
||||||
res?;
|
res?;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if n > 0 {
|
||||||
|
// Consume a unit of I/O budget.
|
||||||
|
throttle::bump();
|
||||||
|
}
|
||||||
|
|
||||||
Poll::Ready(Ok(n))
|
Poll::Ready(Ok(n))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -496,7 +506,16 @@ pub fn writer(writer: impl Write + Send + 'static) -> impl AsyncWrite + Send + U
|
||||||
}
|
}
|
||||||
|
|
||||||
// The writer is busy - write more bytes into the pipe.
|
// The writer is busy - write more bytes into the pipe.
|
||||||
State::Busy(Some(writer), _) => return Pin::new(writer).poll_write(cx, buf),
|
State::Busy(Some(writer), _) => {
|
||||||
|
let res = futures::ready!(Pin::new(writer).poll_write(cx, buf));
|
||||||
|
if let Ok(n) = &res {
|
||||||
|
if *n > 0 {
|
||||||
|
// Consume a unit of I/O budget.
|
||||||
|
throttle::bump();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return Poll::Ready(res);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -524,6 +543,11 @@ pub fn writer(writer: impl Write + Send + 'static) -> impl AsyncWrite + Send + U
|
||||||
let (res, io) = futures::ready!(Pin::new(task).poll(cx));
|
let (res, io) = futures::ready!(Pin::new(task).poll(cx));
|
||||||
// Make sure to move into the idle state before reporting errors.
|
// Make sure to move into the idle state before reporting errors.
|
||||||
*self = State::Idle(Some(io));
|
*self = State::Idle(Some(io));
|
||||||
|
|
||||||
|
if res.is_ok() {
|
||||||
|
// Consume a unit of I/O budget.
|
||||||
|
throttle::bump();
|
||||||
|
}
|
||||||
return Poll::Ready(res);
|
return Poll::Ready(res);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -343,7 +343,13 @@ impl Source {
|
||||||
// Attempt the non-blocking operation.
|
// Attempt the non-blocking operation.
|
||||||
match op() {
|
match op() {
|
||||||
Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
|
Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
|
||||||
res => return Poll::Ready(res),
|
res => {
|
||||||
|
// Consume a unit of I/O budget.
|
||||||
|
if res.is_ok() {
|
||||||
|
throttle::bump();
|
||||||
|
}
|
||||||
|
return Poll::Ready(res);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Lock the waker list and retry the non-blocking operation.
|
// Lock the waker list and retry the non-blocking operation.
|
||||||
|
|
|
@ -1,4 +1,4 @@
|
||||||
//! Throttle tasks if they poll too many I/O operations without yielding.
|
//! Throttle tasks if they perform too many I/O operations without yielding.
|
||||||
//!
|
//!
|
||||||
//! This is used to prevent futures from running forever. Once a certain number of I/O operation is
|
//! This is used to prevent futures from running forever. Once a certain number of I/O operation is
|
||||||
//! hit in a single run, I/O operations will begin returning
|
//! hit in a single run, I/O operations will begin returning
|
||||||
|
@ -35,12 +35,22 @@ pub(crate) fn setup<T>(poll: impl FnOnce() -> T) -> T {
|
||||||
///
|
///
|
||||||
/// [`Poll::Pending`]: `std::task::Poll::Pending`
|
/// [`Poll::Pending`]: `std::task::Poll::Pending`
|
||||||
pub(crate) fn poll(cx: &mut Context<'_>) -> Poll<()> {
|
pub(crate) fn poll(cx: &mut Context<'_>) -> Poll<()> {
|
||||||
// Decrement the budget and check if it was zero.
|
// Check if the budget is zero.
|
||||||
if BUDGET.is_set() && BUDGET.with(|b| b.replace(b.get().saturating_sub(1))) == 0 {
|
if BUDGET.is_set() && BUDGET.with(|b| b.get()) == 0 {
|
||||||
// Make sure to wake the current task. The task is not *really* pending, we're just
|
// Make sure to wake the current task. The task is not *really* pending, we're just
|
||||||
// artificially throttling it to let other tasks be run.
|
// artificially throttling it to let other tasks run.
|
||||||
cx.waker().wake_by_ref();
|
cx.waker().wake_by_ref();
|
||||||
return Poll::Pending;
|
return Poll::Pending;
|
||||||
}
|
}
|
||||||
Poll::Ready(())
|
Poll::Ready(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Consumes a unit of budget.
|
||||||
|
///
|
||||||
|
/// This function is called by successful I/O operation.
|
||||||
|
pub(crate) fn bump() {
|
||||||
|
if BUDGET.is_set() {
|
||||||
|
// Decrement the budget.
|
||||||
|
BUDGET.with(|b| b.replace(b.get().saturating_sub(1)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue