From 60ea36d4600da2a24e3897262967abb1e7c07670 Mon Sep 17 00:00:00 2001
From: = <=>
Date: Sat, 16 Dec 2023 13:35:21 -0600
Subject: [PATCH] io: add budgeting to
 tokio::runtime::io::registration::async_io

Fixes #5946, #4782.

This change adds budgeting to most of the remaining unbudgeted IO operations which can complete instantly, including datagram send/recv operations and listener socket accepts.

This is particularly significant for scenarios in which resource limits are hit, as it can be common for things like listener tasks to spin when receiving errors and just log them, busy looping worker threads which might otherwise be handling existing connections and closing them.

This can also sometimes lead to complex failure scenarios within datagram systems experiencing resource exhaustion.
---
 tokio/src/runtime/io/registration.rs |  8 ++-
 tokio/tests/coop_budget.rs           | 77 ++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 1 deletion(-)
 create mode 100644 tokio/tests/coop_budget.rs

diff --git a/tokio/src/runtime/io/registration.rs b/tokio/src/runtime/io/registration.rs
index 759589863eb..33891173c2d 100644
--- a/tokio/src/runtime/io/registration.rs
+++ b/tokio/src/runtime/io/registration.rs
@@ -219,11 +219,17 @@ impl Registration {
         loop {
             let event = self.readiness(interest).await?;
 
+            let coop = crate::future::poll_fn(|cx| crate::runtime::coop::poll_proceed(cx)).await;
+
             match f() {
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                     self.clear_readiness(event);
                 }
-                x => return x,
+                x => {
+                    coop.made_progress();
+
+                    return x;
+                }
             }
         }
     }
diff --git a/tokio/tests/coop_budget.rs b/tokio/tests/coop_budget.rs
new file mode 100644
index 00000000000..05b5c117692
--- /dev/null
+++ b/tokio/tests/coop_budget.rs
@@ -0,0 +1,77 @@
+#![warn(rust_2018_idioms)]
+#![cfg(all(feature = "full", target_os = "linux"))]
+
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+use tokio::net::UdpSocket;
+
+/// Ensure that UDP sockets have functional budgeting
+///
+/// # Design
+/// Two sockets communicate by spamming packets from one to the other.
+///
+/// In Linux, this packet will be slammed through the entire network stack and into the receiver's buffer during the
+/// send system call because we are using the loopback interface.
+/// This happens because the softirq chain invoked on send when using the loopback interface covers virtually the
+/// entirety of the lifecycle of a packet within the kernel network stack.
+///
+/// As a result, neither socket will ever encounter an EWOULDBLOCK, and the only way for these to yield during the loop
+/// is through budgeting.
+///
+/// A second task runs in the background and increments a counter before yielding, allowing us to know how many times sockets yielded.
+/// Since we are both sending and receiving, that should happen once per 64 packets, because budgets are of size 128
+/// and there are two budget events per packet, a send and a recv.
+#[tokio::test]
+async fn coop_budget_udp_send_recv() {
+    const BUDGET: u64 = 128;
+    const N_ITERATIONS: u64 = 1024;
+
+    const PACKET: &[u8] = b"Hello, world";
+    const PACKET_LEN: usize = 12;
+
+    assert_eq!(
+        PACKET_LEN,
+        PACKET.len(),
+        "Defect in test, programmer can't do math"
+    );
+
+    // bind each socket to a dynamic port, forcing IPv4 addressing on the localhost interface
+    let tx = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+    let rx = UdpSocket::bind("127.0.0.1:0").await.unwrap();
+
+    tx.connect(rx.local_addr().unwrap()).await.unwrap();
+    rx.connect(tx.local_addr().unwrap()).await.unwrap();
+
+    let tracker = Arc::new(AtomicU64::default());
+
+    let tracker_clone = Arc::clone(&tracker);
+
+    tokio::task::yield_now().await;
+
+    tokio::spawn(async move {
+        loop {
+            tracker_clone.fetch_add(1, Ordering::SeqCst);
+
+            tokio::task::yield_now().await;
+        }
+    });
+
+    for _ in 0..N_ITERATIONS {
+        tx.send(PACKET).await.unwrap();
+
+        let mut tmp = [0; PACKET_LEN];
+
+        // ensure that we aren't somehow accumulating other
+        assert_eq!(
+            PACKET_LEN,
+            rx.recv(&mut tmp).await.unwrap(),
+            "Defect in test case, received unexpected result from socket"
+        );
+        assert_eq!(
+            PACKET, &tmp,
+            "Defect in test case, received unexpected result from socket"
+        );
+    }
+
+    assert_eq!(N_ITERATIONS / (BUDGET / 2), tracker.load(Ordering::SeqCst));
+}