pass callback instead of spawning task for durable frameno notify

MarinPostma · MarinPostma · commit 89c7f7dcee71 · 2024-07-25T19:57:18.000+02:00
diff --git a/libsql-wal/src/registry.rs b/libsql-wal/src/registry.rs
@@ -1,4 +1,3 @@
-use std::future::Future;
 use std::num::NonZeroU64;
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -19,7 +18,7 @@ use crate::replication::storage::StorageReplicator;
 use crate::segment::list::SegmentList;
 use crate::segment::{current::CurrentSegment, sealed::SealedSegment};
 use crate::shared_wal::{SharedWal, SwapLog};
-use crate::storage::Storage;
+use crate::storage::{OnStoreCallback, Storage};
 use crate::transaction::TxGuard;
 use libsql_sys::name::NamespaceName;
 
@@ -119,11 +118,13 @@ where
         // where the current log is sealed and it wasn't swapped.
         if let Some(sealed) = current.seal()? {
             // todo: pass config override here
-            let notify = self.storage.store(&shared.namespace, sealed.clone(), None);
             let notifier = self.checkpoint_notifier.clone();
             let namespace = shared.namespace().clone();
             let durable_frame_no = shared.durable_frame_no.clone();
-            tokio::spawn(update_durable(notify, notifier, durable_frame_no, namespace));
+            let cb: OnStoreCallback = Box::new(move |fno| Box::pin(async move {
+                update_durable(fno, notifier, durable_frame_no, namespace).await;
+            }));
+            self.storage.store(&shared.namespace, sealed.clone(), None, cb);
             new.tail().push(sealed);
         }
 
