cloud_sdk/transport/streaming/io/
asynchronous.rs

1use super::{
2    AbortGuard, AsyncStreamSink, AsyncStreamSource, LocalAsyncStreamSink, LocalAsyncStreamSource,
3    ScratchGuard, StreamExecutionError, StreamRead,
4};
5use crate::transport::{
6    StreamAttempt, StreamCompletion, StreamOutcome, StreamPartialState, StreamPolicy,
7    StreamProgressError,
8};
9use core::{
10    future::Future,
11    pin::Pin,
12    task::{Context, Poll},
13};
14
15const MAX_CALLBACKS_BEFORE_YIELD: u16 = 64;
16
17/// Drives one local asynchronous transfer without owning an executor.
18pub async fn drive_local_stream<S, D>(
19    policy: StreamPolicy,
20    source: &mut S,
21    sink: &mut D,
22    scratch: &mut [u8],
23    outcome: &mut StreamOutcome,
24) -> Result<StreamCompletion, StreamExecutionError<S::Error, D::Error>>
25where
26    S: LocalAsyncStreamSource,
27    D: LocalAsyncStreamSink,
28{
29    *outcome = StreamOutcome::new();
30    if scratch.is_empty() {
31        return Err(StreamExecutionError::EmptyScratch);
32    }
33    let mut scratch = ScratchGuard::new(scratch);
34    let mut guard = AbortGuard::new(sink, abort_local::<D>);
35    let mut attempt = StreamAttempt::new(policy, outcome);
36    let mut cooperation = CooperativeBudget::new();
37    let completion = loop {
38        attempt
39            .begin_source_observation()
40            .map_err(StreamExecutionError::Progress)?;
41        let read_limit = core::cmp::min(scratch.bytes().len(), policy.limits().max_chunk_bytes());
42        let Some(output) = scratch.bytes().get_mut(..read_limit) else {
43            attempt.mark_failed();
44            return Err(StreamExecutionError::EmptyScratch);
45        };
46        let read = match source.read_chunk_local(output).await {
47            Ok(read) => read,
48            Err(error) => {
49                attempt.mark_failed();
50                return Err(StreamExecutionError::Source(error));
51            }
52        };
53        match read {
54            StreamRead::End => {
55                let completion = attempt.finish().map_err(StreamExecutionError::Progress)?;
56                cooperation.after_callback().await;
57                break completion;
58            }
59            StreamRead::Wait => {
60                attempt
61                    .observe_wait()
62                    .map_err(StreamExecutionError::Progress)?;
63                cooperation.after_callback().await;
64            }
65            StreamRead::Chunk(len) => {
66                validate_source_length(len, output.len(), &mut attempt)?;
67                attempt
68                    .begin_chunk(len)
69                    .map_err(StreamExecutionError::Progress)?;
70                cooperation.after_callback().await;
71                let mut offset = 0_usize;
72                while offset < len {
73                    let Some(input) = output.get(offset..len) else {
74                        attempt.mark_failed();
75                        return Err(arithmetic_error());
76                    };
77                    attempt
78                        .begin_sink_observation()
79                        .map_err(StreamExecutionError::Progress)?;
80                    guard.record_write_attempt(policy);
81                    let accepted = match guard.sink().write_chunk_local(input).await {
82                        Ok(accepted) => accepted,
83                        Err(error) => {
84                            attempt.mark_failed();
85                            return Err(StreamExecutionError::Sink(error));
86                        }
87                    };
88                    advance::<S::Error, D::Error>(&mut attempt, &mut offset, accepted)?;
89                    cooperation.after_callback().await;
90                }
91            }
92        }
93    };
94    if let Err(error) = guard.sink().commit_local().await {
95        attempt.mark_failed();
96        return Err(StreamExecutionError::Sink(error));
97    }
98    attempt
99        .commit_sink()
100        .map_err(StreamExecutionError::Progress)?;
101    guard.disarm();
102    Ok(completion)
103}
104
105/// Drives one cross-thread asynchronous transfer without owning an executor.
106pub async fn drive_async_stream<S, D>(
107    policy: StreamPolicy,
108    source: &mut S,
109    sink: &mut D,
110    scratch: &mut [u8],
111    outcome: &mut StreamOutcome,
112) -> Result<StreamCompletion, StreamExecutionError<S::Error, D::Error>>
113where
114    S: AsyncStreamSource + Send,
115    D: AsyncStreamSink + Send,
116{
117    *outcome = StreamOutcome::new();
118    if scratch.is_empty() {
119        return Err(StreamExecutionError::EmptyScratch);
120    }
121    let mut scratch = ScratchGuard::new(scratch);
122    let mut guard = AbortGuard::new(sink, abort_async::<D>);
123    let mut attempt = StreamAttempt::new(policy, outcome);
124    let mut cooperation = CooperativeBudget::new();
125    let completion = loop {
126        attempt
127            .begin_source_observation()
128            .map_err(StreamExecutionError::Progress)?;
129        let read_limit = core::cmp::min(scratch.bytes().len(), policy.limits().max_chunk_bytes());
130        let Some(output) = scratch.bytes().get_mut(..read_limit) else {
131            attempt.mark_failed();
132            return Err(StreamExecutionError::EmptyScratch);
133        };
134        let read = match source.read_chunk(output).await {
135            Ok(read) => read,
136            Err(error) => {
137                attempt.mark_failed();
138                return Err(StreamExecutionError::Source(error));
139            }
140        };
141        match read {
142            StreamRead::End => {
143                let completion = attempt.finish().map_err(StreamExecutionError::Progress)?;
144                cooperation.after_callback().await;
145                break completion;
146            }
147            StreamRead::Wait => {
148                attempt
149                    .observe_wait()
150                    .map_err(StreamExecutionError::Progress)?;
151                cooperation.after_callback().await;
152            }
153            StreamRead::Chunk(len) => {
154                validate_source_length(len, output.len(), &mut attempt)?;
155                attempt
156                    .begin_chunk(len)
157                    .map_err(StreamExecutionError::Progress)?;
158                cooperation.after_callback().await;
159                let mut offset = 0_usize;
160                while offset < len {
161                    let Some(input) = output.get(offset..len) else {
162                        attempt.mark_failed();
163                        return Err(arithmetic_error());
164                    };
165                    attempt
166                        .begin_sink_observation()
167                        .map_err(StreamExecutionError::Progress)?;
168                    guard.record_write_attempt(policy);
169                    let accepted = match guard.sink().write_chunk(input).await {
170                        Ok(accepted) => accepted,
171                        Err(error) => {
172                            attempt.mark_failed();
173                            return Err(StreamExecutionError::Sink(error));
174                        }
175                    };
176                    advance::<S::Error, D::Error>(&mut attempt, &mut offset, accepted)?;
177                    cooperation.after_callback().await;
178                }
179            }
180        }
181    };
182    if let Err(error) = guard.sink().commit().await {
183        attempt.mark_failed();
184        return Err(StreamExecutionError::Sink(error));
185    }
186    attempt
187        .commit_sink()
188        .map_err(StreamExecutionError::Progress)?;
189    guard.disarm();
190    Ok(completion)
191}
192
193fn validate_source_length<S, D>(
194    len: usize,
195    capacity: usize,
196    attempt: &mut StreamAttempt<'_>,
197) -> Result<(), StreamExecutionError<S, D>> {
198    if len > capacity {
199        attempt.mark_failed();
200        return Err(StreamExecutionError::InvalidSourceLength);
201    }
202    Ok(())
203}
204
205fn advance<S, E>(
206    attempt: &mut StreamAttempt<'_>,
207    offset: &mut usize,
208    accepted: usize,
209) -> Result<(), StreamExecutionError<S, E>> {
210    attempt
211        .advance(accepted)
212        .map_err(StreamExecutionError::Progress)?;
213    *offset = offset.checked_add(accepted).ok_or_else(|| {
214        attempt.mark_failed();
215        arithmetic_error()
216    })?;
217    Ok(())
218}
219
220const fn arithmetic_error<S, D>() -> StreamExecutionError<S, D> {
221    StreamExecutionError::Progress(StreamProgressError::ArithmeticOverflow)
222}
223
224fn abort_local<S: LocalAsyncStreamSink>(sink: &mut S, state: StreamPartialState) {
225    sink.abort_local(state);
226}
227
228fn abort_async<S: AsyncStreamSink>(sink: &mut S, state: StreamPartialState) {
229    sink.abort(state);
230}
231
232struct CooperativeBudget {
233    completed_callbacks: u16,
234}
235
236impl CooperativeBudget {
237    const fn new() -> Self {
238        Self {
239            completed_callbacks: 0,
240        }
241    }
242
243    async fn after_callback(&mut self) {
244        if self.completed_callbacks == MAX_CALLBACKS_BEFORE_YIELD - 1 {
245            self.completed_callbacks = 0;
246            YieldOnce { yielded: false }.await;
247        } else if let Some(next) = self.completed_callbacks.checked_add(1) {
248            self.completed_callbacks = next;
249        } else {
250            self.completed_callbacks = 0;
251            YieldOnce { yielded: false }.await;
252        }
253    }
254}
255
256struct YieldOnce {
257    yielded: bool,
258}
259
260impl Future for YieldOnce {
261    type Output = ();
262
263    fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
264        if self.yielded {
265            Poll::Ready(())
266        } else {
267            self.yielded = true;
268            context.waker().wake_by_ref();
269            Poll::Pending
270        }
271    }
272}
cloud_sdk/transport/streaming/io/asynchronous.rs

cloud_sdk/transport/streaming/io/
asynchronous.rs
