harmos_sidecar/transport/
connection.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17use std::collections::BTreeMap;
18use std::future::Future;
19use std::io;
20use std::pin::Pin;
21use std::sync::{Arc, Mutex};
22// Third Party Imports
23use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
24use tokio::sync::{mpsc, watch};
25// Local Imports
26use super::credit::{MINIMUM, WINDOW};
27use super::{Consumption, Credit, Packet, ReceiveCredit, Reservation, Role, read, read_data, write, write_with};
28use crate::Frame;
29
30/* === Constants === */
31
32const OUTGOING_DATA: usize = 8;
33const CONTROL: usize = 4;
34const EARLY_CANCELLATIONS: usize = 256;
35const INCOMING_DATA: usize = WINDOW / MINIMUM + 1;
36
37/* === Macros === */
38
39/* === Type Aliases === */
40
41pub type Io = (Box<dyn AsyncRead + Send + Unpin>, Box<dyn AsyncWrite + Send + Unpin>);
42type Outcome = Result<(), Arc<io::Error>>;
43type IoReader = Box<dyn AsyncRead + Send + Unpin>;
44type IoWriter = Box<dyn AsyncWrite + Send + Unpin>;
45type DriverFuture = Pin<Box<dyn Future<Output = Outcome> + Send>>;
46
47/* === Structs === */
48
49pub struct Connection {
50    pub sender: Sender,
51    pub receiver: Receiver,
52    pub driver: Driver,
53}
54
55#[derive(Clone)]
56pub struct Sender {
57    data: mpsc::Sender<Outgoing>,
58    control: mpsc::Sender<Packet>,
59    state: Arc<State>,
60}
61
62pub struct Receiver {
63    delivered: u64,
64    cancellations: BTreeMap<u64, u64>,
65    data: mpsc::Receiver<Received>,
66    control: mpsc::Receiver<Received>,
67    state: Arc<State>,
68}
69
70pub struct Received {
71    pub cancelled: bool,
72    preceding: Option<u64>,
73    pub frame: Frame,
74    pub consumption: Option<Consumption>,
75}
76
77pub struct Driver {
78    work: DriverFuture,
79    result: Option<Outcome>,
80    state: Arc<State>,
81}
82
83struct Outgoing {
84    packet: Packet,
85    reservation: Reservation,
86}
87
88struct State {
89    queued: Mutex<u64>,
90    written: watch::Sender<u64>,
91    sending: Arc<Credit>,
92    receiving: Arc<ReceiveCredit>,
93    stopped: watch::Sender<bool>,
94    failure: Mutex<Option<Arc<io::Error>>>,
95}
96
97/* === Enums === */
98
99/* === Traits === */
100
101/* === Implementations === */
102
103impl Connection {
104    pub fn new(data: Io, control: Io) -> Self {
105        let state = Arc::new(State {
106            queued: Mutex::new(0),
107            written: watch::channel(0).0,
108            sending: Credit::new(),
109            receiving: ReceiveCredit::new(),
110            stopped: watch::channel(false).0,
111            failure: Mutex::new(None),
112        });
113        let (data_tx, data_rx) = mpsc::channel(OUTGOING_DATA);
114        let (control_tx, control_rx) = mpsc::channel(CONTROL);
115        let (incoming_tx, incoming_rx) = mpsc::channel(INCOMING_DATA);
116        let (incoming_control_tx, incoming_control_rx) = mpsc::channel(CONTROL);
117        let owner = state.clone();
118        let work = Box::pin(async move {
119            let result = tokio::select! {
120                biased;
121                _ = stopped(owner.stopped.subscribe()) => owner.outcome().map_err(io::Error::other),
122                _ = stopped(owner.receiving.closed()) => Err(broken("received frame abandoned")),
123                result = receive_data(data.0, incoming_tx, owner.receiving.clone()) => result,
124                result = receive_control(control.0, incoming_control_tx, owner.sending.clone()) => result,
125                result = send_data(data.1, data_rx, owner.written.clone()) => result,
126                result = send_control(control.1, control_rx, owner.receiving.released()) => result,
127            };
128            owner.stop(result.err());
129            owner.outcome()
130        });
131        Self {
132            sender: Sender {
133                data: data_tx,
134                control: control_tx,
135                state: state.clone(),
136            },
137            receiver: Receiver {
138                delivered: 0,
139                cancellations: BTreeMap::new(),
140                data: incoming_rx,
141                control: incoming_control_rx,
142                state: state.clone(),
143            },
144            driver: Driver { work, result: None, state },
145        }
146    }
147}
148
149impl Sender {
150    pub async fn send(&self, frame: Frame) -> io::Result<()> {
151        let packet = Packet::Frame(frame);
152        tokio::select! {
153            biased;
154            _ = stopped(self.state.stopped.subscribe()) => Err(self.state.error()),
155            result = async {
156                if let Packet::Frame(Frame::Cancel { id }) = packet {
157                    let preceding = *self.state.queued.lock().expect("transport queue order lock");
158                    return self.control.send(Packet::Cancellation { id, preceding }).await.map_err(|_| self.state.error());
159                }
160                let reservation = self.state.sending.reserve(&packet).await?;
161                let slot = self.data.reserve().await.map_err(|_| self.state.error())?;
162                let mut queued = self.state.queued.lock().expect("transport queue order lock");
163                *queued = queued.checked_add(1).ok_or_else(|| super::invalid("transport queue counter overflow"))?;
164                slot.send(Outgoing { packet, reservation });
165                Ok(())
166            } => result,
167        }
168    }
169
170    pub async fn drain(&self) -> io::Result<()> {
171        let through = *self.state.queued.lock().expect("transport queue order lock");
172        let mut written = self.state.written.subscribe();
173        loop {
174            if *written.borrow_and_update() >= through {
175                return Ok(());
176            }
177            tokio::select! {
178                biased;
179                _ = stopped(self.state.stopped.subscribe()) => return Err(self.state.error()),
180                result = written.changed() => result.map_err(|_| self.state.error())?,
181            }
182        }
183    }
184
185    pub fn close(&self) {
186        self.state.stop(None);
187    }
188}
189
190impl Receiver {
191    pub async fn receive(&mut self) -> io::Result<Option<Received>> {
192        let mut received = tokio::select! {
193            biased;
194            // Complete Buffered Frames Precede The Terminal Transport Outcome.
195            Some(frame) = self.control.recv() => frame,
196            Some(frame) = self.data.recv() => frame,
197            _ = stopped(self.state.stopped.subscribe()) => return self.state.outcome().map(|()| None).map_err(io::Error::other),
198        };
199        if let Some(preceding) = received.preceding {
200            if preceding > self.delivered {
201                let Frame::Cancel { id } = received.frame else {
202                    unreachable!("only cancellation has a data watermark");
203                };
204                if !self.cancellations.contains_key(&id) && self.cancellations.len() >= EARLY_CANCELLATIONS {
205                    self.state.stop(Some(super::invalid("early cancellation capacity exhausted")));
206                    return Err(self.state.error());
207                }
208                self.cancellations.entry(id).and_modify(|through| *through = (*through).max(preceding)).or_insert(preceding);
209            }
210        } else {
211            self.delivered = self.delivered.checked_add(1).ok_or_else(|| super::invalid("received data counter overflow"))?;
212            if let Frame::Call { id, .. } = &received.frame {
213                received.cancelled = self.cancellations.remove(id).is_some();
214            }
215            self.cancellations.retain(|_, preceding| *preceding > self.delivered);
216        }
217        Ok(Some(received))
218    }
219}
220
221impl Drop for Receiver {
222    fn drop(&mut self) {
223        self.state.stop(None);
224    }
225}
226
227impl Driver {
228    pub async fn run(&mut self) -> io::Result<()> {
229        if self.result.is_none() {
230            self.result = Some(self.work.as_mut().await);
231        }
232        self.result.as_ref().expect("driver completed").clone().map_err(io::Error::other)
233    }
234}
235
236impl Drop for Driver {
237    fn drop(&mut self) {
238        // Dropping The Retained Future Closes All Four Physical Halves Together.
239        self.state.stop(None);
240    }
241}
242
243impl State {
244    fn stop(&self, error: Option<io::Error>) {
245        let mut failure = self.failure.lock().expect("transport failure lock");
246        if *self.stopped.borrow() {
247            return;
248        }
249        *failure = error.map(Arc::new);
250        self.sending.close();
251        // Buffered Frames Retain Their Consumption Receipts After Transport Closure.
252        self.stopped.send_replace(true);
253    }
254
255    fn outcome(&self) -> Outcome {
256        match &*self.failure.lock().expect("transport failure lock") {
257            Some(error) => Err(error.clone()),
258            None => Ok(()),
259        }
260    }
261
262    fn error(&self) -> io::Error {
263        self.outcome().err().map(io::Error::other).unwrap_or_else(|| broken("sidecar connection closed"))
264    }
265}
266
267/* === Public Functions === */
268
269/* === Private Functions === */
270
271async fn stopped(mut receiver: watch::Receiver<bool>) {
272    while !*receiver.borrow_and_update() {
273        if receiver.changed().await.is_err() {
274            break;
275        }
276    }
277}
278
279async fn receive_data(mut reader: IoReader, incoming: mpsc::Sender<Received>, credit: Arc<ReceiveCredit>) -> io::Result<()> {
280    loop {
281        let Some((Packet::Frame(frame), consumption)) = read_data(&mut reader, &credit).await? else {
282            return Err(broken("sidecar data channel ended"));
283        };
284        incoming
285            .try_send(Received {
286                frame,
287                cancelled: false,
288                preceding: None,
289                consumption: Some(consumption),
290            })
291            .map_err(|_| broken("bounded sidecar data receiver unavailable"))?;
292    }
293}
294
295async fn receive_control(mut reader: IoReader, incoming: mpsc::Sender<Received>, credit: Arc<Credit>) -> io::Result<()> {
296    loop {
297        match read(&mut reader, Role::Control).await? {
298            Some(Packet::Credit(consumed)) => credit.returned(consumed)?,
299            // A Full Queue Backpressures The Control Pipe Without Losing Its Packet.
300            Some(Packet::Cancellation { id, preceding }) => incoming
301                .send(Received {
302                    frame: Frame::Cancel { id },
303                    consumption: None,
304                    cancelled: false,
305                    preceding: Some(preceding),
306                })
307                .await
308                .map_err(|_| broken("sidecar control receiver closed"))?,
309            Some(Packet::Frame(_)) => {
310                return Err(super::invalid("cancellation requires a data watermark"));
311            }
312            None => return Err(broken("sidecar control channel ended")),
313        }
314    }
315}
316
317async fn send_data(mut writer: IoWriter, mut outgoing: mpsc::Receiver<Outgoing>, written: watch::Sender<u64>) -> io::Result<()> {
318    // Wire Credits Bound In-Flight Packets; This Separate Driver-Owned Scratch
319    // Holds At Most One Validated Frame, Released When The Driver Exits.
320    // Scalar Small Frames Reuse It; Large Payloads Remain In Their Own Backing.
321    let mut bytes = Vec::new();
322    while let Some(Outgoing { packet, reservation }) = outgoing.recv().await {
323        reservation.issued()?;
324        write_with(&mut writer, &packet, &mut bytes).await?;
325        let span = tracing::trace_span!(target: "harmos::carrier", "sidecar.frame.flush");
326        tracing::Instrument::instrument(writer.flush(), span).await?;
327        let next = (*written.borrow()).checked_add(1).ok_or_else(|| super::invalid("written data counter overflow"))?;
328        written.send_replace(next);
329    }
330    Err(broken("sidecar data sender closed"))
331}
332
333async fn send_control(mut writer: IoWriter, mut outgoing: mpsc::Receiver<Packet>, mut released: watch::Receiver<u64>) -> io::Result<()> {
334    let initial = *released.borrow_and_update();
335    if initial != 0 {
336        write(&mut writer, &Packet::Credit(initial)).await?;
337        writer.flush().await?;
338    }
339    loop {
340        let packet = tokio::select! {
341            packet = outgoing.recv() => packet.ok_or_else(|| broken("sidecar control sender closed"))?,
342            changed = released.changed() => {
343                changed.map_err(|_| broken("sidecar consumption owner closed"))?;
344                Packet::Credit(*released.borrow_and_update())
345            }
346        };
347        write(&mut writer, &packet).await?;
348        writer.flush().await?;
349    }
350}
351
352fn broken(message: &'static str) -> io::Error {
353    io::Error::new(io::ErrorKind::BrokenPipe, message)
354}
355
356/* === Tests === */
harmos_sidecar/transport/connection.rs

harmos_sidecar/transport/
connection.rs
