delete segment after checkpoint

Author: MarinPostma

libsql-wal/src/registry.rs

@@ -36,7 +36,7 @@ enum Slot<IO: Io> {
 
 /// Wal Registry maintains a set of shared Wal, and their respective set of files.
 pub struct WalRegistry<IO: Io, S> {
-    io: IO,
+    io: Arc<IO>,
     path: PathBuf,
     shutdown: AtomicBool,
     opened: DashMap<NamespaceName, Slot<IO>>,
@@ -63,7 +63,7 @@ impl<IO: Io, S> WalRegistry<IO, S> {
     ) -> Result<Self> {
         io.create_dir_all(&path)?;
         let registry = Self {
-            io,
+            io: io.into(),
             path,
             opened: Default::default(),
             shutdown: Default::default(),
@@ -340,6 +340,7 @@ where
             shutdown: false.into(),
             checkpoint_notifier: self.checkpoint_notifier.clone(),
             max_segment_size: 1000.into(),
+            io: self.io.clone(),
         });
 
         self.opened

libsql-wal/src/segment/list.rs

@@ -4,16 +4,16 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
 use std::sync::Arc;
 
 use arc_swap::ArcSwapOption;
-use fst::map::{OpBuilder, Union};
 use fst::raw::IndexedValue;
 use fst::Streamer;
 use roaring::RoaringBitmap;
 use tokio_stream::Stream;
+use uuid::Uuid;
 use zerocopy::FromZeroes;
 
 use crate::error::Result;
 use crate::io::buf::{ZeroCopyBoxIoBuf, ZeroCopyBuf};
-use crate::io::FileExt;
+use crate::io::{FileExt, Io};
 use crate::segment::Frame;
 use crate::{LibsqlFooter, LIBSQL_MAGIC, LIBSQL_PAGE_SIZE, LIBSQL_WAL_VERSION};
 
@@ -78,10 +78,13 @@ where
 
     /// Checkpoints as many segments as possible to the main db file, and return the checkpointed
     /// frame_no, if anything was checkpointed
-    pub async fn checkpoint<F>(&self, db_file: &F, until_frame_no: u64) -> Result<Option<u64>>
-    where
-        F: FileExt,
-    {
+    pub async fn checkpoint<IO: Io>(
+        &self,
+        db_file: &IO::File,
+        until_frame_no: u64,
+        log_id: Uuid,
+        io: &IO,
+    ) -> Result<Option<u64>> {
         struct Guard<'a>(&'a AtomicBool);
         impl<'a> Drop for Guard<'a> {
             fn drop(&mut self) {
@@ -122,24 +125,13 @@ where
 
         let size_after = segs.first().unwrap().size_after();
 
-        let union = segs
-            .iter()
-            .map(|s| s.index())
-            .collect::<OpBuilder>()
-            .union();
-
-        /// Safety: Union contains a Box<dyn trait> that doesn't require Send, to it's not send.
-        /// That's an issue for us, but all the indexes we have are safe to send, so we're good.
-        /// FIXME: we could implement union ourselves.
-        unsafe impl Send for SendUnion<'_> {}
-        unsafe impl Sync for SendUnion<'_> {}
-        struct SendUnion<'a>(Union<'a>);
+        let index_iter = segs.iter().map(|s| s.index());
 
-        let mut union = SendUnion(union);
+        let mut union = send_fst_ops::SendUnion::from_index_iter(index_iter);
 
         let mut buf = ZeroCopyBuf::<Frame>::new_uninit();
         let mut last_replication_index = 0;
-        while let Some((k, v)) = union.0.next() {
+        while let Some((k, v)) = union.next() {
             let page_no = u32::from_be_bytes(k.try_into().unwrap());
             let v = v.iter().min_by_key(|i| i.index).unwrap();
             let offset = v.value as u32;
@@ -163,19 +155,25 @@ where
             magic: LIBSQL_MAGIC.into(),
             version: LIBSQL_WAL_VERSION.into(),
             replication_index: last_replication_index.into(),
+            log_id: log_id.as_u128().into(),
         };
 
+        db_file.set_len(size_after as u64 * LIBSQL_PAGE_SIZE as u64)?;
+
         let footer_offset = size_after as usize * LIBSQL_PAGE_SIZE as usize;
         let (_, ret) = db_file
             .write_all_at_async(ZeroCopyBuf::new_init(footer), footer_offset as u64)
             .await;
         ret?;
 
         // todo: truncate if necessary
-
-        //// todo: make async
+        //// TODO: make async
         db_file.sync_all()?;
 
+        for seg in segs.iter() {
+            seg.destroy(io).await;
+        }
+
         let mut current = self.head.compare_and_swap(&segs[0], None);
         if Arc::ptr_eq(&segs[0], current.as_ref().unwrap()) {
             // nothing to do
@@ -196,8 +194,6 @@ where
 
         self.len.fetch_sub(segs.len(), Ordering::Relaxed);
 
-        db_file.set_len(size_after as u64 * 4096)?;
-
         Ok(Some(last_replication_index))
     }
 
@@ -241,7 +237,8 @@ where
             .max(until_fno);
 
         let stream = async_stream::try_stream! {
-            let mut union = fst::map::OpBuilder::from_iter(segments.iter().map(|s| s.index())).union();
+            let index_iter = segments.iter().map(|s| s.index());
+            let mut union = send_fst_ops::SendUnion::from_index_iter(index_iter);
             while let Some((key_bytes, indexes)) = union.next() {
                 let page_no = u32::from_be_bytes(key_bytes.try_into().unwrap());
                 // we already have a more recent version of this page.
@@ -337,6 +334,46 @@ impl<T> List<T> {
     }
 }
 
+mod send_fst_ops {
+    use std::ops::{Deref, DerefMut};
+    use std::sync::Arc;
+
+    use fst::map::{OpBuilder, Union};
+
+    /// Safety: Union contains a Box<dyn trait> that doesn't require Send, to it's not send.
+    /// That's an issue for us, but all the indexes we have are safe to send, so we're good.
+    /// FIXME: we could implement union ourselves.
+    unsafe impl Send for SendUnion<'_> {}
+    unsafe impl Sync for SendUnion<'_> {}
+
+    #[repr(transparent)]
+    pub(super) struct SendUnion<'a>(Union<'a>);
+
+    impl<'a> SendUnion<'a> {
+        pub fn from_index_iter<I>(iter: I) -> Self
+        where
+            I: Iterator<Item = &'a fst::map::Map<Arc<[u8]>>>,
+        {
+            let op = iter.collect::<OpBuilder>().union();
+            Self(op)
+        }
+    }
+
+    impl<'a> Deref for SendUnion<'a> {
+        type Target = Union<'a>;
+
+        fn deref(&self) -> &Self::Target {
+            &self.0
+        }
+    }
+
+    impl<'a> DerefMut for SendUnion<'a> {
+        fn deref_mut(&mut self) -> &mut Self::Target {
+            &mut self.0
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::io::{Read, Seek, Write};
@@ -393,8 +430,8 @@ mod test {
         *shared.durable_frame_no.lock() = 999999;
         shared.checkpoint().await.unwrap();
         file.seek(std::io::SeekFrom::Start(0)).unwrap();
-        let mut copy_ytes = Vec::new();
-        file.read_to_end(&mut copy_ytes).unwrap();
+        let mut copy_bytes = Vec::new();
+        file.read_to_end(&mut copy_bytes).unwrap();
 
         let mut orig_bytes = Vec::new();
         shared

libsql-wal/src/shared_wal.rs

@@ -53,6 +53,7 @@ pub struct SharedWal<IO: Io> {
     pub(crate) checkpoint_notifier: mpsc::Sender<CheckpointMessage>,
     /// maximum size the segment is allowed to grow
     pub(crate) max_segment_size: AtomicUsize,
+    pub(crate) io: Arc<IO>,
 }
 
 impl<IO: Io> SharedWal<IO> {
@@ -301,7 +302,7 @@ impl<IO: Io> SharedWal<IO> {
             .current
             .load()
             .tail()
-            .checkpoint(&self.db_file, durable_frame_no)
+            .checkpoint(&self.db_file, durable_frame_no, self.log_id(), &self.io)
             .await?;
         if let Some(checkpointed_frame_no) = checkpointed_frame_no {
             self.checkpointed_frame_no