Only consume budget on successful I/O operations

src/blocking.rs

@@ -273,6 +273,11 @@ pub fn iter<T: Send + 'static>(
                         *self = State::Idle(Some(iter));
                     }
 
+                    if opt.is_some() {
+                        // Consume a unit of I/O budget.
+                        throttle::bump();
+                    }
+
                     Poll::Ready(opt)
                 }
             }
@@ -379,6 +384,11 @@ pub fn reader(reader: impl Read + Send + 'static) -> impl AsyncRead + Send + Unp
                         res?;
                     }
 
+                    if n > 0 {
+                        // Consume a unit of I/O budget.
+                        throttle::bump();
+                    }
+
                     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.
-                    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));
                         // Make sure to move into the idle state before reporting errors.
                         *self = State::Idle(Some(io));
+
+                        if res.is_ok() {
+                            // Consume a unit of I/O budget.
+                            throttle::bump();
+                        }
                         return Poll::Ready(res);
                     }
                 }

src/reactor.rs

@@ -343,7 +343,13 @@ impl Source {
             // Attempt the non-blocking operation.
             match op() {
                 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.

src/throttle.rs

@@ -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
 //! 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`
 pub(crate) fn poll(cx: &mut Context<'_>) -> Poll<()> {
-    // Decrement the budget and check if it was zero.
-    if BUDGET.is_set() && BUDGET.with(|b| b.replace(b.get().saturating_sub(1))) == 0 {
+    // Check if the budget is zero.
+    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
-        // artificially throttling it to let other tasks be run.
+        // artificially throttling it to let other tasks run.
         cx.waker().wake_by_ref();
         return Poll::Pending;
     }
     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)));
+    }
+}