YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-03-31 17:55:36 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #68 from greatest-ape/2022-04-09

Browse source 
aquatic_http & aquatic_http_load_test improvements; use idiomatic feature names
This commit is contained in:
Joakim Frostegård 2022-04-10 23:46:08 +02:00 • committed by GitHub
commit 9f9d74e2d8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
21 changed files with 245 additions and 423 deletions
Download patch file Download diff file Expand all files Collapse all files

12
.github/workflows/cargo-build-and-test.yml vendored
View file

@ -15,6 +15,12 @@ jobs:
timeout-minutes: 20
steps:
- uses: actions/checkout@v2
- name: Install latest stable Rust
uses: actions-rs/toolchain@v1
with:
profile: minimal
toolchain: stable
override: true
- name: Install dependencies
run: sudo apt-get update -y && sudo apt-get install libhwloc-dev -y
- name: Build
@ -31,5 +37,11 @@ jobs:
timeout-minutes: 10
steps:
- uses: actions/checkout@v2
- name: Install latest stable Rust
uses: actions-rs/toolchain@v1
with:
profile: minimal
toolchain: stable
override: true
- name: Build
run: cargo build --verbose -p aquatic_udp

2
Cargo.lock generated
View file

@ -90,6 +90,7 @@ dependencies = [
"aquatic_toml_config",
"cfg-if",
"either",
"futures",
"futures-lite",
"futures-rustls",
"glommio",
@ -120,6 +121,7 @@ dependencies = [
"aquatic_http_protocol",
"aquatic_toml_config",
"futures-lite",
"futures-rustls",
"glommio",
"hashbrown 0.12.0",
"log",

6
aquatic_common/Cargo.toml
View file

@ -12,9 +12,7 @@ readme = "../README.md"
name = "aquatic_common"
[features]
with-glommio = ["glommio"]
with-hwloc = ["hwloc"]
rustls-config = ["rustls", "rustls-pemfile"]
rustls = ["dep:rustls", "rustls-pemfile"]
[dependencies]
aquatic_toml_config = { version = "0.2.0", path = "../aquatic_toml_config" }
@ -38,7 +36,5 @@ toml = "0.5"
# Optional
glommio = { version = "0.7", optional = true }
hwloc = { version = "0.5", optional = true }
# rustls-config
rustls = { version = "0.20", optional = true }
rustls-pemfile = { version = "0.3", optional = true }

16
aquatic_common/src/cpu_pinning.rs
View file

@ -16,7 +16,7 @@ impl Default for CpuPinningDirection {
}
}
#[cfg(feature = "with-glommio")]
#[cfg(feature = "glommio")]
#[derive(Clone, Copy, Debug, PartialEq, TomlConfig, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum HyperThreadMapping {
@ -25,7 +25,7 @@ pub enum HyperThreadMapping {
Split,
}
#[cfg(feature = "with-glommio")]
#[cfg(feature = "glommio")]
impl Default for HyperThreadMapping {
fn default() -> Self {
Self::System
@ -35,7 +35,7 @@ impl Default for HyperThreadMapping {
pub trait CpuPinningConfig {
fn active(&self) -> bool;
fn direction(&self) -> CpuPinningDirection;
#[cfg(feature = "with-glommio")]
#[cfg(feature = "glommio")]
fn hyperthread(&self) -> HyperThreadMapping;
fn core_offset(&self) -> usize;
}
@ -54,7 +54,7 @@ pub mod mod_name {
pub struct struct_name {
pub active: bool,
pub direction: CpuPinningDirection,
#[cfg(feature = "with-glommio")]
#[cfg(feature = "glommio")]
pub hyperthread: HyperThreadMapping,
pub core_offset: usize,
}
@ -64,7 +64,7 @@ pub mod mod_name {
Self {
active: false,
direction: cpu_pinning_direction,
#[cfg(feature = "with-glommio")]
#[cfg(feature = "glommio")]
hyperthread: Default::default(),
core_offset: 0,
}
@ -77,7 +77,7 @@ pub mod mod_name {
fn direction(&self) -> CpuPinningDirection {
self.direction
}
#[cfg(feature = "with-glommio")]
#[cfg(feature = "glommio")]
fn hyperthread(&self) -> HyperThreadMapping {
self.hyperthread
}
@ -119,7 +119,7 @@ impl WorkerIndex {
}
}
#[cfg(feature = "with-glommio")]
#[cfg(feature = "glommio")]
pub mod glommio {
use ::glommio::{CpuSet, Placement};
@ -271,7 +271,7 @@ pub mod glommio {
/// Pin current thread to a suitable core
///
/// Requires hwloc (`apt-get install libhwloc-dev`)
#[cfg(feature = "with-hwloc")]
#[cfg(feature = "hwloc")]
pub fn pin_current_if_configured_to<C: CpuPinningConfig>(
config: &C,
socket_workers: usize,

2
aquatic_common/src/lib.rs
View file

@ -10,7 +10,7 @@ pub mod access_list;
pub mod cli;
pub mod cpu_pinning;
pub mod privileges;
#[cfg(feature = "rustls-config")]
#[cfg(feature = "rustls")]
pub mod rustls_config;
/// Amortized IndexMap using AHash hasher

3
aquatic_http/Cargo.toml
View file

@ -16,13 +16,14 @@ name = "aquatic_http"
name = "aquatic_http"
[dependencies]
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["rustls-config", "with-glommio"] }
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["rustls", "glommio"] }
aquatic_http_protocol = { version = "0.2.0", path = "../aquatic_http_protocol" }
aquatic_toml_config = { version = "0.2.0", path = "../aquatic_toml_config" }
anyhow = "1"
cfg-if = "1"
either = "1"
futures = "0.3"
futures-lite = "1"
futures-rustls = "0.22"
glommio = "0.7"

36
aquatic_http/src/common.rs
View file

@ -9,6 +9,7 @@ use aquatic_http_protocol::{
request::{AnnounceRequest, ScrapeRequest},
response::{AnnounceResponse, ScrapeResponse},
};
use glommio::channels::shared_channel::SharedSender;
#[derive(Copy, Clone, Debug)]
pub struct ConsumerId(pub usize);
@ -21,46 +22,15 @@ pub enum ChannelRequest {
Announce {
request: AnnounceRequest,
peer_addr: CanonicalSocketAddr,
connection_id: ConnectionId,
response_consumer_id: ConsumerId,
response_sender: SharedSender<AnnounceResponse>,
},
Scrape {
request: ScrapeRequest,
peer_addr: CanonicalSocketAddr,
connection_id: ConnectionId,
response_consumer_id: ConsumerId,
response_sender: SharedSender<ScrapeResponse>,
},
}
#[derive(Debug)]
pub enum ChannelResponse {
Announce {
response: AnnounceResponse,
peer_addr: CanonicalSocketAddr,
connection_id: ConnectionId,
},
Scrape {
response: ScrapeResponse,
peer_addr: CanonicalSocketAddr,
connection_id: ConnectionId,
},
}
impl ChannelResponse {
pub fn get_connection_id(&self) -> ConnectionId {
match self {
Self::Announce { connection_id, .. } => *connection_id,
Self::Scrape { connection_id, .. } => *connection_id,
}
}
pub fn get_peer_addr(&self) -> CanonicalSocketAddr {
match self {
Self::Announce { peer_addr, .. } => *peer_addr,
Self::Scrape { peer_addr, .. } => *peer_addr,
}
}
}
#[derive(Default, Clone)]
pub struct State {
pub access_list: Arc<AccessListArcSwap>,

14
aquatic_http/src/lib.rs
View file

@ -37,7 +37,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
let num_peers = config.socket_workers + config.request_workers;
let request_mesh_builder = MeshBuilder::partial(num_peers, SHARED_CHANNEL_SIZE);
let response_mesh_builder = MeshBuilder::partial(num_peers, SHARED_CHANNEL_SIZE);
let (sentinel_watcher, sentinel) = PanicSentinelWatcher::create_with_sentinel();
let priv_dropper = PrivilegeDropper::new(config.privileges.clone(), config.socket_workers);
@ -55,7 +54,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
let state = state.clone();
let tls_config = tls_config.clone();
let request_mesh_builder = request_mesh_builder.clone();
let response_mesh_builder = response_mesh_builder.clone();
let priv_dropper = priv_dropper.clone();
let placement = get_worker_placement(
@ -74,7 +72,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
state,
tls_config,
request_mesh_builder,
response_mesh_builder,
priv_dropper,
)
.await
@ -89,7 +86,6 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
let config = config.clone();
let state = state.clone();
let request_mesh_builder = request_mesh_builder.clone();
let response_mesh_builder = response_mesh_builder.clone();
let placement = get_worker_placement(
&config.cpu_pinning,
@ -101,14 +97,8 @@ pub fn run(config: Config) -> ::anyhow::Result<()> {
let executor = builder
.spawn(move || async move {
workers::request::run_request_worker(
sentinel,
config,
state,
request_mesh_builder,
response_mesh_builder,
)
.await
workers::request::run_request_worker(sentinel, config, state, request_mesh_builder)
.await
})
.map_err(|err| anyhow::anyhow!("Spawning executor failed: {:#}", err))?;

136
aquatic_http/src/workers/request.rs
View file

@ -8,7 +8,7 @@ use std::time::Instant;
use either::Either;
use futures_lite::{Stream, StreamExt};
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders};
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role};
use glommio::timer::TimerActionRepeat;
use glommio::{enclose, prelude::*};
use rand::prelude::SmallRng;
@ -33,23 +33,6 @@ pub trait Ip: ::std::fmt::Debug + Copy + Eq + ::std::hash::Hash {}
impl Ip for Ipv4Addr {}
impl Ip for Ipv6Addr {}
#[derive(Clone, Copy, Debug)]
pub struct ConnectionMeta {
/// Index of socket worker responsible for this connection. Required for
/// sending back response through correct channel to correct worker.
pub response_consumer_id: ConsumerId,
pub peer_addr: CanonicalSocketAddr,
/// Connection id local to socket worker
pub connection_id: ConnectionId,
}
#[derive(Clone, Copy, Debug)]
pub struct PeerConnectionMeta<I: Ip> {
pub response_consumer_id: ConsumerId,
pub connection_id: ConnectionId,
pub peer_ip_address: I,
}
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum PeerStatus {
Seeding,
@ -75,7 +58,7 @@ impl PeerStatus {
#[derive(Debug, Clone, Copy)]
pub struct Peer<I: Ip> {
pub connection_meta: PeerConnectionMeta<I>,
pub ip_address: I,
pub port: u16,
pub status: PeerStatus,
pub valid_until: ValidUntil,
@ -84,7 +67,7 @@ pub struct Peer<I: Ip> {
impl<I: Ip> Peer<I> {
pub fn to_response_peer(&self) -> ResponsePeer<I> {
ResponsePeer {
ip_address: self.connection_meta.peer_ip_address,
ip_address: self.ip_address,
port: self.port,
}
}
@ -179,12 +162,8 @@ pub async fn run_request_worker(
config: Config,
state: State,
request_mesh_builder: MeshBuilder<ChannelRequest, Partial>,
response_mesh_builder: MeshBuilder<ChannelResponse, Partial>,
) {
let (_, mut request_receivers) = request_mesh_builder.join(Role::Consumer).await.unwrap();
let (response_senders, _) = response_mesh_builder.join(Role::Producer).await.unwrap();
let response_senders = Rc::new(response_senders);
let torrents = Rc::new(RefCell::new(TorrentMaps::default()));
let access_list = state.access_list;
@ -198,13 +177,25 @@ pub async fn run_request_worker(
})()
}));
let max_peer_age = config.cleaning.max_peer_age;
let peer_valid_until = Rc::new(RefCell::new(ValidUntil::new(max_peer_age)));
// Periodically update peer_valid_until
TimerActionRepeat::repeat(enclose!((peer_valid_until) move || {
enclose!((peer_valid_until) move || async move {
*peer_valid_until.borrow_mut() = ValidUntil::new(max_peer_age);
Some(Duration::from_secs(1))
})()
}));
let mut handles = Vec::new();
for (_, receiver) in request_receivers.streams() {
let handle = spawn_local(handle_request_stream(
config.clone(),
torrents.clone(),
response_senders.clone(),
peer_valid_until.clone(),
receiver,
))
.detach();
@ -220,87 +211,46 @@ pub async fn run_request_worker(
async fn handle_request_stream<S>(
config: Config,
torrents: Rc<RefCell<TorrentMaps>>,
response_senders: Rc<Senders<ChannelResponse>>,
peer_valid_until: Rc<RefCell<ValidUntil>>,
mut stream: S,
) where
S: Stream<Item = ChannelRequest> + ::std::marker::Unpin,
{
let mut rng = SmallRng::from_entropy();
let max_peer_age = config.cleaning.max_peer_age;
let peer_valid_until = Rc::new(RefCell::new(ValidUntil::new(max_peer_age)));
TimerActionRepeat::repeat(enclose!((peer_valid_until) move || {
enclose!((peer_valid_until) move || async move {
*peer_valid_until.borrow_mut() = ValidUntil::new(max_peer_age);
Some(Duration::from_secs(1))
})()
}));
while let Some(channel_request) = stream.next().await {
let (response, consumer_id) = match channel_request {
match channel_request {
ChannelRequest::Announce {
request,
peer_addr,
response_consumer_id,
connection_id,
response_sender,
} => {
let meta = ConnectionMeta {
response_consumer_id,
connection_id,
peer_addr,
};
let response = handle_announce_request(
&config,
&mut rng,
&mut torrents.borrow_mut(),
peer_valid_until.borrow().to_owned(),
meta,
peer_addr,
request,
);
let response = ChannelResponse::Announce {
response,
peer_addr,
connection_id,
};
(response, response_consumer_id)
if let Err(err) = response_sender.connect().await.send(response).await {
::log::error!("request worker could not send announce response: {:#}", err);
}
}
ChannelRequest::Scrape {
request,
peer_addr,
response_consumer_id,
connection_id,
response_sender,
} => {
let meta = ConnectionMeta {
response_consumer_id,
connection_id,
peer_addr,
};
let response =
handle_scrape_request(&config, &mut torrents.borrow_mut(), meta, request);
handle_scrape_request(&config, &mut torrents.borrow_mut(), peer_addr, request);
let response = ChannelResponse::Scrape {
response,
peer_addr,
connection_id,
};
(response, response_consumer_id)
if let Err(err) = response_sender.connect().await.send(response).await {
::log::error!("request worker could not send scrape response: {:#}", err);
}
}
};
::log::debug!("preparing to send response to channel: {:?}", response);
if let Err(err) = response_senders.try_send_to(consumer_id.0, response) {
::log::warn!("response_sender.try_send: {:?}", err);
}
yield_if_needed().await;
}
}
@ -309,24 +259,18 @@ pub fn handle_announce_request(
rng: &mut impl Rng,
torrent_maps: &mut TorrentMaps,
valid_until: ValidUntil,
meta: ConnectionMeta,
peer_addr: CanonicalSocketAddr,
request: AnnounceRequest,
) -> AnnounceResponse {
match meta.peer_addr.get().ip() {
match peer_addr.get().ip() {
IpAddr::V4(peer_ip_address) => {
let torrent_data: &mut TorrentData<Ipv4Addr> =
torrent_maps.ipv4.entry(request.info_hash).or_default();
let peer_connection_meta = PeerConnectionMeta {
response_consumer_id: meta.response_consumer_id,
connection_id: meta.connection_id,
peer_ip_address,
};
let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers(
config,
rng,
peer_connection_meta,
peer_ip_address,
torrent_data,
request,
valid_until,
@ -347,16 +291,10 @@ pub fn handle_announce_request(
let torrent_data: &mut TorrentData<Ipv6Addr> =
torrent_maps.ipv6.entry(request.info_hash).or_default();
let peer_connection_meta = PeerConnectionMeta {
response_consumer_id: meta.response_consumer_id,
connection_id: meta.connection_id,
peer_ip_address,
};
let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers(
config,
rng,
peer_connection_meta,
peer_ip_address,
torrent_data,
request,
valid_until,
@ -380,7 +318,7 @@ pub fn handle_announce_request(
pub fn upsert_peer_and_get_response_peers<I: Ip>(
config: &Config,
rng: &mut impl Rng,
request_sender_meta: PeerConnectionMeta<I>,
peer_ip_address: I,
torrent_data: &mut TorrentData<I>,
request: AnnounceRequest,
valid_until: ValidUntil,
@ -391,7 +329,7 @@ pub fn upsert_peer_and_get_response_peers<I: Ip>(
PeerStatus::from_event_and_bytes_left(request.event, Some(request.bytes_left));
let peer = Peer {
connection_meta: request_sender_meta,
ip_address: peer_ip_address,
port: request.port,
status: peer_status,
valid_until,
@ -402,7 +340,7 @@ pub fn upsert_peer_and_get_response_peers<I: Ip>(
let ip_or_key = request
.key
.map(Either::Right)
.unwrap_or_else(|| Either::Left(request_sender_meta.peer_ip_address));
.unwrap_or_else(|| Either::Left(peer_ip_address));
let peer_map_key = PeerMapKey {
peer_id: request.peer_id,
@ -462,7 +400,7 @@ pub fn upsert_peer_and_get_response_peers<I: Ip>(
pub fn handle_scrape_request(
config: &Config,
torrent_maps: &mut TorrentMaps,
meta: ConnectionMeta,
peer_addr: CanonicalSocketAddr,
request: ScrapeRequest,
) -> ScrapeResponse {
let num_to_take = request
@ -474,7 +412,7 @@ pub fn handle_scrape_request(
files: BTreeMap::new(),
};
let peer_ip = meta.peer_addr.get().ip();
let peer_ip = peer_addr.get().ip();
// If request.info_hashes is empty, don't return scrape for all
// torrents, even though reference server does it. It is too expensive.

135
aquatic_http/src/workers/socket.rs
View file

@ -16,12 +16,12 @@ use aquatic_http_protocol::response::{
FailureResponse, Response, ScrapeResponse, ScrapeStatistics,
};
use either::Either;
use futures::stream::FuturesUnordered;
use futures_lite::{AsyncReadExt, AsyncWriteExt, StreamExt};
use futures_rustls::server::TlsStream;
use futures_rustls::TlsAcceptor;
use glommio::channels::channel_mesh::{MeshBuilder, Partial, Role, Senders};
use glommio::channels::local_channel::{new_bounded, LocalReceiver, LocalSender};
use glommio::channels::shared_channel::ConnectedReceiver;
use glommio::channels::shared_channel::{self, SharedReceiver};
use glommio::net::{TcpListener, TcpStream};
use glommio::task::JoinHandle;
use glommio::timer::TimerActionRepeat;
@ -49,7 +49,6 @@ struct PendingScrapeResponse {
struct ConnectionReference {
task_handle: Option<JoinHandle<()>>,
response_sender: LocalSender<ChannelResponse>,
valid_until: ValidUntil,
}
@ -59,7 +58,6 @@ pub async fn run_socket_worker(
state: State,
tls_config: Arc<RustlsConfig>,
request_mesh_builder: MeshBuilder<ChannelRequest, Partial>,
response_mesh_builder: MeshBuilder<ChannelResponse, Partial>,
priv_dropper: PrivilegeDropper,
) {
let config = Rc::new(config);
@ -70,9 +68,6 @@ pub async fn run_socket_worker(
let (request_senders, _) = request_mesh_builder.join(Role::Producer).await.unwrap();
let request_senders = Rc::new(request_senders);
let (_, mut response_receivers) = response_mesh_builder.join(Role::Consumer).await.unwrap();
let response_consumer_id = ConsumerId(response_receivers.consumer_id().unwrap());
let connection_slab = Rc::new(RefCell::new(Slab::new()));
TimerActionRepeat::repeat(enclose!((config, connection_slab) move || {
@ -82,24 +77,13 @@ pub async fn run_socket_worker(
)
}));
for (_, response_receiver) in response_receivers.streams() {
spawn_local(receive_responses(
response_receiver,
connection_slab.clone(),
))
.detach();
}
let mut incoming = listener.incoming();
while let Some(stream) = incoming.next().await {
match stream {
Ok(stream) => {
let (response_sender, response_receiver) = new_bounded(config.request_workers);
let key = connection_slab.borrow_mut().insert(ConnectionReference {
task_handle: None,
response_sender,
valid_until: ValidUntil::new(config.cleaning.max_connection_idle),
});
@ -108,8 +92,6 @@ pub async fn run_socket_worker(
config,
access_list,
request_senders,
response_receiver,
response_consumer_id,
ConnectionId(key),
tls_config,
connection_slab.clone(),
@ -158,28 +140,10 @@ async fn clean_connections(
))
}
async fn receive_responses(
mut response_receiver: ConnectedReceiver<ChannelResponse>,
connection_references: Rc<RefCell<Slab<ConnectionReference>>>,
) {
while let Some(channel_response) = response_receiver.next().await {
if let Some(reference) = connection_references
.borrow()
.get(channel_response.get_connection_id().0)
{
if let Err(err) = reference.response_sender.try_send(channel_response) {
::log::error!("Couldn't send response to local receiver: {:?}", err);
}
}
}
}
struct Connection {
config: Rc<Config>,
access_list_cache: AccessListCache,
request_senders: Rc<Senders<ChannelRequest>>,
response_receiver: LocalReceiver<ChannelResponse>,
response_consumer_id: ConsumerId,
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
stream: TlsStream<TcpStream>,
peer_addr: CanonicalSocketAddr,
@ -194,8 +158,6 @@ impl Connection {
config: Rc<Config>,
access_list: Arc<AccessListArcSwap>,
request_senders: Rc<Senders<ChannelRequest>>,
response_receiver: LocalReceiver<ChannelResponse>,
response_consumer_id: ConsumerId,
connection_id: ConnectionId,
tls_config: Arc<RustlsConfig>,
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
@ -217,8 +179,6 @@ impl Connection {
config: config.clone(),
access_list_cache: create_access_list_cache(&access_list),
request_senders: request_senders.clone(),
response_receiver,
response_consumer_id,
connection_slab,
stream,
peer_addr,
@ -293,8 +253,8 @@ impl Connection {
}
Err(RequestParseError::NeedMoreData) => {
::log::debug!(
"need more request data. current data: {:?}",
std::str::from_utf8(&self.request_buffer[..self.request_buffer_position])
"need more request data. current data: {}",
&self.request_buffer[..self.request_buffer_position].escape_ascii()
);
}
}
@ -324,11 +284,12 @@ impl Connection {
.load()
.allows(self.config.access_list.mode, &info_hash.0)
{
let (response_sender, response_receiver) = shared_channel::new_bounded(1);
let request = ChannelRequest::Announce {
request,
connection_id: self.connection_id,
response_consumer_id: self.response_consumer_id,
peer_addr: self.peer_addr,
response_sender,
};
let consumer_index = calculate_request_consumer_index(&self.config, info_hash);
@ -339,7 +300,13 @@ impl Connection {
.await
.unwrap();
self.wait_for_response(None).await
response_receiver
.connect()
.await
.recv()
.await
.ok_or_else(|| anyhow::anyhow!("response sender closed"))
.map(Response::Announce)
} else {
let response = Response::Failure(FailureResponse {
failure_reason: "Info hash not allowed".into(),
@ -360,13 +327,17 @@ impl Connection {
}
let pending_worker_responses = info_hashes_by_worker.len();
let mut response_receivers = Vec::with_capacity(pending_worker_responses);
for (consumer_index, info_hashes) in info_hashes_by_worker {
let (response_sender, response_receiver) = shared_channel::new_bounded(1);
response_receivers.push(response_receiver);
let request = ChannelRequest::Scrape {
request: ScrapeRequest { info_hashes },
peer_addr: self.peer_addr,
response_consumer_id: self.response_consumer_id,
connection_id: self.connection_id,
response_sender,
};
// Only fails when receiver is closed
@ -381,53 +352,47 @@ impl Connection {
stats: Default::default(),
};
self.wait_for_response(Some(pending_scrape_response)).await
self.wait_for_scrape_responses(response_receivers, pending_scrape_response)
.await
}
}
}
/// Wait for announce response or partial scrape responses to arrive,
/// Wait for partial scrape responses to arrive,
/// return full response
async fn wait_for_response(
async fn wait_for_scrape_responses(
&self,
mut opt_pending_scrape_response: Option<PendingScrapeResponse>,
response_receivers: Vec<SharedReceiver<ScrapeResponse>>,
mut pending: PendingScrapeResponse,
) -> anyhow::Result<Response> {
let mut responses = response_receivers
.into_iter()
.map(|receiver| async { receiver.connect().await.recv().await })
.collect::<FuturesUnordered<_>>();
loop {
let channel_response = self
.response_receiver
.recv()
let response = responses
.next()
.await
.expect("wait_for_response: can't receive response, sender is closed");
.ok_or_else(|| {
anyhow::anyhow!("stream ended before all partial scrape responses received")
})?
.ok_or_else(|| {
anyhow::anyhow!(
"wait_for_scrape_response: can't receive response, sender is closed"
)
})?;
if channel_response.get_peer_addr() != self.peer_addr {
return Err(anyhow::anyhow!("peer addresses didn't match"));
pending.stats.extend(response.files);
pending.pending_worker_responses -= 1;
if pending.pending_worker_responses == 0 {
let response = Response::Scrape(ScrapeResponse {
files: pending.stats,
});
break Ok(response);
}
match channel_response {
ChannelResponse::Announce { response, .. } => {
break Ok(Response::Announce(response));
}
ChannelResponse::Scrape { response, .. } => {
if let Some(mut pending) = opt_pending_scrape_response.take() {
pending.stats.extend(response.files);
pending.pending_worker_responses -= 1;
if pending.pending_worker_responses == 0 {
let response = Response::Scrape(ScrapeResponse {
files: pending.stats,
});
break Ok(response);
} else {
opt_pending_scrape_response = Some(pending);
}
} else {
return Err(anyhow::anyhow!(
"received channel scrape response without pending scrape response"
));
}
}
};
}
}

3
aquatic_http_load_test/Cargo.toml
View file

@ -13,12 +13,13 @@ readme = "../README.md"
name = "aquatic_http_load_test"
[dependencies]
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["with-glommio"] }
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["glommio"] }
aquatic_http_protocol = { version = "0.2.0", path = "../aquatic_http_protocol" }
aquatic_toml_config = { version = "0.2.0", path = "../aquatic_toml_config" }
anyhow = "1"
futures-lite = "1"
futures-rustls = "0.22"
hashbrown = "0.12"
glommio = "0.7"
log = "0.4"

2
aquatic_http_load_test/src/config.rs
View file

@ -22,6 +22,7 @@ pub struct Config {
/// Announce/scrape url suffix. Use `/my_token/` to get `/announce/my_token/`
pub url_suffix: String,
pub duration: usize,
pub keep_alive: bool,
pub torrents: TorrentConfig,
pub cpu_pinning: CpuPinningConfigDesc,
}
@ -60,6 +61,7 @@ impl Default for Config {
connection_creation_interval_ms: 10,
url_suffix: "".into(),
duration: 0,
keep_alive: true,
torrents: TorrentConfig::default(),
cpu_pinning: Default::default(),
}

276
aquatic_http_load_test/src/network.rs
View file

@ -1,7 +1,7 @@
use std::{
cell::RefCell,
convert::TryInto,
io::{Cursor, ErrorKind, Read},
io::Cursor,
rc::Rc,
sync::{atomic::Ordering, Arc},
time::Duration,
@ -9,10 +9,10 @@ use std::{
use aquatic_http_protocol::response::Response;
use futures_lite::{AsyncReadExt, AsyncWriteExt};
use futures_rustls::{client::TlsStream, TlsConnector};
use glommio::net::TcpStream;
use glommio::{prelude::*, timer::TimerActionRepeat};
use rand::{prelude::SmallRng, SeedableRng};
use rustls::ClientConnection;
use crate::{common::LoadTestState, config::Config, utils::create_random_request};
@ -23,6 +23,7 @@ pub async fn run_socket_thread(
) -> anyhow::Result<()> {
let config = Rc::new(config);
let num_active_connections = Rc::new(RefCell::new(0usize));
let rng = Rc::new(RefCell::new(SmallRng::from_entropy()));
let interval = config.connection_creation_interval_ms;
@ -34,6 +35,7 @@ pub async fn run_socket_thread(
tls_config.clone(),
load_test_state.clone(),
num_active_connections.clone(),
rng.clone(),
)
.await
{
@ -51,6 +53,7 @@ pub async fn run_socket_thread(
tls_config.clone(),
load_test_state.clone(),
num_active_connections.clone(),
rng.clone(),
)
});
}
@ -66,11 +69,18 @@ async fn periodically_open_connections(
tls_config: Arc<rustls::ClientConfig>,
load_test_state: LoadTestState,
num_active_connections: Rc<RefCell<usize>>,
rng: Rc<RefCell<SmallRng>>,
) -> Option<Duration> {
if *num_active_connections.borrow() < config.num_connections {
spawn_local(async move {
if let Err(err) =
Connection::run(config, tls_config, load_test_state, num_active_connections).await
if let Err(err) = Connection::run(
config,
tls_config,
load_test_state,
num_active_connections,
rng.clone(),
)
.await
{
::log::error!("connection creation error: {:?}", err);
}
@ -84,13 +94,9 @@ async fn periodically_open_connections(
struct Connection {
config: Rc<Config>,
load_test_state: LoadTestState,
rng: SmallRng,
stream: TcpStream,
tls: ClientConnection,
response_buffer: [u8; 2048],
response_buffer_position: usize,
send_new_request: bool,
queued_responses: usize,
rng: Rc<RefCell<SmallRng>>,
stream: TlsStream<TcpStream>,
buffer: [u8; 2048],
}
impl Connection {
@ -99,23 +105,22 @@ impl Connection {
tls_config: Arc<rustls::ClientConfig>,
load_test_state: LoadTestState,
num_active_connections: Rc<RefCell<usize>>,
rng: Rc<RefCell<SmallRng>>,
) -> anyhow::Result<()> {
let stream = TcpStream::connect(config.server_address)
.await
.map_err(|err| anyhow::anyhow!("connect: {:?}", err))?;
let tls = ClientConnection::new(tls_config, "example.com".try_into().unwrap()).unwrap();
let rng = SmallRng::from_entropy();
let stream = TlsConnector::from(tls_config)
.connect("example.com".try_into().unwrap(), stream)
.await?;
let mut connection = Connection {
config,
load_test_state,
rng,
stream,
tls,
response_buffer: [0; 2048],
response_buffer_position: 0,
send_new_request: true,
queued_responses: 0,
buffer: [0; 2048],
};
*num_active_connections.borrow_mut() += 1;
@ -131,169 +136,112 @@ impl Connection {
async fn run_connection_loop(&mut self) -> anyhow::Result<()> {
loop {
if self.send_new_request {
let request =
create_random_request(&self.config, &self.load_test_state, &mut self.rng);
self.send_request().await?;
self.read_response().await?;
request.write(&mut self.tls.writer(), self.config.url_suffix.as_bytes())?;
self.queued_responses += 1;
self.send_new_request = false;
if !self.config.keep_alive {
break Ok(());
}
self.write_tls().await?;
self.read_tls().await?;
}
}
async fn read_tls(&mut self) -> anyhow::Result<()> {
loop {
let mut buf = [0u8; 1024];
async fn send_request(&mut self) -> anyhow::Result<()> {
let request = create_random_request(
&self.config,
&self.load_test_state,
&mut self.rng.borrow_mut(),
);
let bytes_read = self.stream.read(&mut buf).await?;
let mut cursor = Cursor::new(&mut self.buffer[..]);
if bytes_read == 0 {
return Err(anyhow::anyhow!("Peer has closed connection"));
}
request.write(&mut cursor, self.config.url_suffix.as_bytes())?;
self.load_test_state
.statistics
.bytes_received
.fetch_add(bytes_read, Ordering::SeqCst);
let cursor_position = cursor.position() as usize;
let _ = self.tls.read_tls(&mut &buf[..bytes_read]).unwrap();
let bytes_sent = self
.stream
.write(&cursor.into_inner()[..cursor_position])
.await?;
let io_state = self.tls.process_new_packets()?;
let mut added_plaintext = false;
if io_state.plaintext_bytes_to_read() != 0 {
loop {
match self.tls.reader().read(&mut buf) {
Ok(0) => {
break;
}
Ok(amt) => {
let end = self.response_buffer_position + amt;
if end > self.response_buffer.len() {
return Err(anyhow::anyhow!("response too large"));
} else {
let response_buffer_slice =
&mut self.response_buffer[self.response_buffer_position..end];
response_buffer_slice.copy_from_slice(&buf[..amt]);
self.response_buffer_position = end;
added_plaintext = true;
}
}
Err(err) if err.kind() == ErrorKind::WouldBlock => {
break;
}
Err(err) => {
panic!("tls.reader().read: {}", err);
}
}
}
}
if added_plaintext {
let interesting_bytes = &self.response_buffer[..self.response_buffer_position];
let mut opt_body_start_index = None;
for (i, chunk) in interesting_bytes.windows(4).enumerate() {
if chunk == b"\r\n\r\n" {
opt_body_start_index = Some(i + 4);
break;
}
}
if let Some(body_start_index) = opt_body_start_index {
match Response::from_bytes(&interesting_bytes[body_start_index..]) {
Ok(response) => {
match response {
Response::Announce(_) => {
self.load_test_state
.statistics
.responses_announce
.fetch_add(1, Ordering::SeqCst);
}
Response::Scrape(_) => {
self.load_test_state
.statistics
.responses_scrape
.fetch_add(1, Ordering::SeqCst);
}
Response::Failure(response) => {
self.load_test_state
.statistics
.responses_failure
.fetch_add(1, Ordering::SeqCst);
println!(
"failure response: reason: {}",
response.failure_reason
);
}
}
self.response_buffer_position = 0;
self.send_new_request = true;
break;
}
Err(err) => {
eprintln!(
"deserialize response error with {} bytes read: {:?}, text: {}",
self.response_buffer_position,
err,
String::from_utf8_lossy(interesting_bytes)
);
}
}
}
}
if self.tls.wants_write() {
break;
}
}
Ok(())
}
async fn write_tls(&mut self) -> anyhow::Result<()> {
if !self.tls.wants_write() {
return Ok(());
}
let mut buf = Vec::new();
let mut buf = Cursor::new(&mut buf);
while self.tls.wants_write() {
self.tls.write_tls(&mut buf).unwrap();
}
let len = buf.get_ref().len();
self.stream.write_all(&buf.into_inner()).await?;
self.stream.flush().await?;
self.load_test_state
.statistics
.bytes_sent
.fetch_add(len, Ordering::SeqCst);
.fetch_add(bytes_sent, Ordering::Relaxed);
if self.queued_responses != 0 {
self.load_test_state
.statistics
.requests
.fetch_add(self.queued_responses, Ordering::SeqCst);
self.load_test_state
.statistics
.requests
.fetch_add(1, Ordering::Relaxed);
self.queued_responses = 0;
Ok(())
}
async fn read_response(&mut self) -> anyhow::Result<()> {
let mut buffer_position = 0;
loop {
let bytes_read = self
.stream
.read(&mut self.buffer[buffer_position..])
.await?;
if bytes_read == 0 {
break;
}
buffer_position += bytes_read;
let interesting_bytes = &self.buffer[..buffer_position];
let mut opt_body_start_index = None;
for (i, chunk) in interesting_bytes.windows(4).enumerate() {
if chunk == b"\r\n\r\n" {
opt_body_start_index = Some(i + 4);
break;
}
}
if let Some(body_start_index) = opt_body_start_index {
match Response::from_bytes(&interesting_bytes[body_start_index..]) {
Ok(response) => {
match response {
Response::Announce(_) => {
self.load_test_state
.statistics
.responses_announce
.fetch_add(1, Ordering::Relaxed);
}
Response::Scrape(_) => {
self.load_test_state
.statistics
.responses_scrape
.fetch_add(1, Ordering::Relaxed);
}
Response::Failure(response) => {
self.load_test_state
.statistics
.responses_failure
.fetch_add(1, Ordering::Relaxed);
println!("failure response: reason: {}", response.failure_reason);
}
}
break;
}
Err(err) => {
eprintln!(
"deserialize response error with {} bytes read: {:?}, text: {}",
buffer_position,
err,
interesting_bytes.escape_ascii()
);
}
}
}
}
Ok(())

2
aquatic_http_load_test/src/utils.rs
View file

@ -10,7 +10,7 @@ use crate::config::*;
pub fn create_random_request(
config: &Config,
state: &LoadTestState,
rng: &mut impl Rng,
rng: &mut SmallRng,
) -> Request {
let weights = [
config.torrents.weight_announce as u32,

4
aquatic_http_private/Cargo.toml
View file

@ -14,8 +14,8 @@ name = "aquatic_http_private"
name = "aquatic_http_private"
[dependencies]
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["rustls-config"] }
aquatic_http_protocol = { version = "0.2.0", path = "../aquatic_http_protocol", features = ["with-axum"] }
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["rustls"] }
aquatic_http_protocol = { version = "0.2.0", path = "../aquatic_http_protocol", features = ["axum"] }
aquatic_toml_config = { version = "0.2.0", path = "../aquatic_toml_config" }
anyhow = "1"

3
aquatic_http_protocol/Cargo.toml
View file

@ -22,9 +22,6 @@ name = "bench_announce_response_to_bytes"
path = "benches/bench_announce_response_to_bytes.rs"
harness = false
[features]
with-axum = ["axum"]
[dependencies]
anyhow = "1"
axum = { version = "0.5", optional = true, default-features = false }

8
aquatic_http_protocol/src/response.rs
View file

@ -117,7 +117,7 @@ impl AnnounceResponse {
}
}
#[cfg(feature = "with-axum")]
#[cfg(feature = "axum")]
impl axum::response::IntoResponse for AnnounceResponse {
fn into_response(self) -> axum::response::Response {
let mut response_bytes = Vec::with_capacity(128);
@ -158,7 +158,7 @@ impl ScrapeResponse {
}
}
#[cfg(feature = "with-axum")]
#[cfg(feature = "axum")]
impl axum::response::IntoResponse for ScrapeResponse {
fn into_response(self) -> axum::response::Response {
let mut response_bytes = Vec::with_capacity(128);
@ -197,7 +197,7 @@ impl FailureResponse {
}
}
#[cfg(feature = "with-axum")]
#[cfg(feature = "axum")]
impl axum::response::IntoResponse for FailureResponse {
fn into_response(self) -> axum::response::Response {
let mut response_bytes = Vec::with_capacity(64);
@ -229,7 +229,7 @@ impl Response {
}
}
#[cfg(feature = "with-axum")]
#[cfg(feature = "axum")]
impl axum::response::IntoResponse for Response {
fn into_response(self) -> axum::response::Response {
match self {

2
aquatic_udp/Cargo.toml
View file

@ -16,7 +16,7 @@ name = "aquatic_udp"
name = "aquatic_udp"
[features]
cpu-pinning = ["aquatic_common/with-hwloc"]
cpu-pinning = ["aquatic_common/hwloc"]
[dependencies]
aquatic_common = { version = "0.2.0", path = "../aquatic_common" }

2
aquatic_udp_load_test/Cargo.toml
View file

@ -10,7 +10,7 @@ keywords = ["udp", "benchmark", "peer-to-peer", "torrent", "bittorrent"]
readme = "../README.md"
[features]
cpu-pinning = ["aquatic_common/with-hwloc"]
cpu-pinning = ["aquatic_common/hwloc"]
[[bin]]
name = "aquatic_udp_load_test"

2
aquatic_ws/Cargo.toml
View file

@ -16,7 +16,7 @@ name = "aquatic_ws"
name = "aquatic_ws"
[dependencies]
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["rustls-config", "with-glommio"] }
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["rustls", "glommio"] }
aquatic_toml_config = { version = "0.2.0", path = "../aquatic_toml_config" }
aquatic_ws_protocol = { version = "0.2.0", path = "../aquatic_ws_protocol" }

2
aquatic_ws_load_test/Cargo.toml
View file

@ -13,7 +13,7 @@ readme = "../README.md"
name = "aquatic_ws_load_test"
[dependencies]
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["with-glommio"] }
aquatic_common = { version = "0.2.0", path = "../aquatic_common", features = ["glommio"] }
aquatic_toml_config = { version = "0.2.0", path = "../aquatic_toml_config" }
aquatic_ws_protocol = { version = "0.2.0", path = "../aquatic_ws_protocol" }