Add wait_idle api call.

Author: rklaehn

src/api.rs

@@ -30,6 +30,7 @@ pub mod downloader;
 pub mod proto;
 pub mod remote;
 pub mod tags;
+use crate::api::proto::WaitIdleRequest;
 pub use crate::{store::util::Tag, util::temp_tag::TempTag};
 
 pub(crate) type ApiClient = irpc::Client<proto::Command, proto::Request, proto::StoreService>;
@@ -314,6 +315,7 @@ impl Store {
                     Request::ClearProtected(msg) => local.send((msg, tx)),
                     Request::SyncDb(msg) => local.send((msg, tx)),
                     Request::Shutdown(msg) => local.send((msg, tx)),
+                    Request::WaitIdle(msg) => local.send((msg, tx)),
                 }
             })
         });
@@ -332,6 +334,19 @@ impl Store {
         Ok(())
     }
 
+    /// Waits for the store to become completely idle.
+    ///
+    /// This is mostly useful for tests, where you want to check that e.g. the
+    /// store has written all data to disk.
+    ///
+    /// Note that a store is not guaranteed to become idle, if it is being
+    /// interacted with concurrently. So this might wait forever.
+    pub async fn wait_idle(&self) -> irpc::Result<()> {
+        let msg = WaitIdleRequest;
+        self.client.rpc(msg).await?;
+        Ok(())
+    }
+
     pub(crate) fn from_sender(client: ApiClient) -> Self {
         Self { client }
     }

src/api/proto.rs

@@ -134,11 +134,16 @@ pub enum Request {
     #[rpc(tx = oneshot::Sender<super::Result<()>>)]
     SyncDb(SyncDbRequest),
     #[rpc(tx = oneshot::Sender<()>)]
+    WaitIdle(WaitIdleRequest),
+    #[rpc(tx = oneshot::Sender<()>)]
     Shutdown(ShutdownRequest),
     #[rpc(tx = oneshot::Sender<super::Result<()>>)]
     ClearProtected(ClearProtectedRequest),
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct WaitIdleRequest;
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct SyncDbRequest;
 

src/store/fs.rs

@@ -243,6 +243,8 @@ struct Actor {
     handles: EntityManagerState<EmParams>,
     // temp tags
     temp_tags: TempTags,
+    // waiters for idle state.
+    idle_waiters: Vec<irpc::channel::oneshot::Sender<()>>,
     // our private tokio runtime. It has to live somewhere.
     _rt: RtWrapper,
 }
@@ -456,6 +458,16 @@ impl Actor {
                 trace!("{cmd:?}");
                 self.db().send(cmd.into()).await.ok();
             }
+            Command::WaitIdle(cmd) => {
+                trace!("{cmd:?}");
+                if self.tasks.is_empty() {
+                    // we are currently idle
+                    cmd.tx.send(()).await.ok();
+                } else {
+                    // wait for idle state
+                    self.idle_waiters.push(cmd.tx);
+                }
+            }
             Command::Shutdown(cmd) => {
                 trace!("{cmd:?}");
                 self.db().send(cmd.into()).await.ok();
@@ -599,6 +611,11 @@ impl Actor {
                 }
                 Some(res) = self.tasks.join_next(), if !self.tasks.is_empty() => {
                     Self::log_task_result(res);
+                    if self.tasks.is_empty() {
+                        for tx in self.idle_waiters.drain(..) {
+                            tx.send(()).await.ok();
+                        }
+                    }
                 }
             }
         }
@@ -648,6 +665,7 @@ impl Actor {
             tasks: JoinSet::new(),
             handles: EntityManagerState::new(slot_context, 1024, 32, 32, 2),
             temp_tags: Default::default(),
+            idle_waiters: Vec::new(),
             _rt: rt,
         })
     }
@@ -818,7 +836,6 @@ async fn handle_batch(cmd: BatchMsg, id: Scope, scope: Arc<TempTagScope>, ctx: A
     if let Err(cause) = handle_batch_impl(cmd, id, &scope).await {
         error!("batch failed: {cause}");
     }
-    println!("batch done, clearing scope {}", id);
     ctx.clear_scope(id).await;
 }
 
