Change list blobs rpc to have an explicit end case

Author: rklaehn

src/api/blobs.rs

@@ -56,7 +56,7 @@ use super::{
     ApiClient, RequestResult, Tags,
 };
 use crate::{
-    api::proto::{BatchRequest, ImportByteStreamUpdate},
+    api::proto::{BatchRequest, ImportByteStreamUpdate, ListBlobsItem},
     provider::StreamContext,
     store::IROH_BLOCK_SIZE,
     util::temp_tag::TempTag,
@@ -835,32 +835,53 @@ impl ImportBaoHandle {
 
 /// A progress handle for a blobs list operation.
 pub struct BlobsListProgress {
-    inner: future::Boxed<irpc::Result<mpsc::Receiver<super::Result<Hash>>>>,
+    inner: future::Boxed<irpc::Result<mpsc::Receiver<ListBlobsItem>>>,
 }
 
 impl BlobsListProgress {
     fn new(
-        fut: impl Future<Output = irpc::Result<mpsc::Receiver<super::Result<Hash>>>> + Send + 'static,
+        fut: impl Future<Output = irpc::Result<mpsc::Receiver<ListBlobsItem>>> + Send + 'static,
     ) -> Self {
         Self {
             inner: Box::pin(fut),
         }
     }
 
     pub async fn hashes(self) -> RequestResult<Vec<Hash>> {
-        let mut rx: mpsc::Receiver<Result<Hash, super::Error>> = self.inner.await?;
+        let mut rx = self.inner.await?;
         let mut hashes = Vec::new();
-        while let Some(item) = rx.recv().await? {
-            hashes.push(item?);
+        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()),
+            }
         }
         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 {
-            while let Ok(Some(item)) = rx.recv().await {
-                co.yield_(item).await;
+            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;
+                    }
+                }
             }
         }))
     }

src/api/proto.rs

@@ -89,7 +89,7 @@ impl HashSpecific for CreateTagMsg {
 #[rpc_requests(message = Command, alias = "Msg")]
 #[derive(Debug, Serialize, Deserialize)]
 pub enum Request {
-    #[rpc(tx = mpsc::Sender<super::Result<Hash>>)]
+    #[rpc(tx = mpsc::Sender<ListBlobsItem>)]
     ListBlobs(ListRequest),
     #[rpc(tx = oneshot::Sender<Scope>, rx = mpsc::Receiver<BatchResponse>)]
     Batch(BatchRequest),
@@ -351,6 +351,13 @@ pub struct TagInfo {
     pub hash: Hash,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub enum ListBlobsItem {
+    Item(Hash),
+    Done,
+    Error(super::Error),
+}
+
 impl From<TagInfo> for HashAndFormat {
     fn from(tag_info: TagInfo) -> Self {
         HashAndFormat {

src/store/fs/meta.rs

@@ -23,8 +23,8 @@ use crate::{
         blobs::BlobStatus,
         proto::{
             BlobDeleteRequest, BlobStatusMsg, BlobStatusRequest, ClearProtectedMsg,
-            CreateTagRequest, DeleteBlobsMsg, DeleteTagsRequest, ListBlobsMsg, ListRequest,
-            ListTagsRequest, RenameTagRequest, SetTagRequest, ShutdownMsg, SyncDbMsg,
+            CreateTagRequest, DeleteBlobsMsg, DeleteTagsRequest, ListBlobsItem, ListBlobsMsg,
+            ListRequest, ListTagsRequest, RenameTagRequest, SetTagRequest, ShutdownMsg, SyncDbMsg,
         },
         tags::TagInfo,
     },
@@ -887,20 +887,21 @@ pub async fn list_blobs(snapshot: ReadOnlyTables, cmd: ListBlobsMsg) {
         Ok(()) => {}
         Err(e) => {
             error!("error listing blobs: {}", e);
-            tx.send(Err(e)).await.ok();
+            tx.send(ListBlobsItem::Error(e)).await.ok();
         }
     }
 }
 
 async fn list_blobs_impl(
     snapshot: ReadOnlyTables,
     _cmd: ListRequest,
-    tx: &mut mpsc::Sender<api::Result<Hash>>,
+    tx: &mut mpsc::Sender<ListBlobsItem>,
 ) -> api::Result<()> {
     for item in snapshot.blobs.iter().map_err(api::Error::other)? {
         let (k, _) = item.map_err(api::Error::other)?;
         let k = k.value();
-        tx.send(Ok(k)).await.ok();
+        tx.send(ListBlobsItem::Item(k)).await.ok();
     }
+    tx.send(ListBlobsItem::Done).await.ok();
     Ok(())
 }

src/store/mem.rs

@@ -330,10 +330,15 @@ impl Actor {
                 let blobs = self.state.data.keys().cloned().collect::<Vec<Hash>>();
                 self.spawn(async move {
                     for blob in blobs {
-                        if tx.send(Ok(blob)).await.is_err() {
+                        if tx
+                            .send(api::proto::ListBlobsItem::Item(blob))
+                            .await
+                            .is_err()
+                        {
                             break;
                         }
                     }
+                    tx.send(api::proto::ListBlobsItem::Done).await.ok();
                 });
             }
             Command::BlobStatus(cmd) => {

src/store/readonly_mem.rs

@@ -36,8 +36,8 @@ use crate::{
         proto::{
             self, BlobStatus, Command, ExportBaoMsg, ExportBaoRequest, ExportPathMsg,
             ExportPathRequest, ExportRangesItem, ExportRangesMsg, ExportRangesRequest,
-            ImportBaoMsg, ImportByteStreamMsg, ImportBytesMsg, ImportPathMsg, ObserveMsg,
-            ObserveRequest, WaitIdleMsg,
+            ImportBaoMsg, ImportByteStreamMsg, ImportBytesMsg, ImportPathMsg, ListBlobsItem,
+            ObserveMsg, ObserveRequest, WaitIdleMsg,
         },
         ApiClient, TempTag,
     },
@@ -178,8 +178,9 @@ impl Actor {
                 let hashes: Vec<Hash> = self.data.keys().cloned().collect();
                 self.tasks.spawn(async move {
                     for hash in hashes {
-                        cmd.tx.send(Ok(hash)).await.ok();
+                        cmd.tx.send(ListBlobsItem::Item(hash)).await.ok();
                     }
+                    cmd.tx.send(ListBlobsItem::Done).await.ok();
                 });
             }
             Command::BlobStatus(cmd) => {