好的，我想我已经为你想出了一个解决方案，它使用了人造丝并行迭代器。
技巧是在rayon迭代器中使用Results，如果设置了取消标志，则返回Err。
我首先创建了一个实用程序类型来创建一个可取消的线程，你可以在其中执行人造丝迭代器。你可以通过传入线程闭包来使用它，线程闭包将原子取消令牌作为参数。然后你必须检查取消令牌是否为true，如果是，就提前退出。

use std::sync::Arc;
use std::sync::atomic::{Ordering, AtomicBool};
use std::thread::JoinHandle;

fn collate(batch: &[Computed]) -> Batch {
    batch.iter().map(|&x| i128::from(x)).sum()
}

# [derive(Debug)]

struct Cancelled;

struct CancellableThread<Output: Send + 'static> {
    cancel_token: Arc<AtomicBool>,
    thread: Option<JoinHandle<Result<Output, Cancelled>>>,
}
impl<Output: Send + 'static> CancellableThread<Output> {
    fn new<F: FnOnce(Arc<AtomicBool>) -> Result<Output, Cancelled> + Send + 'static>(init: F) -> Self {
        let cancel_token = Arc::new(AtomicBool::new(false));
        let thread_cancel_token = Arc::clone(&cancel_token);

        CancellableThread {
            thread: Some(std::thread::spawn(move || init(thread_cancel_token))),
            cancel_token,
        }
    }

    fn output(mut self) -> Output {
        self.thread.take().unwrap().join().unwrap().unwrap()
    }
}
impl<Output: Send + 'static> Drop for CancellableThread<Output> {
    fn drop(&mut self) {
        self.cancel_token.store(true, Ordering::Relaxed);

        if let Some(thread) = self.thread.take() {
            let _ = thread.join().unwrap();
        }
    }
}

我发现创建一个返回Result<(), Cancelled>的闭包很有用，这样我就可以使用try操作符（?）提前退出。

CancellableThread::new(move |cancel_token| {
    let cancelled = || if cancel_token.load(Ordering::Relaxed) {
        Err(Cancelled)
    } else {
        Ok(())
    };

    loop {
        // was the thread dropped?
        // if so, stop what we're doing
        cancelled?;

        // do stuff and 
        // eventually return a result
    }
});

然后，我在DataLoader中使用了CancellableThread抽象。不需要为它创建特殊的Drop impl，因为默认情况下，它将在每个字段上调用drop，这将处理取消操作。

type Data = Vec<u8>;
type Dataset = Vec<Data>;
type Computed = u64;
type Batch = i128;

use rayon::prelude::*;
use crossbeam::channel::{unbounded, Receiver};

struct DataLoader {
    example_worker: CancellableThread<()>,
    collate_worker: CancellableThread<()>,
    receiver: Receiver<Batch>,
    length: usize,
}

我用的是unbounded频道，因为它少了一件麻烦的事情。换用bounded频道应该不难。

impl DataLoader {
    fn new(dataset: Dataset, batch_size: usize) -> Self {
        let (example_sender, collate_receiver) = unbounded();
        let (batch_sender, final_receiver) = unbounded();

我不确定您是否总能保证数据集中的项数是batch_size的倍数，所以我决定显式地处理它。

let length = if dataset.len() % batch_size == 0 {
            dataset.len() / batch_size
        } else {
            dataset.len() / batch_size + 1
        };

我首先创建了排序工作者，尽管这可能不是必需的。正如您所看到的，我不得不复制一点来处理部分批处理。

let collate_worker = CancellableThread::new(move |cancel_token| {
            let cancelled = || if cancel_token.load(Ordering::Relaxed) {
                Err(Cancelled)
            } else {
                Ok(())
            };

            'outer: loop {
                let mut batch = Vec::with_capacity(batch_size);
                for _ in 0..batch_size {
                    cancelled()?;

                    if let Ok(data) = collate_receiver.recv() {
                        batch.push(data);
                    } else {
                        if !batch.is_empty() {
                            // handle the last batch, if there
                            // weren't enough items to fill it
                            let collated = collate(&batch);
                            cancelled()?;
                            batch_sender.send(collated).unwrap();
                        }

                        break 'outer;
                    }
                }

                let collated = collate(&batch);
                cancelled()?;
                batch_sender.send(collated).unwrap();
            }

            Ok(())
        });

在示例worker中，事情变得非常简单，因为我们可以只使用rayon并行迭代器。正如你所看到的，我们在每次繁重的计算之前检查取消。

let example_worker = CancellableThread::new(move |cancel_token| {
            let cancelled = || if cancel_token.load(Ordering::Relaxed) {
                Err(Cancelled)
            } else {
                Ok(())
            };

            let heavy_compute = |data: Data| -> Result<Computed, Cancelled> {
                cancelled()?;

                Ok(data.iter().map(|&x| u64::from(x)).product())
            };

            dataset
                .into_par_iter()
                .map(heavy_compute)
                .try_for_each(|computed| {
                    example_sender.send(computed?).unwrap();

                    Ok(())
                })
        });

然后我们构造DataLoader。你可以看到Python的实现是相同的：

DataLoader {
            example_worker,
            collate_worker,
            receiver: final_receiver,
            length,
        }
    }
}

// #[pymethods]
impl DataLoader {
    fn __iter__(this: Self /* PyRef<Self> */) -> Self /* PyRef<Self> */ { this }

    fn __next__(&mut self) -> Option<Batch> {
        self.receiver.recv().ok() 
    }

    fn __len__(&self) -> usize {
        self.length
    }
}

playground