use super::message::{MessageError, SyncMessage};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ReceiveResult {
NewState {
sender_version: u64,
acked_version: u64,
base_version: u64,
},
AckOnly {
sender_version: u64,
acked_version: u64,
},
Duplicate {
version: u64,
},
Stale {
received: u64,
current: u64,
},
}
#[derive(Debug, Clone)]
pub struct SyncReceiver {
highest_received: u64,
last_acked_to_peer: u64,
}
impl SyncReceiver {
pub fn new() -> Self {
Self {
highest_received: 0,
last_acked_to_peer: 0,
}
}
pub fn highest_received(&self) -> u64 {
self.highest_received
}
pub fn last_acked_to_peer(&self) -> u64 {
self.last_acked_to_peer
}
pub fn needs_ack(&self) -> bool {
self.highest_received > self.last_acked_to_peer
}
pub fn mark_acked(&mut self, version: u64) {
if version > self.last_acked_to_peer {
self.last_acked_to_peer = version;
}
}
pub fn receive_raw(&mut self, data: &[u8]) -> Result<(ReceiveResult, SyncMessage), MessageError> {
let msg = SyncMessage::decode(data)?;
let result = self.receive(&msg);
Ok((result, msg))
}
pub fn receive(&mut self, msg: &SyncMessage) -> ReceiveResult {
let sender_version = msg.sender_state_num;
if sender_version < self.highest_received {
return ReceiveResult::Stale {
received: sender_version,
current: self.highest_received,
};
}
if sender_version == self.highest_received && sender_version > 0 {
return ReceiveResult::Duplicate {
version: sender_version,
};
}
self.highest_received = sender_version;
if msg.is_ack_only() {
ReceiveResult::AckOnly {
sender_version,
acked_version: msg.acked_state_num,
}
} else {
ReceiveResult::NewState {
sender_version,
acked_version: msg.acked_state_num,
base_version: msg.base_state_num,
}
}
}
pub fn reset(&mut self) {
self.highest_received = 0;
self.last_acked_to_peer = 0;
}
}
impl Default for SyncReceiver {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct OrderedReceiver {
inner: SyncReceiver,
received_bitmap: u64,
}
const WINDOW_SIZE: u64 = 64;
impl OrderedReceiver {
pub fn new() -> Self {
Self {
inner: SyncReceiver::new(),
received_bitmap: 0,
}
}
pub fn highest_received(&self) -> u64 {
self.inner.highest_received
}
pub fn needs_ack(&self) -> bool {
self.inner.needs_ack()
}
pub fn mark_acked(&mut self, version: u64) {
self.inner.mark_acked(version);
}
fn version_to_bit_index(&self, version: u64) -> Option<usize> {
if version == 0 || version > self.inner.highest_received {
return None;
}
let offset = self.inner.highest_received - version;
if offset >= WINDOW_SIZE {
return None; }
Some((63 - offset) as usize)
}
pub fn has_received(&self, version: u64) -> bool {
if version > self.inner.highest_received {
return false;
}
if version == 0 {
return true; }
let offset = self.inner.highest_received - version;
if offset >= WINDOW_SIZE {
return true;
}
match self.version_to_bit_index(version) {
Some(bit_index) => (self.received_bitmap & (1u64 << bit_index)) != 0,
None => true, }
}
pub fn receive(&mut self, msg: &SyncMessage) -> ReceiveResult {
let sender_version = msg.sender_state_num;
if self.has_received(sender_version) && sender_version > 0 {
return ReceiveResult::Duplicate {
version: sender_version,
};
}
if sender_version > self.inner.highest_received {
let shift = sender_version - self.inner.highest_received;
if shift >= WINDOW_SIZE {
self.received_bitmap = 1u64 << 63;
} else {
self.received_bitmap >>= shift;
self.received_bitmap |= 1u64 << 63;
}
} else if sender_version > 0 {
if let Some(bit_index) = self.version_to_bit_index(sender_version) {
self.received_bitmap |= 1u64 << bit_index;
}
}
self.inner.receive(msg)
}
pub fn reset(&mut self) {
self.inner.reset();
self.received_bitmap = 0;
}
}
impl Default for OrderedReceiver {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_state_msg(version: u64) -> SyncMessage {
SyncMessage::new(version, 0, 0, vec![1, 2, 3])
}
fn create_ack_msg(sender_version: u64, acked_version: u64) -> SyncMessage {
SyncMessage::ack_only(sender_version, acked_version)
}
mod sync_receiver {
use super::*;
#[test]
fn test_new_receiver() {
let receiver = SyncReceiver::new();
assert_eq!(receiver.highest_received(), 0);
assert!(!receiver.needs_ack());
}
#[test]
fn test_receive_new_state() {
let mut receiver = SyncReceiver::new();
let result = receiver.receive(&create_state_msg(1));
assert!(matches!(result, ReceiveResult::NewState { sender_version: 1, .. }));
assert_eq!(receiver.highest_received(), 1);
assert!(receiver.needs_ack());
}
#[test]
fn test_receive_ack_only() {
let mut receiver = SyncReceiver::new();
let result = receiver.receive(&create_ack_msg(1, 5));
assert!(matches!(
result,
ReceiveResult::AckOnly { sender_version: 1, acked_version: 5 }
));
}
#[test]
fn test_duplicate_detection() {
let mut receiver = SyncReceiver::new();
receiver.receive(&create_state_msg(5));
let result = receiver.receive(&create_state_msg(5));
assert!(matches!(result, ReceiveResult::Duplicate { version: 5 }));
}
#[test]
fn test_stale_detection() {
let mut receiver = SyncReceiver::new();
receiver.receive(&create_state_msg(10));
let result = receiver.receive(&create_state_msg(5));
assert!(matches!(
result,
ReceiveResult::Stale { received: 5, current: 10 }
));
}
#[test]
fn test_needs_ack() {
let mut receiver = SyncReceiver::new();
assert!(!receiver.needs_ack());
receiver.receive(&create_state_msg(1));
assert!(receiver.needs_ack());
receiver.mark_acked(1);
assert!(!receiver.needs_ack());
receiver.receive(&create_state_msg(2));
assert!(receiver.needs_ack());
}
#[test]
fn test_reset() {
let mut receiver = SyncReceiver::new();
receiver.receive(&create_state_msg(5));
receiver.mark_acked(5);
receiver.reset();
assert_eq!(receiver.highest_received(), 0);
assert_eq!(receiver.last_acked_to_peer(), 0);
}
}
mod ordered_receiver {
use super::*;
#[test]
fn test_out_of_order_duplicate() {
let mut receiver = OrderedReceiver::new();
receiver.receive(&create_state_msg(1));
receiver.receive(&create_state_msg(2));
receiver.receive(&create_state_msg(3));
let result = receiver.receive(&create_state_msg(2));
assert!(matches!(result, ReceiveResult::Duplicate { version: 2 }));
}
#[test]
fn test_has_received() {
let mut receiver = OrderedReceiver::new();
receiver.receive(&create_state_msg(5));
receiver.receive(&create_state_msg(10));
receiver.receive(&create_state_msg(7));
assert!(receiver.has_received(5));
assert!(receiver.has_received(7));
assert!(receiver.has_received(10));
assert!(!receiver.has_received(6));
assert!(!receiver.has_received(8));
}
#[test]
fn test_window_sliding() {
let mut receiver = OrderedReceiver::new();
receiver.receive(&create_state_msg(1));
assert!(receiver.has_received(1));
receiver.receive(&create_state_msg(100));
assert!(receiver.has_received(1));
}
#[test]
fn test_reset() {
let mut receiver = OrderedReceiver::new();
receiver.receive(&create_state_msg(5));
receiver.reset();
assert_eq!(receiver.highest_received(), 0);
assert!(!receiver.has_received(5));
}
}
}