Simplify remote store load_bytes API

src/rust/engine/fs/store/src/lib.rs

@@ -58,8 +58,6 @@ use fs::{
 };
 use futures::future::{self, BoxFuture, Either, FutureExt, TryFutureExt};
 use grpc_util::prost::MessageExt;
-use grpc_util::retry::{retry_call, status_is_retryable};
-use grpc_util::status_to_str;
 use hashing::Digest;
 use parking_lot::Mutex;
 use prost::Message;
@@ -71,8 +69,6 @@ use sharded_lmdb::DEFAULT_LEASE_TIME;
 use tryfuture::try_future;
 use workunit_store::{in_workunit, Level, Metric};
 
-use crate::remote::ByteStoreError;
-
 const KILOBYTES: usize = 1024;
 const MEGABYTES: usize = 1024 * KILOBYTES;
 const GIGABYTES: usize = 1024 * MEGABYTES;
@@ -254,23 +250,7 @@ impl RemoteStore {
     let remote_store = self.store.clone();
     self
       .maybe_download(digest, async move {
-        // TODO(#17065): Now that we always copy from the remote store to the local store before
-        // executing the caller's logic against the local store,
-        // `remote::ByteStore::load_bytes_with` no longer needs to accept a function.
-        let bytes = retry_call(
-          remote_store,
-          |remote_store| async move { remote_store.load_bytes_with(digest, Ok).await },
-          |err| match err {
-            ByteStoreError::Grpc(status) => status_is_retryable(status),
-            _ => false,
-          },
-        )
-        .await
-        .map_err(|err| match err {
-          ByteStoreError::Grpc(status) => status_to_str(status),
-          ByteStoreError::Other(msg) => msg,
-        })?
-        .ok_or_else(|| {
+        let bytes = remote_store.load_bytes(digest).await?.ok_or_else(|| {
           StoreError::MissingDigest(
             "Was not present in either the local or remote store".to_owned(),
             digest,
@@ -1076,35 +1056,14 @@ impl Store {
       return Err("Cannot load Trees from a remote without a remote".to_owned());
     };
 
-    let tree_opt = retry_call(
-      remote,
-      |remote| async move {
-        remote
-          .store
-          .load_bytes_with(tree_digest, |b| {
-            let tree = Tree::decode(b).map_err(|e| format!("protobuf decode error: {:?}", e))?;
-            Ok(tree)
-          })
-          .await
-      },
-      |err| match err {
-        ByteStoreError::Grpc(status) => status_is_retryable(status),
-        _ => false,
-      },
-    )
-    .await
-    .map_err(|err| match err {
-      ByteStoreError::Grpc(status) => status_to_str(status),
-      ByteStoreError::Other(msg) => msg,
-    })?;
-
-    let tree = match tree_opt {
-      Some(t) => t,
-      None => return Ok(None),
-    };
-
-    let trie = DigestTrie::try_from(tree)?;
-    Ok(Some(trie.into()))
+    match remote.store.load_bytes(tree_digest).await? {
+      Some(b) => {
+        let tree = Tree::decode(b).map_err(|e| format!("protobuf decode error: {:?}", e))?;
+        let trie = DigestTrie::try_from(tree)?;
+        Ok(Some(trie.into()))
+      }
+      None => Ok(None),
+    }
   }
 
   pub async fn lease_all_recursively<'a, Ds: Iterator<Item = &'a Digest>>(

src/rust/engine/fs/store/src/remote.rs

@@ -49,7 +49,7 @@ impl fmt::Debug for ByteStore {
 
 /// Represents an error from accessing a remote bytestore.
 #[derive(Debug)]
-pub enum ByteStoreError {
+enum ByteStoreError {
   /// gRPC error
   Grpc(Status),
 
@@ -337,14 +337,7 @@ impl ByteStore {
     .await
   }
 
-  pub async fn load_bytes_with<
-    T: Send + 'static,
-    F: Fn(Bytes) -> Result<T, String> + Send + Sync + Clone + 'static,
-  >(
-    &self,
-    digest: Digest,
-    f: F,
-  ) -> Result<Option<T>, ByteStoreError> {
+  pub async fn load_bytes(&self, digest: Digest) -> Result<Option<Bytes>, String> {
     let start = Instant::now();
     let store = self.clone();
     let instance_name = store.instance_name.clone().unwrap_or_default();
@@ -356,80 +349,71 @@ impl ByteStore {
       digest.size_bytes
     );
     let workunit_desc = format!("Loading bytes at: {resource_name}");
-    let f = f.clone();
-
-    let mut client = self.byte_stream_client.as_ref().clone();
-
-    let result_future = async move {
-      let mut start_opt = Some(Instant::now());
 
-      let stream_result = client
-        .read({
-          protos::gen::google::bytestream::ReadRequest {
-            resource_name,
-            read_offset: 0,
-            // 0 means no limit.
-            read_limit: 0,
+    let request = protos::gen::google::bytestream::ReadRequest {
+      resource_name,
+      read_offset: 0,
+      // 0 means no limit.
+      read_limit: 0,
+    };
+    let client = self.byte_stream_client.as_ref().clone();
+
+    let result_future = retry_call(
+      (client, request),
+      move |(mut client, request)| async move {
+        let mut start_opt = Some(Instant::now());
+        let stream_result = client.read(request).await;
+
+        let mut stream = match stream_result {
+          Ok(response) => response.into_inner(),
+          Err(status) => {
+            return match status.code() {
+              Code::NotFound => Ok(None),
+              _ => Err(status),
+            }
           }
-        })
-        .await;
+        };
 
-      let mut stream = match stream_result {
-        Ok(response) => response.into_inner(),
-        Err(status) => {
-          return match status.code() {
-            Code::NotFound => Ok(None),
-            _ => Err(ByteStoreError::Grpc(status)),
-          }
-        }
-      };
-
-      let read_result_closure = async {
-        let mut buf = BytesMut::with_capacity(digest.size_bytes);
-        while let Some(response) = stream.next().await {
-          // Record the observed time to receive the first response for this read.
-          if let Some(start) = start_opt.take() {
-            if let Some(workunit_store_handle) = workunit_store::get_workunit_store_handle() {
-              let timing: Result<u64, _> =
-                Instant::now().duration_since(start).as_micros().try_into();
-              if let Ok(obs) = timing {
-                workunit_store_handle
-                  .store
-                  .record_observation(ObservationMetric::RemoteStoreTimeToFirstByteMicros, obs);
+        let read_result_closure = async {
+          let mut buf = BytesMut::with_capacity(digest.size_bytes);
+          while let Some(response) = stream.next().await {
+            // Record the observed time to receive the first response for this read.
+            if let Some(start) = start_opt.take() {
+              if let Some(workunit_store_handle) = workunit_store::get_workunit_store_handle() {
+                let timing: Result<u64, _> =
+                  Instant::now().duration_since(start).as_micros().try_into();
+                if let Ok(obs) = timing {
+                  workunit_store_handle
+                    .store
+                    .record_observation(ObservationMetric::RemoteStoreTimeToFirstByteMicros, obs);
+                }
               }
             }
-          }
 
-          buf.extend_from_slice(&(response?).data);
-        }
-        Ok(buf.freeze())
-      };
-
-      let read_result: Result<Bytes, tonic::Status> = read_result_closure.await;
-
-      let maybe_bytes = match read_result {
-        Ok(bytes) => Some(bytes),
-        Err(status) => {
-          if status.code() == tonic::Code::NotFound {
-            None
-          } else {
-            return Err(ByteStoreError::Grpc(status));
+            buf.extend_from_slice(&(response?).data);
           }
-        }
-      };
+          Ok(buf.freeze())
+        };
 
-      match maybe_bytes {
-        Some(b) => f(b).map(Some).map_err(ByteStoreError::Other),
-        None => Ok(None),
-      }
-    };
+        let read_result: Result<Bytes, tonic::Status> = read_result_closure.await;
+
+        match read_result {
+          Ok(bytes) => Ok(Some(bytes)),
+          Err(status) => match status.code() {
+            Code::NotFound => Ok(None),
+            _ => Err(status),
+          },
+        }
+      },
+      status_is_retryable,
+    );
 
     in_workunit!(
-      "load_bytes_with",
+      "load_bytes",
       Level::Trace,
       desc = Some(workunit_desc),
       |workunit| async move {
-        let result = result_future.await;
+        let result = result_future.await.map_err(status_to_str);
         workunit.record_observation(
           ObservationMetric::RemoteStoreReadBlobTimeMicros,
           start.elapsed().as_micros() as u64,

src/rust/engine/fs/store/src/remote_tests.rs

@@ -336,8 +336,5 @@ pub async fn load_directory_proto_bytes(
 }
 
 async fn load_bytes(store: &ByteStore, digest: Digest) -> Result<Option<Bytes>, String> {
-  store
-    .load_bytes_with(digest, |b| Ok(b))
-    .await
-    .map_err(|err| format!("{}", err))
+  store.load_bytes(digest).await
 }