1use std::fmt;
28use std::io;
29use std::str::FromStr;
30use std::sync::atomic::{AtomicU32, Ordering};
31use std::time::{Duration, Instant};
32use std::sync::{Arc, Mutex, MutexGuard};
33
34use bytes::{Bytes, BytesMut};
35use orbit_core::{Fleet, NetId64};
36
37mod layout;
38#[cfg(feature = "tokio")]
39mod poll;
40mod table;
41mod wake;
42
43use layout::{
44 Direction, FLAG_FIN, FLAG_READER_GONE, FLAG_RESET, GENERATION_MASK, SIDE_CLAIMED, SIDE_FREE,
45 SIDE_RELEASED, SLOT_BITS, SLOT_EMPTY, SLOT_LIVE, SLOT_MASK, Slot,
46};
47use orbit_core::NodeId;
48use table::Table;
49#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
50pub use orbit_core::readiness::Readiness;
51pub use table::{segment_size, segment_size_for};
52use wake::Interest;
53
54pub const STREAM_KIND: u8 = 246;
57pub const STREAM_LANE_CAPACITY: usize =
65 orbit_core::compile::usize_from_env(option_env!("ORBIT_STREAM_LANE_CAPACITY"), 256);
66pub const STREAM_BUFFER_BYTES: usize =
73 orbit_core::compile::usize_from_env(option_env!("ORBIT_STREAM_BUFFER_BYTES"), 64 * 1024);
74
75#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
87pub struct StreamSpec {
88 pub kind: u8,
90 pub lane_capacity: usize,
93 pub buffer_bytes: usize,
95}
96
97impl StreamSpec {
98 pub const DEFAULT: Self = Self::new(STREAM_KIND, STREAM_LANE_CAPACITY, STREAM_BUFFER_BYTES);
99
100 pub const fn new(kind: u8, lane_capacity: usize, buffer_bytes: usize) -> Self {
101 Self { kind, lane_capacity, buffer_bytes }
102 }
103
104 pub const fn lane_bytes(&self) -> usize {
108 self.lane_capacity * 2 * self.buffer_bytes
109 }
110
111 fn validate(self) -> Result<()> {
114 if self.lane_capacity == 0
115 || self.buffer_bytes == 0
116 || !self.lane_capacity.is_power_of_two()
117 || !self.buffer_bytes.is_power_of_two()
118 || self.lane_capacity > 1 << SLOT_BITS
119 || self.buffer_bytes > u32::MAX as usize
120 {
121 return Err(Error::Malformed(format!(
122 "stream spec kind={} lane_capacity={} buffer_bytes={}: both are powers of two, a lane holds at most {} streams and a ring at most {} bytes",
123 self.kind,
124 self.lane_capacity,
125 self.buffer_bytes,
126 1_usize << SLOT_BITS,
127 u32::MAX
128 )));
129 }
130 Ok(())
131 }
132}
133
134impl Default for StreamSpec {
135 fn default() -> Self {
136 Self::DEFAULT
137 }
138}
139
140pub type Result<T, E = Error> = std::result::Result<T, E>;
141
142#[non_exhaustive]
146#[derive(Debug)]
147pub enum Error {
148 Stale(StreamId),
150 AlreadyClaimed(Ticket),
152 Full {
154 capacity: usize,
155 },
156 Malformed(String),
158 Closed,
160 Reset,
162 PeerGone,
164 WouldBlock,
166 Io(io::Error),
167}
168
169impl fmt::Display for Error {
170 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
171 match self {
172 Self::Stale(id) => write!(f, "stream {id} has ended"),
173 Self::AlreadyClaimed(ticket) => write!(f, "stream side {ticket} is already held"),
174 Self::Full { capacity } => write!(f, "stream lane is full: capacity={capacity}"),
175 Self::Malformed(text) => write!(f, "not a stream address: {text:?}"),
176 Self::Closed => f.write_str("stream write half is shut down"),
177 Self::Reset => f.write_str("stream direction was reset by the peer"),
178 Self::PeerGone => f.write_str("stream peer reader is gone"),
179 Self::WouldBlock => f.write_str("stream operation would block"),
180 Self::Io(error) => write!(f, "Orbit stream io error: {error}"),
181 }
182 }
183}
184
185impl std::error::Error for Error {
186 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
187 match self {
188 Self::Io(error) => Some(error),
189 _ => None,
190 }
191 }
192}
193
194impl From<io::Error> for Error {
195 fn from(value: io::Error) -> Self {
196 Self::Io(value)
197 }
198}
199
200impl From<Error> for io::Error {
201 fn from(value: Error) -> Self {
202 let kind = match &value {
203 Error::Stale(_) | Error::AlreadyClaimed(_) => io::ErrorKind::NotConnected,
204 Error::Full { .. } => io::ErrorKind::OutOfMemory,
205 Error::Malformed(_) => io::ErrorKind::InvalidInput,
206 Error::Closed => io::ErrorKind::BrokenPipe,
207 Error::Reset => io::ErrorKind::ConnectionReset,
208 Error::PeerGone => io::ErrorKind::BrokenPipe,
209 Error::WouldBlock => io::ErrorKind::WouldBlock,
210 Error::Io(error) => error.kind(),
211 };
212 io::Error::new(kind, value)
213 }
214}
215
216#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
219pub struct StreamId(NetId64);
220
221impl StreamId {
222 pub const fn from_net_id(id: NetId64) -> Self {
223 Self(id)
224 }
225
226 pub const fn net_id(self) -> NetId64 {
227 self.0
228 }
229
230 pub const fn kind(self) -> u8 {
231 self.0.kind()
232 }
233
234 pub const fn node(self) -> u16 {
236 self.0.node()
237 }
238
239 pub const fn slot(self) -> u32 {
240 (self.0.counter() & SLOT_MASK) as u32
241 }
242
243 pub const fn generation(self) -> u32 {
244 (self.0.counter() >> SLOT_BITS) as u32
245 }
246
247 fn make(kind: u8, node: u16, slot: u32, generation: u32) -> Self {
248 Self(NetId64::make(
249 kind,
250 node,
251 ((generation as u64) << SLOT_BITS) | (slot as u64 & SLOT_MASK),
252 ))
253 }
254}
255
256impl fmt::Display for StreamId {
257 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
258 self.0.fmt(f)
259 }
260}
261
262impl FromStr for StreamId {
263 type Err = Error;
264
265 fn from_str(text: &str) -> Result<Self> {
266 text.parse::<NetId64>()
267 .map(Self)
268 .map_err(|_| Error::Malformed(text.to_owned()))
269 }
270}
271
272#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
274pub enum Side {
275 A,
276 B,
277}
278
279impl Side {
280 const fn index(self) -> usize {
281 match self {
282 Self::A => 0,
283 Self::B => 1,
284 }
285 }
286
287 const fn write_direction(self) -> usize {
289 self.index()
290 }
291
292 const fn read_direction(self) -> usize {
294 1 - self.index()
295 }
296
297 const fn letter(self) -> char {
298 match self {
299 Self::A => 'a',
300 Self::B => 'b',
301 }
302 }
303}
304
305#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
314pub struct Ticket {
315 pub id: StreamId,
316 pub side: Side,
317 pub epoch: u64,
318}
319
320impl fmt::Display for Ticket {
321 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
322 write!(f, "{}/{}/{}", self.id, self.side.letter(), self.epoch)
323 }
324}
325
326impl FromStr for Ticket {
327 type Err = Error;
328
329 fn from_str(text: &str) -> Result<Self> {
330 let malformed = || Error::Malformed(text.to_owned());
331 let (rest, epoch) = text.rsplit_once('/').ok_or_else(malformed)?;
332 let (id, side) = rest.rsplit_once('/').ok_or_else(malformed)?;
333 let side = match side {
334 "a" => Side::A,
335 "b" => Side::B,
336 _ => return Err(malformed()),
337 };
338 Ok(Self {
339 id: id.parse().map_err(|_| malformed())?,
340 side,
341 epoch: epoch.parse().map_err(|_| malformed())?,
342 })
343 }
344}
345
346#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
352pub struct Incarnation(u64);
353
354impl Incarnation {
355 pub const fn new(value: u64) -> Self {
356 Self(value)
357 }
358
359 pub const fn get(self) -> u64 {
360 self.0
361 }
362}
363
364#[derive(Clone)]
367pub struct Streams {
368 table: Arc<Table>,
369}
370
371impl Streams {
372 pub fn new(fleet: Arc<Fleet>, incarnation: Incarnation) -> Result<Self> {
376 Self::with_spec(fleet, incarnation, StreamSpec::DEFAULT)
377 }
378
379 pub fn with_spec(
384 fleet: Arc<Fleet>,
385 incarnation: Incarnation,
386 spec: StreamSpec,
387 ) -> Result<Self> {
388 spec.validate()?;
389 Ok(Self {
390 table: table::open(&fleet, incarnation, spec)?,
391 })
392 }
393
394 pub fn kind(&self) -> u8 {
396 self.table.kind()
397 }
398
399 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
411 pub fn readiness(&self) -> Result<Readiness> {
412 self.table.take_readiness()
413 }
414
415 pub fn epoch(&self) -> u64 {
417 self.table.epoch()
418 }
419
420 pub fn create(&self) -> Result<(Endpoint, Ticket)> {
423 let (index, generation) = self.table.allocate()?;
424 let id = StreamId::make(
425 self.table.kind(),
426 self.table.node(),
427 (index % self.table.geometry().lane_capacity) as u32,
428 generation,
429 );
430 let handle = Arc::new(Handle {
431 table: Arc::clone(&self.table),
432 index,
433 id,
434 side: Side::A,
435 });
436 Ok((
437 Endpoint::new(handle),
438 Ticket {
439 id,
440 side: Side::B,
441 epoch: self.table.epoch(),
442 },
443 ))
444 }
445
446 pub fn open(&self, ticket: Ticket) -> Result<Endpoint> {
448 let index = self.locate(ticket.id)?;
449 if ticket.epoch != self.table.epoch() {
450 return Err(Error::Stale(ticket.id));
451 }
452 let slot = &self.table.slots()[index];
453 if !slot.is(ticket.id.generation()) {
454 return Err(Error::Stale(ticket.id));
455 }
456 let claimed = &slot.claimed[ticket.side.index()];
457 if claimed
458 .compare_exchange(SIDE_FREE, SIDE_CLAIMED, Ordering::SeqCst, Ordering::SeqCst)
459 .is_err()
460 {
461 return Err(Error::AlreadyClaimed(ticket));
462 }
463 if !slot.is(ticket.id.generation()) {
466 let _ = claimed.compare_exchange(
467 SIDE_CLAIMED,
468 SIDE_FREE,
469 Ordering::SeqCst,
470 Ordering::SeqCst,
471 );
472 return Err(Error::Stale(ticket.id));
473 }
474 slot.incarnation[ticket.side.index()].store(self.table.incarnation(), Ordering::Release);
475 slot.node[ticket.side.index()].store(self.table.node(), Ordering::Release);
476 for direction in &slot.directions {
478 self.table.notify(index, slot, direction, true);
479 }
480 Ok(Endpoint::new(Arc::new(Handle {
481 table: Arc::clone(&self.table),
482 index,
483 id: ticket.id,
484 side: ticket.side,
485 })))
486 }
487
488 pub fn is_live(&self, id: StreamId) -> bool {
490 self.locate(id)
491 .map(|index| self.table.slots()[index].is(id.generation()))
492 .unwrap_or(false)
493 }
494
495 pub fn offer(&self, ticket: Ticket, to: NodeId) -> Result<()> {
501 let index = self.locate(ticket.id)?;
502 if ticket.epoch != self.table.epoch()
503 || !self.table.slots()[index].is(ticket.id.generation())
504 {
505 return Err(Error::Stale(ticket.id));
506 }
507 self.table.offer(index, to.get())
508 }
509
510 pub fn take_offer(&self) -> Option<Ticket> {
513 while let Some(index) = self.table.take_offer() {
514 if let Some(ticket) = self.ticket_for(index) {
515 return Some(ticket);
516 }
517 }
518 None
519 }
520
521 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
525 pub fn take_offer_timeout(&self, timeout: Duration) -> Result<Option<Ticket>> {
526 let deadline = Instant::now() + timeout;
527 loop {
528 if let Some(ticket) = self.take_offer() {
529 return Ok(Some(ticket));
530 }
531 let left = deadline.saturating_duration_since(Instant::now());
532 if left.is_zero() {
533 return Ok(self.take_offer());
534 }
535 if !self.table.wait_offer_timeout(left)? {
536 return Ok(self.take_offer());
537 }
538 }
539 }
540
541 pub fn blocking_take_offer(&self) -> Result<Ticket> {
543 loop {
544 if let Some(ticket) = self.take_offer() {
545 return Ok(ticket);
546 }
547 self.table.wait_offer()?;
548 }
549 }
550
551 pub fn poll_take_offer(
554 &self,
555 cx: &mut std::task::Context<'_>,
556 ) -> std::task::Poll<Result<Ticket>> {
557 if let Some(ticket) = self.take_offer() {
558 return std::task::Poll::Ready(Ok(ticket));
559 }
560 if let Err(error) = self.table.register_offer(cx.waker()) {
561 return std::task::Poll::Ready(Err(error));
562 }
563 match self.take_offer() {
564 Some(ticket) => std::task::Poll::Ready(Ok(ticket)),
565 None => std::task::Poll::Pending,
566 }
567 }
568
569 pub fn owners(&self) -> Vec<(NodeId, Incarnation)> {
575 let mut seen: Vec<(NodeId, Incarnation)> = Vec::new();
576 for slot in self.table.slots() {
577 if slot.state.load(Ordering::Acquire) != SLOT_LIVE {
578 continue;
579 }
580 for side in 0..2 {
581 if slot.claimed[side].load(Ordering::Acquire) != SIDE_CLAIMED {
582 continue;
583 }
584 let holder = (
585 NodeId::new(slot.node[side].load(Ordering::Acquire)),
586 Incarnation::new(slot.incarnation[side].load(Ordering::Acquire)),
587 );
588 if !seen.contains(&holder) {
589 seen.push(holder);
590 }
591 }
592 }
593 seen.sort_by_key(|(node, incarnation)| (node.get(), incarnation.get()));
594 seen
595 }
596
597 pub fn node_dead(&self, node: NodeId, incarnation: Incarnation) {
604 self.table.node_dead(node.get(), incarnation.get());
605 }
606
607 pub fn reset_all(&self) {
610 self.table.reset_all();
611 }
612
613 fn ticket_for(&self, index: usize) -> Option<Ticket> {
614 let slot = &self.table.slots()[index];
615 let generation = slot.generation.load(Ordering::Acquire);
616 if !slot.is(generation) {
617 return None;
618 }
619 let lane_capacity = self.table.geometry().lane_capacity;
620 let id = StreamId::make(
621 self.table.kind(),
622 (index / lane_capacity) as u16,
623 (index % lane_capacity) as u32,
624 generation,
625 );
626 Some(Ticket {
627 id,
628 side: Side::B,
629 epoch: self.table.epoch(),
630 })
631 }
632
633 #[cfg(unix)]
636 pub fn unlink(&self) -> Result<()> {
637 self.table.unlink()
638 }
639
640 fn locate(&self, id: StreamId) -> Result<usize> {
641 let geometry = self.table.geometry();
642 if id.kind() != self.table.kind()
643 || usize::from(id.node()) >= geometry.fleet_capacity
644 || id.slot() as usize >= geometry.lane_capacity
645 || id.generation() == 0
646 || id.generation() > GENERATION_MASK
647 {
648 return Err(Error::Malformed(id.to_string()));
649 }
650 Ok(usize::from(id.node()) * geometry.lane_capacity + id.slot() as usize)
651 }
652}
653
654struct Handle {
656 table: Arc<Table>,
657 index: usize,
658 id: StreamId,
659 side: Side,
660}
661
662impl Handle {
663 fn slot(&self) -> Result<&Slot> {
664 let slot = &self.table.slots()[self.index];
665 if slot.is(self.id.generation()) {
666 Ok(slot)
667 } else {
668 Err(Error::Stale(self.id))
669 }
670 }
671
672 fn direction(slot: &Slot, direction: usize) -> &Direction {
673 &slot.directions[direction]
674 }
675
676 fn try_write(&self, buf: &[u8]) -> Result<usize> {
678 let slot = self.slot()?;
679 let direction = Self::direction(slot, self.side.write_direction());
680 let flags = direction.flags();
681 if flags & FLAG_RESET != 0 {
682 return Err(Error::Reset);
683 }
684 if flags & FLAG_FIN != 0 {
685 return Err(Error::Closed);
686 }
687 if flags & FLAG_READER_GONE != 0 {
688 return Err(Error::PeerGone);
689 }
690 let buffer_bytes = self.table.geometry().buffer_bytes;
691 let head = direction.head.load(Ordering::Relaxed);
692 let tail = direction.tail.load(Ordering::Acquire);
693 let free = buffer_bytes - (head - tail) as usize;
694 if buf.is_empty() {
695 return Ok(0);
696 }
697 if free == 0 {
698 return Err(Error::WouldBlock);
699 }
700 let len = buf.len().min(free);
701 let base = self.table.buffer(self.index, self.side.write_direction());
702 let start = (head as usize) & (buffer_bytes - 1);
703 let first = len.min(buffer_bytes - start);
704 unsafe {
707 std::ptr::copy_nonoverlapping(buf.as_ptr(), base.add(start), first);
708 std::ptr::copy_nonoverlapping(buf.as_ptr().add(first), base, len - first);
709 }
710 direction.head.store(head + len as u64, Ordering::SeqCst);
711 std::sync::atomic::fence(Ordering::SeqCst);
721 let drained = direction.tail.load(Ordering::SeqCst) >= head;
722 self.table.notify(self.index, slot, direction, drained);
723 Ok(len)
724 }
725
726 fn try_read(&self, buf: &mut [u8]) -> Result<usize> {
729 let slot = self.slot()?;
730 let direction = Self::direction(slot, self.side.read_direction());
731 let flags = direction.flags();
732 if flags & FLAG_RESET != 0 {
733 return Err(Error::Reset);
734 }
735 let head = direction.head.load(Ordering::Acquire);
736 let tail = direction.tail.load(Ordering::Relaxed);
737 let available = (head - tail) as usize;
738 if buf.is_empty() {
739 return Ok(0);
740 }
741 if available == 0 {
742 if flags & FLAG_FIN != 0 {
743 return Ok(0);
744 }
745 return Err(Error::WouldBlock);
746 }
747 let len = buf.len().min(available);
748 let base = self.table.buffer(self.index, self.side.read_direction());
749 let buffer_bytes = self.table.geometry().buffer_bytes;
750 let start = (tail as usize) & (buffer_bytes - 1);
751 let first = len.min(buffer_bytes - start);
752 unsafe {
756 std::ptr::copy_nonoverlapping(base.add(start), buf.as_mut_ptr(), first);
757 std::ptr::copy_nonoverlapping(base, buf.as_mut_ptr().add(first), len - first);
758 }
759 direction.tail.store(tail + len as u64, Ordering::SeqCst);
760 std::sync::atomic::fence(Ordering::SeqCst);
764 let head_now = direction.head.load(Ordering::SeqCst);
765 let was_full = head_now - tail >= buffer_bytes as u64;
766 self.table.notify(self.index, slot, direction, was_full);
767 Ok(len)
768 }
769
770 fn set_flag(&self, direction_index: usize, flag: u8) {
771 if let Ok(slot) = self.slot() {
772 let direction = Self::direction(slot, direction_index);
773 if direction.flags.fetch_or(flag, Ordering::SeqCst) & flag == 0 {
774 self.table.notify(self.index, slot, direction, true);
775 }
776 }
777 }
778
779 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
782 fn wait_until_timeout(
783 &self,
784 direction_index: usize,
785 ready: impl Fn(&Direction) -> bool,
786 timeout: Duration,
787 ) -> Result<bool> {
788 let deadline = Instant::now() + timeout;
789 loop {
790 let slot = self.slot()?;
791 let direction = Self::direction(slot, direction_index);
792 if ready(direction) {
793 return Ok(true);
794 }
795 let left = deadline.saturating_duration_since(Instant::now());
796 if left.is_zero() {
797 return Ok(false);
798 }
799 direction.waiters.fetch_add(1, Ordering::SeqCst);
800 let since = direction.changes.load(Ordering::SeqCst);
801 let outcome = if ready(direction) {
802 Ok(true)
803 } else {
804 wait_on_timeout(&direction.changes, since, left)
805 };
806 direction.waiters.fetch_sub(1, Ordering::SeqCst);
807 if !outcome? {
808 let slot = self.slot()?;
811 return Ok(ready(Self::direction(slot, direction_index)));
812 }
813 }
814 }
815
816 fn wait_until(&self, direction_index: usize, ready: impl Fn(&Direction) -> bool) -> Result<()> {
817 loop {
818 let slot = self.slot()?;
819 let direction = Self::direction(slot, direction_index);
820 if ready(direction) {
821 return Ok(());
822 }
823 direction.waiters.fetch_add(1, Ordering::SeqCst);
824 let since = direction.changes.load(Ordering::SeqCst);
825 let outcome = if ready(direction) {
826 Ok(())
827 } else {
828 wait_on(&direction.changes, since)
829 };
830 direction.waiters.fetch_sub(1, Ordering::SeqCst);
831 outcome?;
832 }
833 }
834
835 fn poll_ready(
838 &self,
839 interest: Interest,
840 ready: impl Fn(&Direction) -> bool,
841 cx: &mut std::task::Context<'_>,
842 ) -> std::task::Poll<Result<()>> {
843 let slot = match self.slot() {
844 Ok(slot) => slot,
845 Err(error) => return std::task::Poll::Ready(Err(error)),
846 };
847 let direction = Self::direction(slot, interest.direction);
848 if ready(direction) {
849 return std::task::Poll::Ready(Ok(()));
850 }
851 if let Err(error) = self.table.register(self.index, interest, cx.waker()) {
852 return std::task::Poll::Ready(Err(error));
853 }
854 if ready(direction) {
855 std::task::Poll::Ready(Ok(()))
856 } else {
857 std::task::Poll::Pending
858 }
859 }
860}
861
862fn readable(direction: &Direction) -> bool {
863 direction.head.load(Ordering::Acquire) != direction.tail.load(Ordering::Relaxed)
864 || direction.flags() & (FLAG_FIN | FLAG_RESET) != 0
865}
866
867fn writable(direction: &Direction, buffer_bytes: usize) -> bool {
868 let queued =
869 (direction.head.load(Ordering::Relaxed) - direction.tail.load(Ordering::Acquire)) as usize;
870 queued < buffer_bytes
871 || direction.flags() & (FLAG_FIN | FLAG_RESET | FLAG_READER_GONE) != 0
872}
873
874impl Drop for Handle {
875 fn drop(&mut self) {
879 let slot = &self.table.slots()[self.index];
880 if !slot.is(self.id.generation()) {
881 return;
882 }
883 if slot.claimed[self.side.index()]
886 .compare_exchange(
887 SIDE_CLAIMED,
888 SIDE_RELEASED,
889 Ordering::SeqCst,
890 Ordering::SeqCst,
891 )
892 .is_err()
893 {
894 return;
895 }
896 let other = slot.claimed[1 - self.side.index()].load(Ordering::SeqCst);
900 if other != SIDE_CLAIMED {
901 let _ = slot.state.compare_exchange(
902 SLOT_LIVE,
903 SLOT_EMPTY,
904 Ordering::SeqCst,
905 Ordering::SeqCst,
906 );
907 for direction in &slot.directions {
908 direction.changes.fetch_add(1, Ordering::SeqCst);
911 if direction.waiters.load(Ordering::SeqCst) > 0 {
912 wake_on(&direction.changes);
913 }
914 }
915 self.table.registry.wake(self.index);
916 }
917 }
918}
919
920pub struct Endpoint {
924 read: ReadHalf,
925 write: WriteHalf,
926}
927
928impl Endpoint {
929 fn new(handle: Arc<Handle>) -> Self {
930 Self {
931 read: ReadHalf {
932 handle: Arc::clone(&handle),
933 },
934 write: WriteHalf {
935 handle,
936 shut: false,
937 },
938 }
939 }
940
941 pub fn id(&self) -> StreamId {
942 self.read.handle.id
943 }
944
945 pub fn side(&self) -> Side {
946 self.read.handle.side
947 }
948
949 pub fn split(self) -> (ReadHalf, WriteHalf) {
950 (self.read, self.write)
951 }
952
953 pub fn try_read(&self, buf: &mut [u8]) -> Result<usize> {
954 self.read.try_read(buf)
955 }
956
957 pub fn blocking_read(&self, buf: &mut [u8]) -> Result<usize> {
958 self.read.blocking_read(buf)
959 }
960
961 pub fn blocking_read_chunk(&self, max: usize) -> Result<Bytes> {
962 self.read.blocking_read_chunk(max)
963 }
964
965 pub fn try_write(&self, buf: &[u8]) -> Result<usize> {
966 self.write.try_write(buf)
967 }
968
969 pub fn blocking_write(&self, buf: &[u8]) -> Result<usize> {
970 self.write.blocking_write(buf)
971 }
972
973 pub fn blocking_write_all(&self, buf: &[u8]) -> Result<()> {
974 self.write.blocking_write_all(buf)
975 }
976
977 pub fn finish(&mut self) -> Result<()> {
978 self.write.finish()
979 }
980
981 pub fn reset(&mut self) {
982 self.write.reset()
983 }
984}
985
986pub struct ReadHalf {
989 handle: Arc<Handle>,
990}
991
992impl ReadHalf {
993 pub fn id(&self) -> StreamId {
994 self.handle.id
995 }
996
997 pub fn try_read(&self, buf: &mut [u8]) -> Result<usize> {
1000 self.handle.try_read(buf)
1001 }
1002
1003 pub fn blocking_read(&self, buf: &mut [u8]) -> Result<usize> {
1006 loop {
1007 match self.handle.try_read(buf) {
1008 Err(Error::WouldBlock) => self.wait_readable()?,
1009 other => return other,
1010 }
1011 }
1012 }
1013
1014 pub fn blocking_read_chunk(&self, max: usize) -> Result<Bytes> {
1016 let mut chunk = BytesMut::zeroed(max);
1017 let len = self.blocking_read(&mut chunk)?;
1018 chunk.truncate(len);
1019 Ok(chunk.freeze())
1020 }
1021
1022 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1024 pub fn wait_readable_timeout(&self, timeout: Duration) -> Result<bool> {
1025 self.handle
1026 .wait_until_timeout(self.handle.side.read_direction(), readable, timeout)
1027 }
1028
1029 pub fn wait_readable(&self) -> Result<()> {
1031 self.handle
1032 .wait_until(self.handle.side.read_direction(), readable)
1033 }
1034
1035 pub fn poll_readable(&self, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<()>> {
1040 self.handle.poll_ready(
1041 Interest {
1042 direction: self.handle.side.read_direction(),
1043 writer: false,
1044 },
1045 readable,
1046 cx,
1047 )
1048 }
1049}
1050
1051impl Drop for ReadHalf {
1052 fn drop(&mut self) {
1053 self.handle
1054 .set_flag(self.handle.side.read_direction(), FLAG_READER_GONE);
1055 }
1056}
1057
1058pub struct WriteHalf {
1061 handle: Arc<Handle>,
1062 shut: bool,
1063}
1064
1065impl WriteHalf {
1066 pub fn id(&self) -> StreamId {
1067 self.handle.id
1068 }
1069
1070 pub fn try_write(&self, buf: &[u8]) -> Result<usize> {
1072 self.handle.try_write(buf)
1073 }
1074
1075 pub fn blocking_write(&self, buf: &[u8]) -> Result<usize> {
1077 loop {
1078 match self.handle.try_write(buf) {
1079 Err(Error::WouldBlock) => self.wait_writable()?,
1080 other => return other,
1081 }
1082 }
1083 }
1084
1085 pub fn blocking_write_all(&self, mut buf: &[u8]) -> Result<()> {
1088 while !buf.is_empty() {
1089 let written = self.blocking_write(buf)?;
1090 buf = &buf[written..];
1091 }
1092 Ok(())
1093 }
1094
1095 pub fn finish(&mut self) -> Result<()> {
1099 self.shut = true;
1100 self.handle
1101 .set_flag(self.handle.side.write_direction(), FLAG_FIN);
1102 Ok(())
1103 }
1104
1105 pub fn reset(&mut self) {
1108 self.shut = true;
1109 self.handle
1110 .set_flag(self.handle.side.write_direction(), FLAG_RESET);
1111 }
1112
1113 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1117 pub fn wait_writable_timeout(&self, timeout: Duration) -> Result<bool> {
1118 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1119 self.handle.wait_until_timeout(
1120 self.handle.side.write_direction(),
1121 |direction| writable(direction, buffer_bytes),
1122 timeout,
1123 )
1124 }
1125
1126 pub fn wait_writable(&self) -> Result<()> {
1128 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1129 self.handle.wait_until(self.handle.side.write_direction(), |direction| {
1130 writable(direction, buffer_bytes)
1131 })
1132 }
1133
1134 pub fn poll_writable(&self, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<()>> {
1139 let buffer_bytes = self.handle.table.geometry().buffer_bytes;
1140 self.handle.poll_ready(
1141 Interest {
1142 direction: self.handle.side.write_direction(),
1143 writer: true,
1144 },
1145 |direction| writable(direction, buffer_bytes),
1146 cx,
1147 )
1148 }
1149}
1150
1151impl Drop for WriteHalf {
1152 fn drop(&mut self) {
1153 if !self.shut {
1154 self.handle
1155 .set_flag(self.handle.side.write_direction(), FLAG_RESET);
1156 }
1157 }
1158}
1159
1160#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1164pub(crate) fn waits_supported() -> bool {
1165 orbit_core::sync::supported()
1166}
1167
1168#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
1169pub(crate) fn waits_supported() -> bool {
1170 false
1171}
1172
1173#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1177#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1180pub(crate) fn wait_on_timeout(word: &AtomicU32, expected: u32, timeout: Duration) -> Result<bool> {
1181 orbit_core::sync::wait_word_timeout(word, expected, timeout).map_err(Error::Io)
1182}
1183
1184pub(crate) fn wait_on(word: &AtomicU32, expected: u32) -> Result<()> {
1185 orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
1186}
1187
1188#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
1189pub(crate) fn wait_on(_word: &AtomicU32, _expected: u32) -> Result<()> {
1190 Err(Error::Io(std::io::Error::new(std::io::ErrorKind::Unsupported, "orbit-stream needs a platform that can wait on a shared word")))
1191}
1192
1193pub(crate) fn wake_on(word: &AtomicU32) {
1195 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
1196 let _ = orbit_core::sync::wake_word(word);
1197 #[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
1198 let _ = word;
1199}
1200
1201pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
1202 mutex.lock().unwrap_or_else(|error| error.into_inner())
1203}
1204
1205#[cfg(test)]
1206mod tests {
1207 use std::sync::Arc;
1208 use std::time::Duration;
1209
1210 use orbit_core::{Fleet, NodeId};
1211
1212 use super::{
1213 Error, Incarnation, STREAM_BUFFER_BYTES, STREAM_LANE_CAPACITY, Side, Streams, Ticket,
1214 };
1215
1216 fn streams(name: &'static str) -> Streams {
1217 Streams::new(Arc::new(Fleet::join(name, 2).unwrap()), Incarnation::new(1)).unwrap()
1218 }
1219
1220 #[test]
1221 fn bytes_cross_in_order_and_end_cleanly() {
1222 let streams = streams("stream-order");
1223 let (mut a, ticket) = streams.create().unwrap();
1224 assert_eq!(ticket.side, Side::B);
1225 assert_eq!(ticket.id.node(), 0);
1226 assert_eq!(ticket.epoch, streams.epoch());
1227 let b = streams.open(ticket).unwrap();
1228
1229 a.blocking_write_all(b"hello, ").unwrap();
1230 a.blocking_write_all(b"world").unwrap();
1231 a.finish().unwrap();
1232
1233 let mut out = Vec::new();
1234 let mut buf = [0_u8; 4];
1235 loop {
1236 let n = b.blocking_read(&mut buf).unwrap();
1237 if n == 0 {
1238 break;
1239 }
1240 out.extend_from_slice(&buf[..n]);
1241 }
1242 assert_eq!(out, b"hello, world");
1243 b.blocking_write_all(b"back").unwrap();
1245 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"back");
1246 }
1247
1248 #[test]
1249 fn a_ticket_prints_and_parses() {
1250 let streams = streams("stream-ticket");
1251 let (_a, ticket) = streams.create().unwrap();
1252 let text = ticket.to_string();
1253 assert!(text.contains("/b/"), "{text}");
1254 assert_eq!(text.parse::<Ticket>().unwrap(), ticket);
1255 assert!(matches!("nope".parse::<Ticket>(), Err(Error::Malformed(_))));
1256 assert!(matches!("x/b".parse::<Ticket>(), Err(Error::Malformed(_))));
1257 }
1258
1259 #[test]
1260 fn a_side_is_held_once_and_a_released_stream_goes_stale() {
1261 let streams = streams("stream-claim");
1262 let (a, ticket) = streams.create().unwrap();
1263 let b = streams.open(ticket).unwrap();
1264 assert!(matches!(
1265 streams.open(ticket),
1266 Err(Error::AlreadyClaimed(_))
1267 ));
1268 assert!(streams.is_live(ticket.id));
1269
1270 drop(a);
1271 let mut buf = [0_u8; 1];
1273 assert!(matches!(b.blocking_read(&mut buf), Err(Error::Reset)));
1274 drop(b);
1275 assert!(!streams.is_live(ticket.id));
1276 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
1277
1278 for _ in 0..STREAM_LANE_CAPACITY - 1 {
1281 drop(streams.create().unwrap());
1282 }
1283 let (_a2, ticket2) = streams.create().unwrap();
1284 assert_eq!(ticket2.id.slot(), ticket.id.slot());
1285 assert_ne!(ticket2.id.generation(), ticket.id.generation());
1286 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
1287 }
1288
1289 #[test]
1290 fn a_reset_starts_a_new_epoch_and_refuses_old_tickets() {
1291 let streams = streams("stream-epoch");
1292 let (_a, ticket) = streams.create().unwrap();
1293 let before = streams.epoch();
1294 streams.reset_all();
1295 assert!(streams.epoch() > before);
1296 let (_a2, ticket2) = streams.create().unwrap();
1299 assert_eq!(ticket2.id, ticket.id);
1300 assert_ne!(ticket2.epoch, ticket.epoch);
1301 assert!(matches!(streams.open(ticket), Err(Error::Stale(_))));
1302 assert!(streams.open(ticket2).is_ok());
1303 }
1304
1305 #[test]
1306 fn a_slot_whose_generation_runs_out_is_not_reused() {
1307 let streams = streams("stream-exhaust");
1310 let limit = super::layout::GENERATION_LIMIT as usize;
1311 for _ in 0..STREAM_LANE_CAPACITY * limit {
1312 drop(streams.create().unwrap());
1313 }
1314 assert!(matches!(streams.create(), Err(Error::Full { .. })));
1315 streams.reset_all();
1316 let (_a, ticket) = streams.create().unwrap();
1317 assert_eq!(ticket.id.generation(), 1);
1318 }
1319
1320 #[test]
1321 fn a_full_ring_blocks_the_writer_and_wraps_correctly() {
1322 let streams = streams("stream-wrap");
1323 let (a, ticket) = streams.create().unwrap();
1324 let b = streams.open(ticket).unwrap();
1325
1326 let filler = vec![7_u8; STREAM_BUFFER_BYTES];
1327 assert_eq!(a.try_write(&filler).unwrap(), STREAM_BUFFER_BYTES);
1328 assert!(matches!(a.try_write(b"x"), Err(Error::WouldBlock)));
1329
1330 let mut buf = vec![0_u8; STREAM_BUFFER_BYTES - 5];
1331 assert_eq!(b.blocking_read(&mut buf).unwrap(), STREAM_BUFFER_BYTES - 5);
1332 let tail = (0..STREAM_BUFFER_BYTES)
1334 .map(|i| i as u8)
1335 .collect::<Vec<_>>();
1336 assert_eq!(a.try_write(&tail).unwrap(), STREAM_BUFFER_BYTES - 5);
1337
1338 let mut rest = vec![0_u8; 5];
1339 assert_eq!(b.blocking_read(&mut rest).unwrap(), 5);
1340 assert_eq!(rest, vec![7_u8; 5]);
1341 let mut wrapped = vec![0_u8; STREAM_BUFFER_BYTES];
1342 let mut got = 0;
1343 while got < STREAM_BUFFER_BYTES - 5 {
1344 got += b.blocking_read(&mut wrapped[got..]).unwrap();
1345 }
1346 assert_eq!(&wrapped[..got], &tail[..got]);
1347 }
1348
1349 #[test]
1350 fn a_blocked_reader_is_woken_by_a_write_in_another_thread() {
1351 let streams = streams("stream-wake");
1352 let (a, ticket) = streams.create().unwrap();
1353 let b = streams.open(ticket).unwrap();
1354 let reader = std::thread::spawn(move || {
1355 let mut buf = [0_u8; 8];
1356 let n = b.blocking_read(&mut buf).unwrap();
1357 buf[..n].to_vec()
1358 });
1359 std::thread::sleep(Duration::from_millis(30));
1360 a.blocking_write_all(b"ping").unwrap();
1361 assert_eq!(reader.join().unwrap(), b"ping");
1362 }
1363
1364 #[test]
1365 fn a_reader_gone_fails_the_writer_and_a_lane_fills_up() {
1366 let streams = streams("stream-gone");
1367 let (a, ticket) = streams.create().unwrap();
1368 let b = streams.open(ticket).unwrap();
1369 let (b_read, _b_write) = b.split();
1370 drop(b_read);
1371 assert!(matches!(a.blocking_write(b"nobody"), Err(Error::PeerGone)));
1372
1373 let mut held = vec![a];
1374 for _ in 1..STREAM_LANE_CAPACITY {
1375 held.push(streams.create().unwrap().0);
1376 }
1377 assert!(matches!(streams.create(), Err(Error::Full { .. })));
1378 }
1379
1380 #[test]
1381 fn a_death_report_ends_that_incarnations_sides_only() {
1382 let streams = streams("stream-dead");
1383 let (a, ticket) = streams.create().unwrap();
1384 let b = streams.open(ticket).unwrap();
1385 b.blocking_write_all(b"partial").unwrap();
1386
1387 streams.node_dead(NodeId::ZERO, Incarnation::new(2));
1389 assert_eq!(a.blocking_read_chunk(16).unwrap().as_ref(), b"partial");
1390
1391 let reader = std::thread::spawn(move || {
1397 let mut buf = [0_u8; 8];
1398 a.blocking_read(&mut buf).map(|_| ())
1399 });
1400 std::thread::sleep(Duration::from_millis(30));
1401 streams.node_dead(NodeId::ZERO, Incarnation::new(1));
1402 assert!(matches!(
1403 reader.join().unwrap(),
1404 Err(Error::Reset | Error::Stale(_))
1405 ));
1406 assert!(matches!(
1407 b.try_write(b"x"),
1408 Err(Error::Reset | Error::Stale(_))
1409 ));
1410 assert!(matches!(
1411 streams.open(ticket),
1412 Err(Error::AlreadyClaimed(_) | Error::Stale(_))
1413 ));
1414 drop(b);
1415 assert!(!streams.is_live(ticket.id));
1416 }
1417
1418 #[test]
1419 fn an_offer_is_taken_once_and_a_stale_offer_is_skipped() {
1420 let streams = streams("stream-offer");
1421 assert!(streams.take_offer().is_none());
1422 let (a, ticket) = streams.create().unwrap();
1423 streams.offer(ticket, NodeId::ZERO).unwrap();
1424 assert_eq!(streams.take_offer(), Some(ticket));
1425 assert!(streams.take_offer().is_none());
1426
1427 let (gone, stale) = streams.create().unwrap();
1428 streams.offer(stale, NodeId::ZERO).unwrap();
1429 drop(gone);
1430 assert!(streams.take_offer().is_none());
1431
1432 let taker = streams.clone();
1433 let waiter = std::thread::spawn(move || taker.blocking_take_offer());
1434 std::thread::sleep(Duration::from_millis(30));
1435 streams.offer(ticket, NodeId::ZERO).unwrap();
1436 assert_eq!(waiter.join().unwrap().unwrap(), ticket);
1437 drop(a);
1438 }
1439
1440 #[test]
1441 fn a_contested_claim_goes_to_exactly_one_opener() {
1442 let streams = streams("stream-race");
1443 for _ in 0..50 {
1444 let (_a, ticket) = streams.create().unwrap();
1445 let openers = (0..4)
1446 .map(|_| {
1447 let streams = streams.clone();
1448 std::thread::spawn(move || streams.open(ticket).is_ok())
1449 })
1450 .collect::<Vec<_>>();
1451 let won = openers
1452 .into_iter()
1453 .map(|opener| opener.join().unwrap())
1454 .filter(|won| *won)
1455 .count();
1456 assert_eq!(won, 1);
1457 }
1458 }
1459
1460 #[test]
1461 fn random_sizes_concatenate_exactly() {
1462 let streams = streams("stream-random");
1463 let (mut a, ticket) = streams.create().unwrap();
1464 let b = streams.open(ticket).unwrap();
1465 let payload = (0..300_000_u32)
1466 .map(|i| (i.wrapping_mul(2654435761) >> 13) as u8)
1467 .collect::<Vec<_>>();
1468 let expected = payload.clone();
1469 let writer = std::thread::spawn(move || {
1470 let mut rest = &payload[..];
1471 let mut step = 1_usize;
1472 while !rest.is_empty() {
1473 let len = step.min(rest.len());
1474 a.blocking_write_all(&rest[..len]).unwrap();
1475 rest = &rest[len..];
1476 step = (step * 7 + 13) % 9_001 + 1;
1477 }
1478 a.finish().unwrap();
1479 });
1480 let mut out = Vec::new();
1481 let mut step = 3_usize;
1482 loop {
1483 let chunk = b.blocking_read_chunk(step).unwrap();
1484 if chunk.is_empty() {
1485 break;
1486 }
1487 out.extend_from_slice(&chunk);
1488 step = (step * 5 + 11) % 4_099 + 1;
1489 }
1490 writer.join().unwrap();
1491 assert_eq!(out, expected);
1492 }
1493}