fix(fast-steal): 修复任务窃取中 min_chunk_size 含义错误的问题

Author: share121

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

@@ -4,7 +4,7 @@ use crate::{DownloadResult, Event, ProgressEntry, RandPuller, RandPusher, Total,
 use alloc::{sync::Arc, vec::Vec};
 use bytes::Bytes;
 use core::{
-    num::{NonZero, NonZeroU64, NonZeroUsize},
+    num::{NonZeroU64, NonZeroUsize},
     sync::atomic::{AtomicUsize, Ordering},
     time::Duration,
 };
@@ -98,13 +98,12 @@ where
     type Handle = TokioHandle;
     fn execute(self: Arc<Self>, task: Arc<Task>, task_list: Arc<TaskList<Self>>) -> Self::Handle {
         let id = self.id.fetch_add(1, Ordering::SeqCst);
-        let steal_min_chunk_size = NonZero::new(2 * self.min_chunk_size.get()).unwrap();
         let mut puller = self.puller.clone();
         let handle = tokio::spawn(async move {
             'steal_task: loop {
                 let mut start = task.start();
                 if start >= task.end() {
-                    if task_list.steal(&task, steal_min_chunk_size) {
+                    if task_list.steal(&task, self.min_chunk_size) {
                         continue;
                     } else {
                         break;
@@ -165,6 +164,7 @@ mod tests {
         mock::{MockPuller, build_mock_data},
     };
     use alloc::vec;
+    use core::num::NonZero;
     use std::dbg;
 
     #[tokio::test]

crates/fast-steal/README.md

@@ -49,7 +49,7 @@ impl Executor for TokioExecutor {
                     println!("task: {i} = {res}");
                     self.tx.send((i, res)).unwrap();
                 }
-                if !task_list.steal(&task, NonZero::new(2).unwrap()) {
+                if !task_list.steal(&task, NonZero::new(1).unwrap()) {
                     break;
                 }
             }
@@ -81,8 +81,8 @@ async fn main() {
     let (tx, mut rx) = mpsc::unbounded_channel();
     let executor = TokioExecutor { tx };
     let pre_data = [1..20, 41..48];
-    let task_list = TaskList::run(NonZero::new(8).unwrap(), NonZero::new(2).unwrap(), &pre_data[..], executor);
-    let handles: Arc<[_]> = task_list.handles(|it| it.collect());
+    let task_list = TaskList::run(NonZero::new(8).unwrap(), NonZero::new(1).unwrap(), &pre_data[..], executor);
+    let handles: Arc<[_]> = task_list.handles(|it| it.map(|h| h.clone()).collect());
     drop(task_list);
     for handle in handles.iter() {
         handle.0.lock().await.take().unwrap().await.unwrap();

crates/fast-steal/src/lib.rs

@@ -50,7 +50,7 @@
 //!                     println!("task: {i} = {res}");
 //!                     self.tx.send((i, res)).unwrap();
 //!                 }
-//!                 if !task_list.steal(&task, NonZero::new(2).unwrap()) {
+//!                 if !task_list.steal(&task, NonZero::new(1).unwrap()) {
 //!                     break;
 //!                 }
 //!             }
@@ -85,8 +85,8 @@
 //!     let task_list = Arc::new(TaskList::run(&pre_data[..], executor));
 //!     task_list
 //!         .clone()
-//!         .set_threads(NonZero::new(8).unwrap(), NonZero::new(2).unwrap());
-//!     let handles: Arc<[_]> = task_list.handles(|it| it.map(|task| task.clone()).collect());
+//!         .set_threads(NonZero::new(8).unwrap(), NonZero::new(1).unwrap());
+//!     let handles: Arc<[_]> = task_list.handles(|it| it.map(|h| h.clone()).collect());
 //!     drop(task_list);
 //!     for handle in handles.iter() {
 //!         handle.0.lock().await.take().unwrap().await.unwrap();

crates/fast-steal/src/task_list.rs

@@ -51,7 +51,7 @@ impl<E: Executor> TaskList<E> {
             return true;
         }
         if let Some(steal_task) = guard.running.iter().map(|w| &w.0).max()
-            && steal_task.remain() >= min_chunk_size
+            && steal_task.remain() >= min_chunk_size * 2
         {
             let (start, end) = steal_task.split_two();
             task.set_end(end);
@@ -78,7 +78,7 @@ impl<E: Executor> TaskList<E> {
             guard.running.extend(temp);
             while guard.running.len() < threads
                 && let Some(steal_task) = guard.running.iter().map(|w| &w.0).max()
-                && steal_task.remain() >= min_chunk_size
+                && steal_task.remain() >= min_chunk_size * 2
             {
                 let (start, end) = steal_task.split_two();
                 let task = Arc::new(Task::new(start, end));
@@ -147,7 +147,7 @@ mod tests {
                         println!("task: {i} = {res}");
                         self.tx.send((i, res)).unwrap();
                     }
-                    if !task_list.steal(&task, NonZero::new(2).unwrap()) {
+                    if !task_list.steal(&task, NonZero::new(1).unwrap()) {
                         break;
                     }
                 }
@@ -182,8 +182,8 @@ mod tests {
         let task_list = Arc::new(TaskList::run(&pre_data[..], executor));
         task_list
             .clone()
-            .set_threads(NonZero::new(8).unwrap(), NonZero::new(2).unwrap());
-        let handles: Arc<[_]> = task_list.handles(|it| it.map(|task| task.clone()).collect());
+            .set_threads(NonZero::new(8).unwrap(), NonZero::new(1).unwrap());
+        let handles: Arc<[_]> = task_list.handles(|it| it.map(|h| h.clone()).collect());
         drop(task_list);
         for handle in handles.iter() {
             handle.0.lock().await.take().unwrap().await.unwrap();