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

Refactor `MessageWriter`

feature_link
Anton Tarasenko 3 years ago
parent
5bfdd61248
commit
034b7860d6
1 changed files with 127 additions and 81 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 206
      src/link/writer.rs

206
src/link/writer.rs
Unescape 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: [
let message_as_utf8 = [
&(message_as_utf8.len() as u32).to_be_bytes(),
message_as_utf8,
]
.concat(),
bytes_sent: 0,
});
.concat();
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.should_send_now() {
if self.is_empty() {
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],
]
.concat();
self.sent_bytes += bytes_to_send.len() as u64;
if next_message.bytes_sent >= next_message.contents.len() {
self.pending_messages.pop_front();
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);
}

Write Preview
Loading…
Cancel
Save