refactor(tvix/castore/blobservice): use io::Result in trait r/7207

For all these calls, the caller has enough context about what it did, so
it should be fine to use io::Result here.

We pretty much only constructed crate::Error::StorageError before
anyways, so this conveys *more* information.

Change-Id: I5cabb3769c9c2314bab926d34dda748fda9d3ccc
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10328
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Tested-by: BuildkiteCI
Autosubmit: flokli <flokli@flokli.de>

5 files changed, 39 insertions, 40 deletions

diff --git a/tvix/castore/src/blobservice/grpc.rs b/tvix/castore/src/blobservice/grpc.rs
index 9cc4f340db..f90240d884 100644
--- a/tvix/castore/src/blobservice/grpc.rs
+++ b/tvix/castore/src/blobservice/grpc.rs
@@ -1,7 +1,6 @@
 use super::{naive_seeker::NaiveSeeker, BlobReader, BlobService, BlobWriter};
 use crate::{proto, B3Digest};
 use futures::sink::SinkExt;
-use futures::TryFutureExt;
 use std::{
     collections::VecDeque,
     io::{self},
@@ -39,7 +38,7 @@ impl GRPCBlobService {
 #[async_trait]
 impl BlobService for GRPCBlobService {
     #[instrument(skip(self, digest), fields(blob.digest=%digest))]
-    async fn has(&self, digest: &B3Digest) -> Result<bool, crate::Error> {
+    async fn has(&self, digest: &B3Digest) -> io::Result<bool> {
         let mut grpc_client = self.grpc_client.clone();
         let resp = grpc_client
             .stat(proto::StatBlobRequest {
@@ -51,16 +50,13 @@ impl BlobService for GRPCBlobService {
         match resp {
             Ok(_blob_meta) => Ok(true),
             Err(e) if e.code() == Code::NotFound => Ok(false),
-            Err(e) => Err(crate::Error::StorageError(e.to_string())),
+            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
         }
     }
 
     // On success, this returns a Ok(Some(io::Read)), which can be used to read
     // the contents of the Blob, identified by the digest.
-    async fn open_read(
-        &self,
-        digest: &B3Digest,
-    ) -> Result<Option<Box<dyn BlobReader>>, crate::Error> {
+    async fn open_read(&self, digest: &B3Digest) -> io::Result<Option<Box<dyn BlobReader>>> {
         // Get a stream of [proto::BlobChunk], or return an error if the blob
         // doesn't exist.
         let resp = self
@@ -90,7 +86,7 @@ impl BlobService for GRPCBlobService {
                 Ok(Some(Box::new(NaiveSeeker::new(data_reader))))
             }
             Err(e) if e.code() == Code::NotFound => Ok(None),
-            Err(e) => Err(crate::Error::StorageError(e.to_string())),
+            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
         }
     }
 
@@ -141,25 +137,20 @@ pub struct GRPCBlobWriter<W: tokio::io::AsyncWrite> {
 
 #[async_trait]
 impl<W: tokio::io::AsyncWrite + Send + Sync + Unpin + 'static> BlobWriter for GRPCBlobWriter<W> {
-    async fn close(&mut self) -> Result<B3Digest, crate::Error> {
+    async fn close(&mut self) -> io::Result<B3Digest> {
         if self.task_and_writer.is_none() {
             // if we're already closed, return the b3 digest, which must exist.
             // If it doesn't, we already closed and failed once, and didn't handle the error.
             match &self.digest {
                 Some(digest) => Ok(digest.clone()),
-                None => Err(crate::Error::StorageError(
-                    "previously closed with error".to_string(),
-                )),
+                None => Err(io::Error::new(io::ErrorKind::BrokenPipe, "already closed")),
             }
         } else {
             let (task, mut writer) = self.task_and_writer.take().unwrap();
 
             // invoke shutdown, so the inner writer closes its internal tx side of
             // the channel.
-            writer
-                .shutdown()
-                .map_err(|e| crate::Error::StorageError(e.to_string()))
-                .await?;
+            writer.shutdown().await?;
 
             // block on the RPC call to return.
             // This ensures all chunks are sent out, and have been received by the
@@ -168,15 +159,17 @@ impl<W: tokio::io::AsyncWrite + Send + Sync + Unpin + 'static> BlobWriter for GR
             match task.await? {
                 Ok(resp) => {
                     // return the digest from the response, and store it in self.digest for subsequent closes.
+                    let digest_len = resp.digest.len();
                     let digest: B3Digest = resp.digest.try_into().map_err(|_| {
-                        crate::Error::StorageError(
-                            "invalid root digest length in response".to_string(),
+                        io::Error::new(
+                            io::ErrorKind::Other,
+                            format!("invalid root digest length {} in response", digest_len),
                         )
                     })?;
                     self.digest = Some(digest.clone());
                     Ok(digest)
                 }
-                Err(e) => Err(crate::Error::StorageError(e.to_string())),
+                Err(e) => Err(io::Error::new(io::ErrorKind::Other, e.to_string())),
             }
         }
     }
diff --git a/tvix/castore/src/blobservice/memory.rs b/tvix/castore/src/blobservice/memory.rs
index 0954f0adb5..25eec334de 100644
--- a/tvix/castore/src/blobservice/memory.rs
+++ b/tvix/castore/src/blobservice/memory.rs
@@ -8,7 +8,7 @@ use tonic::async_trait;
 use tracing::instrument;
 
 use super::{BlobReader, BlobService, BlobWriter};
-use crate::{B3Digest, Error};
+use crate::B3Digest;
 
 #[derive(Clone, Default)]
 pub struct MemoryBlobService {
@@ -18,13 +18,13 @@ pub struct MemoryBlobService {
 #[async_trait]
 impl BlobService for MemoryBlobService {
     #[instrument(skip_all, ret, err, fields(blob.digest=%digest))]
-    async fn has(&self, digest: &B3Digest) -> Result<bool, Error> {
+    async fn has(&self, digest: &B3Digest) -> io::Result<bool> {
         let db = self.db.read().unwrap();
         Ok(db.contains_key(digest))
     }
 
     #[instrument(skip_all, err, fields(blob.digest=%digest))]
-    async fn open_read(&self, digest: &B3Digest) -> Result<Option<Box<dyn BlobReader>>, Error> {
+    async fn open_read(&self, digest: &B3Digest) -> io::Result<Option<Box<dyn BlobReader>>> {
         let db = self.db.read().unwrap();
 
         match db.get(digest).map(|x| Cursor::new(x.clone())) {
@@ -100,13 +100,11 @@ impl tokio::io::AsyncWrite for MemoryBlobWriter {
 
 #[async_trait]
 impl BlobWriter for MemoryBlobWriter {
-    async fn close(&mut self) -> Result<B3Digest, Error> {
+    async fn close(&mut self) -> io::Result<B3Digest> {
         if self.writers.is_none() {
             match &self.digest {
                 Some(digest) => Ok(digest.clone()),
-                None => Err(crate::Error::StorageError(
-                    "previously closed with error".to_string(),
-                )),
+                None => Err(io::Error::new(io::ErrorKind::BrokenPipe, "already closed")),
             }
         } else {
             let (buf, hasher) = self.writers.take().unwrap();
@@ -115,13 +113,17 @@ impl BlobWriter for MemoryBlobWriter {
             let digest: B3Digest = hasher.finalize().as_bytes().into();
 
             // Only insert if the blob doesn't already exist.
-            let db = self.db.read()?;
+            let db = self.db.read().map_err(|e| {
+                io::Error::new(io::ErrorKind::BrokenPipe, format!("RwLock poisoned: {}", e))
+            })?;
             if !db.contains_key(&digest) {
                 // drop the read lock, so we can open for writing.
                 drop(db);
 
                 // open the database for writing.
-                let mut db = self.db.write()?;
+                let mut db = self.db.write().map_err(|e| {
+                    io::Error::new(io::ErrorKind::BrokenPipe, format!("RwLock poisoned: {}", e))
+                })?;
 
                 // and put buf in there. This will move buf out.
                 db.insert(digest.clone(), buf);
diff --git a/tvix/castore/src/blobservice/mod.rs b/tvix/castore/src/blobservice/mod.rs
index 61d5583b99..faaf94f037 100644
--- a/tvix/castore/src/blobservice/mod.rs
+++ b/tvix/castore/src/blobservice/mod.rs
@@ -1,7 +1,7 @@
 use std::io;
 use tonic::async_trait;
 
-use crate::{B3Digest, Error};
+use crate::B3Digest;
 
 mod from_addr;
 mod grpc;
@@ -25,10 +25,10 @@ pub use self::sled::SledBlobService;
 #[async_trait]
 pub trait BlobService: Send + Sync {
     /// Check if the service has the blob, by its content hash.
-    async fn has(&self, digest: &B3Digest) -> Result<bool, Error>;
+    async fn has(&self, digest: &B3Digest) -> io::Result<bool>;
 
     /// Request a blob from the store, by its content hash.
-    async fn open_read(&self, digest: &B3Digest) -> Result<Option<Box<dyn BlobReader>>, Error>;
+    async fn open_read(&self, digest: &B3Digest) -> io::Result<Option<Box<dyn BlobReader>>>;
 
     /// Insert a new blob into the store. Returns a [BlobWriter], which
     /// implements [io::Write] and a [BlobWriter::close].
@@ -43,7 +43,7 @@ pub trait BlobWriter: tokio::io::AsyncWrite + Send + Sync + Unpin + 'static {
     /// contents written.
     ///
     /// Closing a already-closed BlobWriter is a no-op.
-    async fn close(&mut self) -> Result<B3Digest, Error>;
+    async fn close(&mut self) -> io::Result<B3Digest>;
 }
 
 /// A [tokio::io::AsyncRead] that also allows seeking.
diff --git a/tvix/castore/src/blobservice/sled.rs b/tvix/castore/src/blobservice/sled.rs
index f7bf33e8c5..3dd4bff7bc 100644
--- a/tvix/castore/src/blobservice/sled.rs
+++ b/tvix/castore/src/blobservice/sled.rs
@@ -34,19 +34,19 @@ impl SledBlobService {
 #[async_trait]
 impl BlobService for SledBlobService {
     #[instrument(skip(self), fields(blob.digest=%digest))]
-    async fn has(&self, digest: &B3Digest) -> Result<bool, Error> {
+    async fn has(&self, digest: &B3Digest) -> io::Result<bool> {
         match self.db.contains_key(digest.as_slice()) {
             Ok(has) => Ok(has),
-            Err(e) => Err(Error::StorageError(e.to_string())),
+            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e.to_string())),
         }
     }
 
     #[instrument(skip(self), fields(blob.digest=%digest))]
-    async fn open_read(&self, digest: &B3Digest) -> Result<Option<Box<dyn BlobReader>>, Error> {
+    async fn open_read(&self, digest: &B3Digest) -> io::Result<Option<Box<dyn BlobReader>>> {
         match self.db.get(digest.as_slice()) {
             Ok(None) => Ok(None),
             Ok(Some(data)) => Ok(Some(Box::new(Cursor::new(data[..].to_vec())))),
-            Err(e) => Err(Error::StorageError(e.to_string())),
+            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e.to_string())),
         }
     }
 
@@ -118,12 +118,13 @@ impl tokio::io::AsyncWrite for SledBlobWriter {
 
 #[async_trait]
 impl BlobWriter for SledBlobWriter {
-    async fn close(&mut self) -> Result<B3Digest, Error> {
+    async fn close(&mut self) -> io::Result<B3Digest> {
         if self.writers.is_none() {
             match &self.digest {
                 Some(digest) => Ok(digest.clone()),
-                None => Err(crate::Error::StorageError(
-                    "previously closed with error".to_string(),
+                None => Err(io::Error::new(
+                    io::ErrorKind::NotConnected,
+                    "already closed",
                 )),
             }
         } else {
diff --git a/tvix/castore/src/import.rs b/tvix/castore/src/import.rs
index 675dd0ec8b..a31bb22a62 100644
--- a/tvix/castore/src/import.rs
+++ b/tvix/castore/src/import.rs
@@ -117,7 +117,10 @@ async fn process_entry<'a>(
             return Err(Error::UnableToRead(entry.path().to_path_buf(), e));
         };
 
-        let digest = writer.close().await?;
+        let digest = writer
+            .close()
+            .await
+            .map_err(|e| Error::UnableToRead(entry.path().to_path_buf(), e))?;
 
         return Ok(Node::File(FileNode {
             name: entry.file_name().as_bytes().to_vec().into(),