Generate data in parallel

Author: alamb

Diff for: tpchgen-cli/Cargo.toml

@@ -7,3 +7,6 @@ authors = ["clflushopt", "alamb"]
 [dependencies]
 clap = { version = "4.5.32", features = ["derive"] }
 tpchgen = { path = "../tpchgen" }
+tokio = { version = "1.44.1", features = ["full"]}
+futures = "0.3.31"
+num_cpus = "1.0"

Diff for: tpchgen-cli/src/generate.rs

@@ -0,0 +1,125 @@
+//! Parallel data generation support: [`Source`] and [`Sink`] and [`generate_in_chunks`]
+
+use futures::StreamExt;
+use std::collections::VecDeque;
+use std::io;
+use std::sync::{Arc, Mutex};
+use tokio::task::JoinSet;
+
+/// Something that knows how to generate data into a buffer
+///
+/// For example, this is implemented for the different generators in the tpchgen
+/// crate
+pub trait Source: Send {
+    /// generates the data for this generator into the buffer, returning the buffer.
+    fn create(self, buffer: Vec<u8>) -> Vec<u8>;
+}
+
+/// Something that can write the contents of a buffer somewhere
+///
+/// For example, this is implemented for a file writer
+pub trait Sink: Send {
+    /// Write all data from the buffer to the sink
+    fn sink(&mut self, buffer: &[u8]);
+
+    fn finish(self) -> Result<(), io::Error>;
+}
+
+/// Creates data from the Generators in parallel and invokes the provided
+/// function on each generated buffer
+///
+///
+/// G: Generator
+/// I: Iterator<Item = G>
+/// S: Sink that writes buffers somewhere
+pub async fn generate_in_chunks<G, I, S>(mut sink: S, generators: I) -> Result<(), io::Error>
+where
+    G: Source + 'static,
+    I: Iterator<Item = G>,
+    S: Sink + 'static,
+{
+    let recycler = BufferRecycler::new();
+
+    // use all cores to make data
+    let num_tasks = num_cpus::get();
+    println!("Using {num_tasks} parallel tasks");
+
+    // create a channel to communicate between the generator tasks and the writer task
+    let (tx, mut rx) = tokio::sync::mpsc::channel(num_tasks);
+
+    let generators_and_recyclers = generators.map(|generator| (generator, recycler.clone()));
+
+    // convert to an async stream to run on tokio
+    let mut stream = futures::stream::iter(generators_and_recyclers)
+        // each generator writes to a buffer
+        .map(async |(generator, recycler)| {
+            let buffer = recycler.new_buffer(1024 * 1024 * 8);
+            // do the work in a task (on a different thread)
+            let mut join_set = JoinSet::new();
+            join_set.spawn(async move { generator.create(buffer) });
+            // wait for the task to be done and return the result
+            let buffer = join_set
+                .join_next()
+                .await
+                .expect("had one item")
+                .expect("join_next join is infallible unless task panics");
+            // send the buffer to the writer task, ignoring error (if the writer errored)
+            tx.send(buffer).await.ok();
+        })
+        // run in parallel
+        .buffered(num_tasks);
+
+    let captured_recycler = recycler.clone();
+    let writer_task = tokio::task::spawn_blocking(move || {
+        while let Some(buffer) = rx.blocking_recv() {
+            //println!("writing buffer with {} bytes", buffer.len());
+            sink.sink(&buffer);
+            captured_recycler.return_buffer(buffer);
+        }
+        sink.finish()
+    });
+
+    // drive the stream to completion
+    while stream.next().await.is_some() {}
+    println!("stream done, dropping");
+    drop(stream); // drop any stream references
+    drop(tx); // drop last tx reference to stop the writer
+
+    // wait for writer to finish
+    println!("waiting on writer task");
+    writer_task.await.expect("writer task panicked")
+}
+
+/// A simple buffer recycler to avoid allocating new buffers for each part
+///
+/// Clones share the same undrlying recycler, so it is not thread safe
+#[derive(Debug, Clone)]
+struct BufferRecycler {
+    buffers: Arc<Mutex<VecDeque<Vec<u8>>>>,
+}
+
+impl BufferRecycler {
+    fn new() -> Self {
+        Self {
+            buffers: Arc::new(Mutex::new(VecDeque::new())),
+        }
+    }
+    /// return a new empty buffer, with size bytes capacity
+    fn new_buffer(&self, size: usize) -> Vec<u8> {
+        let mut buffers = self.buffers.lock().unwrap();
+        if let Some(mut buffer) = buffers.pop_front() {
+            buffer.clear();
+            if size > buffer.capacity() {
+                buffer.reserve(size - buffer.capacity());
+            }
+            buffer
+        } else {
+            Vec::with_capacity(size)
+        }
+    }
+
+    fn return_buffer(&self, buffer: Vec<u8>) {
+        let mut buffers = self.buffers.lock().unwrap();
+        buffers.push_back(buffer);
+    }
+}