laminar_core/shuffle/transport/grpc/
flow.rs

1//! Outbound byte reservation and enqueue, and what they report: how long sends to each peer
2//! waited, what they carried, and how much of each node byte budget is held. An optional stall
3//! timeout fails a data send whose wait makes no progress.
4
5use std::time::{Duration, Instant};
6
7use tokio::sync::TryAcquireError;
8
9use super::*;
10
11/// Totals of one peer's outbound shuffle stream since its sender was created. They survive
12/// reconnects to the peer.
13#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
14pub struct ShufflePeerFlow {
15    /// The peer's node id.
16    pub peer: ShufflePeerId,
17    /// Time sends waited for the peer's or the node's outbound byte budget, or for a slot in the
18    /// peer's stream queue. Sends to one peer take turns and the turn is not counted, so this
19    /// tracks how long the peer's stream was blocked. A send cancelled while waiting still counts.
20    pub send_wait: Duration,
21    /// Arrow IPC payload bytes of data frames queued to the peer's stream; control frames are
22    /// excluded.
23    pub sent_bytes: u64,
24    /// Rows of those data frames.
25    pub sent_rows: u64,
26}
27
28#[derive(Default)]
29pub(super) struct PeerSendCounters {
30    waited_ns: AtomicU64,
31    sent_bytes: AtomicU64,
32    sent_rows: AtomicU64,
33}
34
35/// Adds its lifetime to the peer's wait on drop, so every exit of a wait is counted.
36struct WaitClock<'a> {
37    counters: &'a PeerSendCounters,
38    since: Instant,
39}
40
41impl Drop for WaitClock<'_> {
42    fn drop(&mut self) {
43        let waited = u64::try_from(self.since.elapsed().as_nanos()).unwrap_or(u64::MAX);
44        self.counters.waited_ns.fetch_add(waited, Ordering::Relaxed);
45    }
46}
47
48impl PeerSendCounters {
49    fn start_wait(&self) -> WaitClock<'_> {
50        WaitClock {
51            counters: self,
52            since: Instant::now(),
53        }
54    }
55}
56
57#[derive(Default)]
58pub(super) struct SendFlow {
59    peers: Mutex<FxHashMap<ShufflePeerId, Arc<PeerSendCounters>>>,
60    stall_timeout: Option<Duration>,
61}
62
63#[derive(Debug)]
64struct SendStalled {
65    peer: ShufflePeerId,
66    waiting_for: &'static str,
67    timeout: Duration,
68}
69
70impl std::fmt::Display for SendStalled {
71    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72        write!(
73            f,
74            "shuffle send to peer {} waited {:?} for {} without progress",
75            self.peer, self.timeout, self.waiting_for
76        )
77    }
78}
79
80impl std::error::Error for SendStalled {}
81
82/// Whether a send failed because one of its waits lasted the stall timeout set with
83/// [`ShuffleSender::with_send_stall_timeout`].
84#[must_use]
85pub fn is_shuffle_send_stalled(error: &io::Error) -> bool {
86    error
87        .get_ref()
88        .and_then(|source| source.downcast_ref::<SendStalled>())
89        .is_some()
90}
91
92/// What one wait of a send is for, and how long it may last.
93#[derive(Clone, Copy)]
94struct StallLimit {
95    peer: ShufflePeerId,
96    waiting_for: &'static str,
97    timeout: Option<Duration>,
98}
99
100impl StallLimit {
101    /// Resolves with the stall error once the timeout passes, and never without one.
102    async fn elapsed(self) -> io::Error {
103        let Some(timeout) = self.timeout else {
104            return std::future::pending().await;
105        };
106        tokio::time::sleep(timeout).await;
107        io::Error::new(
108            io::ErrorKind::TimedOut,
109            SendStalled {
110                peer: self.peer,
111                waiting_for: self.waiting_for,
112                timeout,
113            },
114        )
115    }
116}
117
118pub(super) struct OutboundReservation {
119    pub(super) peer: OwnedSemaphorePermit,
120    pub(super) node: OwnedSemaphorePermit,
121}
122
123impl OutboundReservation {
124    pub(super) fn shrink_to(&mut self, retained_bytes: usize) -> io::Result<()> {
125        let peer_bytes = self.peer.num_permits();
126        let node_bytes = self.node.num_permits();
127        if peer_bytes != node_bytes || retained_bytes > peer_bytes {
128            return Err(io::Error::other(
129                "shuffle outbound reservation accounting mismatch",
130            ));
131        }
132        let release = peer_bytes - retained_bytes;
133        if release != 0 {
134            let peer = self
135                .peer
136                .split(release)
137                .expect("validated outbound peer reservation split");
138            let node = self
139                .node
140                .split(release)
141                .expect("validated outbound node reservation split");
142            drop((peer, node));
143        }
144        Ok(())
145    }
146}
147
148/// Only a grant that cannot be made at once reads the clock.
149async fn acquire(
150    budget: &Arc<Semaphore>,
151    permits: u32,
152    cancel: &CancellationToken,
153    counters: &PeerSendCounters,
154    stall: StallLimit,
155    closed: &'static str,
156) -> io::Result<OwnedSemaphorePermit> {
157    if cancel.is_cancelled() {
158        return Err(scope_cancelled_io());
159    }
160    match Arc::clone(budget).try_acquire_many_owned(permits) {
161        Ok(permit) => return Ok(permit),
162        Err(TryAcquireError::Closed) => {
163            return Err(io::Error::new(io::ErrorKind::BrokenPipe, closed));
164        }
165        Err(TryAcquireError::NoPermits) => {}
166    }
167    let _wait = counters.start_wait();
168    tokio::select! {
169        biased;
170        () = cancel.cancelled() => Err(scope_cancelled_io()),
171        permit = Arc::clone(budget).acquire_many_owned(permits) => {
172            permit.map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, closed))
173        }
174        error = stall.elapsed() => Err(error),
175    }
176}
177
178impl ShuffleSender {
179    /// Fail a data or frontier send when one of its waits lasts `timeout`: the wait for the
180    /// peer's or the node's outbound byte budget, or for a slot in the peer's stream queue. A
181    /// wait ends once the streams have drained room for one frame, so a wait that long means
182    /// they stopped draining. [`is_shuffle_send_stalled`] recognises the error. Checkpoint
183    /// barriers are bounded by their checkpoint deadline instead. Unset, a send waits for as
184    /// long as its peer stays slow.
185    #[must_use]
186    pub fn with_send_stall_timeout(mut self, timeout: Option<Duration>) -> Self {
187        self.flow.stall_timeout = timeout;
188        self
189    }
190
191    fn stall_limit(
192        &self,
193        peer: ShufflePeerId,
194        msg: &ShuffleMessage,
195        waiting_for: &'static str,
196    ) -> StallLimit {
197        StallLimit {
198            peer,
199            waiting_for,
200            timeout: self
201                .flow
202                .stall_timeout
203                .filter(|_| !matches!(msg, ShuffleMessage::Barrier(_))),
204        }
205    }
206
207    /// Per-peer totals, ascending by peer.
208    #[must_use]
209    pub fn peer_flow(&self) -> Vec<ShufflePeerFlow> {
210        let mut flows = self
211            .flow
212            .peers
213            .lock()
214            .iter()
215            .map(|(&peer, counters)| ShufflePeerFlow {
216                peer,
217                send_wait: Duration::from_nanos(counters.waited_ns.load(Ordering::Relaxed)),
218                sent_bytes: counters.sent_bytes.load(Ordering::Relaxed),
219                sent_rows: counters.sent_rows.load(Ordering::Relaxed),
220            })
221            .collect::<Vec<_>>();
222        flows.sort_unstable_by_key(|flow| flow.peer);
223        flows
224    }
225
226    /// Bytes of the node's outbound data budget held by frames being encoded or queued,
227    /// including the share already granted to a send still waiting for the rest. The separate
228    /// checkpoint-barrier budget is excluded.
229    #[must_use]
230    pub fn outbound_budget_used_bytes(&self) -> usize {
231        OUTBOUND_NODE_BUDGET_BYTES.saturating_sub(self.node_budget.available_permits())
232    }
233
234    pub(super) fn counters_for(&self, peer: ShufflePeerId) -> Arc<PeerSendCounters> {
235        Arc::clone(self.flow.peers.lock().entry(peer).or_default())
236    }
237
238    /// Reserve the peer's and then the node's byte budget: the control budgets for a barrier,
239    /// the data budgets otherwise.
240    pub(super) async fn reserve_outbound(
241        &self,
242        conn: &PeerConn,
243        peer: ShufflePeerId,
244        msg: &ShuffleMessage,
245        admission_bytes: usize,
246        cancel: &CancellationToken,
247    ) -> io::Result<OutboundReservation> {
248        let permits = u32::try_from(admission_bytes).map_err(|_| {
249            io::Error::new(
250                io::ErrorKind::InvalidInput,
251                "shuffle admission is too large",
252            )
253        })?;
254        let (peer_budget, node_budget) = if matches!(msg, ShuffleMessage::Barrier(_)) {
255            (
256                &conn.control_byte_budget,
257                &self.checkpointed_control_node_budget,
258            )
259        } else {
260            (&conn.byte_budget, &self.node_budget)
261        };
262        let counters = &conn.counters;
263        let peer_permit = acquire(
264            peer_budget,
265            permits,
266            cancel,
267            counters,
268            self.stall_limit(peer, msg, "its outbound byte budget"),
269            "shuffle byte budget closed",
270        )
271        .await?;
272        let node_permit = acquire(
273            node_budget,
274            permits,
275            cancel,
276            counters,
277            self.stall_limit(peer, msg, "the node's outbound byte budget"),
278            "shuffle node byte budget closed",
279        )
280        .await?;
281        Ok(OutboundReservation {
282            peer: peer_permit,
283            node: node_permit,
284        })
285    }
286
287    /// Queue `out` on the peer's ordered stream, then count what it carried.
288    pub(super) async fn enqueue(
289        &self,
290        conn: &PeerConn,
291        peer: ShufflePeerId,
292        msg: &ShuffleMessage,
293        out: Outbound,
294        cancel: &CancellationToken,
295    ) -> io::Result<()> {
296        let sent_bytes = match &out.msg {
297            PreparedMessage::Data { arrow_ipc, .. } => arrow_ipc.len(),
298            PreparedMessage::Barrier(_) | PreparedMessage::Frontier { .. } => 0,
299        };
300        match conn.tx.try_send(out) {
301            Ok(()) => {}
302            Err(crossfire::TrySendError::Full(out)) => {
303                let _wait = conn.counters.start_wait();
304                tokio::select! {
305                    biased;
306                    () = self.process_lease.wait_until_lost() => {
307                        return Err(process_lease_expired_io());
308                    }
309                    () = cancel.cancelled() => return Err(scope_cancelled_io()),
310                    result = conn.tx.send(out) => result.map_err(|_| io::Error::new(
311                        io::ErrorKind::BrokenPipe,
312                        format!("shuffle stream to peer {peer} closed"),
313                    ))?,
314                    error = self.stall_limit(peer, msg, "a slot in its stream queue").elapsed() => {
315                        return Err(error);
316                    }
317                }
318            }
319            Err(crossfire::TrySendError::Disconnected(_)) => {
320                return Err(io::Error::new(
321                    io::ErrorKind::BrokenPipe,
322                    format!("shuffle stream to peer {peer} closed"),
323                ));
324            }
325        }
326        if let ShuffleMessage::Data { batch, .. } = msg {
327            let counters = &conn.counters;
328            counters
329                .sent_bytes
330                .fetch_add(sent_bytes as u64, Ordering::Relaxed);
331            counters
332                .sent_rows
333                .fetch_add(batch.num_rows() as u64, Ordering::Relaxed);
334        }
335        Ok(())
336    }
337}
338
339impl ShuffleReceiver {
340    /// Bytes of the node's inbound budget held by received frames until the operator reading
341    /// them has consumed them.
342    #[must_use]
343    pub fn inbound_budget_used_bytes(&self) -> usize {
344        INBOUND_NODE_BUDGET_BYTES.saturating_sub(self.inbound_budget.node.available_permits())
345    }
346}
laminar_core/shuffle/transport/grpc/flow.rs

laminar_core/shuffle/transport/grpc/
flow.rs
