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;
47mod capacity;
48pub mod exchange;
49mod layout;
50pub use capacity::CapacityChanged;
51#[cfg(feature = "tokio")]
52mod poll;
53mod table;
54mod wake;
55
56use layout::{
57 Direction, FLAG_CLOSE_WATCH, FLAG_FIN, FLAG_READER_GONE, FLAG_RESET,
58 GENERATION_MASK, SIDE_CLAIMED, SIDE_DEAD, SIDE_FREE, SIDE_RELEASED,
59 SLOT_BITS, SLOT_EMPTY, SLOT_LIVE, SLOT_MASK, Slot,
60};
61use orbit_core::NodeId;
62#[cfg(any(
63 target_os = "linux",
64 target_os = "freebsd",
65 target_os = "macos"
66))]
67pub use orbit_core::readiness::Readiness;
68use table::Table;
69pub use table::{segment_size, segment_size_for};
70use wake::Interest;
71
72pub const STREAM_KIND: u8 = 246;
75pub const STREAM_LANE_CAPACITY: usize = orbit_core::compile::usize_from_env(
83 option_env!("ORBIT_STREAM_LANE_CAPACITY"),
84 256,
85);
86pub const STREAM_BUFFER_BYTES: usize = orbit_core::compile::usize_from_env(
93 option_env!("ORBIT_STREAM_BUFFER_BYTES"),
94 64 * 1024,
95);
96
97#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
109pub struct StreamSpec {
110 pub kind: u8,
112 pub lane_capacity: usize,
115 pub buffer_bytes: usize,
117}
118
119impl StreamSpec {
120 pub const DEFAULT: Self =
121 Self::new(STREAM_KIND, STREAM_LANE_CAPACITY, STREAM_BUFFER_BYTES);
122
123 pub const fn new(
124 kind: u8,
125 lane_capacity: usize,
126 buffer_bytes: usize,
127 ) -> Self {
128 Self { kind, lane_capacity, buffer_bytes }
129 }
130
131 pub const fn lane_bytes(&self) -> usize {
135 self.lane_capacity * 2 * self.buffer_bytes
136 }
137
138 fn validate(self) -> Result<()> {
141 if self.lane_capacity == 0
142 || self.buffer_bytes == 0
143 || !self.lane_capacity.is_power_of_two()
144 || !self.buffer_bytes.is_power_of_two()
145 || self.lane_capacity > 1 << SLOT_BITS
146 || self.buffer_bytes > u32::MAX as usize
147 {
148 return Err(Error::Malformed(format!(
149 "stream spec kind={} lane_capacity={} buffer_bytes={}: both are powers of two, a lane holds at most {} streams and a ring at most {} bytes",
150 self.kind,
151 self.lane_capacity,
152 self.buffer_bytes,
153 1_usize << SLOT_BITS,
154 u32::MAX
155 )));
156 }
157 Ok(())
158 }
159}
160
161impl Default for StreamSpec {
162 fn default() -> Self {
163 Self::DEFAULT
164 }
165}
166
167pub type Result<T, E = Error> = std::result::Result<T, E>;
168
169#[non_exhaustive]
173#[derive(Debug)]
174pub enum Error {
175 Stale(StreamId),
177 AlreadyClaimed(Ticket),
179 Full {
181 capacity: usize,
182 },
183 PayloadFull {
185 requested_slots: usize,
186 },
187 PayloadTooLarge {
189 len: usize,
190 capacity: usize,
191 },
192 Malformed(String),
194 Closed,
196 Reset,
198 PeerGone,
200 WouldBlock,
202 Io(io::Error),
203}
204
205impl fmt::Display for Error {
206 fn fmt(
207 &self,
208 f: &mut fmt::Formatter<'_>,
209 ) -> fmt::Result {
210 match self {
211 Self::Stale(id) => write!(f, "stream {id} has ended"),
212 Self::AlreadyClaimed(ticket) => {
213 write!(f, "stream side {ticket} is already held")
214 }
215 Self::Full { capacity } => {
216 write!(f, "stream lane is full: capacity={capacity}")
217 }
218 Self::PayloadFull { requested_slots } => {
219 write!(
220 f,
221 "payload arena has no run of {requested_slots} free slots"
222 )
223 }
224 Self::PayloadTooLarge { len, capacity } => {
225 write!(
226 f,
227 "payload chunk is {len} bytes; arena lane capacity is {capacity}"
228 )
229 }
230 Self::Malformed(text) => {
231 write!(f, "not a stream address: {text:?}")
232 }
233 Self::Closed => f.write_str("stream write half is shut down"),
234 Self::Reset => {
235 f.write_str("stream direction was reset by the peer")
236 }
237 Self::PeerGone => f.write_str("stream peer reader is gone"),
238 Self::WouldBlock => f.write_str("stream operation would block"),
239 Self::Io(error) => write!(f, "Orbit stream io error: {error}"),
240 }
241 }
242}
243
244impl std::error::Error for Error {
245 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
246 match self {
247 Self::Io(error) => Some(error),
248 _ => None,
249 }
250 }
251}
252
253impl From<io::Error> for Error {
254 fn from(value: io::Error) -> Self {
255 Self::Io(value)
256 }
257}
258
259impl From<Error> for io::Error {
260 fn from(value: Error) -> Self {
261 let kind = match &value {
262 Error::Stale(_) | Error::AlreadyClaimed(_) => {
263 io::ErrorKind::NotConnected
264 }
265 Error::Full { .. }
266 | Error::PayloadFull { .. }
267 | Error::PayloadTooLarge { .. } => io::ErrorKind::OutOfMemory,
268 Error::Malformed(_) => io::ErrorKind::InvalidInput,
269 Error::Closed => io::ErrorKind::BrokenPipe,
270 Error::Reset => io::ErrorKind::ConnectionReset,
271 Error::PeerGone => io::ErrorKind::BrokenPipe,
272 Error::WouldBlock => io::ErrorKind::WouldBlock,
273 Error::Io(error) => error.kind(),
274 };
275 io::Error::new(kind, value)
276 }
277}
278
279#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
282pub struct StreamId(NetId64);
283
284impl StreamId {
285 pub const fn from_net_id(id: NetId64) -> Self {
286 Self(id)
287 }
288
289 pub const fn net_id(self) -> NetId64 {
290 self.0
291 }
292
293 pub const fn kind(self) -> u8 {
294 self.0.kind()
295 }
296
297 pub const fn node(self) -> u16 {
299 self.0.node()
300 }
301
302 pub const fn slot(self) -> u32 {
303 (self.0.counter() & SLOT_MASK) as u32
304 }
305
306 pub const fn generation(self) -> u32 {
307 (self.0.counter() >> SLOT_BITS) as u32
308 }
309
310 fn make(
311 kind: u8,
312 node: u16,
313 slot: u32,
314 generation: u32,
315 ) -> Self {
316 Self(NetId64::make(
317 kind,
318 node,
319 ((generation as u64) << SLOT_BITS) | (slot as u64 & SLOT_MASK),
320 ))
321 }
322}
323
324impl fmt::Display for StreamId {
325 fn fmt(
326 &self,
327 f: &mut fmt::Formatter<'_>,
328 ) -> fmt::Result {
329 self.0.fmt(f)
330 }
331}
332
333impl FromStr for StreamId {
334 type Err = Error;
335
336 fn from_str(text: &str) -> Result<Self> {
337 text.parse::<NetId64>()
338 .map(Self)
339 .map_err(|_| Error::Malformed(text.to_owned()))
340 }
341}
342
343#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
345pub enum Side {
346 A,
347 B,
348}
349
350impl Side {
351 const fn index(self) -> usize {
352 match self {
353 Self::A => 0,
354 Self::B => 1,
355 }
356 }
357
358 const fn write_direction(self) -> usize {
360 self.index()
361 }
362
363 const fn read_direction(self) -> usize {
365 1 - self.index()
366 }
367
368 const fn letter(self) -> char {
369 match self {
370 Self::A => 'a',
371 Self::B => 'b',
372 }
373 }
374}
375
376#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
385pub struct Ticket {
386 pub id: StreamId,
387 pub side: Side,
388 pub epoch: u64,
389}
390
391impl fmt::Display for Ticket {
392 fn fmt(
393 &self,
394 f: &mut fmt::Formatter<'_>,
395 ) -> fmt::Result {
396 write!(f, "{}/{}/{}", self.id, self.side.letter(), self.epoch)
397 }
398}
399
400impl FromStr for Ticket {
401 type Err = Error;
402
403 fn from_str(text: &str) -> Result<Self> {
404 let malformed = || Error::Malformed(text.to_owned());
405 let (rest, epoch) = text.rsplit_once('/').ok_or_else(malformed)?;
406 let (id, side) = rest.rsplit_once('/').ok_or_else(malformed)?;
407 let side = match side {
408 "a" => Side::A,
409 "b" => Side::B,
410 _ => return Err(malformed()),
411 };
412 Ok(Self {
413 id: id.parse().map_err(|_| malformed())?,
414 side,
415 epoch: epoch.parse().map_err(|_| malformed())?,
416 })
417 }
418}
419
420#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
426pub struct Incarnation(u64);
427
428impl Incarnation {
429 pub const fn new(value: u64) -> Self {
430 Self(value)
431 }
432
433 pub const fn get(self) -> u64 {
434 self.0
435 }
436}
437
438#[derive(Clone)]
441pub struct Streams {
442 table: Arc<Table>,
443}
444
445impl Streams {
446 #[cfg(unix)]
449 pub fn with_group(
450 name: &str,
451 capacity: u16,
452 node: NodeId,
453 incarnation: Incarnation,
454 spec: StreamSpec,
455 access: orbit_core::shm::GroupAccess,
456 ) -> Result<Self> {
457 spec.validate()?;
458 Ok(Self {
459 table: table::open_group(
460 name,
461 capacity,
462 node,
463 incarnation,
464 spec,
465 access,
466 false,
467 )?,
468 })
469 }
470
471 #[cfg(unix)]
476 pub fn persisted_spec(
477 name: &str,
478 kind: u8,
479 ) -> Result<Option<(u16, StreamSpec)>> {
480 table::persisted_spec(name, kind)
481 }
482
483 #[cfg(unix)]
485 pub fn prepare_group(
486 name: &str,
487 capacity: u16,
488 spec: StreamSpec,
489 access: orbit_core::shm::GroupAccess,
490 ) -> Result<()> {
491 spec.validate()?;
492 table::open_group(
493 name,
494 capacity,
495 NodeId::ZERO,
496 Incarnation::new(0),
497 spec,
498 access,
499 true,
500 )?;
501 Ok(())
502 }
503
504 #[cfg(unix)]
506 pub fn node_dead_in_group(
507 name: &str,
508 capacity: u16,
509 spec: StreamSpec,
510 node: NodeId,
511 incarnation: Incarnation,
512 access: orbit_core::shm::GroupAccess,
513 ) -> Result<()> {
514 spec.validate()?;
515 if node.get() >= capacity {
516 return Err(Error::Malformed(
517 "recovery node is outside the fleet".into(),
518 ));
519 }
520 let table = table::open_group(
521 name,
522 capacity,
523 NodeId::ZERO,
524 Incarnation::new(0),
525 spec,
526 access,
527 true,
528 )?;
529 table.node_dead(node.get(), incarnation.get());
530 Ok(())
531 }
532
533 #[cfg(unix)]
537 pub fn node_dead_in_shm(
538 name: &str,
539 fleet_capacity: u16,
540 spec: StreamSpec,
541 node: NodeId,
542 incarnation: Incarnation,
543 ) -> Result<bool> {
544 spec.validate()?;
545 if node.get() >= fleet_capacity {
546 return Err(Error::Malformed(
547 "recovery node is outside the fleet".into(),
548 ));
549 }
550 let segment = orbit_core::shm::ring_segment_name(name, spec.kind);
551 match orbit_core::shm::ShmRegion::validate_existing(
552 &segment,
553 segment_size_for(fleet_capacity, spec),
554 )? {
555 orbit_core::shm::ShmValidation::Missing => Ok(false),
556 orbit_core::shm::ShmValidation::Valid { .. } => {
557 table::recover_shm(
558 name,
559 fleet_capacity,
560 spec,
561 node.get(),
562 incarnation.get(),
563 )?;
564 Ok(true)
565 }
566 }
567 }
568
569 pub fn new(
573 fleet: Arc<Fleet>,
574 incarnation: Incarnation,
575 ) -> Result<Self> {
576 Self::with_spec(fleet, incarnation, StreamSpec::DEFAULT)
577 }
578
579 pub fn with_spec(
584 fleet: Arc<Fleet>,
585 incarnation: Incarnation,
586 spec: StreamSpec,
587 ) -> Result<Self> {
588 spec.validate()?;
589 Ok(Self { table: table::open(&fleet, incarnation, spec)? })
590 }
591
592 pub fn kind(&self) -> u8 {
594 self.table.kind()
595 }
596
597 #[cfg(any(
609 target_os = "linux",
610 target_os = "freebsd",
611 target_os = "macos"
612 ))]
613 pub fn readiness(&self) -> Result<Readiness> {
614 self.table.take_readiness()
615 }
616
617 pub fn epoch(&self) -> u64 {
619 self.table.epoch()
620 }
621
622 pub fn create(&self) -> Result<(Endpoint, Ticket)> {
625 let (index, generation) = self.table.allocate()?;
626 let id = StreamId::make(
627 self.table.kind(),
628 self.table.node(),
629 (index % self.table.geometry().lane_capacity) as u32,
630 generation,
631 );
632 let handle = Arc::new(Handle {
633 table: Arc::clone(&self.table),
634 index,
635 id,
636 side: Side::A,
637 });
638 Ok((
639 Endpoint::new(handle),
640 Ticket { id, side: Side::B, epoch: self.table.epoch() },
641 ))
642 }
643
644 pub fn capacity_changed(&self) -> CapacityChanged {
648 CapacityChanged::new(Arc::clone(&self.table))
649 }
650
651 pub fn open(
653 &self,
654 ticket: Ticket,
655 ) -> Result<Endpoint> {
656 let index = self.locate(ticket.id)?;
657 if ticket.epoch != self.table.epoch() {
658 return Err(Error::Stale(ticket.id));
659 }
660 let slot = &self.table.slots()[index];
661 if !slot.is(ticket.id.generation()) {
662 return Err(Error::Stale(ticket.id));
663 }
664 let claimed = &slot.claimed[ticket.side.index()];
665 if claimed
666 .compare_exchange(
667 SIDE_FREE,
668 SIDE_CLAIMED,
669 Ordering::SeqCst,
670 Ordering::SeqCst,
671 )
672 .is_err()
673 {
674 return Err(Error::AlreadyClaimed(ticket));
675 }
676 if !slot.is(ticket.id.generation()) {
679 let _ = claimed.compare_exchange(
680 SIDE_CLAIMED,
681 SIDE_FREE,
682 Ordering::SeqCst,
683 Ordering::SeqCst,
684 );
685 return Err(Error::Stale(ticket.id));
686 }
687 slot.incarnation[ticket.side.index()]
688 .store(self.table.incarnation(), Ordering::Release);
689 slot.node[ticket.side.index()]
690 .store(self.table.node(), Ordering::Release);
691 for direction in &slot.directions {
698 self.table.notify(index, slot, direction, table::Ring::Nobody);
699 }
700 Ok(Endpoint::new(Arc::new(Handle {
701 table: Arc::clone(&self.table),
702 index,
703 id: ticket.id,
704 side: ticket.side,
705 })))
706 }
707
708 pub fn rearm(
744 &self,
745 id: StreamId,
746 ) -> Result<()> {
747 let slot = self.rearmable_slot(id)?;
748 let far = Side::B.index();
749 for direction in &slot.directions {
750 direction.clear();
751 }
752 slot.node[far].store(u16::MAX, Ordering::Relaxed);
753 slot.incarnation[far].store(0, Ordering::Relaxed);
754 slot.claimed[far]
758 .compare_exchange(
759 SIDE_RELEASED,
760 SIDE_FREE,
761 Ordering::SeqCst,
762 Ordering::SeqCst,
763 )
764 .map_err(|_| Error::Stale(id))?;
765 Ok(())
766 }
767
768 pub fn can_create(&self) -> bool {
773 self.table.has_capacity()
774 }
775
776 pub fn can_rearm(
782 &self,
783 id: StreamId,
784 ) -> Result<bool> {
785 match self.rearmable_slot(id) {
786 Ok(_) => Ok(true),
787 Err(Error::WouldBlock) => Ok(false),
788 Err(error) => Err(error),
789 }
790 }
791
792 fn rearmable_slot(
793 &self,
794 id: StreamId,
795 ) -> Result<&Slot> {
796 let index = self.locate(id)?;
797 let slot = &self.table.slots()[index];
798 if !slot.is(id.generation()) {
799 return Err(Error::Stale(id));
800 }
801 let far = Side::B.index();
802 match slot.claimed[far].load(Ordering::SeqCst) {
803 SIDE_RELEASED => {}
804 SIDE_CLAIMED => return Err(Error::WouldBlock),
805 _ => return Err(Error::Stale(id)),
806 }
807 Ok(slot)
808 }
809
810 pub fn is_live(
812 &self,
813 id: StreamId,
814 ) -> bool {
815 self.locate(id)
816 .map(|index| self.table.slots()[index].is(id.generation()))
817 .unwrap_or(false)
818 }
819
820 pub fn offer(
826 &self,
827 ticket: Ticket,
828 to: NodeId,
829 ) -> Result<()> {
830 let index = self.locate(ticket.id)?;
831 if ticket.epoch != self.table.epoch()
832 || !self.table.slots()[index].is(ticket.id.generation())
833 {
834 return Err(Error::Stale(ticket.id));
835 }
836 self.table.offer(index, to.get())
837 }
838
839 pub fn take_offer(&self) -> Option<Ticket> {
842 while let Some(index) = self.table.take_offer() {
843 if let Some(ticket) = self.ticket_for(index) {
844 return Some(ticket);
845 }
846 }
847 None
848 }
849
850 #[cfg(any(
854 target_os = "linux",
855 target_os = "freebsd",
856 target_os = "macos"
857 ))]
858 pub fn take_offer_timeout(
859 &self,
860 timeout: Duration,
861 ) -> Result<Option<Ticket>> {
862 let deadline = Instant::now() + timeout;
863 loop {
864 if let Some(ticket) = self.take_offer() {
865 return Ok(Some(ticket));
866 }
867 let left = deadline.saturating_duration_since(Instant::now());
868 if left.is_zero() {
869 return Ok(self.take_offer());
870 }
871 if !self.table.wait_offer_timeout(left)? {
872 return Ok(self.take_offer());
873 }
874 }
875 }
876
877 pub fn blocking_take_offer(&self) -> Result<Ticket> {
879 loop {
880 if let Some(ticket) = self.take_offer() {
881 return Ok(ticket);
882 }
883 self.table.wait_offer()?;
884 }
885 }
886
887 pub fn poll_take_offer(
890 &self,
891 cx: &mut std::task::Context<'_>,
892 ) -> std::task::Poll<Result<Ticket>> {
893 if let Some(ticket) = self.take_offer() {
894 return std::task::Poll::Ready(Ok(ticket));
895 }
896 if let Err(error) = self.table.register_offer(cx.waker()) {
897 return std::task::Poll::Ready(Err(error));
898 }
899 match self.take_offer() {
900 Some(ticket) => std::task::Poll::Ready(Ok(ticket)),
901 None => std::task::Poll::Pending,
902 }
903 }
904
905 pub fn owners(&self) -> Vec<(NodeId, Incarnation)> {
911 let mut seen: Vec<(NodeId, Incarnation)> = Vec::new();
912 for slot in self.table.slots() {
913 if slot.state.load(Ordering::Acquire) != SLOT_LIVE {
914 continue;
915 }
916 for side in 0..2 {
917 if slot.claimed[side].load(Ordering::Acquire) != SIDE_CLAIMED {
918 continue;
919 }
920 let holder = (
921 NodeId::new(slot.node[side].load(Ordering::Acquire)),
922 Incarnation::new(
923 slot.incarnation[side].load(Ordering::Acquire),
924 ),
925 );
926 if !seen.contains(&holder) {
927 seen.push(holder);
928 }
929 }
930 }
931 seen.sort_by_key(|(node, incarnation)| (node.get(), incarnation.get()));
932 seen
933 }
934
935 pub fn node_dead(
942 &self,
943 node: NodeId,
944 incarnation: Incarnation,
945 ) {
946 self.table.node_dead(node.get(), incarnation.get());
947 }
948
949 fn exchange_readers(
950 &self,
951 node: NodeId,
952 incarnation: Incarnation,
953 ) -> impl Iterator<Item = (StreamId, Side)> + '_ {
954 let geometry = self.table.geometry();
955 self.table.slots().iter().enumerate().flat_map(move |(index, slot)| {
956 let generation = slot.generation.load(Ordering::Acquire);
957 [Side::A, Side::B].into_iter().filter_map(move |side| {
958 let index_side = side.index();
959 (slot.claimed[index_side].load(Ordering::SeqCst)
960 == SIDE_CLAIMED
961 && slot.node[index_side].load(Ordering::Acquire)
962 == node.get()
963 && slot.incarnation[index_side].load(Ordering::Acquire)
964 == incarnation.get()
965 && slot.is(generation))
966 .then(|| {
967 (
968 StreamId::make(
969 self.table.kind(),
970 (index / geometry.lane_capacity) as u16,
971 (index % geometry.lane_capacity) as u32,
972 generation,
973 ),
974 side,
975 )
976 })
977 })
978 })
979 }
980
981 pub fn supersede(&self) {
993 self.table.supersede();
994 }
995
996 pub fn reset_all(&self) {
999 self.table.reset_all();
1000 }
1001
1002 fn ticket_for(
1003 &self,
1004 index: usize,
1005 ) -> Option<Ticket> {
1006 let slot = &self.table.slots()[index];
1007 let generation = slot.generation.load(Ordering::Acquire);
1008 if !slot.is(generation) {
1009 return None;
1010 }
1011 let lane_capacity = self.table.geometry().lane_capacity;
1012 let id = StreamId::make(
1013 self.table.kind(),
1014 (index / lane_capacity) as u16,
1015 (index % lane_capacity) as u32,
1016 generation,
1017 );
1018 Some(Ticket { id, side: Side::B, epoch: self.table.epoch() })
1019 }
1020
1021 #[cfg(unix)]
1024 pub fn unlink(&self) -> Result<()> {
1025 self.table.unlink()
1026 }
1027
1028 fn locate(
1029 &self,
1030 id: StreamId,
1031 ) -> Result<usize> {
1032 let geometry = self.table.geometry();
1033 if id.kind() != self.table.kind()
1034 || usize::from(id.node()) >= geometry.fleet_capacity
1035 || id.slot() as usize >= geometry.lane_capacity
1036 || id.generation() == 0
1037 || id.generation() > GENERATION_MASK
1038 {
1039 return Err(Error::Malformed(id.to_string()));
1040 }
1041 Ok(usize::from(id.node()) * geometry.lane_capacity + id.slot() as usize)
1042 }
1043}
1044
1045struct Handle {
1047 table: Arc<Table>,
1048 index: usize,
1049 id: StreamId,
1050 side: Side,
1051}
1052
1053impl Handle {
1054 fn peer_if(
1058 &self,
1059 ring: bool,
1060 ) -> table::Ring {
1061 if ring {
1062 table::Ring::Side(1 - self.side.index())
1063 } else {
1064 table::Ring::Nobody
1065 }
1066 }
1067
1068 fn slot(&self) -> Result<&Slot> {
1069 let slot = &self.table.slots()[self.index];
1070 if slot.is(self.id.generation()) {
1071 Ok(slot)
1072 } else {
1073 Err(Error::Stale(self.id))
1074 }
1075 }
1076
1077 fn peer_is_unclaimed(&self) -> bool {
1078 self.slot().is_ok_and(|slot| {
1079 slot.claimed[1 - self.side.index()].load(Ordering::SeqCst)
1080 == SIDE_FREE
1081 })
1082 }
1083
1084 fn close_unclaimed_peer(&self) -> bool {
1085 self.slot().is_ok_and(|slot| {
1086 slot.claimed[1 - self.side.index()]
1087 .compare_exchange(
1088 SIDE_FREE,
1089 SIDE_RELEASED,
1090 Ordering::SeqCst,
1091 Ordering::SeqCst,
1092 )
1093 .is_ok()
1094 })
1095 }
1096
1097 fn exchange_reset_code(&self) -> Result<Option<u32>> {
1098 let slot = self.slot()?;
1099 let direction = Self::direction(slot, self.side.read_direction());
1100 let encoded = direction.exchange_reset_code.load(Ordering::Acquire);
1101 Ok((encoded != 0).then_some(encoded as u32))
1102 }
1103
1104 fn direction(
1105 slot: &Slot,
1106 direction: usize,
1107 ) -> &Direction {
1108 &slot.directions[direction]
1109 }
1110
1111 fn try_write(
1113 &self,
1114 buf: &[u8],
1115 ) -> Result<usize> {
1116 let slot = self.slot()?;
1117 let direction = Self::direction(slot, self.side.write_direction());
1118 clear_exact_wait(direction);
1119 self.write_bytes(slot, direction, buf)
1120 }
1121
1122 #[cfg(feature = "tokio")]
1125 fn try_write_ready(
1126 &self,
1127 buf: &[u8],
1128 ) -> Result<usize> {
1129 let slot = self.slot()?;
1130 let direction = Self::direction(slot, self.side.write_direction());
1131 self.write_bytes(slot, direction, buf)
1132 }
1133
1134 fn write_bytes(
1135 &self,
1136 slot: &Slot,
1137 direction: &Direction,
1138 buf: &[u8],
1139 ) -> Result<usize> {
1140 let flags = direction.flags();
1141 if flags & FLAG_RESET != 0 {
1142 return Err(Error::Reset);
1143 }
1144 if flags & FLAG_FIN != 0 {
1145 return Err(Error::Closed);
1146 }
1147 if flags & FLAG_READER_GONE != 0 {
1148 return Err(Error::PeerGone);
1149 }
1150 let buffer_bytes = self.table.geometry().buffer_bytes;
1151 let head = direction.head.load(Ordering::Relaxed);
1152 let tail = direction.tail.load(Ordering::Acquire);
1153 let free = buffer_bytes - (head - tail) as usize;
1154 if buf.is_empty() {
1155 return Ok(0);
1156 }
1157 if free == 0 {
1158 return Err(Error::WouldBlock);
1159 }
1160 let len = buf.len().min(free);
1161 let base = self.table.buffer(self.index, self.side.write_direction());
1162 let start = (head as usize) & (buffer_bytes - 1);
1163 let first = len.min(buffer_bytes - start);
1164 unsafe {
1167 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
1168 std::ptr::copy_nonoverlapping(
1169 buf.as_ptr().add(first),
1170 base,
1171 len - first,
1172 );
1173 }
1174 direction.head.store(head + len as u64, Ordering::SeqCst);
1175 std::sync::atomic::fence(Ordering::SeqCst);
1185 let drained = direction.tail.load(Ordering::SeqCst) >= head;
1186 self.table.notify(self.index, slot, direction, self.peer_if(drained));
1187 Ok(len)
1188 }
1189
1190 fn try_write_exact(
1192 &self,
1193 buf: &[u8],
1194 ) -> Result<()> {
1195 let slot = self.slot()?;
1196 let direction = Self::direction(slot, self.side.write_direction());
1197 clear_exact_wait(direction);
1198 let flags = direction.flags();
1199 if flags & FLAG_RESET != 0 {
1200 return Err(Error::Reset);
1201 }
1202 if flags & FLAG_FIN != 0 {
1203 return Err(Error::Closed);
1204 }
1205 if flags & FLAG_READER_GONE != 0 {
1206 return Err(Error::PeerGone);
1207 }
1208 let buffer_bytes = self.table.geometry().buffer_bytes;
1209 if buf.len() > buffer_bytes {
1210 return Err(Error::Malformed(format!(
1211 "atomic stream write is {} bytes; ring capacity is {buffer_bytes}",
1212 buf.len()
1213 )));
1214 }
1215 if buf.is_empty() {
1216 return Ok(());
1217 }
1218 let head = direction.head.load(Ordering::Relaxed);
1219 let tail = direction.tail.load(Ordering::Acquire);
1220 if buffer_bytes - ((head - tail) as usize) < buf.len() {
1221 return Err(Error::WouldBlock);
1222 }
1223 let base = self.table.buffer(self.index, self.side.write_direction());
1224 let start = (head as usize) & (buffer_bytes - 1);
1225 let first = buf.len().min(buffer_bytes - start);
1226 unsafe {
1228 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
1229 std::ptr::copy_nonoverlapping(
1230 buf.as_ptr().add(first),
1231 base,
1232 buf.len() - first,
1233 );
1234 }
1235 direction.head.store(head + buf.len() as u64, Ordering::SeqCst);
1236 std::sync::atomic::fence(Ordering::SeqCst);
1237 let drained = direction.tail.load(Ordering::SeqCst) >= head;
1238 self.table.notify(self.index, slot, direction, self.peer_if(drained));
1239 Ok(())
1240 }
1241
1242 fn try_read(
1245 &self,
1246 buf: &mut [u8],
1247 ) -> Result<usize> {
1248 let slot = self.slot()?;
1249 let direction = Self::direction(slot, self.side.read_direction());
1250 let flags = direction.flags();
1251 if flags & FLAG_RESET != 0 {
1252 return Err(Error::Reset);
1253 }
1254 let head = direction.head.load(Ordering::Acquire);
1255 let tail = direction.tail.load(Ordering::Relaxed);
1256 let available = (head - tail) as usize;
1257 if buf.is_empty() {
1258 return Ok(0);
1259 }
1260 if available == 0 {
1261 if flags & FLAG_FIN != 0 {
1262 return Ok(0);
1263 }
1264 return Err(Error::WouldBlock);
1265 }
1266 let len = buf.len().min(available);
1267 let base = self.table.buffer(self.index, self.side.read_direction());
1268 let buffer_bytes = self.table.geometry().buffer_bytes;
1269 let start = (tail as usize) & (buffer_bytes - 1);
1270 let first = len.min(buffer_bytes - start);
1271 unsafe {
1275 std::ptr::copy_nonoverlapping(
1276 base.add(start),
1277 buf.as_mut_ptr(),
1278 first,
1279 );
1280 std::ptr::copy_nonoverlapping(
1281 base,
1282 buf.as_mut_ptr().add(first),
1283 len - first,
1284 );
1285 }
1286 direction.tail.store(tail + len as u64, Ordering::SeqCst);
1287 std::sync::atomic::fence(Ordering::SeqCst);
1290 let wake = writer_gained_capacity(
1291 direction,
1292 buffer_bytes,
1293 tail,
1294 tail + len as u64,
1295 );
1296 self.table.notify(self.index, slot, direction, self.peer_if(wake));
1297 Ok(len)
1298 }
1299
1300 fn try_read_exact(
1302 &self,
1303 buf: &mut [u8],
1304 head_order: Ordering,
1305 ) -> Result<()> {
1306 let slot = self.slot()?;
1307 let direction = Self::direction(slot, self.side.read_direction());
1308 let flags = direction.flags();
1309 if flags & FLAG_RESET != 0 {
1310 return Err(Error::Reset);
1311 }
1312 let buffer_bytes = self.table.geometry().buffer_bytes;
1313 if buf.len() > buffer_bytes {
1314 return Err(Error::Malformed(format!(
1315 "atomic stream read is {} bytes; ring capacity is {buffer_bytes}",
1316 buf.len()
1317 )));
1318 }
1319 if buf.is_empty() {
1320 return Ok(());
1321 }
1322 let head = direction.head.load(head_order);
1323 let tail = direction.tail.load(Ordering::Relaxed);
1324 let available = (head - tail) as usize;
1325 if available < buf.len() {
1326 if flags & FLAG_FIN != 0 {
1327 return Err(Error::Malformed(format!(
1328 "stream ended with {available} bytes of a {} byte frame",
1329 buf.len()
1330 )));
1331 }
1332 return Err(Error::WouldBlock);
1333 }
1334 let base = self.table.buffer(self.index, self.side.read_direction());
1335 let start = (tail as usize) & (buffer_bytes - 1);
1336 let first = buf.len().min(buffer_bytes - start);
1337 unsafe {
1339 std::ptr::copy_nonoverlapping(
1340 base.add(start),
1341 buf.as_mut_ptr(),
1342 first,
1343 );
1344 std::ptr::copy_nonoverlapping(
1345 base,
1346 buf.as_mut_ptr().add(first),
1347 buf.len() - first,
1348 );
1349 }
1350 direction.tail.store(tail + buf.len() as u64, Ordering::SeqCst);
1351 std::sync::atomic::fence(Ordering::SeqCst);
1352 let wake = writer_gained_capacity(
1353 direction,
1354 buffer_bytes,
1355 tail,
1356 tail + buf.len() as u64,
1357 );
1358 self.table.notify(self.index, slot, direction, self.peer_if(wake));
1359 Ok(())
1360 }
1361
1362 fn set_flag(
1363 &self,
1364 direction_index: usize,
1365 flag: u8,
1366 ) {
1367 if let Ok(slot) = self.slot() {
1368 let direction = Self::direction(slot, direction_index);
1369 let before = direction.flags.fetch_or(flag, Ordering::SeqCst);
1370 if before & flag == 0 {
1371 let ring = flag != FLAG_READER_GONE
1372 || before & FLAG_CLOSE_WATCH != 0
1373 || writer_may_park(direction, before);
1374 self.table.notify(
1375 self.index,
1376 slot,
1377 direction,
1378 self.peer_if(ring),
1379 );
1380 }
1381 }
1382 }
1383
1384 fn set_flag_quietly(
1386 &self,
1387 direction_index: usize,
1388 flag: u8,
1389 ) {
1390 if let Ok(slot) = self.slot() {
1391 let direction = Self::direction(slot, direction_index);
1392 if direction.flags.fetch_or(flag, Ordering::SeqCst) & flag == 0 {
1393 self.table.notify(
1394 self.index,
1395 slot,
1396 direction,
1397 table::Ring::Nobody,
1398 );
1399 }
1400 }
1401 }
1402
1403 #[cfg(any(
1406 target_os = "linux",
1407 target_os = "freebsd",
1408 target_os = "macos"
1409 ))]
1410 fn wait_until_timeout(
1411 &self,
1412 direction_index: usize,
1413 ready: impl Fn(&Direction) -> bool,
1414 timeout: Duration,
1415 ) -> Result<bool> {
1416 let deadline = Instant::now() + timeout;
1417 loop {
1418 let slot = self.slot()?;
1419 let direction = Self::direction(slot, direction_index);
1420 if ready(direction) {
1421 return Ok(true);
1422 }
1423 let left = deadline.saturating_duration_since(Instant::now());
1424 if left.is_zero() {
1425 return Ok(false);
1426 }
1427 direction.waiters.fetch_add(1, Ordering::SeqCst);
1428 let since = direction.changes.load(Ordering::SeqCst);
1429 let outcome = if ready(direction) {
1430 Ok(true)
1431 } else {
1432 wait_on_timeout(&direction.changes, since, left)
1433 };
1434 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1435 if !outcome? {
1436 let slot = self.slot()?;
1439 return Ok(ready(Self::direction(slot, direction_index)));
1440 }
1441 }
1442 }
1443
1444 fn wait_until(
1445 &self,
1446 direction_index: usize,
1447 ready: impl Fn(&Direction) -> bool,
1448 ) -> Result<()> {
1449 loop {
1450 let slot = self.slot()?;
1451 let direction = Self::direction(slot, direction_index);
1452 if ready(direction) {
1453 return Ok(());
1454 }
1455 direction.waiters.fetch_add(1, Ordering::SeqCst);
1456 let since = direction.changes.load(Ordering::SeqCst);
1457 let outcome = if ready(direction) {
1458 Ok(())
1459 } else {
1460 wait_on(&direction.changes, since)
1461 };
1462 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1463 outcome?;
1464 }
1465 }
1466
1467 fn poll_ready(
1470 &self,
1471 interest: Interest,
1472 ready: impl Fn(&Direction) -> bool,
1473 cx: &mut std::task::Context<'_>,
1474 ) -> std::task::Poll<Result<()>> {
1475 let slot = match self.slot() {
1476 Ok(slot) => slot,
1477 Err(error) => return std::task::Poll::Ready(Err(error)),
1478 };
1479 let direction = Self::direction(slot, interest.direction);
1480 if ready(direction) {
1481 return std::task::Poll::Ready(Ok(()));
1482 }
1483 if let Err(error) =
1484 self.table.register(self.index, interest, cx.waker())
1485 {
1486 return std::task::Poll::Ready(Err(error));
1487 }
1488 std::sync::atomic::fence(Ordering::SeqCst);
1491 if ready(direction) {
1492 std::task::Poll::Ready(Ok(()))
1493 } else {
1494 std::task::Poll::Pending
1495 }
1496 }
1497}
1498
1499fn clear_exact_wait(direction: &Direction) {
1500 if direction.exact_needed.load(Ordering::Relaxed) != 0 {
1503 direction.exact_needed.store(0, Ordering::SeqCst);
1504 }
1505}
1506
1507fn writer_gained_capacity(
1508 direction: &Direction,
1509 capacity: usize,
1510 before_tail: u64,
1511 after_tail: u64,
1512) -> bool {
1513 let head = direction.head.load(Ordering::SeqCst);
1514 let needed = direction.exact_needed.load(Ordering::SeqCst) as usize;
1515 if needed == 0 {
1516 return head - before_tail >= capacity as u64;
1517 }
1518 if direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0 {
1519 return false;
1520 }
1521 let before = capacity.saturating_sub((head - before_tail) as usize);
1522 let after = capacity.saturating_sub((head - after_tail) as usize);
1523 before < needed && after >= needed
1524}
1525
1526fn writer_may_park(
1540 direction: &Direction,
1541 before: u8,
1542) -> bool {
1543 if before & (FLAG_FIN | FLAG_RESET) != 0 {
1544 return false;
1545 }
1546 std::sync::atomic::fence(Ordering::SeqCst);
1547 direction.head.load(Ordering::SeqCst)
1548 != direction.tail.load(Ordering::SeqCst)
1549}
1550
1551fn readable(direction: &Direction) -> bool {
1552 direction.head.load(Ordering::Acquire)
1553 != direction.tail.load(Ordering::Relaxed)
1554 || direction.flags() & (FLAG_FIN | FLAG_RESET) != 0
1555}
1556
1557fn writable(
1558 direction: &Direction,
1559 buffer_bytes: usize,
1560) -> bool {
1561 let queued = (direction.head.load(Ordering::Relaxed)
1562 - direction.tail.load(Ordering::Acquire)) as usize;
1563 queued < buffer_bytes
1564 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1565}
1566
1567fn writable_exact(
1568 direction: &Direction,
1569 buffer_bytes: usize,
1570 needed: usize,
1571) -> bool {
1572 let queued = (direction.head.load(Ordering::Relaxed)
1573 - direction.tail.load(Ordering::Acquire)) as usize;
1574 buffer_bytes.saturating_sub(queued) >= needed
1575 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1576}
1577
1578impl Drop for Handle {
1579 fn drop(&mut self) {
1583 let slot = &self.table.slots()[self.index];
1584 if !slot.is(self.id.generation()) {
1585 return;
1586 }
1587 if slot.claimed[self.side.index()]
1590 .compare_exchange(
1591 SIDE_CLAIMED,
1592 SIDE_RELEASED,
1593 Ordering::SeqCst,
1594 Ordering::SeqCst,
1595 )
1596 .is_err()
1597 {
1598 return;
1599 }
1600 let other = slot.claimed[1 - self.side.index()].load(Ordering::SeqCst);
1604 if other != SIDE_CLAIMED {
1605 let _ = slot.state.compare_exchange(
1606 SLOT_LIVE,
1607 SLOT_EMPTY,
1608 Ordering::SeqCst,
1609 Ordering::SeqCst,
1610 );
1611 for direction in &slot.directions {
1612 direction.changes.fetch_add(1, Ordering::SeqCst);
1615 if direction.waiters.load(Ordering::SeqCst) > 0 {
1616 wake_on(&direction.changes);
1617 }
1618 }
1619 self.table.registry.wake(self.index);
1620 }
1621 if self.side == Side::B || other != SIDE_CLAIMED {
1624 self.table.notify_capacity(self.index);
1625 }
1626 }
1627}
1628
1629pub struct Endpoint {
1633 read: ReadHalf,
1634 write: WriteHalf,
1635}
1636
1637impl Endpoint {
1638 fn new(handle: Arc<Handle>) -> Self {
1639 Self {
1640 read: ReadHalf { handle: Arc::clone(&handle) },
1641 write: WriteHalf { handle, shut: false },
1642 }
1643 }
1644
1645 pub fn id(&self) -> StreamId {
1646 self.read.handle.id
1647 }
1648
1649 pub fn side(&self) -> Side {
1650 self.read.handle.side
1651 }
1652
1653 pub fn split(self) -> (ReadHalf, WriteHalf) {
1654 (self.read, self.write)
1655 }
1656
1657 pub fn rejoin(
1668 read: ReadHalf,
1669 write: WriteHalf,
1670 ) -> Result<Self> {
1671 if read.id() != write.id() {
1672 return Err(Error::Stale(read.id()));
1673 }
1674 Ok(Self { read, write })
1675 }
1676
1677 pub fn rearm(
1690 &mut self,
1691 streams: &Streams,
1692 ) -> Result<()> {
1693 if self.side() != Side::A {
1694 return Err(Error::Malformed(
1695 "only the creator side of a stream can be rearmed".to_owned(),
1696 ));
1697 }
1698 streams.rearm(self.id())?;
1699 self.write.shut = false;
1700 Ok(())
1701 }
1702
1703 pub fn try_read(
1704 &self,
1705 buf: &mut [u8],
1706 ) -> Result<usize> {
1707 self.read.try_read(buf)
1708 }
1709
1710 pub fn try_read_exact(
1711 &self,
1712 buf: &mut [u8],
1713 ) -> Result<()> {
1714 self.read.try_read_exact(buf)
1715 }
1716
1717 pub fn blocking_read(
1718 &self,
1719 buf: &mut [u8],
1720 ) -> Result<usize> {
1721 self.read.blocking_read(buf)
1722 }
1723
1724 pub fn blocking_read_chunk(
1725 &self,
1726 max: usize,
1727 ) -> Result<Bytes> {
1728 self.read.blocking_read_chunk(max)
1729 }
1730
1731 pub fn try_write(
1732 &self,
1733 buf: &[u8],
1734 ) -> Result<usize> {
1735 self.write.try_write(buf)
1736 }
1737
1738 pub fn try_write_exact(
1739 &self,
1740 buf: &[u8],
1741 ) -> Result<()> {
1742 self.write.try_write_exact(buf)
1743 }
1744
1745 pub fn blocking_write(
1746 &self,
1747 buf: &[u8],
1748 ) -> Result<usize> {
1749 self.write.blocking_write(buf)
1750 }
1751
1752 pub fn blocking_write_all(
1753 &self,
1754 buf: &[u8],
1755 ) -> Result<()> {
1756 self.write.blocking_write_all(buf)
1757 }
1758
1759 pub fn finish(&mut self) -> Result<()> {
1760 self.write.finish()
1761 }
1762
1763 pub fn reset(&mut self) {
1764 self.write.reset()
1765 }
1766}
1767
1768pub struct ReadHalf {
1771 handle: Arc<Handle>,
1772}
1773
1774pub struct ResetSignal {
1778 handle: Arc<Handle>,
1779}
1780
1781impl ResetSignal {
1782 pub fn exchange_reset_code(&self) -> Result<Option<u32>> {
1784 self.handle.exchange_reset_code()
1785 }
1786
1787 pub fn poll_reset(
1790 &self,
1791 cx: &mut std::task::Context<'_>,
1792 ) -> std::task::Poll<Result<()>> {
1793 self.handle.poll_ready(
1794 Interest {
1795 direction: self.handle.side.read_direction(),
1796 writer: false,
1797 reset: true,
1798 },
1799 |direction| direction.flags() & FLAG_RESET != 0,
1800 cx,
1801 )
1802 }
1803}
1804
1805impl ReadHalf {
1806 pub fn reset_signal(&self) -> ResetSignal {
1809 ResetSignal { handle: Arc::clone(&self.handle) }
1810 }
1811
1812 pub fn id(&self) -> StreamId {
1813 self.handle.id
1814 }
1815
1816 pub fn try_read(
1819 &self,
1820 buf: &mut [u8],
1821 ) -> Result<usize> {
1822 self.handle.try_read(buf)
1823 }
1824
1825 pub fn try_read_exact(
1828 &self,
1829 buf: &mut [u8],
1830 ) -> Result<()> {
1831 self.handle.try_read_exact(buf, Ordering::Acquire)
1832 }
1833
1834 pub(crate) fn try_drain_exchange_frame(
1835 &self,
1836 frame: &mut [u8],
1837 ) -> Result<()> {
1838 self.handle.try_read_exact(frame, Ordering::SeqCst)
1840 }
1841
1842 pub(crate) fn exchange_peer_node(&self) -> Result<u16> {
1843 let slot = self.handle.slot()?;
1844 Ok(slot.node[1 - self.handle.side.index()].load(Ordering::Acquire))
1845 }
1846
1847 pub(crate) fn exchange_reset_code(&self) -> Result<Option<u32>> {
1848 self.handle.exchange_reset_code()
1849 }
1850
1851 pub fn blocking_read(
1854 &self,
1855 buf: &mut [u8],
1856 ) -> Result<usize> {
1857 loop {
1858 match self.handle.try_read(buf) {
1859 Err(Error::WouldBlock) => self.wait_readable()?,
1860 other => return other,
1861 }
1862 }
1863 }
1864
1865 pub fn stop_reading(&self) {
1873 self.handle
1874 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1875 }
1876
1877 pub fn blocking_read_chunk(
1879 &self,
1880 max: usize,
1881 ) -> Result<Bytes> {
1882 let mut chunk = BytesMut::zeroed(max);
1883 let len = self.blocking_read(&mut chunk)?;
1884 chunk.truncate(len);
1885 Ok(chunk.freeze())
1886 }
1887
1888 #[cfg(any(
1890 target_os = "linux",
1891 target_os = "freebsd",
1892 target_os = "macos"
1893 ))]
1894 pub fn wait_readable_timeout(
1895 &self,
1896 timeout: Duration,
1897 ) -> Result<bool> {
1898 self.handle.wait_until_timeout(
1899 self.handle.side.read_direction(),
1900 readable,
1901 timeout,
1902 )
1903 }
1904
1905 pub fn wait_readable(&self) -> Result<()> {
1907 self.handle.wait_until(self.handle.side.read_direction(), readable)
1908 }
1909
1910 pub fn poll_readable(
1915 &self,
1916 cx: &mut std::task::Context<'_>,
1917 ) -> std::task::Poll<Result<()>> {
1918 self.handle.poll_ready(
1919 Interest {
1920 direction: self.handle.side.read_direction(),
1921 writer: false,
1922 reset: false,
1923 },
1924 readable,
1925 cx,
1926 )
1927 }
1928}
1929
1930impl Drop for ReadHalf {
1931 fn drop(&mut self) {
1932 self.handle
1933 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1934 }
1935}
1936
1937pub struct WriteHalf {
1940 handle: Arc<Handle>,
1941 shut: bool,
1942}
1943
1944struct CloseWatch<'a>(&'a Handle);
1946
1947impl Drop for CloseWatch<'_> {
1948 fn drop(&mut self) {
1949 if let Ok(slot) = self.0.slot() {
1950 Handle::direction(slot, self.0.side.write_direction())
1951 .flags
1952 .fetch_and(!FLAG_CLOSE_WATCH, Ordering::SeqCst);
1953 }
1954 }
1955}
1956
1957impl Handle {
1958 fn exchange_write_status(&self) -> Result<()> {
1960 let slot = self.slot()?;
1961 let direction = Handle::direction(slot, self.side.write_direction());
1962 let flags = direction.flags.load(Ordering::SeqCst);
1963 if flags & FLAG_READER_GONE != 0 {
1964 Err(Error::PeerGone)
1965 } else if flags & FLAG_RESET != 0 {
1966 Err(Error::Reset)
1967 } else if flags & FLAG_FIN != 0 {
1968 Err(Error::Closed)
1969 } else {
1970 Ok(())
1971 }
1972 }
1973}
1974
1975#[derive(Clone)]
1978pub(crate) struct WriteStatus(Arc<Handle>);
1979
1980impl WriteStatus {
1981 pub(crate) fn is_ended(&self) -> bool {
1984 self.0.exchange_write_status().is_err()
1985 }
1986}
1987
1988impl WriteHalf {
1989 pub(crate) fn exchange_write_status(&self) -> Result<()> {
1990 self.handle.exchange_write_status()
1991 }
1992
1993 pub(crate) fn write_status(&self) -> WriteStatus {
1994 WriteStatus(Arc::clone(&self.handle))
1995 }
1996
1997 pub(crate) fn reset_exchange(
1998 &mut self,
1999 code: u32,
2000 ) -> Result<()> {
2001 let slot = self.handle.slot()?;
2002 let direction =
2003 Handle::direction(slot, self.handle.side.write_direction());
2004 direction
2005 .exchange_reset_code
2006 .store((1_u64 << 32) | u64::from(code), Ordering::Release);
2007 self.reset();
2008 Ok(())
2009 }
2010
2011 pub async fn peer_closed(&mut self) -> Result<()> {
2017 let slot = self.handle.slot()?;
2018 let direction =
2019 Handle::direction(slot, self.handle.side.write_direction());
2020 direction.flags.fetch_or(FLAG_CLOSE_WATCH, Ordering::SeqCst);
2021 let _watch = CloseWatch(&self.handle);
2022 std::future::poll_fn(|cx| {
2023 self.handle.poll_ready(
2024 Interest {
2025 direction: self.handle.side.write_direction(),
2026 writer: true,
2027 reset: false,
2028 },
2029 |direction| {
2030 direction.flags() & (FLAG_RESET | FLAG_READER_GONE) != 0
2031 },
2032 cx,
2033 )
2034 })
2035 .await
2036 }
2037
2038 pub fn poll_peer_dead(
2044 &self,
2045 cx: &mut std::task::Context<'_>,
2046 ) -> std::task::Poll<Result<()>> {
2047 let slot = match self.handle.slot() {
2048 Ok(slot) => slot,
2049 Err(error) => return std::task::Poll::Ready(Err(error)),
2050 };
2051 let peer = &slot.claimed[1 - self.handle.side.index()];
2052 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
2053 return std::task::Poll::Ready(Ok(()));
2054 }
2055 let interest = Interest {
2056 direction: self.handle.side.write_direction(),
2057 writer: true,
2058 reset: false,
2059 };
2060 if let Err(error) =
2061 self.handle.table.register(self.handle.index, interest, cx.waker())
2062 {
2063 return std::task::Poll::Ready(Err(error));
2064 }
2065 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
2068 std::task::Poll::Ready(Ok(()))
2069 } else {
2070 std::task::Poll::Pending
2071 }
2072 }
2073
2074 pub fn id(&self) -> StreamId {
2075 self.handle.id
2076 }
2077
2078 pub fn try_write(
2080 &self,
2081 buf: &[u8],
2082 ) -> Result<usize> {
2083 self.handle.try_write(buf)
2084 }
2085
2086 pub fn try_write_exact(
2089 &self,
2090 buf: &[u8],
2091 ) -> Result<()> {
2092 self.handle.try_write_exact(buf)
2093 }
2094
2095 pub fn blocking_write(
2097 &self,
2098 buf: &[u8],
2099 ) -> Result<usize> {
2100 loop {
2101 match self.handle.try_write(buf) {
2102 Err(Error::WouldBlock) => self.wait_writable()?,
2103 other => return other,
2104 }
2105 }
2106 }
2107
2108 pub fn blocking_write_all(
2111 &self,
2112 mut buf: &[u8],
2113 ) -> Result<()> {
2114 while !buf.is_empty() {
2115 let written = self.blocking_write(buf)?;
2116 buf = &buf[written..];
2117 }
2118 Ok(())
2119 }
2120
2121 pub fn finish(&mut self) -> Result<()> {
2125 self.shut = true;
2126 self.handle.set_flag(self.handle.side.write_direction(), FLAG_FIN);
2127 Ok(())
2128 }
2129
2130 pub fn finish_quietly(&mut self) {
2137 self.shut = true;
2138 self.handle
2139 .set_flag_quietly(self.handle.side.write_direction(), FLAG_FIN);
2140 }
2141
2142 pub fn reset(&mut self) {
2145 self.shut = true;
2146 self.handle.set_flag(self.handle.side.write_direction(), FLAG_RESET);
2147 }
2148
2149 #[cfg(any(
2153 target_os = "linux",
2154 target_os = "freebsd",
2155 target_os = "macos"
2156 ))]
2157 pub fn wait_writable_timeout(
2158 &self,
2159 timeout: Duration,
2160 ) -> Result<bool> {
2161 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2162 self.handle.wait_until_timeout(
2163 self.handle.side.write_direction(),
2164 |direction| {
2165 clear_exact_wait(direction);
2166 writable(direction, buffer_bytes)
2167 },
2168 timeout,
2169 )
2170 }
2171
2172 pub fn wait_writable(&self) -> Result<()> {
2174 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2175 self.handle.wait_until(
2176 self.handle.side.write_direction(),
2177 |direction| {
2178 clear_exact_wait(direction);
2179 writable(direction, buffer_bytes)
2180 },
2181 )
2182 }
2183
2184 pub fn wait_writable_exact(
2186 &self,
2187 needed: usize,
2188 ) -> Result<()> {
2189 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2190 if needed > buffer_bytes {
2191 return Err(Error::Malformed(format!(
2192 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
2193 )));
2194 }
2195 self.handle.wait_until(
2196 self.handle.side.write_direction(),
2197 |direction| {
2198 clear_exact_wait(direction);
2199 writable_exact(direction, buffer_bytes, needed)
2200 },
2201 )
2202 }
2203
2204 pub fn poll_writable(
2209 &self,
2210 cx: &mut std::task::Context<'_>,
2211 ) -> std::task::Poll<Result<()>> {
2212 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2213 self.handle.poll_ready(
2214 Interest {
2215 direction: self.handle.side.write_direction(),
2216 writer: true,
2217 reset: false,
2218 },
2219 |direction| {
2220 clear_exact_wait(direction);
2221 writable(direction, buffer_bytes)
2222 },
2223 cx,
2224 )
2225 }
2226
2227 pub fn poll_writable_exact(
2229 &self,
2230 needed: usize,
2231 cx: &mut std::task::Context<'_>,
2232 ) -> std::task::Poll<Result<()>> {
2233 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2234 let slot = match self.handle.slot() {
2235 Ok(slot) => slot,
2236 Err(error) => return std::task::Poll::Ready(Err(error)),
2237 };
2238 let direction =
2239 Handle::direction(slot, self.handle.side.write_direction());
2240 if needed > buffer_bytes {
2241 clear_exact_wait(direction);
2242 return std::task::Poll::Ready(Err(Error::Malformed(format!(
2243 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
2244 ))));
2245 }
2246 direction.exact_needed.store(needed as u32, Ordering::SeqCst);
2250 let result = self.handle.poll_ready(
2251 Interest {
2252 direction: self.handle.side.write_direction(),
2253 writer: true,
2254 reset: false,
2255 },
2256 |direction| writable_exact(direction, buffer_bytes, needed),
2257 cx,
2258 );
2259 if result.is_ready() {
2260 clear_exact_wait(direction);
2261 }
2262 result
2263 }
2264}
2265
2266impl Drop for WriteHalf {
2267 fn drop(&mut self) {
2268 if !self.shut {
2269 self.handle
2270 .set_flag(self.handle.side.write_direction(), FLAG_RESET);
2271 }
2272 }
2273}
2274
2275#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2279pub(crate) fn waits_supported() -> bool {
2280 orbit_core::sync::supported()
2281}
2282
2283#[cfg(not(any(
2284 target_os = "linux",
2285 target_os = "freebsd",
2286 target_os = "macos"
2287)))]
2288pub(crate) fn waits_supported() -> bool {
2289 false
2290}
2291
2292#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2296#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2299pub(crate) fn wait_on_timeout(
2300 word: &AtomicU32,
2301 expected: u32,
2302 timeout: Duration,
2303) -> Result<bool> {
2304 orbit_core::sync::wait_word_timeout(word, expected, timeout)
2305 .map_err(Error::Io)
2306}
2307
2308pub(crate) fn wait_on(
2309 word: &AtomicU32,
2310 expected: u32,
2311) -> Result<()> {
2312 orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
2313}
2314
2315#[cfg(not(any(
2316 target_os = "linux",
2317 target_os = "freebsd",
2318 target_os = "macos"
2319)))]
2320pub(crate) fn wait_on(
2321 _word: &AtomicU32,
2322 _expected: u32,
2323) -> Result<()> {
2324 Err(Error::Io(std::io::Error::new(
2325 std::io::ErrorKind::Unsupported,
2326 "orbit-stream needs a platform that can wait on a shared word",
2327 )))
2328}
2329
2330pub(crate) fn wake_on(word: &AtomicU32) {
2332 #[cfg(any(
2333 target_os = "linux",
2334 target_os = "freebsd",
2335 target_os = "macos"
2336 ))]
2337 let _ = orbit_core::sync::wake_word(word);
2338 #[cfg(not(any(
2339 target_os = "linux",
2340 target_os = "freebsd",
2341 target_os = "macos"
2342 )))]
2343 let _ = word;
2344}
2345
2346pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
2347 mutex.lock().unwrap_or_else(|error| error.into_inner())
2348}
2349
2350#[cfg(test)]
2351mod tests {
2352 use std::sync::Arc;
2353 use std::time::Duration;
2354
2355 use orbit_core::{Fleet, NodeId};
2356
2357 use super::{
2358 Error, Incarnation, STREAM_BUFFER_BYTES, STREAM_LANE_CAPACITY, Side,
2359 Streams, Ticket,
2360 };
2361
2362 fn streams(name: &'static str) -> Streams {
2363 Streams::new(
2364 Arc::new(Fleet::join(name, 2).unwrap()),
2365 Incarnation::new(1),
2366 )
2367 .unwrap()
2368 }
2369
2370 #[test]
2371 fn bytes_cross_in_order_and_end_cleanly() {
2372 let streams = streams("stream-order");
2373 let (mut a, ticket) = streams.create().unwrap();
2374 assert_eq!(ticket.side, Side::B);
2375 assert_eq!(ticket.id.node(), 0);
2376 assert_eq!(ticket.epoch, streams.epoch());
2377 let b = streams.open(ticket).unwrap();
2378
2379 a.blocking_write_all(b"hello, ").unwrap();
2380 a.blocking_write_all(b"world").unwrap();
2381 a.finish().unwrap();
2382
2383 let mut out = Vec::new();
2384 let mut buf = [0_u8; 4];
2385 loop {
2386 let n = b.blocking_read(&mut buf).unwrap();
2387 if n == 0 {
2388 break;
2389 }
2390 out.extend_from_slice(&buf[..n]);
2391 }
2392 assert_eq!(out, b"hello, world");
2393 b.blocking_write_all(b"back").unwrap();
2395 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"back");
2396 }
2397
2398 #[test]
2399 fn a_ticket_prints_and_parses() {
2400 let streams = streams("stream-ticket");
2401 let (_a, ticket) = streams.create().unwrap();
2402 let text = ticket.to_string();
2403 assert!(text.contains("/b/"), "{text}");
2404 assert_eq!(text.parse::<Ticket>().unwrap(), ticket);
2405 assert!(matches!("nope".parse::<Ticket>(), Err(Error::Malformed(_))));
2406 assert!(matches!("x/b".parse::<Ticket>(), Err(Error::Malformed(_))));
2407 }
2408
2409 #[test]
2410 fn a_side_is_held_once_and_a_released_stream_goes_stale() {
2411 let streams = streams("stream-claim");
2412 let (a, ticket) = streams.create().unwrap();
2413 let b = streams.open(ticket).unwrap();
2414 assert!(matches!(streams.open(ticket), Err(Error::AlreadyClaimed(_))));
2415 assert!(streams.is_live(ticket.id));
2416
2417 drop(a);
2418 let mut buf = [0_u8; 1];
2420 assert!(matches!(b.blocking_read(&mut buf), Err(Error::Reset)));
2421 drop(b);
2422 assert!(!streams.is_live(ticket.id));
2423 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2424
2425 for _ in 0..STREAM_LANE_CAPACITY - 1 {
2428 drop(streams.create().unwrap());
2429 }
2430 let (_a2, ticket2) = streams.create().unwrap();
2431 assert_eq!(ticket2.id.slot(), ticket.id.slot());
2432 assert_ne!(ticket2.id.generation(), ticket.id.generation());
2433 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2434 }
2435
2436 #[test]
2437 fn a_reset_starts_a_new_epoch_and_refuses_old_tickets() {
2438 let streams = streams("stream-epoch");
2439 let (_a, ticket) = streams.create().unwrap();
2440 let before = streams.epoch();
2441 streams.reset_all();
2442 assert!(streams.epoch() > before);
2443 let (_a2, ticket2) = streams.create().unwrap();
2446 assert_eq!(ticket2.id, ticket.id);
2447 assert_ne!(ticket2.epoch, ticket.epoch);
2448 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2449 assert!(streams.open(ticket2).is_ok());
2450 }
2451
2452 #[test]
2453 fn a_slot_whose_generation_runs_out_is_not_reused() {
2454 let streams = streams("stream-exhaust");
2457 let limit = super::layout::GENERATION_LIMIT as usize;
2458 for _ in 0..STREAM_LANE_CAPACITY * limit {
2459 drop(streams.create().unwrap());
2460 }
2461 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2462 streams.reset_all();
2463 let (_a, ticket) = streams.create().unwrap();
2464 assert_eq!(ticket.id.generation(), 1);
2465 }
2466
2467 #[test]
2468 fn a_full_ring_blocks_the_writer_and_wraps_correctly() {
2469 let streams = streams("stream-wrap");
2470 let (a, ticket) = streams.create().unwrap();
2471 let b = streams.open(ticket).unwrap();
2472
2473 let filler = vec![7_u8; STREAM_BUFFER_BYTES];
2474 assert_eq!(a.try_write(&filler).unwrap(), STREAM_BUFFER_BYTES);
2475 assert!(matches!(a.try_write(b"x"), Err(Error::WouldBlock)));
2476
2477 let mut buf = vec![0_u8; STREAM_BUFFER_BYTES - 5];
2478 assert_eq!(b.blocking_read(&mut buf).unwrap(), STREAM_BUFFER_BYTES - 5);
2479 let tail =
2481 (0..STREAM_BUFFER_BYTES).map(|i| i as u8).collect::<Vec<_>>();
2482 assert_eq!(a.try_write(&tail).unwrap(), STREAM_BUFFER_BYTES - 5);
2483
2484 let mut rest = vec![0_u8; 5];
2485 assert_eq!(b.blocking_read(&mut rest).unwrap(), 5);
2486 assert_eq!(rest, vec![7_u8; 5]);
2487 let mut wrapped = vec![0_u8; STREAM_BUFFER_BYTES];
2488 let mut got = 0;
2489 while got < STREAM_BUFFER_BYTES - 5 {
2490 got += b.blocking_read(&mut wrapped[got..]).unwrap();
2491 }
2492 assert_eq!(&wrapped[..got], &tail[..got]);
2493 }
2494
2495 #[test]
2496 fn exact_frames_never_publish_or_consume_a_prefix() {
2497 let streams = streams("stream-exact-frame");
2498 let (a, ticket) = streams.create().unwrap();
2499 let b = streams.open(ticket).unwrap();
2500
2501 let filler = vec![7_u8; STREAM_BUFFER_BYTES - 32];
2502 a.try_write_exact(&filler).unwrap();
2503 assert!(matches!(
2504 a.try_write_exact(&[9_u8; 64]),
2505 Err(Error::WouldBlock)
2506 ));
2507 let mut consumed = vec![0_u8; filler.len()];
2508 b.try_read_exact(&mut consumed).unwrap();
2509 assert_eq!(consumed, filler);
2510 assert!(matches!(b.try_read(&mut [0_u8; 1]), Err(Error::WouldBlock)));
2511
2512 a.try_write_exact(&[3_u8; 32]).unwrap();
2513 assert!(matches!(
2514 b.try_read_exact(&mut [0_u8; 64]),
2515 Err(Error::WouldBlock)
2516 ));
2517 let mut frame = [0_u8; 32];
2518 b.try_read_exact(&mut frame).unwrap();
2519 assert_eq!(frame, [3_u8; 32]);
2520 }
2521
2522 #[test]
2523 fn a_blocked_reader_is_woken_by_a_write_in_another_thread() {
2524 let streams = streams("stream-wake");
2525 let (a, ticket) = streams.create().unwrap();
2526 let b = streams.open(ticket).unwrap();
2527 let reader = std::thread::spawn(move || {
2528 let mut buf = [0_u8; 8];
2529 let n = b.blocking_read(&mut buf).unwrap();
2530 buf[..n].to_vec()
2531 });
2532 std::thread::sleep(Duration::from_millis(30));
2533 a.blocking_write_all(b"ping").unwrap();
2534 assert_eq!(reader.join().unwrap(), b"ping");
2535 }
2536
2537 #[test]
2538 fn a_reader_gone_fails_the_writer_and_a_lane_fills_up() {
2539 let streams = streams("stream-gone");
2540 let (a, ticket) = streams.create().unwrap();
2541 let b = streams.open(ticket).unwrap();
2542 let (b_read, _b_write) = b.split();
2543 drop(b_read);
2544 assert!(matches!(a.blocking_write(b"nobody"), Err(Error::PeerGone)));
2545
2546 let mut held = vec![a];
2547 for _ in 1..STREAM_LANE_CAPACITY {
2548 held.push(streams.create().unwrap().0);
2549 }
2550 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2551 }
2552
2553 #[test]
2554 fn a_death_report_ends_that_incarnations_sides_only() {
2555 let streams = streams("stream-dead");
2556 let (a, ticket) = streams.create().unwrap();
2557 let b = streams.open(ticket).unwrap();
2558 b.blocking_write_all(b"partial").unwrap();
2559
2560 streams.node_dead(NodeId::ZERO, Incarnation::new(2));
2562 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"partial");
2563
2564 let reader = std::thread::spawn(move || {
2570 let mut buf = [0_u8; 8];
2571 a.blocking_read(&mut buf).map(|_| ())
2572 });
2573 std::thread::sleep(Duration::from_millis(30));
2574 streams.node_dead(NodeId::ZERO, Incarnation::new(1));
2575 assert!(matches!(
2576 reader.join().unwrap(),
2577 Err(Error::Reset | Error::Stale(_))
2578 ));
2579 assert!(matches!(
2580 b.try_write(b"x"),
2581 Err(Error::Reset | Error::Stale(_))
2582 ));
2583 assert!(matches!(
2584 streams.open(ticket),
2585 Err(Error::AlreadyClaimed(_) | Error::Stale(_))
2586 ));
2587 drop(b);
2588 assert!(!streams.is_live(ticket.id));
2589 }
2590
2591 #[test]
2592 fn an_offer_is_taken_once_and_a_stale_offer_is_skipped() {
2593 let streams = streams("stream-offer");
2594 assert!(streams.take_offer().is_none());
2595 let (a, ticket) = streams.create().unwrap();
2596 streams.offer(ticket, NodeId::ZERO).unwrap();
2597 assert_eq!(streams.take_offer(), Some(ticket));
2598 assert!(streams.take_offer().is_none());
2599
2600 let (gone, stale) = streams.create().unwrap();
2601 streams.offer(stale, NodeId::ZERO).unwrap();
2602 drop(gone);
2603 assert!(streams.take_offer().is_none());
2604
2605 let taker = streams.clone();
2606 let waiter = std::thread::spawn(move || taker.blocking_take_offer());
2607 std::thread::sleep(Duration::from_millis(30));
2608 streams.offer(ticket, NodeId::ZERO).unwrap();
2609 assert_eq!(waiter.join().unwrap().unwrap(), ticket);
2610 drop(a);
2611 }
2612
2613 #[test]
2614 fn a_contested_claim_goes_to_exactly_one_opener() {
2615 let streams = streams("stream-race");
2616 for _ in 0..50 {
2617 let (_a, ticket) = streams.create().unwrap();
2618 let openers = (0..4)
2619 .map(|_| {
2620 let streams = streams.clone();
2621 std::thread::spawn(move || streams.open(ticket).is_ok())
2622 })
2623 .collect::<Vec<_>>();
2624 let won = openers
2625 .into_iter()
2626 .map(|opener| opener.join().unwrap())
2627 .filter(|won| *won)
2628 .count();
2629 assert_eq!(won, 1);
2630 }
2631 }
2632
2633 #[test]
2634 fn random_sizes_concatenate_exactly() {
2635 let streams = streams("stream-random");
2636 let (mut a, ticket) = streams.create().unwrap();
2637 let b = streams.open(ticket).unwrap();
2638 let payload = (0..300_000_u32)
2639 .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
2640 .collect::<Vec<_>>();
2641 let expected = payload.clone();
2642 let writer = std::thread::spawn(move || {
2643 let mut rest = &payload[..];
2644 let mut step = 1_usize;
2645 while !rest.is_empty() {
2646 let len = step.min(rest.len());
2647 a.blocking_write_all(&rest[..len]).unwrap();
2648 rest = &rest[len..];
2649 step = (step * 7 + 13) % 9_001 + 1;
2650 }
2651 a.finish().unwrap();
2652 });
2653 let mut out = Vec::new();
2654 let mut step = 3_usize;
2655 loop {
2656 let chunk = b.blocking_read_chunk(step).unwrap();
2657 if chunk.is_empty() {
2658 break;
2659 }
2660 out.extend_from_slice(&chunk);
2661 step = (step * 5 + 11) % 4_099 + 1;
2662 }
2663 writer.join().unwrap();
2664 assert_eq!(out, expected);
2665 }
2666}