chore(fast-pull,fast-steal): 整理代码

Author: share121

Cargo.lock

@@ -16,6 +16,7 @@ homepage.workspace = true
 tokio = { version = "1.47.1", features = [
     "macros",
     "io-util",
+    "time",
 ], default-features = false }
 kanal.workspace = true
 fast-steal.workspace = true
@@ -26,28 +27,12 @@ mmap-io = { version = "0.8.0", features = [
     "async",
     "hugepages",
 ], optional = true }
-reqwest = { version = "0.12.22", features = [
-    "stream",
-], default-features = false, optional = true }
-content_disposition = { version = "0.4.0", optional = true }
-sanitize-filename = { version = "0.6.0", optional = true }
-urlencoding = { version = "2.1.3", optional = true }
-url = { workspace = true, optional = true }
 
 [dev-dependencies]
-mockito = "1.7.0"
 tempfile = "3.20.0"
 
 [features]
-default = ["file", "reqwest"]
+default = ["file"]
 
 file = ["mmap-io"]
 mmap-io = ["dep:mmap-io"]
-
-reqwest = [
-    "dep:url",
-    "dep:reqwest",
-    "dep:content_disposition",
-    "dep:sanitize-filename",
-    "dep:urlencoding",
-]

crates/fast-pull/Cargo.toml

@@ -4,8 +4,6 @@ mod progress;
 mod puller;
 mod pusher;
 mod total;
-#[cfg(feature = "reqwest")]
-pub(crate) mod url;
 
 pub use event::*;
 pub use merge_progress::*;

crates/fast-pull/src/base/mod.rs

@@ -2,9 +2,6 @@ use crate::ProgressEntry;
 use bytes::Bytes;
 use core::future;
 
-/// Random Pusher
-///
-/// Implementation should only overwrite content in `range`, even if `content.len()` is larger than `range.total()`
 pub trait RandPusher: Send {
     type Error: Send;
     fn push(

crates/fast-pull/src/base/pusher.rs

@@ -1,6 +1,6 @@
 extern crate alloc;
 use super::macros::poll_ok;
-use crate::{DownloadResult, Event, ProgressEntry, RandPuller, RandPusher, WorkerId};
+use crate::{DownloadResult, Event, ProgressEntry, RandPuller, RandPusher, Total, WorkerId};
 use alloc::{sync::Arc, vec::Vec};
 use bytes::Bytes;
 use core::{
@@ -51,20 +51,18 @@ where
             options.retry_gap
         );
     });
-    let executor: TokioExecutor<R, W> = TokioExecutor {
+    let executor: Arc<TokioExecutor<R, W>> = Arc::new(TokioExecutor {
         tx,
         tx_push,
         puller,
         retry_gap: options.retry_gap,
         id: Arc::new(AtomicUsize::new(0)),
         min_chunk_size: options.min_chunk_size,
-    };
-    let task_list = TaskList::run(
-        options.concurrent,
-        options.min_chunk_size,
-        &options.download_chunks[..],
-        executor,
-    );
+    });
+    let task_list = Arc::new(TaskList::run(&options.download_chunks[..], executor));
+    task_list
+        .clone()
+        .set_threads(options.concurrent, options.min_chunk_size);
     DownloadResult::new(
         event_chain,
         push_handle,
@@ -117,7 +115,7 @@ where
                 let mut stream = puller.pull(&download_range);
                 loop {
                     match stream.try_next().await {
-                        Ok(Some(chunk)) => {
+                        Ok(Some(mut chunk)) => {
                             let len = chunk.len() as u64;
                             task.fetch_add_start(len);
                             let range_start = start;
@@ -127,14 +125,12 @@ where
                                 continue 'steal_task;
                             }
                             let span = range_start..range_end;
+                            chunk.truncate(span.total() as usize);
                             self.tx
                                 .send(Event::PullProgress(id, span.clone()))
                                 .await
                                 .unwrap();
-                            self.tx_push
-                                .send((id, span, chunk))
-                                .await
-                                .unwrap();
+                            self.tx_push.send((id, span, chunk)).await.unwrap();
                         }
                         Ok(None) => break,
                         Err(e) => {
@@ -144,8 +140,8 @@ where
                     }
                 }
             }
-            self.tx.send(Event::Finished(id)).await.unwrap();
             task_list.remove(&task);
+            self.tx.send(Event::Finished(id)).await.unwrap();
         });
         TokioHandle(handle.abort_handle())
     }

