Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Only upload large files once. #750

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Only upload large files once. #750

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6441ebc
Only upload large files once.
Feb 29, 2024
5824439
Import `dashmap::DashMap`.
Aug 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions remote_execution/oss/re_grpc/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ version = "0.1.0"
[dependencies]
anyhow = { workspace = true }
dupe = { workspace = true }
dashmap = { workspace = true }
futures = { workspace = true }
gazebo = { workspace = true }
http = { workspace = true }
Expand Down
143 changes: 106 additions & 37 deletions remote_execution/oss/re_grpc/src/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ use std::time::Instant;
use anyhow::Context;
use buck2_re_configuration::Buck2OssReConfiguration;
use buck2_re_configuration::HttpHeader;
use dashmap::DashMap;
use dupe::Dupe;
use futures::future::BoxFuture;
use futures::future::Future;
Expand Down Expand Up @@ -478,13 +479,21 @@ impl FindMissingCache {
}
}

#[derive(Clone)]
enum OngoingUploadStatus {
Active(tokio::sync::watch::Receiver<Result<(), ()>>),
Done,
Error,
}

pub struct REClient {
runtime_opts: RERuntimeOpts,
grpc_clients: GRPCClients,
capabilities: RECapabilities,
instance_name: InstanceName,
// buck2 calls find_missing for same blobs
find_missing_cache: Mutex<FindMissingCache>,
prev_uploads: DashMap<TDigest, OngoingUploadStatus>,
}

impl Drop for REClient {
Expand Down Expand Up @@ -559,6 +568,7 @@ impl REClient {
ttl: Duration::from_secs(12 * 60 * 60), // 12 hours TODO: Tune this parameter
last_check: Instant::now(),
}),
prev_uploads: DashMap::new(),
}
}

