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_CLOSE_WATCH, FLAG_FIN, FLAG_READER_GONE, FLAG_RESET,
56 GENERATION_MASK, SIDE_CLAIMED, SIDE_DEAD, SIDE_FREE, SIDE_RELEASED,
57 SLOT_BITS, SLOT_EMPTY, 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 || before & FLAG_CLOSE_WATCH != 0
1229 || writer_may_park(direction, before);
1230 self.table.notify(
1231 self.index,
1232 slot,
1233 direction,
1234 self.peer_if(ring),
1235 );
1236 }
1237 }
1238 }
1239
1240 #[cfg(any(
1243 target_os = "linux",
1244 target_os = "freebsd",
1245 target_os = "macos"
1246 ))]
1247 fn wait_until_timeout(
1248 &self,
1249 direction_index: usize,
1250 ready: impl Fn(&Direction) -> bool,
1251 timeout: Duration,
1252 ) -> Result<bool> {
1253 let deadline = Instant::now() + timeout;
1254 loop {
1255 let slot = self.slot()?;
1256 let direction = Self::direction(slot, direction_index);
1257 if ready(direction) {
1258 return Ok(true);
1259 }
1260 let left = deadline.saturating_duration_since(Instant::now());
1261 if left.is_zero() {
1262 return Ok(false);
1263 }
1264 direction.waiters.fetch_add(1, Ordering::SeqCst);
1265 let since = direction.changes.load(Ordering::SeqCst);
1266 let outcome = if ready(direction) {
1267 Ok(true)
1268 } else {
1269 wait_on_timeout(&direction.changes, since, left)
1270 };
1271 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1272 if !outcome? {
1273 let slot = self.slot()?;
1276 return Ok(ready(Self::direction(slot, direction_index)));
1277 }
1278 }
1279 }
1280
1281 fn wait_until(
1282 &self,
1283 direction_index: usize,
1284 ready: impl Fn(&Direction) -> bool,
1285 ) -> Result<()> {
1286 loop {
1287 let slot = self.slot()?;
1288 let direction = Self::direction(slot, direction_index);
1289 if ready(direction) {
1290 return Ok(());
1291 }
1292 direction.waiters.fetch_add(1, Ordering::SeqCst);
1293 let since = direction.changes.load(Ordering::SeqCst);
1294 let outcome = if ready(direction) {
1295 Ok(())
1296 } else {
1297 wait_on(&direction.changes, since)
1298 };
1299 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1300 outcome?;
1301 }
1302 }
1303
1304 fn poll_ready(
1307 &self,
1308 interest: Interest,
1309 ready: impl Fn(&Direction) -> bool,
1310 cx: &mut std::task::Context<'_>,
1311 ) -> std::task::Poll<Result<()>> {
1312 let slot = match self.slot() {
1313 Ok(slot) => slot,
1314 Err(error) => return std::task::Poll::Ready(Err(error)),
1315 };
1316 let direction = Self::direction(slot, interest.direction);
1317 if ready(direction) {
1318 return std::task::Poll::Ready(Ok(()));
1319 }
1320 if let Err(error) =
1321 self.table.register(self.index, interest, cx.waker())
1322 {
1323 return std::task::Poll::Ready(Err(error));
1324 }
1325 std::sync::atomic::fence(Ordering::SeqCst);
1328 if ready(direction) {
1329 std::task::Poll::Ready(Ok(()))
1330 } else {
1331 std::task::Poll::Pending
1332 }
1333 }
1334}
1335
1336fn writer_may_park(
1350 direction: &Direction,
1351 before: u8,
1352) -> bool {
1353 if before & (FLAG_FIN | FLAG_RESET) != 0 {
1354 return false;
1355 }
1356 std::sync::atomic::fence(Ordering::SeqCst);
1357 direction.head.load(Ordering::SeqCst)
1358 != direction.tail.load(Ordering::SeqCst)
1359}
1360
1361fn readable(direction: &Direction) -> bool {
1362 direction.head.load(Ordering::Acquire)
1363 != direction.tail.load(Ordering::Relaxed)
1364 || direction.flags() & (FLAG_FIN | FLAG_RESET) != 0
1365}
1366
1367fn writable(
1368 direction: &Direction,
1369 buffer_bytes: usize,
1370) -> bool {
1371 let queued = (direction.head.load(Ordering::Relaxed)
1372 - direction.tail.load(Ordering::Acquire)) as usize;
1373 queued < buffer_bytes
1374 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1375}
1376
1377fn writable_exact(
1378 direction: &Direction,
1379 buffer_bytes: usize,
1380 needed: usize,
1381) -> bool {
1382 let queued = (direction.head.load(Ordering::Relaxed)
1383 - direction.tail.load(Ordering::Acquire)) as usize;
1384 buffer_bytes.saturating_sub(queued) >= needed
1385 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1386}
1387
1388impl Drop for Handle {
1389 fn drop(&mut self) {
1393 let slot = &self.table.slots()[self.index];
1394 if !slot.is(self.id.generation()) {
1395 return;
1396 }
1397 if slot.claimed[self.side.index()]
1400 .compare_exchange(
1401 SIDE_CLAIMED,
1402 SIDE_RELEASED,
1403 Ordering::SeqCst,
1404 Ordering::SeqCst,
1405 )
1406 .is_err()
1407 {
1408 return;
1409 }
1410 let other = slot.claimed[1 - self.side.index()].load(Ordering::SeqCst);
1414 if other != SIDE_CLAIMED {
1415 let _ = slot.state.compare_exchange(
1416 SLOT_LIVE,
1417 SLOT_EMPTY,
1418 Ordering::SeqCst,
1419 Ordering::SeqCst,
1420 );
1421 for direction in &slot.directions {
1422 direction.changes.fetch_add(1, Ordering::SeqCst);
1425 if direction.waiters.load(Ordering::SeqCst) > 0 {
1426 wake_on(&direction.changes);
1427 }
1428 }
1429 self.table.registry.wake(self.index);
1430 }
1431 }
1432}
1433
1434pub struct Endpoint {
1438 read: ReadHalf,
1439 write: WriteHalf,
1440}
1441
1442impl Endpoint {
1443 fn new(handle: Arc<Handle>) -> Self {
1444 Self {
1445 read: ReadHalf { handle: Arc::clone(&handle) },
1446 write: WriteHalf { handle, shut: false },
1447 }
1448 }
1449
1450 pub fn id(&self) -> StreamId {
1451 self.read.handle.id
1452 }
1453
1454 pub fn side(&self) -> Side {
1455 self.read.handle.side
1456 }
1457
1458 pub fn split(self) -> (ReadHalf, WriteHalf) {
1459 (self.read, self.write)
1460 }
1461
1462 pub fn rejoin(
1473 read: ReadHalf,
1474 write: WriteHalf,
1475 ) -> Result<Self> {
1476 if read.id() != write.id() {
1477 return Err(Error::Stale(read.id()));
1478 }
1479 Ok(Self { read, write })
1480 }
1481
1482 pub fn rearm(
1495 &mut self,
1496 streams: &Streams,
1497 ) -> Result<()> {
1498 if self.side() != Side::A {
1499 return Err(Error::Malformed(
1500 "only the creator side of a stream can be rearmed".to_owned(),
1501 ));
1502 }
1503 streams.rearm(self.id())?;
1504 self.write.shut = false;
1505 Ok(())
1506 }
1507
1508 pub fn try_read(
1509 &self,
1510 buf: &mut [u8],
1511 ) -> Result<usize> {
1512 self.read.try_read(buf)
1513 }
1514
1515 pub fn try_read_exact(
1516 &self,
1517 buf: &mut [u8],
1518 ) -> Result<()> {
1519 self.read.try_read_exact(buf)
1520 }
1521
1522 pub fn blocking_read(
1523 &self,
1524 buf: &mut [u8],
1525 ) -> Result<usize> {
1526 self.read.blocking_read(buf)
1527 }
1528
1529 pub fn blocking_read_chunk(
1530 &self,
1531 max: usize,
1532 ) -> Result<Bytes> {
1533 self.read.blocking_read_chunk(max)
1534 }
1535
1536 pub fn try_write(
1537 &self,
1538 buf: &[u8],
1539 ) -> Result<usize> {
1540 self.write.try_write(buf)
1541 }
1542
1543 pub fn try_write_exact(
1544 &self,
1545 buf: &[u8],
1546 ) -> Result<()> {
1547 self.write.try_write_exact(buf)
1548 }
1549
1550 pub fn blocking_write(
1551 &self,
1552 buf: &[u8],
1553 ) -> Result<usize> {
1554 self.write.blocking_write(buf)
1555 }
1556
1557 pub fn blocking_write_all(
1558 &self,
1559 buf: &[u8],
1560 ) -> Result<()> {
1561 self.write.blocking_write_all(buf)
1562 }
1563
1564 pub fn finish(&mut self) -> Result<()> {
1565 self.write.finish()
1566 }
1567
1568 pub fn reset(&mut self) {
1569 self.write.reset()
1570 }
1571}
1572
1573pub struct ReadHalf {
1576 handle: Arc<Handle>,
1577}
1578
1579pub struct ResetSignal {
1583 handle: Arc<Handle>,
1584}
1585
1586impl ResetSignal {
1587 pub fn poll_reset(
1590 &self,
1591 cx: &mut std::task::Context<'_>,
1592 ) -> std::task::Poll<Result<()>> {
1593 self.handle.poll_ready(
1594 Interest {
1595 direction: self.handle.side.read_direction(),
1596 writer: false,
1597 reset: true,
1598 },
1599 |direction| direction.flags() & FLAG_RESET != 0,
1600 cx,
1601 )
1602 }
1603}
1604
1605impl ReadHalf {
1606 pub fn reset_signal(&self) -> ResetSignal {
1609 ResetSignal { handle: Arc::clone(&self.handle) }
1610 }
1611
1612 pub fn id(&self) -> StreamId {
1613 self.handle.id
1614 }
1615
1616 pub fn try_read(
1619 &self,
1620 buf: &mut [u8],
1621 ) -> Result<usize> {
1622 self.handle.try_read(buf)
1623 }
1624
1625 pub fn try_read_exact(
1628 &self,
1629 buf: &mut [u8],
1630 ) -> Result<()> {
1631 self.handle.try_read_exact(buf)
1632 }
1633
1634 pub fn blocking_read(
1637 &self,
1638 buf: &mut [u8],
1639 ) -> Result<usize> {
1640 loop {
1641 match self.handle.try_read(buf) {
1642 Err(Error::WouldBlock) => self.wait_readable()?,
1643 other => return other,
1644 }
1645 }
1646 }
1647
1648 pub fn stop_reading(&self) {
1656 self.handle
1657 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1658 }
1659
1660 pub fn blocking_read_chunk(
1662 &self,
1663 max: usize,
1664 ) -> Result<Bytes> {
1665 let mut chunk = BytesMut::zeroed(max);
1666 let len = self.blocking_read(&mut chunk)?;
1667 chunk.truncate(len);
1668 Ok(chunk.freeze())
1669 }
1670
1671 #[cfg(any(
1673 target_os = "linux",
1674 target_os = "freebsd",
1675 target_os = "macos"
1676 ))]
1677 pub fn wait_readable_timeout(
1678 &self,
1679 timeout: Duration,
1680 ) -> Result<bool> {
1681 self.handle.wait_until_timeout(
1682 self.handle.side.read_direction(),
1683 readable,
1684 timeout,
1685 )
1686 }
1687
1688 pub fn wait_readable(&self) -> Result<()> {
1690 self.handle.wait_until(self.handle.side.read_direction(), readable)
1691 }
1692
1693 pub fn poll_readable(
1698 &self,
1699 cx: &mut std::task::Context<'_>,
1700 ) -> std::task::Poll<Result<()>> {
1701 self.handle.poll_ready(
1702 Interest {
1703 direction: self.handle.side.read_direction(),
1704 writer: false,
1705 reset: false,
1706 },
1707 readable,
1708 cx,
1709 )
1710 }
1711}
1712
1713impl Drop for ReadHalf {
1714 fn drop(&mut self) {
1715 self.handle
1716 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1717 }
1718}
1719
1720pub struct WriteHalf {
1723 handle: Arc<Handle>,
1724 shut: bool,
1725}
1726
1727struct CloseWatch<'a>(&'a Handle);
1729
1730impl Drop for CloseWatch<'_> {
1731 fn drop(&mut self) {
1732 if let Ok(slot) = self.0.slot() {
1733 Handle::direction(slot, self.0.side.write_direction())
1734 .flags
1735 .fetch_and(!FLAG_CLOSE_WATCH, Ordering::SeqCst);
1736 }
1737 }
1738}
1739
1740impl WriteHalf {
1741 pub async fn peer_closed(&mut self) -> Result<()> {
1747 let slot = self.handle.slot()?;
1748 let direction =
1749 Handle::direction(slot, self.handle.side.write_direction());
1750 direction.flags.fetch_or(FLAG_CLOSE_WATCH, Ordering::SeqCst);
1751 let _watch = CloseWatch(&self.handle);
1752 std::future::poll_fn(|cx| {
1753 self.handle.poll_ready(
1754 Interest {
1755 direction: self.handle.side.write_direction(),
1756 writer: true,
1757 reset: false,
1758 },
1759 |direction| {
1760 direction.flags() & (FLAG_RESET | FLAG_READER_GONE) != 0
1761 },
1762 cx,
1763 )
1764 })
1765 .await
1766 }
1767
1768 pub fn poll_peer_dead(
1774 &self,
1775 cx: &mut std::task::Context<'_>,
1776 ) -> std::task::Poll<Result<()>> {
1777 let slot = match self.handle.slot() {
1778 Ok(slot) => slot,
1779 Err(error) => return std::task::Poll::Ready(Err(error)),
1780 };
1781 let peer = &slot.claimed[1 - self.handle.side.index()];
1782 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1783 return std::task::Poll::Ready(Ok(()));
1784 }
1785 let interest = Interest {
1786 direction: self.handle.side.write_direction(),
1787 writer: true,
1788 reset: false,
1789 };
1790 if let Err(error) =
1791 self.handle.table.register(self.handle.index, interest, cx.waker())
1792 {
1793 return std::task::Poll::Ready(Err(error));
1794 }
1795 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1798 std::task::Poll::Ready(Ok(()))
1799 } else {
1800 std::task::Poll::Pending
1801 }
1802 }
1803
1804 pub fn id(&self) -> StreamId {
1805 self.handle.id
1806 }
1807
1808 pub fn try_write(
1810 &self,
1811 buf: &[u8],
1812 ) -> Result<usize> {
1813 self.handle.try_write(buf)
1814 }
1815
1816 pub fn try_write_exact(
1819 &self,
1820 buf: &[u8],
1821 ) -> Result<()> {
1822 self.handle.try_write_exact(buf)
1823 }
1824
1825 pub fn blocking_write(
1827 &self,
1828 buf: &[u8],
1829 ) -> Result<usize> {
1830 loop {
1831 match self.handle.try_write(buf) {
1832 Err(Error::WouldBlock) => self.wait_writable()?,
1833 other => return other,
1834 }
1835 }
1836 }
1837
1838 pub fn blocking_write_all(
1841 &self,
1842 mut buf: &[u8],
1843 ) -> Result<()> {
1844 while !buf.is_empty() {
1845 let written = self.blocking_write(buf)?;
1846 buf = &buf[written..];
1847 }
1848 Ok(())
1849 }
1850
1851 pub fn finish(&mut self) -> Result<()> {
1855 self.shut = true;
1856 self.handle.set_flag(self.handle.side.write_direction(), FLAG_FIN);
1857 Ok(())
1858 }
1859
1860 pub fn reset(&mut self) {
1863 self.shut = true;
1864 self.handle.set_flag(self.handle.side.write_direction(), FLAG_RESET);
1865 }
1866
1867 #[cfg(any(
1871 target_os = "linux",
1872 target_os = "freebsd",
1873 target_os = "macos"
1874 ))]
1875 pub fn wait_writable_timeout(
1876 &self,
1877 timeout: Duration,
1878 ) -> Result<bool> {
1879 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1880 self.handle.wait_until_timeout(
1881 self.handle.side.write_direction(),
1882 |direction| writable(direction, buffer_bytes),
1883 timeout,
1884 )
1885 }
1886
1887 pub fn wait_writable(&self) -> Result<()> {
1889 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1890 self.handle
1891 .wait_until(self.handle.side.write_direction(), |direction| {
1892 writable(direction, buffer_bytes)
1893 })
1894 }
1895
1896 pub fn wait_writable_exact(
1898 &self,
1899 needed: usize,
1900 ) -> Result<()> {
1901 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1902 if needed > buffer_bytes {
1903 return Err(Error::Malformed(format!(
1904 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
1905 )));
1906 }
1907 self.handle
1908 .wait_until(self.handle.side.write_direction(), |direction| {
1909 writable_exact(direction, buffer_bytes, needed)
1910 })
1911 }
1912
1913 pub fn poll_writable(
1918 &self,
1919 cx: &mut std::task::Context<'_>,
1920 ) -> std::task::Poll<Result<()>> {
1921 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1922 self.handle.poll_ready(
1923 Interest {
1924 direction: self.handle.side.write_direction(),
1925 writer: true,
1926 reset: false,
1927 },
1928 |direction| writable(direction, buffer_bytes),
1929 cx,
1930 )
1931 }
1932
1933 pub fn poll_writable_exact(
1935 &self,
1936 needed: usize,
1937 cx: &mut std::task::Context<'_>,
1938 ) -> std::task::Poll<Result<()>> {
1939 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1940 if needed > buffer_bytes {
1941 return std::task::Poll::Ready(Err(Error::Malformed(format!(
1942 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
1943 ))));
1944 }
1945 self.handle.poll_ready(
1946 Interest {
1947 direction: self.handle.side.write_direction(),
1948 writer: true,
1949 reset: false,
1950 },
1951 |direction| writable_exact(direction, buffer_bytes, needed),
1952 cx,
1953 )
1954 }
1955}
1956
1957impl Drop for WriteHalf {
1958 fn drop(&mut self) {
1959 if !self.shut {
1960 self.handle
1961 .set_flag(self.handle.side.write_direction(), FLAG_RESET);
1962 }
1963 }
1964}
1965
1966#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1970pub(crate) fn waits_supported() -> bool {
1971 orbit_core::sync::supported()
1972}
1973
1974#[cfg(not(any(
1975 target_os = "linux",
1976 target_os = "freebsd",
1977 target_os = "macos"
1978)))]
1979pub(crate) fn waits_supported() -> bool {
1980 false
1981}
1982
1983#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1987#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1990pub(crate) fn wait_on_timeout(
1991 word: &AtomicU32,
1992 expected: u32,
1993 timeout: Duration,
1994) -> Result<bool> {
1995 orbit_core::sync::wait_word_timeout(word, expected, timeout)
1996 .map_err(Error::Io)
1997}
1998
1999pub(crate) fn wait_on(
2000 word: &AtomicU32,
2001 expected: u32,
2002) -> Result<()> {
2003 orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
2004}
2005
2006#[cfg(not(any(
2007 target_os = "linux",
2008 target_os = "freebsd",
2009 target_os = "macos"
2010)))]
2011pub(crate) fn wait_on(
2012 _word: &AtomicU32,
2013 _expected: u32,
2014) -> Result<()> {
2015 Err(Error::Io(std::io::Error::new(
2016 std::io::ErrorKind::Unsupported,
2017 "orbit-stream needs a platform that can wait on a shared word",
2018 )))
2019}
2020
2021pub(crate) fn wake_on(word: &AtomicU32) {
2023 #[cfg(any(
2024 target_os = "linux",
2025 target_os = "freebsd",
2026 target_os = "macos"
2027 ))]
2028 let _ = orbit_core::sync::wake_word(word);
2029 #[cfg(not(any(
2030 target_os = "linux",
2031 target_os = "freebsd",
2032 target_os = "macos"
2033 )))]
2034 let _ = word;
2035}
2036
2037pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
2038 mutex.lock().unwrap_or_else(|error| error.into_inner())
2039}
2040
2041#[cfg(test)]
2042mod tests {
2043 use std::sync::Arc;
2044 use std::time::Duration;
2045
2046 use orbit_core::{Fleet, NodeId};
2047
2048 use super::{
2049 Error, Incarnation, STREAM_BUFFER_BYTES, STREAM_LANE_CAPACITY, Side,
2050 Streams, Ticket,
2051 };
2052
2053 fn streams(name: &'static str) -> Streams {
2054 Streams::new(
2055 Arc::new(Fleet::join(name, 2).unwrap()),
2056 Incarnation::new(1),
2057 )
2058 .unwrap()
2059 }
2060
2061 #[test]
2062 fn bytes_cross_in_order_and_end_cleanly() {
2063 let streams = streams("stream-order");
2064 let (mut a, ticket) = streams.create().unwrap();
2065 assert_eq!(ticket.side, Side::B);
2066 assert_eq!(ticket.id.node(), 0);
2067 assert_eq!(ticket.epoch, streams.epoch());
2068 let b = streams.open(ticket).unwrap();
2069
2070 a.blocking_write_all(b"hello, ").unwrap();
2071 a.blocking_write_all(b"world").unwrap();
2072 a.finish().unwrap();
2073
2074 let mut out = Vec::new();
2075 let mut buf = [0_u8; 4];
2076 loop {
2077 let n = b.blocking_read(&mut buf).unwrap();
2078 if n == 0 {
2079 break;
2080 }
2081 out.extend_from_slice(&buf[..n]);
2082 }
2083 assert_eq!(out, b"hello, world");
2084 b.blocking_write_all(b"back").unwrap();
2086 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"back");
2087 }
2088
2089 #[test]
2090 fn a_ticket_prints_and_parses() {
2091 let streams = streams("stream-ticket");
2092 let (_a, ticket) = streams.create().unwrap();
2093 let text = ticket.to_string();
2094 assert!(text.contains("/b/"), "{text}");
2095 assert_eq!(text.parse::<Ticket>().unwrap(), ticket);
2096 assert!(matches!("nope".parse::<Ticket>(), Err(Error::Malformed(_))));
2097 assert!(matches!("x/b".parse::<Ticket>(), Err(Error::Malformed(_))));
2098 }
2099
2100 #[test]
2101 fn a_side_is_held_once_and_a_released_stream_goes_stale() {
2102 let streams = streams("stream-claim");
2103 let (a, ticket) = streams.create().unwrap();
2104 let b = streams.open(ticket).unwrap();
2105 assert!(matches!(streams.open(ticket), Err(Error::AlreadyClaimed(_))));
2106 assert!(streams.is_live(ticket.id));
2107
2108 drop(a);
2109 let mut buf = [0_u8; 1];
2111 assert!(matches!(b.blocking_read(&mut buf), Err(Error::Reset)));
2112 drop(b);
2113 assert!(!streams.is_live(ticket.id));
2114 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2115
2116 for _ in 0..STREAM_LANE_CAPACITY - 1 {
2119 drop(streams.create().unwrap());
2120 }
2121 let (_a2, ticket2) = streams.create().unwrap();
2122 assert_eq!(ticket2.id.slot(), ticket.id.slot());
2123 assert_ne!(ticket2.id.generation(), ticket.id.generation());
2124 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2125 }
2126
2127 #[test]
2128 fn a_reset_starts_a_new_epoch_and_refuses_old_tickets() {
2129 let streams = streams("stream-epoch");
2130 let (_a, ticket) = streams.create().unwrap();
2131 let before = streams.epoch();
2132 streams.reset_all();
2133 assert!(streams.epoch() > before);
2134 let (_a2, ticket2) = streams.create().unwrap();
2137 assert_eq!(ticket2.id, ticket.id);
2138 assert_ne!(ticket2.epoch, ticket.epoch);
2139 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2140 assert!(streams.open(ticket2).is_ok());
2141 }
2142
2143 #[test]
2144 fn a_slot_whose_generation_runs_out_is_not_reused() {
2145 let streams = streams("stream-exhaust");
2148 let limit = super::layout::GENERATION_LIMIT as usize;
2149 for _ in 0..STREAM_LANE_CAPACITY * limit {
2150 drop(streams.create().unwrap());
2151 }
2152 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2153 streams.reset_all();
2154 let (_a, ticket) = streams.create().unwrap();
2155 assert_eq!(ticket.id.generation(), 1);
2156 }
2157
2158 #[test]
2159 fn a_full_ring_blocks_the_writer_and_wraps_correctly() {
2160 let streams = streams("stream-wrap");
2161 let (a, ticket) = streams.create().unwrap();
2162 let b = streams.open(ticket).unwrap();
2163
2164 let filler = vec![7_u8; STREAM_BUFFER_BYTES];
2165 assert_eq!(a.try_write(&filler).unwrap(), STREAM_BUFFER_BYTES);
2166 assert!(matches!(a.try_write(b"x"), Err(Error::WouldBlock)));
2167
2168 let mut buf = vec![0_u8; STREAM_BUFFER_BYTES - 5];
2169 assert_eq!(b.blocking_read(&mut buf).unwrap(), STREAM_BUFFER_BYTES - 5);
2170 let tail =
2172 (0..STREAM_BUFFER_BYTES).map(|i| i as u8).collect::<Vec<_>>();
2173 assert_eq!(a.try_write(&tail).unwrap(), STREAM_BUFFER_BYTES - 5);
2174
2175 let mut rest = vec![0_u8; 5];
2176 assert_eq!(b.blocking_read(&mut rest).unwrap(), 5);
2177 assert_eq!(rest, vec![7_u8; 5]);
2178 let mut wrapped = vec![0_u8; STREAM_BUFFER_BYTES];
2179 let mut got = 0;
2180 while got < STREAM_BUFFER_BYTES - 5 {
2181 got += b.blocking_read(&mut wrapped[got..]).unwrap();
2182 }
2183 assert_eq!(&wrapped[..got], &tail[..got]);
2184 }
2185
2186 #[test]
2187 fn exact_frames_never_publish_or_consume_a_prefix() {
2188 let streams = streams("stream-exact-frame");
2189 let (a, ticket) = streams.create().unwrap();
2190 let b = streams.open(ticket).unwrap();
2191
2192 let filler = vec![7_u8; STREAM_BUFFER_BYTES - 32];
2193 a.try_write_exact(&filler).unwrap();
2194 assert!(matches!(
2195 a.try_write_exact(&[9_u8; 64]),
2196 Err(Error::WouldBlock)
2197 ));
2198 let mut consumed = vec![0_u8; filler.len()];
2199 b.try_read_exact(&mut consumed).unwrap();
2200 assert_eq!(consumed, filler);
2201 assert!(matches!(b.try_read(&mut [0_u8; 1]), Err(Error::WouldBlock)));
2202
2203 a.try_write_exact(&[3_u8; 32]).unwrap();
2204 assert!(matches!(
2205 b.try_read_exact(&mut [0_u8; 64]),
2206 Err(Error::WouldBlock)
2207 ));
2208 let mut frame = [0_u8; 32];
2209 b.try_read_exact(&mut frame).unwrap();
2210 assert_eq!(frame, [3_u8; 32]);
2211 }
2212
2213 #[test]
2214 fn a_blocked_reader_is_woken_by_a_write_in_another_thread() {
2215 let streams = streams("stream-wake");
2216 let (a, ticket) = streams.create().unwrap();
2217 let b = streams.open(ticket).unwrap();
2218 let reader = std::thread::spawn(move || {
2219 let mut buf = [0_u8; 8];
2220 let n = b.blocking_read(&mut buf).unwrap();
2221 buf[..n].to_vec()
2222 });
2223 std::thread::sleep(Duration::from_millis(30));
2224 a.blocking_write_all(b"ping").unwrap();
2225 assert_eq!(reader.join().unwrap(), b"ping");
2226 }
2227
2228 #[test]
2229 fn a_reader_gone_fails_the_writer_and_a_lane_fills_up() {
2230 let streams = streams("stream-gone");
2231 let (a, ticket) = streams.create().unwrap();
2232 let b = streams.open(ticket).unwrap();
2233 let (b_read, _b_write) = b.split();
2234 drop(b_read);
2235 assert!(matches!(a.blocking_write(b"nobody"), Err(Error::PeerGone)));
2236
2237 let mut held = vec![a];
2238 for _ in 1..STREAM_LANE_CAPACITY {
2239 held.push(streams.create().unwrap().0);
2240 }
2241 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2242 }
2243
2244 #[test]
2245 fn a_death_report_ends_that_incarnations_sides_only() {
2246 let streams = streams("stream-dead");
2247 let (a, ticket) = streams.create().unwrap();
2248 let b = streams.open(ticket).unwrap();
2249 b.blocking_write_all(b"partial").unwrap();
2250
2251 streams.node_dead(NodeId::ZERO, Incarnation::new(2));
2253 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"partial");
2254
2255 let reader = std::thread::spawn(move || {
2261 let mut buf = [0_u8; 8];
2262 a.blocking_read(&mut buf).map(|_| ())
2263 });
2264 std::thread::sleep(Duration::from_millis(30));
2265 streams.node_dead(NodeId::ZERO, Incarnation::new(1));
2266 assert!(matches!(
2267 reader.join().unwrap(),
2268 Err(Error::Reset | Error::Stale(_))
2269 ));
2270 assert!(matches!(
2271 b.try_write(b"x"),
2272 Err(Error::Reset | Error::Stale(_))
2273 ));
2274 assert!(matches!(
2275 streams.open(ticket),
2276 Err(Error::AlreadyClaimed(_) | Error::Stale(_))
2277 ));
2278 drop(b);
2279 assert!(!streams.is_live(ticket.id));
2280 }
2281
2282 #[test]
2283 fn an_offer_is_taken_once_and_a_stale_offer_is_skipped() {
2284 let streams = streams("stream-offer");
2285 assert!(streams.take_offer().is_none());
2286 let (a, ticket) = streams.create().unwrap();
2287 streams.offer(ticket, NodeId::ZERO).unwrap();
2288 assert_eq!(streams.take_offer(), Some(ticket));
2289 assert!(streams.take_offer().is_none());
2290
2291 let (gone, stale) = streams.create().unwrap();
2292 streams.offer(stale, NodeId::ZERO).unwrap();
2293 drop(gone);
2294 assert!(streams.take_offer().is_none());
2295
2296 let taker = streams.clone();
2297 let waiter = std::thread::spawn(move || taker.blocking_take_offer());
2298 std::thread::sleep(Duration::from_millis(30));
2299 streams.offer(ticket, NodeId::ZERO).unwrap();
2300 assert_eq!(waiter.join().unwrap().unwrap(), ticket);
2301 drop(a);
2302 }
2303
2304 #[test]
2305 fn a_contested_claim_goes_to_exactly_one_opener() {
2306 let streams = streams("stream-race");
2307 for _ in 0..50 {
2308 let (_a, ticket) = streams.create().unwrap();
2309 let openers = (0..4)
2310 .map(|_| {
2311 let streams = streams.clone();
2312 std::thread::spawn(move || streams.open(ticket).is_ok())
2313 })
2314 .collect::<Vec<_>>();
2315 let won = openers
2316 .into_iter()
2317 .map(|opener| opener.join().unwrap())
2318 .filter(|won| *won)
2319 .count();
2320 assert_eq!(won, 1);
2321 }
2322 }
2323
2324 #[test]
2325 fn random_sizes_concatenate_exactly() {
2326 let streams = streams("stream-random");
2327 let (mut a, ticket) = streams.create().unwrap();
2328 let b = streams.open(ticket).unwrap();
2329 let payload = (0..300_000_u32)
2330 .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
2331 .collect::<Vec<_>>();
2332 let expected = payload.clone();
2333 let writer = std::thread::spawn(move || {
2334 let mut rest = &payload[..];
2335 let mut step = 1_usize;
2336 while !rest.is_empty() {
2337 let len = step.min(rest.len());
2338 a.blocking_write_all(&rest[..len]).unwrap();
2339 rest = &rest[len..];
2340 step = (step * 7 + 13) % 9_001 + 1;
2341 }
2342 a.finish().unwrap();
2343 });
2344 let mut out = Vec::new();
2345 let mut step = 3_usize;
2346 loop {
2347 let chunk = b.blocking_read_chunk(step).unwrap();
2348 if chunk.is_empty() {
2349 break;
2350 }
2351 out.extend_from_slice(&chunk);
2352 step = (step * 5 + 11) % 4_099 + 1;
2353 }
2354 writer.join().unwrap();
2355 assert_eq!(out, expected);
2356 }
2357}