Stream large blobs to remote cache directly from local cache file

src/rust/engine/Cargo.toml

@@ -244,7 +244,6 @@ lmdb-rkv = { git = "https://github.com/pantsbuild/lmdb-rs.git", rev = "6ae7a552a
 log = "0.4.17"
 madvise = "0.1"
 maplit = "1.0.1"
-memmap = "0.7"
 nails = "0.13"
 nix = "0.25"
 notify = { git = "https://github.com/pantsbuild/notify", rev = "276af0f3c5f300bfd534941386ba2f3b3a022df7" }

src/rust/engine/fs/store/Cargo.toml

@@ -25,7 +25,6 @@ itertools = { workspace = true }
 lmdb-rkv = { workspace = true }
 log = { workspace = true }
 madvise = { workspace = true }
-memmap = { workspace = true }
 parking_lot = { workspace = true }
 prost = { workspace = true }
 prost-types = { workspace = true }
@@ -36,6 +35,7 @@ task_executor = { path = "../../task_executor" }
 tempfile = { workspace = true }
 tokio-rustls = { workspace = true }
 tokio = { workspace = true, features = ["fs"] }
+tokio-util = { workspace = true, features = ["io"] }
 tonic = { workspace = true, features = ["transport", "codegen", "tls", "tls-roots", "prost"] }
 tower-service = { workspace = true }
 tryfuture = { path = "../../tryfuture" }

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

@@ -806,15 +806,16 @@ impl Store {
               remote_store
                 .clone()
                 .maybe_upload(digest, async move {
-                  // TODO(John Sirois): Consider allowing configuration of when to buffer large blobs
-                  // to disk to be independent of the remote store wire chunk size.
-                  if digest.size_bytes > remote_store.store.chunk_size_bytes() {
-                    Self::store_large_blob_remote(local, remote_store.store, entry_type, digest)
-                      .await
-                  } else {
-                    Self::store_small_blob_remote(local, remote_store.store, entry_type, digest)
-                      .await
-                  }
+                  match local.load_from_fs(digest).await? {
+                    Some(path) => {
+                      Self::store_fsdb_blob_remote(remote_store.store, digest, path).await?
+                    }
+                    None => {
+                      Self::store_lmdb_blob_remote(local, remote_store.store, entry_type, digest)
+                        .await?
+                    }
+                  };
+                  Ok(())
                 })
                 .await
             }
@@ -837,15 +838,17 @@ impl Store {
     .boxed()
   }
 
-  async fn store_small_blob_remote(
+  async fn store_lmdb_blob_remote(
     local: local::ByteStore,
     remote: remote::ByteStore,
     entry_type: EntryType,
     digest: Digest,
   ) -> Result<(), StoreError> {
     // We need to copy the bytes into memory so that they may be used safely in an async
-    // future. While this unfortunately increases memory consumption, we prioritize
-    // being able to run `remote.store_bytes()` as async.
+    // future. While this unfortunately increases memory consumption, we prioritize being
+    // able to run `remote.store_bytes()` as async. In addition, this is only used for
+    // blobs in the LMDB store, most of which are small: large blobs end up in the FSDB
+    // store.
     //
     // See https://github.com/pantsbuild/pants/pull/9793 for an earlier implementation
     // that used `Executor.block_on`, which avoided the clone but was blocking.
@@ -863,31 +866,16 @@ impl Store {
     }
   }
 
-  async fn store_large_blob_remote(
-    local: local::ByteStore,
+  async fn store_fsdb_blob_remote(
     remote: remote::ByteStore,
-    entry_type: EntryType,
     digest: Digest,
+    path: PathBuf,
   ) -> Result<(), StoreError> {
-    remote
-      .store_buffered(digest, |mut buffer| async {
-        let result = local
-          .load_bytes_with(entry_type, digest, move |bytes| {
-            buffer.write_all(bytes).map_err(|e| {
-              format!("Failed to write {entry_type:?} {digest:?} to temporary buffer: {e}")
-            })
-          })
-          .await?;
-        match result {
-          None => Err(StoreError::MissingDigest(
-            format!("Failed to upload {entry_type:?}: Not found in local store",),
-            digest,
-          )),
-          Some(Err(err)) => Err(err.into()),
-          Some(Ok(())) => Ok(()),
-        }
-      })
+    let file = tokio::fs::File::open(&path)
       .await
+      .map_err(|e| format!("failed to read {digest:?} from {path:?}: {e}"))?;
+    remote.store(digest, Box::new(file)).await?;
+    Ok(())
   }
 
   ///

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

@@ -2,7 +2,6 @@
 // Licensed under the Apache License, Version 2.0 (see LICENSE).
 use std::collections::{BTreeMap, HashSet};
 use std::fmt;
-use std::ops::Range;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
@@ -14,33 +13,37 @@ use hashing::Digest;
 use log::Level;
 use protos::gen::build::bazel::remote::execution::v2 as remexec;
 use remexec::ServerCapabilities;
-use tokio::io::{AsyncSeekExt, AsyncWrite};
+use tokio::io::{AsyncRead, AsyncSeekExt, AsyncWrite};
 use workunit_store::{in_workunit, ObservationMetric};
 
-use crate::StoreError;
-
 mod reapi;
 #[cfg(test)]
 mod reapi_tests;
 
-pub type ByteSource = Arc<(dyn Fn(Range<usize>) -> Bytes + Send + Sync + 'static)>;
+pub type StoreSource = Box<dyn AsyncRead + Send + Sync + Unpin + 'static>;
 
 #[async_trait]
 pub trait ByteStoreProvider: Sync + Send + 'static {
-  async fn store_bytes(&self, digest: Digest, bytes: ByteSource) -> Result<(), String>;
+  /// Store the bytes readable from `source` into the remote store
+  async fn store(&self, digest: Digest, source: StoreSource) -> Result<(), String>;
+
+  /// Store the bytes in `bytes` into the remote store, as an optimisation of `store` when the bytes
+  /// are already in memory
+  async fn store_bytes(&self, digest: Digest, bytes: Bytes) -> Result<(), String>;
 
+  /// Load the data stored (if any) in the remote store for `digest` into `destination`. Returns
+  /// true when found, false when not.
   async fn load(
     &self,
     digest: Digest,
     destination: &mut dyn LoadDestination,
   ) -> Result<bool, String>;
 
+  /// Return any digests from `digests` that are not (currently) available in the remote store.
   async fn list_missing_digests(
     &self,
     digests: &mut (dyn Iterator<Item = Digest> + Send),
   ) -> Result<HashSet<Digest>, String>;
-
-  fn chunk_size_bytes(&self) -> usize;
 }
 
 // TODO: Consider providing `impl Default`, similar to `super::LocalOptions`.
@@ -110,74 +113,40 @@ impl ByteStore {
     Ok(ByteStore::new(instance_name, provider))
   }
 
-  pub(crate) fn chunk_size_bytes(&self) -> usize {
-    self.provider.chunk_size_bytes()
-  }
-
-  pub async fn store_buffered<WriteToBuffer, WriteResult>(
-    &self,
-    digest: Digest,
-    write_to_buffer: WriteToBuffer,
-  ) -> Result<(), StoreError>
-  where
-    WriteToBuffer: FnOnce(std::fs::File) -> WriteResult,
-    WriteResult: Future<Output = Result<(), StoreError>>,
-  {
-    let write_buffer = tempfile::tempfile().map_err(|e| {
-      format!("Failed to create a temporary blob upload buffer for {digest:?}: {e}")
-    })?;
-    let read_buffer = write_buffer.try_clone().map_err(|e| {
-      format!("Failed to create a read handle for the temporary upload buffer for {digest:?}: {e}")
-    })?;
-    write_to_buffer(write_buffer).await?;
-
-    // Unsafety: Mmap presents an immutable slice of bytes, but the underlying file that is mapped
-    // could be mutated by another process. We guard against this by creating an anonymous
-    // temporary file and ensuring it is written to and closed via the only other handle to it in
-    // the code just above.
-    let mmap = Arc::new(unsafe {
-      let mapping = memmap::Mmap::map(&read_buffer).map_err(|e| {
-        format!("Failed to memory map the temporary file buffer for {digest:?}: {e}")
-      })?;
-      if let Err(err) = madvise::madvise(
-        mapping.as_ptr(),
-        mapping.len(),
-        madvise::AccessPattern::Sequential,
-      ) {
-        log::warn!(
-          "Failed to madvise(MADV_SEQUENTIAL) for the memory map of the temporary file buffer for \
-          {digest:?}. Continuing with possible reduced performance: {err}",
-          digest = digest,
-          err = err
-        )
-      }
-      Ok(mapping) as Result<memmap::Mmap, String>
-    }?);
-
+  /// Store the bytes readable from `source` into the remote store
+  pub async fn store(&self, digest: Digest, source: StoreSource) -> Result<(), String> {
     self
-      .store_bytes_source(
-        digest,
-        Arc::new(move |range| Bytes::copy_from_slice(&mmap[range])),
-      )
-      .await?;
-
-    Ok(())
+      .store_tracking("store", digest, || self.provider.store(digest, source))
+      .await
   }
 
+  /// Store the bytes in `bytes` into the remote store, as an optimisation of `store` when the bytes
+  /// are already in memory
   pub async fn store_bytes(&self, bytes: Bytes) -> Result<(), String> {
     let digest = Digest::of_bytes(&bytes);
     self
-      .store_bytes_source(digest, Arc::new(move |range| bytes.slice(range)))
+      .store_tracking("store_bytes", digest, || {
+        self.provider.store_bytes(digest, bytes)
+      })
       .await
   }
 
-  async fn store_bytes_source(&self, digest: Digest, bytes: ByteSource) -> Result<(), String> {
+  async fn store_tracking<DoStore, Fut>(
+    &self,
+    workunit: &'static str,
+    digest: Digest,
+    do_store: DoStore,
+  ) -> Result<(), String>
+  where
+    DoStore: FnOnce() -> Fut + Send,
+    Fut: Future<Output = Result<(), String>> + Send,
+  {
     in_workunit!(
-      "store_bytes",
+      workunit,
       Level::Trace,
       desc = Some(format!("Storing {digest:?}")),
       |workunit| async move {
-        let result = self.provider.store_bytes(digest, bytes).await;
+        let result = do_store().await;
 
         if result.is_ok() {
           workunit.record_observation(