fast_pull/core/
mod.rs

1extern crate alloc;
2use crate::Event;
3use alloc::sync::{Arc, Weak};
4use core::num::{NonZeroU64, NonZeroUsize};
5use fast_steal::{Executor, TaskList};
6use kanal::AsyncReceiver;
7use tokio::{
8    sync::Mutex,
9    task::{AbortHandle, JoinError, JoinHandle},
10};
11
12mod macros;
13pub mod mock;
14pub mod multi;
15pub mod single;
16
17#[derive(Clone)]
18pub struct DownloadResult<E: Executor, PullError, PushError> {
19    pub event_chain: AsyncReceiver<Event<PullError, PushError>>,
20    handle: Arc<Mutex<Option<JoinHandle<()>>>>,
21    abort_handles: Arc<[AbortHandle]>,
22    task_list: Option<Weak<TaskList<E>>>,
23}
24
25impl<E: Executor, PullError, PushError> DownloadResult<E, PullError, PushError> {
26    pub fn new(
27        event_chain: AsyncReceiver<Event<PullError, PushError>>,
28        handle: JoinHandle<()>,
29        abort_handles: &[AbortHandle],
30        task_list: Option<Weak<TaskList<E>>>,
31    ) -> Self {
32        Self {
33            event_chain,
34            abort_handles: Arc::from(abort_handles),
35            handle: Arc::new(Mutex::new(Some(handle))),
36            task_list,
37        }
38    }
39
40    /// 只有第一次调用有效
41    pub async fn join(&self) -> Result<(), JoinError> {
42        if let Some(handle) = self.handle.lock().await.take() {
43            handle.await?
44        }
45        Ok(())
46    }
47
48    pub fn abort(&self) {
49        for abort_handle in self.abort_handles.iter() {
50            abort_handle.abort();
51        }
52    }
53
54    pub fn set_threads(&self, threads: NonZeroUsize, min_chunk_size: NonZeroU64) {
55        if let Some(task_list) = &self.task_list
56            && let Some(task_list) = task_list.upgrade()
57        {
58            task_list.set_threads(threads, min_chunk_size);
59        }
60    }
61}