Skip to main content

orbit_stream/exchange/
pair.rs

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/// How an exchange budgets payload memory for its two directions.
23///
24/// The shared form is the compact generic default. Directional arenas are for
25/// adapters whose two flows have deliberately different budgets (for example
26/// an upload and a streamed response) or must not consume each other's
27/// credit.
28#[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    /// Physical payload arena used by one byte-flow direction.
36    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/// The physical resources used by an exchange table.
49#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
50pub struct ExchangeSpec {
51    /// One lossless, bounded control table. Its two byte directions carry
52    /// fixed-size typed events, never application payload bytes.
53    pub control: StreamSpec,
54    /// Payload memory may be shared by both directions or split into two
55    /// independently bounded arenas.
56    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    /// Physical payload arena used by one byte-flow direction.
79    pub const fn payload_for(
80        self,
81        flow: Flow,
82    ) -> PayloadArenaSpec {
83        self.payload.for_flow(flow)
84    }
85
86    /// Physical chunk geometry used by one byte-flow direction.
87    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/// The client representative's capability to open the other side of an
127/// exchange. It may travel through an Orbit offer or any setup channel.
128#[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/// One fleet's exchange transport: one duplex control table and either one
220/// shared payload arena or one independently bounded arena per direction.
221#[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    /// Bound outstanding payload independently in each exchange direction.
256    /// By default a direction may use its producer's entire arena lane.
257    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    /// Create side A. Its sender writes A-to-B and its receiver reads B-to-A.
285    /// The ticket opens side B.
286    pub fn create(&self) -> Result<(ExchangeEndpoint, ExchangeTicket)> {
287        let (endpoint, ticket) = self.control.create()?;
288        Ok((self.endpoint(endpoint), ExchangeTicket(ticket)))
289    }
290
291    /// Whether side A can currently reserve its first flow payload.
292    /// Session transports use this before reserving a remote resource so
293    /// payload backpressure does not become an abandoned resource lease.
294    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    /// Task readiness for side A's first flow payload credit.
302    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    /// Open side B. Its sender writes B-to-A and its receiver reads A-to-B.
311    /// Both directions have independent lifecycle and may start in either order.
312    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    /// Park until an exchange is offered to this node.
369    pub fn blocking_take_offer(&self) -> Result<ExchangeTicket> {
370        self.control.blocking_take_offer().map(ExchangeTicket)
371    }
372
373    /// Wait at most `timeout` for an offered exchange.
374    #[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    /// Task readiness for the next offered exchange.
389    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    /// Apply a confirmed process-death report to control ownership and every
406    /// payload arena. This is never inferred from a replacement attaching.
407    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    /// Quiescent-owner maintenance only. Attaching a replacement process
428    /// never resets an exchange table or either payload arena implicitly.
429    pub fn reset_all(&self) {
430        self.control.reset_all();
431        self.payload.reset_all();
432    }
433
434    /// Remove the control and payload SHM names. Existing mappings remain
435    /// valid.
436    #[cfg(unix)]
437    pub fn unlink(&self) -> Result<()> {
438        self.control.unlink()?;
439        self.payload.unlink()
440    }
441}
442
443/// One held exchange side. A worker may hold endpoints on either side of many
444/// exchanges at once; sender and receiver are capabilities, not worker roles.
445pub 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
475// Both local halves keep the ticket openable. If the writer finishes before
476// the peer opens, the last local half seals an unclaimed ticket and reclaims
477// the unread chunks. A claimed peer owns draining its own direction.
478struct 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    /// The exchange status a credit driver reads while this sender is parked.
514    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    /// Task readiness for the payload credit of the next `len` bytes. The
660    /// registration carries what this sender waits for, so the credit driver
661    /// wakes it only when that is there.
662    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        // Either the departing reader saw the committed descriptor while
698        // draining, or this check sees its close and returns the unread extent.
699        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        // A parked registration ends with its sender.
805        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
826/// A flow-start payload reserved directly in the shared SHM arena. The bytes
827/// become immutable and visible only when `commit` sends the matching start
828/// event. Dropping this value returns the reserved slots.
829pub 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    /// Commit only the initialized prefix of a larger reservation. Any whole
844    /// trailing slots that prefix does not need return to the arena first.
845    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
868/// A data chunk reserved directly in the shared SHM arena. One commit publishes
869/// the complete variable-sized chunk and one descriptor.
870pub 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    /// Commit only the initialized prefix of a larger reservation. This is
885    /// the direct-read path for producers that learn the byte count from IO.
886    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                    // Claim with the same generation as a regular read. A
1030                    // concurrent writer abort can win, but a held chunk cannot.
1031                    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
1061/// The event program seen by a flow consumer.
1062pub enum FlowEvent {
1063    Start { metadata: Option<PayloadChunk> },
1064    Data(PayloadChunk),
1065    Fin,
1066    Reset(ResetCode),
1067}
1068
1069/// Typed callbacks for one flow. A handler decides what its events mean; the
1070/// transport does not encode application sequencing policy.
1071pub 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
1114/// The capability to publish one endpoint's outbound flow.
1115pub 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    /// Effective chunk geometry, including this direction's credit limit.
1127    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    /// Plan a known payload using the default chunk target.
1133    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    /// Plan a known payload with a runtime-selected number of slots per chunk.
1147    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    /// Reserve start metadata directly in the shared payload arena. No start
1166    /// event is visible until the returned guard is committed.
1167    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    /// Reserve one variable-sized data chunk directly in the payload arena.
1182    /// One commit emits one data descriptor.
1183    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    /// Close an unused direction. The application already knows from its
1195    /// metadata that there is nothing to receive from this direction.
1196    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    /// Observe cancellation while waiting on application work instead of a
1215    /// control write or a payload reservation. Clean FIN is not cancellation.
1216    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    /// Wait for one control frame and any payload credit needed by the next
1232    /// publication. Use zero for FIN or a start without metadata.
1233    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    /// Task readiness for the same combined control and payload capacity.
1245    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
1272/// The capability to consume the peer endpoint's outbound flow.
1273pub 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    /// Stop this direction and return published chunks that were not taken.
1285    /// Previously returned chunks keep their own credit until dropped.
1286    pub fn cancel(&mut self) {
1287        self.0.cancel();
1288    }
1289
1290    /// Observe a reset independently of the task reading body descriptors.
1291    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    /// Counts its wakes and says so on a channel, so a test waits for a wake
1385    /// rather than for a time.
1386    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    /// Park `sender` on payload credit for one more slot.
1400    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    // The driver takes a pass's records out under its lock and wakes them
1409    // after it: once A's wake arrives, the pass that woke it has decided about
1410    // B too, so B's registration still standing means B was not chosen.
1411    #[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        // Both directions share one arena, but their flow lifecycles are
1687        // independent, so the peer can start before the first flow finishes.
1688        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}