dkanus/Avarice
2
1
Fork 1
Code Issues 9 Pull Requests 2 Releases Wiki Activity

Refactor MessageWriter

Browse Source
This commit is contained in:
Anton Tarasenko 2021-07-24 02:26:43 +07:00
parent 5bfdd61248
commit 034b7860d6
1 changed files with 130 additions and 84 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

214
src/link/writer.rs
View File

@ -1,79 +1,86 @@
use std::cmp::{max, min};
use std::collections::VecDeque;
use std::convert::TryFrom;
use std::iter::Extend;
const LENGTH_FIELD_SIZE: usize = 4;
// Defines how many bytes is used to encode "LENGTH" field in the chunk sent to ue-server
const CHUNK_LENGTH_FIELD: usize = 2;
// Maximum amount of bytes ue-server is able to receive at once
const UE_INPUT_BUFFER: usize = 4095;
struct PendingMessage {
contents: Vec<u8>,
bytes_sent: usize,
}
// Minimal payload size (in bytes) to send, unless there is not enough data left
const MIN_PAYLOAD_SIZE: usize = 50;
pub struct MessageWriter {
sent_bytes: u64,
ue_received_bytes: u64,
pending_messages: VecDeque<PendingMessage>,
}
impl PendingMessage {
fn bytes_left(&self) -> usize {
max(0, self.contents.len() - self.bytes_sent)
}
pending_data: VecDeque<u8>,
}
/// For converting byte stream that is expected from the ue-server into actual messages.
/// Conversion process has two steps:
/// 1. Every string message is converted into it's utf8 representation and is pre-pended with
/// it's own length in format:
/// [MESSAGE_LENGTH: length of the message in utf8 encoding | 4 bytes: u32 BE]
/// [MESSAGE: utf8-encoded string | `MESSAGE_LENGTH` bytes]
/// Resulting byte sequences from all the messages then concatenated, in order, into
/// a single data stream.
/// 2. Resulting data stream is then separated into "chunks" that can be accepted by
/// the ue-server (each no longer than `UE_INPUT_BUFFER` in total) and are sent in a format:
/// [LENGTH: length of the chunk | 2 bytes: u16 BE]
/// [PAYLOAD: utf8-encoded string | `LENGTH` bytes]
/// Use `push()` to input string messages and `try_pop()` to retrieve next chunk, if ue-server
/// can accept it. Call `update_ue_received_bytes()` to update `MessageWriter`'s information about
/// how many bytes ue-server has received so far.
/// NOTE: `try_pop()` can return `None` even if not all message data has been transferred,
/// in case ue-server's buffer does not have enough space. Use `is_empty()` call to check for
/// whether `MessageWriter` has depleted all it's data.
impl MessageWriter {
pub fn new() -> MessageWriter {
MessageWriter {
sent_bytes: 0,
ue_received_bytes: 0,
pending_messages: VecDeque::with_capacity(100),
// This value should be more than enough for typical use:
// it will take at least one second to send this much data to a 30 tick rate ue-server.
pending_data: VecDeque::with_capacity(30 * UE_INPUT_BUFFER),
}
}
pub fn push(&mut self, message: &str) {
let message_as_utf8 = message.as_bytes();
self.pending_messages.push_back(PendingMessage {
contents: [
&(message_as_utf8.len() as u32).to_be_bytes(),
message_as_utf8,
]
.concat(),
bytes_sent: 0,
});
}
pub fn try_pop(&mut self) -> Option<Vec<u8>> {
if !self.should_send_now() {
return None;
}
let max_output_bytes = self.available_ue_buffer_capacity();
let next_payload_size = max_output_bytes - 2;
let next_message = match self.pending_messages.get_mut(0) {
Some(message) => message,
_ => return None,
};
let first_index = next_message.bytes_sent;
let last_index = min(
next_message.bytes_sent + next_payload_size - 1,
next_message.contents.len() - 1,
);
next_message.bytes_sent = last_index + 1;
let chunk_length = last_index - first_index + 1;
let bytes_to_send = [
&(chunk_length as u16).to_be_bytes(),
&next_message.contents[first_index..last_index + 1],
let message_as_utf8 = [
&(message_as_utf8.len() as u32).to_be_bytes(),
message_as_utf8,
]
.concat();
self.sent_bytes += bytes_to_send.len() as u64;
if next_message.bytes_sent >= next_message.contents.len() {
self.pending_messages.pop_front();
self.pending_data.extend(message_as_utf8.into_iter());
}
/// This method will always return chunk if all remaining data will fit inside it, otherwise it
/// will wait until ue-server's buffer has enough space for at least `MIN_PAYLOAD_SIZE` bytes.
pub fn try_pop(&mut self) -> Option<Vec<u8>> {
if self.is_empty() {
return None;
}
let required_payload_size = min(self.pending_data.len(), MIN_PAYLOAD_SIZE);
let available_payload_space = self
.available_ue_buffer_capacity()
.checked_sub(CHUNK_LENGTH_FIELD)
.unwrap_or_default();
if required_payload_size > available_payload_space {
return None;
}
let payload_size = min(available_payload_space, self.pending_data.len());
let mut bytes_to_send = Vec::with_capacity(CHUNK_LENGTH_FIELD + payload_size);
bytes_to_send.extend((payload_size as u16).to_be_bytes().iter());
for next_byte in self.pending_data.drain(..payload_size) {
bytes_to_send.push(next_byte);
}
self.sent_bytes += bytes_to_send.len() as u64;
Some(bytes_to_send)
}
pub fn is_empty(&self) -> bool {
self.pending_messages.is_empty()
self.pending_data.is_empty()
}
pub fn update_ue_received_bytes(&mut self, ue_received_bytes: u64) {
@ -86,39 +93,27 @@ impl MessageWriter {
_ => 0,
}
}
fn should_send_now(&self) -> bool {
if self.pending_messages.is_empty() {
return false;
}
let bytes_left_in_message = match self.pending_messages.get(0) {
Some(message) => message.bytes_left(),
_ => return false,
};
let max_output_bytes = self.available_ue_buffer_capacity();
if bytes_left_in_message <= max_output_bytes {
return true;
};
let max_bytes_in_payload = UE_INPUT_BUFFER - LENGTH_FIELD_SIZE;
let next_payload_size = max_output_bytes.checked_sub(4).unwrap_or_default();
let chunks_in_message = (bytes_left_in_message / max_bytes_in_payload) + 1;
let chunks_in_message_after_sending =
((bytes_left_in_message - next_payload_size) / max_bytes_in_payload) + 1;
chunks_in_message_after_sending < chunks_in_message
}
}
#[test]
fn new_writer_is_empty() {
let mut writer = MessageWriter::new();
fn writer_contents_after_creation() {
let writer = MessageWriter::new();
assert_eq!(writer.is_empty(), true);
assert_eq!(writer.available_ue_buffer_capacity(), UE_INPUT_BUFFER);
}
#[test]
fn single_short_message() {
fn writer_content_after_push() {
let mut writer = MessageWriter::new();
writer.push("Hello, world!");
assert_eq!(writer.is_empty(), false);
assert_eq!(writer.available_ue_buffer_capacity(), UE_INPUT_BUFFER);
}
#[test]
fn writing_single_short_message() {
let mut writer = MessageWriter::new();
writer.push("Hello, world!");
let resulting_bytes = writer.try_pop().unwrap();
let expected_bytes = [
0, 17, // Bytes in the chunk = message length (4 bytes) + message (13 bytes)
@ -126,44 +121,95 @@ fn single_short_message() {
b'H', b'e', b'l', b'l', b'o', b',', b' ', b'w', b'o', b'r', b'l', b'd', b'!',
];
assert_eq!(writer.is_empty(), true);
assert_eq!(
writer.available_ue_buffer_capacity(),
UE_INPUT_BUFFER - "Hello, world!".len() - 2 - 4
);
assert_eq!(resulting_bytes, expected_bytes);
assert_eq!(writer.sent_bytes, expected_bytes.len() as u64);
}
#[test]
fn single_long_message() {
fn writing_first_chunk_of_single_long_message() {
let mut writer = MessageWriter::new();
// Because we also have to pass lengths, this will go over the sending limit
let long_message = "Q".repeat(UE_INPUT_BUFFER);
writer.push(&long_message);
assert_eq!(writer.is_empty(), false);
let resulting_bytes = writer.try_pop().unwrap();
assert_eq!(writer.is_empty(), false);
assert_eq!(resulting_bytes.len(), 4095);
assert_eq!(resulting_bytes.len(), UE_INPUT_BUFFER);
assert_eq!(writer.available_ue_buffer_capacity(), 0);
// Bytes in the chunk = 4095 - 2 = 4093 = 0x0ffd
assert_eq!(resulting_bytes[0], 0x0f);
assert_eq!(resulting_bytes[1], 0x0fd);
assert_eq!(resulting_bytes[1], 0xfd);
// Bytes in message = 4095 = 0x0fff
assert_eq!(resulting_bytes[2], 0);
assert_eq!(resulting_bytes[3], 0);
assert_eq!(resulting_bytes[4], 0x0f);
assert_eq!(resulting_bytes[5], 0x0ff);
assert_eq!(resulting_bytes[5], 0xff);
for &byte in resulting_bytes[6..].iter() {
assert_eq!(byte, b'Q');
}
assert_eq!(writer.try_pop(), None);
assert_eq!(writer.is_empty(), false);
writer.update_ue_received_bytes(UE_INPUT_BUFFER as u64);
}
#[test]
fn writing_second_chunk_of_single_long_message() {
let mut writer = MessageWriter::new();
// Because we also have to pass lengths, this will go over the sending limit
let long_message = "Q".repeat(UE_INPUT_BUFFER);
writer.push(&long_message);
// This pops all but 6 bytes of `long_message`, that were required to encode lengths of
// message and first chunk
let first_bytes = writer.try_pop().unwrap();
writer.update_ue_received_bytes(first_bytes.len() as u64);
let resulting_bytes = writer.try_pop().unwrap();
assert_eq!(
writer.available_ue_buffer_capacity(),
UE_INPUT_BUFFER - resulting_bytes.len()
);
assert_eq!(writer.is_empty(), true);
// Last time we have popped all but 6 bytes from `long_message`
// Bytes in the chunk = 6 = 0x06
// Bytes in the chunk = 6
assert_eq!(resulting_bytes[0], 0);
assert_eq!(resulting_bytes[1], 0x06);
assert_eq!(resulting_bytes[1], 6);
assert_eq!(resulting_bytes[2..], [b'Q', b'Q', b'Q', b'Q', b'Q', b'Q'])
}
// TODO: test that `update_ue_received_bytes()` cannot reduce that value
// TODO: test several messages, message only partially fitting into remaining buffer
// TODO: test giving just barely enough free buffer space
#[test]
fn will_write_small_chunks_if_no_more_data() {
let mut writer = MessageWriter::new();
// Because we also have to pass lengths (of chunk `CHUNK_LENGTH_FIELD` amd of message `4`),
// sending this will leave us with exactly 10 free bytes in the buffer
let long_message = "Q".repeat(UE_INPUT_BUFFER / 2);
writer.push(&long_message);
writer.try_pop();
let short_message = "Hello, world!";
writer.push(&short_message);
let expected_bytes = [
0, 17, // Bytes in the chunk = message length (4 bytes) + message (13 bytes)
0, 0, 0, 13, // Bytes in the message
b'H', b'e', b'l', b'l', b'o', b',', b' ', b'w', b'o', b'r', b'l', b'd', b'!',
];
// There should be enough space in the ue-server buffer to send `short_message`
let resulting_bytes = writer.try_pop().unwrap();
assert_eq!(resulting_bytes, expected_bytes);
assert_eq!(writer.try_pop(), None);
assert_eq!(writer.is_empty(), true);
}
#[test]
fn will_not_write_small_chunks_if_more_data_remains() {
let mut writer = MessageWriter::new();
// Because we also have to pass lengths (of chunk `CHUNK_LENGTH_FIELD` amd of message `4`),
// sending this will leave us with exactly 10 free bytes in the buffer
let long_message = "Q".repeat(UE_INPUT_BUFFER - CHUNK_LENGTH_FIELD - 4 - 10);
writer.push(&long_message);
writer.try_pop();
let short_message = "Hello, world!";
writer.push(&short_message);
// `MessageWriter` can neither send full message, nor a chunk of size 10
// (because it is too short)
assert_eq!(writer.try_pop(), None);
assert_eq!(writer.is_empty(), false);
}