YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-04-01 10:15:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Move all crates to new crates dir

Browse source
This commit is contained in:
Joakim Frostegård 2023-10-18 23:53:41 +02:00
parent 3835da22ac
commit 9b032f7e24
128 changed files with 27 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

548
aquatic_udp/src/workers/socket/uring/mod.rs
View file

@ -1,548 +0,0 @@
mod buf_ring;
mod recv_helper;
mod send_buffers;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::net::UdpSocket;
use std::ops::DerefMut;
use std::os::fd::AsRawFd;
use std::sync::atomic::Ordering;
use anyhow::Context;
use aquatic_common::access_list::AccessListCache;
use aquatic_common::ServerStartInstant;
use crossbeam_channel::Receiver;
use io_uring::opcode::Timeout;
use io_uring::types::{Fixed, Timespec};
use io_uring::{IoUring, Probe};
use aquatic_common::{
access_list::create_access_list_cache, privileges::PrivilegeDropper, CanonicalSocketAddr,
PanicSentinel, ValidUntil,
};
use aquatic_udp_protocol::*;
use crate::common::*;
use crate::config::Config;
use self::buf_ring::BufRing;
use self::recv_helper::RecvHelper;
use self::send_buffers::{ResponseType, SendBuffers};
use super::storage::PendingScrapeResponseSlab;
use super::validator::ConnectionValidator;
use super::{create_socket, EXTRA_PACKET_SIZE_IPV4, EXTRA_PACKET_SIZE_IPV6};
/// Size of each request buffer
///
/// Enough for scrape request with 20 info hashes
const REQUEST_BUF_LEN: usize = 256;
/// Size of each response buffer
///
/// Enough for:
/// - IPv6 announce response with 112 peers
/// - scrape response for 170 info hashes
const RESPONSE_BUF_LEN: usize = 2048;
const USER_DATA_RECV: u64 = u64::MAX;
const USER_DATA_PULSE_TIMEOUT: u64 = u64::MAX - 1;
const USER_DATA_CLEANING_TIMEOUT: u64 = u64::MAX - 2;
const SOCKET_IDENTIFIER: Fixed = Fixed(0);
thread_local! {
/// Store IoUring instance here so that it can be accessed in BufRing::drop
pub static CURRENT_RING: CurrentRing = CurrentRing(RefCell::new(None));
}
pub struct CurrentRing(RefCell<Option<IoUring>>);
impl CurrentRing {
fn with<F, T>(mut f: F) -> T
where
F: FnMut(&mut IoUring) -> T,
{
CURRENT_RING.with(|r| {
let mut opt_ring = r.0.borrow_mut();
f(Option::as_mut(opt_ring.deref_mut()).expect("IoUring not set"))
})
}
}
pub struct SocketWorker {
config: Config,
shared_state: State,
request_sender: ConnectedRequestSender,
response_receiver: Receiver<(ConnectedResponse, CanonicalSocketAddr)>,
access_list_cache: AccessListCache,
validator: ConnectionValidator,
server_start_instant: ServerStartInstant,
#[allow(dead_code)]
socket: UdpSocket,
pending_scrape_responses: PendingScrapeResponseSlab,
buf_ring: BufRing,
send_buffers: SendBuffers,
recv_helper: RecvHelper,
local_responses: VecDeque<(Response, CanonicalSocketAddr)>,
resubmittable_sqe_buf: Vec<io_uring::squeue::Entry>,
recv_sqe: io_uring::squeue::Entry,
pulse_timeout_sqe: io_uring::squeue::Entry,
cleaning_timeout_sqe: io_uring::squeue::Entry,
pending_scrape_valid_until: ValidUntil,
}
impl SocketWorker {
pub fn run(
_sentinel: PanicSentinel,
shared_state: State,
config: Config,
validator: ConnectionValidator,
server_start_instant: ServerStartInstant,
request_sender: ConnectedRequestSender,
response_receiver: Receiver<(ConnectedResponse, CanonicalSocketAddr)>,
priv_dropper: PrivilegeDropper,
) {
let ring_entries = config.network.ring_size.next_power_of_two();
// Try to fill up the ring with send requests
let send_buffer_entries = ring_entries;
let socket = create_socket(&config, priv_dropper).expect("create socket");
let access_list_cache = create_access_list_cache(&shared_state.access_list);
let send_buffers = SendBuffers::new(&config, send_buffer_entries as usize);
let recv_helper = RecvHelper::new(&config);
let ring = IoUring::builder()
.setup_coop_taskrun()
.setup_single_issuer()
.setup_submit_all()
.build(ring_entries.into())
.unwrap();
ring.submitter()
.register_files(&[socket.as_raw_fd()])
.unwrap();
// Store ring in thread local storage before creating BufRing
CURRENT_RING.with(|r| *r.0.borrow_mut() = Some(ring));
let buf_ring = buf_ring::Builder::new(0)
.ring_entries(ring_entries)
.buf_len(REQUEST_BUF_LEN)
.build()
.unwrap();
let recv_sqe = recv_helper.create_entry(buf_ring.bgid().try_into().unwrap());
// This timeout enables regular updates of pending_scrape_valid_until
// and wakes the main loop to send any pending responses in the case
// of no incoming requests
let pulse_timeout_sqe = {
let timespec_ptr = Box::into_raw(Box::new(Timespec::new().sec(1))) as *const _;
Timeout::new(timespec_ptr)
.build()
.user_data(USER_DATA_PULSE_TIMEOUT)
};
let cleaning_timeout_sqe = {
let timespec_ptr = Box::into_raw(Box::new(
Timespec::new().sec(config.cleaning.pending_scrape_cleaning_interval),
)) as *const _;
Timeout::new(timespec_ptr)
.build()
.user_data(USER_DATA_CLEANING_TIMEOUT)
};
let resubmittable_sqe_buf = vec![
recv_sqe.clone(),
pulse_timeout_sqe.clone(),
cleaning_timeout_sqe.clone(),
];
let pending_scrape_valid_until =
ValidUntil::new(server_start_instant, config.cleaning.max_pending_scrape_age);
let mut worker = Self {
config,
shared_state,
validator,
server_start_instant,
request_sender,
response_receiver,
access_list_cache,
pending_scrape_responses: Default::default(),
send_buffers,
recv_helper,
local_responses: Default::default(),
buf_ring,
recv_sqe,
pulse_timeout_sqe,
cleaning_timeout_sqe,
resubmittable_sqe_buf,
socket,
pending_scrape_valid_until,
};
CurrentRing::with(|ring| worker.run_inner(ring));
}
fn run_inner(&mut self, ring: &mut IoUring) {
loop {
for sqe in self.resubmittable_sqe_buf.drain(..) {
unsafe { ring.submission().push(&sqe).unwrap() };
}
let sq_space = {
let sq = ring.submission();
sq.capacity() - sq.len()
};
let mut num_send_added = 0;
// Enqueue local responses
for _ in 0..sq_space {
if let Some((response, addr)) = self.local_responses.pop_front() {
match self.send_buffers.prepare_entry(&response, addr) {
Ok(entry) => {
unsafe { ring.submission().push(&entry).unwrap() };
num_send_added += 1;
}
Err(send_buffers::Error::NoBuffers) => {
self.local_responses.push_front((response, addr));
break;
}
Err(send_buffers::Error::SerializationFailed(err)) => {
::log::error!("Failed serializing response: {:#}", err);
}
}
} else {
break;
}
}
// Enqueue swarm worker responses
for _ in 0..(sq_space - num_send_added) {
if let Some((response, addr)) = self.get_next_swarm_response() {
match self.send_buffers.prepare_entry(&response, addr) {
Ok(entry) => {
unsafe { ring.submission().push(&entry).unwrap() };
num_send_added += 1;
}
Err(send_buffers::Error::NoBuffers) => {
self.local_responses.push_back((response, addr));
break;
}
Err(send_buffers::Error::SerializationFailed(err)) => {
::log::error!("Failed serializing response: {:#}", err);
}
}
} else {
break;
}
}
// Wait for all sendmsg entries to complete. If none were added,
// wait for at least one recvmsg or timeout in order to avoid
// busy-polling if there is no incoming data.
ring.submitter()
.submit_and_wait(num_send_added.max(1))
.unwrap();
for cqe in ring.completion() {
self.handle_cqe(cqe);
}
self.send_buffers.reset_likely_next_free_index();
}
}
fn handle_cqe(&mut self, cqe: io_uring::cqueue::Entry) {
match cqe.user_data() {
USER_DATA_RECV => {
self.handle_recv_cqe(&cqe);
if !io_uring::cqueue::more(cqe.flags()) {
self.resubmittable_sqe_buf.push(self.recv_sqe.clone());
}
}
USER_DATA_PULSE_TIMEOUT => {
self.pending_scrape_valid_until = ValidUntil::new(
self.server_start_instant,
self.config.cleaning.max_pending_scrape_age,
);
::log::info!(
"pending responses: {} local, {} swarm",
self.local_responses.len(),
self.response_receiver.len()
);
self.resubmittable_sqe_buf
.push(self.pulse_timeout_sqe.clone());
}
USER_DATA_CLEANING_TIMEOUT => {
self.pending_scrape_responses
.clean(self.server_start_instant.seconds_elapsed());
self.resubmittable_sqe_buf
.push(self.cleaning_timeout_sqe.clone());
}
send_buffer_index => {
let result = cqe.result();
if result < 0 {
::log::error!(
"Couldn't send response: {:#}",
::std::io::Error::from_raw_os_error(-result)
);
} else if self.config.statistics.active() {
let send_buffer_index = send_buffer_index as usize;
let (response_type, receiver_is_ipv4) =
self.send_buffers.response_type_and_ipv4(send_buffer_index);
let (statistics, extra_bytes) = if receiver_is_ipv4 {
(&self.shared_state.statistics_ipv4, EXTRA_PACKET_SIZE_IPV4)
} else {
(&self.shared_state.statistics_ipv6, EXTRA_PACKET_SIZE_IPV6)
};
statistics
.bytes_sent
.fetch_add(result as usize + extra_bytes, Ordering::Relaxed);
let response_counter = match response_type {
ResponseType::Connect => &statistics.responses_sent_connect,
ResponseType::Announce => &statistics.responses_sent_announce,
ResponseType::Scrape => &statistics.responses_sent_scrape,
ResponseType::Error => &statistics.responses_sent_error,
};
response_counter.fetch_add(1, Ordering::Relaxed);
}
// Safety: OK because cqe using buffer has been returned and
// contents will no longer be accessed by kernel
unsafe {
self.send_buffers
.mark_buffer_as_free(send_buffer_index as usize);
}
}
}
}
fn handle_recv_cqe(&mut self, cqe: &io_uring::cqueue::Entry) {
let result = cqe.result();
if result < 0 {
if -result == libc::ENOBUFS {
::log::info!("recv failed due to lack of buffers. If increasing ring size doesn't help, get faster hardware");
} else {
::log::warn!(
"recv failed: {:#}",
::std::io::Error::from_raw_os_error(-result)
);
}
return;
}
let buffer = unsafe {
match self.buf_ring.get_buf(result as u32, cqe.flags()) {
Ok(Some(buffer)) => buffer,
Ok(None) => {
::log::error!("Couldn't get recv buffer");
return;
}
Err(err) => {
::log::error!("Couldn't get recv buffer: {:#}", err);
return;
}
}
};
let buffer = buffer.as_slice();
let addr = match self.recv_helper.parse(buffer) {
Ok((request, addr)) => {
self.handle_request(request, addr);
addr
}
Err(self::recv_helper::Error::RequestParseError(err, addr)) => {
match err {
RequestParseError::Sendable {
connection_id,
transaction_id,
err,
} => {
::log::debug!("Couldn't parse request from {:?}: {}", addr, err);
if self.validator.connection_id_valid(addr, connection_id) {
let response = ErrorResponse {
transaction_id,
message: err.right_or("Parse error").into(),
};
self.local_responses.push_back((response.into(), addr));
}
}
RequestParseError::Unsendable { err } => {
::log::debug!("Couldn't parse request from {:?}: {}", addr, err);
}
}
addr
}
Err(self::recv_helper::Error::InvalidSocketAddress) => {
::log::debug!("Ignored request claiming to be from port 0");
return;
}
Err(self::recv_helper::Error::RecvMsgParseError) => {
::log::error!("RecvMsgOut::parse failed");
return;
}
};
if self.config.statistics.active() {
let (statistics, extra_bytes) = if addr.is_ipv4() {
(&self.shared_state.statistics_ipv4, EXTRA_PACKET_SIZE_IPV4)
} else {
(&self.shared_state.statistics_ipv6, EXTRA_PACKET_SIZE_IPV6)
};
statistics
.bytes_received
.fetch_add(buffer.len() + extra_bytes, Ordering::Relaxed);
statistics.requests_received.fetch_add(1, Ordering::Relaxed);
}
}
fn handle_request(&mut self, request: Request, src: CanonicalSocketAddr) {
let access_list_mode = self.config.access_list.mode;
match request {
Request::Connect(request) => {
let connection_id = self.validator.create_connection_id(src);
let response = Response::Connect(ConnectResponse {
connection_id,
transaction_id: request.transaction_id,
});
self.local_responses.push_back((response, src));
}
Request::Announce(request) => {
if self
.validator
.connection_id_valid(src, request.connection_id)
{
if self
.access_list_cache
.load()
.allows(access_list_mode, &request.info_hash.0)
{
let worker_index =
SwarmWorkerIndex::from_info_hash(&self.config, request.info_hash);
self.request_sender.try_send_to(
worker_index,
ConnectedRequest::Announce(request),
src,
);
} else {
let response = Response::Error(ErrorResponse {
transaction_id: request.transaction_id,
message: "Info hash not allowed".into(),
});
self.local_responses.push_back((response, src))
}
}
}
Request::Scrape(request) => {
if self
.validator
.connection_id_valid(src, request.connection_id)
{
let split_requests = self.pending_scrape_responses.prepare_split_requests(
&self.config,
request,
self.pending_scrape_valid_until,
);
for (swarm_worker_index, request) in split_requests {
self.request_sender.try_send_to(
swarm_worker_index,
ConnectedRequest::Scrape(request),
src,
);
}
}
}
}
}
fn get_next_swarm_response(&mut self) -> Option<(Response, CanonicalSocketAddr)> {
loop {
match self.response_receiver.try_recv() {
Ok((ConnectedResponse::AnnounceIpv4(response), addr)) => {
return Some((Response::AnnounceIpv4(response), addr));
}
Ok((ConnectedResponse::AnnounceIpv6(response), addr)) => {
return Some((Response::AnnounceIpv6(response), addr));
}
Ok((ConnectedResponse::Scrape(response), addr)) => {
if let Some(response) =
self.pending_scrape_responses.add_and_get_finished(response)
{
return Some((Response::Scrape(response), addr));
}
}
Err(_) => {
return None;
}
}
}
}
}
pub fn supported_on_current_kernel() -> anyhow::Result<()> {
let opcodes = [
// We can't probe for RecvMsgMulti, so we probe for SendZc, which was
// also introduced in Linux 6.0
io_uring::opcode::SendZc::CODE,
];
let ring = IoUring::new(1).with_context(|| "create ring")?;
let mut probe = Probe::new();
ring.submitter()
.register_probe(&mut probe)
.with_context(|| "register probe")?;
for opcode in opcodes {
if !probe.is_supported(opcode) {
return Err(anyhow::anyhow!(
"io_uring opcode {:b} not supported",
opcode
));
}
}
Ok(())
}