1use std::str::FromStr;
37use std::sync::atomic::{AtomicU32, Ordering};
38use std::sync::{Arc, Mutex, MutexGuard};
39use std::time::{Duration, Instant};
40use std::{fmt, io};
41
42use bytes::{Bytes, BytesMut};
43use orbit_core::{Fleet, NetId64};
44
45#[cfg(unix)]
46mod bell;
47pub mod exchange;
48mod layout;
49#[cfg(feature = "tokio")]
50mod poll;
51mod table;
52mod wake;
53
54use layout::{
55 Direction, FLAG_FIN, FLAG_READER_GONE, FLAG_RESET, GENERATION_MASK,
56 SIDE_CLAIMED, SIDE_DEAD, SIDE_FREE, SIDE_RELEASED, SLOT_BITS, SLOT_EMPTY,
57 SLOT_LIVE, SLOT_MASK, Slot,
58};
59use orbit_core::NodeId;
60#[cfg(any(
61 target_os = "linux",
62 target_os = "freebsd",
63 target_os = "macos"
64))]
65pub use orbit_core::readiness::Readiness;
66use table::Table;
67pub use table::{segment_size, segment_size_for};
68use wake::Interest;
69
70pub const STREAM_KIND: u8 = 246;
73pub const STREAM_LANE_CAPACITY: usize = orbit_core::compile::usize_from_env(
81 option_env!("ORBIT_STREAM_LANE_CAPACITY"),
82 256,
83);
84pub const STREAM_BUFFER_BYTES: usize = orbit_core::compile::usize_from_env(
91 option_env!("ORBIT_STREAM_BUFFER_BYTES"),
92 64 * 1024,
93);
94
95#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
107pub struct StreamSpec {
108 pub kind: u8,
110 pub lane_capacity: usize,
113 pub buffer_bytes: usize,
115}
116
117impl StreamSpec {
118 pub const DEFAULT: Self =
119 Self::new(STREAM_KIND, STREAM_LANE_CAPACITY, STREAM_BUFFER_BYTES);
120
121 pub const fn new(
122 kind: u8,
123 lane_capacity: usize,
124 buffer_bytes: usize,
125 ) -> Self {
126 Self { kind, lane_capacity, buffer_bytes }
127 }
128
129 pub const fn lane_bytes(&self) -> usize {
133 self.lane_capacity * 2 * self.buffer_bytes
134 }
135
136 fn validate(self) -> Result<()> {
139 if self.lane_capacity == 0
140 || self.buffer_bytes == 0
141 || !self.lane_capacity.is_power_of_two()
142 || !self.buffer_bytes.is_power_of_two()
143 || self.lane_capacity > 1 << SLOT_BITS
144 || self.buffer_bytes > u32::MAX as usize
145 {
146 return Err(Error::Malformed(format!(
147 "stream spec kind={} lane_capacity={} buffer_bytes={}: both are powers of two, a lane holds at most {} streams and a ring at most {} bytes",
148 self.kind,
149 self.lane_capacity,
150 self.buffer_bytes,
151 1_usize << SLOT_BITS,
152 u32::MAX
153 )));
154 }
155 Ok(())
156 }
157}
158
159impl Default for StreamSpec {
160 fn default() -> Self {
161 Self::DEFAULT
162 }
163}
164
165pub type Result<T, E = Error> = std::result::Result<T, E>;
166
167#[non_exhaustive]
171#[derive(Debug)]
172pub enum Error {
173 Stale(StreamId),
175 AlreadyClaimed(Ticket),
177 Full {
179 capacity: usize,
180 },
181 PayloadFull {
183 requested_slots: usize,
184 },
185 PayloadTooLarge {
187 len: usize,
188 capacity: usize,
189 },
190 Malformed(String),
192 Closed,
194 Reset,
196 PeerGone,
198 WouldBlock,
200 Io(io::Error),
201}
202
203impl fmt::Display for Error {
204 fn fmt(
205 &self,
206 f: &mut fmt::Formatter<'_>,
207 ) -> fmt::Result {
208 match self {
209 Self::Stale(id) => write!(f, "stream {id} has ended"),
210 Self::AlreadyClaimed(ticket) => {
211 write!(f, "stream side {ticket} is already held")
212 }
213 Self::Full { capacity } => {
214 write!(f, "stream lane is full: capacity={capacity}")
215 }
216 Self::PayloadFull { requested_slots } => {
217 write!(
218 f,
219 "payload arena has no run of {requested_slots} free slots"
220 )
221 }
222 Self::PayloadTooLarge { len, capacity } => {
223 write!(
224 f,
225 "payload chunk is {len} bytes; arena lane capacity is {capacity}"
226 )
227 }
228 Self::Malformed(text) => {
229 write!(f, "not a stream address: {text:?}")
230 }
231 Self::Closed => f.write_str("stream write half is shut down"),
232 Self::Reset => {
233 f.write_str("stream direction was reset by the peer")
234 }
235 Self::PeerGone => f.write_str("stream peer reader is gone"),
236 Self::WouldBlock => f.write_str("stream operation would block"),
237 Self::Io(error) => write!(f, "Orbit stream io error: {error}"),
238 }
239 }
240}
241
242impl std::error::Error for Error {
243 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
244 match self {
245 Self::Io(error) => Some(error),
246 _ => None,
247 }
248 }
249}
250
251impl From<io::Error> for Error {
252 fn from(value: io::Error) -> Self {
253 Self::Io(value)
254 }
255}
256
257impl From<Error> for io::Error {
258 fn from(value: Error) -> Self {
259 let kind = match &value {
260 Error::Stale(_) | Error::AlreadyClaimed(_) => {
261 io::ErrorKind::NotConnected
262 }
263 Error::Full { .. }
264 | Error::PayloadFull { .. }
265 | Error::PayloadTooLarge { .. } => io::ErrorKind::OutOfMemory,
266 Error::Malformed(_) => io::ErrorKind::InvalidInput,
267 Error::Closed => io::ErrorKind::BrokenPipe,
268 Error::Reset => io::ErrorKind::ConnectionReset,
269 Error::PeerGone => io::ErrorKind::BrokenPipe,
270 Error::WouldBlock => io::ErrorKind::WouldBlock,
271 Error::Io(error) => error.kind(),
272 };
273 io::Error::new(kind, value)
274 }
275}
276
277#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
280pub struct StreamId(NetId64);
281
282impl StreamId {
283 pub const fn from_net_id(id: NetId64) -> Self {
284 Self(id)
285 }
286
287 pub const fn net_id(self) -> NetId64 {
288 self.0
289 }
290
291 pub const fn kind(self) -> u8 {
292 self.0.kind()
293 }
294
295 pub const fn node(self) -> u16 {
297 self.0.node()
298 }
299
300 pub const fn slot(self) -> u32 {
301 (self.0.counter() & SLOT_MASK) as u32
302 }
303
304 pub const fn generation(self) -> u32 {
305 (self.0.counter() >> SLOT_BITS) as u32
306 }
307
308 fn make(
309 kind: u8,
310 node: u16,
311 slot: u32,
312 generation: u32,
313 ) -> Self {
314 Self(NetId64::make(
315 kind,
316 node,
317 ((generation as u64) << SLOT_BITS) | (slot as u64 & SLOT_MASK),
318 ))
319 }
320}
321
322impl fmt::Display for StreamId {
323 fn fmt(
324 &self,
325 f: &mut fmt::Formatter<'_>,
326 ) -> fmt::Result {
327 self.0.fmt(f)
328 }
329}
330
331impl FromStr for StreamId {
332 type Err = Error;
333
334 fn from_str(text: &str) -> Result<Self> {
335 text.parse::<NetId64>()
336 .map(Self)
337 .map_err(|_| Error::Malformed(text.to_owned()))
338 }
339}
340
341#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
343pub enum Side {
344 A,
345 B,
346}
347
348impl Side {
349 const fn index(self) -> usize {
350 match self {
351 Self::A => 0,
352 Self::B => 1,
353 }
354 }
355
356 const fn write_direction(self) -> usize {
358 self.index()
359 }
360
361 const fn read_direction(self) -> usize {
363 1 - self.index()
364 }
365
366 const fn letter(self) -> char {
367 match self {
368 Self::A => 'a',
369 Self::B => 'b',
370 }
371 }
372}
373
374#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
383pub struct Ticket {
384 pub id: StreamId,
385 pub side: Side,
386 pub epoch: u64,
387}
388
389impl fmt::Display for Ticket {
390 fn fmt(
391 &self,
392 f: &mut fmt::Formatter<'_>,
393 ) -> fmt::Result {
394 write!(f, "{}/{}/{}", self.id, self.side.letter(), self.epoch)
395 }
396}
397
398impl FromStr for Ticket {
399 type Err = Error;
400
401 fn from_str(text: &str) -> Result<Self> {
402 let malformed = || Error::Malformed(text.to_owned());
403 let (rest, epoch) = text.rsplit_once('/').ok_or_else(malformed)?;
404 let (id, side) = rest.rsplit_once('/').ok_or_else(malformed)?;
405 let side = match side {
406 "a" => Side::A,
407 "b" => Side::B,
408 _ => return Err(malformed()),
409 };
410 Ok(Self {
411 id: id.parse().map_err(|_| malformed())?,
412 side,
413 epoch: epoch.parse().map_err(|_| malformed())?,
414 })
415 }
416}
417
418#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
424pub struct Incarnation(u64);
425
426impl Incarnation {
427 pub const fn new(value: u64) -> Self {
428 Self(value)
429 }
430
431 pub const fn get(self) -> u64 {
432 self.0
433 }
434}
435
436#[derive(Clone)]
439pub struct Streams {
440 table: Arc<Table>,
441}
442
443impl Streams {
444 #[cfg(unix)]
447 pub fn with_group(
448 name: &str,
449 capacity: u16,
450 node: NodeId,
451 incarnation: Incarnation,
452 spec: StreamSpec,
453 access: orbit_core::shm::GroupAccess,
454 ) -> Result<Self> {
455 spec.validate()?;
456 Ok(Self {
457 table: table::open_group(
458 name,
459 capacity,
460 node,
461 incarnation,
462 spec,
463 access,
464 false,
465 )?,
466 })
467 }
468
469 #[cfg(unix)]
471 pub fn prepare_group(
472 name: &str,
473 capacity: u16,
474 spec: StreamSpec,
475 access: orbit_core::shm::GroupAccess,
476 ) -> Result<()> {
477 spec.validate()?;
478 table::open_group(
479 name,
480 capacity,
481 NodeId::ZERO,
482 Incarnation::new(0),
483 spec,
484 access,
485 true,
486 )?;
487 Ok(())
488 }
489
490 #[cfg(unix)]
492 pub fn node_dead_in_group(
493 name: &str,
494 capacity: u16,
495 spec: StreamSpec,
496 node: NodeId,
497 incarnation: Incarnation,
498 access: orbit_core::shm::GroupAccess,
499 ) -> Result<()> {
500 spec.validate()?;
501 if node.get() >= capacity {
502 return Err(Error::Malformed(
503 "recovery node is outside the fleet".into(),
504 ));
505 }
506 let table = table::open_group(
507 name,
508 capacity,
509 NodeId::ZERO,
510 Incarnation::new(0),
511 spec,
512 access,
513 true,
514 )?;
515 table.node_dead(node.get(), incarnation.get());
516 Ok(())
517 }
518
519 #[cfg(unix)]
523 pub fn node_dead_in_shm(
524 name: &str,
525 fleet_capacity: u16,
526 spec: StreamSpec,
527 node: NodeId,
528 incarnation: Incarnation,
529 ) -> Result<bool> {
530 spec.validate()?;
531 if node.get() >= fleet_capacity {
532 return Err(Error::Malformed(
533 "recovery node is outside the fleet".into(),
534 ));
535 }
536 let segment = orbit_core::shm::ring_segment_name(name, spec.kind);
537 match orbit_core::shm::ShmRegion::validate_existing(
538 &segment,
539 segment_size_for(fleet_capacity, spec),
540 )? {
541 orbit_core::shm::ShmValidation::Missing => Ok(false),
542 orbit_core::shm::ShmValidation::Valid { .. } => {
543 table::recover_shm(
544 name,
545 fleet_capacity,
546 spec,
547 node.get(),
548 incarnation.get(),
549 )?;
550 Ok(true)
551 }
552 }
553 }
554
555 pub fn new(
559 fleet: Arc<Fleet>,
560 incarnation: Incarnation,
561 ) -> Result<Self> {
562 Self::with_spec(fleet, incarnation, StreamSpec::DEFAULT)
563 }
564
565 pub fn with_spec(
570 fleet: Arc<Fleet>,
571 incarnation: Incarnation,
572 spec: StreamSpec,
573 ) -> Result<Self> {
574 spec.validate()?;
575 Ok(Self { table: table::open(&fleet, incarnation, spec)? })
576 }
577
578 pub fn kind(&self) -> u8 {
580 self.table.kind()
581 }
582
583 #[cfg(any(
595 target_os = "linux",
596 target_os = "freebsd",
597 target_os = "macos"
598 ))]
599 pub fn readiness(&self) -> Result<Readiness> {
600 self.table.take_readiness()
601 }
602
603 pub fn epoch(&self) -> u64 {
605 self.table.epoch()
606 }
607
608 pub fn create(&self) -> Result<(Endpoint, Ticket)> {
611 let (index, generation) = self.table.allocate()?;
612 let id = StreamId::make(
613 self.table.kind(),
614 self.table.node(),
615 (index % self.table.geometry().lane_capacity) as u32,
616 generation,
617 );
618 let handle = Arc::new(Handle {
619 table: Arc::clone(&self.table),
620 index,
621 id,
622 side: Side::A,
623 });
624 Ok((
625 Endpoint::new(handle),
626 Ticket { id, side: Side::B, epoch: self.table.epoch() },
627 ))
628 }
629
630 pub fn open(
632 &self,
633 ticket: Ticket,
634 ) -> Result<Endpoint> {
635 let index = self.locate(ticket.id)?;
636 if ticket.epoch != self.table.epoch() {
637 return Err(Error::Stale(ticket.id));
638 }
639 let slot = &self.table.slots()[index];
640 if !slot.is(ticket.id.generation()) {
641 return Err(Error::Stale(ticket.id));
642 }
643 let claimed = &slot.claimed[ticket.side.index()];
644 if claimed
645 .compare_exchange(
646 SIDE_FREE,
647 SIDE_CLAIMED,
648 Ordering::SeqCst,
649 Ordering::SeqCst,
650 )
651 .is_err()
652 {
653 return Err(Error::AlreadyClaimed(ticket));
654 }
655 if !slot.is(ticket.id.generation()) {
658 let _ = claimed.compare_exchange(
659 SIDE_CLAIMED,
660 SIDE_FREE,
661 Ordering::SeqCst,
662 Ordering::SeqCst,
663 );
664 return Err(Error::Stale(ticket.id));
665 }
666 slot.incarnation[ticket.side.index()]
667 .store(self.table.incarnation(), Ordering::Release);
668 slot.node[ticket.side.index()]
669 .store(self.table.node(), Ordering::Release);
670 for direction in &slot.directions {
677 self.table.notify(index, slot, direction, table::Ring::Nobody);
678 }
679 Ok(Endpoint::new(Arc::new(Handle {
680 table: Arc::clone(&self.table),
681 index,
682 id: ticket.id,
683 side: ticket.side,
684 })))
685 }
686
687 pub fn rearm(
723 &self,
724 id: StreamId,
725 ) -> Result<()> {
726 let index = self.locate(id)?;
727 let slot = &self.table.slots()[index];
728 if !slot.is(id.generation()) {
729 return Err(Error::Stale(id));
730 }
731 let far = Side::B.index();
732 match slot.claimed[far].load(Ordering::SeqCst) {
733 SIDE_RELEASED => {}
734 SIDE_CLAIMED => return Err(Error::WouldBlock),
735 _ => return Err(Error::Stale(id)),
736 }
737 for direction in &slot.directions {
738 direction.clear();
739 }
740 slot.node[far].store(u16::MAX, Ordering::Relaxed);
741 slot.incarnation[far].store(0, Ordering::Relaxed);
742 slot.claimed[far]
746 .compare_exchange(
747 SIDE_RELEASED,
748 SIDE_FREE,
749 Ordering::SeqCst,
750 Ordering::SeqCst,
751 )
752 .map_err(|_| Error::Stale(id))?;
753 Ok(())
754 }
755
756 pub fn is_live(
758 &self,
759 id: StreamId,
760 ) -> bool {
761 self.locate(id)
762 .map(|index| self.table.slots()[index].is(id.generation()))
763 .unwrap_or(false)
764 }
765
766 pub fn offer(
772 &self,
773 ticket: Ticket,
774 to: NodeId,
775 ) -> Result<()> {
776 let index = self.locate(ticket.id)?;
777 if ticket.epoch != self.table.epoch()
778 || !self.table.slots()[index].is(ticket.id.generation())
779 {
780 return Err(Error::Stale(ticket.id));
781 }
782 self.table.offer(index, to.get())
783 }
784
785 pub fn take_offer(&self) -> Option<Ticket> {
788 while let Some(index) = self.table.take_offer() {
789 if let Some(ticket) = self.ticket_for(index) {
790 return Some(ticket);
791 }
792 }
793 None
794 }
795
796 #[cfg(any(
800 target_os = "linux",
801 target_os = "freebsd",
802 target_os = "macos"
803 ))]
804 pub fn take_offer_timeout(
805 &self,
806 timeout: Duration,
807 ) -> Result<Option<Ticket>> {
808 let deadline = Instant::now() + timeout;
809 loop {
810 if let Some(ticket) = self.take_offer() {
811 return Ok(Some(ticket));
812 }
813 let left = deadline.saturating_duration_since(Instant::now());
814 if left.is_zero() {
815 return Ok(self.take_offer());
816 }
817 if !self.table.wait_offer_timeout(left)? {
818 return Ok(self.take_offer());
819 }
820 }
821 }
822
823 pub fn blocking_take_offer(&self) -> Result<Ticket> {
825 loop {
826 if let Some(ticket) = self.take_offer() {
827 return Ok(ticket);
828 }
829 self.table.wait_offer()?;
830 }
831 }
832
833 pub fn poll_take_offer(
836 &self,
837 cx: &mut std::task::Context<'_>,
838 ) -> std::task::Poll<Result<Ticket>> {
839 if let Some(ticket) = self.take_offer() {
840 return std::task::Poll::Ready(Ok(ticket));
841 }
842 if let Err(error) = self.table.register_offer(cx.waker()) {
843 return std::task::Poll::Ready(Err(error));
844 }
845 match self.take_offer() {
846 Some(ticket) => std::task::Poll::Ready(Ok(ticket)),
847 None => std::task::Poll::Pending,
848 }
849 }
850
851 pub fn owners(&self) -> Vec<(NodeId, Incarnation)> {
857 let mut seen: Vec<(NodeId, Incarnation)> = Vec::new();
858 for slot in self.table.slots() {
859 if slot.state.load(Ordering::Acquire) != SLOT_LIVE {
860 continue;
861 }
862 for side in 0..2 {
863 if slot.claimed[side].load(Ordering::Acquire) != SIDE_CLAIMED {
864 continue;
865 }
866 let holder = (
867 NodeId::new(slot.node[side].load(Ordering::Acquire)),
868 Incarnation::new(
869 slot.incarnation[side].load(Ordering::Acquire),
870 ),
871 );
872 if !seen.contains(&holder) {
873 seen.push(holder);
874 }
875 }
876 }
877 seen.sort_by_key(|(node, incarnation)| (node.get(), incarnation.get()));
878 seen
879 }
880
881 pub fn node_dead(
888 &self,
889 node: NodeId,
890 incarnation: Incarnation,
891 ) {
892 self.table.node_dead(node.get(), incarnation.get());
893 }
894
895 pub fn supersede(&self) {
907 self.table.supersede();
908 }
909
910 pub fn reset_all(&self) {
913 self.table.reset_all();
914 }
915
916 fn ticket_for(
917 &self,
918 index: usize,
919 ) -> Option<Ticket> {
920 let slot = &self.table.slots()[index];
921 let generation = slot.generation.load(Ordering::Acquire);
922 if !slot.is(generation) {
923 return None;
924 }
925 let lane_capacity = self.table.geometry().lane_capacity;
926 let id = StreamId::make(
927 self.table.kind(),
928 (index / lane_capacity) as u16,
929 (index % lane_capacity) as u32,
930 generation,
931 );
932 Some(Ticket { id, side: Side::B, epoch: self.table.epoch() })
933 }
934
935 #[cfg(unix)]
938 pub fn unlink(&self) -> Result<()> {
939 self.table.unlink()
940 }
941
942 fn locate(
943 &self,
944 id: StreamId,
945 ) -> Result<usize> {
946 let geometry = self.table.geometry();
947 if id.kind() != self.table.kind()
948 || usize::from(id.node()) >= geometry.fleet_capacity
949 || id.slot() as usize >= geometry.lane_capacity
950 || id.generation() == 0
951 || id.generation() > GENERATION_MASK
952 {
953 return Err(Error::Malformed(id.to_string()));
954 }
955 Ok(usize::from(id.node()) * geometry.lane_capacity + id.slot() as usize)
956 }
957}
958
959struct Handle {
961 table: Arc<Table>,
962 index: usize,
963 id: StreamId,
964 side: Side,
965}
966
967impl Handle {
968 fn peer_if(
972 &self,
973 ring: bool,
974 ) -> table::Ring {
975 if ring {
976 table::Ring::Side(1 - self.side.index())
977 } else {
978 table::Ring::Nobody
979 }
980 }
981
982 fn slot(&self) -> Result<&Slot> {
983 let slot = &self.table.slots()[self.index];
984 if slot.is(self.id.generation()) {
985 Ok(slot)
986 } else {
987 Err(Error::Stale(self.id))
988 }
989 }
990
991 fn direction(
992 slot: &Slot,
993 direction: usize,
994 ) -> &Direction {
995 &slot.directions[direction]
996 }
997
998 fn try_write(
1000 &self,
1001 buf: &[u8],
1002 ) -> Result<usize> {
1003 let slot = self.slot()?;
1004 let direction = Self::direction(slot, self.side.write_direction());
1005 let flags = direction.flags();
1006 if flags & FLAG_RESET != 0 {
1007 return Err(Error::Reset);
1008 }
1009 if flags & FLAG_FIN != 0 {
1010 return Err(Error::Closed);
1011 }
1012 if flags & FLAG_READER_GONE != 0 {
1013 return Err(Error::PeerGone);
1014 }
1015 let buffer_bytes = self.table.geometry().buffer_bytes;
1016 let head = direction.head.load(Ordering::Relaxed);
1017 let tail = direction.tail.load(Ordering::Acquire);
1018 let free = buffer_bytes - (head - tail) as usize;
1019 if buf.is_empty() {
1020 return Ok(0);
1021 }
1022 if free == 0 {
1023 return Err(Error::WouldBlock);
1024 }
1025 let len = buf.len().min(free);
1026 let base = self.table.buffer(self.index, self.side.write_direction());
1027 let start = (head as usize) & (buffer_bytes - 1);
1028 let first = len.min(buffer_bytes - start);
1029 unsafe {
1032 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
1033 std::ptr::copy_nonoverlapping(
1034 buf.as_ptr().add(first),
1035 base,
1036 len - first,
1037 );
1038 }
1039 direction.head.store(head + len as u64, Ordering::SeqCst);
1040 std::sync::atomic::fence(Ordering::SeqCst);
1050 let drained = direction.tail.load(Ordering::SeqCst) >= head;
1051 self.table.notify(self.index, slot, direction, self.peer_if(drained));
1052 Ok(len)
1053 }
1054
1055 fn try_write_exact(
1057 &self,
1058 buf: &[u8],
1059 ) -> Result<()> {
1060 let slot = self.slot()?;
1061 let direction = Self::direction(slot, self.side.write_direction());
1062 let flags = direction.flags();
1063 if flags & FLAG_RESET != 0 {
1064 return Err(Error::Reset);
1065 }
1066 if flags & FLAG_FIN != 0 {
1067 return Err(Error::Closed);
1068 }
1069 if flags & FLAG_READER_GONE != 0 {
1070 return Err(Error::PeerGone);
1071 }
1072 let buffer_bytes = self.table.geometry().buffer_bytes;
1073 if buf.len() > buffer_bytes {
1074 return Err(Error::Malformed(format!(
1075 "atomic stream write is {} bytes; ring capacity is {buffer_bytes}",
1076 buf.len()
1077 )));
1078 }
1079 if buf.is_empty() {
1080 return Ok(());
1081 }
1082 let head = direction.head.load(Ordering::Relaxed);
1083 let tail = direction.tail.load(Ordering::Acquire);
1084 if buffer_bytes - ((head - tail) as usize) < buf.len() {
1085 return Err(Error::WouldBlock);
1086 }
1087 let base = self.table.buffer(self.index, self.side.write_direction());
1088 let start = (head as usize) & (buffer_bytes - 1);
1089 let first = buf.len().min(buffer_bytes - start);
1090 unsafe {
1092 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
1093 std::ptr::copy_nonoverlapping(
1094 buf.as_ptr().add(first),
1095 base,
1096 buf.len() - first,
1097 );
1098 }
1099 direction.head.store(head + buf.len() as u64, Ordering::SeqCst);
1100 std::sync::atomic::fence(Ordering::SeqCst);
1101 let drained = direction.tail.load(Ordering::SeqCst) >= head;
1102 self.table.notify(self.index, slot, direction, self.peer_if(drained));
1103 Ok(())
1104 }
1105
1106 fn try_read(
1109 &self,
1110 buf: &mut [u8],
1111 ) -> Result<usize> {
1112 let slot = self.slot()?;
1113 let direction = Self::direction(slot, self.side.read_direction());
1114 let flags = direction.flags();
1115 if flags & FLAG_RESET != 0 {
1116 return Err(Error::Reset);
1117 }
1118 let head = direction.head.load(Ordering::Acquire);
1119 let tail = direction.tail.load(Ordering::Relaxed);
1120 let available = (head - tail) as usize;
1121 if buf.is_empty() {
1122 return Ok(0);
1123 }
1124 if available == 0 {
1125 if flags & FLAG_FIN != 0 {
1126 return Ok(0);
1127 }
1128 return Err(Error::WouldBlock);
1129 }
1130 let len = buf.len().min(available);
1131 let base = self.table.buffer(self.index, self.side.read_direction());
1132 let buffer_bytes = self.table.geometry().buffer_bytes;
1133 let start = (tail as usize) & (buffer_bytes - 1);
1134 let first = len.min(buffer_bytes - start);
1135 unsafe {
1139 std::ptr::copy_nonoverlapping(
1140 base.add(start),
1141 buf.as_mut_ptr(),
1142 first,
1143 );
1144 std::ptr::copy_nonoverlapping(
1145 base,
1146 buf.as_mut_ptr().add(first),
1147 len - first,
1148 );
1149 }
1150 direction.tail.store(tail + len as u64, Ordering::SeqCst);
1151 std::sync::atomic::fence(Ordering::SeqCst);
1155 let head_now = direction.head.load(Ordering::SeqCst);
1156 let was_full = head_now - tail >= buffer_bytes as u64;
1157 self.table.notify(self.index, slot, direction, self.peer_if(was_full));
1158 Ok(len)
1159 }
1160
1161 fn try_read_exact(
1163 &self,
1164 buf: &mut [u8],
1165 ) -> Result<()> {
1166 let slot = self.slot()?;
1167 let direction = Self::direction(slot, self.side.read_direction());
1168 let flags = direction.flags();
1169 if flags & FLAG_RESET != 0 {
1170 return Err(Error::Reset);
1171 }
1172 let buffer_bytes = self.table.geometry().buffer_bytes;
1173 if buf.len() > buffer_bytes {
1174 return Err(Error::Malformed(format!(
1175 "atomic stream read is {} bytes; ring capacity is {buffer_bytes}",
1176 buf.len()
1177 )));
1178 }
1179 if buf.is_empty() {
1180 return Ok(());
1181 }
1182 let head = direction.head.load(Ordering::Acquire);
1183 let tail = direction.tail.load(Ordering::Relaxed);
1184 let available = (head - tail) as usize;
1185 if available < buf.len() {
1186 if flags & FLAG_FIN != 0 {
1187 return Err(Error::Malformed(format!(
1188 "stream ended with {available} bytes of a {} byte frame",
1189 buf.len()
1190 )));
1191 }
1192 return Err(Error::WouldBlock);
1193 }
1194 let base = self.table.buffer(self.index, self.side.read_direction());
1195 let start = (tail as usize) & (buffer_bytes - 1);
1196 let first = buf.len().min(buffer_bytes - start);
1197 unsafe {
1199 std::ptr::copy_nonoverlapping(
1200 base.add(start),
1201 buf.as_mut_ptr(),
1202 first,
1203 );
1204 std::ptr::copy_nonoverlapping(
1205 base,
1206 buf.as_mut_ptr().add(first),
1207 buf.len() - first,
1208 );
1209 }
1210 direction.tail.store(tail + buf.len() as u64, Ordering::SeqCst);
1211 std::sync::atomic::fence(Ordering::SeqCst);
1212 let head_now = direction.head.load(Ordering::SeqCst);
1213 let was_full = head_now - tail >= buffer_bytes as u64;
1214 self.table.notify(self.index, slot, direction, self.peer_if(was_full));
1215 Ok(())
1216 }
1217
1218 fn set_flag(
1219 &self,
1220 direction_index: usize,
1221 flag: u8,
1222 ) {
1223 if let Ok(slot) = self.slot() {
1224 let direction = Self::direction(slot, direction_index);
1225 let before = direction.flags.fetch_or(flag, Ordering::SeqCst);
1226 if before & flag == 0 {
1227 let ring = flag != FLAG_READER_GONE
1228 || writer_may_park(direction, before);
1229 self.table.notify(
1230 self.index,
1231 slot,
1232 direction,
1233 self.peer_if(ring),
1234 );
1235 }
1236 }
1237 }
1238
1239 #[cfg(any(
1242 target_os = "linux",
1243 target_os = "freebsd",
1244 target_os = "macos"
1245 ))]
1246 fn wait_until_timeout(
1247 &self,
1248 direction_index: usize,
1249 ready: impl Fn(&Direction) -> bool,
1250 timeout: Duration,
1251 ) -> Result<bool> {
1252 let deadline = Instant::now() + timeout;
1253 loop {
1254 let slot = self.slot()?;
1255 let direction = Self::direction(slot, direction_index);
1256 if ready(direction) {
1257 return Ok(true);
1258 }
1259 let left = deadline.saturating_duration_since(Instant::now());
1260 if left.is_zero() {
1261 return Ok(false);
1262 }
1263 direction.waiters.fetch_add(1, Ordering::SeqCst);
1264 let since = direction.changes.load(Ordering::SeqCst);
1265 let outcome = if ready(direction) {
1266 Ok(true)
1267 } else {
1268 wait_on_timeout(&direction.changes, since, left)
1269 };
1270 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1271 if !outcome? {
1272 let slot = self.slot()?;
1275 return Ok(ready(Self::direction(slot, direction_index)));
1276 }
1277 }
1278 }
1279
1280 fn wait_until(
1281 &self,
1282 direction_index: usize,
1283 ready: impl Fn(&Direction) -> bool,
1284 ) -> Result<()> {
1285 loop {
1286 let slot = self.slot()?;
1287 let direction = Self::direction(slot, direction_index);
1288 if ready(direction) {
1289 return Ok(());
1290 }
1291 direction.waiters.fetch_add(1, Ordering::SeqCst);
1292 let since = direction.changes.load(Ordering::SeqCst);
1293 let outcome = if ready(direction) {
1294 Ok(())
1295 } else {
1296 wait_on(&direction.changes, since)
1297 };
1298 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1299 outcome?;
1300 }
1301 }
1302
1303 fn poll_ready(
1306 &self,
1307 interest: Interest,
1308 ready: impl Fn(&Direction) -> bool,
1309 cx: &mut std::task::Context<'_>,
1310 ) -> std::task::Poll<Result<()>> {
1311 let slot = match self.slot() {
1312 Ok(slot) => slot,
1313 Err(error) => return std::task::Poll::Ready(Err(error)),
1314 };
1315 let direction = Self::direction(slot, interest.direction);
1316 if ready(direction) {
1317 return std::task::Poll::Ready(Ok(()));
1318 }
1319 if let Err(error) =
1320 self.table.register(self.index, interest, cx.waker())
1321 {
1322 return std::task::Poll::Ready(Err(error));
1323 }
1324 std::sync::atomic::fence(Ordering::SeqCst);
1327 if ready(direction) {
1328 std::task::Poll::Ready(Ok(()))
1329 } else {
1330 std::task::Poll::Pending
1331 }
1332 }
1333}
1334
1335fn writer_may_park(
1349 direction: &Direction,
1350 before: u8,
1351) -> bool {
1352 if before & (FLAG_FIN | FLAG_RESET) != 0 {
1353 return false;
1354 }
1355 std::sync::atomic::fence(Ordering::SeqCst);
1356 direction.head.load(Ordering::SeqCst)
1357 != direction.tail.load(Ordering::SeqCst)
1358}
1359
1360fn readable(direction: &Direction) -> bool {
1361 direction.head.load(Ordering::Acquire)
1362 != direction.tail.load(Ordering::Relaxed)
1363 || direction.flags() & (FLAG_FIN | FLAG_RESET) != 0
1364}
1365
1366fn writable(
1367 direction: &Direction,
1368 buffer_bytes: usize,
1369) -> bool {
1370 let queued = (direction.head.load(Ordering::Relaxed)
1371 - direction.tail.load(Ordering::Acquire)) as usize;
1372 queued < buffer_bytes
1373 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1374}
1375
1376fn writable_exact(
1377 direction: &Direction,
1378 buffer_bytes: usize,
1379 needed: usize,
1380) -> bool {
1381 let queued = (direction.head.load(Ordering::Relaxed)
1382 - direction.tail.load(Ordering::Acquire)) as usize;
1383 buffer_bytes.saturating_sub(queued) >= needed
1384 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1385}
1386
1387impl Drop for Handle {
1388 fn drop(&mut self) {
1392 let slot = &self.table.slots()[self.index];
1393 if !slot.is(self.id.generation()) {
1394 return;
1395 }
1396 if slot.claimed[self.side.index()]
1399 .compare_exchange(
1400 SIDE_CLAIMED,
1401 SIDE_RELEASED,
1402 Ordering::SeqCst,
1403 Ordering::SeqCst,
1404 )
1405 .is_err()
1406 {
1407 return;
1408 }
1409 let other = slot.claimed[1 - self.side.index()].load(Ordering::SeqCst);
1413 if other != SIDE_CLAIMED {
1414 let _ = slot.state.compare_exchange(
1415 SLOT_LIVE,
1416 SLOT_EMPTY,
1417 Ordering::SeqCst,
1418 Ordering::SeqCst,
1419 );
1420 for direction in &slot.directions {
1421 direction.changes.fetch_add(1, Ordering::SeqCst);
1424 if direction.waiters.load(Ordering::SeqCst) > 0 {
1425 wake_on(&direction.changes);
1426 }
1427 }
1428 self.table.registry.wake(self.index);
1429 }
1430 }
1431}
1432
1433pub struct Endpoint {
1437 read: ReadHalf,
1438 write: WriteHalf,
1439}
1440
1441impl Endpoint {
1442 fn new(handle: Arc<Handle>) -> Self {
1443 Self {
1444 read: ReadHalf { handle: Arc::clone(&handle) },
1445 write: WriteHalf { handle, shut: false },
1446 }
1447 }
1448
1449 pub fn id(&self) -> StreamId {
1450 self.read.handle.id
1451 }
1452
1453 pub fn side(&self) -> Side {
1454 self.read.handle.side
1455 }
1456
1457 pub fn split(self) -> (ReadHalf, WriteHalf) {
1458 (self.read, self.write)
1459 }
1460
1461 pub fn rejoin(
1472 read: ReadHalf,
1473 write: WriteHalf,
1474 ) -> Result<Self> {
1475 if read.id() != write.id() {
1476 return Err(Error::Stale(read.id()));
1477 }
1478 Ok(Self { read, write })
1479 }
1480
1481 pub fn rearm(
1494 &mut self,
1495 streams: &Streams,
1496 ) -> Result<()> {
1497 if self.side() != Side::A {
1498 return Err(Error::Malformed(
1499 "only the creator side of a stream can be rearmed".to_owned(),
1500 ));
1501 }
1502 streams.rearm(self.id())?;
1503 self.write.shut = false;
1504 Ok(())
1505 }
1506
1507 pub fn try_read(
1508 &self,
1509 buf: &mut [u8],
1510 ) -> Result<usize> {
1511 self.read.try_read(buf)
1512 }
1513
1514 pub fn try_read_exact(
1515 &self,
1516 buf: &mut [u8],
1517 ) -> Result<()> {
1518 self.read.try_read_exact(buf)
1519 }
1520
1521 pub fn blocking_read(
1522 &self,
1523 buf: &mut [u8],
1524 ) -> Result<usize> {
1525 self.read.blocking_read(buf)
1526 }
1527
1528 pub fn blocking_read_chunk(
1529 &self,
1530 max: usize,
1531 ) -> Result<Bytes> {
1532 self.read.blocking_read_chunk(max)
1533 }
1534
1535 pub fn try_write(
1536 &self,
1537 buf: &[u8],
1538 ) -> Result<usize> {
1539 self.write.try_write(buf)
1540 }
1541
1542 pub fn try_write_exact(
1543 &self,
1544 buf: &[u8],
1545 ) -> Result<()> {
1546 self.write.try_write_exact(buf)
1547 }
1548
1549 pub fn blocking_write(
1550 &self,
1551 buf: &[u8],
1552 ) -> Result<usize> {
1553 self.write.blocking_write(buf)
1554 }
1555
1556 pub fn blocking_write_all(
1557 &self,
1558 buf: &[u8],
1559 ) -> Result<()> {
1560 self.write.blocking_write_all(buf)
1561 }
1562
1563 pub fn finish(&mut self) -> Result<()> {
1564 self.write.finish()
1565 }
1566
1567 pub fn reset(&mut self) {
1568 self.write.reset()
1569 }
1570}
1571
1572pub struct ReadHalf {
1575 handle: Arc<Handle>,
1576}
1577
1578impl ReadHalf {
1579 pub fn id(&self) -> StreamId {
1580 self.handle.id
1581 }
1582
1583 pub fn try_read(
1586 &self,
1587 buf: &mut [u8],
1588 ) -> Result<usize> {
1589 self.handle.try_read(buf)
1590 }
1591
1592 pub fn try_read_exact(
1595 &self,
1596 buf: &mut [u8],
1597 ) -> Result<()> {
1598 self.handle.try_read_exact(buf)
1599 }
1600
1601 pub fn blocking_read(
1604 &self,
1605 buf: &mut [u8],
1606 ) -> Result<usize> {
1607 loop {
1608 match self.handle.try_read(buf) {
1609 Err(Error::WouldBlock) => self.wait_readable()?,
1610 other => return other,
1611 }
1612 }
1613 }
1614
1615 pub fn stop_reading(&self) {
1623 self.handle
1624 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1625 }
1626
1627 pub fn blocking_read_chunk(
1629 &self,
1630 max: usize,
1631 ) -> Result<Bytes> {
1632 let mut chunk = BytesMut::zeroed(max);
1633 let len = self.blocking_read(&mut chunk)?;
1634 chunk.truncate(len);
1635 Ok(chunk.freeze())
1636 }
1637
1638 #[cfg(any(
1640 target_os = "linux",
1641 target_os = "freebsd",
1642 target_os = "macos"
1643 ))]
1644 pub fn wait_readable_timeout(
1645 &self,
1646 timeout: Duration,
1647 ) -> Result<bool> {
1648 self.handle.wait_until_timeout(
1649 self.handle.side.read_direction(),
1650 readable,
1651 timeout,
1652 )
1653 }
1654
1655 pub fn wait_readable(&self) -> Result<()> {
1657 self.handle.wait_until(self.handle.side.read_direction(), readable)
1658 }
1659
1660 pub fn poll_readable(
1665 &self,
1666 cx: &mut std::task::Context<'_>,
1667 ) -> std::task::Poll<Result<()>> {
1668 self.handle.poll_ready(
1669 Interest {
1670 direction: self.handle.side.read_direction(),
1671 writer: false,
1672 },
1673 readable,
1674 cx,
1675 )
1676 }
1677}
1678
1679impl Drop for ReadHalf {
1680 fn drop(&mut self) {
1681 self.handle
1682 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1683 }
1684}
1685
1686pub struct WriteHalf {
1689 handle: Arc<Handle>,
1690 shut: bool,
1691}
1692
1693impl WriteHalf {
1694 pub fn poll_peer_dead(
1700 &self,
1701 cx: &mut std::task::Context<'_>,
1702 ) -> std::task::Poll<Result<()>> {
1703 let slot = match self.handle.slot() {
1704 Ok(slot) => slot,
1705 Err(error) => return std::task::Poll::Ready(Err(error)),
1706 };
1707 let peer = &slot.claimed[1 - self.handle.side.index()];
1708 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1709 return std::task::Poll::Ready(Ok(()));
1710 }
1711 let interest = Interest {
1712 direction: self.handle.side.write_direction(),
1713 writer: true,
1714 };
1715 if let Err(error) =
1716 self.handle.table.register(self.handle.index, interest, cx.waker())
1717 {
1718 return std::task::Poll::Ready(Err(error));
1719 }
1720 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1723 std::task::Poll::Ready(Ok(()))
1724 } else {
1725 std::task::Poll::Pending
1726 }
1727 }
1728
1729 pub fn id(&self) -> StreamId {
1730 self.handle.id
1731 }
1732
1733 pub fn try_write(
1735 &self,
1736 buf: &[u8],
1737 ) -> Result<usize> {
1738 self.handle.try_write(buf)
1739 }
1740
1741 pub fn try_write_exact(
1744 &self,
1745 buf: &[u8],
1746 ) -> Result<()> {
1747 self.handle.try_write_exact(buf)
1748 }
1749
1750 pub fn blocking_write(
1752 &self,
1753 buf: &[u8],
1754 ) -> Result<usize> {
1755 loop {
1756 match self.handle.try_write(buf) {
1757 Err(Error::WouldBlock) => self.wait_writable()?,
1758 other => return other,
1759 }
1760 }
1761 }
1762
1763 pub fn blocking_write_all(
1766 &self,
1767 mut buf: &[u8],
1768 ) -> Result<()> {
1769 while !buf.is_empty() {
1770 let written = self.blocking_write(buf)?;
1771 buf = &buf[written..];
1772 }
1773 Ok(())
1774 }
1775
1776 pub fn finish(&mut self) -> Result<()> {
1780 self.shut = true;
1781 self.handle.set_flag(self.handle.side.write_direction(), FLAG_FIN);
1782 Ok(())
1783 }
1784
1785 pub fn reset(&mut self) {
1788 self.shut = true;
1789 self.handle.set_flag(self.handle.side.write_direction(), FLAG_RESET);
1790 }
1791
1792 #[cfg(any(
1796 target_os = "linux",
1797 target_os = "freebsd",
1798 target_os = "macos"
1799 ))]
1800 pub fn wait_writable_timeout(
1801 &self,
1802 timeout: Duration,
1803 ) -> Result<bool> {
1804 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1805 self.handle.wait_until_timeout(
1806 self.handle.side.write_direction(),
1807 |direction| writable(direction, buffer_bytes),
1808 timeout,
1809 )
1810 }
1811
1812 pub fn wait_writable(&self) -> Result<()> {
1814 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1815 self.handle
1816 .wait_until(self.handle.side.write_direction(), |direction| {
1817 writable(direction, buffer_bytes)
1818 })
1819 }
1820
1821 pub fn wait_writable_exact(
1823 &self,
1824 needed: usize,
1825 ) -> Result<()> {
1826 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1827 if needed > buffer_bytes {
1828 return Err(Error::Malformed(format!(
1829 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
1830 )));
1831 }
1832 self.handle
1833 .wait_until(self.handle.side.write_direction(), |direction| {
1834 writable_exact(direction, buffer_bytes, needed)
1835 })
1836 }
1837
1838 pub fn poll_writable(
1843 &self,
1844 cx: &mut std::task::Context<'_>,
1845 ) -> std::task::Poll<Result<()>> {
1846 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1847 self.handle.poll_ready(
1848 Interest {
1849 direction: self.handle.side.write_direction(),
1850 writer: true,
1851 },
1852 |direction| writable(direction, buffer_bytes),
1853 cx,
1854 )
1855 }
1856
1857 pub fn poll_writable_exact(
1859 &self,
1860 needed: usize,
1861 cx: &mut std::task::Context<'_>,
1862 ) -> std::task::Poll<Result<()>> {
1863 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1864 if needed > buffer_bytes {
1865 return std::task::Poll::Ready(Err(Error::Malformed(format!(
1866 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
1867 ))));
1868 }
1869 self.handle.poll_ready(
1870 Interest {
1871 direction: self.handle.side.write_direction(),
1872 writer: true,
1873 },
1874 |direction| writable_exact(direction, buffer_bytes, needed),
1875 cx,
1876 )
1877 }
1878}
1879
1880impl Drop for WriteHalf {
1881 fn drop(&mut self) {
1882 if !self.shut {
1883 self.handle
1884 .set_flag(self.handle.side.write_direction(), FLAG_RESET);
1885 }
1886 }
1887}
1888
1889#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1893pub(crate) fn waits_supported() -> bool {
1894 orbit_core::sync::supported()
1895}
1896
1897#[cfg(not(any(
1898 target_os = "linux",
1899 target_os = "freebsd",
1900 target_os = "macos"
1901)))]
1902pub(crate) fn waits_supported() -> bool {
1903 false
1904}
1905
1906#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1910#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1913pub(crate) fn wait_on_timeout(
1914 word: &AtomicU32,
1915 expected: u32,
1916 timeout: Duration,
1917) -> Result<bool> {
1918 orbit_core::sync::wait_word_timeout(word, expected, timeout)
1919 .map_err(Error::Io)
1920}
1921
1922pub(crate) fn wait_on(
1923 word: &AtomicU32,
1924 expected: u32,
1925) -> Result<()> {
1926 orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
1927}
1928
1929#[cfg(not(any(
1930 target_os = "linux",
1931 target_os = "freebsd",
1932 target_os = "macos"
1933)))]
1934pub(crate) fn wait_on(
1935 _word: &AtomicU32,
1936 _expected: u32,
1937) -> Result<()> {
1938 Err(Error::Io(std::io::Error::new(
1939 std::io::ErrorKind::Unsupported,
1940 "orbit-stream needs a platform that can wait on a shared word",
1941 )))
1942}
1943
1944pub(crate) fn wake_on(word: &AtomicU32) {
1946 #[cfg(any(
1947 target_os = "linux",
1948 target_os = "freebsd",
1949 target_os = "macos"
1950 ))]
1951 let _ = orbit_core::sync::wake_word(word);
1952 #[cfg(not(any(
1953 target_os = "linux",
1954 target_os = "freebsd",
1955 target_os = "macos"
1956 )))]
1957 let _ = word;
1958}
1959
1960pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
1961 mutex.lock().unwrap_or_else(|error| error.into_inner())
1962}
1963
1964#[cfg(test)]
1965mod tests {
1966 use std::sync::Arc;
1967 use std::time::Duration;
1968
1969 use orbit_core::{Fleet, NodeId};
1970
1971 use super::{
1972 Error, Incarnation, STREAM_BUFFER_BYTES, STREAM_LANE_CAPACITY, Side,
1973 Streams, Ticket,
1974 };
1975
1976 fn streams(name: &'static str) -> Streams {
1977 Streams::new(
1978 Arc::new(Fleet::join(name, 2).unwrap()),
1979 Incarnation::new(1),
1980 )
1981 .unwrap()
1982 }
1983
1984 #[test]
1985 fn bytes_cross_in_order_and_end_cleanly() {
1986 let streams = streams("stream-order");
1987 let (mut a, ticket) = streams.create().unwrap();
1988 assert_eq!(ticket.side, Side::B);
1989 assert_eq!(ticket.id.node(), 0);
1990 assert_eq!(ticket.epoch, streams.epoch());
1991 let b = streams.open(ticket).unwrap();
1992
1993 a.blocking_write_all(b"hello, ").unwrap();
1994 a.blocking_write_all(b"world").unwrap();
1995 a.finish().unwrap();
1996
1997 let mut out = Vec::new();
1998 let mut buf = [0_u8; 4];
1999 loop {
2000 let n = b.blocking_read(&mut buf).unwrap();
2001 if n == 0 {
2002 break;
2003 }
2004 out.extend_from_slice(&buf[..n]);
2005 }
2006 assert_eq!(out, b"hello, world");
2007 b.blocking_write_all(b"back").unwrap();
2009 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"back");
2010 }
2011
2012 #[test]
2013 fn a_ticket_prints_and_parses() {
2014 let streams = streams("stream-ticket");
2015 let (_a, ticket) = streams.create().unwrap();
2016 let text = ticket.to_string();
2017 assert!(text.contains("/b/"), "{text}");
2018 assert_eq!(text.parse::<Ticket>().unwrap(), ticket);
2019 assert!(matches!("nope".parse::<Ticket>(), Err(Error::Malformed(_))));
2020 assert!(matches!("x/b".parse::<Ticket>(), Err(Error::Malformed(_))));
2021 }
2022
2023 #[test]
2024 fn a_side_is_held_once_and_a_released_stream_goes_stale() {
2025 let streams = streams("stream-claim");
2026 let (a, ticket) = streams.create().unwrap();
2027 let b = streams.open(ticket).unwrap();
2028 assert!(matches!(streams.open(ticket), Err(Error::AlreadyClaimed(_))));
2029 assert!(streams.is_live(ticket.id));
2030
2031 drop(a);
2032 let mut buf = [0_u8; 1];
2034 assert!(matches!(b.blocking_read(&mut buf), Err(Error::Reset)));
2035 drop(b);
2036 assert!(!streams.is_live(ticket.id));
2037 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2038
2039 for _ in 0..STREAM_LANE_CAPACITY - 1 {
2042 drop(streams.create().unwrap());
2043 }
2044 let (_a2, ticket2) = streams.create().unwrap();
2045 assert_eq!(ticket2.id.slot(), ticket.id.slot());
2046 assert_ne!(ticket2.id.generation(), ticket.id.generation());
2047 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2048 }
2049
2050 #[test]
2051 fn a_reset_starts_a_new_epoch_and_refuses_old_tickets() {
2052 let streams = streams("stream-epoch");
2053 let (_a, ticket) = streams.create().unwrap();
2054 let before = streams.epoch();
2055 streams.reset_all();
2056 assert!(streams.epoch() > before);
2057 let (_a2, ticket2) = streams.create().unwrap();
2060 assert_eq!(ticket2.id, ticket.id);
2061 assert_ne!(ticket2.epoch, ticket.epoch);
2062 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2063 assert!(streams.open(ticket2).is_ok());
2064 }
2065
2066 #[test]
2067 fn a_slot_whose_generation_runs_out_is_not_reused() {
2068 let streams = streams("stream-exhaust");
2071 let limit = super::layout::GENERATION_LIMIT as usize;
2072 for _ in 0..STREAM_LANE_CAPACITY * limit {
2073 drop(streams.create().unwrap());
2074 }
2075 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2076 streams.reset_all();
2077 let (_a, ticket) = streams.create().unwrap();
2078 assert_eq!(ticket.id.generation(), 1);
2079 }
2080
2081 #[test]
2082 fn a_full_ring_blocks_the_writer_and_wraps_correctly() {
2083 let streams = streams("stream-wrap");
2084 let (a, ticket) = streams.create().unwrap();
2085 let b = streams.open(ticket).unwrap();
2086
2087 let filler = vec![7_u8; STREAM_BUFFER_BYTES];
2088 assert_eq!(a.try_write(&filler).unwrap(), STREAM_BUFFER_BYTES);
2089 assert!(matches!(a.try_write(b"x"), Err(Error::WouldBlock)));
2090
2091 let mut buf = vec![0_u8; STREAM_BUFFER_BYTES - 5];
2092 assert_eq!(b.blocking_read(&mut buf).unwrap(), STREAM_BUFFER_BYTES - 5);
2093 let tail =
2095 (0..STREAM_BUFFER_BYTES).map(|i| i as u8).collect::<Vec<_>>();
2096 assert_eq!(a.try_write(&tail).unwrap(), STREAM_BUFFER_BYTES - 5);
2097
2098 let mut rest = vec![0_u8; 5];
2099 assert_eq!(b.blocking_read(&mut rest).unwrap(), 5);
2100 assert_eq!(rest, vec![7_u8; 5]);
2101 let mut wrapped = vec![0_u8; STREAM_BUFFER_BYTES];
2102 let mut got = 0;
2103 while got < STREAM_BUFFER_BYTES - 5 {
2104 got += b.blocking_read(&mut wrapped[got..]).unwrap();
2105 }
2106 assert_eq!(&wrapped[..got], &tail[..got]);
2107 }
2108
2109 #[test]
2110 fn exact_frames_never_publish_or_consume_a_prefix() {
2111 let streams = streams("stream-exact-frame");
2112 let (a, ticket) = streams.create().unwrap();
2113 let b = streams.open(ticket).unwrap();
2114
2115 let filler = vec![7_u8; STREAM_BUFFER_BYTES - 32];
2116 a.try_write_exact(&filler).unwrap();
2117 assert!(matches!(
2118 a.try_write_exact(&[9_u8; 64]),
2119 Err(Error::WouldBlock)
2120 ));
2121 let mut consumed = vec![0_u8; filler.len()];
2122 b.try_read_exact(&mut consumed).unwrap();
2123 assert_eq!(consumed, filler);
2124 assert!(matches!(b.try_read(&mut [0_u8; 1]), Err(Error::WouldBlock)));
2125
2126 a.try_write_exact(&[3_u8; 32]).unwrap();
2127 assert!(matches!(
2128 b.try_read_exact(&mut [0_u8; 64]),
2129 Err(Error::WouldBlock)
2130 ));
2131 let mut frame = [0_u8; 32];
2132 b.try_read_exact(&mut frame).unwrap();
2133 assert_eq!(frame, [3_u8; 32]);
2134 }
2135
2136 #[test]
2137 fn a_blocked_reader_is_woken_by_a_write_in_another_thread() {
2138 let streams = streams("stream-wake");
2139 let (a, ticket) = streams.create().unwrap();
2140 let b = streams.open(ticket).unwrap();
2141 let reader = std::thread::spawn(move || {
2142 let mut buf = [0_u8; 8];
2143 let n = b.blocking_read(&mut buf).unwrap();
2144 buf[..n].to_vec()
2145 });
2146 std::thread::sleep(Duration::from_millis(30));
2147 a.blocking_write_all(b"ping").unwrap();
2148 assert_eq!(reader.join().unwrap(), b"ping");
2149 }
2150
2151 #[test]
2152 fn a_reader_gone_fails_the_writer_and_a_lane_fills_up() {
2153 let streams = streams("stream-gone");
2154 let (a, ticket) = streams.create().unwrap();
2155 let b = streams.open(ticket).unwrap();
2156 let (b_read, _b_write) = b.split();
2157 drop(b_read);
2158 assert!(matches!(a.blocking_write(b"nobody"), Err(Error::PeerGone)));
2159
2160 let mut held = vec![a];
2161 for _ in 1..STREAM_LANE_CAPACITY {
2162 held.push(streams.create().unwrap().0);
2163 }
2164 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2165 }
2166
2167 #[test]
2168 fn a_death_report_ends_that_incarnations_sides_only() {
2169 let streams = streams("stream-dead");
2170 let (a, ticket) = streams.create().unwrap();
2171 let b = streams.open(ticket).unwrap();
2172 b.blocking_write_all(b"partial").unwrap();
2173
2174 streams.node_dead(NodeId::ZERO, Incarnation::new(2));
2176 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"partial");
2177
2178 let reader = std::thread::spawn(move || {
2184 let mut buf = [0_u8; 8];
2185 a.blocking_read(&mut buf).map(|_| ())
2186 });
2187 std::thread::sleep(Duration::from_millis(30));
2188 streams.node_dead(NodeId::ZERO, Incarnation::new(1));
2189 assert!(matches!(
2190 reader.join().unwrap(),
2191 Err(Error::Reset | Error::Stale(_))
2192 ));
2193 assert!(matches!(
2194 b.try_write(b"x"),
2195 Err(Error::Reset | Error::Stale(_))
2196 ));
2197 assert!(matches!(
2198 streams.open(ticket),
2199 Err(Error::AlreadyClaimed(_) | Error::Stale(_))
2200 ));
2201 drop(b);
2202 assert!(!streams.is_live(ticket.id));
2203 }
2204
2205 #[test]
2206 fn an_offer_is_taken_once_and_a_stale_offer_is_skipped() {
2207 let streams = streams("stream-offer");
2208 assert!(streams.take_offer().is_none());
2209 let (a, ticket) = streams.create().unwrap();
2210 streams.offer(ticket, NodeId::ZERO).unwrap();
2211 assert_eq!(streams.take_offer(), Some(ticket));
2212 assert!(streams.take_offer().is_none());
2213
2214 let (gone, stale) = streams.create().unwrap();
2215 streams.offer(stale, NodeId::ZERO).unwrap();
2216 drop(gone);
2217 assert!(streams.take_offer().is_none());
2218
2219 let taker = streams.clone();
2220 let waiter = std::thread::spawn(move || taker.blocking_take_offer());
2221 std::thread::sleep(Duration::from_millis(30));
2222 streams.offer(ticket, NodeId::ZERO).unwrap();
2223 assert_eq!(waiter.join().unwrap().unwrap(), ticket);
2224 drop(a);
2225 }
2226
2227 #[test]
2228 fn a_contested_claim_goes_to_exactly_one_opener() {
2229 let streams = streams("stream-race");
2230 for _ in 0..50 {
2231 let (_a, ticket) = streams.create().unwrap();
2232 let openers = (0..4)
2233 .map(|_| {
2234 let streams = streams.clone();
2235 std::thread::spawn(move || streams.open(ticket).is_ok())
2236 })
2237 .collect::<Vec<_>>();
2238 let won = openers
2239 .into_iter()
2240 .map(|opener| opener.join().unwrap())
2241 .filter(|won| *won)
2242 .count();
2243 assert_eq!(won, 1);
2244 }
2245 }
2246
2247 #[test]
2248 fn random_sizes_concatenate_exactly() {
2249 let streams = streams("stream-random");
2250 let (mut a, ticket) = streams.create().unwrap();
2251 let b = streams.open(ticket).unwrap();
2252 let payload = (0..300_000_u32)
2253 .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
2254 .collect::<Vec<_>>();
2255 let expected = payload.clone();
2256 let writer = std::thread::spawn(move || {
2257 let mut rest = &payload[..];
2258 let mut step = 1_usize;
2259 while !rest.is_empty() {
2260 let len = step.min(rest.len());
2261 a.blocking_write_all(&rest[..len]).unwrap();
2262 rest = &rest[len..];
2263 step = (step * 7 + 13) % 9_001 + 1;
2264 }
2265 a.finish().unwrap();
2266 });
2267 let mut out = Vec::new();
2268 let mut step = 3_usize;
2269 loop {
2270 let chunk = b.blocking_read_chunk(step).unwrap();
2271 if chunk.is_empty() {
2272 break;
2273 }
2274 out.extend_from_slice(&chunk);
2275 step = (step * 5 + 11) % 4_099 + 1;
2276 }
2277 writer.join().unwrap();
2278 assert_eq!(out, expected);
2279 }
2280}