buck2_execute: implement OSS upload_blob for local_only cache upl…

…oads

Forward-port of patch 4 in <#477>,
providing a clear piece of missing functionality: in the event that stdout
or stderr were more than 50KiB of output when caching `local_only` actions,
then this dead path was taken, and so stdout/stderr would not be uploaded
successfully in the cache.

Co-authored-by: Austin Seipp <aseipp@pobox.com>

app/buck2_execute/src/execute/output.rs

@@ -14,10 +14,12 @@ use anyhow::Context;
 use buck2_common::file_ops::FileDigest;
 use buck2_core::execution_types::executor_config::RemoteExecutorUseCase;
 use futures::future;
+use remote_execution::InlinedBlobWithDigest;
 use remote_execution::TDigest;
 
 use crate::digest::CasDigestConversionResultExt;
 use crate::digest::CasDigestFromReExt;
+use crate::digest::CasDigestToReExt;
 use crate::digest_config::DigestConfig;
 use crate::re::manager::ManagedRemoteExecutionClient;
 use crate::re::streams::RemoteCommandStdStreams;
@@ -238,12 +240,13 @@ impl CommandStdStreams {
  self,
  client: &ManagedRemoteExecutionClient,
  use_case: RemoteExecutorUseCase,
+ digest_config: DigestConfig,
  ) -> anyhow::Result<StdStreamPair<ReStdStream>> {
  match self {
  Self::Local { stdout, stderr } => {
  let (stdout, stderr) = future::try_join(
- maybe_upload_to_re(client, use_case, stdout),
- maybe_upload_to_re(client, use_case, stderr),
+ maybe_upload_to_re(client, use_case, stdout, digest_config),
+ maybe_upload_to_re(client, use_case, stderr, digest_config),
  )
  .await?;
 
@@ -276,11 +279,17 @@ async fn maybe_upload_to_re(
  client: &ManagedRemoteExecutionClient,
  use_case: RemoteExecutorUseCase,
  bytes: Vec<u8>,
+ digest_config: DigestConfig,
 ) -> anyhow::Result<ReStdStream> {
  const MIN_STREAM_UPLOAD_SIZE: usize = 50 * 1024; // Same as RE
  if bytes.len() < MIN_STREAM_UPLOAD_SIZE {
  return Ok(ReStdStream::Raw(bytes));
  }
- let digest = client.upload_blob(bytes, use_case).await?;
+ let inline_blob = InlinedBlobWithDigest {
+ digest: FileDigest::from_content(&bytes, digest_config.cas_digest_config()).to_re(),
+ blob: bytes,
+ ..Default::default()
+ };
+ let digest = client.upload_blob(inline_blob, use_case).await?;
  Ok(ReStdStream::Digest(digest))
 }

app/buck2_execute/src/re/client.rs

@@ -318,7 +318,7 @@ impl RemoteExecutionClient {
 
  pub async fn upload_blob(
  &self,
- blob: Vec<u8>,
+ blob: InlinedBlobWithDigest,
  use_case: RemoteExecutorUseCase,
  ) -> anyhow::Result<TDigest> {
  self.data
@@ -1151,17 +1151,19 @@ impl RemoteExecutionClientImpl {
 
  pub async fn upload_blob(
  &self,
- blob: Vec<u8>,
+ blob: InlinedBlobWithDigest,
  use_case: RemoteExecutorUseCase,
  ) -> anyhow::Result<TDigest> {
+ let digest = blob.digest.clone();
  with_error_handler(
  "upload_blob",
  self.get_session_id(),
  self.client()
  .upload_blob(blob, use_case.metadata(None))
  .await,
  )
- .await
+ .await?;
+ Ok(digest)
  }
 
  async fn materialize_files(

app/buck2_execute/src/re/manager.rs

@@ -460,7 +460,7 @@ impl ManagedRemoteExecutionClient {
 
  pub async fn upload_blob(
  &self,
- blob: Vec<u8>,
+ blob: InlinedBlobWithDigest,
  use_case: RemoteExecutorUseCase,
  ) -> anyhow::Result<TDigest> {
  self.lock()?.get().await?.upload_blob(blob, use_case).await

app/buck2_execute_impl/src/executors/caching.rs

@@ -416,7 +416,7 @@ impl CacheUploader {
  .report
  .std_streams
  .clone()
- .into_re(&self.re_client, self.re_use_case)
+ .into_re(&self.re_client, self.re_use_case, digest_config)
  .await
  .context("Error accessing std_streams")
  };

remote_execution/oss/re_grpc/src/client.rs

@@ -758,11 +758,21 @@ impl REClient {
 
  pub async fn upload_blob(
  &self,
- _blob: Vec<u8>,
- _metadata: RemoteExecutionMetadata,
- ) -> anyhow::Result<TDigest> {
- // TODO(aloiscochard)
- Err(anyhow::anyhow!("Not implemented (RE upload_blob)"))
+ blob: InlinedBlobWithDigest,
+ metadata: RemoteExecutionMetadata,
+ ) -> anyhow::Result<()> {
+ self.upload(
+ metadata,
+ UploadRequest {
+ inlined_blobs_with_digest: Some(vec![blob]),
+ files_with_digest: None,
+ directories: None,
+ upload_only_missing: false,
+ ..Default::default()
+ },
+ )
+ .await?;
+ Ok(())
  }
 
  pub async fn download(