Anton Tarasenko
3 years ago
2 changed files with 171 additions and 0 deletions
@ -1,2 +1,4 @@
|
||||
mod reader; |
||||
mod writer; |
||||
pub use reader::MessageReader; |
||||
pub use writer::MessageWriter; |
||||
|
@ -0,0 +1,169 @@
|
||||
use std::cmp::{max, min}; |
||||
use std::collections::VecDeque; |
||||
use std::convert::TryFrom; |
||||
|
||||
const LENGTH_FIELD_SIZE: usize = 4; |
||||
const UE_INPUT_BUFFER: usize = 4095; |
||||
|
||||
struct PendingMessage { |
||||
contents: Vec<u8>, |
||||
bytes_sent: usize, |
||||
} |
||||
|
||||
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) |
||||
} |
||||
} |
||||
|
||||
impl MessageWriter { |
||||
pub fn new() -> MessageWriter { |
||||
MessageWriter { |
||||
sent_bytes: 0, |
||||
ue_received_bytes: 0, |
||||
pending_messages: VecDeque::with_capacity(100), |
||||
} |
||||
} |
||||
|
||||
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], |
||||
] |
||||
.concat(); |
||||
self.sent_bytes += bytes_to_send.len() as u64; |
||||
if next_message.bytes_sent >= next_message.contents.len() { |
||||
self.pending_messages.pop_front(); |
||||
} |
||||
Some(bytes_to_send) |
||||
} |
||||
|
||||
pub fn is_empty(&self) -> bool { |
||||
self.pending_messages.is_empty() |
||||
} |
||||
|
||||
pub fn update_ue_received_bytes(&mut self, ue_received_bytes: u64) { |
||||
self.ue_received_bytes = max(ue_received_bytes, self.ue_received_bytes); |
||||
} |
||||
|
||||
fn available_ue_buffer_capacity(&self) -> usize { |
||||
match usize::try_from(self.sent_bytes - self.ue_received_bytes).ok() { |
||||
Some(ue_buffered_bytes) => max(0, UE_INPUT_BUFFER - ue_buffered_bytes), |
||||
_ => 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(); |
||||
assert_eq!(writer.is_empty(), true); |
||||
} |
||||
|
||||
#[test] |
||||
fn single_short_message() { |
||||
let mut writer = MessageWriter::new(); |
||||
writer.push("Hello, world!"); |
||||
assert_eq!(writer.is_empty(), false); |
||||
let resulting_bytes = writer.try_pop().unwrap(); |
||||
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'!', |
||||
]; |
||||
assert_eq!(writer.is_empty(), true); |
||||
assert_eq!(resulting_bytes, expected_bytes); |
||||
assert_eq!(writer.sent_bytes, expected_bytes.len() as u64); |
||||
} |
||||
|
||||
#[test] |
||||
fn 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); |
||||
// Bytes in the chunk = 4095 - 2 = 4093 = 0x0ffd
|
||||
assert_eq!(resulting_bytes[0], 0x0f); |
||||
assert_eq!(resulting_bytes[1], 0x0fd); |
||||
// 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); |
||||
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); |
||||
let resulting_bytes = writer.try_pop().unwrap(); |
||||
assert_eq!(writer.is_empty(), true); |
||||
// Last time we have popped all but 6 bytes from `long_message`
|
||||
// Bytes in the chunk = 6 = 0x06
|
||||
assert_eq!(resulting_bytes[0], 0); |
||||
assert_eq!(resulting_bytes[1], 0x06); |
||||
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
|
Loading…
Reference in new issue