use std::cmp::Ordering;
use std::cmp::Reverse;
use std::collections::BinaryHeap;
use crate::adaptive_ring::{AdaptiveRing, PinnedRing};
use crate::ordering::{OrderingMode, StampKind, STAMPED_PAYLOAD_BYTES};
pub const DEFAULT_FLOOR: usize = 8;
pub const DEFAULT_CAP: usize = 1024;
#[derive(Clone)]
struct Entry {
stamp: u64,
seq: u64,
len: usize,
payload: [u8; STAMPED_PAYLOAD_BYTES],
}
impl PartialEq for Entry {
fn eq(&self, other: &Self) -> bool {
(self.stamp, self.seq) == (other.stamp, other.seq)
}
}
impl Eq for Entry {}
impl PartialOrd for Entry {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for Entry {
fn cmp(&self, other: &Self) -> Ordering {
(self.stamp, self.seq).cmp(&(other.stamp, other.seq))
}
}
pub struct ReorderBuffer {
heap: BinaryHeap<Reverse<Entry>>,
window: usize,
cap: usize,
last_emitted: u64,
have_emitted: bool,
next_seq: u64,
corrections: u64,
}
impl Default for ReorderBuffer {
fn default() -> Self {
Self::with_window(DEFAULT_FLOOR, DEFAULT_CAP)
}
}
impl ReorderBuffer {
pub fn new() -> Self {
Self::default()
}
pub fn with_window(floor: usize, cap: usize) -> Self {
Self {
heap: BinaryHeap::new(),
window: floor,
cap: cap.max(floor),
last_emitted: 0,
have_emitted: false,
next_seq: 0,
corrections: 0,
}
}
pub fn widen_to(&mut self, min_window: usize) {
if min_window > self.window {
self.window = min_window;
self.cap = self.cap.max(min_window);
}
}
pub fn push(&mut self, stamp: u64, payload: &[u8]) {
let len = payload.len().min(STAMPED_PAYLOAD_BYTES);
let mut buf = [0u8; STAMPED_PAYLOAD_BYTES];
buf[..len].copy_from_slice(&payload[..len]);
let seq = self.next_seq;
self.next_seq += 1;
self.heap.push(Reverse(Entry { stamp, seq, len, payload: buf }));
}
pub fn try_take(&mut self, out: &mut [u8]) -> Option<(u64, usize)> {
if self.heap.len() > self.window {
self.release(out)
} else {
None
}
}
pub fn flush_one(&mut self, out: &mut [u8]) -> Option<(u64, usize)> {
self.release(out)
}
fn release(&mut self, out: &mut [u8]) -> Option<(u64, usize)> {
let Reverse(entry) = self.heap.pop()?;
if self.have_emitted && entry.stamp < self.last_emitted {
self.corrections += 1;
self.window = self.window.saturating_mul(2).min(self.cap).max(1);
}
self.last_emitted = entry.stamp;
self.have_emitted = true;
let n = entry.len.min(out.len());
out[..n].copy_from_slice(&entry.payload[..n]);
Some((entry.stamp, n))
}
pub fn len(&self) -> usize {
self.heap.len()
}
pub fn is_empty(&self) -> bool {
self.heap.is_empty()
}
pub fn window(&self) -> usize {
self.window
}
pub fn corrections(&self) -> u64 {
self.corrections
}
}
pub struct ReorderingReceiver<'a> {
pin: &'a PinnedRing<'a>,
consumer_id: usize,
buf: ReorderBuffer,
scratch: [u8; STAMPED_PAYLOAD_BYTES],
}
impl<'a> ReorderingReceiver<'a> {
pub fn new(pin: &'a PinnedRing<'a>, consumer_id: usize) -> Self {
Self::with_window(pin, consumer_id, DEFAULT_FLOOR, DEFAULT_CAP)
}
pub fn with_window(
pin: &'a PinnedRing<'a>,
consumer_id: usize,
floor: usize,
cap: usize,
) -> Self {
Self {
pin,
consumer_id,
buf: ReorderBuffer::with_window(floor, cap),
scratch: [0u8; STAMPED_PAYLOAD_BYTES],
}
}
pub fn try_recv(&mut self, out: &mut [u8]) -> Option<(usize, u64)> {
if let Ok((n, stamp)) =
self.pin.ordered_try_pop_with_stamp(self.consumer_id, &mut self.scratch)
{
self.buf.push(stamp, &self.scratch[..n]);
}
self.buf.try_take(out).map(|(stamp, len)| (len, stamp))
}
pub fn flush(&mut self, out: &mut [u8]) -> Option<(usize, u64)> {
self.buf.flush_one(out).map(|(stamp, len)| (len, stamp))
}
pub fn corrections(&self) -> u64 {
self.buf.corrections()
}
pub fn window(&self) -> usize {
self.buf.window()
}
}
pub const REORDER_PRODUCER_CAP: usize = 256;
enum ExactMode {
Reorder(ReorderBuffer),
Strict,
Direct,
}
pub struct AdaptiveOrderedReceiver<'a> {
ring: &'a AdaptiveRing,
consumer_id: usize,
mode: ExactMode,
}
impl<'a> AdaptiveOrderedReceiver<'a> {
pub fn new(ring: &'a AdaptiveRing, consumer_id: usize) -> Self {
let producers = ring.published_producers().max(ring.max_producers());
let mode = match ring.stamp_kind() {
Some(StampKind::SharedCounter) => {
if producers <= REORDER_PRODUCER_CAP {
ring.set_ordering_mode(OrderingMode::MergeByStamp).ok();
let window = producers.max(DEFAULT_FLOOR);
ExactMode::Reorder(ReorderBuffer::with_window(
window,
window.max(DEFAULT_CAP),
))
} else {
ring.set_ordering_mode(OrderingMode::MergeStrict).ok();
ExactMode::Strict
}
}
_ => ExactMode::Direct,
};
Self { ring, consumer_id, mode }
}
pub fn try_recv(&mut self, out: &mut [u8]) -> Option<(usize, u64)> {
let pin = self.ring.pin_current_shape();
let cid = self.consumer_id;
match &mut self.mode {
ExactMode::Reorder(rb) => {
let producers = self.ring.published_producers();
if producers > rb.window() {
if producers > REORDER_PRODUCER_CAP {
self.ring
.set_ordering_mode(OrderingMode::MergeStrict)
.ok();
}
rb.widen_to(producers.min(REORDER_PRODUCER_CAP));
}
let mut scratch = [0u8; STAMPED_PAYLOAD_BYTES];
if let Ok((n, stamp)) = pin.ordered_try_pop_with_stamp(cid, &mut scratch) {
rb.push(stamp, &scratch[..n]);
}
rb.try_take(out).map(|(stamp, len)| (len, stamp))
}
ExactMode::Strict | ExactMode::Direct => {
match pin.ordered_try_pop_with_stamp(cid, out) {
Ok((n, stamp)) => Some((n, stamp)),
Err(_) => None,
}
}
}
}
pub fn flush(&mut self, out: &mut [u8]) -> Option<(usize, u64)> {
match &mut self.mode {
ExactMode::Reorder(rb) => rb.flush_one(out).map(|(stamp, len)| (len, stamp)),
_ => None,
}
}
pub fn strategy(&self) -> &'static str {
match self.mode {
ExactMode::Reorder(_) => "reorder",
ExactMode::Strict => "strict",
ExactMode::Direct => "direct",
}
}
pub fn corrections(&self) -> u64 {
match &self.mode {
ExactMode::Reorder(rb) => rb.corrections(),
_ => 0,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn displaced(n: u64, d: usize) -> Vec<u64> {
let mut v = Vec::with_capacity(n as usize);
let block = d as u64 + 1;
let mut base = 1u64;
while base <= n {
let end = (base + block - 1).min(n);
for s in (base..=end).rev() {
v.push(s);
}
base = end + 1;
}
v
}
fn drive(stream: &[u64], floor: usize) -> (Vec<u64>, u64) {
let mut rb = ReorderBuffer::with_window(floor, floor); let mut out = [0u8; STAMPED_PAYLOAD_BYTES];
let mut emitted = Vec::new();
for &s in stream {
rb.push(s, &s.to_le_bytes());
if let Some((stamp, _)) = rb.try_take(&mut out) {
emitted.push(stamp);
}
}
while let Some((stamp, _)) = rb.flush_one(&mut out) {
emitted.push(stamp);
}
(emitted, rb.corrections())
}
fn is_monotone(v: &[u64]) -> bool {
v.windows(2).all(|w| w[1] >= w[0])
}
#[test]
fn in_order_stream_is_untouched() {
let stream: Vec<u64> = (1..=1000).collect();
let (emitted, corr) = drive(&stream, 8);
assert_eq!(emitted, stream);
assert_eq!(corr, 0);
}
#[test]
fn window_at_or_above_displacement_is_exact() {
for d in [1usize, 3, 8] {
let stream = displaced(2000, d);
let (emitted, corr) = drive(&stream, d); assert!(is_monotone(&emitted), "d={d}: not monotone");
assert_eq!(corr, 0, "d={d}: unexpected corrections at window==d");
assert_eq!(emitted.len(), stream.len());
}
}
#[test]
fn window_below_displacement_slips_and_is_flagged() {
let stream = displaced(2000, 4);
let (emitted, corr) = drive(&stream, 1);
assert!(corr > 0, "expected corrections when window < displacement");
assert_eq!(emitted.len(), stream.len());
}
#[test]
fn adaptive_window_grows_toward_displacement() {
let stream = displaced(4000, 8);
let mut rb = ReorderBuffer::with_window(2, 64);
let mut out = [0u8; STAMPED_PAYLOAD_BYTES];
for &s in &stream {
rb.push(s, &s.to_le_bytes());
while rb.try_take(&mut out).is_some() {}
}
while rb.flush_one(&mut out).is_some() {}
assert!(rb.window() > 2, "window should have grown, got {}", rb.window());
assert!(rb.corrections() > 0, "growth is driven by caught slips");
assert!(rb.window() >= 8, "window should reach the displacement, got {}", rb.window());
}
#[test]
fn holes_do_not_stall_and_stay_monotone() {
let mut stream = Vec::new();
let mut s = 1u64;
while s < 300 {
stream.push(s + 1);
stream.push(s);
s += 3; }
let (emitted, _) = drive(&stream, 8);
assert!(is_monotone(&emitted), "delivery must be monotone across holes");
assert_eq!(emitted.len(), stream.len(), "no item dropped");
}
#[test]
fn payload_round_trips() {
let mut rb = ReorderBuffer::with_window(1, 1);
rb.push(10, b"hello");
rb.push(9, b"world");
let mut out = [0u8; STAMPED_PAYLOAD_BYTES];
let (stamp, n) = rb.try_take(&mut out).expect("release");
assert_eq!(stamp, 9);
assert_eq!(&out[..n], b"world");
let (stamp, n) = rb.flush_one(&mut out).expect("flush");
assert_eq!(stamp, 10);
assert_eq!(&out[..n], b"hello");
}
#[test]
fn adaptive_receiver_selects_strategy_by_config() {
use crate::adaptive_ring::{AdaptiveRing, RingShape};
let ring = AdaptiveRing::create_anon(4, 1, 256)
.unwrap()
.with_ordering_stamps_kind(StampKind::SharedCounter)
.unwrap();
ring.morph_to(RingShape::Mpsc).unwrap();
let rx = AdaptiveOrderedReceiver::new(&ring, 0);
assert_eq!(rx.strategy(), "reorder");
assert_eq!(ring.ordering_mode(), Some(OrderingMode::MergeByStamp));
let big = AdaptiveRing::create_anon(REORDER_PRODUCER_CAP + 1, 1, 64)
.unwrap()
.with_ordering_stamps_kind(StampKind::SharedCounter)
.unwrap();
big.morph_to(RingShape::Mpsc).unwrap();
let rx = AdaptiveOrderedReceiver::new(&big, 0);
assert_eq!(rx.strategy(), "strict");
assert_eq!(big.ordering_mode(), Some(OrderingMode::MergeStrict));
let tsc = AdaptiveRing::create_anon(4, 1, 256)
.unwrap()
.with_ordering_stamps_kind(StampKind::Tsc)
.unwrap();
tsc.morph_to(RingShape::Mpsc).unwrap();
let rx = AdaptiveOrderedReceiver::new(&tsc, 0);
assert_eq!(rx.strategy(), "direct");
}
}