Give the ListTempTags rpc call the stream treatment as well!

Author: rklaehn

src/api/proto.rs

@@ -14,6 +14,7 @@
 //! The file system store is quite complex and optimized, so to get started take a look at
 //! the much simpler memory store.
 use std::{
+    collections::HashSet,
     fmt::{self, Debug},
     future::{Future, IntoFuture},
     io,
@@ -131,7 +132,7 @@ pub enum Request {
     RenameTag(RenameTagRequest),
     #[rpc(tx = oneshot::Sender<super::Result<Tag>>)]
     CreateTag(CreateTagRequest),
-    #[rpc(tx = oneshot::Sender<Vec<HashAndFormat>>)]
+    #[rpc(tx = mpsc::Sender<ListTempTagsItem>)]
     ListTempTags(ListTempTagsRequest),
     #[rpc(tx = oneshot::Sender<TempTag>)]
     CreateTempTag(CreateTempTagRequest),
@@ -406,6 +407,34 @@ impl IrpcStreamItem for ListTagsItem {
     }
 }
 
+pub struct ListTempTagsProgress {
+    inner: future::Boxed<irpc::Result<mpsc::Receiver<ListTempTagsItem>>>,
+}
+
+impl IntoFuture for ListTempTagsProgress {
+    fn into_future(self) -> Self::IntoFuture {
+        Box::pin(self.inner.try_collect())
+    }
+
+    type IntoFuture = future::Boxed<Self::Output>;
+
+    type Output = super::Result<HashSet<HashAndFormat>>;
+}
+
+impl ListTempTagsProgress {
+    pub(super) fn new(
+        fut: impl Future<Output = irpc::Result<mpsc::Receiver<ListTempTagsItem>>> + Send + 'static,
+    ) -> Self {
+        Self {
+            inner: Box::pin(fut),
+        }
+    }
+
+    pub fn stream(self) -> impl Stream<Item = super::Result<HashAndFormat>> {
+        self.inner.into_stream()
+    }
+}
+
 pub struct ListTagsProgress {
     inner: future::Boxed<irpc::Result<mpsc::Receiver<ListTagsItem>>>,
 }
@@ -493,6 +522,37 @@ pub struct CreateTempTagRequest {
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ListTempTagsRequest;
 
+#[derive(Debug, Serialize, Deserialize)]
+pub enum ListTempTagsItem {
+    Item(HashAndFormat),
+    Error(super::Error),
+    Done,
+}
+
+impl IrpcStreamItem for ListTempTagsItem {
+    type Error = super::Error;
+    type Item = HashAndFormat;
+
+    fn into_result_opt(self) -> Option<Result<HashAndFormat, super::Error>> {
+        match self {
+            ListTempTagsItem::Item(item) => Some(Ok(item)),
+            ListTempTagsItem::Done => None,
+            ListTempTagsItem::Error(err) => Some(Err(err)),
+        }
+    }
+
+    fn from_result(item: std::result::Result<HashAndFormat, super::Error>) -> Self {
+        match item {
+            Ok(i) => Self::Item(i),
+            Err(e) => Self::Error(e),
+        }
+    }
+
+    fn done() -> Self {
+        Self::Done
+    }
+}
+
 /// Rename a tag atomically
 #[derive(Debug, Serialize, Deserialize)]
 pub struct RenameTagRequest {

src/api/tags.rs

@@ -3,7 +3,7 @@
 //! The main entry point is the [`Tags`] struct.
 use std::ops::RangeBounds;
 
-use n0_future::{Stream, StreamExt};
+use n0_future::StreamExt;
 use ref_cast::RefCast;
 use tracing::trace;
 
@@ -17,7 +17,7 @@ use super::{
     ApiClient, Tag, TempTag,
 };
 use crate::{
-    api::proto::{ListTagsProgress, ListTempTagsRequest},
+    api::proto::{ListTagsProgress, ListTempTagsProgress, ListTempTagsRequest},
     HashAndFormat,
 };
 
@@ -33,11 +33,10 @@ impl Tags {
         Self::ref_cast(sender)
     }
 
-    pub async fn list_temp_tags(&self) -> irpc::Result<impl Stream<Item = HashAndFormat>> {
+    pub fn list_temp_tags(&self) -> ListTempTagsProgress {
         let options = ListTempTagsRequest;
         trace!("{:?}", options);
-        let res = self.client.rpc(options).await?;
-        Ok(n0_future::stream::iter(res))
+        ListTempTagsProgress::new(self.client.server_streaming(options, 32))
     }
 
     /// List all tags with options.

src/store/fs.rs

@@ -123,6 +123,7 @@ use crate::{
     },
     util::{
         channel::oneshot,
+        irpc::MpscSenderExt,
         temp_tag::{TagDrop, TempTag, TempTagScope, TempTags},
         ChunkRangesExt,
     },
@@ -520,7 +521,7 @@ impl Actor {
             Command::ListTempTags(cmd) => {
                 trace!("{cmd:?}");
                 let tts = self.temp_tags.list();
-                cmd.tx.send(tts).await.ok();
+                cmd.tx.forward_iter(tts.into_iter().map(Ok)).await.ok();
             }
             Command::ImportBytes(cmd) => {
                 trace!("{cmd:?}");
@@ -1417,13 +1418,13 @@ impl FsStore {
 #[cfg(test)]
 pub mod tests {
     use core::panic;
-    use std::collections::{HashMap, HashSet};
+    use std::{collections::HashMap, future::IntoFuture};
 
     use bao_tree::{
         io::{outboard::PreOrderMemOutboard, round_up_to_chunks_groups},
         ChunkRanges,
     };
-    use n0_future::{stream, Stream, StreamExt};
+    use n0_future::{stream, Stream};
     use testresult::TestResult;
     use walkdir::WalkDir;
 
@@ -1973,23 +1974,13 @@ pub mod tests {
         let batch = store.blobs().batch().await?;
         let tt1 = batch.temp_tag(Hash::new("foo")).await?;
         let tt2 = batch.add_slice("boo").await?;
-        let tts = store
-            .tags()
-            .list_temp_tags()
-            .await?
-            .collect::<HashSet<_>>()
-            .await;
+        let tts = store.tags().list_temp_tags().into_future().await?;
         assert!(tts.contains(tt1.hash_and_format()));
         assert!(tts.contains(tt2.hash_and_format()));
         drop(batch);
         store.sync_db().await?;
         store.wait_idle().await?;
-        let tts = store
-            .tags()
-            .list_temp_tags()
-            .await?
-            .collect::<HashSet<_>>()
-            .await;
+        let tts = store.tags().list_temp_tags().await?;
         // temp tag went out of scope, so it does not work anymore
         assert!(!tts.contains(tt1.hash_and_format()));
         assert!(!tts.contains(tt2.hash_and_format()));

src/store/fs/gc.rs

@@ -55,10 +55,10 @@ pub(super) async fn gc_mark_task(
         roots.insert(info.hash_and_format());
     }
     trace!("traversing temp roots");
-    let mut tts = store.tags().list_temp_tags().await?;
+    let mut tts = store.tags().list_temp_tags().stream();
     while let Some(tt) = tts.next().await {
         trace!("adding temp root {:?}", tt);
-        roots.insert(tt);
+        roots.insert(tt?);
     }
     for HashAndFormat { hash, format } in roots {
         // we need to do this for all formats except raw

src/store/mem.rs

@@ -324,7 +324,7 @@ impl Actor {
             Command::ListTempTags(cmd) => {
                 trace!("{cmd:?}");
                 let tts = self.temp_tags.list();
-                cmd.tx.send(tts).await.ok();
+                cmd.tx.forward_iter(tts.into_iter().map(Ok)).await.ok();
             }
             Command::ListBlobs(cmd) => {
                 let ListBlobsMsg { tx, .. } = cmd;

src/store/readonly_mem.rs

@@ -205,7 +205,7 @@ impl Actor {
                     .ok();
             }
             Command::ListTempTags(cmd) => {
-                cmd.tx.send(Vec::new()).await.ok();
+                cmd.tx.forward_iter(std::iter::empty()).await.ok();
             }
             Command::SyncDb(cmd) => {
                 cmd.tx.send(Ok(())).await.ok();