1use crate::{
13 ObjectDescriptor, TransferControl, TransferError as E, TransferResult, TransferSession,
14 VerifiedObject,
15};
16use std::{
17 collections::{HashSet, VecDeque},
18 io::Write,
19 time::Duration,
20};
21pub trait BlockSource: Send + Sync {
25 fn read_block(
28 &self,
29 descriptor: &ObjectDescriptor,
30 index: usize,
31 control: &TransferControl,
32 sink: &mut dyn Write,
33 ) -> TransferResult<()>;
34}
35pub struct SourceBinding<'a> {
37 id: String,
38 source: &'a dyn BlockSource,
39}
40impl<'a> SourceBinding<'a> {
41 pub fn approved(id: impl Into<String>, source: &'a dyn BlockSource) -> TransferResult<Self> {
44 let id = id.into();
45 if id.is_empty() || id.len() > 128 || id.chars().any(char::is_control) {
46 return Err(E::InvalidContract);
47 }
48 Ok(Self { id, source })
49 }
50 pub fn id(&self) -> &str {
52 &self.id
53 }
54}
55#[derive(Debug, Clone, Copy)]
57pub struct TransferOptions {
58 pub max_parallel: usize,
60 pub max_attempts: usize,
62 pub retry_delay: Duration,
64}
65impl Default for TransferOptions {
66 fn default() -> Self {
67 Self {
68 max_parallel: 4,
69 max_attempts: 8,
70 retry_delay: Duration::from_millis(20),
71 }
72 }
73}
74#[derive(Debug)]
76pub struct TransferReport {
77 pub object: VerifiedObject,
79 pub resumed_blocks: usize,
81 pub attempts: usize,
83 pub quarantined_sources: usize,
85 pub source_bytes: Vec<u64>,
87}
88pub fn transfer(
92 mut session: TransferSession<'_>,
93 sources: &[SourceBinding<'_>],
94 options: TransferOptions,
95 control: &TransferControl,
96) -> TransferResult<TransferReport> {
97 validate_sources(sources, options)?;
98 control.checkpoint()?;
99 let descriptor = session.descriptor().clone();
100 let missing = session.missing();
101 let resumed_blocks = descriptor.block_count() - missing.len();
102 let mut pending: VecDeque<_> = missing.into();
103 let mut attempts = vec![0; descriptor.block_count()];
104 let mut quarantined = vec![false; sources.len()];
105 let mut source_bytes = vec![0; sources.len()];
106 let mut cursor = 0;
107 while !pending.is_empty() {
108 control.checkpoint()?;
109 let jobs = select_jobs(
110 &mut pending,
111 &mut attempts,
112 &quarantined,
113 &mut cursor,
114 options,
115 )?;
116 let results = fetch_wave(&descriptor, sources, &jobs, control);
117 let mut retry = false;
118 for ((index, source), result) in jobs.into_iter().zip(results) {
119 control.checkpoint()?;
120 match result {
121 Ok(bytes) => {
122 session.accept(index, &bytes)?;
123 source_bytes[source] += bytes.len() as u64;
124 }
125 Err(E::Cancelled) => return Err(E::Cancelled),
126 Err(E::Deadline) => return Err(E::Deadline),
127 Err(error) => {
128 if matches!(error, E::Integrity | E::Limit | E::InvalidContract) {
129 quarantined[source] = true;
130 }
131 pending.push_back(index);
132 retry = true;
133 }
134 }
135 }
136 if retry {
137 let span = options.retry_delay.as_nanos() as u64;
138 let seed = attempts.iter().sum::<usize>() as u64
139 ^ cursor as u64
140 ^ u64::from(descriptor.object_hash().as_bytes()[0]);
141 let jitter = Duration::from_nanos(seed.wrapping_mul(6364136223846793005) % (span + 1));
142 control.wait(options.retry_delay + jitter)?;
143 }
144 }
145 control.checkpoint()?;
146 let object = session.finish_with_control(control)?;
147 control.checkpoint()?;
148 Ok(TransferReport {
149 object,
150 resumed_blocks,
151 attempts: attempts.iter().sum(),
152 quarantined_sources: quarantined.iter().filter(|v| **v).count(),
153 source_bytes,
154 })
155}
156fn validate_sources(sources: &[SourceBinding<'_>], options: TransferOptions) -> TransferResult<()> {
157 if sources.is_empty()
158 || sources.len() > 16
159 || options.max_parallel == 0
160 || options.max_parallel > 8
161 || options.max_attempts == 0
162 || options.max_attempts > 32
163 || options.retry_delay > Duration::from_secs(1)
164 {
165 return Err(E::InvalidContract);
166 }
167 let mut ids = HashSet::new();
168 if sources.iter().any(|s| !ids.insert(&s.id)) {
169 return Err(E::InvalidContract);
170 }
171 Ok(())
172}
173fn select_jobs(
174 pending: &mut VecDeque<usize>,
175 attempts: &mut [usize],
176 quarantined: &[bool],
177 cursor: &mut usize,
178 options: TransferOptions,
179) -> TransferResult<Vec<(usize, usize)>> {
180 let mut used = vec![false; quarantined.len()];
181 let mut jobs = Vec::new();
182 for _ in 0..options.max_parallel.min(pending.len()) {
183 let source = (0..quarantined.len())
184 .map(|offset| (*cursor + offset) % quarantined.len())
185 .find(|i| !used[*i] && !quarantined[*i]);
186 let Some(source) = source else { break };
187 let index = pending.pop_front().ok_or(E::InvalidContract)?;
188 if attempts[index] >= options.max_attempts {
189 return Err(E::SourcesExhausted);
190 }
191 attempts[index] += 1;
192 used[source] = true;
193 *cursor = (source + 1) % quarantined.len();
194 jobs.push((index, source));
195 }
196 if jobs.is_empty() {
197 return Err(E::SourcesExhausted);
198 }
199 Ok(jobs)
200}
201fn fetch_wave(
202 descriptor: &ObjectDescriptor,
203 sources: &[SourceBinding<'_>],
204 jobs: &[(usize, usize)],
205 control: &TransferControl,
206) -> Vec<TransferResult<Vec<u8>>> {
207 std::thread::scope(|scope| {
208 let handles: Vec<_> = jobs
209 .iter()
210 .map(|(index, source)| {
211 std::thread::Builder::new().spawn_scoped(scope, move || {
212 let mut sink = BlockSink {
213 bytes: Vec::with_capacity(descriptor.block_len(*index)?),
214 limit: descriptor.block_len(*index)?,
215 failure: None,
216 control,
217 };
218 let result = sources[*source]
219 .source
220 .read_block(descriptor, *index, control, &mut sink);
221 if let Some(error) = sink.failure {
222 return Err(error);
223 }
224 result?;
225 control.checkpoint()?;
226 descriptor.verify_block(*index, &sink.bytes)?;
227 Ok(sink.bytes)
228 })
229 })
230 .collect();
231 handles
232 .into_iter()
233 .map(|h| match h {
234 Ok(h) => h.join().unwrap_or(Err(E::InvalidContract)),
235 Err(_) => Err(E::Source),
236 })
237 .collect()
238 })
239}
240struct BlockSink<'a> {
241 bytes: Vec<u8>,
242 limit: usize,
243 failure: Option<E>,
244 control: &'a TransferControl,
245}
246impl Write for BlockSink<'_> {
247 fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
248 let result = self.control.checkpoint().and_then(|()| {
249 if self
250 .bytes
251 .len()
252 .checked_add(bytes.len())
253 .is_none_or(|n| n > self.limit)
254 {
255 Err(E::Limit)
256 } else {
257 Ok(())
258 }
259 });
260 if let Err(error) = result {
261 self.failure = Some(error);
262 return Err(std::io::Error::other(error));
263 }
264 self.bytes.extend_from_slice(bytes);
265 Ok(bytes.len())
266 }
267 fn flush(&mut self) -> std::io::Result<()> {
268 Ok(())
269 }
270}