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