buck2_execute: implement OSS upload_blob to handle large stdout/`…

…stderr` Forward-port of patch 4 in <facebook#477> since this is a clear piece of missing functionality: in the event that stdout/stderr were more than 50KiB, then this dead path was taken, and so stdout/stderr would not be uploaded successfully. Co-authored-by: Austin Seipp <[email protected]>
thoughtpolice · Sep 3, 2024 · 00da7c1 · 00da7c1
1 parent 2fea973
commit 00da7c1
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 13 deletions.
diff --git a/app/buck2_execute/src/execute/output.rs b/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))
 }
diff --git a/app/buck2_execute/src/re/client.rs b/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
@@ -1131,17 +1131,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(

diff --git a/app/buck2_execute/src/re/manager.rs b/app/buck2_execute/src/re/manager.rs
@@ -455,7 +455,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

diff --git a/app/buck2_execute_impl/src/executors/caching.rs b/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")
         };

diff --git a/remote_execution/oss/re_grpc/src/client.rs b/remote_execution/oss/re_grpc/src/client.rs
@@ -796,11 +796,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(