Expand Down Expand Up @@ -726,6 +736,7 @@ impl REClient {
request,
self.capabilities.max_total_batch_size,
self.runtime_opts.max_concurrent_uploads_per_action,
&self.prev_uploads,
|re_request| async {
let metadata = metadata.clone();
let mut cas_client = self.grpc_clients.cas_client.clone();
Expand Down Expand Up @@ -1186,6 +1197,7 @@ async fn upload_impl<Byt, Cas>(
request: UploadRequest,
max_total_batch_size: usize,
max_concurrent_uploads: Option<usize>,
prev_uploads: &DashMap<TDigest, OngoingUploadStatus>,
cas_f: impl Fn(BatchUpdateBlobsRequest) -> Cas + Sync + Send + Copy,
bystream_fut: impl Fn(Vec<WriteRequest>) -> Byt + Sync + Send + Copy,
) -> anyhow::Result<UploadResponse>
Expand Down Expand Up @@ -1246,10 +1258,9 @@ where

// Create futures for any files that needs uploading.
for file in request.files_with_digest.unwrap_or_default() {
let hash = file.digest.hash.clone();
let size = file.digest.size_in_bytes;
let digest = file.digest.clone();
let name = file.name.clone();
if size < max_total_batch_size as i64 {
if digest.size_in_bytes < max_total_batch_size as i64 {
batched_blob_updates.push(BatchUploadRequest::File(file));
continue;
}
Expand All @@ -1258,45 +1269,96 @@ where
"{}uploads/{}/blobs/{}/{}",
instance_name.as_resource_prefix(),
client_uuid,
hash.clone(),
size
file.digest.hash,
file.digest.size_in_bytes
);

enum UploadStatus {
New(tokio::sync::watch::Sender<Result<(), ()>>),
Ongoing(OngoingUploadStatus),
}

let upload_status = match prev_uploads.entry(digest.clone()) {
dashmap::mapref::entry::Entry::Occupied(o) => UploadStatus::Ongoing(o.get().clone()),
dashmap::mapref::entry::Entry::Vacant(v) => {
let (tx, rx) = tokio::sync::watch::channel(Err(()));
v.insert(OngoingUploadStatus::Active(rx));
UploadStatus::New(tx)
}
};
let fut = async move {
let mut file = tokio::fs::File::open(&name)
.await
.with_context(|| format!("Opening `{name}` for reading failed"))?;
let mut data = vec![0; max_total_batch_size];

let mut write_offset = 0;
let mut upload_segments = Vec::new();
loop {
let length = file
.read(&mut data)
.await
.with_context(|| format!("Error reading from {name}"))?;
if length == 0 {
break;
match upload_status {
UploadStatus::Ongoing(OngoingUploadStatus::Active(mut rx)) => {
// Another task was already uploading this artifact, wait for it complete and report result.
rx.changed().await?;
rx.borrow_and_update().as_ref().map_err(|_e| {
anyhow::anyhow!("Upload queued for previous action failed.")
})?;
}
upload_segments.push(WriteRequest {
resource_name: resource_name.to_owned(),
write_offset,
finish_write: false,
data: data[..length].to_owned(),
});
write_offset += length as i64;
}
upload_segments
.last_mut()
.with_context(|| format!("Read no segments from `{name} "))?
.finish_write = true;
UploadStatus::Ongoing(OngoingUploadStatus::Done) => {
// Another task has already completed the upload of this artifact, no need to do any work.
}
UploadStatus::Ongoing(OngoingUploadStatus::Error) => {
// Another task tried to perform the transmission, but failed.
anyhow::bail!("Upload queued for previous action failed.")
}
UploadStatus::New(tx) => {
let mut file = tokio::fs::File::open(&name)
.await
.with_context(|| format!("Opening `{name}` for reading failed"))?;
let mut data = vec![0; max_total_batch_size];

let mut write_offset = 0;
let mut upload_segments = Vec::new();
loop {
let length = file
.read(&mut data)
.await
.with_context(|| format!("Error reading from {name}"))?;
if length == 0 {
break;
}
upload_segments.push(WriteRequest {
resource_name: resource_name.to_owned(),
write_offset,
finish_write: false,
data: data[..length].to_owned(),
});
write_offset += length as i64;
}
upload_segments
.last_mut()
.with_context(|| format!("Read no segments from `{name} "))?
.finish_write = true;

let upload_ret = bystream_fut(upload_segments)
.await
.and_then(|resp| {
if resp.committed_size != digest.size_in_bytes {
Err(anyhow::anyhow!(
"Failed to upload `{name}`: invalid committed_size from WriteResponse"
))
}
else {
Ok(())
}
});

let resp = bystream_fut(upload_segments).await?;
if resp.committed_size != size {
return Err(anyhow::anyhow!(
"Failed to upload `{name}`: invalid committed_size from WriteResponse"
));
// Mark artifact as uploaded and notify other potentially waiting tasks.
if upload_ret.is_ok() {
prev_uploads.alter(&digest, |_, _| OngoingUploadStatus::Done);
let _ = tx.send(upload_ret.as_ref().map_err(|_| ()).cloned());
} else {
prev_uploads.alter(&digest, |_, _| OngoingUploadStatus::Error);
let _ = tx.send(Err(()));
}

// Only propage errors _after_ notifying other waiting tasks that this task is complete.
upload_ret?;
}
}
Ok(vec![hash])

Ok(vec![digest.hash])
};
upload_futures.push(Box::pin(fut));
}
Expand Down Expand Up @@ -2071,6 +2133,7 @@ mod tests {
req,
10000,
None,
&DashMap::new(),
|req| {
let res = res.clone();
let digest1 = digest1.clone();
Expand Down Expand Up @@ -2154,6 +2217,7 @@ mod tests {
req,
10, // kept small to simulate a large file upload
None,
&DashMap::new(),
|req| {
let res = res.clone();
let digest1 = digest1.clone();
Expand Down Expand Up @@ -2228,6 +2292,7 @@ mod tests {
req,
10, // kept small to simulate a large inlined upload
None,
&DashMap::new(),
|req| {
let res = res.clone();
let digest1 = digest1.clone();
Expand Down Expand Up @@ -2289,6 +2354,7 @@ mod tests {
req,
10,
None,
&DashMap::new(),
|_req| async move {
panic!("This should not be called as there are no blobs to upload in batch");
},
Expand Down Expand Up @@ -2350,6 +2416,7 @@ mod tests {
req,
3,
None,
&DashMap::new(),
|_req| async move {
panic!("Not called");
},
Expand Down Expand Up @@ -2391,6 +2458,7 @@ mod tests {
req,
0,
None,
&DashMap::new(),
|_req| async move {
panic!("Not called");
},
Expand Down Expand Up @@ -2437,6 +2505,7 @@ mod tests {
req,
1,
None,
&DashMap::new(),
|_req| async move {
panic!("Not called");
},
Expand Down