@@ -135,14 +136,13 @@ where
 }
 
 async fn update_durable(
-    notify: impl Future<Output = u64>,
+    new_durable: u64,
     notifier: mpsc::Sender<CheckpointMessage>,
-    durable_frame_no: Arc<Mutex<u64>>,
+    durable_frame_no_slot: Arc<Mutex<u64>>,
     namespace: NamespaceName,
     ) {
-    let new_durable = notify.await;
     {
-        let mut g = durable_frame_no.lock();
+        let mut g = durable_frame_no_slot.lock();
         if *g < new_durable {
             *g = new_durable;
         }
@@ -253,12 +253,14 @@ where
             if let Some(sealed) =
                 SealedSegment::open(file.into(), entry.path().to_path_buf(), Default::default())?
             {
-                // TODO: pass config override here
-                let notify = self.storage.store(&namespace, sealed.clone(), None);
                 let notifier = self.checkpoint_notifier.clone();
-                let namespace = namespace.clone(); 
+                let ns = namespace.clone(); 
                 let durable_frame_no = durable_frame_no.clone();
-                tokio::spawn(update_durable(notify, notifier, durable_frame_no, namespace));
+                let cb: OnStoreCallback = Box::new(move |fno| Box::pin(async move {
+                    update_durable(fno, notifier, durable_frame_no, ns).await;
+                }));
+                // TODO: pass config override here
+                self.storage.store(&namespace, sealed.clone(), None, cb);
                 tail.push(sealed);
             }
         }
diff --git a/libsql-wal/src/storage/async_storage.rs b/libsql-wal/src/storage/async_storage.rs
@@ -1,7 +1,6 @@
 //! `AsyncStorage` is a `Storage` implementation that defer storage to a background thread. The
 //! durable frame_no is notified asynchronously.
 
-use std::future::Future;
 use std::sync::Arc;
 
 use chrono::Utc;
@@ -15,7 +14,7 @@ use crate::segment::Segment;
 
 use super::backend::Backend;
 use super::scheduler::Scheduler;
-use super::{RestoreOptions, Storage, StoreSegmentRequest};
+use super::{OnStoreCallback, RestoreOptions, Storage, StoreSegmentRequest};
 
 /// Background loop task state.
 ///
@@ -86,8 +85,8 @@ where
                 }
                 msg = self.receiver.recv(), if !shutting_down => {
                     match msg {
-                        Some(StorageLoopMessage::StoreReq(req, ret)) => {
-                            self.scheduler.register(req, ret);
+                        Some(StorageLoopMessage::StoreReq(req)) => {
+                            self.scheduler.register(req);
                         }
                         Some(StorageLoopMessage::DurableFrameNoReq { namespace, ret, config_override }) => {
                             self.fetch_durable_frame_no_async(namespace, ret, config_override);
@@ -134,7 +133,7 @@ pub struct BottomlessConfig<C> {
 }
 
 enum StorageLoopMessage<C, S> {
-    StoreReq(StoreSegmentRequest<C, S>, oneshot::Sender<u64>),
+    StoreReq(StoreSegmentRequest<C, S>),
     DurableFrameNoReq {
         namespace: NamespaceName,
         config_override: Option<Arc<C>>,
@@ -164,22 +163,19 @@ where
         namespace: &NamespaceName,
         segment: Self::Segment,
         config_override: Option<Arc<Self::Config>>,
-    ) -> impl Future<Output = u64>  + Send + Sync + 'static{
+        on_store_callback: OnStoreCallback,
+    ){
         let req = StoreSegmentRequest {
             namespace: namespace.clone(),
             segment,
             created_at: Utc::now(),
             storage_config_override: config_override,
+            on_store_callback,
         };
 
-        let (sender, receiver) = oneshot::channel();
         self.job_sender
-            .send(StorageLoopMessage::StoreReq(req, sender))
+            .send(StorageLoopMessage::StoreReq(req))
             .expect("bottomless loop was closed before the handle was dropped");
-
-        async move {
-            receiver.await.unwrap()
-        }
     }
 
     async fn durable_frame_no(
diff --git a/libsql-wal/src/storage/job.rs b/libsql-wal/src/storage/job.rs
@@ -1,7 +1,5 @@
 use std::ops::Deref;
 
-use tokio::sync::oneshot;
-
 use super::backend::Backend;
 use super::backend::SegmentMeta;
 use super::Result;
@@ -14,7 +12,6 @@ use crate::segment::Segment;
 pub(crate) struct IndexedRequest<C, T> {
     pub(crate) request: StoreSegmentRequest<C, T>,
     pub(crate) id: u64,
-    pub(crate) ret: oneshot::Sender<u64>,
 }
 
 impl<C, T> Deref for IndexedRequest<C, T> {
@@ -102,6 +99,7 @@ pub(crate) struct JobResult<C, S> {
 
 #[cfg(test)]
 mod test {
+    use std::future::ready;
     // use std::fs::File;
     // use std::io::Write;
     // use std::mem::size_of;
@@ -489,7 +487,7 @@ mod test {
                 _config: &Self::Config,
                 _namespace: &NamespaceName,
                 _key: &SegmentKey,
-                _file: &impl FileExt
+                _file: &impl FileExt,
             ) -> Result<()> {
                 todo!()
             }
@@ -521,9 +519,9 @@ mod test {
                     segment: TestSegment,
                     created_at: Utc::now(),
                     storage_config_override: None,
+                    on_store_callback: Box::new(|_| Box::pin(ready(()))),
                 },
                 id: 0,
-                ret: oneshot::channel().0,
             },
         };
 
diff --git a/libsql-wal/src/storage/mod.rs b/libsql-wal/src/storage/mod.rs
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
 use std::fmt;
 use std::future::Future;
 use std::path::PathBuf;
+use std::pin::Pin;
 use std::str::FromStr;
 use std::sync::Arc;
 
@@ -121,6 +122,14 @@ impl fmt::Display for SegmentKey {
     }
 }
 
+/// takes the new durable frame_no and returns a future
+pub type OnStoreCallback = Box<
+    dyn FnOnce(u64) -> Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>
+        + Send
+        + Sync
+        + 'static,
+>;
+
 pub trait Storage: Send + Sync + 'static {
     type Segment: Segment;
     type Config;
@@ -133,7 +142,8 @@ pub trait Storage: Send + Sync + 'static {
         namespace: &NamespaceName,
         seg: Self::Segment,
         config_override: Option<Arc<Self::Config>>,
-    ) -> impl Future<Output = u64> + Send + Sync + 'static;
+        on_store: OnStoreCallback,
+    );
 
     fn durable_frame_no_sync(
         &self,
@@ -190,8 +200,8 @@ impl Storage for NoStorage {
         _namespace: &NamespaceName,
         _seg: Self::Segment,
         _config: Option<Arc<Self::Config>>,
-    ) -> impl Future<Output = u64>  + Send + Sync + 'static{
-        std::future::ready(u64::MAX)
+        _on_store: OnStoreCallback,
+    ) {
     }
 
     async fn durable_frame_no(
@@ -286,12 +296,15 @@ impl<IO: Io> TestStorage<IO> {
     pub fn new_io(store: bool, io: IO) -> Self {
         let dir = tempdir().unwrap();
         Self {
-            inner: Arc::new(TestStorageInner {
-                dir,
-                stored: Default::default(),
-                io,
-                store,
-            }.into()),
+            inner: Arc::new(
+                TestStorageInner {
+                    dir,
+                    stored: Default::default(),
+                    io,
+                    store,
+                }
+                .into(),
+            ),
         }
     }
 }
@@ -305,14 +318,17 @@ impl<IO: Io> Storage for TestStorage<IO> {
         namespace: &NamespaceName,
         seg: Self::Segment,
         _config: Option<Arc<Self::Config>>,
-    ) -> impl Future<Output = u64>  + Send + Sync + 'static{
+        on_store: OnStoreCallback,
+    ) {
         let mut inner = self.inner.lock();
         if inner.store {
             let id = uuid::Uuid::new_v4();
             let out_path = inner.dir.path().join(id.to_string());
             let out_file = inner.io.open(true, true, true, &out_path).unwrap();
-            let index = tokio::runtime::Handle::current().block_on(seg.compact(&out_file, id)).unwrap();
-            let end_frame_no =  seg.header().last_committed();
+            let index = tokio::runtime::Handle::current()
+                .block_on(seg.compact(&out_file, id))
+                .unwrap();
+            let end_frame_no = seg.header().last_committed();
             let key = SegmentKey {
                 start_frame_no: seg.header().start_frame_no.get(),
                 end_frame_no,
@@ -323,9 +339,7 @@ impl<IO: Io> Storage for TestStorage<IO> {
                 .entry(namespace.clone())
                 .or_default()
                 .insert(key, (out_path, index));
-            std::future::ready(end_frame_no)
-        } else {
-            std::future::ready(u64::MAX)
+            tokio::runtime::Handle::current().block_on(on_store(end_frame_no))
         }
     }
 
@@ -354,12 +368,10 @@ impl<IO: Io> Storage for TestStorage<IO> {
     ) -> u64 {
         let inner = self.inner.lock();
         if inner.store {
-            let Some(segs) = inner.stored.get(namespace) else { return 0 };
-            segs
-                .keys()
-                .map(|k| k.end_frame_no)
-                .max()
-                .unwrap_or(0)
+            let Some(segs) = inner.stored.get(namespace) else {
+                return 0;
+            };
+            segs.keys().map(|k| k.end_frame_no).max().unwrap_or(0)
         } else {
             u64::MAX
         }
@@ -374,9 +386,10 @@ impl<IO: Io> Storage for TestStorage<IO> {
         let inner = self.inner.lock();
         if inner.store {
             if let Some(segs) = inner.stored.get(namespace) {
-                let Some((key, _path)) = segs.iter().find(|(k, _)| k.includes(frame_no))
-                    else { return Err(Error::FrameNotFound(frame_no)) };
-                    return Ok(*key)
+                let Some((key, _path)) = segs.iter().find(|(k, _)| k.includes(frame_no)) else {
+                    return Err(Error::FrameNotFound(frame_no));
+                };
+                return Ok(*key);
             } else {
                 panic!("namespace not found");
             }
@@ -394,12 +407,9 @@ impl<IO: Io> Storage for TestStorage<IO> {
         let inner = self.inner.lock();
         if inner.store {
             match inner.stored.get(namespace) {
-                Some(segs) => {
-                    Ok(segs.get(&key).unwrap().1.clone())
-                }
+                Some(segs) => Ok(segs.get(&key).unwrap().1.clone()),
                 None => panic!("unknown namespace"),
             }
-
         } else {
             panic!("not storing")
         }
@@ -421,14 +431,12 @@ impl<IO: Io> Storage for TestStorage<IO> {
                 }
                 None => panic!("unknown namespace"),
             }
-
         } else {
             panic!("not storing")
         }
     }
 }
 
-#[derive(Debug)]
 pub struct StoreSegmentRequest<C, S> {
     namespace: NamespaceName,
     /// Path to the segment. Read-only for bottomless
@@ -439,4 +447,21 @@ pub struct StoreSegmentRequest<C, S> {
     /// alternative configuration to use with the storage layer.
     /// e.g: S3 overrides
     storage_config_override: Option<Arc<C>>,
+    /// Called after the segment was stored, with the new durable index
+    on_store_callback: OnStoreCallback,
+}
+
+impl<C, S> fmt::Debug for StoreSegmentRequest<C, S>
+where
+    C: fmt::Debug,
+    S: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("StoreSegmentRequest")
+            .field("namespace", &self.namespace)
+            .field("segment", &self.segment)
+            .field("created_at", &self.created_at)
+            .field("storage_config_override", &self.storage_config_override)
+            .finish()
+    }
 }
diff --git a/libsql-wal/src/storage/scheduler.rs b/libsql-wal/src/storage/scheduler.rs
