mirror of
https://github.com/YGGverse/aquatic.git
synced 2026-04-02 10:45:30 +00:00
Run rustfmt, clean up aquatic_http_protocol/Cargo.toml
This commit is contained in:
Joakim Frostegård
parent
0cc312a78d
commit
d0e716f80b
65 changed files with 1754 additions and 2590 deletions
|
|
@ -1,26 +1,24 @@
|
|||
use std::io::ErrorKind;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
use std::io::ErrorKind;
|
||||
|
||||
use hashbrown::HashMap;
|
||||
use mio::{net::TcpStream, Events, Poll, Interest, Token};
|
||||
use rand::{rngs::SmallRng, prelude::*};
|
||||
use tungstenite::{WebSocket, HandshakeError, ClientHandshake, handshake::MidHandshake};
|
||||
use mio::{net::TcpStream, Events, Interest, Poll, Token};
|
||||
use rand::{prelude::*, rngs::SmallRng};
|
||||
use tungstenite::{handshake::MidHandshake, ClientHandshake, HandshakeError, WebSocket};
|
||||
|
||||
use crate::common::*;
|
||||
use crate::config::*;
|
||||
use crate::utils::create_random_request;
|
||||
|
||||
|
||||
// Allow large enum variant WebSocket because it should be very common
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
pub enum ConnectionState {
|
||||
TcpStream(TcpStream),
|
||||
WebSocket(WebSocket<TcpStream>),
|
||||
MidHandshake(MidHandshake<ClientHandshake<TcpStream>>)
|
||||
MidHandshake(MidHandshake<ClientHandshake<TcpStream>>),
|
||||
}
|
||||
|
||||
|
||||
impl ConnectionState {
|
||||
fn advance(self, config: &Config) -> Option<Self> {
|
||||
match self {
|
||||
|
|
@ -31,33 +29,27 @@ impl ConnectionState {
|
|||
config.server_address.port()
|
||||
);
|
||||
|
||||
match ::tungstenite::client(req, stream){
|
||||
Ok((ws, _)) => {
|
||||
Some(ConnectionState::WebSocket(ws))
|
||||
},
|
||||
match ::tungstenite::client(req, stream) {
|
||||
Ok((ws, _)) => Some(ConnectionState::WebSocket(ws)),
|
||||
Err(HandshakeError::Interrupted(handshake)) => {
|
||||
Some(ConnectionState::MidHandshake(handshake))
|
||||
},
|
||||
}
|
||||
Err(HandshakeError::Failure(err)) => {
|
||||
eprintln!("handshake error: {:?}", err);
|
||||
|
||||
None
|
||||
}
|
||||
}
|
||||
},
|
||||
Self::MidHandshake(handshake) => {
|
||||
match handshake.handshake() {
|
||||
Ok((ws, _)) => {
|
||||
Some(ConnectionState::WebSocket(ws))
|
||||
},
|
||||
Err(HandshakeError::Interrupted(handshake)) => {
|
||||
Some(ConnectionState::MidHandshake(handshake))
|
||||
},
|
||||
Err(HandshakeError::Failure(err)) => {
|
||||
eprintln!("handshake error: {:?}", err);
|
||||
}
|
||||
Self::MidHandshake(handshake) => match handshake.handshake() {
|
||||
Ok((ws, _)) => Some(ConnectionState::WebSocket(ws)),
|
||||
Err(HandshakeError::Interrupted(handshake)) => {
|
||||
Some(ConnectionState::MidHandshake(handshake))
|
||||
}
|
||||
Err(HandshakeError::Failure(err)) => {
|
||||
eprintln!("handshake error: {:?}", err);
|
||||
|
||||
None
|
||||
}
|
||||
None
|
||||
}
|
||||
},
|
||||
Self::WebSocket(ws) => Some(Self::WebSocket(ws)),
|
||||
|
|
@ -65,7 +57,6 @@ impl ConnectionState {
|
|||
}
|
||||
}
|
||||
|
||||
|
||||
pub struct Connection {
|
||||
stream: ConnectionState,
|
||||
peer_id: PeerId,
|
||||
|
|
@ -73,7 +64,6 @@ pub struct Connection {
|
|||
send_answer: Option<(PeerId, OfferId)>,
|
||||
}
|
||||
|
||||
|
||||
impl Connection {
|
||||
pub fn create_and_register(
|
||||
config: &Config,
|
||||
|
|
@ -83,27 +73,31 @@ impl Connection {
|
|||
token_counter: &mut usize,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut stream = TcpStream::connect(config.server_address)?;
|
||||
|
||||
|
||||
poll.registry()
|
||||
.register(&mut stream, Token(*token_counter), Interest::READABLE | Interest::WRITABLE)
|
||||
.register(
|
||||
&mut stream,
|
||||
Token(*token_counter),
|
||||
Interest::READABLE | Interest::WRITABLE,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let connection = Connection {
|
||||
let connection = Connection {
|
||||
stream: ConnectionState::TcpStream(stream),
|
||||
peer_id: PeerId(rng.gen()),
|
||||
can_send: false,
|
||||
send_answer: None,
|
||||
};
|
||||
|
||||
|
||||
connections.insert(*token_counter, connection);
|
||||
|
||||
*token_counter += 1;
|
||||
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn advance(self, config: &Config) -> Option<Self> {
|
||||
if let Some(stream) = self.stream.advance(config){
|
||||
if let Some(stream) = self.stream.advance(config) {
|
||||
let can_send = matches!(stream, ConnectionState::WebSocket(_));
|
||||
|
||||
Some(Self {
|
||||
|
|
@ -117,52 +111,53 @@ impl Connection {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn read_responses(
|
||||
&mut self,
|
||||
state: &LoadTestState,
|
||||
) -> bool { // bool = drop connection
|
||||
pub fn read_responses(&mut self, state: &LoadTestState) -> bool {
|
||||
// bool = drop connection
|
||||
if let ConnectionState::WebSocket(ref mut ws) = self.stream {
|
||||
loop {
|
||||
match ws.read_message(){
|
||||
Ok(message) => {
|
||||
match OutMessage::from_ws_message(message){
|
||||
Ok(OutMessage::Offer(offer)) => {
|
||||
state.statistics.responses_offer
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
self.send_answer = Some((
|
||||
offer.peer_id,
|
||||
offer.offer_id
|
||||
));
|
||||
match ws.read_message() {
|
||||
Ok(message) => match OutMessage::from_ws_message(message) {
|
||||
Ok(OutMessage::Offer(offer)) => {
|
||||
state
|
||||
.statistics
|
||||
.responses_offer
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
self.can_send = true;
|
||||
},
|
||||
Ok(OutMessage::Answer(_)) => {
|
||||
state.statistics.responses_answer
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
self.send_answer = Some((offer.peer_id, offer.offer_id));
|
||||
|
||||
self.can_send = true;
|
||||
},
|
||||
Ok(OutMessage::AnnounceResponse(_)) => {
|
||||
state.statistics.responses_announce
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
self.can_send = true;
|
||||
}
|
||||
Ok(OutMessage::Answer(_)) => {
|
||||
state
|
||||
.statistics
|
||||
.responses_answer
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
self.can_send = true;
|
||||
},
|
||||
Ok(OutMessage::ScrapeResponse(_)) => {
|
||||
state.statistics.responses_scrape
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
self.can_send = true;
|
||||
}
|
||||
Ok(OutMessage::AnnounceResponse(_)) => {
|
||||
state
|
||||
.statistics
|
||||
.responses_announce
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
self.can_send = true;
|
||||
},
|
||||
Err(err) => {
|
||||
eprintln!("error deserializing offer: {:?}", err);
|
||||
}
|
||||
self.can_send = true;
|
||||
}
|
||||
Ok(OutMessage::ScrapeResponse(_)) => {
|
||||
state
|
||||
.statistics
|
||||
.responses_scrape
|
||||
.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
self.can_send = true;
|
||||
}
|
||||
Err(err) => {
|
||||
eprintln!("error deserializing offer: {:?}", err);
|
||||
}
|
||||
},
|
||||
Err(tungstenite::Error::Io(err)) if err.kind() == ErrorKind::WouldBlock => {
|
||||
return false;
|
||||
},
|
||||
}
|
||||
Err(_) => {
|
||||
return true;
|
||||
}
|
||||
|
|
@ -178,18 +173,14 @@ impl Connection {
|
|||
config: &Config,
|
||||
state: &LoadTestState,
|
||||
rng: &mut impl Rng,
|
||||
) -> bool { // bool = remove connection
|
||||
) -> bool {
|
||||
// bool = remove connection
|
||||
if !self.can_send {
|
||||
return false;
|
||||
}
|
||||
|
||||
if let ConnectionState::WebSocket(ref mut ws) = self.stream {
|
||||
let request = create_random_request(
|
||||
&config,
|
||||
&state,
|
||||
rng,
|
||||
self.peer_id
|
||||
);
|
||||
let request = create_random_request(&config, &state, rng, self.peer_id);
|
||||
|
||||
// If self.send_answer is set and request is announce request, make
|
||||
// the request an offer answer
|
||||
|
|
@ -211,20 +202,16 @@ impl Connection {
|
|||
request
|
||||
};
|
||||
|
||||
match ws.write_message(request.to_ws_message()){
|
||||
match ws.write_message(request.to_ws_message()) {
|
||||
Ok(()) => {
|
||||
state.statistics.requests.fetch_add(1, Ordering::SeqCst);
|
||||
|
||||
self.can_send = false;
|
||||
|
||||
false
|
||||
},
|
||||
Err(tungstenite::Error::Io(err)) if err.kind() == ErrorKind::WouldBlock => {
|
||||
false
|
||||
}
|
||||
Err(_) => {
|
||||
true
|
||||
}
|
||||
Err(tungstenite::Error::Io(err)) if err.kind() == ErrorKind::WouldBlock => false,
|
||||
Err(_) => true,
|
||||
}
|
||||
} else {
|
||||
println!("send request can't send to non-ws stream");
|
||||
|
|
@ -234,14 +221,9 @@ impl Connection {
|
|||
}
|
||||
}
|
||||
|
||||
|
||||
pub type ConnectionMap = HashMap<usize, Connection>;
|
||||
|
||||
|
||||
pub fn run_socket_thread(
|
||||
config: &Config,
|
||||
state: LoadTestState,
|
||||
) {
|
||||
pub fn run_socket_thread(config: &Config, state: LoadTestState) {
|
||||
let timeout = Duration::from_micros(config.network.poll_timeout_microseconds);
|
||||
let create_conn_interval = 2 ^ config.network.connection_creation_interval;
|
||||
|
||||
|
|
@ -259,11 +241,11 @@ pub fn run_socket_thread(
|
|||
poll.poll(&mut events, Some(timeout))
|
||||
.expect("failed polling");
|
||||
|
||||
for event in events.iter(){
|
||||
for event in events.iter() {
|
||||
let token = event.token();
|
||||
|
||||
if event.is_readable(){
|
||||
if let Some(connection) = connections.get_mut(&token.0){
|
||||
if event.is_readable() {
|
||||
if let Some(connection) = connections.get_mut(&token.0) {
|
||||
if let ConnectionState::WebSocket(_) = connection.stream {
|
||||
let drop_connection = connection.read_responses(&state);
|
||||
|
||||
|
|
@ -276,26 +258,22 @@ pub fn run_socket_thread(
|
|||
}
|
||||
}
|
||||
|
||||
if let Some(connection) = connections.remove(&token.0){
|
||||
if let Some(connection) = connection.advance(config){
|
||||
if let Some(connection) = connections.remove(&token.0) {
|
||||
if let Some(connection) = connection.advance(config) {
|
||||
connections.insert(token.0, connection);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (k, connection) in connections.iter_mut(){
|
||||
let drop_connection = connection.send_request(
|
||||
config,
|
||||
&state,
|
||||
&mut rng,
|
||||
);
|
||||
for (k, connection) in connections.iter_mut() {
|
||||
let drop_connection = connection.send_request(config, &state, &mut rng);
|
||||
|
||||
if drop_connection {
|
||||
drop_keys.push(*k)
|
||||
}
|
||||
}
|
||||
|
||||
for k in drop_keys.drain(..){
|
||||
for k in drop_keys.drain(..) {
|
||||
connections.remove(&k);
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue