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 pub fn supersede(&self) {
961 self.table.supersede();
962 }
963
964 pub fn reset_all(&self) {
967 self.table.reset_all();
968 }
969
970 fn ticket_for(
971 &self,
972 index: usize,
973 ) -> Option<Ticket> {
974 let slot = &self.table.slots()[index];
975 let generation = slot.generation.load(Ordering::Acquire);
976 if !slot.is(generation) {
977 return None;
978 }
979 let lane_capacity = self.table.geometry().lane_capacity;
980 let id = StreamId::make(
981 self.table.kind(),
982 (index / lane_capacity) as u16,
983 (index % lane_capacity) as u32,
984 generation,
985 );
986 Some(Ticket { id, side: Side::B, epoch: self.table.epoch() })
987 }
988
989 #[cfg(unix)]
992 pub fn unlink(&self) -> Result<()> {
993 self.table.unlink()
994 }
995
996 fn locate(
997 &self,
998 id: StreamId,
999 ) -> Result<usize> {
1000 let geometry = self.table.geometry();
1001 if id.kind() != self.table.kind()
1002 || usize::from(id.node()) >= geometry.fleet_capacity
1003 || id.slot() as usize >= geometry.lane_capacity
1004 || id.generation() == 0
1005 || id.generation() > GENERATION_MASK
1006 {
1007 return Err(Error::Malformed(id.to_string()));
1008 }
1009 Ok(usize::from(id.node()) * geometry.lane_capacity + id.slot() as usize)
1010 }
1011}
1012
1013struct Handle {
1015 table: Arc<Table>,
1016 index: usize,
1017 id: StreamId,
1018 side: Side,
1019}
1020
1021impl Handle {
1022 fn peer_if(
1026 &self,
1027 ring: bool,
1028 ) -> table::Ring {
1029 if ring {
1030 table::Ring::Side(1 - self.side.index())
1031 } else {
1032 table::Ring::Nobody
1033 }
1034 }
1035
1036 fn slot(&self) -> Result<&Slot> {
1037 let slot = &self.table.slots()[self.index];
1038 if slot.is(self.id.generation()) {
1039 Ok(slot)
1040 } else {
1041 Err(Error::Stale(self.id))
1042 }
1043 }
1044
1045 fn direction(
1046 slot: &Slot,
1047 direction: usize,
1048 ) -> &Direction {
1049 &slot.directions[direction]
1050 }
1051
1052 fn try_write(
1054 &self,
1055 buf: &[u8],
1056 ) -> Result<usize> {
1057 let slot = self.slot()?;
1058 let direction = Self::direction(slot, self.side.write_direction());
1059 clear_exact_wait(direction);
1060 self.write_bytes(slot, direction, buf)
1061 }
1062
1063 #[cfg(feature = "tokio")]
1066 fn try_write_ready(
1067 &self,
1068 buf: &[u8],
1069 ) -> Result<usize> {
1070 let slot = self.slot()?;
1071 let direction = Self::direction(slot, self.side.write_direction());
1072 self.write_bytes(slot, direction, buf)
1073 }
1074
1075 fn write_bytes(
1076 &self,
1077 slot: &Slot,
1078 direction: &Direction,
1079 buf: &[u8],
1080 ) -> Result<usize> {
1081 let flags = direction.flags();
1082 if flags & FLAG_RESET != 0 {
1083 return Err(Error::Reset);
1084 }
1085 if flags & FLAG_FIN != 0 {
1086 return Err(Error::Closed);
1087 }
1088 if flags & FLAG_READER_GONE != 0 {
1089 return Err(Error::PeerGone);
1090 }
1091 let buffer_bytes = self.table.geometry().buffer_bytes;
1092 let head = direction.head.load(Ordering::Relaxed);
1093 let tail = direction.tail.load(Ordering::Acquire);
1094 let free = buffer_bytes - (head - tail) as usize;
1095 if buf.is_empty() {
1096 return Ok(0);
1097 }
1098 if free == 0 {
1099 return Err(Error::WouldBlock);
1100 }
1101 let len = buf.len().min(free);
1102 let base = self.table.buffer(self.index, self.side.write_direction());
1103 let start = (head as usize) & (buffer_bytes - 1);
1104 let first = len.min(buffer_bytes - start);
1105 unsafe {
1108 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
1109 std::ptr::copy_nonoverlapping(
1110 buf.as_ptr().add(first),
1111 base,
1112 len - first,
1113 );
1114 }
1115 direction.head.store(head + len as u64, Ordering::SeqCst);
1116 std::sync::atomic::fence(Ordering::SeqCst);
1126 let drained = direction.tail.load(Ordering::SeqCst) >= head;
1127 self.table.notify(self.index, slot, direction, self.peer_if(drained));
1128 Ok(len)
1129 }
1130
1131 fn try_write_exact(
1133 &self,
1134 buf: &[u8],
1135 ) -> Result<()> {
1136 let slot = self.slot()?;
1137 let direction = Self::direction(slot, self.side.write_direction());
1138 clear_exact_wait(direction);
1139 let flags = direction.flags();
1140 if flags & FLAG_RESET != 0 {
1141 return Err(Error::Reset);
1142 }
1143 if flags & FLAG_FIN != 0 {
1144 return Err(Error::Closed);
1145 }
1146 if flags & FLAG_READER_GONE != 0 {
1147 return Err(Error::PeerGone);
1148 }
1149 let buffer_bytes = self.table.geometry().buffer_bytes;
1150 if buf.len() > buffer_bytes {
1151 return Err(Error::Malformed(format!(
1152 "atomic stream write is {} bytes; ring capacity is {buffer_bytes}",
1153 buf.len()
1154 )));
1155 }
1156 if buf.is_empty() {
1157 return Ok(());
1158 }
1159 let head = direction.head.load(Ordering::Relaxed);
1160 let tail = direction.tail.load(Ordering::Acquire);
1161 if buffer_bytes - ((head - tail) as usize) < buf.len() {
1162 return Err(Error::WouldBlock);
1163 }
1164 let base = self.table.buffer(self.index, self.side.write_direction());
1165 let start = (head as usize) & (buffer_bytes - 1);
1166 let first = buf.len().min(buffer_bytes - start);
1167 unsafe {
1169 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
1170 std::ptr::copy_nonoverlapping(
1171 buf.as_ptr().add(first),
1172 base,
1173 buf.len() - first,
1174 );
1175 }
1176 direction.head.store(head + buf.len() as u64, Ordering::SeqCst);
1177 std::sync::atomic::fence(Ordering::SeqCst);
1178 let drained = direction.tail.load(Ordering::SeqCst) >= head;
1179 self.table.notify(self.index, slot, direction, self.peer_if(drained));
1180 Ok(())
1181 }
1182
1183 fn try_read(
1186 &self,
1187 buf: &mut [u8],
1188 ) -> Result<usize> {
1189 let slot = self.slot()?;
1190 let direction = Self::direction(slot, self.side.read_direction());
1191 let flags = direction.flags();
1192 if flags & FLAG_RESET != 0 {
1193 return Err(Error::Reset);
1194 }
1195 let head = direction.head.load(Ordering::Acquire);
1196 let tail = direction.tail.load(Ordering::Relaxed);
1197 let available = (head - tail) as usize;
1198 if buf.is_empty() {
1199 return Ok(0);
1200 }
1201 if available == 0 {
1202 if flags & FLAG_FIN != 0 {
1203 return Ok(0);
1204 }
1205 return Err(Error::WouldBlock);
1206 }
1207 let len = buf.len().min(available);
1208 let base = self.table.buffer(self.index, self.side.read_direction());
1209 let buffer_bytes = self.table.geometry().buffer_bytes;
1210 let start = (tail as usize) & (buffer_bytes - 1);
1211 let first = len.min(buffer_bytes - start);
1212 unsafe {
1216 std::ptr::copy_nonoverlapping(
1217 base.add(start),
1218 buf.as_mut_ptr(),
1219 first,
1220 );
1221 std::ptr::copy_nonoverlapping(
1222 base,
1223 buf.as_mut_ptr().add(first),
1224 len - first,
1225 );
1226 }
1227 direction.tail.store(tail + len as u64, Ordering::SeqCst);
1228 std::sync::atomic::fence(Ordering::SeqCst);
1231 let wake = writer_gained_capacity(
1232 direction,
1233 buffer_bytes,
1234 tail,
1235 tail + len as u64,
1236 );
1237 self.table.notify(self.index, slot, direction, self.peer_if(wake));
1238 Ok(len)
1239 }
1240
1241 fn try_read_exact(
1243 &self,
1244 buf: &mut [u8],
1245 ) -> Result<()> {
1246 let slot = self.slot()?;
1247 let direction = Self::direction(slot, self.side.read_direction());
1248 let flags = direction.flags();
1249 if flags & FLAG_RESET != 0 {
1250 return Err(Error::Reset);
1251 }
1252 let buffer_bytes = self.table.geometry().buffer_bytes;
1253 if buf.len() > buffer_bytes {
1254 return Err(Error::Malformed(format!(
1255 "atomic stream read is {} bytes; ring capacity is {buffer_bytes}",
1256 buf.len()
1257 )));
1258 }
1259 if buf.is_empty() {
1260 return Ok(());
1261 }
1262 let head = direction.head.load(Ordering::Acquire);
1263 let tail = direction.tail.load(Ordering::Relaxed);
1264 let available = (head - tail) as usize;
1265 if available < buf.len() {
1266 if flags & FLAG_FIN != 0 {
1267 return Err(Error::Malformed(format!(
1268 "stream ended with {available} bytes of a {} byte frame",
1269 buf.len()
1270 )));
1271 }
1272 return Err(Error::WouldBlock);
1273 }
1274 let base = self.table.buffer(self.index, self.side.read_direction());
1275 let start = (tail as usize) & (buffer_bytes - 1);
1276 let first = buf.len().min(buffer_bytes - start);
1277 unsafe {
1279 std::ptr::copy_nonoverlapping(
1280 base.add(start),
1281 buf.as_mut_ptr(),
1282 first,
1283 );
1284 std::ptr::copy_nonoverlapping(
1285 base,
1286 buf.as_mut_ptr().add(first),
1287 buf.len() - first,
1288 );
1289 }
1290 direction.tail.store(tail + buf.len() as u64, Ordering::SeqCst);
1291 std::sync::atomic::fence(Ordering::SeqCst);
1292 let wake = writer_gained_capacity(
1293 direction,
1294 buffer_bytes,
1295 tail,
1296 tail + buf.len() as u64,
1297 );
1298 self.table.notify(self.index, slot, direction, self.peer_if(wake));
1299 Ok(())
1300 }
1301
1302 fn set_flag(
1303 &self,
1304 direction_index: usize,
1305 flag: u8,
1306 ) {
1307 if let Ok(slot) = self.slot() {
1308 let direction = Self::direction(slot, direction_index);
1309 let before = direction.flags.fetch_or(flag, Ordering::SeqCst);
1310 if before & flag == 0 {
1311 let ring = flag != FLAG_READER_GONE
1312 || before & FLAG_CLOSE_WATCH != 0
1313 || writer_may_park(direction, before);
1314 self.table.notify(
1315 self.index,
1316 slot,
1317 direction,
1318 self.peer_if(ring),
1319 );
1320 }
1321 }
1322 }
1323
1324 fn set_flag_quietly(
1326 &self,
1327 direction_index: usize,
1328 flag: u8,
1329 ) {
1330 if let Ok(slot) = self.slot() {
1331 let direction = Self::direction(slot, direction_index);
1332 if direction.flags.fetch_or(flag, Ordering::SeqCst) & flag == 0 {
1333 self.table.notify(
1334 self.index,
1335 slot,
1336 direction,
1337 table::Ring::Nobody,
1338 );
1339 }
1340 }
1341 }
1342
1343 #[cfg(any(
1346 target_os = "linux",
1347 target_os = "freebsd",
1348 target_os = "macos"
1349 ))]
1350 fn wait_until_timeout(
1351 &self,
1352 direction_index: usize,
1353 ready: impl Fn(&Direction) -> bool,
1354 timeout: Duration,
1355 ) -> Result<bool> {
1356 let deadline = Instant::now() + timeout;
1357 loop {
1358 let slot = self.slot()?;
1359 let direction = Self::direction(slot, direction_index);
1360 if ready(direction) {
1361 return Ok(true);
1362 }
1363 let left = deadline.saturating_duration_since(Instant::now());
1364 if left.is_zero() {
1365 return Ok(false);
1366 }
1367 direction.waiters.fetch_add(1, Ordering::SeqCst);
1368 let since = direction.changes.load(Ordering::SeqCst);
1369 let outcome = if ready(direction) {
1370 Ok(true)
1371 } else {
1372 wait_on_timeout(&direction.changes, since, left)
1373 };
1374 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1375 if !outcome? {
1376 let slot = self.slot()?;
1379 return Ok(ready(Self::direction(slot, direction_index)));
1380 }
1381 }
1382 }
1383
1384 fn wait_until(
1385 &self,
1386 direction_index: usize,
1387 ready: impl Fn(&Direction) -> bool,
1388 ) -> Result<()> {
1389 loop {
1390 let slot = self.slot()?;
1391 let direction = Self::direction(slot, direction_index);
1392 if ready(direction) {
1393 return Ok(());
1394 }
1395 direction.waiters.fetch_add(1, Ordering::SeqCst);
1396 let since = direction.changes.load(Ordering::SeqCst);
1397 let outcome = if ready(direction) {
1398 Ok(())
1399 } else {
1400 wait_on(&direction.changes, since)
1401 };
1402 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1403 outcome?;
1404 }
1405 }
1406
1407 fn poll_ready(
1410 &self,
1411 interest: Interest,
1412 ready: impl Fn(&Direction) -> bool,
1413 cx: &mut std::task::Context<'_>,
1414 ) -> std::task::Poll<Result<()>> {
1415 let slot = match self.slot() {
1416 Ok(slot) => slot,
1417 Err(error) => return std::task::Poll::Ready(Err(error)),
1418 };
1419 let direction = Self::direction(slot, interest.direction);
1420 if ready(direction) {
1421 return std::task::Poll::Ready(Ok(()));
1422 }
1423 if let Err(error) =
1424 self.table.register(self.index, interest, cx.waker())
1425 {
1426 return std::task::Poll::Ready(Err(error));
1427 }
1428 std::sync::atomic::fence(Ordering::SeqCst);
1431 if ready(direction) {
1432 std::task::Poll::Ready(Ok(()))
1433 } else {
1434 std::task::Poll::Pending
1435 }
1436 }
1437}
1438
1439fn clear_exact_wait(direction: &Direction) {
1440 if direction.exact_needed.load(Ordering::Relaxed) != 0 {
1443 direction.exact_needed.store(0, Ordering::SeqCst);
1444 }
1445}
1446
1447fn writer_gained_capacity(
1448 direction: &Direction,
1449 capacity: usize,
1450 before_tail: u64,
1451 after_tail: u64,
1452) -> bool {
1453 let head = direction.head.load(Ordering::SeqCst);
1454 let needed = direction.exact_needed.load(Ordering::SeqCst) as usize;
1455 if needed == 0 {
1456 return head - before_tail >= capacity as u64;
1457 }
1458 if direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0 {
1459 return false;
1460 }
1461 let before = capacity.saturating_sub((head - before_tail) as usize);
1462 let after = capacity.saturating_sub((head - after_tail) as usize);
1463 before < needed && after >= needed
1464}
1465
1466fn writer_may_park(
1480 direction: &Direction,
1481 before: u8,
1482) -> bool {
1483 if before & (FLAG_FIN | FLAG_RESET) != 0 {
1484 return false;
1485 }
1486 std::sync::atomic::fence(Ordering::SeqCst);
1487 direction.head.load(Ordering::SeqCst)
1488 != direction.tail.load(Ordering::SeqCst)
1489}
1490
1491fn readable(direction: &Direction) -> bool {
1492 direction.head.load(Ordering::Acquire)
1493 != direction.tail.load(Ordering::Relaxed)
1494 || direction.flags() & (FLAG_FIN | FLAG_RESET) != 0
1495}
1496
1497fn writable(
1498 direction: &Direction,
1499 buffer_bytes: usize,
1500) -> bool {
1501 let queued = (direction.head.load(Ordering::Relaxed)
1502 - direction.tail.load(Ordering::Acquire)) as usize;
1503 queued < buffer_bytes
1504 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1505}
1506
1507fn writable_exact(
1508 direction: &Direction,
1509 buffer_bytes: usize,
1510 needed: usize,
1511) -> bool {
1512 let queued = (direction.head.load(Ordering::Relaxed)
1513 - direction.tail.load(Ordering::Acquire)) as usize;
1514 buffer_bytes.saturating_sub(queued) >= needed
1515 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1516}
1517
1518impl Drop for Handle {
1519 fn drop(&mut self) {
1523 let slot = &self.table.slots()[self.index];
1524 if !slot.is(self.id.generation()) {
1525 return;
1526 }
1527 if slot.claimed[self.side.index()]
1530 .compare_exchange(
1531 SIDE_CLAIMED,
1532 SIDE_RELEASED,
1533 Ordering::SeqCst,
1534 Ordering::SeqCst,
1535 )
1536 .is_err()
1537 {
1538 return;
1539 }
1540 let other = slot.claimed[1 - self.side.index()].load(Ordering::SeqCst);
1544 if other != SIDE_CLAIMED {
1545 let _ = slot.state.compare_exchange(
1546 SLOT_LIVE,
1547 SLOT_EMPTY,
1548 Ordering::SeqCst,
1549 Ordering::SeqCst,
1550 );
1551 for direction in &slot.directions {
1552 direction.changes.fetch_add(1, Ordering::SeqCst);
1555 if direction.waiters.load(Ordering::SeqCst) > 0 {
1556 wake_on(&direction.changes);
1557 }
1558 }
1559 self.table.registry.wake(self.index);
1560 }
1561 if self.side == Side::B || other != SIDE_CLAIMED {
1564 self.table.notify_capacity(self.index);
1565 }
1566 }
1567}
1568
1569pub struct Endpoint {
1573 read: ReadHalf,
1574 write: WriteHalf,
1575}
1576
1577impl Endpoint {
1578 fn new(handle: Arc<Handle>) -> Self {
1579 Self {
1580 read: ReadHalf { handle: Arc::clone(&handle) },
1581 write: WriteHalf { handle, shut: false },
1582 }
1583 }
1584
1585 pub fn id(&self) -> StreamId {
1586 self.read.handle.id
1587 }
1588
1589 pub fn side(&self) -> Side {
1590 self.read.handle.side
1591 }
1592
1593 pub fn split(self) -> (ReadHalf, WriteHalf) {
1594 (self.read, self.write)
1595 }
1596
1597 pub fn rejoin(
1608 read: ReadHalf,
1609 write: WriteHalf,
1610 ) -> Result<Self> {
1611 if read.id() != write.id() {
1612 return Err(Error::Stale(read.id()));
1613 }
1614 Ok(Self { read, write })
1615 }
1616
1617 pub fn rearm(
1630 &mut self,
1631 streams: &Streams,
1632 ) -> Result<()> {
1633 if self.side() != Side::A {
1634 return Err(Error::Malformed(
1635 "only the creator side of a stream can be rearmed".to_owned(),
1636 ));
1637 }
1638 streams.rearm(self.id())?;
1639 self.write.shut = false;
1640 Ok(())
1641 }
1642
1643 pub fn try_read(
1644 &self,
1645 buf: &mut [u8],
1646 ) -> Result<usize> {
1647 self.read.try_read(buf)
1648 }
1649
1650 pub fn try_read_exact(
1651 &self,
1652 buf: &mut [u8],
1653 ) -> Result<()> {
1654 self.read.try_read_exact(buf)
1655 }
1656
1657 pub fn blocking_read(
1658 &self,
1659 buf: &mut [u8],
1660 ) -> Result<usize> {
1661 self.read.blocking_read(buf)
1662 }
1663
1664 pub fn blocking_read_chunk(
1665 &self,
1666 max: usize,
1667 ) -> Result<Bytes> {
1668 self.read.blocking_read_chunk(max)
1669 }
1670
1671 pub fn try_write(
1672 &self,
1673 buf: &[u8],
1674 ) -> Result<usize> {
1675 self.write.try_write(buf)
1676 }
1677
1678 pub fn try_write_exact(
1679 &self,
1680 buf: &[u8],
1681 ) -> Result<()> {
1682 self.write.try_write_exact(buf)
1683 }
1684
1685 pub fn blocking_write(
1686 &self,
1687 buf: &[u8],
1688 ) -> Result<usize> {
1689 self.write.blocking_write(buf)
1690 }
1691
1692 pub fn blocking_write_all(
1693 &self,
1694 buf: &[u8],
1695 ) -> Result<()> {
1696 self.write.blocking_write_all(buf)
1697 }
1698
1699 pub fn finish(&mut self) -> Result<()> {
1700 self.write.finish()
1701 }
1702
1703 pub fn reset(&mut self) {
1704 self.write.reset()
1705 }
1706}
1707
1708pub struct ReadHalf {
1711 handle: Arc<Handle>,
1712}
1713
1714pub struct ResetSignal {
1718 handle: Arc<Handle>,
1719}
1720
1721impl ResetSignal {
1722 pub fn poll_reset(
1725 &self,
1726 cx: &mut std::task::Context<'_>,
1727 ) -> std::task::Poll<Result<()>> {
1728 self.handle.poll_ready(
1729 Interest {
1730 direction: self.handle.side.read_direction(),
1731 writer: false,
1732 reset: true,
1733 },
1734 |direction| direction.flags() & FLAG_RESET != 0,
1735 cx,
1736 )
1737 }
1738}
1739
1740impl ReadHalf {
1741 pub fn reset_signal(&self) -> ResetSignal {
1744 ResetSignal { handle: Arc::clone(&self.handle) }
1745 }
1746
1747 pub fn id(&self) -> StreamId {
1748 self.handle.id
1749 }
1750
1751 pub fn try_read(
1754 &self,
1755 buf: &mut [u8],
1756 ) -> Result<usize> {
1757 self.handle.try_read(buf)
1758 }
1759
1760 pub fn try_read_exact(
1763 &self,
1764 buf: &mut [u8],
1765 ) -> Result<()> {
1766 self.handle.try_read_exact(buf)
1767 }
1768
1769 pub fn blocking_read(
1772 &self,
1773 buf: &mut [u8],
1774 ) -> Result<usize> {
1775 loop {
1776 match self.handle.try_read(buf) {
1777 Err(Error::WouldBlock) => self.wait_readable()?,
1778 other => return other,
1779 }
1780 }
1781 }
1782
1783 pub fn stop_reading(&self) {
1791 self.handle
1792 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1793 }
1794
1795 pub fn blocking_read_chunk(
1797 &self,
1798 max: usize,
1799 ) -> Result<Bytes> {
1800 let mut chunk = BytesMut::zeroed(max);
1801 let len = self.blocking_read(&mut chunk)?;
1802 chunk.truncate(len);
1803 Ok(chunk.freeze())
1804 }
1805
1806 #[cfg(any(
1808 target_os = "linux",
1809 target_os = "freebsd",
1810 target_os = "macos"
1811 ))]
1812 pub fn wait_readable_timeout(
1813 &self,
1814 timeout: Duration,
1815 ) -> Result<bool> {
1816 self.handle.wait_until_timeout(
1817 self.handle.side.read_direction(),
1818 readable,
1819 timeout,
1820 )
1821 }
1822
1823 pub fn wait_readable(&self) -> Result<()> {
1825 self.handle.wait_until(self.handle.side.read_direction(), readable)
1826 }
1827
1828 pub fn poll_readable(
1833 &self,
1834 cx: &mut std::task::Context<'_>,
1835 ) -> std::task::Poll<Result<()>> {
1836 self.handle.poll_ready(
1837 Interest {
1838 direction: self.handle.side.read_direction(),
1839 writer: false,
1840 reset: false,
1841 },
1842 readable,
1843 cx,
1844 )
1845 }
1846}
1847
1848impl Drop for ReadHalf {
1849 fn drop(&mut self) {
1850 self.handle
1851 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1852 }
1853}
1854
1855pub struct WriteHalf {
1858 handle: Arc<Handle>,
1859 shut: bool,
1860}
1861
1862struct CloseWatch<'a>(&'a Handle);
1864
1865impl Drop for CloseWatch<'_> {
1866 fn drop(&mut self) {
1867 if let Ok(slot) = self.0.slot() {
1868 Handle::direction(slot, self.0.side.write_direction())
1869 .flags
1870 .fetch_and(!FLAG_CLOSE_WATCH, Ordering::SeqCst);
1871 }
1872 }
1873}
1874
1875impl WriteHalf {
1876 pub async fn peer_closed(&mut self) -> Result<()> {
1882 let slot = self.handle.slot()?;
1883 let direction =
1884 Handle::direction(slot, self.handle.side.write_direction());
1885 direction.flags.fetch_or(FLAG_CLOSE_WATCH, Ordering::SeqCst);
1886 let _watch = CloseWatch(&self.handle);
1887 std::future::poll_fn(|cx| {
1888 self.handle.poll_ready(
1889 Interest {
1890 direction: self.handle.side.write_direction(),
1891 writer: true,
1892 reset: false,
1893 },
1894 |direction| {
1895 direction.flags() & (FLAG_RESET | FLAG_READER_GONE) != 0
1896 },
1897 cx,
1898 )
1899 })
1900 .await
1901 }
1902
1903 pub fn poll_peer_dead(
1909 &self,
1910 cx: &mut std::task::Context<'_>,
1911 ) -> std::task::Poll<Result<()>> {
1912 let slot = match self.handle.slot() {
1913 Ok(slot) => slot,
1914 Err(error) => return std::task::Poll::Ready(Err(error)),
1915 };
1916 let peer = &slot.claimed[1 - self.handle.side.index()];
1917 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1918 return std::task::Poll::Ready(Ok(()));
1919 }
1920 let interest = Interest {
1921 direction: self.handle.side.write_direction(),
1922 writer: true,
1923 reset: false,
1924 };
1925 if let Err(error) =
1926 self.handle.table.register(self.handle.index, interest, cx.waker())
1927 {
1928 return std::task::Poll::Ready(Err(error));
1929 }
1930 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1933 std::task::Poll::Ready(Ok(()))
1934 } else {
1935 std::task::Poll::Pending
1936 }
1937 }
1938
1939 pub fn id(&self) -> StreamId {
1940 self.handle.id
1941 }
1942
1943 pub fn try_write(
1945 &self,
1946 buf: &[u8],
1947 ) -> Result<usize> {
1948 self.handle.try_write(buf)
1949 }
1950
1951 pub fn try_write_exact(
1954 &self,
1955 buf: &[u8],
1956 ) -> Result<()> {
1957 self.handle.try_write_exact(buf)
1958 }
1959
1960 pub fn blocking_write(
1962 &self,
1963 buf: &[u8],
1964 ) -> Result<usize> {
1965 loop {
1966 match self.handle.try_write(buf) {
1967 Err(Error::WouldBlock) => self.wait_writable()?,
1968 other => return other,
1969 }
1970 }
1971 }
1972
1973 pub fn blocking_write_all(
1976 &self,
1977 mut buf: &[u8],
1978 ) -> Result<()> {
1979 while !buf.is_empty() {
1980 let written = self.blocking_write(buf)?;
1981 buf = &buf[written..];
1982 }
1983 Ok(())
1984 }
1985
1986 pub fn finish(&mut self) -> Result<()> {
1990 self.shut = true;
1991 self.handle.set_flag(self.handle.side.write_direction(), FLAG_FIN);
1992 Ok(())
1993 }
1994
1995 pub fn finish_quietly(&mut self) {
2002 self.shut = true;
2003 self.handle
2004 .set_flag_quietly(self.handle.side.write_direction(), FLAG_FIN);
2005 }
2006
2007 pub fn reset(&mut self) {
2010 self.shut = true;
2011 self.handle.set_flag(self.handle.side.write_direction(), FLAG_RESET);
2012 }
2013
2014 #[cfg(any(
2018 target_os = "linux",
2019 target_os = "freebsd",
2020 target_os = "macos"
2021 ))]
2022 pub fn wait_writable_timeout(
2023 &self,
2024 timeout: Duration,
2025 ) -> Result<bool> {
2026 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2027 self.handle.wait_until_timeout(
2028 self.handle.side.write_direction(),
2029 |direction| {
2030 clear_exact_wait(direction);
2031 writable(direction, buffer_bytes)
2032 },
2033 timeout,
2034 )
2035 }
2036
2037 pub fn wait_writable(&self) -> Result<()> {
2039 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2040 self.handle.wait_until(
2041 self.handle.side.write_direction(),
2042 |direction| {
2043 clear_exact_wait(direction);
2044 writable(direction, buffer_bytes)
2045 },
2046 )
2047 }
2048
2049 pub fn wait_writable_exact(
2051 &self,
2052 needed: usize,
2053 ) -> Result<()> {
2054 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2055 if needed > buffer_bytes {
2056 return Err(Error::Malformed(format!(
2057 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
2058 )));
2059 }
2060 self.handle.wait_until(
2061 self.handle.side.write_direction(),
2062 |direction| {
2063 clear_exact_wait(direction);
2064 writable_exact(direction, buffer_bytes, needed)
2065 },
2066 )
2067 }
2068
2069 pub fn poll_writable(
2074 &self,
2075 cx: &mut std::task::Context<'_>,
2076 ) -> std::task::Poll<Result<()>> {
2077 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2078 self.handle.poll_ready(
2079 Interest {
2080 direction: self.handle.side.write_direction(),
2081 writer: true,
2082 reset: false,
2083 },
2084 |direction| {
2085 clear_exact_wait(direction);
2086 writable(direction, buffer_bytes)
2087 },
2088 cx,
2089 )
2090 }
2091
2092 pub fn poll_writable_exact(
2094 &self,
2095 needed: usize,
2096 cx: &mut std::task::Context<'_>,
2097 ) -> std::task::Poll<Result<()>> {
2098 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2099 let slot = match self.handle.slot() {
2100 Ok(slot) => slot,
2101 Err(error) => return std::task::Poll::Ready(Err(error)),
2102 };
2103 let direction =
2104 Handle::direction(slot, self.handle.side.write_direction());
2105 if needed > buffer_bytes {
2106 clear_exact_wait(direction);
2107 return std::task::Poll::Ready(Err(Error::Malformed(format!(
2108 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
2109 ))));
2110 }
2111 direction.exact_needed.store(needed as u32, Ordering::SeqCst);
2115 let result = self.handle.poll_ready(
2116 Interest {
2117 direction: self.handle.side.write_direction(),
2118 writer: true,
2119 reset: false,
2120 },
2121 |direction| writable_exact(direction, buffer_bytes, needed),
2122 cx,
2123 );
2124 if result.is_ready() {
2125 clear_exact_wait(direction);
2126 }
2127 result
2128 }
2129}
2130
2131impl Drop for WriteHalf {
2132 fn drop(&mut self) {
2133 if !self.shut {
2134 self.handle
2135 .set_flag(self.handle.side.write_direction(), FLAG_RESET);
2136 }
2137 }
2138}
2139
2140#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2144pub(crate) fn waits_supported() -> bool {
2145 orbit_core::sync::supported()
2146}
2147
2148#[cfg(not(any(
2149 target_os = "linux",
2150 target_os = "freebsd",
2151 target_os = "macos"
2152)))]
2153pub(crate) fn waits_supported() -> bool {
2154 false
2155}
2156
2157#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2161#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2164pub(crate) fn wait_on_timeout(
2165 word: &AtomicU32,
2166 expected: u32,
2167 timeout: Duration,
2168) -> Result<bool> {
2169 orbit_core::sync::wait_word_timeout(word, expected, timeout)
2170 .map_err(Error::Io)
2171}
2172
2173pub(crate) fn wait_on(
2174 word: &AtomicU32,
2175 expected: u32,
2176) -> Result<()> {
2177 orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
2178}
2179
2180#[cfg(not(any(
2181 target_os = "linux",
2182 target_os = "freebsd",
2183 target_os = "macos"
2184)))]
2185pub(crate) fn wait_on(
2186 _word: &AtomicU32,
2187 _expected: u32,
2188) -> Result<()> {
2189 Err(Error::Io(std::io::Error::new(
2190 std::io::ErrorKind::Unsupported,
2191 "orbit-stream needs a platform that can wait on a shared word",
2192 )))
2193}
2194
2195pub(crate) fn wake_on(word: &AtomicU32) {
2197 #[cfg(any(
2198 target_os = "linux",
2199 target_os = "freebsd",
2200 target_os = "macos"
2201 ))]
2202 let _ = orbit_core::sync::wake_word(word);
2203 #[cfg(not(any(
2204 target_os = "linux",
2205 target_os = "freebsd",
2206 target_os = "macos"
2207 )))]
2208 let _ = word;
2209}
2210
2211pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
2212 mutex.lock().unwrap_or_else(|error| error.into_inner())
2213}
2214
2215#[cfg(test)]
2216mod tests {
2217 use std::sync::Arc;
2218 use std::time::Duration;
2219
2220 use orbit_core::{Fleet, NodeId};
2221
2222 use super::{
2223 Error, Incarnation, STREAM_BUFFER_BYTES, STREAM_LANE_CAPACITY, Side,
2224 Streams, Ticket,
2225 };
2226
2227 fn streams(name: &'static str) -> Streams {
2228 Streams::new(
2229 Arc::new(Fleet::join(name, 2).unwrap()),
2230 Incarnation::new(1),
2231 )
2232 .unwrap()
2233 }
2234
2235 #[test]
2236 fn bytes_cross_in_order_and_end_cleanly() {
2237 let streams = streams("stream-order");
2238 let (mut a, ticket) = streams.create().unwrap();
2239 assert_eq!(ticket.side, Side::B);
2240 assert_eq!(ticket.id.node(), 0);
2241 assert_eq!(ticket.epoch, streams.epoch());
2242 let b = streams.open(ticket).unwrap();
2243
2244 a.blocking_write_all(b"hello, ").unwrap();
2245 a.blocking_write_all(b"world").unwrap();
2246 a.finish().unwrap();
2247
2248 let mut out = Vec::new();
2249 let mut buf = [0_u8; 4];
2250 loop {
2251 let n = b.blocking_read(&mut buf).unwrap();
2252 if n == 0 {
2253 break;
2254 }
2255 out.extend_from_slice(&buf[..n]);
2256 }
2257 assert_eq!(out, b"hello, world");
2258 b.blocking_write_all(b"back").unwrap();
2260 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"back");
2261 }
2262
2263 #[test]
2264 fn a_ticket_prints_and_parses() {
2265 let streams = streams("stream-ticket");
2266 let (_a, ticket) = streams.create().unwrap();
2267 let text = ticket.to_string();
2268 assert!(text.contains("/b/"), "{text}");
2269 assert_eq!(text.parse::<Ticket>().unwrap(), ticket);
2270 assert!(matches!("nope".parse::<Ticket>(), Err(Error::Malformed(_))));
2271 assert!(matches!("x/b".parse::<Ticket>(), Err(Error::Malformed(_))));
2272 }
2273
2274 #[test]
2275 fn a_side_is_held_once_and_a_released_stream_goes_stale() {
2276 let streams = streams("stream-claim");
2277 let (a, ticket) = streams.create().unwrap();
2278 let b = streams.open(ticket).unwrap();
2279 assert!(matches!(streams.open(ticket), Err(Error::AlreadyClaimed(_))));
2280 assert!(streams.is_live(ticket.id));
2281
2282 drop(a);
2283 let mut buf = [0_u8; 1];
2285 assert!(matches!(b.blocking_read(&mut buf), Err(Error::Reset)));
2286 drop(b);
2287 assert!(!streams.is_live(ticket.id));
2288 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2289
2290 for _ in 0..STREAM_LANE_CAPACITY - 1 {
2293 drop(streams.create().unwrap());
2294 }
2295 let (_a2, ticket2) = streams.create().unwrap();
2296 assert_eq!(ticket2.id.slot(), ticket.id.slot());
2297 assert_ne!(ticket2.id.generation(), ticket.id.generation());
2298 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2299 }
2300
2301 #[test]
2302 fn a_reset_starts_a_new_epoch_and_refuses_old_tickets() {
2303 let streams = streams("stream-epoch");
2304 let (_a, ticket) = streams.create().unwrap();
2305 let before = streams.epoch();
2306 streams.reset_all();
2307 assert!(streams.epoch() > before);
2308 let (_a2, ticket2) = streams.create().unwrap();
2311 assert_eq!(ticket2.id, ticket.id);
2312 assert_ne!(ticket2.epoch, ticket.epoch);
2313 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2314 assert!(streams.open(ticket2).is_ok());
2315 }
2316
2317 #[test]
2318 fn a_slot_whose_generation_runs_out_is_not_reused() {
2319 let streams = streams("stream-exhaust");
2322 let limit = super::layout::GENERATION_LIMIT as usize;
2323 for _ in 0..STREAM_LANE_CAPACITY * limit {
2324 drop(streams.create().unwrap());
2325 }
2326 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2327 streams.reset_all();
2328 let (_a, ticket) = streams.create().unwrap();
2329 assert_eq!(ticket.id.generation(), 1);
2330 }
2331
2332 #[test]
2333 fn a_full_ring_blocks_the_writer_and_wraps_correctly() {
2334 let streams = streams("stream-wrap");
2335 let (a, ticket) = streams.create().unwrap();
2336 let b = streams.open(ticket).unwrap();
2337
2338 let filler = vec![7_u8; STREAM_BUFFER_BYTES];
2339 assert_eq!(a.try_write(&filler).unwrap(), STREAM_BUFFER_BYTES);
2340 assert!(matches!(a.try_write(b"x"), Err(Error::WouldBlock)));
2341
2342 let mut buf = vec![0_u8; STREAM_BUFFER_BYTES - 5];
2343 assert_eq!(b.blocking_read(&mut buf).unwrap(), STREAM_BUFFER_BYTES - 5);
2344 let tail =
2346 (0..STREAM_BUFFER_BYTES).map(|i| i as u8).collect::<Vec<_>>();
2347 assert_eq!(a.try_write(&tail).unwrap(), STREAM_BUFFER_BYTES - 5);
2348
2349 let mut rest = vec![0_u8; 5];
2350 assert_eq!(b.blocking_read(&mut rest).unwrap(), 5);
2351 assert_eq!(rest, vec![7_u8; 5]);
2352 let mut wrapped = vec![0_u8; STREAM_BUFFER_BYTES];
2353 let mut got = 0;
2354 while got < STREAM_BUFFER_BYTES - 5 {
2355 got += b.blocking_read(&mut wrapped[got..]).unwrap();
2356 }
2357 assert_eq!(&wrapped[..got], &tail[..got]);
2358 }
2359
2360 #[test]
2361 fn exact_frames_never_publish_or_consume_a_prefix() {
2362 let streams = streams("stream-exact-frame");
2363 let (a, ticket) = streams.create().unwrap();
2364 let b = streams.open(ticket).unwrap();
2365
2366 let filler = vec![7_u8; STREAM_BUFFER_BYTES - 32];
2367 a.try_write_exact(&filler).unwrap();
2368 assert!(matches!(
2369 a.try_write_exact(&[9_u8; 64]),
2370 Err(Error::WouldBlock)
2371 ));
2372 let mut consumed = vec![0_u8; filler.len()];
2373 b.try_read_exact(&mut consumed).unwrap();
2374 assert_eq!(consumed, filler);
2375 assert!(matches!(b.try_read(&mut [0_u8; 1]), Err(Error::WouldBlock)));
2376
2377 a.try_write_exact(&[3_u8; 32]).unwrap();
2378 assert!(matches!(
2379 b.try_read_exact(&mut [0_u8; 64]),
2380 Err(Error::WouldBlock)
2381 ));
2382 let mut frame = [0_u8; 32];
2383 b.try_read_exact(&mut frame).unwrap();
2384 assert_eq!(frame, [3_u8; 32]);
2385 }
2386
2387 #[test]
2388 fn a_blocked_reader_is_woken_by_a_write_in_another_thread() {
2389 let streams = streams("stream-wake");
2390 let (a, ticket) = streams.create().unwrap();
2391 let b = streams.open(ticket).unwrap();
2392 let reader = std::thread::spawn(move || {
2393 let mut buf = [0_u8; 8];
2394 let n = b.blocking_read(&mut buf).unwrap();
2395 buf[..n].to_vec()
2396 });
2397 std::thread::sleep(Duration::from_millis(30));
2398 a.blocking_write_all(b"ping").unwrap();
2399 assert_eq!(reader.join().unwrap(), b"ping");
2400 }
2401
2402 #[test]
2403 fn a_reader_gone_fails_the_writer_and_a_lane_fills_up() {
2404 let streams = streams("stream-gone");
2405 let (a, ticket) = streams.create().unwrap();
2406 let b = streams.open(ticket).unwrap();
2407 let (b_read, _b_write) = b.split();
2408 drop(b_read);
2409 assert!(matches!(a.blocking_write(b"nobody"), Err(Error::PeerGone)));
2410
2411 let mut held = vec![a];
2412 for _ in 1..STREAM_LANE_CAPACITY {
2413 held.push(streams.create().unwrap().0);
2414 }
2415 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2416 }
2417
2418 #[test]
2419 fn a_death_report_ends_that_incarnations_sides_only() {
2420 let streams = streams("stream-dead");
2421 let (a, ticket) = streams.create().unwrap();
2422 let b = streams.open(ticket).unwrap();
2423 b.blocking_write_all(b"partial").unwrap();
2424
2425 streams.node_dead(NodeId::ZERO, Incarnation::new(2));
2427 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"partial");
2428
2429 let reader = std::thread::spawn(move || {
2435 let mut buf = [0_u8; 8];
2436 a.blocking_read(&mut buf).map(|_| ())
2437 });
2438 std::thread::sleep(Duration::from_millis(30));
2439 streams.node_dead(NodeId::ZERO, Incarnation::new(1));
2440 assert!(matches!(
2441 reader.join().unwrap(),
2442 Err(Error::Reset | Error::Stale(_))
2443 ));
2444 assert!(matches!(
2445 b.try_write(b"x"),
2446 Err(Error::Reset | Error::Stale(_))
2447 ));
2448 assert!(matches!(
2449 streams.open(ticket),
2450 Err(Error::AlreadyClaimed(_) | Error::Stale(_))
2451 ));
2452 drop(b);
2453 assert!(!streams.is_live(ticket.id));
2454 }
2455
2456 #[test]
2457 fn an_offer_is_taken_once_and_a_stale_offer_is_skipped() {
2458 let streams = streams("stream-offer");
2459 assert!(streams.take_offer().is_none());
2460 let (a, ticket) = streams.create().unwrap();
2461 streams.offer(ticket, NodeId::ZERO).unwrap();
2462 assert_eq!(streams.take_offer(), Some(ticket));
2463 assert!(streams.take_offer().is_none());
2464
2465 let (gone, stale) = streams.create().unwrap();
2466 streams.offer(stale, NodeId::ZERO).unwrap();
2467 drop(gone);
2468 assert!(streams.take_offer().is_none());
2469
2470 let taker = streams.clone();
2471 let waiter = std::thread::spawn(move || taker.blocking_take_offer());
2472 std::thread::sleep(Duration::from_millis(30));
2473 streams.offer(ticket, NodeId::ZERO).unwrap();
2474 assert_eq!(waiter.join().unwrap().unwrap(), ticket);
2475 drop(a);
2476 }
2477
2478 #[test]
2479 fn a_contested_claim_goes_to_exactly_one_opener() {
2480 let streams = streams("stream-race");
2481 for _ in 0..50 {
2482 let (_a, ticket) = streams.create().unwrap();
2483 let openers = (0..4)
2484 .map(|_| {
2485 let streams = streams.clone();
2486 std::thread::spawn(move || streams.open(ticket).is_ok())
2487 })
2488 .collect::<Vec<_>>();
2489 let won = openers
2490 .into_iter()
2491 .map(|opener| opener.join().unwrap())
2492 .filter(|won| *won)
2493 .count();
2494 assert_eq!(won, 1);
2495 }
2496 }
2497
2498 #[test]
2499 fn random_sizes_concatenate_exactly() {
2500 let streams = streams("stream-random");
2501 let (mut a, ticket) = streams.create().unwrap();
2502 let b = streams.open(ticket).unwrap();
2503 let payload = (0..300_000_u32)
2504 .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
2505 .collect::<Vec<_>>();
2506 let expected = payload.clone();
2507 let writer = std::thread::spawn(move || {
2508 let mut rest = &payload[..];
2509 let mut step = 1_usize;
2510 while !rest.is_empty() {
2511 let len = step.min(rest.len());
2512 a.blocking_write_all(&rest[..len]).unwrap();
2513 rest = &rest[len..];
2514 step = (step * 7 + 13) % 9_001 + 1;
2515 }
2516 a.finish().unwrap();
2517 });
2518 let mut out = Vec::new();
2519 let mut step = 3_usize;
2520 loop {
2521 let chunk = b.blocking_read_chunk(step).unwrap();
2522 if chunk.is_empty() {
2523 break;
2524 }
2525 out.extend_from_slice(&chunk);
2526 step = (step * 5 + 11) % 4_099 + 1;
2527 }
2528 writer.join().unwrap();
2529 assert_eq!(out, expected);
2530 }
2531}