cloud_sdk/transport/streaming/io/
blocking.rs

1use super::{
2    AbortGuard, BlockingStreamSink, BlockingStreamSource, ScratchGuard, StreamExecutionError,
3    StreamRead,
4};
5use crate::transport::{
6    StreamAttempt, StreamCompletion, StreamOutcome, StreamPartialState, StreamPolicy,
7    StreamProgressError,
8};
9
10/// Drives one finite blocking transfer with caller-owned scratch storage.
11pub fn drive_blocking_stream<S, D>(
12    policy: StreamPolicy,
13    source: &mut S,
14    sink: &mut D,
15    scratch: &mut [u8],
16    outcome: &mut StreamOutcome,
17) -> Result<StreamCompletion, StreamExecutionError<S::Error, D::Error>>
18where
19    S: BlockingStreamSource,
20    D: BlockingStreamSink,
21{
22    *outcome = StreamOutcome::new();
23    if scratch.is_empty() {
24        return Err(StreamExecutionError::EmptyScratch);
25    }
26    let mut scratch = ScratchGuard::new(scratch);
27    let mut guard = AbortGuard::new(sink, abort_blocking::<D>);
28    let mut attempt = StreamAttempt::new(policy, outcome);
29    let completion = pump(policy, source, &mut guard, &mut scratch, &mut attempt)?;
30    if let Err(error) = guard.sink().commit() {
31        attempt.mark_failed();
32        return Err(StreamExecutionError::Sink(error));
33    }
34    attempt
35        .commit_sink()
36        .map_err(StreamExecutionError::Progress)?;
37    guard.disarm();
38    Ok(completion)
39}
40
41fn pump<S, D>(
42    policy: StreamPolicy,
43    source: &mut S,
44    sink: &mut AbortGuard<'_, D>,
45    scratch: &mut ScratchGuard<'_>,
46    attempt: &mut StreamAttempt<'_>,
47) -> Result<StreamCompletion, StreamExecutionError<S::Error, D::Error>>
48where
49    S: BlockingStreamSource,
50    D: BlockingStreamSink,
51{
52    loop {
53        attempt
54            .begin_source_observation()
55            .map_err(StreamExecutionError::Progress)?;
56        let read_limit = core::cmp::min(scratch.bytes().len(), policy.limits().max_chunk_bytes());
57        let Some(output) = scratch.bytes().get_mut(..read_limit) else {
58            attempt.mark_failed();
59            return Err(StreamExecutionError::EmptyScratch);
60        };
61        let read = source.read_chunk(output).map_err(|error| {
62            attempt.mark_failed();
63            StreamExecutionError::Source(error)
64        })?;
65        match read {
66            StreamRead::End => return attempt.finish().map_err(StreamExecutionError::Progress),
67            StreamRead::Wait => attempt
68                .observe_wait()
69                .map_err(StreamExecutionError::Progress)?,
70            StreamRead::Chunk(len) => {
71                if len > output.len() {
72                    attempt.mark_failed();
73                    return Err(StreamExecutionError::InvalidSourceLength);
74                }
75                transfer(policy, sink, output, len, attempt)?;
76            }
77        }
78    }
79}
80
81fn transfer<S, D>(
82    policy: StreamPolicy,
83    sink: &mut AbortGuard<'_, D>,
84    chunk: &[u8],
85    len: usize,
86    attempt: &mut StreamAttempt<'_>,
87) -> Result<(), StreamExecutionError<S, D::Error>>
88where
89    D: BlockingStreamSink,
90{
91    attempt
92        .begin_chunk(len)
93        .map_err(StreamExecutionError::Progress)?;
94    let mut offset = 0_usize;
95    while offset < len {
96        let Some(input) = chunk.get(offset..len) else {
97            attempt.mark_failed();
98            return Err(StreamExecutionError::Progress(
99                StreamProgressError::ArithmeticOverflow,
100            ));
101        };
102        attempt
103            .begin_sink_observation()
104            .map_err(StreamExecutionError::Progress)?;
105        sink.record_write_attempt(policy);
106        let accepted = sink.sink().write_chunk(input).map_err(|error| {
107            attempt.mark_failed();
108            StreamExecutionError::Sink(error)
109        })?;
110        attempt
111            .advance(accepted)
112            .map_err(StreamExecutionError::Progress)?;
113        offset = offset.checked_add(accepted).ok_or_else(|| {
114            attempt.mark_failed();
115            StreamExecutionError::Progress(StreamProgressError::ArithmeticOverflow)
116        })?;
117    }
118    Ok(())
119}
120
121fn abort_blocking<S: BlockingStreamSink>(sink: &mut S, state: StreamPartialState) {
122    sink.abort(state);
123}
cloud_sdk/transport/streaming/io/blocking.rs

cloud_sdk/transport/streaming/io/
blocking.rs
