add support for upload speed / remaining in the cache upload step

vercel · May 22, 2024 · 6213948 · 6213948
1 parent d52a6d2
commit 6213948
Show file tree

Hide file tree

Showing 17 changed files with 418 additions and 60 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/turborepo-api-client/Cargo.toml b/crates/turborepo-api-client/Cargo.toml
@@ -21,15 +21,18 @@ workspace = true
 
 [dependencies]
 anyhow = { workspace = true }
+bytes.workspace = true
 chrono = { workspace = true, features = ["serde"] }
 lazy_static = { workspace = true }
 regex = { workspace = true }
-reqwest = { workspace = true, features = ["json"] }
+reqwest = { workspace = true, features = ["json", "stream"] }
 rustc_version_runtime = "0.2.1"
 serde = { workspace = true }
 serde_json = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
+tokio-stream = "0.1.15"
+tokio-util = { version = "0.7.10", features = ["codec"] }
 tracing = { workspace = true }
 turbopath = { workspace = true }
 turborepo-ci = { workspace = true }

diff --git a/crates/turborepo-api-client/src/error.rs b/crates/turborepo-api-client/src/error.rs
@@ -7,6 +7,8 @@ use crate::CachingStatus;
 
 #[derive(Debug, Error)]
 pub enum Error {
+    #[error("Error reading from disk: {0}")]
+    ReadError(#[from] std::io::Error),
     #[error("Error making HTTP request: {0}")]
     ReqwestError(#[from] reqwest::Error),
     #[error("skipping HTTP Request, too many failures have occurred.\nLast error: {0}")]

diff --git a/crates/turborepo-api-client/src/lib.rs b/crates/turborepo-api-client/src/lib.rs
@@ -8,7 +8,7 @@ use std::{backtrace::Backtrace, env, future::Future, time::Duration};
 use lazy_static::lazy_static;
 use regex::Regex;
 pub use reqwest::Response;
-use reqwest::{Method, RequestBuilder, StatusCode};
+use reqwest::{Body, Method, RequestBuilder, StatusCode};
 use serde::Deserialize;
 use turborepo_ci::{is_ci, Vendor};
 use turborepo_vercel_api::{
@@ -26,6 +26,9 @@ mod retry;
 pub mod spaces;
 pub mod telemetry;
 
+pub use bytes::Bytes;
+pub use tokio_stream::Stream;
+
 lazy_static! {
     static ref AUTHORIZATION_REGEX: Regex =
         Regex::new(r"(?i)(?:^|,) *authorization *(?:,|$)").unwrap();
@@ -74,7 +77,7 @@ pub trait CacheClient {
     fn put_artifact(
         &self,
         hash: &str,
-        artifact_body: &[u8],
+        artifact_body: impl tokio_stream::Stream<Item = Result<bytes::Bytes>> + Send + Sync + 'static,
         duration: u64,
         tag: Option<&str>,
         token: &str,
@@ -358,7 +361,7 @@ impl CacheClient for APIClient {
     async fn put_artifact(
         &self,
         hash: &str,
-        artifact_body: &[u8],
+        artifact_body: impl tokio_stream::Stream<Item = Result<bytes::Bytes>> + Send + Sync + 'static,
         duration: u64,
         tag: Option<&str>,
         token: &str,
@@ -382,13 +385,15 @@ impl CacheClient for APIClient {
             request_url = preflight_response.location.clone();
         }
 
+        let stream = Body::wrap_stream(artifact_body);
+
         let mut request_builder = self
             .cache_client
             .put(request_url)
             .header("Content-Type", "application/octet-stream")
             .header("x-artifact-duration", duration.to_string())
             .header("User-Agent", self.user_agent.clone())
-            .body(artifact_body.to_vec());
+            .body(stream);
 
         if allow_auth {
             request_builder = request_builder.header("Authorization", format!("Bearer {}", token));

diff --git a/crates/turborepo-auth/src/auth/login.rs b/crates/turborepo-auth/src/auth/login.rs
@@ -313,7 +313,11 @@ mod tests {
         async fn put_artifact(
             &self,
             _hash: &str,
-            _artifact_body: &[u8],
+            _artifact_body: impl turborepo_api_client::Stream<
+                    Item = Result<turborepo_api_client::Bytes, turborepo_api_client::Error>,
+                > + Send
+                + Sync
+                + 'static,
             _duration: u64,
             _tag: Option<&str>,
             _token: &str,

diff --git a/crates/turborepo-auth/src/auth/sso.rs b/crates/turborepo-auth/src/auth/sso.rs
@@ -310,7 +310,11 @@ mod tests {
         async fn put_artifact(
             &self,
             _hash: &str,
-            _artifact_body: &[u8],
+            _artifact_body: impl turborepo_api_client::Stream<
+                    Item = Result<turborepo_api_client::Bytes, turborepo_api_client::Error>,
+                > + Send
+                + Sync
+                + 'static,
             _duration: u64,
             _tag: Option<&str>,
             _token: &str,

diff --git a/crates/turborepo-auth/src/lib.rs b/crates/turborepo-auth/src/lib.rs
@@ -420,7 +420,11 @@ mod tests {
         async fn put_artifact(
             &self,
             _hash: &str,
-            _artifact_body: &[u8],
+            _artifact_body: impl turborepo_api_client::Stream<
+                    Item = Result<turborepo_api_client::Bytes, turborepo_api_client::Error>,
+                > + Send
+                + Sync
+                + 'static,
             _duration: u64,
             _tag: Option<&str>,
             _token: &str,