instead of having constantly running thread, segment creator v2 spans… (#98)

* instead of having constantly running thread, segment creator v2 spans threads on demand

* fmt

* linter fixes

* windows clippy

* Use sort_by_key

---------

Co-authored-by: Tim Visée <tim@visee.me>

Author: generall

src/lib.rs

@@ -1,6 +1,6 @@
-use crossbeam_channel::{Receiver, Sender};
+use crate::segment_creator::SegmentCreatorV2;
 use fs4::fs_std::FileExt;
-use log::{debug, info, trace, warn};
+use log::{debug, info, trace};
 pub use segment::{Entry, Segment};
 use std::cmp::Ordering;
 use std::collections::HashMap;
@@ -17,6 +17,7 @@ use std::thread;
 
 mod mmap_view_sync;
 mod segment;
+mod segment_creator;
 pub mod test_utils;
 
 #[derive(Debug)]
@@ -75,7 +76,7 @@ pub struct Wal {
     /// The segment currently being appended to.
     open_segment: OpenSegment,
     closed_segments: Vec<ClosedSegment>,
-    creator: SegmentCreator,
+    creator: SegmentCreatorV2,
 
     /// The number of closed segments to retain.
     retain_closed: NonZeroUsize,
@@ -226,7 +227,7 @@ impl Wal {
             }
         }
 
-        let mut creator = SegmentCreator::new(
+        let mut creator = SegmentCreatorV2::new(
             &path,
             unused_segments,
             options.segment_capacity,
@@ -604,122 +605,15 @@ fn open_dir_entry(entry: fs::DirEntry) -> Result<Option<WalSegment>> {
     }
 }
 
-struct SegmentCreator {
-    /// Receive channel for new segments.
-    rx: Option<Receiver<OpenSegment>>,
-    /// The segment creator thread.
-    ///
-    /// Used for retrieving error upon failure.
-    thread: Option<thread::JoinHandle<Result<()>>>,
-}
-
-impl SegmentCreator {
-    /// Creates a new segment creator.
-    ///
-    /// The segment creator must be started before new segments will be created.
-    pub fn new<P>(
-        dir: P,
-        existing: Vec<OpenSegment>,
-        segment_capacity: usize,
-        segment_queue_len: usize,
-    ) -> SegmentCreator
-    where
-        P: AsRef<Path>,
-    {
-        let (tx, rx) = crossbeam_channel::bounded(segment_queue_len);
-
-        let dir = dir.as_ref().to_path_buf();
-        let thread = thread::Builder::new()
-            .name("wal-segment-creator".to_string())
-            .spawn(move || create_loop(tx, dir, segment_capacity, existing))
-            .unwrap();
-        SegmentCreator {
-            rx: Some(rx),
-            thread: Some(thread),
-        }
-    }
-
-    /// Retrieves the next segment.
-    pub fn next(&mut self) -> Result<OpenSegment> {
-        self.rx.as_mut().unwrap().recv().map_err(|_| {
-            match self.thread.take().map(|join_handle| join_handle.join()) {
-                Some(Ok(Err(error))) => error,
-                None => Error::other("segment creator thread already failed"),
-                Some(Ok(Ok(()))) => unreachable!(
-                    "segment creator thread finished without an error,
-                                                  but the segment creator is still live"
-                ),
-                Some(Err(_)) => unreachable!("segment creator thread panicked"),
-            }
-        })
-    }
-}
-
-impl Drop for SegmentCreator {
-    fn drop(&mut self) {
-        drop(self.rx.take());
-        if let Some(join_handle) = self.thread.take()
-            && let Err(error) = join_handle.join()
-        {
-            warn!("Error while shutting down segment creator: {error:?}");
-        }
-    }
-}
-
-fn create_loop(
-    tx: Sender<OpenSegment>,
-    mut path: PathBuf,
-    capacity: usize,
-    mut existing_segments: Vec<OpenSegment>,
-) -> Result<()> {
-    // Ensure the existing segments are in ID order.
-    existing_segments.sort_by(|a, b| a.id.cmp(&b.id));
-
-    let mut cont = true;
-    let mut id = 0;
-
-    for segment in existing_segments {
-        id = segment.id;
-        if tx.send(segment).is_err() {
-            cont = false;
-            break;
-        }
-    }
-
-    // Directory being a file only applies to Linux
-    #[cfg(not(target_os = "windows"))]
-    let dir = File::open(&path)?;
-
-    while cont {
-        id += 1;
-        path.push(format!("open-{id}"));
-        let segment = OpenSegment {
-            id,
-            segment: Segment::create(&path, capacity)?,
-        };
-        path.pop();
-        // Sync the directory, guaranteeing that the segment file is durably
-        // stored on the filesystem.
-        #[cfg(not(target_os = "windows"))]
-        dir.sync_all()?;
-        cont = tx.send(segment).is_ok();
-    }
-
-    info!("segment creator shutting down");
-    Ok(())
-}
-
 #[cfg(test)]
 mod test {
+    use crate::test_utils::EntryGenerator;
     use log::trace;
     use quickcheck::TestResult;
     use std::{io::Write, num::NonZeroUsize};
     use tempfile::Builder;
 
-    use crate::segment::Segment;
-    use crate::test_utils::EntryGenerator;
-
-    use super::{OpenSegment, SegmentCreator, Wal, WalOptions};
+    use super::{Wal, WalOptions};
 
     fn init_logger() {
         let _ = env_logger::builder().is_test(true).try_init();
@@ -1330,22 +1224,6 @@ mod test {
         assert!(Wal::open(dir.path()).is_ok());
     }
 
-    #[test]
-    fn test_segment_creator() {
-        init_logger();
-        let dir = Builder::new().prefix("segment").tempdir().unwrap();
-
-        let segments = vec![OpenSegment {
-            id: 3,
-            segment: Segment::create(dir.path().join("open-3"), 1024).unwrap(),
-        }];
-
-        let mut creator = SegmentCreator::new(dir.path(), segments, 1024, 1);
-        for i in 3..10 {
-            assert_eq!(i, creator.next().unwrap().id);
-        }
-    }
-
     #[test]
     fn test_record_id_preserving() {
         init_logger();

src/segment_creator.rs

@@ -0,0 +1,196 @@
+use crate::{OpenSegment, Segment};
+use log::warn;
+#[cfg(not(target_os = "windows"))]
+use std::fs::File;
+use std::io::Error;
+use std::path::{Path, PathBuf};
+use std::thread;
+
+pub struct SegmentCreatorV2 {
+    /// Where to create segments
+    dir: PathBuf,
+    /// Segments, which are already created, but not yet handed out.
+    pending_segments: Vec<OpenSegment>,
+
+    /// Capacity of segments to create.
+    segment_capacity: usize,
+
+    /// How many segments to create ahead of time.
+    segment_queue_len: usize,
+
+    /// Current id to create
+    current_id: u64,
+
+    /// The segment creator thread.
+    /// Used to create segments asynchronously.
+    thread: Option<thread::JoinHandle<std::io::Result<Vec<OpenSegment>>>>,
+}
+
+fn create_segments(
+    path: PathBuf,
+    start_id: u64,
+    count: usize,
+    segment_capacity: usize,
+) -> std::io::Result<Vec<OpenSegment>> {
+    // Directory being a file only applies to Linux
+    #[cfg(not(target_os = "windows"))]
+    let dir = File::open(&path)?;
+
+    let mut segments = Vec::with_capacity(count);
+
+    for i in 0..count {
+        let segment_id = start_id + i as u64;
+        let segment_path = path.join(format!("open-{segment_id}"));
+        let segment = OpenSegment {
+            id: segment_id,
+            segment: Segment::create(segment_path, segment_capacity)?,
+        };
+        segments.push(segment);
+    }
+
+    #[cfg(not(target_os = "windows"))]
+    dir.sync_all()?;
+
+    Ok(segments)
+}
+
+impl SegmentCreatorV2 {
+    /// Creates a new segment creator.
+    ///
+    /// The segment creator must be started before new segments will be created.
+    pub fn new<P>(
+        dir: P,
+        mut existing: Vec<OpenSegment>,
+        segment_capacity: usize,
+        segment_queue_len: usize,
+    ) -> SegmentCreatorV2
+    where
+        P: AsRef<Path>,
+    {
+        let dir = dir.as_ref().to_path_buf();
+
+        existing.sort_by_key(|segment| segment.id);
+
+        let current_id = existing.last().map_or(1, |s| s.id + 1);
+
+        let mut result = Self {
+            dir,
+            pending_segments: existing,
+            segment_capacity,
+            segment_queue_len: segment_queue_len + 1, // Always create at least one segment
+            current_id,
+            thread: None,
+        };
+
+        result.schedule_creation();
+
+        result
+    }
+
+    fn schedule_creation(&mut self) {
+        if self.thread.is_none() && self.pending_segments.len() < self.segment_queue_len {
+            let count = self
+                .segment_queue_len
+                .saturating_sub(self.pending_segments.len());
+            let dir = self.dir.clone();
+            let start_id = self.current_id;
+            let segment_capacity = self.segment_capacity;
+
+            self.thread = Some(
+                thread::Builder::new()
+                    .name("wal-segment-creator".to_string())
+                    .spawn(move || create_segments(dir, start_id, count, segment_capacity))
+                    .unwrap(),
+            );
+        }
+    }
+
+    fn refill_pending(&mut self) -> std::io::Result<()> {
+        if let Some(thread) = self.thread.take() {
+            if self.pending_segments.is_empty() || thread.is_finished() {
+                let new_segments = thread
+                    .join()
+                    .map_err(|_| Error::other("segment creation thread panicked"))??;
+                self.current_id += new_segments.len() as u64;
+                self.pending_segments.extend(new_segments);
+            } else {
+                // Put the thread back if it's not finished, and we don't need segments yet.
+                self.thread = Some(thread);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn next(&mut self) -> std::io::Result<OpenSegment> {
+        self.refill_pending()?;
+
+        if self.pending_segments.is_empty() {
+            return Err(Error::other("no segments available"));
+        }
+
+        let segment = self.pending_segments.remove(0);
+        self.schedule_creation();
+        Ok(segment)
+    }
+}
+
+impl Drop for SegmentCreatorV2 {
+    fn drop(&mut self) {
+        if let Some(thread) = self.thread.take() {
+            match thread.join() {
+                Ok(Ok(_)) => {}
+                Ok(Err(err)) => {
+                    warn!("Error while shutting down segment creator: {err:?}");
+                }
+                Err(err) => {
+                    warn!("Segment creator thread panicked: {err:?}");
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tempfile::Builder;
+
+    fn init_logger() {
+        let _ = env_logger::builder().is_test(true).try_init();
+    }
+
+    #[test]
+    fn test_segment_creator_v2() {
+        init_logger();
+        let dir = Builder::new().prefix("segment").tempdir().unwrap();
+
+        let segments = vec![OpenSegment {
+            id: 3,
+            segment: Segment::create(dir.path().join("open-3"), 1024).unwrap(),
+        }];
+
+        let mut creator = SegmentCreatorV2::new(dir.path(), segments, 1024, 1);
+        for i in 3..10 {
+            assert_eq!(i, creator.next().unwrap().id);
+        }
+
+        // This awaits for the thread to finish
+        drop(creator);
+
+        // List directory contents
+        let mut entries: Vec<_> = std::fs::read_dir(dir.path())
+            .unwrap()
+            .map(|res| res.map(|e| e.file_name()))
+            .collect::<Result<_, std::io::Error>>()
+            .unwrap();
+
+        entries.sort();
+
+        for entry in &entries {
+            eprintln!("{:?}", entry);
+        }
+
+        assert_eq!(entries.len(), 10 - 3 + 1); // open-3 existed, open-4 to open-9 created + open-10 created ahead
+    }
+}