crates/fast-pull/src/base/url.rs

@@ -2,14 +2,13 @@ extern crate std;
 use crate::{ProgressEntry, RandPusher, SeqPusher, Total};
 use bytes::Bytes;
 use mmap_io::{MemoryMappedFile, MmapIoError, MmapMode, flush::FlushPolicy};
-use std::{path::Path, vec::Vec};
-use thiserror::Error;
+use std::{collections::VecDeque, path::Path};
 use tokio::{
     fs::{File, OpenOptions},
     io::{self, AsyncSeekExt, AsyncWriteExt, BufWriter, SeekFrom},
 };
 
-#[derive(Error, Debug)]
+#[derive(thiserror::Error, Debug)]
 pub enum FilePusherError {
     #[error(transparent)]
     MmapIo(#[from] MmapIoError),
@@ -76,7 +75,7 @@ impl RandPusher for RandFilePusherMmap {
         self.mmap
             .as_slice_mut(range.start, range.total())?
             .as_mut()
-            .copy_from_slice(&bytes[0..(range.total() as usize)]);
+            .copy_from_slice(&bytes);
         self.downloaded += bytes.len();
         if self.downloaded >= self.buffer_size {
             self.mmap.flush_async().await?;
@@ -93,7 +92,7 @@ impl RandPusher for RandFilePusherMmap {
 #[derive(Debug)]
 pub struct RandFilePusherStd {
     buffer: BufWriter<File>,
-    cache: Vec<(u64, Bytes)>,
+    cache: VecDeque<(u64, Bytes)>,
     p: u64,
     cache_size: usize,
     buffer_size: usize,
@@ -103,7 +102,7 @@ impl RandFilePusherStd {
         file.set_len(size).await?;
         Ok(Self {
             buffer: BufWriter::with_capacity(buffer_size, file),
-            cache: Vec::new(),
+            cache: VecDeque::new(),
             p: 0,
             cache_size: 0,
             buffer_size,
@@ -112,26 +111,26 @@ impl RandFilePusherStd {
 }
 impl RandPusher for RandFilePusherStd {
     type Error = FilePusherError;
-    async fn push(&mut self, range: ProgressEntry, mut bytes: Bytes) -> Result<(), Self::Error> {
+    async fn push(&mut self, range: ProgressEntry, bytes: Bytes) -> Result<(), Self::Error> {
         let pos = self.cache.partition_point(|(i, _)| i < &range.start);
         self.cache_size += bytes.len();
-        drop(bytes.split_off(range.total() as usize));
         self.cache.insert(pos, (range.start, bytes));
         if self.cache_size >= self.buffer_size {
             self.flush().await?;
         }
         Ok(())
     }
     async fn flush(&mut self) -> Result<(), Self::Error> {
-        for (start, bytes) in self.cache.drain(..) {
+        while let Some((start, bytes)) = self.cache.front() {
             let len = bytes.len();
             self.cache_size -= len;
-            if start != self.p {
-                self.buffer.seek(SeekFrom::Start(start)).await?;
-                self.p = start;
+            if *start != self.p {
+                self.buffer.seek(SeekFrom::Start(*start)).await?;
+                self.p = *start;
             }
-            self.buffer.write_all(&bytes).await?;
+            self.buffer.write_all(bytes).await?;
             self.p += len as u64;
+            self.cache.pop_front();
         }
         self.buffer.flush().await?;
         Ok(())
@@ -142,6 +141,7 @@ impl RandPusher for RandFilePusherStd {
 mod tests {
     use super::*;
     use bytes::Bytes;
+    use std::vec::Vec;
     use tempfile::NamedTempFile;
     use tokio::io::AsyncReadExt;
 

crates/fast-pull/src/core/multi.rs

@@ -3,7 +3,5 @@ mod base;
 mod core;
 #[cfg(feature = "file")]
 pub mod file;
-#[cfg(feature = "reqwest")]
-pub mod reqwest;
 pub use base::*;
 pub use core::*;