fast_pull/core/mod.rs
1//! Top-level download orchestration: session handle plus single- and
2//! multi-threaded entry points.
3//!
4//! [`download_single`](crate::single::download_single) runs a sequential pull,
5//! while [`download_multi`](crate::multi::download_multi) splits the work across
6//! concurrent workers with work-stealing. Both return a [`DownloadResult`], a
7//! cheaply cloneable handle that keeps the session alive until the last clone is
8//! dropped (or [`DownloadResult::abort`] is called).
9
10use crate::Event;
11use crossfire::{MAsyncRx, mpmc};
12use fast_steal::{Executor, TaskQueue};
13use force_send::ForceSend;
14use std::fmt;
15use std::sync::Arc;
16use std::sync::atomic::{AtomicU64, Ordering};
17use tokio_util::sync::CancellationToken;
18
19mod force_send;
20pub mod mock;
21pub mod multi;
22pub mod single;
23
24/// Shared state of an active download session.
25///
26/// Owned inside an `Arc` by [`DownloadResult`]. Because all clones of a
27/// `DownloadResult` share the **same** `Arc<DownloadResultInner>`, the `Drop`
28/// impl below runs exactly once — when the last clone is dropped. That is where
29/// cancellation happens, giving `DownloadResult` `Arc`-style "last owner gone →
30/// release" semantics: the download keeps running as long as any handle is
31/// alive, and is cancelled only when the final one is dropped.
32struct DownloadResultInner<E, PullError, PushError>
33where
34 E: Executor + Send + Sync,
35 PullError: Send + Unpin + 'static,
36 PushError: Send + Unpin + 'static,
37{
38 event_chain: MAsyncRx<mpmc::List<Event<PullError, PushError>>>,
39 /// The work-stealing queue of a multi-threaded session together with the
40 /// executor that spawns its workers, or `None` for a single-threaded one.
41 ///
42 /// The executor is kept for the entire lifetime of the handle so
43 /// [`set_threads`](Self::set_threads) can spawn additional workers while the
44 /// download is running. It reaches the session's channels through a weak
45 /// reference, so retaining it here never keeps a finished session open.
46 task_queue: Option<(E, TaskQueue<E::Handle>)>,
47 /// Live lower bound on how finely a worker may split a range when it steals,
48 /// shared with the session's executor. `None` for a single-threaded session,
49 /// which has no work-stealing and so no granularity to change.
50 min_chunk_size: Option<Arc<AtomicU64>>,
51 /// Session-wide cancellation token, shared (as a clone) with every worker
52 /// and with the `spawn_blocking` push driver. [`abort`](Self::abort) cancels
53 /// this root token, which broadcasts to all linked child tokens — including
54 /// workers spawned *after* the cancel call — so a late worker observes
55 /// cancellation on its next poll instead of needing a separate one-shot
56 /// notify per handle. Cancellation is terminal: once cancelled it never
57 /// clears, so `is_aborted` stays `true` for the rest of the session.
58 abort_token: CancellationToken,
59}
60
61impl<E, PullError, PushError> fmt::Debug for DownloadResultInner<E, PullError, PushError>
62where
63 E: Executor + Send + Sync,
64 PullError: Send + Unpin + 'static,
65 PushError: Send + Unpin + 'static,
66{
67 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
68 f.debug_struct("DownloadResultInner")
69 .field("event_chain", &self.event_chain)
70 .field("is_aborted", &self.abort_token.is_cancelled())
71 .finish_non_exhaustive()
72 }
73}
74
75impl<E, PullError, PushError> DownloadResultInner<E, PullError, PushError>
76where
77 E: Executor + Send + Sync,
78 PullError: Send + Unpin + 'static,
79 PushError: Send + Unpin + 'static,
80{
81 /// Cancel all workers immediately.
82 ///
83 /// Safe to call multiple times and safe to call while other clones of the
84 /// owning [`DownloadResult`] are still alive. The implicit drop-based
85 /// cancellation (on the last clone) becomes a no-op once this has run.
86 pub fn abort(&self) {
87 self.abort_token.cancel();
88 }
89
90 pub fn set_threads(&self, threads: usize) -> Option<()> {
91 let (executor, task_queue) = self.task_queue.as_ref()?;
92 let min = self
93 .min_chunk_size
94 .as_ref()
95 .map_or(1, |live| live.load(Ordering::Relaxed));
96 task_queue.set_threads(threads, min, Some(executor))
97 }
98
99 /// Set the live minimum chunk size used when a worker next steals work.
100 ///
101 /// A no-op for single-threaded sessions, which have no worker pool.
102 pub fn set_min_chunk_size(&self, min_chunk_size: u64) {
103 if let Some(live) = &self.min_chunk_size {
104 live.store(min_chunk_size.max(1), Ordering::Relaxed);
105 }
106 }
107
108 #[must_use]
109 pub fn is_aborted(&self) -> bool {
110 self.abort_token.is_cancelled()
111 }
112}
113
114impl<E, PullError, PushError> Drop for DownloadResultInner<E, PullError, PushError>
115where
116 E: Executor + Send + Sync,
117 PullError: Send + Unpin + 'static,
118 PushError: Send + Unpin + 'static,
119{
120 fn drop(&mut self) {
121 self.abort();
122 }
123}
124
125/// Handle to an active download session.
126///
127/// Cheaply cloneable shared handle. The underlying download keeps running as
128/// long as **any** clone is alive, and is cancelled only once the last clone is
129/// dropped. An explicit [`abort`](Self::abort) cancels immediately.
130///
131/// `DownloadResult` wraps `Arc<DownloadResultInner>` and exposes the session
132/// methods (`abort`, `set_threads`, `set_min_chunk_size`, `is_aborted`) and
133/// [`event_chain`](Self::event_chain) directly; each delegates to the inner
134/// value. There is intentionally **no** `Deref` impl — `DownloadResultInner`
135/// is private, so callers reach session state only through these methods.
136///
137/// The inner value is held in a `ForceSend`, a deliberately isolated
138/// workaround so the handle is `Send` even when it is built inside an `async`
139/// block (see that type's module docs).
140///
141/// Completion is observed by draining the event stream: once the last sender is
142/// dropped (the download finished or was aborted) the receiver disconnects, so
143/// `while result.recv().await.is_ok() {}` awaits the session end. [`recv`](Self::recv)
144/// is the `Send`-safe form of that call; prefer it over
145/// `result.event_chain().recv()` in a `tokio::spawn`ed driver.
146pub struct DownloadResult<E, PullError, PushError>
147where
148 E: Executor + Send + Sync,
149 PullError: Send + Unpin + 'static,
150 PushError: Send + Unpin + 'static,
151{
152 inner: ForceSend<Arc<DownloadResultInner<E, PullError, PushError>>>,
153}
154
155impl<E, PullError, PushError> fmt::Debug for DownloadResult<E, PullError, PushError>
156where
157 E: Executor + Send + Sync,
158 PullError: Send + Unpin + 'static,
159 PushError: Send + Unpin + 'static,
160{
161 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
162 f.debug_struct("DownloadResult")
163 .field("inner", &*self.inner)
164 .finish()
165 }
166}
167
168impl<E, PullError, PushError> Clone for DownloadResult<E, PullError, PushError>
169where
170 E: Executor + Send + Sync,
171 PullError: Send + Unpin + 'static,
172 PushError: Send + Unpin + 'static,
173{
174 fn clone(&self) -> Self {
175 Self {
176 inner: self.inner.clone(),
177 }
178 }
179}
180
181impl<E, PullError, PushError> DownloadResult<E, PullError, PushError>
182where
183 E: Executor + Send + Sync,
184 PullError: Send + Unpin + 'static,
185 PushError: Send + Unpin + 'static,
186{
187 /// Construct a [`DownloadResult`] from the raw session pieces.
188 ///
189 /// This is an internal constructor used by
190 /// [`download_single`](crate::single::download_single) and
191 /// [`download_multi`](crate::multi::download_multi); prefer those entry
192 /// points instead of calling this directly.
193 ///
194 /// `min_chunk_size` is the session's live steal-granularity cell, or `None`
195 /// for a single-threaded session (see [`set_min_chunk_size`](Self::set_min_chunk_size)).
196 ///
197 /// Not part of the stable surface: kept reachable only so the allocation
198 /// tests can build a session without going through an entry point that
199 /// spawns tasks.
200 #[doc(hidden)]
201 pub fn new(
202 event_chain: MAsyncRx<mpmc::List<Event<PullError, PushError>>>,
203 task_queue: Option<(E, TaskQueue<E::Handle>)>,
204 abort_token: CancellationToken,
205 min_chunk_size: Option<Arc<AtomicU64>>,
206 ) -> Self {
207 Self {
208 inner: ForceSend(Arc::new(DownloadResultInner {
209 event_chain,
210 task_queue,
211 min_chunk_size,
212 abort_token,
213 })),
214 }
215 }
216
217 /// Access the stream of [`Event`]s emitted during the session.
218 ///
219 /// The receiver closes once the last clone of this handle is dropped or the
220 /// session is aborted, so draining it is a natural way to observe progress.
221 #[must_use]
222 pub fn event_chain(&self) -> &MAsyncRx<mpmc::List<Event<PullError, PushError>>> {
223 &self.inner.event_chain
224 }
225
226 /// Await the next event, or `Err` once the session has ended.
227 ///
228 /// Equivalent to `self.event_chain().recv()`, but returns
229 /// `impl Future + Send`, so an enclosing `async` block that awaits it stays
230 /// `Send`. Prefer this over `event_chain().recv()` in a `tokio::spawn`ed
231 /// driver: the underlying `RecvFuture` is only `Send` through a hand-written
232 /// impl that the `async fn` auto-trait leak check does not see through, so
233 /// awaiting it inline reintroduces a spurious `higher-ranked lifetime error`.
234 /// No allocation.
235 pub fn recv(
236 &self,
237 ) -> impl std::future::Future<Output = Result<Event<PullError, PushError>, crossfire::RecvError>>
238 + Send
239 + '_ {
240 self.inner.event_chain.recv()
241 }
242
243 /// Cancel all workers immediately.
244 ///
245 /// Safe to call multiple times and safe to call while other clones of the
246 /// owning [`DownloadResult`] are still alive. The implicit drop-based
247 /// cancellation (on the last clone) becomes a no-op once this has run.
248 pub fn abort(&self) {
249 self.inner.abort();
250 }
251
252 /// Adjust the worker thread count of a running multi-threaded session.
253 ///
254 /// The minimum chunk size used while growing is the session's live value,
255 /// set independently with [`set_min_chunk_size`](Self::set_min_chunk_size)
256 /// (or the configured value until then).
257 ///
258 /// Growing spawns workers for ranges still waiting in the queue, or splits a
259 /// range off the busiest running worker when nothing is waiting. Shrinking
260 /// aborts the surplus workers and returns their ranges to the queue for the
261 /// survivors to steal. No-op for single-threaded sessions, which have no
262 /// task queue.
263 ///
264 /// A session ends when its last worker exits, and it cannot be restarted:
265 /// growing afterwards spawns nothing, because the workers' shared channels
266 /// are closed at that point. Shrinking is clamped to a minimum of one worker
267 /// by the scheduler, so `set_threads(0)` keeps a single worker alive rather
268 /// than ending the session. The session terminates only when that last
269 /// (clamped) worker exits on its own.
270 ///
271 /// This never touches the [`abort`](Self::abort) token: abort is terminal and
272 /// cannot be undone by resizing, so [`is_aborted`](Self::is_aborted) stays
273 /// `true` across any later `set_threads` call, and any worker spawned by such
274 /// a call observes the cancelled token and exits without pulling.
275 pub fn set_threads(&self, threads: usize) {
276 self.inner.set_threads(threads);
277 }
278
279 /// Set the minimum chunk size used when a worker next steals work.
280 ///
281 /// Bounds how finely a running worker may split a range to take a share of a
282 /// peer's work: a range is only halved when at least `min_chunk_size * 2`
283 /// bytes remain. The new value takes effect on the next steal — there is no
284 /// need to also call [`set_threads`](Self::set_threads) — and it is also the
285 /// granularity used to split ranges the next time the pool grows.
286 ///
287 /// A no-op for single-threaded sessions, which have no worker pool.
288 pub fn set_min_chunk_size(&self, min_chunk_size: u64) {
289 self.inner.set_min_chunk_size(min_chunk_size);
290 }
291
292 /// Whether the session has been (or is being) cancelled.
293 #[must_use]
294 pub fn is_aborted(&self) -> bool {
295 self.inner.is_aborted()
296 }
297}
298
299#[cfg(test)]
300mod tests {
301 #![allow(clippy::unwrap_used)]
302 #![allow(clippy::cast_possible_truncation)]
303 use crate::MemPusher;
304 use crate::mock::{MockPuller, build_mock_data};
305 use crate::multi::{DownloadOptions, download_multi};
306 use crate::{Event, ProgressEntry, PullResult, PullStream, Puller};
307 use bytes::Bytes;
308 use futures::{StreamExt, stream};
309 use std::collections::BTreeSet;
310 use std::sync::Arc;
311 use tokio::time::{Duration, sleep, timeout};
312
313 /// A [`Puller`] that stalls for `delay` before yielding a range in one piece,
314 /// so a test can resize the worker pool while the download is still running.
315 #[derive(Debug, Clone)]
316 struct SlowPuller {
317 data: Arc<[u8]>,
318 delay: Duration,
319 }
320 impl Puller for SlowPuller {
321 type Error = std::convert::Infallible;
322 fn pull(
323 &mut self,
324 range: Option<&ProgressEntry>,
325 ) -> impl Future<Output = PullResult<impl PullStream<Self::Error>, Self::Error>> + Send
326 {
327 type PullItem = PullResult<Bytes, std::convert::Infallible>;
328 let owned: Vec<u8> = match range {
329 Some(r) => self.data[r.start as usize..r.end as usize].to_vec(),
330 None => self.data.to_vec(),
331 };
332 let delay = self.delay;
333 async move {
334 sleep(delay).await;
335 let items: Vec<PullItem> = vec![Ok(Bytes::from(owned))];
336 Ok(stream::iter(items))
337 }
338 }
339 }
340
341 /// A [`Puller`] that yields its range in small pieces with a pause between
342 /// each, so a shrink that aborts a worker lands *mid-range* instead of on a
343 /// clean boundary. That is the shape which exercises cursor hand-off: the
344 /// reclaimed range must resume from the advanced cursor, without dropping or
345 /// re-delivering the bytes the aborted worker had already pushed.
346 #[derive(Debug, Clone)]
347 struct ChunkedPuller {
348 data: Arc<[u8]>,
349 piece: usize,
350 delay: Duration,
351 }
352 impl Puller for ChunkedPuller {
353 type Error = std::convert::Infallible;
354 fn pull(
355 &mut self,
356 range: Option<&ProgressEntry>,
357 ) -> impl Future<Output = PullResult<impl PullStream<Self::Error>, Self::Error>> + Send
358 {
359 let owned: Vec<u8> = match range {
360 Some(r) => self.data[r.start as usize..r.end as usize].to_vec(),
361 None => self.data.to_vec(),
362 };
363 let piece = self.piece;
364 let delay = self.delay;
365 async move {
366 Ok(
367 stream::unfold(Bytes::from(owned), move |mut buf| async move {
368 if buf.is_empty() {
369 return None;
370 }
371 let next = buf.split_to(piece.min(buf.len()));
372 sleep(delay).await;
373 Some((Ok(next), buf))
374 })
375 .boxed(),
376 )
377 }
378 }
379 }
380
381 /// Deterministic xorshift64. A churn schedule driven by a fixed seed keeps a
382 /// failure reproducible instead of turning the test into a lottery.
383 fn next_rand(state: &mut u64) -> u64 {
384 let mut x = *state;
385 x ^= x << 13;
386 x ^= x >> 7;
387 x ^= x << 17;
388 *state = x;
389 x
390 }
391
392 /// Eight equally sized ranges over `size` bytes, so a one-worker session
393 /// leaves seven of them waiting in the queue.
394 fn eight_chunks(size: u64) -> Vec<ProgressEntry> {
395 let step = size / 8;
396 (0..8).map(|i| i * step..(i + 1) * step).collect()
397 }
398
399 #[tokio::test(flavor = "multi_thread")]
400 async fn download_result_debug_and_set_threads() {
401 let mock_data = build_mock_data(1024);
402 let puller = MockPuller::new(&mock_data);
403 let pusher = MemPusher::with_capacity(mock_data.len());
404 let receive = pusher.receive.clone();
405 #[allow(clippy::single_range_in_vec_init)]
406 let download_chunks = [0..mock_data.len() as u64];
407 let result = download_multi(
408 puller,
409 pusher,
410 DownloadOptions {
411 concurrent: 4,
412 retry_gap: Duration::from_secs(1),
413 push_queue_cap: 1024,
414 download_chunks: download_chunks.iter().cloned(),
415 pull_timeout: Duration::from_secs(5),
416 min_chunk_size: 1,
417 max_speculative: 3,
418 },
419 );
420 // `Debug` of `DownloadResultInner` is reached through `DownloadResult`'s
421 // own `Debug` impl, which forwards to the inner value.
422 let _ = format!("{result:?}");
423 // Lines 234-236 (forwarding) and 111-123 (inner task-queue adjustment).
424 result.set_threads(4);
425 // Await completion by draining `event_chain` (it disconnects once the
426 // push driver drops its sender) — this replaces `join()`.
427 while result.event_chain().recv().await.is_ok() {}
428 assert_eq!(&**receive.lock(), mock_data);
429 }
430
431 #[tokio::test(flavor = "multi_thread")]
432 async fn download_result_clone_and_set_threads_no_queue() {
433 use crate::single::download_single;
434 let mock_data = build_mock_data(1024);
435 let puller = MockPuller::new(&mock_data);
436 let pusher = MemPusher::with_capacity(mock_data.len());
437 let receive = pusher.receive.clone();
438 // Single-threaded sessions have no task queue, so `set_threads` is a no-op
439 // (covers the `if let` else path, line 122 of `DownloadResultInner`).
440 let result = download_single(
441 puller,
442 pusher,
443 crate::single::DownloadOptions {
444 retry_gap: Duration::from_secs(1),
445 push_queue_cap: 1024,
446 },
447 );
448 // Lines 167-171: `DownloadResult` is `Clone`.
449 let _clone = result.clone();
450 result.set_threads(4);
451 while result.event_chain().recv().await.is_ok() {}
452 assert_eq!(&**receive.lock(), mock_data);
453 }
454
455 // Pins the `is_aborted` interaction with `set_threads`: a *live*
456 // (never-aborted) session keeps `is_aborted() == false` after `set_threads`,
457 // while an *already-aborted* session stays aborted — the flag is a one-way
458 // latch, so `is_aborted()` cannot report a false "live" state for a cancelled
459 // download.
460 #[tokio::test(flavor = "multi_thread")]
461 async fn set_threads_does_not_unabort_an_aborted_session() {
462 let mock_data = build_mock_data(1024);
463 let puller = MockPuller::new(&mock_data);
464 let pusher = MemPusher::with_capacity(mock_data.len());
465 let result = download_multi(
466 puller,
467 pusher,
468 DownloadOptions {
469 concurrent: 4,
470 retry_gap: Duration::from_secs(1),
471 push_queue_cap: 1024,
472 download_chunks: std::iter::once(0..mock_data.len() as u64),
473 pull_timeout: Duration::from_secs(5),
474 min_chunk_size: 1,
475 max_speculative: 3,
476 },
477 );
478 // Live session: flag starts false and must stay false after a resize.
479 assert!(!result.is_aborted());
480 result.set_threads(4);
481 assert!(!result.is_aborted());
482 // Now abort: flag is true and must *stay* true across a later resize.
483 result.abort();
484 assert!(result.is_aborted());
485 result.set_threads(4);
486 assert!(result.is_aborted());
487 while result.event_chain().recv().await.is_ok() {}
488 }
489
490 // Growing the pool of a *running* session must actually put more workers to
491 // work. The session starts with a single worker and seven ranges waiting;
492 // if the growth were a no-op, every `Pulling` event would carry worker id 0.
493 #[tokio::test(flavor = "multi_thread")]
494 async fn set_threads_growth_spawns_additional_workers() {
495 let mock_data = build_mock_data(8 * 1024);
496 let download_chunks = eight_chunks(mock_data.len() as u64);
497 let puller = SlowPuller {
498 data: Arc::from(mock_data.as_slice()),
499 delay: Duration::from_millis(80),
500 };
501 let pusher = MemPusher::with_capacity(mock_data.len());
502 let receive = pusher.receive.clone();
503 let result = download_multi(
504 puller,
505 pusher,
506 DownloadOptions {
507 concurrent: 1,
508 retry_gap: Duration::from_secs(1),
509 push_queue_cap: 1024,
510 download_chunks: download_chunks.iter().cloned(),
511 pull_timeout: Duration::from_secs(5),
512 min_chunk_size: 1,
513 max_speculative: 1,
514 },
515 );
516
517 // Resize while the lone worker is still stalled on its first range.
518 let grower = result.clone();
519 tokio::spawn(async move {
520 sleep(Duration::from_millis(30)).await;
521 grower.set_threads(8);
522 });
523
524 let mut pulling_ids = BTreeSet::new();
525 while let Ok(e) = result.event_chain().recv().await {
526 if let Event::Pulling(id) = e {
527 pulling_ids.insert(id);
528 }
529 }
530 assert!(
531 pulling_ids.len() > 1,
532 "growth spawned no additional worker (pulling ids: {pulling_ids:?})"
533 );
534 assert_eq!(&**receive.lock(), mock_data);
535 }
536
537 // The mirror image of the test above: a worker spawned by a growth that
538 // races an `abort` must observe the abort latch and exit without pulling, so
539 // the session still finalizes promptly instead of resuming.
540 #[tokio::test(flavor = "multi_thread")]
541 async fn set_threads_growth_after_abort_does_not_resume() {
542 let mock_data = build_mock_data(8 * 1024);
543 let download_chunks = eight_chunks(mock_data.len() as u64);
544 let puller = SlowPuller {
545 data: Arc::from(mock_data.as_slice()),
546 delay: Duration::from_millis(150),
547 };
548 let pusher = MemPusher::with_capacity(mock_data.len());
549 let receive = pusher.receive.clone();
550 let result = download_multi(
551 puller,
552 pusher,
553 DownloadOptions {
554 concurrent: 1,
555 retry_gap: Duration::from_secs(1),
556 push_queue_cap: 1024,
557 download_chunks: download_chunks.iter().cloned(),
558 pull_timeout: Duration::from_secs(5),
559 min_chunk_size: 1,
560 max_speculative: 1,
561 },
562 );
563
564 result.abort();
565 result.set_threads(8);
566 // Await termination by draining `event_chain` (replaces `join()`); the
567 // timeout guards against a hang.
568 timeout(Duration::from_secs(10), async {
569 while result.event_chain().recv().await.is_ok() {}
570 })
571 .await
572 .expect("join() hung after abort followed by growth");
573 assert!(
574 receive.lock().len() < mock_data.len(),
575 "an aborted session must not be resumed by a later resize"
576 );
577 }
578
579 // Repeatedly resizing a *running* pool in both directions must never corrupt
580 // the download. Each shrink aborts live workers and reclaims their ranges
581 // mid-flight; each growth hands those ranges to fresh workers. A cursor
582 // hand-off that is off by even one piece shows up here as missing or
583 // duplicated bytes, which a single grow-once test cannot catch. The pieces
584 // are small and paced so the churn lands inside a range rather than on a
585 // boundary, and the schedule is seeded so any failure reproduces exactly.
586 #[tokio::test(flavor = "multi_thread")]
587 async fn set_threads_random_churn_preserves_all_bytes() {
588 let mock_data = build_mock_data(64 * 1024);
589 let download_chunks = eight_chunks(mock_data.len() as u64);
590 let puller = ChunkedPuller {
591 data: Arc::from(mock_data.as_slice()),
592 piece: 512,
593 delay: Duration::from_millis(2),
594 };
595 let pusher = MemPusher::with_capacity(mock_data.len());
596 let receive = pusher.receive.clone();
597 let result = download_multi(
598 puller,
599 pusher,
600 DownloadOptions {
601 concurrent: 4,
602 retry_gap: Duration::from_millis(10),
603 push_queue_cap: 1024,
604 download_chunks: download_chunks.iter().cloned(),
605 pull_timeout: Duration::from_secs(5),
606 min_chunk_size: 1,
607 max_speculative: 3,
608 },
609 );
610
611 // Never resize to zero: that collapses the pool and finalizes the
612 // session, which is a separate contract from mid-flight churn.
613 let churner = result.clone();
614 let probe = receive.clone();
615 let total = mock_data.len();
616 let churn = tokio::spawn(async move {
617 let mut state = 0x2545_F491_4F6C_DD1D_u64;
618 let mut seen = BTreeSet::new();
619 // Resizes landing after the last byte is written prove nothing, so
620 // count the ones that actually hit a still-running download.
621 let mut inflight = 0usize;
622 for _ in 0..40 {
623 let threads = (next_rand(&mut state) % 8 + 1) as usize;
624 seen.insert(threads);
625 if probe.lock().len() < total {
626 inflight += 1;
627 }
628 churner.set_threads(threads);
629 sleep(Duration::from_millis(3)).await;
630 }
631 // Leave a healthy pool behind so the remaining ranges drain.
632 churner.set_threads(8);
633 (seen, inflight)
634 });
635
636 let chunk_count = download_chunks.len();
637 let mut pulling_total = 0usize;
638 while let Ok(e) = result.event_chain().recv().await {
639 if matches!(e, Event::Pulling(_)) {
640 pulling_total += 1;
641 }
642 }
643 let (seen, inflight) = churn.await.unwrap();
644 assert!(
645 seen.len() > 2,
646 "churn never varied the pool size, so nothing was exercised: {seen:?}"
647 );
648 // Without this the test could silently degrade into a no-op: a download
649 // that outran the churn loop would take every resize on a dead session.
650 assert!(
651 inflight > 0,
652 "every resize landed after the download finished, so no running pool was churned"
653 );
654 // Reclaimed ranges are handed out again, so a pool that really churned
655 // pulls far more often than once per chunk.
656 assert!(
657 pulling_total > chunk_count,
658 "ranges were never redistributed ({pulling_total} pulls for {chunk_count} chunks)"
659 );
660 assert_eq!(
661 &**receive.lock(),
662 mock_data,
663 "repeated resizing corrupted the downloaded bytes"
664 );
665 }
666
667 /// A high-frequency churn of *both* live setters at once (pool size and
668 /// granularity) must leave the download intact. Unlike the threads-only
669 /// churn, this stresses the granularity cell being written concurrently with
670 /// every worker's steal. The schedule is seeded and the granularity is
671 /// varied across the `split_two` boundary (`>= 2 * piece`) and beyond, and
672 /// reset small at the end so the remaining ranges can drain.
673 #[tokio::test(flavor = "multi_thread")]
674 async fn set_threads_and_min_chunk_size_churn_preserves_all_bytes() {
675 let mock_data = build_mock_data(64 * 1024);
676 let download_chunks = eight_chunks(mock_data.len() as u64);
677 let puller = ChunkedPuller {
678 data: Arc::from(mock_data.as_slice()),
679 piece: 512,
680 delay: Duration::from_millis(2),
681 };
682 let pusher = MemPusher::with_capacity(mock_data.len());
683 let receive = pusher.receive.clone();
684 let result = download_multi(
685 puller,
686 pusher,
687 DownloadOptions {
688 concurrent: 4,
689 retry_gap: Duration::from_millis(10),
690 push_queue_cap: 1024,
691 download_chunks: download_chunks.iter().cloned(),
692 pull_timeout: Duration::from_secs(5),
693 min_chunk_size: 1,
694 max_speculative: 3,
695 },
696 );
697
698 let churner = result.clone();
699 let probe = receive.clone();
700 let total = mock_data.len();
701 let churn = tokio::spawn(async move {
702 let mut state = 0x9E37_79B9_7F4A_7C15_u64;
703 let mut threads_seen = BTreeSet::new();
704 let mut mins_seen = BTreeSet::new();
705 let mut inflight = 0usize;
706 for _ in 0..40 {
707 let threads = (next_rand(&mut state) % 8 + 1) as usize;
708 // Straddle the split boundary (`2 * piece == 1024`) and go well
709 // past it, so both "can split" and "cannot split" are exercised.
710 let min = match next_rand(&mut state) % 4 {
711 0 => 1,
712 1 => 512,
713 2 => 1024,
714 _ => 4096,
715 };
716 threads_seen.insert(threads);
717 mins_seen.insert(min);
718 if probe.lock().len() < total {
719 inflight += 1;
720 }
721 churner.set_threads(threads);
722 churner.set_min_chunk_size(min);
723 sleep(Duration::from_millis(3)).await;
724 }
725 // Leave a healthy pool and a permissive granularity behind so the
726 // remaining ranges can still be split and drained.
727 churner.set_threads(8);
728 churner.set_min_chunk_size(1);
729 (threads_seen, mins_seen, inflight)
730 });
731
732 let mut pulling_total = 0usize;
733 while let Ok(e) = result.event_chain().recv().await {
734 if matches!(e, Event::Pulling(_)) {
735 pulling_total += 1;
736 }
737 }
738 let (threads_seen, mins_seen, inflight) = churn.await.unwrap();
739 assert!(
740 threads_seen.len() > 2,
741 "churn never varied the pool size: {threads_seen:?}"
742 );
743 assert!(
744 mins_seen.len() > 2,
745 "churn never varied the granularity: {mins_seen:?}"
746 );
747 assert!(
748 inflight > 0,
749 "every churn step landed after the download finished"
750 );
751 assert!(
752 pulling_total >= download_chunks.len(),
753 "ranges were never redistributed ({pulling_total} pulls)"
754 );
755 assert_eq!(
756 &**receive.lock(),
757 mock_data,
758 "churning threads and granularity together corrupted the bytes"
759 );
760 }
761}