Move the utilities for dealing with item/error/done into utils.

Author: rklaehn

examples/expiring-tags.rs

@@ -73,7 +73,7 @@ async fn delete_expired_tags(blobs: &Store, prefix: &str, bulk: bool) -> anyhow:
         // find tags to delete one by one and then delete them
         //
         // this allows us to print the tags before deleting them
-        let mut tags = blobs.tags().list().await?;
+        let mut tags = blobs.tags().list().stream();
         let mut to_delete = Vec::new();
         while let Some(tag) = tags.next().await {
             let tag = tag?.name;
@@ -102,7 +102,7 @@ async fn delete_expired_tags(blobs: &Store, prefix: &str, bulk: bool) -> anyhow:
 
 async fn print_store_info(store: &Store) -> anyhow::Result<()> {
     let now = chrono::Utc::now();
-    let mut tags = store.tags().list().await?;
+    let mut tags = store.tags().list().stream();
     println!(
         "Current time: {}",
         now.to_rfc3339_opts(chrono::SecondsFormat::Secs, true)
@@ -112,7 +112,7 @@ async fn print_store_info(store: &Store) -> anyhow::Result<()> {
         let tag = tag?;
         println!("  {tag:?}");
     }
-    let mut blobs = store.list().stream().await?;
+    let mut blobs = store.list().stream();
     println!("Blobs:");
     while let Some(item) = blobs.next().await {
         println!("  {}", item?);

src/api/blobs.rs

@@ -59,7 +59,10 @@ use crate::{
     api::proto::{BatchRequest, ImportByteStreamUpdate, ListBlobsItem},
     provider::StreamContext,
     store::IROH_BLOCK_SIZE,
-    util::temp_tag::TempTag,
+    util::{
+        irpc::{IrpcReceiverFutExt, IrpcStreamItem},
+        temp_tag::TempTag,
+    },
     BlobFormat, Hash, HashAndFormat,
 };
 
@@ -847,43 +850,36 @@ impl BlobsListProgress {
         }
     }
 
-    pub async fn hashes(self) -> RequestResult<Vec<Hash>> {
-        let mut rx = self.inner.await?;
-        let mut hashes = Vec::new();
-        loop {
-            match rx.recv().await? {
-                Some(ListBlobsItem::Item(hash)) => hashes.push(hash),
-                Some(ListBlobsItem::Error(cause)) => return Err(cause.into()),
-                Some(ListBlobsItem::Done) => break,
-                None => return Err(super::Error::other("unexpected end of stream").into()),
-            }
+    pub async fn hashes(self) -> super::Result<Vec<Hash>> {
+        self.inner.try_collect().await
+    }
+
+    pub fn stream(self) -> impl Stream<Item = super::Result<Hash>> {
+        self.inner.into_stream()
+    }
+}
+
+impl IrpcStreamItem for ListBlobsItem {
+    type Error = super::Error;
+    type Item = Hash;
+
+    fn into_result_opt(self) -> Option<Result<Hash, super::Error>> {
+        match self {
+            Self::Item(hash) => Some(Ok(hash)),
+            Self::Error(e) => Some(Err(e.into())),
+            Self::Done => None,
         }
-        Ok(hashes)
     }
 
-    pub async fn stream(self) -> irpc::Result<impl Stream<Item = super::Result<Hash>>> {
-        let mut rx = self.inner.await?;
-        Ok(Gen::new(|co| async move {
-            loop {
-                match rx.recv().await {
-                    Ok(Some(ListBlobsItem::Item(hash))) => co.yield_(Ok(hash)).await,
-                    Ok(Some(ListBlobsItem::Error(cause))) => {
-                        co.yield_(Err(cause)).await;
-                        break;
-                    }
-                    Ok(Some(ListBlobsItem::Done)) => break,
-                    Ok(None) => {
-                        co.yield_(Err(super::Error::other("unexpected end of stream").into()))
-                            .await;
-                        break;
-                    }
-                    Err(cause) => {
-                        co.yield_(Err(cause.into())).await;
-                        break;
-                    }
-                }
-            }
-        }))
+    fn from_result(item: std::result::Result<Hash, super::Error>) -> Self {
+        match item {
+            Ok(hash) => Self::Item(hash),
+            Err(e) => Self::Error(e.into()),
+        }
+    }
+
+    fn done() -> Self {
+        Self::Done
     }
 }
 

src/api/proto.rs

@@ -15,6 +15,7 @@
 //! the much simpler memory store.
 use std::{
     fmt::{self, Debug},
+    future::{Future, IntoFuture},
     io,
     num::NonZeroU64,
     ops::{Bound, RangeBounds},
@@ -28,17 +29,25 @@ use bao_tree::{
     ChunkRanges,
 };
 use bytes::Bytes;
+use genawaiter::sync::Gen;
 use irpc::{
     channel::{mpsc, oneshot},
     rpc_requests,
 };
-use n0_future::Stream;
+use n0_future::{future, Stream};
 use range_collections::RangeSet2;
 use serde::{Deserialize, Serialize};
 pub(crate) mod bitfield;
 pub use bitfield::Bitfield;
 
-use crate::{store::util::Tag, util::temp_tag::TempTag, BlobFormat, Hash, HashAndFormat};
+use crate::{
+    store::util::Tag,
+    util::{
+        irpc::{IrpcReceiverFutExt, IrpcStreamItem},
+        temp_tag::TempTag,
+    },
+    BlobFormat, Hash, HashAndFormat,
+};
 
 pub(crate) trait HashSpecific {
     fn hash(&self) -> Hash;
@@ -113,7 +122,7 @@ pub enum Request {
     ImportPath(ImportPathRequest),
     #[rpc(tx = mpsc::Sender<ExportProgressItem>)]
     ExportPath(ExportPathRequest),
-    #[rpc(tx = oneshot::Sender<Vec<super::Result<TagInfo>>>)]
+    #[rpc(tx = mpsc::Sender<ListTagsItem>)]
     ListTags(ListTagsRequest),
     #[rpc(tx = oneshot::Sender<super::Result<()>>)]
     SetTag(SetTagRequest),
@@ -354,8 +363,110 @@ pub struct TagInfo {
 #[derive(Debug, Serialize, Deserialize)]
 pub enum ListBlobsItem {
     Item(Hash),
+    Error(super::Error),
     Done,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub enum ListTagsItem {
+    Item(TagInfo),
     Error(super::Error),
+    Done,
+}
+
+impl From<std::result::Result<TagInfo, super::Error>> for ListTagsItem {
+    fn from(item: std::result::Result<TagInfo, super::Error>) -> Self {
+        match item {
+            Ok(item) => ListTagsItem::Item(item),
+            Err(err) => ListTagsItem::Error(err),
+        }
+    }
+}
+
+impl IrpcStreamItem for ListTagsItem {
+    type Error = super::Error;
+    type Item = TagInfo;
+
+    fn into_result_opt(self) -> Option<Result<TagInfo, super::Error>> {
+        match self {
+            ListTagsItem::Item(item) => Some(Ok(item)),
+            ListTagsItem::Done => None,
+            ListTagsItem::Error(err) => Some(Err(err)),
+        }
+    }
+
+    fn from_result(item: std::result::Result<TagInfo, super::Error>) -> Self {
+        match item {
+            Ok(i) => Self::Item(i),
+            Err(e) => Self::Error(e.into()),
+        }
+    }
+
+    fn done() -> Self {
+        Self::Done
+    }
+}
+
+pub struct ListTagsProgress {
+    inner: future::Boxed<irpc::Result<mpsc::Receiver<ListTagsItem>>>,
+}
+
+impl IntoFuture for ListTagsProgress {
+    fn into_future(self) -> Self::IntoFuture {
+        Box::pin(self.inner.try_collect())
+    }
+
+    type IntoFuture = future::Boxed<Self::Output>;
+
+    type Output = super::Result<Vec<TagInfo>>;
+}
+
+impl ListTagsProgress {
+    pub(super) fn new(
+        fut: impl Future<Output = irpc::Result<mpsc::Receiver<ListTagsItem>>> + Send + 'static,
+    ) -> Self {
+        Self {
+            inner: Box::pin(fut),
+        }
+    }
+
+    pub fn stream(self) -> impl Stream<Item = super::Result<TagInfo>> {
+        Gen::new(|co| async move {
+            let mut rx = match self.inner.await {
+                Ok(rx) => rx,
+                Err(err) => {
+                    co.yield_(Err(super::Error::from(err))).await;
+                    return;
+                }
+            };
+            loop {
+                match rx.recv().await {
+                    Ok(Some(ListTagsItem::Item(item))) => {
+                        co.yield_(Ok(item)).await;
+                    }
+                    Ok(Some(ListTagsItem::Done)) => {
+                        break;
+                    }
+                    Ok(Some(ListTagsItem::Error(err))) => {
+                        co.yield_(Err(err.into())).await;
+                        break;
+                    }
+                    Ok(None) => {
+                        co.yield_(Err(super::Error::Io(io::Error::new(
+                            io::ErrorKind::UnexpectedEof,
+                            "stream ended",
+                        ))))
+                        .await;
+                        break;
+                    }
+                    Err(cause) => {
+                        co.yield_(Err(super::Error::from(cause))).await;
+                        break;
+                    }
+                }
+            }
+        })
+    }
 }
 
 impl From<TagInfo> for HashAndFormat {

src/api/tags.rs

@@ -16,7 +16,10 @@ use super::{
     proto::{CreateTempTagRequest, Scope},
     ApiClient, Tag, TempTag,
 };
-use crate::{api::proto::ListTempTagsRequest, HashAndFormat};
+use crate::{
+    api::proto::{ListTagsProgress, ListTempTagsRequest},
+    HashAndFormat,
+};
 
 /// The API for interacting with tags and temp tags.
 #[derive(Debug, Clone, ref_cast::RefCast)]
@@ -41,21 +44,15 @@ impl Tags {
     ///
     /// This is the most flexible way to list tags. All the other list methods are just convenience
     /// methods that call this one with the appropriate options.
-    pub async fn list_with_opts(
-        &self,
-        options: ListOptions,
-    ) -> irpc::Result<impl Stream<Item = super::Result<TagInfo>>> {
+    pub fn list_with_opts(&self, options: ListOptions) -> ListTagsProgress {
         trace!("{:?}", options);
-        let res = self.client.rpc(options).await?;
-        Ok(n0_future::stream::iter(res))
+        ListTagsProgress::new(self.client.server_streaming(options, 32))
     }
 
     /// Get the value of a single tag
     pub async fn get(&self, name: impl AsRef<[u8]>) -> super::RequestResult<Option<TagInfo>> {
-        let mut stream = self
-            .list_with_opts(ListOptions::single(name.as_ref()))
-            .await?;
-        Ok(stream.next().await.transpose()?)
+        let progress = self.list_with_opts(ListOptions::single(name.as_ref()));
+        Ok(progress.stream().next().await.transpose()?)
     }
 
     pub async fn set_with_opts(&self, options: SetOptions) -> super::RequestResult<()> {
@@ -77,34 +74,27 @@ impl Tags {
     }
 
     /// List a range of tags
-    pub async fn list_range<R, E>(
-        &self,
-        range: R,
-    ) -> irpc::Result<impl Stream<Item = super::Result<TagInfo>>>
+    pub fn list_range<R, E>(&self, range: R) -> ListTagsProgress
     where
         R: RangeBounds<E>,
         E: AsRef<[u8]>,
     {
-        self.list_with_opts(ListOptions::range(range)).await
+        self.list_with_opts(ListOptions::range(range))
     }
 
     /// Lists all tags with the given prefix.
-    pub async fn list_prefix(
-        &self,
-        prefix: impl AsRef<[u8]>,
-    ) -> irpc::Result<impl Stream<Item = super::Result<TagInfo>>> {
+    pub fn list_prefix(&self, prefix: impl AsRef<[u8]>) -> ListTagsProgress {
         self.list_with_opts(ListOptions::prefix(prefix.as_ref()))
-            .await
     }
 
     /// Lists all tags.
-    pub async fn list(&self) -> irpc::Result<impl Stream<Item = super::Result<TagInfo>>> {
-        self.list_with_opts(ListOptions::all()).await
+    pub fn list(&self) -> ListTagsProgress {
+        self.list_with_opts(ListOptions::all())
     }
 
     /// Lists all tags with a hash_seq format.
-    pub async fn list_hash_seq(&self) -> irpc::Result<impl Stream<Item = super::Result<TagInfo>>> {
-        self.list_with_opts(ListOptions::hash_seq()).await
+    pub fn list_hash_seq(&self) -> ListTagsProgress {
+        self.list_with_opts(ListOptions::hash_seq())
     }
 
     /// Deletes a tag.

src/store/fs/gc.rs

@@ -48,7 +48,7 @@ pub(super) async fn gc_mark_task(
     }
     let mut roots = HashSet::new();
     trace!("traversing tags");
-    let mut tags = store.tags().list().await?;
+    let mut tags = store.tags().list().stream();
     while let Some(tag) = tags.next().await {
         let info = tag?;
         trace!("adding root {:?} {:?}", info.name, info.hash_and_format());
@@ -85,7 +85,7 @@ async fn gc_sweep_task(
     live: &HashSet<Hash>,
     co: &Co<GcSweepEvent>,
 ) -> crate::api::Result<()> {
-    let mut blobs = store.blobs().list().stream().await?;
+    let mut blobs = store.blobs().list().stream();
     let mut count = 0;
     let mut batch = Vec::new();
     while let Some(hash) = blobs.next().await {

src/store/fs/meta.rs

@@ -24,7 +24,8 @@ use crate::{
         proto::{
             BlobDeleteRequest, BlobStatusMsg, BlobStatusRequest, ClearProtectedMsg,
             CreateTagRequest, DeleteBlobsMsg, DeleteTagsRequest, ListBlobsItem, ListBlobsMsg,
-            ListRequest, ListTagsRequest, RenameTagRequest, SetTagRequest, ShutdownMsg, SyncDbMsg,
+            ListRequest, ListTagsItem, ListTagsRequest, RenameTagRequest, SetTagRequest,
+            ShutdownMsg, SyncDbMsg,
         },
         tags::TagInfo,
     },
@@ -319,7 +320,20 @@ async fn handle_list_tags(msg: ListTagsMsg, tables: &impl ReadableTables) -> Act
             }
         }
     }
-    tx.send(res).await.ok();
+    for item in res {
+        match item {
+            Ok(tag_info) => {
+                if tx.send(ListTagsItem::Item(tag_info)).await.is_err() {
+                    return Ok(());
+                }
+            }
+            Err(err) => {
+                tx.send(ListTagsItem::Error(err)).await.ok();
+                return Ok(());
+            }
+        }
+    }
+    tx.send(ListTagsItem::Done).await.ok();
     Ok(())
 }