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)]
473 pub fn prepare_group(
474 name: &str,
475 capacity: u16,
476 spec: StreamSpec,
477 access: orbit_core::shm::GroupAccess,
478 ) -> Result<()> {
479 spec.validate()?;
480 table::open_group(
481 name,
482 capacity,
483 NodeId::ZERO,
484 Incarnation::new(0),
485 spec,
486 access,
487 true,
488 )?;
489 Ok(())
490 }
491
492 #[cfg(unix)]
494 pub fn node_dead_in_group(
495 name: &str,
496 capacity: u16,
497 spec: StreamSpec,
498 node: NodeId,
499 incarnation: Incarnation,
500 access: orbit_core::shm::GroupAccess,
501 ) -> Result<()> {
502 spec.validate()?;
503 if node.get() >= capacity {
504 return Err(Error::Malformed(
505 "recovery node is outside the fleet".into(),
506 ));
507 }
508 let table = table::open_group(
509 name,
510 capacity,
511 NodeId::ZERO,
512 Incarnation::new(0),
513 spec,
514 access,
515 true,
516 )?;
517 table.node_dead(node.get(), incarnation.get());
518 Ok(())
519 }
520
521 #[cfg(unix)]
525 pub fn node_dead_in_shm(
526 name: &str,
527 fleet_capacity: u16,
528 spec: StreamSpec,
529 node: NodeId,
530 incarnation: Incarnation,
531 ) -> Result<bool> {
532 spec.validate()?;
533 if node.get() >= fleet_capacity {
534 return Err(Error::Malformed(
535 "recovery node is outside the fleet".into(),
536 ));
537 }
538 let segment = orbit_core::shm::ring_segment_name(name, spec.kind);
539 match orbit_core::shm::ShmRegion::validate_existing(
540 &segment,
541 segment_size_for(fleet_capacity, spec),
542 )? {
543 orbit_core::shm::ShmValidation::Missing => Ok(false),
544 orbit_core::shm::ShmValidation::Valid { .. } => {
545 table::recover_shm(
546 name,
547 fleet_capacity,
548 spec,
549 node.get(),
550 incarnation.get(),
551 )?;
552 Ok(true)
553 }
554 }
555 }
556
557 pub fn new(
561 fleet: Arc<Fleet>,
562 incarnation: Incarnation,
563 ) -> Result<Self> {
564 Self::with_spec(fleet, incarnation, StreamSpec::DEFAULT)
565 }
566
567 pub fn with_spec(
572 fleet: Arc<Fleet>,
573 incarnation: Incarnation,
574 spec: StreamSpec,
575 ) -> Result<Self> {
576 spec.validate()?;
577 Ok(Self { table: table::open(&fleet, incarnation, spec)? })
578 }
579
580 pub fn kind(&self) -> u8 {
582 self.table.kind()
583 }
584
585 #[cfg(any(
597 target_os = "linux",
598 target_os = "freebsd",
599 target_os = "macos"
600 ))]
601 pub fn readiness(&self) -> Result<Readiness> {
602 self.table.take_readiness()
603 }
604
605 pub fn epoch(&self) -> u64 {
607 self.table.epoch()
608 }
609
610 pub fn create(&self) -> Result<(Endpoint, Ticket)> {
613 let (index, generation) = self.table.allocate()?;
614 let id = StreamId::make(
615 self.table.kind(),
616 self.table.node(),
617 (index % self.table.geometry().lane_capacity) as u32,
618 generation,
619 );
620 let handle = Arc::new(Handle {
621 table: Arc::clone(&self.table),
622 index,
623 id,
624 side: Side::A,
625 });
626 Ok((
627 Endpoint::new(handle),
628 Ticket { id, side: Side::B, epoch: self.table.epoch() },
629 ))
630 }
631
632 pub fn capacity_changed(&self) -> CapacityChanged {
636 CapacityChanged::new(Arc::clone(&self.table))
637 }
638
639 pub fn open(
641 &self,
642 ticket: Ticket,
643 ) -> Result<Endpoint> {
644 let index = self.locate(ticket.id)?;
645 if ticket.epoch != self.table.epoch() {
646 return Err(Error::Stale(ticket.id));
647 }
648 let slot = &self.table.slots()[index];
649 if !slot.is(ticket.id.generation()) {
650 return Err(Error::Stale(ticket.id));
651 }
652 let claimed = &slot.claimed[ticket.side.index()];
653 if claimed
654 .compare_exchange(
655 SIDE_FREE,
656 SIDE_CLAIMED,
657 Ordering::SeqCst,
658 Ordering::SeqCst,
659 )
660 .is_err()
661 {
662 return Err(Error::AlreadyClaimed(ticket));
663 }
664 if !slot.is(ticket.id.generation()) {
667 let _ = claimed.compare_exchange(
668 SIDE_CLAIMED,
669 SIDE_FREE,
670 Ordering::SeqCst,
671 Ordering::SeqCst,
672 );
673 return Err(Error::Stale(ticket.id));
674 }
675 slot.incarnation[ticket.side.index()]
676 .store(self.table.incarnation(), Ordering::Release);
677 slot.node[ticket.side.index()]
678 .store(self.table.node(), Ordering::Release);
679 for direction in &slot.directions {
686 self.table.notify(index, slot, direction, table::Ring::Nobody);
687 }
688 Ok(Endpoint::new(Arc::new(Handle {
689 table: Arc::clone(&self.table),
690 index,
691 id: ticket.id,
692 side: ticket.side,
693 })))
694 }
695
696 pub fn rearm(
732 &self,
733 id: StreamId,
734 ) -> Result<()> {
735 let slot = self.rearmable_slot(id)?;
736 let far = Side::B.index();
737 for direction in &slot.directions {
738 direction.clear();
739 }
740 slot.node[far].store(u16::MAX, Ordering::Relaxed);
741 slot.incarnation[far].store(0, Ordering::Relaxed);
742 slot.claimed[far]
746 .compare_exchange(
747 SIDE_RELEASED,
748 SIDE_FREE,
749 Ordering::SeqCst,
750 Ordering::SeqCst,
751 )
752 .map_err(|_| Error::Stale(id))?;
753 Ok(())
754 }
755
756 pub fn can_create(&self) -> bool {
761 self.table.has_capacity()
762 }
763
764 pub fn can_rearm(
770 &self,
771 id: StreamId,
772 ) -> Result<bool> {
773 match self.rearmable_slot(id) {
774 Ok(_) => Ok(true),
775 Err(Error::WouldBlock) => Ok(false),
776 Err(error) => Err(error),
777 }
778 }
779
780 fn rearmable_slot(
781 &self,
782 id: StreamId,
783 ) -> Result<&Slot> {
784 let index = self.locate(id)?;
785 let slot = &self.table.slots()[index];
786 if !slot.is(id.generation()) {
787 return Err(Error::Stale(id));
788 }
789 let far = Side::B.index();
790 match slot.claimed[far].load(Ordering::SeqCst) {
791 SIDE_RELEASED => {}
792 SIDE_CLAIMED => return Err(Error::WouldBlock),
793 _ => return Err(Error::Stale(id)),
794 }
795 Ok(slot)
796 }
797
798 pub fn is_live(
800 &self,
801 id: StreamId,
802 ) -> bool {
803 self.locate(id)
804 .map(|index| self.table.slots()[index].is(id.generation()))
805 .unwrap_or(false)
806 }
807
808 pub fn offer(
814 &self,
815 ticket: Ticket,
816 to: NodeId,
817 ) -> Result<()> {
818 let index = self.locate(ticket.id)?;
819 if ticket.epoch != self.table.epoch()
820 || !self.table.slots()[index].is(ticket.id.generation())
821 {
822 return Err(Error::Stale(ticket.id));
823 }
824 self.table.offer(index, to.get())
825 }
826
827 pub fn take_offer(&self) -> Option<Ticket> {
830 while let Some(index) = self.table.take_offer() {
831 if let Some(ticket) = self.ticket_for(index) {
832 return Some(ticket);
833 }
834 }
835 None
836 }
837
838 #[cfg(any(
842 target_os = "linux",
843 target_os = "freebsd",
844 target_os = "macos"
845 ))]
846 pub fn take_offer_timeout(
847 &self,
848 timeout: Duration,
849 ) -> Result<Option<Ticket>> {
850 let deadline = Instant::now() + timeout;
851 loop {
852 if let Some(ticket) = self.take_offer() {
853 return Ok(Some(ticket));
854 }
855 let left = deadline.saturating_duration_since(Instant::now());
856 if left.is_zero() {
857 return Ok(self.take_offer());
858 }
859 if !self.table.wait_offer_timeout(left)? {
860 return Ok(self.take_offer());
861 }
862 }
863 }
864
865 pub fn blocking_take_offer(&self) -> Result<Ticket> {
867 loop {
868 if let Some(ticket) = self.take_offer() {
869 return Ok(ticket);
870 }
871 self.table.wait_offer()?;
872 }
873 }
874
875 pub fn poll_take_offer(
878 &self,
879 cx: &mut std::task::Context<'_>,
880 ) -> std::task::Poll<Result<Ticket>> {
881 if let Some(ticket) = self.take_offer() {
882 return std::task::Poll::Ready(Ok(ticket));
883 }
884 if let Err(error) = self.table.register_offer(cx.waker()) {
885 return std::task::Poll::Ready(Err(error));
886 }
887 match self.take_offer() {
888 Some(ticket) => std::task::Poll::Ready(Ok(ticket)),
889 None => std::task::Poll::Pending,
890 }
891 }
892
893 pub fn owners(&self) -> Vec<(NodeId, Incarnation)> {
899 let mut seen: Vec<(NodeId, Incarnation)> = Vec::new();
900 for slot in self.table.slots() {
901 if slot.state.load(Ordering::Acquire) != SLOT_LIVE {
902 continue;
903 }
904 for side in 0..2 {
905 if slot.claimed[side].load(Ordering::Acquire) != SIDE_CLAIMED {
906 continue;
907 }
908 let holder = (
909 NodeId::new(slot.node[side].load(Ordering::Acquire)),
910 Incarnation::new(
911 slot.incarnation[side].load(Ordering::Acquire),
912 ),
913 );
914 if !seen.contains(&holder) {
915 seen.push(holder);
916 }
917 }
918 }
919 seen.sort_by_key(|(node, incarnation)| (node.get(), incarnation.get()));
920 seen
921 }
922
923 pub fn node_dead(
930 &self,
931 node: NodeId,
932 incarnation: Incarnation,
933 ) {
934 self.table.node_dead(node.get(), incarnation.get());
935 }
936
937 pub fn supersede(&self) {
949 self.table.supersede();
950 }
951
952 pub fn reset_all(&self) {
955 self.table.reset_all();
956 }
957
958 fn ticket_for(
959 &self,
960 index: usize,
961 ) -> Option<Ticket> {
962 let slot = &self.table.slots()[index];
963 let generation = slot.generation.load(Ordering::Acquire);
964 if !slot.is(generation) {
965 return None;
966 }
967 let lane_capacity = self.table.geometry().lane_capacity;
968 let id = StreamId::make(
969 self.table.kind(),
970 (index / lane_capacity) as u16,
971 (index % lane_capacity) as u32,
972 generation,
973 );
974 Some(Ticket { id, side: Side::B, epoch: self.table.epoch() })
975 }
976
977 #[cfg(unix)]
980 pub fn unlink(&self) -> Result<()> {
981 self.table.unlink()
982 }
983
984 fn locate(
985 &self,
986 id: StreamId,
987 ) -> Result<usize> {
988 let geometry = self.table.geometry();
989 if id.kind() != self.table.kind()
990 || usize::from(id.node()) >= geometry.fleet_capacity
991 || id.slot() as usize >= geometry.lane_capacity
992 || id.generation() == 0
993 || id.generation() > GENERATION_MASK
994 {
995 return Err(Error::Malformed(id.to_string()));
996 }
997 Ok(usize::from(id.node()) * geometry.lane_capacity + id.slot() as usize)
998 }
999}
1000
1001struct Handle {
1003 table: Arc<Table>,
1004 index: usize,
1005 id: StreamId,
1006 side: Side,
1007}
1008
1009impl Handle {
1010 fn peer_if(
1014 &self,
1015 ring: bool,
1016 ) -> table::Ring {
1017 if ring {
1018 table::Ring::Side(1 - self.side.index())
1019 } else {
1020 table::Ring::Nobody
1021 }
1022 }
1023
1024 fn slot(&self) -> Result<&Slot> {
1025 let slot = &self.table.slots()[self.index];
1026 if slot.is(self.id.generation()) {
1027 Ok(slot)
1028 } else {
1029 Err(Error::Stale(self.id))
1030 }
1031 }
1032
1033 fn direction(
1034 slot: &Slot,
1035 direction: usize,
1036 ) -> &Direction {
1037 &slot.directions[direction]
1038 }
1039
1040 fn try_write(
1042 &self,
1043 buf: &[u8],
1044 ) -> Result<usize> {
1045 let slot = self.slot()?;
1046 let direction = Self::direction(slot, self.side.write_direction());
1047 clear_exact_wait(direction);
1048 self.write_bytes(slot, direction, buf)
1049 }
1050
1051 #[cfg(feature = "tokio")]
1054 fn try_write_ready(
1055 &self,
1056 buf: &[u8],
1057 ) -> Result<usize> {
1058 let slot = self.slot()?;
1059 let direction = Self::direction(slot, self.side.write_direction());
1060 self.write_bytes(slot, direction, buf)
1061 }
1062
1063 fn write_bytes(
1064 &self,
1065 slot: &Slot,
1066 direction: &Direction,
1067 buf: &[u8],
1068 ) -> Result<usize> {
1069 let flags = direction.flags();
1070 if flags & FLAG_RESET != 0 {
1071 return Err(Error::Reset);
1072 }
1073 if flags & FLAG_FIN != 0 {
1074 return Err(Error::Closed);
1075 }
1076 if flags & FLAG_READER_GONE != 0 {
1077 return Err(Error::PeerGone);
1078 }
1079 let buffer_bytes = self.table.geometry().buffer_bytes;
1080 let head = direction.head.load(Ordering::Relaxed);
1081 let tail = direction.tail.load(Ordering::Acquire);
1082 let free = buffer_bytes - (head - tail) as usize;
1083 if buf.is_empty() {
1084 return Ok(0);
1085 }
1086 if free == 0 {
1087 return Err(Error::WouldBlock);
1088 }
1089 let len = buf.len().min(free);
1090 let base = self.table.buffer(self.index, self.side.write_direction());
1091 let start = (head as usize) & (buffer_bytes - 1);
1092 let first = len.min(buffer_bytes - start);
1093 unsafe {
1096 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
1097 std::ptr::copy_nonoverlapping(
1098 buf.as_ptr().add(first),
1099 base,
1100 len - first,
1101 );
1102 }
1103 direction.head.store(head + len as u64, Ordering::SeqCst);
1104 std::sync::atomic::fence(Ordering::SeqCst);
1114 let drained = direction.tail.load(Ordering::SeqCst) >= head;
1115 self.table.notify(self.index, slot, direction, self.peer_if(drained));
1116 Ok(len)
1117 }
1118
1119 fn try_write_exact(
1121 &self,
1122 buf: &[u8],
1123 ) -> Result<()> {
1124 let slot = self.slot()?;
1125 let direction = Self::direction(slot, self.side.write_direction());
1126 clear_exact_wait(direction);
1127 let flags = direction.flags();
1128 if flags & FLAG_RESET != 0 {
1129 return Err(Error::Reset);
1130 }
1131 if flags & FLAG_FIN != 0 {
1132 return Err(Error::Closed);
1133 }
1134 if flags & FLAG_READER_GONE != 0 {
1135 return Err(Error::PeerGone);
1136 }
1137 let buffer_bytes = self.table.geometry().buffer_bytes;
1138 if buf.len() > buffer_bytes {
1139 return Err(Error::Malformed(format!(
1140 "atomic stream write is {} bytes; ring capacity is {buffer_bytes}",
1141 buf.len()
1142 )));
1143 }
1144 if buf.is_empty() {
1145 return Ok(());
1146 }
1147 let head = direction.head.load(Ordering::Relaxed);
1148 let tail = direction.tail.load(Ordering::Acquire);
1149 if buffer_bytes - ((head - tail) as usize) < buf.len() {
1150 return Err(Error::WouldBlock);
1151 }
1152 let base = self.table.buffer(self.index, self.side.write_direction());
1153 let start = (head as usize) & (buffer_bytes - 1);
1154 let first = buf.len().min(buffer_bytes - start);
1155 unsafe {
1157 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
1158 std::ptr::copy_nonoverlapping(
1159 buf.as_ptr().add(first),
1160 base,
1161 buf.len() - first,
1162 );
1163 }
1164 direction.head.store(head + buf.len() as u64, Ordering::SeqCst);
1165 std::sync::atomic::fence(Ordering::SeqCst);
1166 let drained = direction.tail.load(Ordering::SeqCst) >= head;
1167 self.table.notify(self.index, slot, direction, self.peer_if(drained));
1168 Ok(())
1169 }
1170
1171 fn try_read(
1174 &self,
1175 buf: &mut [u8],
1176 ) -> Result<usize> {
1177 let slot = self.slot()?;
1178 let direction = Self::direction(slot, self.side.read_direction());
1179 let flags = direction.flags();
1180 if flags & FLAG_RESET != 0 {
1181 return Err(Error::Reset);
1182 }
1183 let head = direction.head.load(Ordering::Acquire);
1184 let tail = direction.tail.load(Ordering::Relaxed);
1185 let available = (head - tail) as usize;
1186 if buf.is_empty() {
1187 return Ok(0);
1188 }
1189 if available == 0 {
1190 if flags & FLAG_FIN != 0 {
1191 return Ok(0);
1192 }
1193 return Err(Error::WouldBlock);
1194 }
1195 let len = buf.len().min(available);
1196 let base = self.table.buffer(self.index, self.side.read_direction());
1197 let buffer_bytes = self.table.geometry().buffer_bytes;
1198 let start = (tail as usize) & (buffer_bytes - 1);
1199 let first = len.min(buffer_bytes - start);
1200 unsafe {
1204 std::ptr::copy_nonoverlapping(
1205 base.add(start),
1206 buf.as_mut_ptr(),
1207 first,
1208 );
1209 std::ptr::copy_nonoverlapping(
1210 base,
1211 buf.as_mut_ptr().add(first),
1212 len - first,
1213 );
1214 }
1215 direction.tail.store(tail + len as u64, Ordering::SeqCst);
1216 std::sync::atomic::fence(Ordering::SeqCst);
1219 let wake = writer_gained_capacity(
1220 direction,
1221 buffer_bytes,
1222 tail,
1223 tail + len as u64,
1224 );
1225 self.table.notify(self.index, slot, direction, self.peer_if(wake));
1226 Ok(len)
1227 }
1228
1229 fn try_read_exact(
1231 &self,
1232 buf: &mut [u8],
1233 ) -> Result<()> {
1234 let slot = self.slot()?;
1235 let direction = Self::direction(slot, self.side.read_direction());
1236 let flags = direction.flags();
1237 if flags & FLAG_RESET != 0 {
1238 return Err(Error::Reset);
1239 }
1240 let buffer_bytes = self.table.geometry().buffer_bytes;
1241 if buf.len() > buffer_bytes {
1242 return Err(Error::Malformed(format!(
1243 "atomic stream read is {} bytes; ring capacity is {buffer_bytes}",
1244 buf.len()
1245 )));
1246 }
1247 if buf.is_empty() {
1248 return Ok(());
1249 }
1250 let head = direction.head.load(Ordering::Acquire);
1251 let tail = direction.tail.load(Ordering::Relaxed);
1252 let available = (head - tail) as usize;
1253 if available < buf.len() {
1254 if flags & FLAG_FIN != 0 {
1255 return Err(Error::Malformed(format!(
1256 "stream ended with {available} bytes of a {} byte frame",
1257 buf.len()
1258 )));
1259 }
1260 return Err(Error::WouldBlock);
1261 }
1262 let base = self.table.buffer(self.index, self.side.read_direction());
1263 let start = (tail as usize) & (buffer_bytes - 1);
1264 let first = buf.len().min(buffer_bytes - start);
1265 unsafe {
1267 std::ptr::copy_nonoverlapping(
1268 base.add(start),
1269 buf.as_mut_ptr(),
1270 first,
1271 );
1272 std::ptr::copy_nonoverlapping(
1273 base,
1274 buf.as_mut_ptr().add(first),
1275 buf.len() - first,
1276 );
1277 }
1278 direction.tail.store(tail + buf.len() as u64, Ordering::SeqCst);
1279 std::sync::atomic::fence(Ordering::SeqCst);
1280 let wake = writer_gained_capacity(
1281 direction,
1282 buffer_bytes,
1283 tail,
1284 tail + buf.len() as u64,
1285 );
1286 self.table.notify(self.index, slot, direction, self.peer_if(wake));
1287 Ok(())
1288 }
1289
1290 fn set_flag(
1291 &self,
1292 direction_index: usize,
1293 flag: u8,
1294 ) {
1295 if let Ok(slot) = self.slot() {
1296 let direction = Self::direction(slot, direction_index);
1297 let before = direction.flags.fetch_or(flag, Ordering::SeqCst);
1298 if before & flag == 0 {
1299 let ring = flag != FLAG_READER_GONE
1300 || before & FLAG_CLOSE_WATCH != 0
1301 || writer_may_park(direction, before);
1302 self.table.notify(
1303 self.index,
1304 slot,
1305 direction,
1306 self.peer_if(ring),
1307 );
1308 }
1309 }
1310 }
1311
1312 fn set_flag_quietly(
1314 &self,
1315 direction_index: usize,
1316 flag: u8,
1317 ) {
1318 if let Ok(slot) = self.slot() {
1319 let direction = Self::direction(slot, direction_index);
1320 if direction.flags.fetch_or(flag, Ordering::SeqCst) & flag == 0 {
1321 self.table.notify(
1322 self.index,
1323 slot,
1324 direction,
1325 table::Ring::Nobody,
1326 );
1327 }
1328 }
1329 }
1330
1331 #[cfg(any(
1334 target_os = "linux",
1335 target_os = "freebsd",
1336 target_os = "macos"
1337 ))]
1338 fn wait_until_timeout(
1339 &self,
1340 direction_index: usize,
1341 ready: impl Fn(&Direction) -> bool,
1342 timeout: Duration,
1343 ) -> Result<bool> {
1344 let deadline = Instant::now() + timeout;
1345 loop {
1346 let slot = self.slot()?;
1347 let direction = Self::direction(slot, direction_index);
1348 if ready(direction) {
1349 return Ok(true);
1350 }
1351 let left = deadline.saturating_duration_since(Instant::now());
1352 if left.is_zero() {
1353 return Ok(false);
1354 }
1355 direction.waiters.fetch_add(1, Ordering::SeqCst);
1356 let since = direction.changes.load(Ordering::SeqCst);
1357 let outcome = if ready(direction) {
1358 Ok(true)
1359 } else {
1360 wait_on_timeout(&direction.changes, since, left)
1361 };
1362 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1363 if !outcome? {
1364 let slot = self.slot()?;
1367 return Ok(ready(Self::direction(slot, direction_index)));
1368 }
1369 }
1370 }
1371
1372 fn wait_until(
1373 &self,
1374 direction_index: usize,
1375 ready: impl Fn(&Direction) -> bool,
1376 ) -> Result<()> {
1377 loop {
1378 let slot = self.slot()?;
1379 let direction = Self::direction(slot, direction_index);
1380 if ready(direction) {
1381 return Ok(());
1382 }
1383 direction.waiters.fetch_add(1, Ordering::SeqCst);
1384 let since = direction.changes.load(Ordering::SeqCst);
1385 let outcome = if ready(direction) {
1386 Ok(())
1387 } else {
1388 wait_on(&direction.changes, since)
1389 };
1390 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1391 outcome?;
1392 }
1393 }
1394
1395 fn poll_ready(
1398 &self,
1399 interest: Interest,
1400 ready: impl Fn(&Direction) -> bool,
1401 cx: &mut std::task::Context<'_>,
1402 ) -> std::task::Poll<Result<()>> {
1403 let slot = match self.slot() {
1404 Ok(slot) => slot,
1405 Err(error) => return std::task::Poll::Ready(Err(error)),
1406 };
1407 let direction = Self::direction(slot, interest.direction);
1408 if ready(direction) {
1409 return std::task::Poll::Ready(Ok(()));
1410 }
1411 if let Err(error) =
1412 self.table.register(self.index, interest, cx.waker())
1413 {
1414 return std::task::Poll::Ready(Err(error));
1415 }
1416 std::sync::atomic::fence(Ordering::SeqCst);
1419 if ready(direction) {
1420 std::task::Poll::Ready(Ok(()))
1421 } else {
1422 std::task::Poll::Pending
1423 }
1424 }
1425}
1426
1427fn clear_exact_wait(direction: &Direction) {
1428 if direction.exact_needed.load(Ordering::Relaxed) != 0 {
1431 direction.exact_needed.store(0, Ordering::SeqCst);
1432 }
1433}
1434
1435fn writer_gained_capacity(
1436 direction: &Direction,
1437 capacity: usize,
1438 before_tail: u64,
1439 after_tail: u64,
1440) -> bool {
1441 let head = direction.head.load(Ordering::SeqCst);
1442 let needed = direction.exact_needed.load(Ordering::SeqCst) as usize;
1443 if needed == 0 {
1444 return head - before_tail >= capacity as u64;
1445 }
1446 if direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0 {
1447 return false;
1448 }
1449 let before = capacity.saturating_sub((head - before_tail) as usize);
1450 let after = capacity.saturating_sub((head - after_tail) as usize);
1451 before < needed && after >= needed
1452}
1453
1454fn writer_may_park(
1468 direction: &Direction,
1469 before: u8,
1470) -> bool {
1471 if before & (FLAG_FIN | FLAG_RESET) != 0 {
1472 return false;
1473 }
1474 std::sync::atomic::fence(Ordering::SeqCst);
1475 direction.head.load(Ordering::SeqCst)
1476 != direction.tail.load(Ordering::SeqCst)
1477}
1478
1479fn readable(direction: &Direction) -> bool {
1480 direction.head.load(Ordering::Acquire)
1481 != direction.tail.load(Ordering::Relaxed)
1482 || direction.flags() & (FLAG_FIN | FLAG_RESET) != 0
1483}
1484
1485fn writable(
1486 direction: &Direction,
1487 buffer_bytes: usize,
1488) -> bool {
1489 let queued = (direction.head.load(Ordering::Relaxed)
1490 - direction.tail.load(Ordering::Acquire)) as usize;
1491 queued < buffer_bytes
1492 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1493}
1494
1495fn writable_exact(
1496 direction: &Direction,
1497 buffer_bytes: usize,
1498 needed: usize,
1499) -> bool {
1500 let queued = (direction.head.load(Ordering::Relaxed)
1501 - direction.tail.load(Ordering::Acquire)) as usize;
1502 buffer_bytes.saturating_sub(queued) >= needed
1503 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1504}
1505
1506impl Drop for Handle {
1507 fn drop(&mut self) {
1511 let slot = &self.table.slots()[self.index];
1512 if !slot.is(self.id.generation()) {
1513 return;
1514 }
1515 if slot.claimed[self.side.index()]
1518 .compare_exchange(
1519 SIDE_CLAIMED,
1520 SIDE_RELEASED,
1521 Ordering::SeqCst,
1522 Ordering::SeqCst,
1523 )
1524 .is_err()
1525 {
1526 return;
1527 }
1528 let other = slot.claimed[1 - self.side.index()].load(Ordering::SeqCst);
1532 if other != SIDE_CLAIMED {
1533 let _ = slot.state.compare_exchange(
1534 SLOT_LIVE,
1535 SLOT_EMPTY,
1536 Ordering::SeqCst,
1537 Ordering::SeqCst,
1538 );
1539 for direction in &slot.directions {
1540 direction.changes.fetch_add(1, Ordering::SeqCst);
1543 if direction.waiters.load(Ordering::SeqCst) > 0 {
1544 wake_on(&direction.changes);
1545 }
1546 }
1547 self.table.registry.wake(self.index);
1548 }
1549 if self.side == Side::B || other != SIDE_CLAIMED {
1552 self.table.notify_capacity(self.index);
1553 }
1554 }
1555}
1556
1557pub struct Endpoint {
1561 read: ReadHalf,
1562 write: WriteHalf,
1563}
1564
1565impl Endpoint {
1566 fn new(handle: Arc<Handle>) -> Self {
1567 Self {
1568 read: ReadHalf { handle: Arc::clone(&handle) },
1569 write: WriteHalf { handle, shut: false },
1570 }
1571 }
1572
1573 pub fn id(&self) -> StreamId {
1574 self.read.handle.id
1575 }
1576
1577 pub fn side(&self) -> Side {
1578 self.read.handle.side
1579 }
1580
1581 pub fn split(self) -> (ReadHalf, WriteHalf) {
1582 (self.read, self.write)
1583 }
1584
1585 pub fn rejoin(
1596 read: ReadHalf,
1597 write: WriteHalf,
1598 ) -> Result<Self> {
1599 if read.id() != write.id() {
1600 return Err(Error::Stale(read.id()));
1601 }
1602 Ok(Self { read, write })
1603 }
1604
1605 pub fn rearm(
1618 &mut self,
1619 streams: &Streams,
1620 ) -> Result<()> {
1621 if self.side() != Side::A {
1622 return Err(Error::Malformed(
1623 "only the creator side of a stream can be rearmed".to_owned(),
1624 ));
1625 }
1626 streams.rearm(self.id())?;
1627 self.write.shut = false;
1628 Ok(())
1629 }
1630
1631 pub fn try_read(
1632 &self,
1633 buf: &mut [u8],
1634 ) -> Result<usize> {
1635 self.read.try_read(buf)
1636 }
1637
1638 pub fn try_read_exact(
1639 &self,
1640 buf: &mut [u8],
1641 ) -> Result<()> {
1642 self.read.try_read_exact(buf)
1643 }
1644
1645 pub fn blocking_read(
1646 &self,
1647 buf: &mut [u8],
1648 ) -> Result<usize> {
1649 self.read.blocking_read(buf)
1650 }
1651
1652 pub fn blocking_read_chunk(
1653 &self,
1654 max: usize,
1655 ) -> Result<Bytes> {
1656 self.read.blocking_read_chunk(max)
1657 }
1658
1659 pub fn try_write(
1660 &self,
1661 buf: &[u8],
1662 ) -> Result<usize> {
1663 self.write.try_write(buf)
1664 }
1665
1666 pub fn try_write_exact(
1667 &self,
1668 buf: &[u8],
1669 ) -> Result<()> {
1670 self.write.try_write_exact(buf)
1671 }
1672
1673 pub fn blocking_write(
1674 &self,
1675 buf: &[u8],
1676 ) -> Result<usize> {
1677 self.write.blocking_write(buf)
1678 }
1679
1680 pub fn blocking_write_all(
1681 &self,
1682 buf: &[u8],
1683 ) -> Result<()> {
1684 self.write.blocking_write_all(buf)
1685 }
1686
1687 pub fn finish(&mut self) -> Result<()> {
1688 self.write.finish()
1689 }
1690
1691 pub fn reset(&mut self) {
1692 self.write.reset()
1693 }
1694}
1695
1696pub struct ReadHalf {
1699 handle: Arc<Handle>,
1700}
1701
1702pub struct ResetSignal {
1706 handle: Arc<Handle>,
1707}
1708
1709impl ResetSignal {
1710 pub fn poll_reset(
1713 &self,
1714 cx: &mut std::task::Context<'_>,
1715 ) -> std::task::Poll<Result<()>> {
1716 self.handle.poll_ready(
1717 Interest {
1718 direction: self.handle.side.read_direction(),
1719 writer: false,
1720 reset: true,
1721 },
1722 |direction| direction.flags() & FLAG_RESET != 0,
1723 cx,
1724 )
1725 }
1726}
1727
1728impl ReadHalf {
1729 pub fn reset_signal(&self) -> ResetSignal {
1732 ResetSignal { handle: Arc::clone(&self.handle) }
1733 }
1734
1735 pub fn id(&self) -> StreamId {
1736 self.handle.id
1737 }
1738
1739 pub fn try_read(
1742 &self,
1743 buf: &mut [u8],
1744 ) -> Result<usize> {
1745 self.handle.try_read(buf)
1746 }
1747
1748 pub fn try_read_exact(
1751 &self,
1752 buf: &mut [u8],
1753 ) -> Result<()> {
1754 self.handle.try_read_exact(buf)
1755 }
1756
1757 pub fn blocking_read(
1760 &self,
1761 buf: &mut [u8],
1762 ) -> Result<usize> {
1763 loop {
1764 match self.handle.try_read(buf) {
1765 Err(Error::WouldBlock) => self.wait_readable()?,
1766 other => return other,
1767 }
1768 }
1769 }
1770
1771 pub fn stop_reading(&self) {
1779 self.handle
1780 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1781 }
1782
1783 pub fn blocking_read_chunk(
1785 &self,
1786 max: usize,
1787 ) -> Result<Bytes> {
1788 let mut chunk = BytesMut::zeroed(max);
1789 let len = self.blocking_read(&mut chunk)?;
1790 chunk.truncate(len);
1791 Ok(chunk.freeze())
1792 }
1793
1794 #[cfg(any(
1796 target_os = "linux",
1797 target_os = "freebsd",
1798 target_os = "macos"
1799 ))]
1800 pub fn wait_readable_timeout(
1801 &self,
1802 timeout: Duration,
1803 ) -> Result<bool> {
1804 self.handle.wait_until_timeout(
1805 self.handle.side.read_direction(),
1806 readable,
1807 timeout,
1808 )
1809 }
1810
1811 pub fn wait_readable(&self) -> Result<()> {
1813 self.handle.wait_until(self.handle.side.read_direction(), readable)
1814 }
1815
1816 pub fn poll_readable(
1821 &self,
1822 cx: &mut std::task::Context<'_>,
1823 ) -> std::task::Poll<Result<()>> {
1824 self.handle.poll_ready(
1825 Interest {
1826 direction: self.handle.side.read_direction(),
1827 writer: false,
1828 reset: false,
1829 },
1830 readable,
1831 cx,
1832 )
1833 }
1834}
1835
1836impl Drop for ReadHalf {
1837 fn drop(&mut self) {
1838 self.handle
1839 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1840 }
1841}
1842
1843pub struct WriteHalf {
1846 handle: Arc<Handle>,
1847 shut: bool,
1848}
1849
1850struct CloseWatch<'a>(&'a Handle);
1852
1853impl Drop for CloseWatch<'_> {
1854 fn drop(&mut self) {
1855 if let Ok(slot) = self.0.slot() {
1856 Handle::direction(slot, self.0.side.write_direction())
1857 .flags
1858 .fetch_and(!FLAG_CLOSE_WATCH, Ordering::SeqCst);
1859 }
1860 }
1861}
1862
1863impl WriteHalf {
1864 pub async fn peer_closed(&mut self) -> Result<()> {
1870 let slot = self.handle.slot()?;
1871 let direction =
1872 Handle::direction(slot, self.handle.side.write_direction());
1873 direction.flags.fetch_or(FLAG_CLOSE_WATCH, Ordering::SeqCst);
1874 let _watch = CloseWatch(&self.handle);
1875 std::future::poll_fn(|cx| {
1876 self.handle.poll_ready(
1877 Interest {
1878 direction: self.handle.side.write_direction(),
1879 writer: true,
1880 reset: false,
1881 },
1882 |direction| {
1883 direction.flags() & (FLAG_RESET | FLAG_READER_GONE) != 0
1884 },
1885 cx,
1886 )
1887 })
1888 .await
1889 }
1890
1891 pub fn poll_peer_dead(
1897 &self,
1898 cx: &mut std::task::Context<'_>,
1899 ) -> std::task::Poll<Result<()>> {
1900 let slot = match self.handle.slot() {
1901 Ok(slot) => slot,
1902 Err(error) => return std::task::Poll::Ready(Err(error)),
1903 };
1904 let peer = &slot.claimed[1 - self.handle.side.index()];
1905 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1906 return std::task::Poll::Ready(Ok(()));
1907 }
1908 let interest = Interest {
1909 direction: self.handle.side.write_direction(),
1910 writer: true,
1911 reset: false,
1912 };
1913 if let Err(error) =
1914 self.handle.table.register(self.handle.index, interest, cx.waker())
1915 {
1916 return std::task::Poll::Ready(Err(error));
1917 }
1918 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1921 std::task::Poll::Ready(Ok(()))
1922 } else {
1923 std::task::Poll::Pending
1924 }
1925 }
1926
1927 pub fn id(&self) -> StreamId {
1928 self.handle.id
1929 }
1930
1931 pub fn try_write(
1933 &self,
1934 buf: &[u8],
1935 ) -> Result<usize> {
1936 self.handle.try_write(buf)
1937 }
1938
1939 pub fn try_write_exact(
1942 &self,
1943 buf: &[u8],
1944 ) -> Result<()> {
1945 self.handle.try_write_exact(buf)
1946 }
1947
1948 pub fn blocking_write(
1950 &self,
1951 buf: &[u8],
1952 ) -> Result<usize> {
1953 loop {
1954 match self.handle.try_write(buf) {
1955 Err(Error::WouldBlock) => self.wait_writable()?,
1956 other => return other,
1957 }
1958 }
1959 }
1960
1961 pub fn blocking_write_all(
1964 &self,
1965 mut buf: &[u8],
1966 ) -> Result<()> {
1967 while !buf.is_empty() {
1968 let written = self.blocking_write(buf)?;
1969 buf = &buf[written..];
1970 }
1971 Ok(())
1972 }
1973
1974 pub fn finish(&mut self) -> Result<()> {
1978 self.shut = true;
1979 self.handle.set_flag(self.handle.side.write_direction(), FLAG_FIN);
1980 Ok(())
1981 }
1982
1983 pub fn finish_quietly(&mut self) {
1990 self.shut = true;
1991 self.handle
1992 .set_flag_quietly(self.handle.side.write_direction(), FLAG_FIN);
1993 }
1994
1995 pub fn reset(&mut self) {
1998 self.shut = true;
1999 self.handle.set_flag(self.handle.side.write_direction(), FLAG_RESET);
2000 }
2001
2002 #[cfg(any(
2006 target_os = "linux",
2007 target_os = "freebsd",
2008 target_os = "macos"
2009 ))]
2010 pub fn wait_writable_timeout(
2011 &self,
2012 timeout: Duration,
2013 ) -> Result<bool> {
2014 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2015 self.handle.wait_until_timeout(
2016 self.handle.side.write_direction(),
2017 |direction| {
2018 clear_exact_wait(direction);
2019 writable(direction, buffer_bytes)
2020 },
2021 timeout,
2022 )
2023 }
2024
2025 pub fn wait_writable(&self) -> Result<()> {
2027 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2028 self.handle.wait_until(
2029 self.handle.side.write_direction(),
2030 |direction| {
2031 clear_exact_wait(direction);
2032 writable(direction, buffer_bytes)
2033 },
2034 )
2035 }
2036
2037 pub fn wait_writable_exact(
2039 &self,
2040 needed: usize,
2041 ) -> Result<()> {
2042 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2043 if needed > buffer_bytes {
2044 return Err(Error::Malformed(format!(
2045 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
2046 )));
2047 }
2048 self.handle.wait_until(
2049 self.handle.side.write_direction(),
2050 |direction| {
2051 clear_exact_wait(direction);
2052 writable_exact(direction, buffer_bytes, needed)
2053 },
2054 )
2055 }
2056
2057 pub fn poll_writable(
2062 &self,
2063 cx: &mut std::task::Context<'_>,
2064 ) -> std::task::Poll<Result<()>> {
2065 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2066 self.handle.poll_ready(
2067 Interest {
2068 direction: self.handle.side.write_direction(),
2069 writer: true,
2070 reset: false,
2071 },
2072 |direction| {
2073 clear_exact_wait(direction);
2074 writable(direction, buffer_bytes)
2075 },
2076 cx,
2077 )
2078 }
2079
2080 pub fn poll_writable_exact(
2082 &self,
2083 needed: usize,
2084 cx: &mut std::task::Context<'_>,
2085 ) -> std::task::Poll<Result<()>> {
2086 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2087 let slot = match self.handle.slot() {
2088 Ok(slot) => slot,
2089 Err(error) => return std::task::Poll::Ready(Err(error)),
2090 };
2091 let direction =
2092 Handle::direction(slot, self.handle.side.write_direction());
2093 if needed > buffer_bytes {
2094 clear_exact_wait(direction);
2095 return std::task::Poll::Ready(Err(Error::Malformed(format!(
2096 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
2097 ))));
2098 }
2099 direction.exact_needed.store(needed as u32, Ordering::SeqCst);
2103 let result = self.handle.poll_ready(
2104 Interest {
2105 direction: self.handle.side.write_direction(),
2106 writer: true,
2107 reset: false,
2108 },
2109 |direction| writable_exact(direction, buffer_bytes, needed),
2110 cx,
2111 );
2112 if result.is_ready() {
2113 clear_exact_wait(direction);
2114 }
2115 result
2116 }
2117}
2118
2119impl Drop for WriteHalf {
2120 fn drop(&mut self) {
2121 if !self.shut {
2122 self.handle
2123 .set_flag(self.handle.side.write_direction(), FLAG_RESET);
2124 }
2125 }
2126}
2127
2128#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2132pub(crate) fn waits_supported() -> bool {
2133 orbit_core::sync::supported()
2134}
2135
2136#[cfg(not(any(
2137 target_os = "linux",
2138 target_os = "freebsd",
2139 target_os = "macos"
2140)))]
2141pub(crate) fn waits_supported() -> bool {
2142 false
2143}
2144
2145#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2149#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2152pub(crate) fn wait_on_timeout(
2153 word: &AtomicU32,
2154 expected: u32,
2155 timeout: Duration,
2156) -> Result<bool> {
2157 orbit_core::sync::wait_word_timeout(word, expected, timeout)
2158 .map_err(Error::Io)
2159}
2160
2161pub(crate) fn wait_on(
2162 word: &AtomicU32,
2163 expected: u32,
2164) -> Result<()> {
2165 orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
2166}
2167
2168#[cfg(not(any(
2169 target_os = "linux",
2170 target_os = "freebsd",
2171 target_os = "macos"
2172)))]
2173pub(crate) fn wait_on(
2174 _word: &AtomicU32,
2175 _expected: u32,
2176) -> Result<()> {
2177 Err(Error::Io(std::io::Error::new(
2178 std::io::ErrorKind::Unsupported,
2179 "orbit-stream needs a platform that can wait on a shared word",
2180 )))
2181}
2182
2183pub(crate) fn wake_on(word: &AtomicU32) {
2185 #[cfg(any(
2186 target_os = "linux",
2187 target_os = "freebsd",
2188 target_os = "macos"
2189 ))]
2190 let _ = orbit_core::sync::wake_word(word);
2191 #[cfg(not(any(
2192 target_os = "linux",
2193 target_os = "freebsd",
2194 target_os = "macos"
2195 )))]
2196 let _ = word;
2197}
2198
2199pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
2200 mutex.lock().unwrap_or_else(|error| error.into_inner())
2201}
2202
2203#[cfg(test)]
2204mod tests {
2205 use std::sync::Arc;
2206 use std::time::Duration;
2207
2208 use orbit_core::{Fleet, NodeId};
2209
2210 use super::{
2211 Error, Incarnation, STREAM_BUFFER_BYTES, STREAM_LANE_CAPACITY, Side,
2212 Streams, Ticket,
2213 };
2214
2215 fn streams(name: &'static str) -> Streams {
2216 Streams::new(
2217 Arc::new(Fleet::join(name, 2).unwrap()),
2218 Incarnation::new(1),
2219 )
2220 .unwrap()
2221 }
2222
2223 #[test]
2224 fn bytes_cross_in_order_and_end_cleanly() {
2225 let streams = streams("stream-order");
2226 let (mut a, ticket) = streams.create().unwrap();
2227 assert_eq!(ticket.side, Side::B);
2228 assert_eq!(ticket.id.node(), 0);
2229 assert_eq!(ticket.epoch, streams.epoch());
2230 let b = streams.open(ticket).unwrap();
2231
2232 a.blocking_write_all(b"hello, ").unwrap();
2233 a.blocking_write_all(b"world").unwrap();
2234 a.finish().unwrap();
2235
2236 let mut out = Vec::new();
2237 let mut buf = [0_u8; 4];
2238 loop {
2239 let n = b.blocking_read(&mut buf).unwrap();
2240 if n == 0 {
2241 break;
2242 }
2243 out.extend_from_slice(&buf[..n]);
2244 }
2245 assert_eq!(out, b"hello, world");
2246 b.blocking_write_all(b"back").unwrap();
2248 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"back");
2249 }
2250
2251 #[test]
2252 fn a_ticket_prints_and_parses() {
2253 let streams = streams("stream-ticket");
2254 let (_a, ticket) = streams.create().unwrap();
2255 let text = ticket.to_string();
2256 assert!(text.contains("/b/"), "{text}");
2257 assert_eq!(text.parse::<Ticket>().unwrap(), ticket);
2258 assert!(matches!("nope".parse::<Ticket>(), Err(Error::Malformed(_))));
2259 assert!(matches!("x/b".parse::<Ticket>(), Err(Error::Malformed(_))));
2260 }
2261
2262 #[test]
2263 fn a_side_is_held_once_and_a_released_stream_goes_stale() {
2264 let streams = streams("stream-claim");
2265 let (a, ticket) = streams.create().unwrap();
2266 let b = streams.open(ticket).unwrap();
2267 assert!(matches!(streams.open(ticket), Err(Error::AlreadyClaimed(_))));
2268 assert!(streams.is_live(ticket.id));
2269
2270 drop(a);
2271 let mut buf = [0_u8; 1];
2273 assert!(matches!(b.blocking_read(&mut buf), Err(Error::Reset)));
2274 drop(b);
2275 assert!(!streams.is_live(ticket.id));
2276 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2277
2278 for _ in 0..STREAM_LANE_CAPACITY - 1 {
2281 drop(streams.create().unwrap());
2282 }
2283 let (_a2, ticket2) = streams.create().unwrap();
2284 assert_eq!(ticket2.id.slot(), ticket.id.slot());
2285 assert_ne!(ticket2.id.generation(), ticket.id.generation());
2286 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2287 }
2288
2289 #[test]
2290 fn a_reset_starts_a_new_epoch_and_refuses_old_tickets() {
2291 let streams = streams("stream-epoch");
2292 let (_a, ticket) = streams.create().unwrap();
2293 let before = streams.epoch();
2294 streams.reset_all();
2295 assert!(streams.epoch() > before);
2296 let (_a2, ticket2) = streams.create().unwrap();
2299 assert_eq!(ticket2.id, ticket.id);
2300 assert_ne!(ticket2.epoch, ticket.epoch);
2301 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2302 assert!(streams.open(ticket2).is_ok());
2303 }
2304
2305 #[test]
2306 fn a_slot_whose_generation_runs_out_is_not_reused() {
2307 let streams = streams("stream-exhaust");
2310 let limit = super::layout::GENERATION_LIMIT as usize;
2311 for _ in 0..STREAM_LANE_CAPACITY * limit {
2312 drop(streams.create().unwrap());
2313 }
2314 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2315 streams.reset_all();
2316 let (_a, ticket) = streams.create().unwrap();
2317 assert_eq!(ticket.id.generation(), 1);
2318 }
2319
2320 #[test]
2321 fn a_full_ring_blocks_the_writer_and_wraps_correctly() {
2322 let streams = streams("stream-wrap");
2323 let (a, ticket) = streams.create().unwrap();
2324 let b = streams.open(ticket).unwrap();
2325
2326 let filler = vec![7_u8; STREAM_BUFFER_BYTES];
2327 assert_eq!(a.try_write(&filler).unwrap(), STREAM_BUFFER_BYTES);
2328 assert!(matches!(a.try_write(b"x"), Err(Error::WouldBlock)));
2329
2330 let mut buf = vec![0_u8; STREAM_BUFFER_BYTES - 5];
2331 assert_eq!(b.blocking_read(&mut buf).unwrap(), STREAM_BUFFER_BYTES - 5);
2332 let tail =
2334 (0..STREAM_BUFFER_BYTES).map(|i| i as u8).collect::<Vec<_>>();
2335 assert_eq!(a.try_write(&tail).unwrap(), STREAM_BUFFER_BYTES - 5);
2336
2337 let mut rest = vec![0_u8; 5];
2338 assert_eq!(b.blocking_read(&mut rest).unwrap(), 5);
2339 assert_eq!(rest, vec![7_u8; 5]);
2340 let mut wrapped = vec![0_u8; STREAM_BUFFER_BYTES];
2341 let mut got = 0;
2342 while got < STREAM_BUFFER_BYTES - 5 {
2343 got += b.blocking_read(&mut wrapped[got..]).unwrap();
2344 }
2345 assert_eq!(&wrapped[..got], &tail[..got]);
2346 }
2347
2348 #[test]
2349 fn exact_frames_never_publish_or_consume_a_prefix() {
2350 let streams = streams("stream-exact-frame");
2351 let (a, ticket) = streams.create().unwrap();
2352 let b = streams.open(ticket).unwrap();
2353
2354 let filler = vec![7_u8; STREAM_BUFFER_BYTES - 32];
2355 a.try_write_exact(&filler).unwrap();
2356 assert!(matches!(
2357 a.try_write_exact(&[9_u8; 64]),
2358 Err(Error::WouldBlock)
2359 ));
2360 let mut consumed = vec![0_u8; filler.len()];
2361 b.try_read_exact(&mut consumed).unwrap();
2362 assert_eq!(consumed, filler);
2363 assert!(matches!(b.try_read(&mut [0_u8; 1]), Err(Error::WouldBlock)));
2364
2365 a.try_write_exact(&[3_u8; 32]).unwrap();
2366 assert!(matches!(
2367 b.try_read_exact(&mut [0_u8; 64]),
2368 Err(Error::WouldBlock)
2369 ));
2370 let mut frame = [0_u8; 32];
2371 b.try_read_exact(&mut frame).unwrap();
2372 assert_eq!(frame, [3_u8; 32]);
2373 }
2374
2375 #[test]
2376 fn a_blocked_reader_is_woken_by_a_write_in_another_thread() {
2377 let streams = streams("stream-wake");
2378 let (a, ticket) = streams.create().unwrap();
2379 let b = streams.open(ticket).unwrap();
2380 let reader = std::thread::spawn(move || {
2381 let mut buf = [0_u8; 8];
2382 let n = b.blocking_read(&mut buf).unwrap();
2383 buf[..n].to_vec()
2384 });
2385 std::thread::sleep(Duration::from_millis(30));
2386 a.blocking_write_all(b"ping").unwrap();
2387 assert_eq!(reader.join().unwrap(), b"ping");
2388 }
2389
2390 #[test]
2391 fn a_reader_gone_fails_the_writer_and_a_lane_fills_up() {
2392 let streams = streams("stream-gone");
2393 let (a, ticket) = streams.create().unwrap();
2394 let b = streams.open(ticket).unwrap();
2395 let (b_read, _b_write) = b.split();
2396 drop(b_read);
2397 assert!(matches!(a.blocking_write(b"nobody"), Err(Error::PeerGone)));
2398
2399 let mut held = vec![a];
2400 for _ in 1..STREAM_LANE_CAPACITY {
2401 held.push(streams.create().unwrap().0);
2402 }
2403 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2404 }
2405
2406 #[test]
2407 fn a_death_report_ends_that_incarnations_sides_only() {
2408 let streams = streams("stream-dead");
2409 let (a, ticket) = streams.create().unwrap();
2410 let b = streams.open(ticket).unwrap();
2411 b.blocking_write_all(b"partial").unwrap();
2412
2413 streams.node_dead(NodeId::ZERO, Incarnation::new(2));
2415 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"partial");
2416
2417 let reader = std::thread::spawn(move || {
2423 let mut buf = [0_u8; 8];
2424 a.blocking_read(&mut buf).map(|_| ())
2425 });
2426 std::thread::sleep(Duration::from_millis(30));
2427 streams.node_dead(NodeId::ZERO, Incarnation::new(1));
2428 assert!(matches!(
2429 reader.join().unwrap(),
2430 Err(Error::Reset | Error::Stale(_))
2431 ));
2432 assert!(matches!(
2433 b.try_write(b"x"),
2434 Err(Error::Reset | Error::Stale(_))
2435 ));
2436 assert!(matches!(
2437 streams.open(ticket),
2438 Err(Error::AlreadyClaimed(_) | Error::Stale(_))
2439 ));
2440 drop(b);
2441 assert!(!streams.is_live(ticket.id));
2442 }
2443
2444 #[test]
2445 fn an_offer_is_taken_once_and_a_stale_offer_is_skipped() {
2446 let streams = streams("stream-offer");
2447 assert!(streams.take_offer().is_none());
2448 let (a, ticket) = streams.create().unwrap();
2449 streams.offer(ticket, NodeId::ZERO).unwrap();
2450 assert_eq!(streams.take_offer(), Some(ticket));
2451 assert!(streams.take_offer().is_none());
2452
2453 let (gone, stale) = streams.create().unwrap();
2454 streams.offer(stale, NodeId::ZERO).unwrap();
2455 drop(gone);
2456 assert!(streams.take_offer().is_none());
2457
2458 let taker = streams.clone();
2459 let waiter = std::thread::spawn(move || taker.blocking_take_offer());
2460 std::thread::sleep(Duration::from_millis(30));
2461 streams.offer(ticket, NodeId::ZERO).unwrap();
2462 assert_eq!(waiter.join().unwrap().unwrap(), ticket);
2463 drop(a);
2464 }
2465
2466 #[test]
2467 fn a_contested_claim_goes_to_exactly_one_opener() {
2468 let streams = streams("stream-race");
2469 for _ in 0..50 {
2470 let (_a, ticket) = streams.create().unwrap();
2471 let openers = (0..4)
2472 .map(|_| {
2473 let streams = streams.clone();
2474 std::thread::spawn(move || streams.open(ticket).is_ok())
2475 })
2476 .collect::<Vec<_>>();
2477 let won = openers
2478 .into_iter()
2479 .map(|opener| opener.join().unwrap())
2480 .filter(|won| *won)
2481 .count();
2482 assert_eq!(won, 1);
2483 }
2484 }
2485
2486 #[test]
2487 fn random_sizes_concatenate_exactly() {
2488 let streams = streams("stream-random");
2489 let (mut a, ticket) = streams.create().unwrap();
2490 let b = streams.open(ticket).unwrap();
2491 let payload = (0..300_000_u32)
2492 .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
2493 .collect::<Vec<_>>();
2494 let expected = payload.clone();
2495 let writer = std::thread::spawn(move || {
2496 let mut rest = &payload[..];
2497 let mut step = 1_usize;
2498 while !rest.is_empty() {
2499 let len = step.min(rest.len());
2500 a.blocking_write_all(&rest[..len]).unwrap();
2501 rest = &rest[len..];
2502 step = (step * 7 + 13) % 9_001 + 1;
2503 }
2504 a.finish().unwrap();
2505 });
2506 let mut out = Vec::new();
2507 let mut step = 3_usize;
2508 loop {
2509 let chunk = b.blocking_read_chunk(step).unwrap();
2510 if chunk.is_empty() {
2511 break;
2512 }
2513 out.extend_from_slice(&chunk);
2514 step = (step * 5 + 11) % 4_099 + 1;
2515 }
2516 writer.join().unwrap();
2517 assert_eq!(out, expected);
2518 }
2519}