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 fn set_flag_quietly(
1272 &self,
1273 direction_index: usize,
1274 flag: u8,
1275 ) {
1276 if let Ok(slot) = self.slot() {
1277 let direction = Self::direction(slot, direction_index);
1278 if direction.flags.fetch_or(flag, Ordering::SeqCst) & flag == 0 {
1279 self.table.notify(
1280 self.index,
1281 slot,
1282 direction,
1283 table::Ring::Nobody,
1284 );
1285 }
1286 }
1287 }
1288
1289 #[cfg(any(
1292 target_os = "linux",
1293 target_os = "freebsd",
1294 target_os = "macos"
1295 ))]
1296 fn wait_until_timeout(
1297 &self,
1298 direction_index: usize,
1299 ready: impl Fn(&Direction) -> bool,
1300 timeout: Duration,
1301 ) -> Result<bool> {
1302 let deadline = Instant::now() + timeout;
1303 loop {
1304 let slot = self.slot()?;
1305 let direction = Self::direction(slot, direction_index);
1306 if ready(direction) {
1307 return Ok(true);
1308 }
1309 let left = deadline.saturating_duration_since(Instant::now());
1310 if left.is_zero() {
1311 return Ok(false);
1312 }
1313 direction.waiters.fetch_add(1, Ordering::SeqCst);
1314 let since = direction.changes.load(Ordering::SeqCst);
1315 let outcome = if ready(direction) {
1316 Ok(true)
1317 } else {
1318 wait_on_timeout(&direction.changes, since, left)
1319 };
1320 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1321 if !outcome? {
1322 let slot = self.slot()?;
1325 return Ok(ready(Self::direction(slot, direction_index)));
1326 }
1327 }
1328 }
1329
1330 fn wait_until(
1331 &self,
1332 direction_index: usize,
1333 ready: impl Fn(&Direction) -> bool,
1334 ) -> Result<()> {
1335 loop {
1336 let slot = self.slot()?;
1337 let direction = Self::direction(slot, direction_index);
1338 if ready(direction) {
1339 return Ok(());
1340 }
1341 direction.waiters.fetch_add(1, Ordering::SeqCst);
1342 let since = direction.changes.load(Ordering::SeqCst);
1343 let outcome = if ready(direction) {
1344 Ok(())
1345 } else {
1346 wait_on(&direction.changes, since)
1347 };
1348 direction.waiters.fetch_sub(1, Ordering::SeqCst);
1349 outcome?;
1350 }
1351 }
1352
1353 fn poll_ready(
1356 &self,
1357 interest: Interest,
1358 ready: impl Fn(&Direction) -> bool,
1359 cx: &mut std::task::Context<'_>,
1360 ) -> std::task::Poll<Result<()>> {
1361 let slot = match self.slot() {
1362 Ok(slot) => slot,
1363 Err(error) => return std::task::Poll::Ready(Err(error)),
1364 };
1365 let direction = Self::direction(slot, interest.direction);
1366 if ready(direction) {
1367 return std::task::Poll::Ready(Ok(()));
1368 }
1369 if let Err(error) =
1370 self.table.register(self.index, interest, cx.waker())
1371 {
1372 return std::task::Poll::Ready(Err(error));
1373 }
1374 std::sync::atomic::fence(Ordering::SeqCst);
1377 if ready(direction) {
1378 std::task::Poll::Ready(Ok(()))
1379 } else {
1380 std::task::Poll::Pending
1381 }
1382 }
1383}
1384
1385fn clear_exact_wait(direction: &Direction) {
1386 if direction.exact_needed.load(Ordering::Relaxed) != 0 {
1389 direction.exact_needed.store(0, Ordering::SeqCst);
1390 }
1391}
1392
1393fn writer_gained_capacity(
1394 direction: &Direction,
1395 capacity: usize,
1396 before_tail: u64,
1397 after_tail: u64,
1398) -> bool {
1399 let head = direction.head.load(Ordering::SeqCst);
1400 let needed = direction.exact_needed.load(Ordering::SeqCst) as usize;
1401 if needed == 0 {
1402 return head - before_tail >= capacity as u64;
1403 }
1404 if direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0 {
1405 return false;
1406 }
1407 let before = capacity.saturating_sub((head - before_tail) as usize);
1408 let after = capacity.saturating_sub((head - after_tail) as usize);
1409 before < needed && after >= needed
1410}
1411
1412fn writer_may_park(
1426 direction: &Direction,
1427 before: u8,
1428) -> bool {
1429 if before & (FLAG_FIN | FLAG_RESET) != 0 {
1430 return false;
1431 }
1432 std::sync::atomic::fence(Ordering::SeqCst);
1433 direction.head.load(Ordering::SeqCst)
1434 != direction.tail.load(Ordering::SeqCst)
1435}
1436
1437fn readable(direction: &Direction) -> bool {
1438 direction.head.load(Ordering::Acquire)
1439 != direction.tail.load(Ordering::Relaxed)
1440 || direction.flags() & (FLAG_FIN | FLAG_RESET) != 0
1441}
1442
1443fn writable(
1444 direction: &Direction,
1445 buffer_bytes: usize,
1446) -> bool {
1447 let queued = (direction.head.load(Ordering::Relaxed)
1448 - direction.tail.load(Ordering::Acquire)) as usize;
1449 queued < buffer_bytes
1450 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1451}
1452
1453fn writable_exact(
1454 direction: &Direction,
1455 buffer_bytes: usize,
1456 needed: usize,
1457) -> bool {
1458 let queued = (direction.head.load(Ordering::Relaxed)
1459 - direction.tail.load(Ordering::Acquire)) as usize;
1460 buffer_bytes.saturating_sub(queued) >= needed
1461 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
1462}
1463
1464impl Drop for Handle {
1465 fn drop(&mut self) {
1469 let slot = &self.table.slots()[self.index];
1470 if !slot.is(self.id.generation()) {
1471 return;
1472 }
1473 if slot.claimed[self.side.index()]
1476 .compare_exchange(
1477 SIDE_CLAIMED,
1478 SIDE_RELEASED,
1479 Ordering::SeqCst,
1480 Ordering::SeqCst,
1481 )
1482 .is_err()
1483 {
1484 return;
1485 }
1486 let other = slot.claimed[1 - self.side.index()].load(Ordering::SeqCst);
1490 if other != SIDE_CLAIMED {
1491 let _ = slot.state.compare_exchange(
1492 SLOT_LIVE,
1493 SLOT_EMPTY,
1494 Ordering::SeqCst,
1495 Ordering::SeqCst,
1496 );
1497 for direction in &slot.directions {
1498 direction.changes.fetch_add(1, Ordering::SeqCst);
1501 if direction.waiters.load(Ordering::SeqCst) > 0 {
1502 wake_on(&direction.changes);
1503 }
1504 }
1505 self.table.registry.wake(self.index);
1506 }
1507 }
1508}
1509
1510pub struct Endpoint {
1514 read: ReadHalf,
1515 write: WriteHalf,
1516}
1517
1518impl Endpoint {
1519 fn new(handle: Arc<Handle>) -> Self {
1520 Self {
1521 read: ReadHalf { handle: Arc::clone(&handle) },
1522 write: WriteHalf { handle, shut: false },
1523 }
1524 }
1525
1526 pub fn id(&self) -> StreamId {
1527 self.read.handle.id
1528 }
1529
1530 pub fn side(&self) -> Side {
1531 self.read.handle.side
1532 }
1533
1534 pub fn split(self) -> (ReadHalf, WriteHalf) {
1535 (self.read, self.write)
1536 }
1537
1538 pub fn rejoin(
1549 read: ReadHalf,
1550 write: WriteHalf,
1551 ) -> Result<Self> {
1552 if read.id() != write.id() {
1553 return Err(Error::Stale(read.id()));
1554 }
1555 Ok(Self { read, write })
1556 }
1557
1558 pub fn rearm(
1571 &mut self,
1572 streams: &Streams,
1573 ) -> Result<()> {
1574 if self.side() != Side::A {
1575 return Err(Error::Malformed(
1576 "only the creator side of a stream can be rearmed".to_owned(),
1577 ));
1578 }
1579 streams.rearm(self.id())?;
1580 self.write.shut = false;
1581 Ok(())
1582 }
1583
1584 pub fn try_read(
1585 &self,
1586 buf: &mut [u8],
1587 ) -> Result<usize> {
1588 self.read.try_read(buf)
1589 }
1590
1591 pub fn try_read_exact(
1592 &self,
1593 buf: &mut [u8],
1594 ) -> Result<()> {
1595 self.read.try_read_exact(buf)
1596 }
1597
1598 pub fn blocking_read(
1599 &self,
1600 buf: &mut [u8],
1601 ) -> Result<usize> {
1602 self.read.blocking_read(buf)
1603 }
1604
1605 pub fn blocking_read_chunk(
1606 &self,
1607 max: usize,
1608 ) -> Result<Bytes> {
1609 self.read.blocking_read_chunk(max)
1610 }
1611
1612 pub fn try_write(
1613 &self,
1614 buf: &[u8],
1615 ) -> Result<usize> {
1616 self.write.try_write(buf)
1617 }
1618
1619 pub fn try_write_exact(
1620 &self,
1621 buf: &[u8],
1622 ) -> Result<()> {
1623 self.write.try_write_exact(buf)
1624 }
1625
1626 pub fn blocking_write(
1627 &self,
1628 buf: &[u8],
1629 ) -> Result<usize> {
1630 self.write.blocking_write(buf)
1631 }
1632
1633 pub fn blocking_write_all(
1634 &self,
1635 buf: &[u8],
1636 ) -> Result<()> {
1637 self.write.blocking_write_all(buf)
1638 }
1639
1640 pub fn finish(&mut self) -> Result<()> {
1641 self.write.finish()
1642 }
1643
1644 pub fn reset(&mut self) {
1645 self.write.reset()
1646 }
1647}
1648
1649pub struct ReadHalf {
1652 handle: Arc<Handle>,
1653}
1654
1655pub struct ResetSignal {
1659 handle: Arc<Handle>,
1660}
1661
1662impl ResetSignal {
1663 pub fn poll_reset(
1666 &self,
1667 cx: &mut std::task::Context<'_>,
1668 ) -> std::task::Poll<Result<()>> {
1669 self.handle.poll_ready(
1670 Interest {
1671 direction: self.handle.side.read_direction(),
1672 writer: false,
1673 reset: true,
1674 },
1675 |direction| direction.flags() & FLAG_RESET != 0,
1676 cx,
1677 )
1678 }
1679}
1680
1681impl ReadHalf {
1682 pub fn reset_signal(&self) -> ResetSignal {
1685 ResetSignal { handle: Arc::clone(&self.handle) }
1686 }
1687
1688 pub fn id(&self) -> StreamId {
1689 self.handle.id
1690 }
1691
1692 pub fn try_read(
1695 &self,
1696 buf: &mut [u8],
1697 ) -> Result<usize> {
1698 self.handle.try_read(buf)
1699 }
1700
1701 pub fn try_read_exact(
1704 &self,
1705 buf: &mut [u8],
1706 ) -> Result<()> {
1707 self.handle.try_read_exact(buf)
1708 }
1709
1710 pub fn blocking_read(
1713 &self,
1714 buf: &mut [u8],
1715 ) -> Result<usize> {
1716 loop {
1717 match self.handle.try_read(buf) {
1718 Err(Error::WouldBlock) => self.wait_readable()?,
1719 other => return other,
1720 }
1721 }
1722 }
1723
1724 pub fn stop_reading(&self) {
1732 self.handle
1733 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1734 }
1735
1736 pub fn blocking_read_chunk(
1738 &self,
1739 max: usize,
1740 ) -> Result<Bytes> {
1741 let mut chunk = BytesMut::zeroed(max);
1742 let len = self.blocking_read(&mut chunk)?;
1743 chunk.truncate(len);
1744 Ok(chunk.freeze())
1745 }
1746
1747 #[cfg(any(
1749 target_os = "linux",
1750 target_os = "freebsd",
1751 target_os = "macos"
1752 ))]
1753 pub fn wait_readable_timeout(
1754 &self,
1755 timeout: Duration,
1756 ) -> Result<bool> {
1757 self.handle.wait_until_timeout(
1758 self.handle.side.read_direction(),
1759 readable,
1760 timeout,
1761 )
1762 }
1763
1764 pub fn wait_readable(&self) -> Result<()> {
1766 self.handle.wait_until(self.handle.side.read_direction(), readable)
1767 }
1768
1769 pub fn poll_readable(
1774 &self,
1775 cx: &mut std::task::Context<'_>,
1776 ) -> std::task::Poll<Result<()>> {
1777 self.handle.poll_ready(
1778 Interest {
1779 direction: self.handle.side.read_direction(),
1780 writer: false,
1781 reset: false,
1782 },
1783 readable,
1784 cx,
1785 )
1786 }
1787}
1788
1789impl Drop for ReadHalf {
1790 fn drop(&mut self) {
1791 self.handle
1792 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1793 }
1794}
1795
1796pub struct WriteHalf {
1799 handle: Arc<Handle>,
1800 shut: bool,
1801}
1802
1803struct CloseWatch<'a>(&'a Handle);
1805
1806impl Drop for CloseWatch<'_> {
1807 fn drop(&mut self) {
1808 if let Ok(slot) = self.0.slot() {
1809 Handle::direction(slot, self.0.side.write_direction())
1810 .flags
1811 .fetch_and(!FLAG_CLOSE_WATCH, Ordering::SeqCst);
1812 }
1813 }
1814}
1815
1816impl WriteHalf {
1817 pub async fn peer_closed(&mut self) -> Result<()> {
1823 let slot = self.handle.slot()?;
1824 let direction =
1825 Handle::direction(slot, self.handle.side.write_direction());
1826 direction.flags.fetch_or(FLAG_CLOSE_WATCH, Ordering::SeqCst);
1827 let _watch = CloseWatch(&self.handle);
1828 std::future::poll_fn(|cx| {
1829 self.handle.poll_ready(
1830 Interest {
1831 direction: self.handle.side.write_direction(),
1832 writer: true,
1833 reset: false,
1834 },
1835 |direction| {
1836 direction.flags() & (FLAG_RESET | FLAG_READER_GONE) != 0
1837 },
1838 cx,
1839 )
1840 })
1841 .await
1842 }
1843
1844 pub fn poll_peer_dead(
1850 &self,
1851 cx: &mut std::task::Context<'_>,
1852 ) -> std::task::Poll<Result<()>> {
1853 let slot = match self.handle.slot() {
1854 Ok(slot) => slot,
1855 Err(error) => return std::task::Poll::Ready(Err(error)),
1856 };
1857 let peer = &slot.claimed[1 - self.handle.side.index()];
1858 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1859 return std::task::Poll::Ready(Ok(()));
1860 }
1861 let interest = Interest {
1862 direction: self.handle.side.write_direction(),
1863 writer: true,
1864 reset: false,
1865 };
1866 if let Err(error) =
1867 self.handle.table.register(self.handle.index, interest, cx.waker())
1868 {
1869 return std::task::Poll::Ready(Err(error));
1870 }
1871 if peer.load(Ordering::SeqCst) == SIDE_DEAD {
1874 std::task::Poll::Ready(Ok(()))
1875 } else {
1876 std::task::Poll::Pending
1877 }
1878 }
1879
1880 pub fn id(&self) -> StreamId {
1881 self.handle.id
1882 }
1883
1884 pub fn try_write(
1886 &self,
1887 buf: &[u8],
1888 ) -> Result<usize> {
1889 self.handle.try_write(buf)
1890 }
1891
1892 pub fn try_write_exact(
1895 &self,
1896 buf: &[u8],
1897 ) -> Result<()> {
1898 self.handle.try_write_exact(buf)
1899 }
1900
1901 pub fn blocking_write(
1903 &self,
1904 buf: &[u8],
1905 ) -> Result<usize> {
1906 loop {
1907 match self.handle.try_write(buf) {
1908 Err(Error::WouldBlock) => self.wait_writable()?,
1909 other => return other,
1910 }
1911 }
1912 }
1913
1914 pub fn blocking_write_all(
1917 &self,
1918 mut buf: &[u8],
1919 ) -> Result<()> {
1920 while !buf.is_empty() {
1921 let written = self.blocking_write(buf)?;
1922 buf = &buf[written..];
1923 }
1924 Ok(())
1925 }
1926
1927 pub fn finish(&mut self) -> Result<()> {
1931 self.shut = true;
1932 self.handle.set_flag(self.handle.side.write_direction(), FLAG_FIN);
1933 Ok(())
1934 }
1935
1936 pub fn finish_quietly(&mut self) {
1943 self.shut = true;
1944 self.handle
1945 .set_flag_quietly(self.handle.side.write_direction(), FLAG_FIN);
1946 }
1947
1948 pub fn reset(&mut self) {
1951 self.shut = true;
1952 self.handle.set_flag(self.handle.side.write_direction(), FLAG_RESET);
1953 }
1954
1955 #[cfg(any(
1959 target_os = "linux",
1960 target_os = "freebsd",
1961 target_os = "macos"
1962 ))]
1963 pub fn wait_writable_timeout(
1964 &self,
1965 timeout: Duration,
1966 ) -> Result<bool> {
1967 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1968 self.handle.wait_until_timeout(
1969 self.handle.side.write_direction(),
1970 |direction| {
1971 clear_exact_wait(direction);
1972 writable(direction, buffer_bytes)
1973 },
1974 timeout,
1975 )
1976 }
1977
1978 pub fn wait_writable(&self) -> Result<()> {
1980 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1981 self.handle.wait_until(
1982 self.handle.side.write_direction(),
1983 |direction| {
1984 clear_exact_wait(direction);
1985 writable(direction, buffer_bytes)
1986 },
1987 )
1988 }
1989
1990 pub fn wait_writable_exact(
1992 &self,
1993 needed: usize,
1994 ) -> Result<()> {
1995 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1996 if needed > buffer_bytes {
1997 return Err(Error::Malformed(format!(
1998 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
1999 )));
2000 }
2001 self.handle.wait_until(
2002 self.handle.side.write_direction(),
2003 |direction| {
2004 clear_exact_wait(direction);
2005 writable_exact(direction, buffer_bytes, needed)
2006 },
2007 )
2008 }
2009
2010 pub fn poll_writable(
2015 &self,
2016 cx: &mut std::task::Context<'_>,
2017 ) -> std::task::Poll<Result<()>> {
2018 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2019 self.handle.poll_ready(
2020 Interest {
2021 direction: self.handle.side.write_direction(),
2022 writer: true,
2023 reset: false,
2024 },
2025 |direction| {
2026 clear_exact_wait(direction);
2027 writable(direction, buffer_bytes)
2028 },
2029 cx,
2030 )
2031 }
2032
2033 pub fn poll_writable_exact(
2035 &self,
2036 needed: usize,
2037 cx: &mut std::task::Context<'_>,
2038 ) -> std::task::Poll<Result<()>> {
2039 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
2040 let slot = match self.handle.slot() {
2041 Ok(slot) => slot,
2042 Err(error) => return std::task::Poll::Ready(Err(error)),
2043 };
2044 let direction =
2045 Handle::direction(slot, self.handle.side.write_direction());
2046 if needed > buffer_bytes {
2047 clear_exact_wait(direction);
2048 return std::task::Poll::Ready(Err(Error::Malformed(format!(
2049 "atomic stream write is {needed} bytes; ring capacity is {buffer_bytes}"
2050 ))));
2051 }
2052 direction.exact_needed.store(needed as u32, Ordering::SeqCst);
2056 let result = self.handle.poll_ready(
2057 Interest {
2058 direction: self.handle.side.write_direction(),
2059 writer: true,
2060 reset: false,
2061 },
2062 |direction| writable_exact(direction, buffer_bytes, needed),
2063 cx,
2064 );
2065 if result.is_ready() {
2066 clear_exact_wait(direction);
2067 }
2068 result
2069 }
2070}
2071
2072impl Drop for WriteHalf {
2073 fn drop(&mut self) {
2074 if !self.shut {
2075 self.handle
2076 .set_flag(self.handle.side.write_direction(), FLAG_RESET);
2077 }
2078 }
2079}
2080
2081#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2085pub(crate) fn waits_supported() -> bool {
2086 orbit_core::sync::supported()
2087}
2088
2089#[cfg(not(any(
2090 target_os = "linux",
2091 target_os = "freebsd",
2092 target_os = "macos"
2093)))]
2094pub(crate) fn waits_supported() -> bool {
2095 false
2096}
2097
2098#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2102#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
2105pub(crate) fn wait_on_timeout(
2106 word: &AtomicU32,
2107 expected: u32,
2108 timeout: Duration,
2109) -> Result<bool> {
2110 orbit_core::sync::wait_word_timeout(word, expected, timeout)
2111 .map_err(Error::Io)
2112}
2113
2114pub(crate) fn wait_on(
2115 word: &AtomicU32,
2116 expected: u32,
2117) -> Result<()> {
2118 orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
2119}
2120
2121#[cfg(not(any(
2122 target_os = "linux",
2123 target_os = "freebsd",
2124 target_os = "macos"
2125)))]
2126pub(crate) fn wait_on(
2127 _word: &AtomicU32,
2128 _expected: u32,
2129) -> Result<()> {
2130 Err(Error::Io(std::io::Error::new(
2131 std::io::ErrorKind::Unsupported,
2132 "orbit-stream needs a platform that can wait on a shared word",
2133 )))
2134}
2135
2136pub(crate) fn wake_on(word: &AtomicU32) {
2138 #[cfg(any(
2139 target_os = "linux",
2140 target_os = "freebsd",
2141 target_os = "macos"
2142 ))]
2143 let _ = orbit_core::sync::wake_word(word);
2144 #[cfg(not(any(
2145 target_os = "linux",
2146 target_os = "freebsd",
2147 target_os = "macos"
2148 )))]
2149 let _ = word;
2150}
2151
2152pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
2153 mutex.lock().unwrap_or_else(|error| error.into_inner())
2154}
2155
2156#[cfg(test)]
2157mod tests {
2158 use std::sync::Arc;
2159 use std::time::Duration;
2160
2161 use orbit_core::{Fleet, NodeId};
2162
2163 use super::{
2164 Error, Incarnation, STREAM_BUFFER_BYTES, STREAM_LANE_CAPACITY, Side,
2165 Streams, Ticket,
2166 };
2167
2168 fn streams(name: &'static str) -> Streams {
2169 Streams::new(
2170 Arc::new(Fleet::join(name, 2).unwrap()),
2171 Incarnation::new(1),
2172 )
2173 .unwrap()
2174 }
2175
2176 #[test]
2177 fn bytes_cross_in_order_and_end_cleanly() {
2178 let streams = streams("stream-order");
2179 let (mut a, ticket) = streams.create().unwrap();
2180 assert_eq!(ticket.side, Side::B);
2181 assert_eq!(ticket.id.node(), 0);
2182 assert_eq!(ticket.epoch, streams.epoch());
2183 let b = streams.open(ticket).unwrap();
2184
2185 a.blocking_write_all(b"hello, ").unwrap();
2186 a.blocking_write_all(b"world").unwrap();
2187 a.finish().unwrap();
2188
2189 let mut out = Vec::new();
2190 let mut buf = [0_u8; 4];
2191 loop {
2192 let n = b.blocking_read(&mut buf).unwrap();
2193 if n == 0 {
2194 break;
2195 }
2196 out.extend_from_slice(&buf[..n]);
2197 }
2198 assert_eq!(out, b"hello, world");
2199 b.blocking_write_all(b"back").unwrap();
2201 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"back");
2202 }
2203
2204 #[test]
2205 fn a_ticket_prints_and_parses() {
2206 let streams = streams("stream-ticket");
2207 let (_a, ticket) = streams.create().unwrap();
2208 let text = ticket.to_string();
2209 assert!(text.contains("/b/"), "{text}");
2210 assert_eq!(text.parse::<Ticket>().unwrap(), ticket);
2211 assert!(matches!("nope".parse::<Ticket>(), Err(Error::Malformed(_))));
2212 assert!(matches!("x/b".parse::<Ticket>(), Err(Error::Malformed(_))));
2213 }
2214
2215 #[test]
2216 fn a_side_is_held_once_and_a_released_stream_goes_stale() {
2217 let streams = streams("stream-claim");
2218 let (a, ticket) = streams.create().unwrap();
2219 let b = streams.open(ticket).unwrap();
2220 assert!(matches!(streams.open(ticket), Err(Error::AlreadyClaimed(_))));
2221 assert!(streams.is_live(ticket.id));
2222
2223 drop(a);
2224 let mut buf = [0_u8; 1];
2226 assert!(matches!(b.blocking_read(&mut buf), Err(Error::Reset)));
2227 drop(b);
2228 assert!(!streams.is_live(ticket.id));
2229 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2230
2231 for _ in 0..STREAM_LANE_CAPACITY - 1 {
2234 drop(streams.create().unwrap());
2235 }
2236 let (_a2, ticket2) = streams.create().unwrap();
2237 assert_eq!(ticket2.id.slot(), ticket.id.slot());
2238 assert_ne!(ticket2.id.generation(), ticket.id.generation());
2239 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2240 }
2241
2242 #[test]
2243 fn a_reset_starts_a_new_epoch_and_refuses_old_tickets() {
2244 let streams = streams("stream-epoch");
2245 let (_a, ticket) = streams.create().unwrap();
2246 let before = streams.epoch();
2247 streams.reset_all();
2248 assert!(streams.epoch() > before);
2249 let (_a2, ticket2) = streams.create().unwrap();
2252 assert_eq!(ticket2.id, ticket.id);
2253 assert_ne!(ticket2.epoch, ticket.epoch);
2254 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
2255 assert!(streams.open(ticket2).is_ok());
2256 }
2257
2258 #[test]
2259 fn a_slot_whose_generation_runs_out_is_not_reused() {
2260 let streams = streams("stream-exhaust");
2263 let limit = super::layout::GENERATION_LIMIT as usize;
2264 for _ in 0..STREAM_LANE_CAPACITY * limit {
2265 drop(streams.create().unwrap());
2266 }
2267 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2268 streams.reset_all();
2269 let (_a, ticket) = streams.create().unwrap();
2270 assert_eq!(ticket.id.generation(), 1);
2271 }
2272
2273 #[test]
2274 fn a_full_ring_blocks_the_writer_and_wraps_correctly() {
2275 let streams = streams("stream-wrap");
2276 let (a, ticket) = streams.create().unwrap();
2277 let b = streams.open(ticket).unwrap();
2278
2279 let filler = vec![7_u8; STREAM_BUFFER_BYTES];
2280 assert_eq!(a.try_write(&filler).unwrap(), STREAM_BUFFER_BYTES);
2281 assert!(matches!(a.try_write(b"x"), Err(Error::WouldBlock)));
2282
2283 let mut buf = vec![0_u8; STREAM_BUFFER_BYTES - 5];
2284 assert_eq!(b.blocking_read(&mut buf).unwrap(), STREAM_BUFFER_BYTES - 5);
2285 let tail =
2287 (0..STREAM_BUFFER_BYTES).map(|i| i as u8).collect::<Vec<_>>();
2288 assert_eq!(a.try_write(&tail).unwrap(), STREAM_BUFFER_BYTES - 5);
2289
2290 let mut rest = vec![0_u8; 5];
2291 assert_eq!(b.blocking_read(&mut rest).unwrap(), 5);
2292 assert_eq!(rest, vec![7_u8; 5]);
2293 let mut wrapped = vec![0_u8; STREAM_BUFFER_BYTES];
2294 let mut got = 0;
2295 while got < STREAM_BUFFER_BYTES - 5 {
2296 got += b.blocking_read(&mut wrapped[got..]).unwrap();
2297 }
2298 assert_eq!(&wrapped[..got], &tail[..got]);
2299 }
2300
2301 #[test]
2302 fn exact_frames_never_publish_or_consume_a_prefix() {
2303 let streams = streams("stream-exact-frame");
2304 let (a, ticket) = streams.create().unwrap();
2305 let b = streams.open(ticket).unwrap();
2306
2307 let filler = vec![7_u8; STREAM_BUFFER_BYTES - 32];
2308 a.try_write_exact(&filler).unwrap();
2309 assert!(matches!(
2310 a.try_write_exact(&[9_u8; 64]),
2311 Err(Error::WouldBlock)
2312 ));
2313 let mut consumed = vec![0_u8; filler.len()];
2314 b.try_read_exact(&mut consumed).unwrap();
2315 assert_eq!(consumed, filler);
2316 assert!(matches!(b.try_read(&mut [0_u8; 1]), Err(Error::WouldBlock)));
2317
2318 a.try_write_exact(&[3_u8; 32]).unwrap();
2319 assert!(matches!(
2320 b.try_read_exact(&mut [0_u8; 64]),
2321 Err(Error::WouldBlock)
2322 ));
2323 let mut frame = [0_u8; 32];
2324 b.try_read_exact(&mut frame).unwrap();
2325 assert_eq!(frame, [3_u8; 32]);
2326 }
2327
2328 #[test]
2329 fn a_blocked_reader_is_woken_by_a_write_in_another_thread() {
2330 let streams = streams("stream-wake");
2331 let (a, ticket) = streams.create().unwrap();
2332 let b = streams.open(ticket).unwrap();
2333 let reader = std::thread::spawn(move || {
2334 let mut buf = [0_u8; 8];
2335 let n = b.blocking_read(&mut buf).unwrap();
2336 buf[..n].to_vec()
2337 });
2338 std::thread::sleep(Duration::from_millis(30));
2339 a.blocking_write_all(b"ping").unwrap();
2340 assert_eq!(reader.join().unwrap(), b"ping");
2341 }
2342
2343 #[test]
2344 fn a_reader_gone_fails_the_writer_and_a_lane_fills_up() {
2345 let streams = streams("stream-gone");
2346 let (a, ticket) = streams.create().unwrap();
2347 let b = streams.open(ticket).unwrap();
2348 let (b_read, _b_write) = b.split();
2349 drop(b_read);
2350 assert!(matches!(a.blocking_write(b"nobody"), Err(Error::PeerGone)));
2351
2352 let mut held = vec![a];
2353 for _ in 1..STREAM_LANE_CAPACITY {
2354 held.push(streams.create().unwrap().0);
2355 }
2356 assert!(matches!(streams.create(), Err(Error::Full { .. })));
2357 }
2358
2359 #[test]
2360 fn a_death_report_ends_that_incarnations_sides_only() {
2361 let streams = streams("stream-dead");
2362 let (a, ticket) = streams.create().unwrap();
2363 let b = streams.open(ticket).unwrap();
2364 b.blocking_write_all(b"partial").unwrap();
2365
2366 streams.node_dead(NodeId::ZERO, Incarnation::new(2));
2368 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"partial");
2369
2370 let reader = std::thread::spawn(move || {
2376 let mut buf = [0_u8; 8];
2377 a.blocking_read(&mut buf).map(|_| ())
2378 });
2379 std::thread::sleep(Duration::from_millis(30));
2380 streams.node_dead(NodeId::ZERO, Incarnation::new(1));
2381 assert!(matches!(
2382 reader.join().unwrap(),
2383 Err(Error::Reset | Error::Stale(_))
2384 ));
2385 assert!(matches!(
2386 b.try_write(b"x"),
2387 Err(Error::Reset | Error::Stale(_))
2388 ));
2389 assert!(matches!(
2390 streams.open(ticket),
2391 Err(Error::AlreadyClaimed(_) | Error::Stale(_))
2392 ));
2393 drop(b);
2394 assert!(!streams.is_live(ticket.id));
2395 }
2396
2397 #[test]
2398 fn an_offer_is_taken_once_and_a_stale_offer_is_skipped() {
2399 let streams = streams("stream-offer");
2400 assert!(streams.take_offer().is_none());
2401 let (a, ticket) = streams.create().unwrap();
2402 streams.offer(ticket, NodeId::ZERO).unwrap();
2403 assert_eq!(streams.take_offer(), Some(ticket));
2404 assert!(streams.take_offer().is_none());
2405
2406 let (gone, stale) = streams.create().unwrap();
2407 streams.offer(stale, NodeId::ZERO).unwrap();
2408 drop(gone);
2409 assert!(streams.take_offer().is_none());
2410
2411 let taker = streams.clone();
2412 let waiter = std::thread::spawn(move || taker.blocking_take_offer());
2413 std::thread::sleep(Duration::from_millis(30));
2414 streams.offer(ticket, NodeId::ZERO).unwrap();
2415 assert_eq!(waiter.join().unwrap().unwrap(), ticket);
2416 drop(a);
2417 }
2418
2419 #[test]
2420 fn a_contested_claim_goes_to_exactly_one_opener() {
2421 let streams = streams("stream-race");
2422 for _ in 0..50 {
2423 let (_a, ticket) = streams.create().unwrap();
2424 let openers = (0..4)
2425 .map(|_| {
2426 let streams = streams.clone();
2427 std::thread::spawn(move || streams.open(ticket).is_ok())
2428 })
2429 .collect::<Vec<_>>();
2430 let won = openers
2431 .into_iter()
2432 .map(|opener| opener.join().unwrap())
2433 .filter(|won| *won)
2434 .count();
2435 assert_eq!(won, 1);
2436 }
2437 }
2438
2439 #[test]
2440 fn random_sizes_concatenate_exactly() {
2441 let streams = streams("stream-random");
2442 let (mut a, ticket) = streams.create().unwrap();
2443 let b = streams.open(ticket).unwrap();
2444 let payload = (0..300_000_u32)
2445 .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
2446 .collect::<Vec<_>>();
2447 let expected = payload.clone();
2448 let writer = std::thread::spawn(move || {
2449 let mut rest = &payload[..];
2450 let mut step = 1_usize;
2451 while !rest.is_empty() {
2452 let len = step.min(rest.len());
2453 a.blocking_write_all(&rest[..len]).unwrap();
2454 rest = &rest[len..];
2455 step = (step * 7 + 13) % 9_001 + 1;
2456 }
2457 a.finish().unwrap();
2458 });
2459 let mut out = Vec::new();
2460 let mut step = 3_usize;
2461 loop {
2462 let chunk = b.blocking_read_chunk(step).unwrap();
2463 if chunk.is_empty() {
2464 break;
2465 }
2466 out.extend_from_slice(&chunk);
2467 step = (step * 5 + 11) % 4_099 + 1;
2468 }
2469 writer.join().unwrap();
2470 assert_eq!(out, expected);
2471 }
2472}