aquatic_http: write to buf ref in request serialization

This commit is contained in:
Joakim Frostegård 2020-07-24 22:27:06 +02:00
parent fd68a5f603
commit cc46c4b72c
3 changed files with 60 additions and 57 deletions

View file

@ -1,6 +1,6 @@
use std::sync::atomic::Ordering; use std::sync::atomic::Ordering;
use std::time::Duration; use std::time::Duration;
use std::io::{Read, Write, ErrorKind}; use std::io::{Read, Write, ErrorKind, Cursor};
use mio::{net::TcpStream, Events, Poll, Interest, Token}; use mio::{net::TcpStream, Events, Poll, Interest, Token};
use rand::{rngs::SmallRng, prelude::*}; use rand::{rngs::SmallRng, prelude::*};
@ -53,6 +53,7 @@ impl Connection {
config: &Config, config: &Config,
state: &LoadTestState, state: &LoadTestState,
rng: &mut impl Rng, rng: &mut impl Rng,
request_buffer: &mut Cursor<&mut [u8]>,
){ ){
loop { loop {
match self.stream.read(&mut self.read_buffer[self.bytes_read..]){ match self.stream.read(&mut self.read_buffer[self.bytes_read..]){
@ -90,7 +91,8 @@ impl Connection {
self.send_request( self.send_request(
config, config,
state, state,
rng rng,
request_buffer,
); );
} }
@ -128,6 +130,7 @@ impl Connection {
config: &Config, config: &Config,
state: &LoadTestState, state: &LoadTestState,
rng: &mut impl Rng, rng: &mut impl Rng,
request_buffer: &mut Cursor<&mut [u8]>,
){ ){
let request = create_random_request( let request = create_random_request(
&config, &config,
@ -135,7 +138,11 @@ impl Connection {
rng rng
); );
match self.send_request_inner(state, &request.as_bytes()){ request_buffer.set_position(0);
request.write(request_buffer).unwrap();
let position = request_buffer.position() as usize;
match self.send_request_inner(state, &request_buffer.get_mut()[..position]){
Ok(_) => { Ok(_) => {
state.statistics.requests.fetch_add(1, Ordering::SeqCst); state.statistics.requests.fetch_add(1, Ordering::SeqCst);
}, },
@ -183,6 +190,8 @@ pub fn run_socket_thread(
let mut poll = Poll::new().expect("create poll"); let mut poll = Poll::new().expect("create poll");
let mut events = Events::with_capacity(config.network.poll_event_capacity); let mut events = Events::with_capacity(config.network.poll_event_capacity);
let mut rng = SmallRng::from_entropy(); let mut rng = SmallRng::from_entropy();
let mut request_buffer = [0u8; 1024];
let mut request_buffer = Cursor::new(&mut request_buffer[..]);
let mut token_counter = 0usize; let mut token_counter = 0usize;
@ -210,7 +219,8 @@ pub fn run_socket_thread(
connection.read_response_and_send_request( connection.read_response_and_send_request(
config, config,
&state, &state,
&mut rng &mut rng,
&mut request_buffer
); );
} else { } else {
eprintln!("connection not found: {:?}", token); eprintln!("connection not found: {:?}", token);
@ -221,7 +231,12 @@ pub fn run_socket_thread(
if !initial_sent { if !initial_sent {
for (_, connection) in connections.iter_mut(){ for (_, connection) in connections.iter_mut(){
if connection.can_send_initial { if connection.can_send_initial {
connection.send_request(config, &state, &mut rng); connection.send_request(
config,
&state,
&mut rng,
&mut request_buffer
);
initial_sent = true; initial_sent = true;
} }

View file

@ -1,3 +1,5 @@
use std::io::Write;
use anyhow::Context; use anyhow::Context;
use smartstring::{SmartString, LazyCompact}; use smartstring::{SmartString, LazyCompact};
@ -20,54 +22,37 @@ pub struct AnnounceRequest {
impl AnnounceRequest { impl AnnounceRequest {
pub fn as_bytes(&self) -> Vec<u8> { fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<()> {
let mut bytes = Vec::with_capacity( output.write(b"GET /announce?info_hash=")?;
24 + urlencode_20_bytes(self.info_hash.0, output)?;
60 +
9 +
60 +
6 +
5 + // high estimate
6 +
2 + // estimate
14 + // FIXME event
9 +
1 +
20 + // numwant bad estimate
20 + // key bad estimate
13
);
bytes.extend_from_slice(b"GET /announce?info_hash="); output.write(b"&peer_id=")?;
urlencode_20_bytes(self.info_hash.0, &mut bytes); urlencode_20_bytes(self.info_hash.0, output)?;
bytes.extend_from_slice(b"&peer_id="); output.write(b"&port=")?;
urlencode_20_bytes(self.info_hash.0, &mut bytes); output.write(itoa::Buffer::new().format(self.port).as_bytes())?;
bytes.extend_from_slice(b"&port="); output.write(b"&left=")?;
let _ = itoa::write(&mut bytes, self.port); output.write(itoa::Buffer::new().format(self.bytes_left).as_bytes())?;
bytes.extend_from_slice(b"&left="); output.write(b"&event=started")?; // FIXME
let _ = itoa::write(&mut bytes, self.bytes_left);
bytes.extend_from_slice(b"&event=started"); // FIXME output.write(b"&compact=")?;
output.write(itoa::Buffer::new().format(self.compact as u8).as_bytes())?;
bytes.extend_from_slice(b"&compact=");
let _ = itoa::write(&mut bytes, self.compact as u8);
if let Some(numwant) = self.numwant { if let Some(numwant) = self.numwant {
bytes.extend_from_slice(b"&numwant="); output.write(b"&numwant=")?;
let _ = itoa::write(&mut bytes, numwant); output.write(itoa::Buffer::new().format(numwant).as_bytes())?;
} }
if let Some(ref key) = self.key { if let Some(ref key) = self.key {
bytes.extend_from_slice(b"&key="); output.write(b"&key=")?;
bytes.extend_from_slice(key.as_str().as_bytes()); output.write(key.as_str().as_bytes())?;
} }
bytes.extend_from_slice(b" HTTP/1.1\r\n\r\n"); output.write(b" HTTP/1.1\r\n\r\n")?;
bytes Ok(())
} }
} }
@ -79,27 +64,25 @@ pub struct ScrapeRequest {
impl ScrapeRequest { impl ScrapeRequest {
pub fn as_bytes(&self) -> Vec<u8> { fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<()> {
let mut bytes = Vec::new(); output.write(b"GET /scrape?")?;
bytes.extend_from_slice(b"GET /scrape?");
let mut first = true; let mut first = true;
for info_hash in self.info_hashes.iter() { for info_hash in self.info_hashes.iter() {
if !first { if !first {
bytes.push(b'&') output.write(b"&")?;
} }
bytes.extend_from_slice(b"info_hash="); output.write(b"info_hash=")?;
urlencode_20_bytes(info_hash.0, &mut bytes); urlencode_20_bytes(info_hash.0, output)?;
first = false; first = false;
} }
bytes.extend_from_slice(b" HTTP/1.1\r\n\r\n"); output.write(b" HTTP/1.1\r\n\r\n")?;
bytes Ok(())
} }
} }
@ -270,10 +253,10 @@ impl Request {
} }
} }
pub fn as_bytes(&self) -> Vec<u8> { pub fn write<W: Write>(&self, output: &mut W) -> ::std::io::Result<()> {
match self { match self {
Self::Announce(r) => r.as_bytes(), Self::Announce(r) => r.write(output),
Self::Scrape(r) => r.as_bytes(), Self::Scrape(r) => r.write(output),
} }
} }
} }

View file

@ -8,15 +8,20 @@ use smartstring::{SmartString, LazyCompact};
use super::response::ResponsePeer; use super::response::ResponsePeer;
pub fn urlencode_20_bytes(input: [u8; 20], output: &mut impl Write){ pub fn urlencode_20_bytes(
input: [u8; 20],
output: &mut impl Write
) -> ::std::io::Result<()> {
let mut tmp = [0u8; 2]; let mut tmp = [0u8; 2];
for i in 0..input.len() { for i in 0..input.len() {
hex::encode_to_slice(&input[i..i+1], &mut tmp).unwrap(); hex::encode_to_slice(&input[i..i+1], &mut tmp).unwrap();
output.write(b"%"); output.write(b"%")?;
output.write(&tmp); output.write(&tmp)?;
} }
Ok(())
} }
@ -160,7 +165,7 @@ mod tests {
let mut output = Vec::new(); let mut output = Vec::new();
urlencode_20_bytes(input, &mut output); urlencode_20_bytes(input, &mut output).unwrap();
assert_eq!(output.len(), 60); assert_eq!(output.len(), 60);
@ -195,7 +200,7 @@ mod tests {
let mut output = Vec::new(); let mut output = Vec::new();
urlencode_20_bytes(input, &mut output); urlencode_20_bytes(input, &mut output).unwrap();
let s = ::std::str::from_utf8(&output).unwrap(); let s = ::std::str::from_utf8(&output).unwrap();