moq-uring 0.0.8

//! A steered thread-per-core group, end to end: two workers on two threads,
//! each adopting its own member of one `SO_REUSEPORT` group and serving an
//! endpoint on it, for clients that dial the shared port. The member carries
//! the slot, so the endpoint issues steering-prefixed connection ids without
//! being told which. Whichever worker the Initial hashes to owns the
//! connection, and the prefix keeps every later packet (handshake
//! continuation included) on that worker; a wrong prefix stalls the
//! handshake, so every dial completing is what proves the steering.
//!
//! Kernel-gated: skips loudly below the Linux 6.12 floor (GitHub-hosted CI),
//! and runs everywhere else.

#![cfg(all(target_os = "linux", feature = "noq"))]

#[path = "support.rs"]
mod support;

use std::net::UdpSocket;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::task::Poll;

use moq_sock::shard::Group;
use moq_uring::{Config, Error, Worker, quic, udp};

fn worker() -> Option<Worker> {
	match Worker::new(Config::default()) {
		Ok(worker) => Some(worker),
		Err(Error::Unsupported(reason)) => {
			eprintln!("skipping io_uring workers test: {reason}");
			None
		}
		Err(err) => panic!("worker setup failed: {err}"),
	}
}

const ALPN: &str = "moq-uring-workers";
const WORKERS: u16 = 2;
/// Enough dials that both workers land some with near certainty: each
/// Initial's placement hashes a random byte, so all landing on one side is a
/// `2^-31` event. Steering cannot be asserted deterministically from here,
/// since the client picks its own connection id; the count is what keeps the
/// false failure below every other source of noise in the suite.
const DIALS: usize = 32;

/// A stop signal a worker parks on, wakeable from another thread.
#[derive(Default)]
struct Stop {
	stopped: AtomicBool,
	waker: Mutex<Option<std::task::Waker>>,
}

impl Stop {
	async fn wait(self: &Arc<Self>) {
		std::future::poll_fn(|cx| {
			if self.stopped.load(Ordering::Acquire) {
				return Poll::Ready(());
			}
			*self.waker.lock().unwrap() = Some(cx.waker().clone());
			Poll::Pending
		})
		.await
	}

	fn stop(&self) {
		self.stopped.store(true, Ordering::Release);
		if let Some(waker) = self.waker.lock().unwrap().take() {
			waker.wake();
		}
	}
}

#[test]
fn a_steered_group_serves_a_shared_port() {
	// Gate on the kernel before spawning anything.
	let Some(client_worker) = worker() else { return };
	let certs = support::certs().expect("certificates");

	// Bind the group up front, in index order: that order is the identity the
	// kernel steers by, and binding before any thread spawns is what
	// guarantees it. The group is held for as long as the sockets are served,
	// since it is what holds the port.
	let mut group = Group::acquire("127.0.0.1:0".parse().expect("addr"), WORKERS).expect("group");
	let mut claims = Vec::new();
	while let Some(member) = group.member() {
		claims.push(member.bind().expect("bind group member"));
	}
	let mut group = group.complete(claims).expect("complete group");
	let mut members = Vec::new();
	while let Some(member) = group.member().expect("clone retained socket") {
		members.push(member);
	}
	let addr = group.addr();

	let accepted: Arc<Vec<AtomicUsize>> = Arc::new((0..WORKERS).map(|_| AtomicUsize::new(0)).collect());
	// One stop each: the wakers are per-thread, so a shared slot would let one
	// worker's registration clobber the other's.
	let stops: Vec<Arc<Stop>> = (0..WORKERS).map(|_| Arc::new(Stop::default())).collect();
	// Each worker reports in once serving, so a setup failure fails here
	// instead of as the dials hashed to it idling out.
	let (ready, started) = std::sync::mpsc::channel();

	let threads: Vec<_> = members
		.into_iter()
		.zip(&stops)
		.map(|(member, stop)| {
			let accepted = accepted.clone();
			let stop = stop.clone();
			let cert = certs.cert.clone();
			let key = certs.key.clone();
			let ready = ready.clone();
			std::thread::spawn(move || {
				let shard = member.shard();
				let mut worker = Worker::new(Config::default()).expect("worker");
				let handle = worker.handle();
				// The member is adopted whole: its slot is what the endpoint
				// steers with, and there is no other way to hand one over.
				let socket = handle.udp(member, udp::Config::default()).expect("socket");

				let mut server = quic::server::Config::new(quic::Identity::open(cert, key).expect("identity"));
				server.alpn = vec![ALPN.to_string()];
				let endpoint = quic::Endpoint::new(socket, quic::endpoint::Config::default().with_server(server))
					.expect("endpoint");

				handle.spawn(async move {
					// Accepted connections are dropped once counted; the
					// client is what closes them.
					while endpoint.accept().await.is_ok() {
						accepted[usize::from(shard.index())].fetch_add(1, Ordering::AcqRel);
					}
				});
				ready.send(()).expect("test alive");
				worker.block_on(stop.wait()).expect("worker loop");
			})
		})
		.collect();
	drop(ready);
	for _ in 0..WORKERS {
		started.recv().expect("a worker thread failed to start");
	}

	// Dial the shared port repeatedly from one client worker. Every handshake
	// completing is the steering assertion (see the module docs).
	let mut client_worker = client_worker;
	let handle = client_worker.handle();
	let mut dial = quic::client::Config::new(addr, "localhost");
	dial.alpn = vec![ALPN.to_string()];
	dial.verify = false;

	client_worker
		.block_on(async {
			for _ in 0..DIALS {
				let socket = handle
					.udp(UdpSocket::bind("127.0.0.1:0").expect("bind"), udp::Config::default())
					.expect("client socket");
				let mut conn = quic::client::connect(socket, &dial).await.expect("connect");
				assert_eq!(
					web_transport_trait::poll::Session::protocol(&conn),
					Some(ALPN),
					"negotiated ALPN"
				);
				web_transport_trait::poll::Session::close(&mut conn, 0, "done");
			}

			// The server side counts a connection when its accept loop takes
			// it, which can trail the client's handshake; wait for the tally.
			let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
			loop {
				let total: usize = accepted.iter().map(|count| count.load(Ordering::Acquire)).sum();
				if total == DIALS {
					break;
				}
				assert!(
					std::time::Instant::now() < deadline,
					"only {total} of {DIALS} dials were accepted"
				);
				moq_uring::Timer::after(&handle, std::time::Duration::from_millis(10))
					.wait()
					.await;
			}
		})
		.expect("client worker");

	// Every member has to have been fed, or the group is steering into a
	// subset and the rest sit idle.
	for (index, count) in accepted.iter().enumerate() {
		assert!(count.load(Ordering::Acquire) > 0, "worker {index} accepted nothing");
	}

	for stop in &stops {
		stop.stop();
	}
	for thread in threads {
		thread.join().expect("worker thread");
	}
}