use std::str::FromStr;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant};
use std::{fmt, io};
use bytes::{Bytes, BytesMut};
use orbit_core::{Fleet, NetId64};
#[cfg(unix)]
mod bell;
mod capacity;
pub mod exchange;
mod layout;
pub use capacity::CapacityChanged;
#[cfg(feature = "tokio")]
mod poll;
mod table;
mod wake;
use layout::{
Direction, FLAG_CLOSE_WATCH, FLAG_FIN, FLAG_READER_GONE, FLAG_RESET,
GENERATION_MASK, SIDE_CLAIMED, SIDE_DEAD, SIDE_FREE, SIDE_RELEASED,
SLOT_BITS, SLOT_EMPTY, SLOT_LIVE, SLOT_MASK, Slot,
};
use orbit_core::NodeId;
#[cfg(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
))]
pub use orbit_core::readiness::Readiness;
use table::Table;
pub use table::{segment_size, segment_size_for};
use wake::Interest;
pub const STREAM_KIND: u8 = 246;
pub const STREAM_LANE_CAPACITY: usize = orbit_core::compile::usize_from_env(
option_env!("ORBIT_STREAM_LANE_CAPACITY"),
256,
);
pub const STREAM_BUFFER_BYTES: usize = orbit_core::compile::usize_from_env(
option_env!("ORBIT_STREAM_BUFFER_BYTES"),
64 * 1024,
);
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct StreamSpec {
pub kind: u8,
pub lane_capacity: usize,
pub buffer_bytes: usize,
}
impl StreamSpec {
pub const DEFAULT: Self =
Self::new(STREAM_KIND, STREAM_LANE_CAPACITY, STREAM_BUFFER_BYTES);
pub const fn new(
kind: u8,
lane_capacity: usize,
buffer_bytes: usize,
) -> Self {
Self { kind, lane_capacity, buffer_bytes }
}
pub const fn lane_bytes(&self) -> usize {
self.lane_capacity * 2 * self.buffer_bytes
}
fn validate(self) -> Result<()> {
if self.lane_capacity == 0
|| self.buffer_bytes == 0
|| !self.lane_capacity.is_power_of_two()
|| !self.buffer_bytes.is_power_of_two()
|| self.lane_capacity > 1 << SLOT_BITS
|| self.buffer_bytes > u32::MAX as usize
{
return Err(Error::Malformed(format!(
"stream spec kind={} lane_capacity={} buffer_bytes={}: both are powers of two, a lane holds at most {} streams and a ring at most {} bytes",
self.kind,
self.lane_capacity,
self.buffer_bytes,
1_usize << SLOT_BITS,
u32::MAX
)));
}
Ok(())
}
}
impl Default for StreamSpec {
fn default() -> Self {
Self::DEFAULT
}
}
pub type Result<T, E = Error> = std::result::Result<T, E>;
#[non_exhaustive]
#[derive(Debug)]
pub enum Error {
Stale(StreamId),
AlreadyClaimed(Ticket),
Full {
capacity: usize,
},
PayloadFull {
requested_slots: usize,
},
PayloadTooLarge {
len: usize,
capacity: usize,
},
Malformed(String),
Closed,
Reset,
PeerGone,
WouldBlock,
Io(io::Error),
}
impl fmt::Display for Error {
fn fmt(
&self,
f: &mut fmt::Formatter<'_>,
) -> fmt::Result {
match self {
Self::Stale(id) => write!(f, "stream {id} has ended"),
Self::AlreadyClaimed(ticket) => {
write!(f, "stream side {ticket} is already held")
}
Self::Full { capacity } => {
write!(f, "stream lane is full: capacity={capacity}")
}
Self::PayloadFull { requested_slots } => {
write!(
f,
"payload arena has no run of {requested_slots} free slots"
)
}
Self::PayloadTooLarge { len, capacity } => {
write!(
f,
"payload chunk is {len} bytes; arena lane capacity is {capacity}"
)
}
Self::Malformed(text) => {
write!(f, "not a stream address: {text:?}")
}
Self::Closed => f.write_str("stream write half is shut down"),
Self::Reset => {
f.write_str("stream direction was reset by the peer")
}
Self::PeerGone => f.write_str("stream peer reader is gone"),
Self::WouldBlock => f.write_str("stream operation would block"),
Self::Io(error) => write!(f, "Orbit stream io error: {error}"),
}
}
}
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(error) => Some(error),
_ => None,
}
}
}
impl From<io::Error> for Error {
fn from(value: io::Error) -> Self {
Self::Io(value)
}
}
impl From<Error> for io::Error {
fn from(value: Error) -> Self {
let kind = match &value {
Error::Stale(_) | Error::AlreadyClaimed(_) => {
io::ErrorKind::NotConnected
}
Error::Full { .. }
| Error::PayloadFull { .. }
| Error::PayloadTooLarge { .. } => io::ErrorKind::OutOfMemory,
Error::Malformed(_) => io::ErrorKind::InvalidInput,
Error::Closed => io::ErrorKind::BrokenPipe,
Error::Reset => io::ErrorKind::ConnectionReset,
Error::PeerGone => io::ErrorKind::BrokenPipe,
Error::WouldBlock => io::ErrorKind::WouldBlock,
Error::Io(error) => error.kind(),
};
io::Error::new(kind, value)
}
}
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct StreamId(NetId64);
impl StreamId {
pub const fn from_net_id(id: NetId64) -> Self {
Self(id)
}
pub const fn net_id(self) -> NetId64 {
self.0
}
pub const fn kind(self) -> u8 {
self.0.kind()
}
pub const fn node(self) -> u16 {
self.0.node()
}
pub const fn slot(self) -> u32 {
(self.0.counter() & SLOT_MASK) as u32
}
pub const fn generation(self) -> u32 {
(self.0.counter() >> SLOT_BITS) as u32
}
fn make(
kind: u8,
node: u16,
slot: u32,
generation: u32,
) -> Self {
Self(NetId64::make(
kind,
node,
((generation as u64) << SLOT_BITS) | (slot as u64 & SLOT_MASK),
))
}
}
impl fmt::Display for StreamId {
fn fmt(
&self,
f: &mut fmt::Formatter<'_>,
) -> fmt::Result {
self.0.fmt(f)
}
}
impl FromStr for StreamId {
type Err = Error;
fn from_str(text: &str) -> Result<Self> {
text.parse::<NetId64>()
.map(Self)
.map_err(|_| Error::Malformed(text.to_owned()))
}
}
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum Side {
A,
B,
}
impl Side {
const fn index(self) -> usize {
match self {
Self::A => 0,
Self::B => 1,
}
}
const fn write_direction(self) -> usize {
self.index()
}
const fn read_direction(self) -> usize {
1 - self.index()
}
const fn letter(self) -> char {
match self {
Self::A => 'a',
Self::B => 'b',
}
}
}
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Ticket {
pub id: StreamId,
pub side: Side,
pub epoch: u64,
}
impl fmt::Display for Ticket {
fn fmt(
&self,
f: &mut fmt::Formatter<'_>,
) -> fmt::Result {
write!(f, "{}/{}/{}", self.id, self.side.letter(), self.epoch)
}
}
impl FromStr for Ticket {
type Err = Error;
fn from_str(text: &str) -> Result<Self> {
let malformed = || Error::Malformed(text.to_owned());
let (rest, epoch) = text.rsplit_once('/').ok_or_else(malformed)?;
let (id, side) = rest.rsplit_once('/').ok_or_else(malformed)?;
let side = match side {
"a" => Side::A,
"b" => Side::B,
_ => return Err(malformed()),
};
Ok(Self {
id: id.parse().map_err(|_| malformed())?,
side,
epoch: epoch.parse().map_err(|_| malformed())?,
})
}
}
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Incarnation(u64);
impl Incarnation {
pub const fn new(value: u64) -> Self {
Self(value)
}
pub const fn get(self) -> u64 {
self.0
}
}
#[derive(Clone)]
pub struct Streams {
table: Arc<Table>,
}
impl Streams {
#[cfg(unix)]
pub fn with_group(
name: &str,
capacity: u16,
node: NodeId,
incarnation: Incarnation,
spec: StreamSpec,
access: orbit_core::shm::GroupAccess,
) -> Result<Self> {
spec.validate()?;
Ok(Self {
table: table::open_group(
name,
capacity,
node,
incarnation,
spec,
access,
false,
)?,
})
}
#[cfg(unix)]
pub fn prepare_group(
name: &str,
capacity: u16,
spec: StreamSpec,
access: orbit_core::shm::GroupAccess,
) -> Result<()> {
spec.validate()?;
table::open_group(
name,
capacity,
NodeId::ZERO,
Incarnation::new(0),
spec,
access,
true,
)?;
Ok(())
}
#[cfg(unix)]
pub fn node_dead_in_group(
name: &str,
capacity: u16,
spec: StreamSpec,
node: NodeId,
incarnation: Incarnation,
access: orbit_core::shm::GroupAccess,
) -> Result<()> {
spec.validate()?;
if node.get() >= capacity {
return Err(Error::Malformed(
"recovery node is outside the fleet".into(),
));
}
let table = table::open_group(
name,
capacity,
NodeId::ZERO,
Incarnation::new(0),
spec,
access,
true,
)?;
table.node_dead(node.get(), incarnation.get());
Ok(())
}
#[cfg(unix)]
pub fn node_dead_in_shm(
name: &str,
fleet_capacity: u16,
spec: StreamSpec,
node: NodeId,
incarnation: Incarnation,
) -> Result<bool> {
spec.validate()?;
if node.get() >= fleet_capacity {
return Err(Error::Malformed(
"recovery node is outside the fleet".into(),
));
}
let segment = orbit_core::shm::ring_segment_name(name, spec.kind);
match orbit_core::shm::ShmRegion::validate_existing(
&segment,
segment_size_for(fleet_capacity, spec),
)? {
orbit_core::shm::ShmValidation::Missing => Ok(false),
orbit_core::shm::ShmValidation::Valid { .. } => {
table::recover_shm(
name,
fleet_capacity,
spec,
node.get(),
incarnation.get(),
)?;
Ok(true)
}
}
}
pub fn new(
fleet: Arc<Fleet>,
incarnation: Incarnation,
) -> Result<Self> {
Self::with_spec(fleet, incarnation, StreamSpec::DEFAULT)
}
pub fn with_spec(
fleet: Arc<Fleet>,
incarnation: Incarnation,
spec: StreamSpec,
) -> Result<Self> {
spec.validate()?;
Ok(Self { table: table::open(&fleet, incarnation, spec)? })
}
pub fn kind(&self) -> u8 {
self.table.kind()
}
#[cfg(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
))]
pub fn readiness(&self) -> Result<Readiness> {
self.table.take_readiness()
}
pub fn epoch(&self) -> u64 {
self.table.epoch()
}
pub fn create(&self) -> Result<(Endpoint, Ticket)> {
let (index, generation) = self.table.allocate()?;
let id = StreamId::make(
self.table.kind(),
self.table.node(),
(index % self.table.geometry().lane_capacity) as u32,
generation,
);
let handle = Arc::new(Handle {
table: Arc::clone(&self.table),
index,
id,
side: Side::A,
});
Ok((
Endpoint::new(handle),
Ticket { id, side: Side::B, epoch: self.table.epoch() },
))
}
pub fn capacity_changed(&self) -> CapacityChanged {
CapacityChanged::new(Arc::clone(&self.table))
}
pub fn open(
&self,
ticket: Ticket,
) -> Result<Endpoint> {
let index = self.locate(ticket.id)?;
if ticket.epoch != self.table.epoch() {
return Err(Error::Stale(ticket.id));
}
let slot = &self.table.slots()[index];
if !slot.is(ticket.id.generation()) {
return Err(Error::Stale(ticket.id));
}
let claimed = &slot.claimed[ticket.side.index()];
if claimed
.compare_exchange(
SIDE_FREE,
SIDE_CLAIMED,
Ordering::SeqCst,
Ordering::SeqCst,
)
.is_err()
{
return Err(Error::AlreadyClaimed(ticket));
}
if !slot.is(ticket.id.generation()) {
let _ = claimed.compare_exchange(
SIDE_CLAIMED,
SIDE_FREE,
Ordering::SeqCst,
Ordering::SeqCst,
);
return Err(Error::Stale(ticket.id));
}
slot.incarnation[ticket.side.index()]
.store(self.table.incarnation(), Ordering::Release);
slot.node[ticket.side.index()]
.store(self.table.node(), Ordering::Release);
for direction in &slot.directions {
self.table.notify(index, slot, direction, table::Ring::Nobody);
}
Ok(Endpoint::new(Arc::new(Handle {
table: Arc::clone(&self.table),
index,
id: ticket.id,
side: ticket.side,
})))
}
pub fn rearm(
&self,
id: StreamId,
) -> Result<()> {
let slot = self.rearmable_slot(id)?;
let far = Side::B.index();
for direction in &slot.directions {
direction.clear();
}
slot.node[far].store(u16::MAX, Ordering::Relaxed);
slot.incarnation[far].store(0, Ordering::Relaxed);
slot.claimed[far]
.compare_exchange(
SIDE_RELEASED,
SIDE_FREE,
Ordering::SeqCst,
Ordering::SeqCst,
)
.map_err(|_| Error::Stale(id))?;
Ok(())
}
pub fn can_create(&self) -> bool {
self.table.has_capacity()
}
pub fn can_rearm(
&self,
id: StreamId,
) -> Result<bool> {
match self.rearmable_slot(id) {
Ok(_) => Ok(true),
Err(Error::WouldBlock) => Ok(false),
Err(error) => Err(error),
}
}
fn rearmable_slot(
&self,
id: StreamId,
) -> Result<&Slot> {
let index = self.locate(id)?;
let slot = &self.table.slots()[index];
if !slot.is(id.generation()) {
return Err(Error::Stale(id));
}
let far = Side::B.index();
match slot.claimed[far].load(Ordering::SeqCst) {
SIDE_RELEASED => {}
SIDE_CLAIMED => return Err(Error::WouldBlock),
_ => return Err(Error::Stale(id)),
}
Ok(slot)
}
pub fn is_live(
&self,
id: StreamId,
) -> bool {
self.locate(id)
.map(|index| self.table.slots()[index].is(id.generation()))
.unwrap_or(false)
}
pub fn offer(
&self,
ticket: Ticket,
to: NodeId,
) -> Result<()> {
let index = self.locate(ticket.id)?;
if ticket.epoch != self.table.epoch()
|| !self.table.slots()[index].is(ticket.id.generation())
{
return Err(Error::Stale(ticket.id));
}
self.table.offer(index, to.get())
}
pub fn take_offer(&self) -> Option<Ticket> {
while let Some(index) = self.table.take_offer() {
if let Some(ticket) = self.ticket_for(index) {
return Some(ticket);
}
}
None
}
#[cfg(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
))]
pub fn take_offer_timeout(
&self,
timeout: Duration,
) -> Result<Option<Ticket>> {
let deadline = Instant::now() + timeout;
loop {
if let Some(ticket) = self.take_offer() {
return Ok(Some(ticket));
}
let left = deadline.saturating_duration_since(Instant::now());
if left.is_zero() {
return Ok(self.take_offer());
}
if !self.table.wait_offer_timeout(left)? {
return Ok(self.take_offer());
}
}
}
pub fn blocking_take_offer(&self) -> Result<Ticket> {
loop {
if let Some(ticket) = self.take_offer() {
return Ok(ticket);
}
self.table.wait_offer()?;
}
}
pub fn poll_take_offer(
&self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<Ticket>> {
if let Some(ticket) = self.take_offer() {
return std::task::Poll::Ready(Ok(ticket));
}
if let Err(error) = self.table.register_offer(cx.waker()) {
return std::task::Poll::Ready(Err(error));
}
match self.take_offer() {
Some(ticket) => std::task::Poll::Ready(Ok(ticket)),
None => std::task::Poll::Pending,
}
}
pub fn owners(&self) -> Vec<(NodeId, Incarnation)> {
let mut seen: Vec<(NodeId, Incarnation)> = Vec::new();
for slot in self.table.slots() {
if slot.state.load(Ordering::Acquire) != SLOT_LIVE {
continue;
}
for side in 0..2 {
if slot.claimed[side].load(Ordering::Acquire) != SIDE_CLAIMED {
continue;
}
let holder = (
NodeId::new(slot.node[side].load(Ordering::Acquire)),
Incarnation::new(
slot.incarnation[side].load(Ordering::Acquire),
),
);
if !seen.contains(&holder) {
seen.push(holder);
}
}
}
seen.sort_by_key(|(node, incarnation)| (node.get(), incarnation.get()));
seen
}
pub fn node_dead(
&self,
node: NodeId,
incarnation: Incarnation,
) {
self.table.node_dead(node.get(), incarnation.get());
}
pub fn supersede(&self) {
self.table.supersede();
}
pub fn reset_all(&self) {
self.table.reset_all();
}
fn ticket_for(
&self,
index: usize,
) -> Option<Ticket> {
let slot = &self.table.slots()[index];
let generation = slot.generation.load(Ordering::Acquire);
if !slot.is(generation) {
return None;
}
let lane_capacity = self.table.geometry().lane_capacity;
let id = StreamId::make(
self.table.kind(),
(index / lane_capacity) as u16,
(index % lane_capacity) as u32,
generation,
);
Some(Ticket { id, side: Side::B, epoch: self.table.epoch() })
}
#[cfg(unix)]
pub fn unlink(&self) -> Result<()> {
self.table.unlink()
}
fn locate(
&self,
id: StreamId,
) -> Result<usize> {
let geometry = self.table.geometry();
if id.kind() != self.table.kind()
|| usize::from(id.node()) >= geometry.fleet_capacity
|| id.slot() as usize >= geometry.lane_capacity
|| id.generation() == 0
|| id.generation() > GENERATION_MASK
{
return Err(Error::Malformed(id.to_string()));
}
Ok(usize::from(id.node()) * geometry.lane_capacity + id.slot() as usize)
}
}
struct Handle {
table: Arc<Table>,
index: usize,
id: StreamId,
side: Side,
}
impl Handle {
fn peer_if(
&self,
ring: bool,
) -> table::Ring {
if ring {
table::Ring::Side(1 - self.side.index())
} else {
table::Ring::Nobody
}
}
fn slot(&self) -> Result<&Slot> {
let slot = &self.table.slots()[self.index];
if slot.is(self.id.generation()) {
Ok(slot)
} else {
Err(Error::Stale(self.id))
}
}
fn direction(
slot: &Slot,
direction: usize,
) -> &Direction {
&slot.directions[direction]
}
fn try_write(
&self,
buf: &[u8],
) -> Result<usize> {
let slot = self.slot()?;
let direction = Self::direction(slot, self.side.write_direction());
clear_exact_wait(direction);
self.write_bytes(slot, direction, buf)
}
#[cfg(feature = "tokio")]
fn try_write_ready(
&self,
buf: &[u8],
) -> Result<usize> {
let slot = self.slot()?;
let direction = Self::direction(slot, self.side.write_direction());
self.write_bytes(slot, direction, buf)
}
fn write_bytes(
&self,
slot: &Slot,
direction: &Direction,
buf: &[u8],
) -> Result<usize> {
let flags = direction.flags();
if flags & FLAG_RESET != 0 {
return Err(Error::Reset);
}
if flags & FLAG_FIN != 0 {
return Err(Error::Closed);
}
if flags & FLAG_READER_GONE != 0 {
return Err(Error::PeerGone);
}
let buffer_bytes = self.table.geometry().buffer_bytes;
let head = direction.head.load(Ordering::Relaxed);
let tail = direction.tail.load(Ordering::Acquire);
let free = buffer_bytes - (head - tail) as usize;
if buf.is_empty() {
return Ok(0);
}
if free == 0 {
return Err(Error::WouldBlock);
}
let len = buf.len().min(free);
let base = self.table.buffer(self.index, self.side.write_direction());
let start = (head as usize) & (buffer_bytes - 1);
let first = len.min(buffer_bytes - start);
unsafe {
std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
std::ptr::copy_nonoverlapping(
buf.as_ptr().add(first),
base,
len - first,
);
}
direction.head.store(head + len as u64, Ordering::SeqCst);
std::sync::atomic::fence(Ordering::SeqCst);
let drained = direction.tail.load(Ordering::SeqCst) >= head;
self.table.notify(self.index, slot, direction, self.peer_if(drained));
Ok(len)
}
fn try_write_exact(
&self,
buf: &[u8],
) -> Result<()> {
let slot = self.slot()?;
let direction = Self::direction(slot, self.side.write_direction());
clear_exact_wait(direction);
let flags = direction.flags();
if flags & FLAG_RESET != 0 {
return Err(Error::Reset);
}
if flags & FLAG_FIN != 0 {
return Err(Error::Closed);
}
if flags & FLAG_READER_GONE != 0 {
return Err(Error::PeerGone);
}
let buffer_bytes = self.table.geometry().buffer_bytes;
if buf.len() > buffer_bytes {
return Err(Error::Malformed(format!(
"atomic stream write is {} bytes; ring capacity is {buffer_bytes}",
buf.len()
)));
}
if buf.is_empty() {
return Ok(());
}
let head = direction.head.load(Ordering::Relaxed);
let tail = direction.tail.load(Ordering::Acquire);
if buffer_bytes - ((head - tail) as usize) < buf.len() {
return Err(Error::WouldBlock);
}
let base = self.table.buffer(self.index, self.side.write_direction());
let start = (head as usize) & (buffer_bytes - 1);
let first = buf.len().min(buffer_bytes - start);
unsafe {
std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
std::ptr::copy_nonoverlapping(
buf.as_ptr().add(first),
base,
buf.len() - first,
);
}
direction.head.store(head + buf.len() as u64, Ordering::SeqCst);
std::sync::atomic::fence(Ordering::SeqCst);
let drained = direction.tail.load(Ordering::SeqCst) >= head;
self.table.notify(self.index, slot, direction, self.peer_if(drained));
Ok(())
}
fn try_read(
&self,
buf: &mut [u8],
) -> Result<usize> {
let slot = self.slot()?;
let direction = Self::direction(slot, self.side.read_direction());
let flags = direction.flags();
if flags & FLAG_RESET != 0 {
return Err(Error::Reset);
}
let head = direction.head.load(Ordering::Acquire);
let tail = direction.tail.load(Ordering::Relaxed);
let available = (head - tail) as usize;
if buf.is_empty() {
return Ok(0);
}
if available == 0 {
if flags & FLAG_FIN != 0 {
return Ok(0);
}
return Err(Error::WouldBlock);
}
let len = buf.len().min(available);
let base = self.table.buffer(self.index, self.side.read_direction());
let buffer_bytes = self.table.geometry().buffer_bytes;
let start = (tail as usize) & (buffer_bytes - 1);
let first = len.min(buffer_bytes - start);
unsafe {
std::ptr::copy_nonoverlapping(
base.add(start),
buf.as_mut_ptr(),
first,
);
std::ptr::copy_nonoverlapping(
base,
buf.as_mut_ptr().add(first),
len - first,
);
}
direction.tail.store(tail + len as u64, Ordering::SeqCst);
std::sync::atomic::fence(Ordering::SeqCst);
let wake = writer_gained_capacity(
direction,
buffer_bytes,
tail,
tail + len as u64,
);
self.table.notify(self.index, slot, direction, self.peer_if(wake));
Ok(len)
}
fn try_read_exact(
&self,
buf: &mut [u8],
) -> Result<()> {
let slot = self.slot()?;
let direction = Self::direction(slot, self.side.read_direction());
let flags = direction.flags();
if flags & FLAG_RESET != 0 {
return Err(Error::Reset);
}
let buffer_bytes = self.table.geometry().buffer_bytes;
if buf.len() > buffer_bytes {
return Err(Error::Malformed(format!(
"atomic stream read is {} bytes; ring capacity is {buffer_bytes}",
buf.len()
)));
}
if buf.is_empty() {
return Ok(());
}
let head = direction.head.load(Ordering::Acquire);
let tail = direction.tail.load(Ordering::Relaxed);
let available = (head - tail) as usize;
if available < buf.len() {
if flags & FLAG_FIN != 0 {
return Err(Error::Malformed(format!(
"stream ended with {available} bytes of a {} byte frame",
buf.len()
)));
}
return Err(Error::WouldBlock);
}
let base = self.table.buffer(self.index, self.side.read_direction());
let start = (tail as usize) & (buffer_bytes - 1);
let first = buf.len().min(buffer_bytes - start);
unsafe {
std::ptr::copy_nonoverlapping(
base.add(start),
buf.as_mut_ptr(),
first,
);
std::ptr::copy_nonoverlapping(
base,
buf.as_mut_ptr().add(first),
buf.len() - first,
);
}
direction.tail.store(tail + buf.len() as u64, Ordering::SeqCst);
std::sync::atomic::fence(Ordering::SeqCst);
let wake = writer_gained_capacity(
direction,
buffer_bytes,
tail,
tail + buf.len() as u64,
);
self.table.notify(self.index, slot, direction, self.peer_if(wake));
Ok(())
}
fn set_flag(
&self,
direction_index: usize,
flag: u8,
) {
if let Ok(slot) = self.slot() {
let direction = Self::direction(slot, direction_index);
let before = direction.flags.fetch_or(flag, Ordering::SeqCst);
if before & flag == 0 {
let ring = flag != FLAG_READER_GONE
|| before & FLAG_CLOSE_WATCH != 0
|| writer_may_park(direction, before);
self.table.notify(
self.index,
slot,
direction,
self.peer_if(ring),
);
}
}
}
fn set_flag_quietly(
&self,
direction_index: usize,
flag: u8,
) {
if let Ok(slot) = self.slot() {
let direction = Self::direction(slot, direction_index);
if direction.flags.fetch_or(flag, Ordering::SeqCst) & flag == 0 {
self.table.notify(
self.index,
slot,
direction,
table::Ring::Nobody,
);
}
}
}
#[cfg(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
))]
fn wait_until_timeout(
&self,
direction_index: usize,
ready: impl Fn(&Direction) -> bool,
timeout: Duration,
) -> Result<bool> {
let deadline = Instant::now() + timeout;
loop {
let slot = self.slot()?;
let direction = Self::direction(slot, direction_index);
if ready(direction) {
return Ok(true);
}
let left = deadline.saturating_duration_since(Instant::now());
if left.is_zero() {
return Ok(false);
}
direction.waiters.fetch_add(1, Ordering::SeqCst);
let since = direction.changes.load(Ordering::SeqCst);
let outcome = if ready(direction) {
Ok(true)
} else {
wait_on_timeout(&direction.changes, since, left)
};
direction.waiters.fetch_sub(1, Ordering::SeqCst);
if !outcome? {
let slot = self.slot()?;
return Ok(ready(Self::direction(slot, direction_index)));
}
}
}
fn wait_until(
&self,
direction_index: usize,
ready: impl Fn(&Direction) -> bool,
) -> Result<()> {
loop {
let slot = self.slot()?;
let direction = Self::direction(slot, direction_index);
if ready(direction) {
return Ok(());
}
direction.waiters.fetch_add(1, Ordering::SeqCst);
let since = direction.changes.load(Ordering::SeqCst);
let outcome = if ready(direction) {
Ok(())
} else {
wait_on(&direction.changes, since)
};
direction.waiters.fetch_sub(1, Ordering::SeqCst);
outcome?;
}
}
fn poll_ready(
&self,
interest: Interest,
ready: impl Fn(&Direction) -> bool,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<()>> {
let slot = match self.slot() {
Ok(slot) => slot,
Err(error) => return std::task::Poll::Ready(Err(error)),
};
let direction = Self::direction(slot, interest.direction);
if ready(direction) {
return std::task::Poll::Ready(Ok(()));
}
if let Err(error) =
self.table.register(self.index, interest, cx.waker())
{
return std::task::Poll::Ready(Err(error));
}
std::sync::atomic::fence(Ordering::SeqCst);
if ready(direction) {
std::task::Poll::Ready(Ok(()))
} else {
std::task::Poll::Pending
}
}
}
fn clear_exact_wait(direction: &Direction) {
if direction.exact_needed.load(Ordering::Relaxed) != 0 {
direction.exact_needed.store(0, Ordering::SeqCst);
}
}
fn writer_gained_capacity(
direction: &Direction,
capacity: usize,
before_tail: u64,
after_tail: u64,
) -> bool {
let head = direction.head.load(Ordering::SeqCst);
let needed = direction.exact_needed.load(Ordering::SeqCst) as usize;
if needed == 0 {
return head - before_tail >= capacity as u64;
}
if direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0 {
return false;
}
let before = capacity.saturating_sub((head - before_tail) as usize);
let after = capacity.saturating_sub((head - after_tail) as usize);
before < needed && after >= needed
}
fn writer_may_park(
direction: &Direction,
before: u8,
) -> bool {
if before & (FLAG_FIN | FLAG_RESET) != 0 {
return false;
}
std::sync::atomic::fence(Ordering::SeqCst);
direction.head.load(Ordering::SeqCst)
!= direction.tail.load(Ordering::SeqCst)
}
fn readable(direction: &Direction) -> bool {
direction.head.load(Ordering::Acquire)
!= direction.tail.load(Ordering::Relaxed)
|| direction.flags() & (FLAG_FIN | FLAG_RESET) != 0
}
fn writable(
direction: &Direction,
buffer_bytes: usize,
) -> bool {
let queued = (direction.head.load(Ordering::Relaxed)
- direction.tail.load(Ordering::Acquire)) as usize;
queued < buffer_bytes
|| direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
}
fn writable_exact(
direction: &Direction,
buffer_bytes: usize,
needed: usize,
) -> bool {
let queued = (direction.head.load(Ordering::Relaxed)
- direction.tail.load(Ordering::Acquire)) as usize;
buffer_bytes.saturating_sub(queued) >= needed
|| direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
}
impl Drop for Handle {
fn drop(&mut self) {
let slot = &self.table.slots()[self.index];
if !slot.is(self.id.generation()) {
return;
}
if slot.claimed[self.side.index()]
.compare_exchange(
SIDE_CLAIMED,
SIDE_RELEASED,
Ordering::SeqCst,
Ordering::SeqCst,
)
.is_err()
{
return;
}
let other = slot.claimed[1 - self.side.index()].load(Ordering::SeqCst);
if other != SIDE_CLAIMED {
let _ = slot.state.compare_exchange(
SLOT_LIVE,
SLOT_EMPTY,
Ordering::SeqCst,
Ordering::SeqCst,
);
for direction in &slot.directions {
direction.changes.fetch_add(1, Ordering::SeqCst);
if direction.waiters.load(Ordering::SeqCst) > 0 {
wake_on(&direction.changes);
}
}
self.table.registry.wake(self.index);
}
if self.side == Side::B || other != SIDE_CLAIMED {
self.table.notify_capacity(self.index);
}
}
}
pub struct Endpoint {
read: ReadHalf,
write: WriteHalf,
}
impl Endpoint {
fn new(handle: Arc<Handle>) -> Self {
Self {
read: ReadHalf { handle: Arc::clone(&handle) },
write: WriteHalf { handle, shut: false },
}
}
pub fn id(&self) -> StreamId {
self.read.handle.id
}
pub fn side(&self) -> Side {
self.read.handle.side
}
pub fn split(self) -> (ReadHalf, WriteHalf) {
(self.read, self.write)
}
pub fn rejoin(
read: ReadHalf,
write: WriteHalf,
) -> Result<Self> {
if read.id() != write.id() {
return Err(Error::Stale(read.id()));
}
Ok(Self { read, write })
}
pub fn rearm(
&mut self,
streams: &Streams,
) -> Result<()> {
if self.side() != Side::A {
return Err(Error::Malformed(
"only the creator side of a stream can be rearmed".to_owned(),
));
}
streams.rearm(self.id())?;
self.write.shut = false;
Ok(())
}
pub fn try_read(
&self,
buf: &mut [u8],
) -> Result<usize> {
self.read.try_read(buf)
}
pub fn try_read_exact(
&self,
buf: &mut [u8],
) -> Result<()> {
self.read.try_read_exact(buf)
}
pub fn blocking_read(
&self,
buf: &mut [u8],
) -> Result<usize> {
self.read.blocking_read(buf)
}
pub fn blocking_read_chunk(
&self,
max: usize,
) -> Result<Bytes> {
self.read.blocking_read_chunk(max)
}
pub fn try_write(
&self,
buf: &[u8],
) -> Result<usize> {
self.write.try_write(buf)
}
pub fn try_write_exact(
&self,
buf: &[u8],
) -> Result<()> {
self.write.try_write_exact(buf)
}
pub fn blocking_write(
&self,
buf: &[u8],
) -> Result<usize> {
self.write.blocking_write(buf)
}
pub fn blocking_write_all(
&self,
buf: &[u8],
) -> Result<()> {
self.write.blocking_write_all(buf)
}
pub fn finish(&mut self) -> Result<()> {
self.write.finish()
}
pub fn reset(&mut self) {
self.write.reset()
}
}
pub struct ReadHalf {
handle: Arc<Handle>,
}
pub struct ResetSignal {
handle: Arc<Handle>,
}
impl ResetSignal {
pub fn poll_reset(
&self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<()>> {
self.handle.poll_ready(
Interest {
direction: self.handle.side.read_direction(),
writer: false,
reset: true,
},
|direction| direction.flags() & FLAG_RESET != 0,
cx,
)
}
}
impl ReadHalf {
pub fn reset_signal(&self) -> ResetSignal {
ResetSignal { handle: Arc::clone(&self.handle) }
}
pub fn id(&self) -> StreamId {
self.handle.id
}
pub fn try_read(
&self,
buf: &mut [u8],
) -> Result<usize> {
self.handle.try_read(buf)
}
pub fn try_read_exact(
&self,
buf: &mut [u8],
) -> Result<()> {
self.handle.try_read_exact(buf)
}
pub fn blocking_read(
&self,
buf: &mut [u8],
) -> Result<usize> {
loop {
match self.handle.try_read(buf) {
Err(Error::WouldBlock) => self.wait_readable()?,
other => return other,
}
}
}
pub fn stop_reading(&self) {
self.handle
.set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
}
pub fn blocking_read_chunk(
&self,
max: usize,
) -> Result<Bytes> {
let mut chunk = BytesMut::zeroed(max);
let len = self.blocking_read(&mut chunk)?;
chunk.truncate(len);
Ok(chunk.freeze())
}
#[cfg(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
))]
pub fn wait_readable_timeout(
&self,
timeout: Duration,
) -> Result<bool> {
self.handle.wait_until_timeout(
self.handle.side.read_direction(),
readable,
timeout,
)
}
pub fn wait_readable(&self) -> Result<()> {
self.handle.wait_until(self.handle.side.read_direction(), readable)
}
pub fn poll_readable(
&self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<()>> {
self.handle.poll_ready(
Interest {
direction: self.handle.side.read_direction(),
writer: false,
reset: false,
},
readable,
cx,
)
}
}
impl Drop for ReadHalf {
fn drop(&mut self) {
self.handle
.set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
}
}
pub struct WriteHalf {
handle: Arc<Handle>,
shut: bool,
}
struct CloseWatch<'a>(&'a Handle);
impl Drop for CloseWatch<'_> {
fn drop(&mut self) {
if let Ok(slot) = self.0.slot() {
Handle::direction(slot, self.0.side.write_direction())
.flags
.fetch_and(!FLAG_CLOSE_WATCH, Ordering::SeqCst);
}
}
}
impl WriteHalf {
pub async fn peer_closed(&mut self) -> Result<()> {
let slot = self.handle.slot()?;
let direction =
Handle::direction(slot, self.handle.side.write_direction());
direction.flags.fetch_or(FLAG_CLOSE_WATCH, Ordering::SeqCst);
let _watch = CloseWatch(&self.handle);
std::future::poll_fn(|cx| {
self.handle.poll_ready(
Interest {
direction: self.handle.side.write_direction(),
writer: true,
reset: false,
},
|direction| {
direction.flags() & (FLAG_RESET | FLAG_READER_GONE) != 0
},
cx,
)
})
.await
}
pub fn poll_peer_dead(
&self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<()>> {
let slot = match self.handle.slot() {
Ok(slot) => slot,
Err(error) => return std::task::Poll::Ready(Err(error)),
};
let peer = &slot.claimed[1 - self.handle.side.index()];
if peer.load(Ordering::SeqCst) == SIDE_DEAD {
return std::task::Poll::Ready(Ok(()));
}
let interest = Interest {
direction: self.handle.side.write_direction(),
writer: true,
reset: false,
};
if let Err(error) =
self.handle.table.register(self.handle.index, interest, cx.waker())
{
return std::task::Poll::Ready(Err(error));
}
if peer.load(Ordering::SeqCst) == SIDE_DEAD {
std::task::Poll::Ready(Ok(()))
} else {
std::task::Poll::Pending
}
}
pub fn id(&self) -> StreamId {
self.handle.id
}
pub fn try_write(
&self,
buf: &[u8],
) -> Result<usize> {
self.handle.try_write(buf)
}
pub fn try_write_exact(
&self,
buf: &[u8],
) -> Result<()> {
self.handle.try_write_exact(buf)
}
pub fn blocking_write(
&self,
buf: &[u8],
) -> Result<usize> {
loop {
match self.handle.try_write(buf) {
Err(Error::WouldBlock) => self.wait_writable()?,
other => return other,
}
}
}
pub fn blocking_write_all(
&self,
mut buf: &[u8],
) -> Result<()> {
while !buf.is_empty() {
let written = self.blocking_write(buf)?;
buf = &buf[written..];
}
Ok(())
}
pub fn finish(&mut self) -> Result<()> {
self.shut = true;
self.handle.set_flag(self.handle.side.write_direction(), FLAG_FIN);
Ok(())
}
pub fn finish_quietly(&mut self) {
self.shut = true;
self.handle
.set_flag_quietly(self.handle.side.write_direction(), FLAG_FIN);
}
pub fn reset(&mut self) {
self.shut = true;
self.handle.set_flag(self.handle.side.write_direction(), FLAG_RESET);
}
#[cfg(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
))]
pub fn wait_writable_timeout(
&self,
timeout: Duration,
) -> Result<bool> {
let buffer_bytes = self.handle.table.geometry().buffer_bytes;
self.handle.wait_until_timeout(
self.handle.side.write_direction(),
|direction| {
clear_exact_wait(direction);
writable(direction, buffer_bytes)
},
timeout,
)
}
pub fn wait_writable(&self) -> Result<()> {
let buffer_bytes = self.handle.table.geometry().buffer_bytes;
self.handle.wait_until(
self.handle.side.write_direction(),
|direction| {
clear_exact_wait(direction);
writable(direction, buffer_bytes)
},
)
}
pub fn wait_writable_exact(
&self,
needed: usize,
) -> Result<()> {
let buffer_bytes = self.handle.table.geometry().buffer_bytes;
if needed > buffer_bytes {
return Err(Error::Malformed(format!(
"atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
)));
}
self.handle.wait_until(
self.handle.side.write_direction(),
|direction| {
clear_exact_wait(direction);
writable_exact(direction, buffer_bytes, needed)
},
)
}
pub fn poll_writable(
&self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<()>> {
let buffer_bytes = self.handle.table.geometry().buffer_bytes;
self.handle.poll_ready(
Interest {
direction: self.handle.side.write_direction(),
writer: true,
reset: false,
},
|direction| {
clear_exact_wait(direction);
writable(direction, buffer_bytes)
},
cx,
)
}
pub fn poll_writable_exact(
&self,
needed: usize,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<()>> {
let buffer_bytes = self.handle.table.geometry().buffer_bytes;
let slot = match self.handle.slot() {
Ok(slot) => slot,
Err(error) => return std::task::Poll::Ready(Err(error)),
};
let direction =
Handle::direction(slot, self.handle.side.write_direction());
if needed > buffer_bytes {
clear_exact_wait(direction);
return std::task::Poll::Ready(Err(Error::Malformed(format!(
"atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
))));
}
direction.exact_needed.store(needed as u32, Ordering::SeqCst);
let result = self.handle.poll_ready(
Interest {
direction: self.handle.side.write_direction(),
writer: true,
reset: false,
},
|direction| writable_exact(direction, buffer_bytes, needed),
cx,
);
if result.is_ready() {
clear_exact_wait(direction);
}
result
}
}
impl Drop for WriteHalf {
fn drop(&mut self) {
if !self.shut {
self.handle
.set_flag(self.handle.side.write_direction(), FLAG_RESET);
}
}
}
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub(crate) fn waits_supported() -> bool {
orbit_core::sync::supported()
}
#[cfg(not(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
)))]
pub(crate) fn waits_supported() -> bool {
false
}
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub(crate) fn wait_on_timeout(
word: &AtomicU32,
expected: u32,
timeout: Duration,
) -> Result<bool> {
orbit_core::sync::wait_word_timeout(word, expected, timeout)
.map_err(Error::Io)
}
pub(crate) fn wait_on(
word: &AtomicU32,
expected: u32,
) -> Result<()> {
orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
}
#[cfg(not(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
)))]
pub(crate) fn wait_on(
_word: &AtomicU32,
_expected: u32,
) -> Result<()> {
Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"orbit-stream needs a platform that can wait on a shared word",
)))
}
pub(crate) fn wake_on(word: &AtomicU32) {
#[cfg(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
))]
let _ = orbit_core::sync::wake_word(word);
#[cfg(not(any(
target_os = "linux",
target_os = "freebsd",
target_os = "macos"
)))]
let _ = word;
}
pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(|error| error.into_inner())
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use orbit_core::{Fleet, NodeId};
use super::{
Error, Incarnation, STREAM_BUFFER_BYTES, STREAM_LANE_CAPACITY, Side,
Streams, Ticket,
};
fn streams(name: &'static str) -> Streams {
Streams::new(
Arc::new(Fleet::join(name, 2).unwrap()),
Incarnation::new(1),
)
.unwrap()
}
#[test]
fn bytes_cross_in_order_and_end_cleanly() {
let streams = streams("stream-order");
let (mut a, ticket) = streams.create().unwrap();
assert_eq!(ticket.side, Side::B);
assert_eq!(ticket.id.node(), 0);
assert_eq!(ticket.epoch, streams.epoch());
let b = streams.open(ticket).unwrap();
a.blocking_write_all(b"hello, ").unwrap();
a.blocking_write_all(b"world").unwrap();
a.finish().unwrap();
let mut out = Vec::new();
let mut buf = [0_u8; 4];
loop {
let n = b.blocking_read(&mut buf).unwrap();
if n == 0 {
break;
}
out.extend_from_slice(&buf[..n]);
}
assert_eq!(out, b"hello, world");
b.blocking_write_all(b"back").unwrap();
assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"back");
}
#[test]
fn a_ticket_prints_and_parses() {
let streams = streams("stream-ticket");
let (_a, ticket) = streams.create().unwrap();
let text = ticket.to_string();
assert!(text.contains("/b/"), "{text}");
assert_eq!(text.parse::<Ticket>().unwrap(), ticket);
assert!(matches!("nope".parse::<Ticket>(), Err(Error::Malformed(_))));
assert!(matches!("x/b".parse::<Ticket>(), Err(Error::Malformed(_))));
}
#[test]
fn a_side_is_held_once_and_a_released_stream_goes_stale() {
let streams = streams("stream-claim");
let (a, ticket) = streams.create().unwrap();
let b = streams.open(ticket).unwrap();
assert!(matches!(streams.open(ticket), Err(Error::AlreadyClaimed(_))));
assert!(streams.is_live(ticket.id));
drop(a);
let mut buf = [0_u8; 1];
assert!(matches!(b.blocking_read(&mut buf), Err(Error::Reset)));
drop(b);
assert!(!streams.is_live(ticket.id));
assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
for _ in 0..STREAM_LANE_CAPACITY - 1 {
drop(streams.create().unwrap());
}
let (_a2, ticket2) = streams.create().unwrap();
assert_eq!(ticket2.id.slot(), ticket.id.slot());
assert_ne!(ticket2.id.generation(), ticket.id.generation());
assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
}
#[test]
fn a_reset_starts_a_new_epoch_and_refuses_old_tickets() {
let streams = streams("stream-epoch");
let (_a, ticket) = streams.create().unwrap();
let before = streams.epoch();
streams.reset_all();
assert!(streams.epoch() > before);
let (_a2, ticket2) = streams.create().unwrap();
assert_eq!(ticket2.id, ticket.id);
assert_ne!(ticket2.epoch, ticket.epoch);
assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
assert!(streams.open(ticket2).is_ok());
}
#[test]
fn a_slot_whose_generation_runs_out_is_not_reused() {
let streams = streams("stream-exhaust");
let limit = super::layout::GENERATION_LIMIT as usize;
for _ in 0..STREAM_LANE_CAPACITY * limit {
drop(streams.create().unwrap());
}
assert!(matches!(streams.create(), Err(Error::Full { .. })));
streams.reset_all();
let (_a, ticket) = streams.create().unwrap();
assert_eq!(ticket.id.generation(), 1);
}
#[test]
fn a_full_ring_blocks_the_writer_and_wraps_correctly() {
let streams = streams("stream-wrap");
let (a, ticket) = streams.create().unwrap();
let b = streams.open(ticket).unwrap();
let filler = vec![7_u8; STREAM_BUFFER_BYTES];
assert_eq!(a.try_write(&filler).unwrap(), STREAM_BUFFER_BYTES);
assert!(matches!(a.try_write(b"x"), Err(Error::WouldBlock)));
let mut buf = vec![0_u8; STREAM_BUFFER_BYTES - 5];
assert_eq!(b.blocking_read(&mut buf).unwrap(), STREAM_BUFFER_BYTES - 5);
let tail =
(0..STREAM_BUFFER_BYTES).map(|i| i as u8).collect::<Vec<_>>();
assert_eq!(a.try_write(&tail).unwrap(), STREAM_BUFFER_BYTES - 5);
let mut rest = vec![0_u8; 5];
assert_eq!(b.blocking_read(&mut rest).unwrap(), 5);
assert_eq!(rest, vec![7_u8; 5]);
let mut wrapped = vec![0_u8; STREAM_BUFFER_BYTES];
let mut got = 0;
while got < STREAM_BUFFER_BYTES - 5 {
got += b.blocking_read(&mut wrapped[got..]).unwrap();
}
assert_eq!(&wrapped[..got], &tail[..got]);
}
#[test]
fn exact_frames_never_publish_or_consume_a_prefix() {
let streams = streams("stream-exact-frame");
let (a, ticket) = streams.create().unwrap();
let b = streams.open(ticket).unwrap();
let filler = vec![7_u8; STREAM_BUFFER_BYTES - 32];
a.try_write_exact(&filler).unwrap();
assert!(matches!(
a.try_write_exact(&[9_u8; 64]),
Err(Error::WouldBlock)
));
let mut consumed = vec![0_u8; filler.len()];
b.try_read_exact(&mut consumed).unwrap();
assert_eq!(consumed, filler);
assert!(matches!(b.try_read(&mut [0_u8; 1]), Err(Error::WouldBlock)));
a.try_write_exact(&[3_u8; 32]).unwrap();
assert!(matches!(
b.try_read_exact(&mut [0_u8; 64]),
Err(Error::WouldBlock)
));
let mut frame = [0_u8; 32];
b.try_read_exact(&mut frame).unwrap();
assert_eq!(frame, [3_u8; 32]);
}
#[test]
fn a_blocked_reader_is_woken_by_a_write_in_another_thread() {
let streams = streams("stream-wake");
let (a, ticket) = streams.create().unwrap();
let b = streams.open(ticket).unwrap();
let reader = std::thread::spawn(move || {
let mut buf = [0_u8; 8];
let n = b.blocking_read(&mut buf).unwrap();
buf[..n].to_vec()
});
std::thread::sleep(Duration::from_millis(30));
a.blocking_write_all(b"ping").unwrap();
assert_eq!(reader.join().unwrap(), b"ping");
}
#[test]
fn a_reader_gone_fails_the_writer_and_a_lane_fills_up() {
let streams = streams("stream-gone");
let (a, ticket) = streams.create().unwrap();
let b = streams.open(ticket).unwrap();
let (b_read, _b_write) = b.split();
drop(b_read);
assert!(matches!(a.blocking_write(b"nobody"), Err(Error::PeerGone)));
let mut held = vec![a];
for _ in 1..STREAM_LANE_CAPACITY {
held.push(streams.create().unwrap().0);
}
assert!(matches!(streams.create(), Err(Error::Full { .. })));
}
#[test]
fn a_death_report_ends_that_incarnations_sides_only() {
let streams = streams("stream-dead");
let (a, ticket) = streams.create().unwrap();
let b = streams.open(ticket).unwrap();
b.blocking_write_all(b"partial").unwrap();
streams.node_dead(NodeId::ZERO, Incarnation::new(2));
assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"partial");
let reader = std::thread::spawn(move || {
let mut buf = [0_u8; 8];
a.blocking_read(&mut buf).map(|_| ())
});
std::thread::sleep(Duration::from_millis(30));
streams.node_dead(NodeId::ZERO, Incarnation::new(1));
assert!(matches!(
reader.join().unwrap(),
Err(Error::Reset | Error::Stale(_))
));
assert!(matches!(
b.try_write(b"x"),
Err(Error::Reset | Error::Stale(_))
));
assert!(matches!(
streams.open(ticket),
Err(Error::AlreadyClaimed(_) | Error::Stale(_))
));
drop(b);
assert!(!streams.is_live(ticket.id));
}
#[test]
fn an_offer_is_taken_once_and_a_stale_offer_is_skipped() {
let streams = streams("stream-offer");
assert!(streams.take_offer().is_none());
let (a, ticket) = streams.create().unwrap();
streams.offer(ticket, NodeId::ZERO).unwrap();
assert_eq!(streams.take_offer(), Some(ticket));
assert!(streams.take_offer().is_none());
let (gone, stale) = streams.create().unwrap();
streams.offer(stale, NodeId::ZERO).unwrap();
drop(gone);
assert!(streams.take_offer().is_none());
let taker = streams.clone();
let waiter = std::thread::spawn(move || taker.blocking_take_offer());
std::thread::sleep(Duration::from_millis(30));
streams.offer(ticket, NodeId::ZERO).unwrap();
assert_eq!(waiter.join().unwrap().unwrap(), ticket);
drop(a);
}
#[test]
fn a_contested_claim_goes_to_exactly_one_opener() {
let streams = streams("stream-race");
for _ in 0..50 {
let (_a, ticket) = streams.create().unwrap();
let openers = (0..4)
.map(|_| {
let streams = streams.clone();
std::thread::spawn(move || streams.open(ticket).is_ok())
})
.collect::<Vec<_>>();
let won = openers
.into_iter()
.map(|opener| opener.join().unwrap())
.filter(|won| *won)
.count();
assert_eq!(won, 1);
}
}
#[test]
fn random_sizes_concatenate_exactly() {
let streams = streams("stream-random");
let (mut a, ticket) = streams.create().unwrap();
let b = streams.open(ticket).unwrap();
let payload = (0..300_000_u32)
.map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
.collect::<Vec<_>>();
let expected = payload.clone();
let writer = std::thread::spawn(move || {
let mut rest = &payload[..];
let mut step = 1_usize;
while !rest.is_empty() {
let len = step.min(rest.len());
a.blocking_write_all(&rest[..len]).unwrap();
rest = &rest[len..];
step = (step * 7 + 13) % 9_001 + 1;
}
a.finish().unwrap();
});
let mut out = Vec::new();
let mut step = 3_usize;
loop {
let chunk = b.blocking_read_chunk(step).unwrap();
if chunk.is_empty() {
break;
}
out.extend_from_slice(&chunk);
step = (step * 5 + 11) % 4_099 + 1;
}
writer.join().unwrap();
assert_eq!(out, expected);
}
}