Repository: tokio-rs/tokio
File: tokio/tests/coop_budget.rs (Rust)

Code before:
#![warn(rust_2018_idioms)]
#![cfg(all(feature = "full", target_os = "linux"))]

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::net::UdpSocket;
use tokio::runtime::coop::has_budget_remaining;
use tokio::task::consume_budget;

/// 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]
#[cfg_attr(miri, ignore)] // No `socket` on miri.
async fn coop_budget_udp_send_recv() {
    const BUDGET: usize = 128;
    const N_ITERATIONS: usize = 1024;

    const PACKET: &[u8] = b"Hello, world";
    const PACKET_LEN: usize = 12;

Review comment (suggestion):
I want the modules to work out so that this becomes.
```rs
use tokio::task::coop::{consume_budget, has_budget_remaining};
```
I know that `consume_budget` is defined directly in `task`, but we can move it as long as we keep a (possibly hidden) re-export.

Code after:
#![warn(rust_2018_idioms)]
#![cfg(all(feature = "full", target_os = "linux"))]

use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::net::UdpSocket;
use tokio::task::coop::{consume_budget, has_budget_remaining};

const BUDGET: usize = 128;

/// 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]
#[cfg_attr(miri, ignore)] // No `socket` on miri.
async fn coop_budget_udp_send_recv() {
    const N_ITERATIONS: usize = 1024;

    const PACKET: &[u8] = b"Hello, world";
    const PACKET_LEN: usize = 12;