1use std::collections::HashSet;
2use std::fmt;
3use std::ops::{Deref, DerefMut};
4use std::str::FromStr;
5use std::sync::{Arc, Mutex};
6use std::time::Duration;
7
8use orbit_core::{Fleet, NodeId};
9
10use super::arena::{FlowKey, FlowWaitSpec, Publication, allowance_admits};
11use super::protocol::CONTROL_FRAME_BYTES;
12use super::{
13 ChunkDescriptor, ChunkGeometry, ChunkPlan, ControlEvent,
14 DEFAULT_CHUNK_BYTES, ExchangeId, Flow, PayloadArena, PayloadArenaSpec,
15 PayloadChunk, ResetCode,
16};
17use crate::{
18 Error, Incarnation, ReadHalf, Result, Side, StreamSpec, Streams, Ticket,
19 WriteHalf, lock_unpoisoned,
20};
21
22#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
29pub enum ExchangePayloadSpec {
30 Shared(PayloadArenaSpec),
31 Directional { a_to_b: PayloadArenaSpec, b_to_a: PayloadArenaSpec },
32}
33
34impl ExchangePayloadSpec {
35 pub const fn for_flow(
37 self,
38 flow: Flow,
39 ) -> PayloadArenaSpec {
40 match (self, flow) {
41 (Self::Shared(payload), _) => payload,
42 (Self::Directional { a_to_b, .. }, Flow::AtoB) => a_to_b,
43 (Self::Directional { b_to_a, .. }, Flow::BtoA) => b_to_a,
44 }
45 }
46}
47
48#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
50pub struct ExchangeSpec {
51 pub control: StreamSpec,
54 pub payload: ExchangePayloadSpec,
57}
58
59impl ExchangeSpec {
60 pub const fn new(
61 control: StreamSpec,
62 payload: PayloadArenaSpec,
63 ) -> Self {
64 Self { control, payload: ExchangePayloadSpec::Shared(payload) }
65 }
66
67 pub const fn directional(
68 control: StreamSpec,
69 a_to_b: PayloadArenaSpec,
70 b_to_a: PayloadArenaSpec,
71 ) -> Self {
72 Self {
73 control,
74 payload: ExchangePayloadSpec::Directional { a_to_b, b_to_a },
75 }
76 }
77
78 pub const fn payload_for(
80 self,
81 flow: Flow,
82 ) -> PayloadArenaSpec {
83 self.payload.for_flow(flow)
84 }
85
86 pub const fn chunk_geometry(
88 self,
89 flow: Flow,
90 ) -> Option<ChunkGeometry> {
91 self.payload_for(flow).chunk_geometry()
92 }
93
94 fn validate(self) -> Result<()> {
95 if self.control.buffer_bytes < CONTROL_FRAME_BYTES {
96 return Err(Error::Malformed(format!(
97 "exchange control ring is {} bytes; one control frame is {CONTROL_FRAME_BYTES}",
98 self.control.buffer_bytes
99 )));
100 }
101 match self.payload {
102 ExchangePayloadSpec::Shared(payload)
103 if self.control.kind == payload.kind =>
104 {
105 return Err(Error::Malformed(
106 "exchange control and payload need distinct SHM kinds"
107 .to_owned(),
108 ));
109 }
110 ExchangePayloadSpec::Directional { a_to_b, b_to_a }
111 if self.control.kind == a_to_b.kind
112 || self.control.kind == b_to_a.kind
113 || a_to_b.kind == b_to_a.kind =>
114 {
115 return Err(Error::Malformed(
116 "exchange control and directional payloads need three distinct SHM kinds"
117 .to_owned()
118 ));
119 }
120 _ => {}
121 }
122 Ok(())
123 }
124}
125
126#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
129pub struct ExchangeTicket(Ticket);
130
131impl ExchangeTicket {
132 pub const fn from_stream_ticket(ticket: Ticket) -> Self {
133 Self(ticket)
134 }
135
136 pub const fn stream_ticket(self) -> Ticket {
137 self.0
138 }
139
140 pub const fn id(self) -> ExchangeId {
141 ExchangeId::from_stream_id(self.0.id)
142 }
143}
144
145impl fmt::Display for ExchangeTicket {
146 fn fmt(
147 &self,
148 f: &mut fmt::Formatter<'_>,
149 ) -> fmt::Result {
150 self.0.fmt(f)
151 }
152}
153
154impl FromStr for ExchangeTicket {
155 type Err = Error;
156
157 fn from_str(text: &str) -> Result<Self> {
158 text.parse().map(Self)
159 }
160}
161
162#[derive(Clone)]
163enum PayloadArenas {
164 Shared(PayloadArena),
165 Directional { a_to_b: PayloadArena, b_to_a: PayloadArena },
166}
167
168impl PayloadArenas {
169 fn for_flow(
170 &self,
171 flow: Flow,
172 ) -> PayloadArena {
173 match (self, flow) {
174 (Self::Shared(arena), _) => arena.clone(),
175 (Self::Directional { a_to_b, .. }, Flow::AtoB) => a_to_b.clone(),
176 (Self::Directional { b_to_a, .. }, Flow::BtoA) => b_to_a.clone(),
177 }
178 }
179
180 fn node_dead(
181 &self,
182 node: NodeId,
183 incarnation: Incarnation,
184 dead_flows: &HashSet<(u64, u16)>,
185 ) {
186 match self {
187 Self::Shared(arena) => {
188 arena.node_dead_with_flows(node, incarnation, dead_flows)
189 }
190 Self::Directional { a_to_b, b_to_a } => {
191 a_to_b.node_dead_with_flows(node, incarnation, dead_flows);
192 b_to_a.node_dead_with_flows(node, incarnation, dead_flows);
193 }
194 }
195 }
196
197 fn reset_all(&self) {
198 match self {
199 Self::Shared(arena) => arena.reset_all(),
200 Self::Directional { a_to_b, b_to_a } => {
201 a_to_b.reset_all();
202 b_to_a.reset_all();
203 }
204 }
205 }
206
207 #[cfg(unix)]
208 fn unlink(&self) -> Result<()> {
209 match self {
210 Self::Shared(arena) => arena.unlink(),
211 Self::Directional { a_to_b, b_to_a } => {
212 a_to_b.unlink()?;
213 b_to_a.unlink()
214 }
215 }
216 }
217}
218
219#[derive(Clone)]
222pub struct Exchanges {
223 control: Streams,
224 payload: PayloadArenas,
225 flow_slots: Option<usize>,
226}
227
228impl Exchanges {
229 pub fn open(
230 fleet: Arc<Fleet>,
231 incarnation: Incarnation,
232 spec: ExchangeSpec,
233 ) -> Result<Self> {
234 spec.validate()?;
235 let control =
236 Streams::with_spec(Arc::clone(&fleet), incarnation, spec.control)?;
237 let payload = match spec.payload {
238 ExchangePayloadSpec::Shared(payload) => PayloadArenas::Shared(
239 PayloadArena::open(fleet, incarnation, payload)?,
240 ),
241 ExchangePayloadSpec::Directional { a_to_b, b_to_a } => {
242 PayloadArenas::Directional {
243 a_to_b: PayloadArena::open(
244 Arc::clone(&fleet),
245 incarnation,
246 a_to_b,
247 )?,
248 b_to_a: PayloadArena::open(fleet, incarnation, b_to_a)?,
249 }
250 }
251 };
252 Ok(Self { control, payload, flow_slots: None })
253 }
254
255 pub fn with_flow_slots(
258 mut self,
259 slots: usize,
260 ) -> Result<Self> {
261 if slots == 0
262 || slots > self.payload.for_flow(Flow::AtoB).slots_per_node()
263 || slots > self.payload.for_flow(Flow::BtoA).slots_per_node()
264 {
265 return Err(Error::Malformed(
266 "exchange flow credit must fit both payload lanes".to_owned(),
267 ));
268 }
269 self.flow_slots = Some(slots);
270 Ok(self)
271 }
272
273 fn creation_arena(
274 &self,
275 len: usize,
276 ) -> Result<PayloadArena> {
277 let arena = self.payload.for_flow(Flow::AtoB);
278 if let Some(slots) = self.flow_slots {
279 flow_slots_for(len, slots, arena.slot_size())?;
280 }
281 Ok(arena)
282 }
283
284 pub fn create(&self) -> Result<(ExchangeEndpoint, ExchangeTicket)> {
287 let (endpoint, ticket) = self.control.create()?;
288 Ok((self.endpoint(endpoint), ExchangeTicket(ticket)))
289 }
290
291 pub fn create_payload_available(
295 &self,
296 payload_len: usize,
297 ) -> Result<bool> {
298 self.creation_arena(payload_len)?.is_available(payload_len)
299 }
300
301 pub fn poll_create_payload_available(
303 &self,
304 payload_len: usize,
305 cx: &mut std::task::Context<'_>,
306 ) -> std::task::Poll<Result<()>> {
307 self.creation_arena(payload_len)?.poll_available(payload_len, cx)
308 }
309
310 pub fn open_peer(
313 &self,
314 ticket: ExchangeTicket,
315 ) -> Result<ExchangeEndpoint> {
316 let endpoint = self.control.open(ticket.0)?;
317 Ok(self.endpoint(endpoint))
318 }
319
320 fn endpoint(
321 &self,
322 endpoint: crate::Endpoint,
323 ) -> ExchangeEndpoint {
324 let id = ExchangeId::from_stream_id(endpoint.id());
325 let (send_flow, receive_flow) = match endpoint.side() {
326 Side::A => (Flow::AtoB, Flow::BtoA),
327 Side::B => (Flow::BtoA, Flow::AtoB),
328 };
329 let (read, write) = endpoint.split();
330 let unopened = (send_flow == Flow::AtoB).then(|| {
331 Arc::new(UnopenedPeer {
332 control: Arc::clone(&write.handle),
333 arena: self.payload.for_flow(send_flow),
334 unread: Mutex::new(Vec::new()),
335 })
336 });
337 ExchangeEndpoint {
338 sender: Sender(Producer::new(
339 id,
340 send_flow,
341 self.payload.for_flow(send_flow),
342 write,
343 self.flow_slots,
344 unopened.clone(),
345 )),
346 receiver: Receiver(Consumer::new(
347 id,
348 receive_flow,
349 self.payload.for_flow(receive_flow),
350 read,
351 unopened,
352 )),
353 }
354 }
355
356 pub fn offer(
357 &self,
358 ticket: ExchangeTicket,
359 to: NodeId,
360 ) -> Result<()> {
361 self.control.offer(ticket.0, to)
362 }
363
364 pub fn take_offer(&self) -> Option<ExchangeTicket> {
365 self.control.take_offer().map(ExchangeTicket)
366 }
367
368 pub fn blocking_take_offer(&self) -> Result<ExchangeTicket> {
370 self.control.blocking_take_offer().map(ExchangeTicket)
371 }
372
373 #[cfg(any(
375 target_os = "linux",
376 target_os = "freebsd",
377 target_os = "macos"
378 ))]
379 pub fn take_offer_timeout(
380 &self,
381 timeout: Duration,
382 ) -> Result<Option<ExchangeTicket>> {
383 self.control
384 .take_offer_timeout(timeout)
385 .map(|ticket| ticket.map(ExchangeTicket))
386 }
387
388 pub fn poll_take_offer(
390 &self,
391 cx: &mut std::task::Context<'_>,
392 ) -> std::task::Poll<Result<ExchangeTicket>> {
393 match self.control.poll_take_offer(cx) {
394 std::task::Poll::Ready(result) => {
395 std::task::Poll::Ready(result.map(ExchangeTicket))
396 }
397 std::task::Poll::Pending => std::task::Poll::Pending,
398 }
399 }
400
401 pub fn control(&self) -> &Streams {
402 &self.control
403 }
404
405 pub fn node_dead(
408 &self,
409 node: NodeId,
410 incarnation: Incarnation,
411 ) {
412 let dead_flows = self
413 .control
414 .exchange_readers(node, incarnation)
415 .map(|(id, side)| {
416 let flow = match side {
417 Side::A => Flow::BtoA,
418 Side::B => Flow::AtoB,
419 };
420 (id.net_id().raw(), u16::from(flow.wire()))
421 })
422 .collect();
423 self.control.node_dead(node, incarnation);
424 self.payload.node_dead(node, incarnation, &dead_flows);
425 }
426
427 pub fn reset_all(&self) {
430 self.control.reset_all();
431 self.payload.reset_all();
432 }
433
434 #[cfg(unix)]
437 pub fn unlink(&self) -> Result<()> {
438 self.control.unlink()?;
439 self.payload.unlink()
440 }
441}
442
443pub struct ExchangeEndpoint {
446 sender: Sender,
447 receiver: Receiver,
448}
449
450impl ExchangeEndpoint {
451 pub fn id(&self) -> ExchangeId {
452 self.sender.id()
453 }
454
455 pub fn sender(&mut self) -> &mut Sender {
456 &mut self.sender
457 }
458
459 pub fn receiver(&mut self) -> &mut Receiver {
460 &mut self.receiver
461 }
462
463 pub fn split(self) -> (Sender, Receiver) {
464 (self.sender, self.receiver)
465 }
466}
467
468#[derive(Clone, Copy, Debug, Eq, PartialEq)]
469enum ProducerState {
470 New,
471 Open,
472 Terminal,
473}
474
475struct UnopenedPeer {
479 control: Arc<crate::Handle>,
480 arena: PayloadArena,
481 unread: Mutex<Vec<ChunkDescriptor>>,
482}
483
484impl Drop for UnopenedPeer {
485 fn drop(&mut self) {
486 let unread =
487 self.unread.get_mut().unwrap_or_else(|error| error.into_inner());
488 if !unread.is_empty() && self.control.close_unclaimed_peer() {
489 for descriptor in unread.drain(..) {
490 self.arena.discard(descriptor);
491 }
492 }
493 }
494}
495
496fn flow_slots_for(
497 len: usize,
498 slots: usize,
499 slot_size: usize,
500) -> Result<usize> {
501 let capacity = slots * slot_size;
502 if len > capacity {
503 return Err(Error::PayloadTooLarge { len, capacity });
504 }
505 Ok(len.div_ceil(slot_size))
506}
507
508struct Producer {
509 exchange: ExchangeId,
510 flow: Flow,
511 arena: PayloadArena,
512 control: WriteHalf,
513 status: crate::WriteStatus,
515 next_chunk: u64,
516 state: ProducerState,
517 flow_slots: usize,
518 published: Vec<ChunkDescriptor>,
519 published_slots: usize,
520 unopened: Option<Arc<UnopenedPeer>>,
521}
522
523impl Producer {
524 fn new(
525 exchange: ExchangeId,
526 flow: Flow,
527 arena: PayloadArena,
528 control: WriteHalf,
529 flow_slots: Option<usize>,
530 unopened: Option<Arc<UnopenedPeer>>,
531 ) -> Self {
532 let flow_slots = flow_slots.unwrap_or_else(|| arena.slots_per_node());
533 let status = control.write_status();
534 Self {
535 exchange,
536 flow,
537 arena,
538 control,
539 status,
540 next_chunk: 1,
541 state: ProducerState::New,
542 flow_slots,
543 published: Vec::new(),
544 published_slots: 0,
545 unopened,
546 }
547 }
548
549 fn id(&self) -> ExchangeId {
550 self.exchange
551 }
552
553 fn start(
554 &mut self,
555 metadata: Option<&[u8]>,
556 ) -> Result<()> {
557 if self.state != ProducerState::New {
558 return Err(Error::Malformed(
559 "flow start was already sent".to_owned(),
560 ));
561 }
562 if let Some(bytes) = metadata.filter(|bytes| !bytes.is_empty()) {
563 let mut pending = self.reserve_start(bytes.len())?;
564 pending.copy_from_slice(bytes);
565 pending.commit()?;
566 } else {
567 self.send(ControlEvent::Start {
568 exchange: self.exchange,
569 flow: self.flow,
570 metadata: None,
571 })?;
572 self.state = ProducerState::Open;
573 }
574 Ok(())
575 }
576
577 fn data(
578 &mut self,
579 payload: &[u8],
580 ) -> Result<ChunkDescriptor> {
581 let mut pending = self.reserve_data(payload.len())?;
582 pending.copy_from_slice(payload);
583 pending.commit()
584 }
585
586 fn reserve_start(
587 &mut self,
588 payload_len: usize,
589 ) -> Result<PendingStart<'_>> {
590 if self.state != ProducerState::New {
591 return Err(Error::Malformed(
592 "flow start was already sent".to_owned(),
593 ));
594 }
595 let publication = self.reserve(payload_len)?;
596 Ok(PendingStart { producer: self, publication })
597 }
598
599 fn reserve_data(
600 &mut self,
601 payload_len: usize,
602 ) -> Result<PendingData<'_>> {
603 if self.state != ProducerState::Open {
604 return Err(Error::Malformed(
605 "flow data requires an open flow".to_owned(),
606 ));
607 }
608 let publication = self.reserve(payload_len)?;
609 Ok(PendingData { producer: self, publication })
610 }
611
612 fn required_slots(
613 &self,
614 len: usize,
615 ) -> Result<usize> {
616 flow_slots_for(len, self.flow_slots, self.arena.slot_size())
617 }
618
619 fn reserve(
620 &mut self,
621 len: usize,
622 ) -> Result<Publication> {
623 self.control.exchange_write_status()?;
624 let count = self.required_slots(len)?;
625 if self.published_slots + count > self.flow_slots {
626 self.published.retain(|descriptor| self.arena.retains(*descriptor));
627 self.published_slots =
628 self.published.iter().map(|d| d.slot_count() as usize).sum();
629 }
630 if self.published_slots + count > self.flow_slots {
631 return Err(Error::PayloadFull { requested_slots: count });
632 }
633 self.arena.reserve(len)
634 }
635
636 fn credit_ready(
637 &self,
638 len: usize,
639 ) -> Result<bool> {
640 self.control.exchange_write_status()?;
641 let count = self.required_slots(len)?;
642 if self.published_slots + count > self.flow_slots
643 && !allowance_admits(
644 &self.published,
645 self.flow_slots,
646 count,
647 |descriptor| self.arena.retains(descriptor),
648 )
649 {
650 return Ok(false);
651 }
652 self.arena.is_available(len)
653 }
654
655 fn credit_key(&self) -> FlowKey {
656 (self.exchange.stream_id().net_id().raw(), u16::from(self.flow.wire()))
657 }
658
659 fn poll_credit(
663 &self,
664 len: usize,
665 cx: &mut std::task::Context<'_>,
666 ) -> std::task::Poll<Result<()>> {
667 let demand = match self.required_slots(len) {
668 Ok(demand) => demand,
669 Err(error) => return std::task::Poll::Ready(Err(error)),
670 };
671 self.arena.poll_flow(
672 FlowWaitSpec {
673 key: self.credit_key(),
674 status: &self.status,
675 demand,
676 limit: self.flow_slots,
677 held: &self.published,
678 },
679 || self.credit_ready(len),
680 cx,
681 )
682 }
683
684 fn publish(
685 &mut self,
686 publication: Publication,
687 descriptor: ChunkDescriptor,
688 event: ControlEvent,
689 ) -> Result<()> {
690 publication.bind(self.exchange, self.flow);
691 publication.make_live();
692 self.send(event)?;
693 publication.mark_published();
694 self.published_slots += descriptor.slot_count() as usize;
695 self.published.push(descriptor);
696
697 if matches!(self.control.exchange_write_status(), Err(Error::PeerGone))
700 {
701 self.discard_unread();
702 }
703 self.next_chunk += 1;
704 Ok(())
705 }
706
707 fn discard_unread(&mut self) {
708 for descriptor in self.published.drain(..) {
709 self.arena.discard(descriptor);
710 }
711 self.published_slots = 0;
712 }
713
714 fn commit_start(
715 &mut self,
716 publication: Publication,
717 ) -> Result<ChunkDescriptor> {
718 let descriptor = self.descriptor(&publication);
719 self.publish(
720 publication,
721 descriptor,
722 ControlEvent::Start {
723 exchange: self.exchange,
724 flow: self.flow,
725 metadata: Some(descriptor),
726 },
727 )?;
728 self.state = ProducerState::Open;
729 Ok(descriptor)
730 }
731
732 fn commit_data(
733 &mut self,
734 publication: Publication,
735 ) -> Result<ChunkDescriptor> {
736 let descriptor = self.descriptor(&publication);
737 self.publish(publication, descriptor, ControlEvent::Data(descriptor))?;
738 Ok(descriptor)
739 }
740
741 fn finish(&mut self) -> Result<()> {
742 if self.state != ProducerState::Open {
743 return Err(Error::Malformed(
744 "flow FIN requires an open flow".to_owned(),
745 ));
746 }
747 self.send(ControlEvent::Fin {
748 exchange: self.exchange,
749 flow: self.flow,
750 })?;
751 self.control.finish()?;
752 self.state = ProducerState::Terminal;
753 Ok(())
754 }
755
756 fn reset(
757 &mut self,
758 code: ResetCode,
759 ) -> Result<()> {
760 if self.state == ProducerState::Terminal {
761 return Err(Error::Closed);
762 }
763 self.control.reset_exchange(code.get())?;
764 self.discard_unread();
765 self.state = ProducerState::Terminal;
766 Ok(())
767 }
768
769 fn send(
770 &self,
771 event: ControlEvent,
772 ) -> Result<()> {
773 self.control.try_write_exact(&event.encode())
774 }
775
776 fn descriptor(
777 &self,
778 publication: &Publication,
779 ) -> ChunkDescriptor {
780 let (
781 generation,
782 first_slot,
783 slot_count,
784 payload_len,
785 arena_kind,
786 owner_node,
787 ) = publication.coordinates();
788 ChunkDescriptor::new(
789 self.exchange,
790 self.flow,
791 self.next_chunk,
792 generation,
793 first_slot,
794 slot_count,
795 payload_len,
796 arena_kind,
797 owner_node,
798 )
799 }
800}
801
802impl Drop for Producer {
803 fn drop(&mut self) {
804 self.arena.forget_flow(self.credit_key());
806 if self.published.is_empty() {
807 return;
808 }
809 if self.state != ProducerState::Terminal {
810 self.control.reset();
811 self.discard_unread();
812 } else if matches!(
813 self.control.exchange_write_status(),
814 Err(Error::PeerGone)
815 ) {
816 self.discard_unread();
817 } else if let Some(unopened) = &self.unopened
818 && unopened.control.peer_is_unclaimed()
819 {
820 *lock_unpoisoned(&unopened.unread) =
821 std::mem::take(&mut self.published);
822 }
823 }
824}
825
826pub struct PendingStart<'a> {
830 producer: &'a mut Producer,
831 publication: Publication,
832}
833
834impl PendingStart<'_> {
835 pub fn as_mut_slice(&mut self) -> &mut [u8] {
836 self.publication.as_mut_slice()
837 }
838
839 pub fn commit(self) -> Result<ChunkDescriptor> {
840 self.producer.commit_start(self.publication)
841 }
842
843 pub fn commit_len(
846 mut self,
847 payload_len: usize,
848 ) -> Result<ChunkDescriptor> {
849 self.publication.truncate(payload_len)?;
850 self.producer.commit_start(self.publication)
851 }
852}
853
854impl Deref for PendingStart<'_> {
855 type Target = [u8];
856
857 fn deref(&self) -> &Self::Target {
858 self.publication.as_slice()
859 }
860}
861
862impl DerefMut for PendingStart<'_> {
863 fn deref_mut(&mut self) -> &mut Self::Target {
864 self.as_mut_slice()
865 }
866}
867
868pub struct PendingData<'a> {
871 producer: &'a mut Producer,
872 publication: Publication,
873}
874
875impl PendingData<'_> {
876 pub fn as_mut_slice(&mut self) -> &mut [u8] {
877 self.publication.as_mut_slice()
878 }
879
880 pub fn commit(self) -> Result<ChunkDescriptor> {
881 self.producer.commit_data(self.publication)
882 }
883
884 pub fn commit_len(
887 mut self,
888 payload_len: usize,
889 ) -> Result<ChunkDescriptor> {
890 self.publication.truncate(payload_len)?;
891 self.producer.commit_data(self.publication)
892 }
893}
894
895impl Deref for PendingData<'_> {
896 type Target = [u8];
897
898 fn deref(&self) -> &Self::Target {
899 self.publication.as_slice()
900 }
901}
902
903impl DerefMut for PendingData<'_> {
904 fn deref_mut(&mut self) -> &mut Self::Target {
905 self.as_mut_slice()
906 }
907}
908
909#[derive(Clone, Copy, Debug, Eq, PartialEq)]
910enum ConsumerState {
911 New,
912 Open,
913 Terminal,
914}
915
916struct Consumer {
917 exchange: ExchangeId,
918 flow: Flow,
919 arena: PayloadArena,
920 control: ReadHalf,
921 next_chunk: u64,
922 state: ConsumerState,
923 _unopened: Option<Arc<UnopenedPeer>>,
924}
925
926impl Consumer {
927 fn new(
928 exchange: ExchangeId,
929 flow: Flow,
930 arena: PayloadArena,
931 control: ReadHalf,
932 unopened: Option<Arc<UnopenedPeer>>,
933 ) -> Self {
934 Self {
935 exchange,
936 flow,
937 arena,
938 control,
939 next_chunk: 1,
940 state: ConsumerState::New,
941 _unopened: unopened,
942 }
943 }
944
945 fn id(&self) -> ExchangeId {
946 self.exchange
947 }
948
949 fn next(&mut self) -> Result<FlowEvent> {
950 if self.state == ConsumerState::Terminal {
951 return Err(Error::Closed);
952 }
953 let mut frame = [0_u8; CONTROL_FRAME_BYTES];
954 if let Err(error) = self.control.try_read_exact(&mut frame) {
955 return self.read_error(error);
956 }
957 let event = ControlEvent::decode(&frame)
958 .map_err(|reason| Error::Malformed(reason.to_owned()))?;
959 if event.exchange() != self.exchange || event.flow() != self.flow {
960 return Err(Error::Malformed(
961 "control event belongs to another exchange flow".to_owned(),
962 ));
963 }
964 match event {
965 ControlEvent::Start { metadata, .. }
966 if self.state == ConsumerState::New =>
967 {
968 let metadata = metadata
969 .map(|descriptor| self.take_chunk(descriptor))
970 .transpose();
971 let metadata = match metadata {
972 Ok(metadata) => metadata,
973 Err(error) => return self.read_error(error),
974 };
975 self.state = ConsumerState::Open;
976 Ok(FlowEvent::Start { metadata })
977 }
978 ControlEvent::Data(descriptor)
979 if self.state == ConsumerState::Open =>
980 {
981 match self.take_chunk(descriptor) {
982 Ok(chunk) => Ok(FlowEvent::Data(chunk)),
983 Err(error) => self.read_error(error),
984 }
985 }
986 ControlEvent::Fin { .. } if self.state == ConsumerState::Open => {
987 self.state = ConsumerState::Terminal;
988 Ok(FlowEvent::Fin)
989 }
990 ControlEvent::Reset { code, .. } => {
991 self.state = ConsumerState::Terminal;
992 Ok(FlowEvent::Reset(code))
993 }
994 _ => Err(Error::Malformed(
995 "control event violates the flow lifecycle".to_owned(),
996 )),
997 }
998 }
999
1000 fn read_error(
1001 &mut self,
1002 error: Error,
1003 ) -> Result<FlowEvent> {
1004 if let Some(code) = self.control.exchange_reset_code()? {
1005 self.state = ConsumerState::Terminal;
1006 return Ok(FlowEvent::Reset(ResetCode::new(code)));
1007 }
1008 Err(error)
1009 }
1010
1011 fn cancel(&mut self) {
1012 if self.state == ConsumerState::Terminal {
1013 return;
1014 }
1015 self.state = ConsumerState::Terminal;
1016 self.control.stop_reading();
1017 if let Ok(node) = self.control.exchange_peer_node() {
1018 self.arena.notify_peer(node);
1019 }
1020 let mut frame = [0; CONTROL_FRAME_BYTES];
1021 while self.control.try_drain_exchange_frame(&mut frame).is_ok() {
1022 let Ok(event) = ControlEvent::decode(&frame) else { continue };
1023 if event.exchange() != self.exchange || event.flow() != self.flow {
1024 continue;
1025 }
1026 match event {
1027 ControlEvent::Start { metadata: Some(descriptor), .. }
1028 | ControlEvent::Data(descriptor) => {
1029 drop(self.arena.read(descriptor));
1032 }
1033 _ => {}
1034 }
1035 }
1036 }
1037
1038 fn take_chunk(
1039 &mut self,
1040 descriptor: ChunkDescriptor,
1041 ) -> Result<PayloadChunk> {
1042 if descriptor.chunk_id() != self.next_chunk {
1043 return Err(Error::Malformed(format!(
1044 "flow expected chunk {}, received {}",
1045 self.next_chunk,
1046 descriptor.chunk_id()
1047 )));
1048 }
1049 let chunk = self.arena.read(descriptor)?;
1050 self.next_chunk += 1;
1051 Ok(chunk)
1052 }
1053}
1054
1055impl Drop for Consumer {
1056 fn drop(&mut self) {
1057 self.cancel();
1058 }
1059}
1060
1061pub enum FlowEvent {
1063 Start { metadata: Option<PayloadChunk> },
1064 Data(PayloadChunk),
1065 Fin,
1066 Reset(ResetCode),
1067}
1068
1069pub trait FlowHandler {
1072 type Error;
1073
1074 fn on_start(
1075 &mut self,
1076 metadata: Option<PayloadChunk>,
1077 ) -> std::result::Result<(), Self::Error>;
1078 fn on_data(
1079 &mut self,
1080 chunk: PayloadChunk,
1081 ) -> std::result::Result<(), Self::Error>;
1082 fn on_fin(&mut self) -> std::result::Result<(), Self::Error>;
1083 fn on_reset(
1084 &mut self,
1085 code: ResetCode,
1086 ) -> std::result::Result<(), Self::Error>;
1087}
1088
1089#[derive(Debug)]
1090pub enum DispatchError<E> {
1091 Transport(Error),
1092 Handler(E),
1093}
1094
1095impl<E> From<Error> for DispatchError<E> {
1096 fn from(value: Error) -> Self {
1097 Self::Transport(value)
1098 }
1099}
1100
1101fn dispatch<H: FlowHandler>(
1102 consumer: &mut Consumer,
1103 handler: &mut H,
1104) -> std::result::Result<(), DispatchError<H::Error>> {
1105 match consumer.next()? {
1106 FlowEvent::Start { metadata } => handler.on_start(metadata),
1107 FlowEvent::Data(chunk) => handler.on_data(chunk),
1108 FlowEvent::Fin => handler.on_fin(),
1109 FlowEvent::Reset(code) => handler.on_reset(code),
1110 }
1111 .map_err(DispatchError::Handler)
1112}
1113
1114pub struct Sender(Producer);
1116
1117impl Sender {
1118 pub fn id(&self) -> ExchangeId {
1119 self.0.id()
1120 }
1121
1122 pub fn flow(&self) -> Flow {
1123 self.0.flow
1124 }
1125
1126 pub fn chunk_geometry(&self) -> ChunkGeometry {
1128 ChunkGeometry::new(self.0.arena.slot_size(), self.0.flow_slots)
1129 .expect("opened exchange has nonzero flow credit")
1130 }
1131
1132 pub fn plan_chunks(
1134 &self,
1135 data_bytes: usize,
1136 ) -> Result<ChunkPlan> {
1137 self.chunk_geometry()
1138 .plan_for_target(data_bytes, DEFAULT_CHUNK_BYTES)
1139 .ok_or_else(|| {
1140 Error::Malformed(
1141 "payload chunk plan exceeds flow credit".to_owned(),
1142 )
1143 })
1144 }
1145
1146 pub fn plan_chunks_with(
1148 &self,
1149 data_bytes: usize,
1150 slots_per_chunk: usize,
1151 ) -> Result<ChunkPlan> {
1152 self.chunk_geometry().plan(data_bytes, slots_per_chunk)
1153 .ok_or_else(|| Error::Malformed(format!(
1154 "chunk slot count {slots_per_chunk} exceeds exchange flow credit {}", self.0.flow_slots,
1155 )))
1156 }
1157
1158 pub fn start(
1159 &mut self,
1160 metadata: Option<&[u8]>,
1161 ) -> Result<()> {
1162 self.0.start(metadata)
1163 }
1164
1165 pub fn reserve_start(
1168 &mut self,
1169 payload_len: usize,
1170 ) -> Result<PendingStart<'_>> {
1171 self.0.reserve_start(payload_len)
1172 }
1173
1174 pub fn data(
1175 &mut self,
1176 payload: &[u8],
1177 ) -> Result<ChunkDescriptor> {
1178 self.0.data(payload)
1179 }
1180
1181 pub fn reserve_data(
1184 &mut self,
1185 payload_len: usize,
1186 ) -> Result<PendingData<'_>> {
1187 self.0.reserve_data(payload_len)
1188 }
1189
1190 pub fn finish(&mut self) -> Result<()> {
1191 self.0.finish()
1192 }
1193
1194 pub fn finish_empty(&mut self) -> Result<()> {
1197 if self.0.state != ProducerState::New {
1198 return Err(Error::Malformed(
1199 "empty FIN requires an unused flow".to_owned(),
1200 ));
1201 }
1202 self.0.control.finish_quietly();
1203 self.0.state = ProducerState::Terminal;
1204 Ok(())
1205 }
1206
1207 pub fn reset(
1208 &mut self,
1209 code: ResetCode,
1210 ) -> Result<()> {
1211 self.0.reset(code)
1212 }
1213
1214 pub async fn peer_closed(&mut self) -> Result<()> {
1217 self.0.control.peer_closed().await
1218 }
1219
1220 pub fn wait_writable(&self) -> Result<()> {
1221 self.0.control.wait_writable_exact(CONTROL_FRAME_BYTES)
1222 }
1223
1224 pub fn poll_writable(
1225 &self,
1226 cx: &mut std::task::Context<'_>,
1227 ) -> std::task::Poll<Result<()>> {
1228 self.0.control.poll_writable_exact(CONTROL_FRAME_BYTES, cx)
1229 }
1230
1231 pub fn wait_ready(
1234 &self,
1235 payload_len: usize,
1236 ) -> Result<()> {
1237 if payload_len != 0 {
1238 self.0.arena.wait_credit(|| self.0.credit_ready(payload_len))?;
1239 }
1240 self.0.control.wait_writable_exact(CONTROL_FRAME_BYTES)?;
1241 self.0.control.exchange_write_status()
1242 }
1243
1244 pub fn poll_ready(
1246 &self,
1247 payload_len: usize,
1248 cx: &mut std::task::Context<'_>,
1249 ) -> std::task::Poll<Result<()>> {
1250 self.0.control.exchange_write_status()?;
1251 let control =
1252 self.0.control.poll_writable_exact(CONTROL_FRAME_BYTES, cx);
1253 let payload = if payload_len == 0 {
1254 std::task::Poll::Ready(Ok(()))
1255 } else {
1256 self.0.poll_credit(payload_len, cx)
1257 };
1258 match (control, payload) {
1259 (std::task::Poll::Ready(Err(error)), _)
1260 | (_, std::task::Poll::Ready(Err(error))) => {
1261 std::task::Poll::Ready(Err(error))
1262 }
1263 (
1264 std::task::Poll::Ready(Ok(())),
1265 std::task::Poll::Ready(Ok(())),
1266 ) => std::task::Poll::Ready(self.0.control.exchange_write_status()),
1267 _ => std::task::Poll::Pending,
1268 }
1269 }
1270}
1271
1272pub struct Receiver(Consumer);
1274
1275impl Receiver {
1276 pub fn id(&self) -> ExchangeId {
1277 self.0.id()
1278 }
1279
1280 pub fn flow(&self) -> Flow {
1281 self.0.flow
1282 }
1283
1284 pub fn cancel(&mut self) {
1287 self.0.cancel();
1288 }
1289
1290 pub fn reset_signal(&self) -> crate::ResetSignal {
1292 self.0.control.reset_signal()
1293 }
1294
1295 pub fn try_next(&mut self) -> Result<FlowEvent> {
1296 self.0.next()
1297 }
1298
1299 pub fn dispatch<H: FlowHandler>(
1300 &mut self,
1301 handler: &mut H,
1302 ) -> std::result::Result<(), DispatchError<H::Error>> {
1303 dispatch(&mut self.0, handler)
1304 }
1305
1306 pub fn wait_readable(&self) -> Result<()> {
1307 if self.0.state == ConsumerState::Terminal {
1308 return Err(Error::Closed);
1309 }
1310 self.0.control.wait_readable()
1311 }
1312
1313 pub fn poll_readable(
1314 &self,
1315 cx: &mut std::task::Context<'_>,
1316 ) -> std::task::Poll<Result<()>> {
1317 if self.0.state == ConsumerState::Terminal {
1318 return std::task::Poll::Ready(Err(Error::Closed));
1319 }
1320 self.0.control.poll_readable(cx)
1321 }
1322}
1323
1324#[cfg(test)]
1325mod tests {
1326 use super::*;
1327
1328 fn exchanges() -> Exchanges {
1329 Exchanges::open(
1330 Arc::new(Fleet::join("exchange-pair-test", 2).expect("fleet")),
1331 Incarnation::new(1),
1332 ExchangeSpec::new(
1333 StreamSpec::new(230, 8, 512),
1334 PayloadArenaSpec::new(231, 8, 64),
1335 ),
1336 )
1337 .expect("exchanges")
1338 }
1339
1340 fn take_data(receiver: &mut Receiver) -> bytes::Bytes {
1341 let Ok(FlowEvent::Data(chunk)) = receiver.try_next() else {
1342 panic!("expected data chunk");
1343 };
1344 bytes::Bytes::from_owner(chunk)
1345 }
1346
1347 #[test]
1348 fn flow_credit_counts_slots_until_the_last_bytes_slice_drops() {
1349 let exchanges = exchanges().with_flow_slots(2).unwrap();
1350 let (mut a, ticket) = exchanges.create().unwrap();
1351 let mut b = exchanges.open_peer(ticket).unwrap();
1352 let (mut sibling, _) = exchanges.create().unwrap();
1353 a.sender().start(None).unwrap();
1354 b.sender().start(None).unwrap();
1355 sibling.sender().start(None).unwrap();
1356 assert!(matches!(b.receiver().try_next(), Ok(FlowEvent::Start { .. })));
1357 a.sender().data(&[7; 65]).unwrap();
1358 let body = take_data(b.receiver());
1359 let slice = body.slice(1..2);
1360 drop(body);
1361
1362 assert!(matches!(
1363 a.sender().data(b"blocked"),
1364 Err(Error::PayloadFull { requested_slots: 1 })
1365 ));
1366 assert!(matches!(
1367 a.sender().data(&[0; 129]),
1368 Err(Error::PayloadTooLarge { capacity: 128, .. })
1369 ));
1370 assert!(a.sender().chunk_geometry().chunk_bytes(3).is_none());
1371 assert!(matches!(
1372 exchanges.create_payload_available(129),
1373 Err(Error::PayloadTooLarge { .. })
1374 ));
1375 sibling.sender().data(b"sibling can send").unwrap();
1376 b.sender().data(b"early response").unwrap();
1377
1378 assert_eq!(&slice[..], &[7]);
1379 drop(slice);
1380 a.sender().data(&[8; 65]).unwrap();
1381 assert_eq!(&take_data(b.receiver())[..], &[8; 65]);
1382 }
1383
1384 struct Signal(Mutex<std::sync::mpsc::Sender<()>>);
1387
1388 impl std::task::Wake for Signal {
1389 fn wake(self: Arc<Self>) {
1390 let _ = lock_unpoisoned(&self.0).send(());
1391 }
1392 }
1393
1394 fn signal() -> (std::task::Waker, std::sync::mpsc::Receiver<()>) {
1395 let (sender, receiver) = std::sync::mpsc::channel();
1396 (Arc::new(Signal(Mutex::new(sender))).into(), receiver)
1397 }
1398
1399 fn park(
1401 sender: &Sender,
1402 waker: &std::task::Waker,
1403 ) {
1404 let mut cx = std::task::Context::from_waker(waker);
1405 assert!(sender.poll_ready(64, &mut cx).is_pending());
1406 }
1407
1408 #[test]
1412 fn a_return_wakes_only_the_flow_it_frees() {
1413 let exchanges = exchanges().with_flow_slots(1).unwrap();
1414 let (mut a, ticket) = exchanges.create().unwrap();
1415 let mut a_peer = exchanges.open_peer(ticket).unwrap();
1416 let (mut b, ticket) = exchanges.create().unwrap();
1417 let mut b_peer = exchanges.open_peer(ticket).unwrap();
1418 for (sender, receiver) in [(&mut a, &mut a_peer), (&mut b, &mut b_peer)]
1419 {
1420 sender.sender().start(None).unwrap();
1421 assert!(matches!(
1422 receiver.receiver().try_next(),
1423 Ok(FlowEvent::Start { .. })
1424 ));
1425 sender.sender().data(&[1; 64]).unwrap();
1426 }
1427 let a_held = take_data(a_peer.receiver());
1428 let b_held = take_data(b_peer.receiver());
1429 let (a_waker, a_woken) = signal();
1430 let (b_waker, _b_woken) = signal();
1431 park(a.sender(), &a_waker);
1432 park(b.sender(), &b_waker);
1433
1434 drop(a_held);
1435 a_woken.recv().unwrap();
1436 let a_key = a.sender().0.credit_key();
1437 let b_key = b.sender().0.credit_key();
1438 let arena = a.sender().0.arena.clone();
1439 assert!(!arena.is_waiting(a_key));
1440 assert!(arena.is_waiting(b_key));
1441 a.sender().data(&[2; 64]).unwrap();
1442 assert_eq!(&b_held[..], &[1; 64]);
1443 }
1444
1445 #[test]
1446 fn a_cancelled_reader_wakes_its_sender_while_its_chunk_is_held() {
1447 let exchanges = exchanges().with_flow_slots(1).unwrap();
1448 let (mut a, ticket) = exchanges.create().unwrap();
1449 let mut peer = exchanges.open_peer(ticket).unwrap();
1450 a.sender().start(None).unwrap();
1451 assert!(matches!(
1452 peer.receiver().try_next(),
1453 Ok(FlowEvent::Start { .. })
1454 ));
1455 a.sender().data(&[1; 64]).unwrap();
1456 let held = take_data(peer.receiver());
1457 let (waker, woken) = signal();
1458 park(a.sender(), &waker);
1459
1460 drop(peer);
1461 woken.recv().unwrap();
1462 let mut cx = std::task::Context::from_waker(&waker);
1463 assert!(matches!(
1464 a.sender().poll_ready(64, &mut cx),
1465 std::task::Poll::Ready(Err(Error::PeerGone))
1466 ));
1467 assert_eq!(&held[..], &[1; 64]);
1468 }
1469
1470 #[test]
1471 fn reset_bypasses_a_full_ring_and_reclaims_only_unread_chunks() {
1472 let exchanges = Exchanges::open(
1473 Arc::new(Fleet::join("exchange-reset-full", 2).unwrap()),
1474 Incarnation::new(1),
1475 ExchangeSpec::new(
1476 StreamSpec::new(224, 8, CONTROL_FRAME_BYTES),
1477 PayloadArenaSpec::new(225, 2, 64),
1478 ),
1479 )
1480 .unwrap();
1481 let (mut a, ticket) = exchanges.create().unwrap();
1482 let mut b = exchanges.open_peer(ticket).unwrap();
1483 a.sender().start(None).unwrap();
1484 assert!(matches!(b.receiver().try_next(), Ok(FlowEvent::Start { .. })));
1485 a.sender().data(&[1; 64]).unwrap();
1486 let held = take_data(b.receiver());
1487 a.sender().data(&[2; 64]).unwrap();
1488 assert!(matches!(a.sender().finish(), Err(Error::WouldBlock)));
1489 a.sender().reset(ResetCode::new(0)).unwrap();
1490 assert!(
1491 matches!(b.receiver().try_next(), Ok(FlowEvent::Reset(code)) if code.get() == 0)
1492 );
1493
1494 let (mut next, ticket) = exchanges.create().unwrap();
1495 let mut peer = exchanges.open_peer(ticket).unwrap();
1496 next.sender().start(Some(&[3; 64])).unwrap();
1497 assert!(matches!(
1498 peer.receiver().try_next(),
1499 Ok(FlowEvent::Start { .. })
1500 ));
1501 assert_eq!(&held[..], &[1; 64]);
1502 }
1503
1504 #[test]
1505 fn receiver_cancel_drains_a_finished_sender_without_invalidating_held_data()
1506 {
1507 let exchanges = exchanges();
1508 let (a, ticket) = exchanges.create().unwrap();
1509 let (mut sender, _reverse) = a.split();
1510 let mut b = exchanges.open_peer(ticket).unwrap();
1511 sender.start(None).unwrap();
1512 assert!(matches!(b.receiver().try_next(), Ok(FlowEvent::Start { .. })));
1513 sender.data(&[1; 64]).unwrap();
1514 let held = take_data(b.receiver());
1515 sender.data(&[2; 448]).unwrap();
1516 sender.finish().unwrap();
1517 drop(sender);
1518
1519 b.receiver().cancel();
1520 assert!(matches!(b.receiver().try_next(), Err(Error::Closed)));
1521 let (mut next, _) = exchanges.create().unwrap();
1522 next.sender().start(Some(&[3; 448])).unwrap();
1523 assert_eq!(&held[..], &[1; 64]);
1524 }
1525
1526 #[test]
1527 fn finish_keeps_a_ticket_open_until_the_last_local_half_drops() {
1528 let exchanges = exchanges();
1529 let (a, ticket) = exchanges.create().unwrap();
1530 let (mut sender, reverse) = a.split();
1531 sender.start(Some(&[1; 512])).unwrap();
1532 sender.finish().unwrap();
1533 drop(sender);
1534 let mut late_peer = exchanges.open_peer(ticket).unwrap();
1535 assert!(
1536 matches!(late_peer.receiver().try_next(), Ok(FlowEvent::Start { metadata: Some(data) }) if data.as_slice() == [1; 512])
1537 );
1538 drop(late_peer);
1539 drop(reverse);
1540
1541 let (mut abandoned, ticket) = exchanges.create().unwrap();
1542 abandoned.sender().start(Some(&[2; 512])).unwrap();
1543 abandoned.sender().finish().unwrap();
1544 drop(abandoned);
1545 assert!(matches!(exchanges.open_peer(ticket), Err(Error::Stale(_))));
1546 assert!(exchanges.create_payload_available(512).unwrap());
1547 }
1548
1549 #[test]
1550 fn cancellation_racing_publication_does_not_lose_credit() {
1551 let exchanges = exchanges();
1552 for _ in 0..24 {
1553 let (a, ticket) = exchanges.create().unwrap();
1554 let b = exchanges.open_peer(ticket).unwrap();
1555 let (mut sender, _reverse) = a.split();
1556 let (_reverse, mut receiver) = b.split();
1557 sender.start(None).unwrap();
1558
1559 let ready = std::sync::Barrier::new(2);
1560 let ready_ref = &ready;
1561 std::thread::scope(|scope| {
1562 let mut pending = sender.reserve_data(512).unwrap();
1563 pending.fill(7);
1564 let cancelled = scope.spawn(move || {
1565 ready_ref.wait();
1566 receiver.cancel();
1567 });
1568 ready.wait();
1569 assert!(matches!(
1570 pending.commit(),
1571 Ok(_) | Err(Error::PeerGone)
1572 ));
1573 cancelled.join().unwrap();
1574 assert!(exchanges.create_payload_available(512).unwrap());
1575 });
1576 }
1577 }
1578
1579 #[tokio::test]
1580 async fn a_credit_wait_wakes_on_the_last_bytes_drop_or_reader_cancel() {
1581 use std::future::poll_fn;
1582
1583 let exchanges = exchanges().with_flow_slots(1).unwrap();
1584 let (mut a, ticket) = exchanges.create().unwrap();
1585 let mut b = exchanges.open_peer(ticket).unwrap();
1586 a.sender().start(None).unwrap();
1587 assert!(matches!(b.receiver().try_next(), Ok(FlowEvent::Start { .. })));
1588
1589 for cancel in [false, true] {
1590 a.sender().data(b"held").unwrap();
1591 let held = take_data(b.receiver());
1592 let (armed, registered) = tokio::sync::oneshot::channel();
1593 let mut armed = Some(armed);
1594 let sender = a.sender();
1595 let ready = async {
1596 poll_fn(|cx| {
1597 let result = sender.poll_ready(1, cx);
1598 if let Some(armed) = armed.take() {
1599 assert!(result.is_pending());
1600 armed.send(()).unwrap();
1601 }
1602 result
1603 })
1604 .await
1605 };
1606 let release = async {
1607 registered.await.unwrap();
1608 if cancel {
1609 b.receiver().cancel();
1610 Some(held)
1611 } else {
1612 drop(held);
1613 None
1614 }
1615 };
1616 let (result, still_held) =
1617 tokio::time::timeout(Duration::from_secs(5), async {
1618 tokio::join!(ready, release)
1619 })
1620 .await
1621 .expect("writer must wake");
1622 if cancel {
1623 assert!(matches!(result, Err(Error::PeerGone)));
1624 assert_eq!(&still_held.unwrap()[..], b"held");
1625 } else {
1626 assert!(matches!(result, Ok(())));
1627 }
1628 }
1629 }
1630
1631 #[tokio::test]
1632 async fn cancellation_wakes_a_sender_waiting_on_application_work() {
1633 let exchanges = exchanges();
1634 let (mut a, ticket) = exchanges.create().unwrap();
1635 let mut b = exchanges.open_peer(ticket).unwrap();
1636 let (armed, registered) = tokio::sync::oneshot::channel();
1637 let mut armed = Some(armed);
1638 let mut closed = std::pin::pin!(a.sender().peer_closed());
1639 let wait = std::future::poll_fn(|cx| {
1640 let result = std::future::Future::poll(closed.as_mut(), cx);
1641 if let Some(armed) = armed.take() {
1642 assert!(result.is_pending());
1643 armed.send(()).unwrap();
1644 }
1645 result
1646 });
1647 let cancel = async {
1648 registered.await.unwrap();
1649 b.receiver().cancel();
1650 assert!(matches!(b.receiver().wait_readable(), Err(Error::Closed)));
1651 };
1652 let (result, ()) =
1653 tokio::time::timeout(Duration::from_secs(5), async {
1654 tokio::join!(wait, cancel)
1655 })
1656 .await
1657 .expect("writer must wake");
1658 result.unwrap();
1659 }
1660
1661 #[test]
1662 fn peer_can_send_before_the_first_flow_finishes() {
1663 let exchanges = exchanges();
1664 let (server, ticket) = exchanges.create().expect("server");
1665 let client = exchanges.open_peer(ticket).expect("client");
1666 let (mut request_out, mut response_in) = server.split();
1667 let (mut response_out, mut request_in) = client.split();
1668
1669 request_out.start(Some(b"request-head")).expect("request start");
1670 let request_descriptor =
1671 request_out.data(&[7_u8; 192]).expect("request data");
1672 let request_metadata =
1673 match request_in.try_next().expect("request start event") {
1674 FlowEvent::Start { metadata: Some(metadata) } => metadata,
1675 _ => panic!("expected request start"),
1676 };
1677 assert_eq!(&*request_metadata, b"request-head");
1678 let request_chunk =
1679 match request_in.try_next().expect("request data event") {
1680 FlowEvent::Data(chunk) => chunk,
1681 _ => panic!("expected request data"),
1682 };
1683 assert_eq!(request_descriptor.arena_kind(), 231);
1684 assert_eq!(request_descriptor.slot_count(), 3);
1685
1686 response_out
1689 .start(Some(b"response-head"))
1690 .expect("early response start");
1691 let response_descriptor =
1692 response_out.data(b"response").expect("response data");
1693 assert_eq!(response_descriptor.arena_kind(), 231);
1694 assert!(matches!(response_in.try_next(), Ok(FlowEvent::Start { .. })));
1695 let response_chunk =
1696 match response_in.try_next().expect("response data event") {
1697 FlowEvent::Data(chunk) => chunk,
1698 _ => panic!("expected response data"),
1699 };
1700 assert_eq!(&*response_chunk, b"response");
1701
1702 drop(request_chunk);
1703 request_out.finish().expect("request fin");
1704 assert!(matches!(request_in.try_next(), Ok(FlowEvent::Fin)));
1705 response_out.finish().expect("response fin");
1706 assert!(matches!(response_in.try_next(), Ok(FlowEvent::Fin)));
1707 }
1708
1709 #[test]
1710 fn directional_payloads_keep_each_flows_credit_and_kind_separate() {
1711 let exchanges = Exchanges::open(
1712 Arc::new(
1713 Fleet::join("exchange-directional-test", 2).expect("fleet"),
1714 ),
1715 Incarnation::new(1),
1716 ExchangeSpec::directional(
1717 StreamSpec::new(220, 8, 512),
1718 PayloadArenaSpec::new(221, 1, 64),
1719 PayloadArenaSpec::new(222, 1, 128),
1720 ),
1721 )
1722 .expect("directional exchanges");
1723 let (mut side_a, ticket) = exchanges.create().expect("side A");
1724 let mut side_b = exchanges.open_peer(ticket).expect("side B");
1725
1726 let request_geometry = side_a.sender().chunk_geometry();
1727 let response_geometry = side_b.sender().chunk_geometry();
1728 assert_eq!(request_geometry.slot_payload_bytes(), 64);
1729 assert_eq!(response_geometry.slot_payload_bytes(), 128);
1730 assert_eq!(request_geometry.chunk_bytes(1), Some(64));
1731 assert_eq!(response_geometry.chunk_bytes(1), Some(128));
1732
1733 side_a.sender().start(None).expect("A-to-B start");
1734 assert!(matches!(
1735 side_b.receiver().try_next(),
1736 Ok(FlowEvent::Start { .. })
1737 ));
1738 side_b.sender().start(None).expect("B-to-A start");
1739 assert!(matches!(
1740 side_a.receiver().try_next(),
1741 Ok(FlowEvent::Start { .. })
1742 ));
1743
1744 let a = side_a.sender().data(&[1; 64]).expect("A-to-B payload");
1745 let b = side_b.sender().data(&[2; 128]).expect("B-to-A payload");
1746 assert_eq!(a.arena_kind(), 221);
1747 assert_eq!(b.arena_kind(), 222);
1748
1749 let a_payload = match side_b.receiver().try_next().expect("A-to-B data")
1750 {
1751 FlowEvent::Data(chunk) => chunk,
1752 _ => panic!("expected A-to-B data"),
1753 };
1754 let b_payload = match side_a.receiver().try_next().expect("B-to-A data")
1755 {
1756 FlowEvent::Data(chunk) => chunk,
1757 _ => panic!("expected B-to-A data"),
1758 };
1759 assert_eq!(&*a_payload, &[1; 64]);
1760 assert_eq!(&*b_payload, &[2; 128]);
1761 }
1762
1763 #[test]
1764 fn a_reserved_chunk_is_one_descriptor_over_many_slots() {
1765 let exchanges = Exchanges::open(
1766 Arc::new(
1767 Fleet::join("exchange-reservation-test", 2).expect("fleet"),
1768 ),
1769 Incarnation::new(1),
1770 ExchangeSpec::new(
1771 StreamSpec::new(233, 8, 512),
1772 PayloadArenaSpec::new(234, 8, 256),
1773 ),
1774 )
1775 .expect("exchanges");
1776 let (mut server, ticket) = exchanges.create().expect("server");
1777 let mut client = exchanges.open_peer(ticket).expect("client");
1778
1779 let expected: Vec<u8> =
1780 (0..777).map(|index| (index % 251) as u8).collect();
1781 let descriptor = {
1782 let mut pending =
1783 server.sender().reserve_start(expected.len()).expect("reserve");
1784 pending.copy_from_slice(&expected);
1785 pending.commit().expect("commit")
1786 };
1787 assert_eq!(descriptor.payload_len(), 777);
1788 assert_eq!(descriptor.slot_count(), 4);
1789 let received = match client.receiver().try_next().expect("start") {
1790 FlowEvent::Start { metadata: Some(metadata) } => metadata,
1791 _ => panic!("expected start metadata"),
1792 };
1793 assert_eq!(&*received, expected);
1794 }
1795
1796 #[test]
1797 fn dropping_an_uncommitted_chunk_returns_credit_without_an_event() {
1798 let exchanges = exchanges();
1799 let (mut server, ticket) = exchanges.create().expect("server");
1800 let mut client = exchanges.open_peer(ticket).expect("client");
1801 server.sender().start(None).expect("start");
1802 assert!(matches!(
1803 client.receiver().try_next(),
1804 Ok(FlowEvent::Start { metadata: None })
1805 ));
1806
1807 let pending =
1808 server.sender().reserve_data(256).expect("reserve all slots");
1809 drop(pending);
1810 assert!(matches!(client.receiver().try_next(), Err(Error::WouldBlock)));
1811 assert!(server.sender().reserve_data(256).is_ok());
1812 }
1813
1814 #[test]
1815 fn a_failed_control_commit_returns_its_payload_credit() {
1816 let exchanges = Exchanges::open(
1817 Arc::new(
1818 Fleet::join("exchange-failed-commit-test", 2).expect("fleet"),
1819 ),
1820 Incarnation::new(1),
1821 ExchangeSpec::new(
1822 StreamSpec::new(236, 8, CONTROL_FRAME_BYTES),
1823 PayloadArenaSpec::new(237, 1, 64),
1824 ),
1825 )
1826 .expect("exchanges");
1827 let (mut server, ticket) = exchanges.create().expect("server");
1828 let mut client = exchanges.open_peer(ticket).expect("client");
1829
1830 server.sender().start(None).expect("start fills control ring");
1831 let mut pending = server
1832 .sender()
1833 .reserve_data(64)
1834 .expect("reserve only payload slot");
1835 pending.fill(7);
1836 assert!(matches!(pending.commit(), Err(Error::WouldBlock)));
1837
1838 assert!(matches!(
1839 client.receiver().try_next(),
1840 Ok(FlowEvent::Start { metadata: None })
1841 ));
1842 let mut retry = server
1843 .sender()
1844 .reserve_data(64)
1845 .expect("failed commit returned credit");
1846 retry.fill(9);
1847 retry.commit().expect("retry commit");
1848 let received = match client.receiver().try_next().expect("retry data") {
1849 FlowEvent::Data(chunk) => chunk,
1850 _ => panic!("expected retry data"),
1851 };
1852 assert_eq!(&*received, &[9; 64]);
1853 }
1854
1855 #[test]
1856 fn a_short_io_commit_returns_unused_trailing_slots() {
1857 let exchanges = exchanges();
1858 let (mut server, ticket) = exchanges.create().expect("server");
1859 let mut client = exchanges.open_peer(ticket).expect("client");
1860 server.sender().start(None).expect("start");
1861 assert!(matches!(
1862 client.receiver().try_next(),
1863 Ok(FlowEvent::Start { metadata: None })
1864 ));
1865
1866 let descriptor = {
1867 let mut pending =
1868 server.sender().reserve_data(256).expect("reserve IO capacity");
1869 pending[..37].fill(5);
1870 pending.commit_len(37).expect("commit bytes read")
1871 };
1872 assert_eq!(descriptor.payload_len(), 37);
1873 assert_eq!(descriptor.slot_count(), 1);
1874
1875 let mut following =
1876 server.sender().reserve_data(192).expect("trailing slots returned");
1877 following.fill(9);
1878 following.commit().expect("commit following chunk");
1879 let first = match client.receiver().try_next().expect("short data") {
1880 FlowEvent::Data(chunk) => chunk,
1881 _ => panic!("expected short data"),
1882 };
1883 assert_eq!(&*first, &[5; 37]);
1884 let second = match client.receiver().try_next().expect("following data")
1885 {
1886 FlowEvent::Data(chunk) => chunk,
1887 _ => panic!("expected following data"),
1888 };
1889 assert_eq!(&*second, &[9; 192]);
1890 }
1891
1892 #[test]
1893 fn reset_is_terminal_for_only_its_own_flow() {
1894 let exchanges = exchanges();
1895 let (server, ticket) = exchanges.create().expect("server");
1896 let client = exchanges.open_peer(ticket).expect("client");
1897 let (mut request_out, mut response_in) = server.split();
1898 let (mut response_out, mut request_in) = client.split();
1899
1900 request_out.reset(ResetCode::new(41)).expect("request reset");
1901 assert!(matches!(
1902 request_in.try_next(),
1903 Ok(FlowEvent::Reset(code)) if code.get() == 41
1904 ));
1905
1906 response_out.start(None).expect("response remains independent");
1907 assert!(matches!(
1908 response_in.try_next(),
1909 Ok(FlowEvent::Start { metadata: None })
1910 ));
1911 }
1912}