@@ -1969,6 +1986,7 @@ pub mod tests {
         println!("dropping batch");
         drop(batch);
         store.sync_db().await?;
+        store.wait_idle().await?;
         println!("reading temp tags after batch drop");
         let tts = store
             .tags()

src/store/mem.rs

@@ -51,7 +51,7 @@ use crate::{
             ImportByteStreamMsg, ImportByteStreamUpdate, ImportBytesMsg, ImportBytesRequest,
             ImportPathMsg, ImportPathRequest, ListBlobsMsg, ListTagsMsg, ListTagsRequest,
             ObserveMsg, ObserveRequest, RenameTagMsg, RenameTagRequest, Scope, SetTagMsg,
-            SetTagRequest, ShutdownMsg, SyncDbMsg,
+            SetTagRequest, ShutdownMsg, SyncDbMsg, WaitIdleMsg,
         },
         tags::TagInfo,
         ApiClient,
@@ -122,6 +122,7 @@ impl MemStore {
                 options: Arc::new(Options::default()),
                 temp_tags: Default::default(),
                 protected: Default::default(),
+                idle_waiters: Default::default(),
             }
             .run(),
         );
@@ -137,6 +138,8 @@ struct Actor {
     options: Arc<Options>,
     // temp tags
     temp_tags: TempTags,
+    // idle waiters
+    idle_waiters: Vec<irpc::channel::oneshot::Sender<()>>,
     protected: HashSet<Hash>,
 }
 
@@ -162,6 +165,16 @@ impl Actor {
                 let entry = self.get_or_create_entry(hash);
                 self.spawn(import_bao(entry, size, data, tx));
             }
+            Command::WaitIdle(WaitIdleMsg { tx, .. }) => {
+                trace!("wait idle");
+                if self.tasks.is_empty() {
+                    // we are currently idle
+                    tx.send(()).await.ok();
+                } else {
+                    // wait for idle state
+                    self.idle_waiters.push(tx);
+                }
+            }
             Command::Observe(ObserveMsg {
                 inner: ObserveRequest { hash },
                 tx,
@@ -485,6 +498,12 @@ impl Actor {
                         }
                         TaskResult::Unit(_) => {}
                     }
+                    if self.tasks.is_empty() {
+                        // we are idle now
+                        for tx in self.idle_waiters.drain(..) {
+                            tx.send(()).await.ok();
+                        }
+                    }
                 }
             }
         };

src/store/readonly_mem.rs

@@ -37,7 +37,7 @@ use crate::{
             self, BlobStatus, Command, ExportBaoMsg, ExportBaoRequest, ExportPathMsg,
             ExportPathRequest, ExportRangesItem, ExportRangesMsg, ExportRangesRequest,
             ImportBaoMsg, ImportByteStreamMsg, ImportBytesMsg, ImportPathMsg, ObserveMsg,
-            ObserveRequest,
+            ObserveRequest, WaitIdleMsg,
         },
         ApiClient, TempTag,
     },
@@ -62,6 +62,7 @@ impl Deref for ReadonlyMemStore {
 struct Actor {
     commands: tokio::sync::mpsc::Receiver<proto::Command>,
     tasks: JoinSet<()>,
+    idle_waiters: Vec<irpc::channel::oneshot::Sender<()>>,
     data: HashMap<Hash, CompleteStorage>,
 }
 
@@ -74,6 +75,7 @@ impl Actor {
             data,
             commands,
             tasks: JoinSet::new(),
+            idle_waiters: Vec::new(),
         }
     }
 
@@ -86,6 +88,15 @@ impl Actor {
                 .await
                 .ok();
             }
+            Command::WaitIdle(WaitIdleMsg { tx, .. }) => {
+                if self.tasks.is_empty() {
+                    // we are currently idle
+                    tx.send(()).await.ok();
+                } else {
+                    // wait for idle state
+                    self.idle_waiters.push(tx);
+                }
+            }
             Command::ImportBytes(ImportBytesMsg { tx, .. }) => {
                 tx.send(io::Error::other("import not supported").into())
                     .await
@@ -226,6 +237,12 @@ impl Actor {
                 },
                 Some(res) = self.tasks.join_next(), if !self.tasks.is_empty() => {
                     self.log_unit_task(res);
+                    if self.tasks.is_empty() {
+                        // we are idle now
+                        for tx in self.idle_waiters.drain(..) {
+                            tx.send(()).await.ok();
+                        }
+                    }
                 },
                 else => break,
             }

src/util.rs

@@ -472,6 +472,7 @@ pub mod sink {
         }
     }
 
+    #[allow(dead_code)]
     pub struct IrpcSenderSink<T>(pub irpc::channel::mpsc::Sender<T>);
 
     impl<T> Sink<T> for IrpcSenderSink<T>