Skip to content

remove find_missing_blobs deduplication to fix bug #1121

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

Closed
bvinc wants to merge 1 commit into from
Closed

remove find_missing_blobs deduplication to fix bug #1121

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
133 changes: 26 additions & 107 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 @@ -82,7 +82,6 @@ use tokio::io::AsyncRead;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::sync::broadcast;
use tokio_util::io::StreamReader;
use tonic::codegen::InterceptedService;
use tonic::metadata;
Expand Down Expand Up @@ -558,20 +557,8 @@ enum DigestRemoteState {
Missing,
}

#[derive(Debug)]
enum FindMissingCacheState {
Finished(DigestRemoteState),
InFlight(broadcast::Receiver<Result<DigestRemoteState, ()>>),
}

#[derive(Debug)]
enum FindMissingStartResult {
InCache(FindMissingCacheState),
Started(broadcast::Sender<Result<DigestRemoteState, ()>>),
}

struct FindMissingCache {
cache: LruCache<TDigest, FindMissingCacheState>,
cache: LruCache<TDigest, DigestRemoteState>,
/// To avoid a situation where we cache that an artifact is available remotely, but the artifact then expires
/// we clear our local cache once every `ttl`.
ttl: Duration,
Expand All @@ -586,32 +573,14 @@ impl FindMissingCache {
}
}

pub fn get_or_start(&mut self, digest: &TDigest) -> FindMissingStartResult {
pub fn get(&mut self, digest: &TDigest) -> Option<DigestRemoteState> {
self.clear_if_ttl_expires();

let mut inserted_sender = None;
let cache_state = self.cache.get_or_insert(digest.clone(), || {
let (tx, rx) = broadcast::channel(1);
inserted_sender = Some(tx);
FindMissingCacheState::InFlight(rx)
});
if let Some(sender) = inserted_sender {
return FindMissingStartResult::Started(sender);
}
match cache_state {
FindMissingCacheState::Finished(drs) => {
FindMissingStartResult::InCache(FindMissingCacheState::Finished(*drs))
}
FindMissingCacheState::InFlight(rx) => {
FindMissingStartResult::InCache(FindMissingCacheState::InFlight(rx.resubscribe()))
}
}
self.cache.get(digest).copied()
}

pub fn put(&mut self, digest: TDigest, state: DigestRemoteState) {
self.clear_if_ttl_expires();
self.cache
.put(digest, FindMissingCacheState::Finished(state));
self.cache.put(digest, state);
}
}

Expand Down Expand Up @@ -998,32 +967,20 @@ impl REClient {
) -> anyhow::Result<GetDigestsTtlResponse> {
let mut cas_client = self.grpc_clients.cas_client.clone();
let mut remote_results: HashMap<TDigest, DigestRemoteState> = HashMap::new();
let mut digests_waiting: HashMap<
TDigest,
broadcast::Receiver<Result<DigestRemoteState, ()>>,
> = HashMap::new();
let mut digests_to_check: Vec<(TDigest, broadcast::Sender<Result<DigestRemoteState, ()>>)> =
Vec::new();
let mut digests_to_check: Vec<TDigest> = Vec::new();

let mut digest_iter = request.digests.iter();
while digest_iter.len() > 0 {
// Sort our blobs based on what action we need to take
{
let mut find_missing_cache = self.find_missing_cache.lock().unwrap();
for digest in digest_iter.by_ref() {
match find_missing_cache.get_or_start(&digest) {
if let Some(rs) = find_missing_cache.get(&digest) {
// We have our final result already cached
FindMissingStartResult::InCache(FindMissingCacheState::Finished(rs)) => {
remote_results.insert(digest.clone(), rs);
}
// This blob is already being requested by another task
FindMissingStartResult::InCache(FindMissingCacheState::InFlight(rx)) => {
digests_waiting.insert(digest.clone(), rx);
}
remote_results.insert(digest.clone(), rs);
} else {
// We can check this blob
FindMissingStartResult::Started(tx) => {
digests_to_check.push((digest.clone(), tx));
}
digests_to_check.push(digest.clone());
}
if digests_to_check.len() >= 100 {
break;
Expand All @@ -1034,72 +991,34 @@ impl REClient {
// Send a request and notify others of the result
if !digests_to_check.is_empty() {
tracing::debug!(num_digests = digests_to_check.len(), "FindMissingBlobs");
match cas_client
let missing_blobs = cas_client
.find_missing_blobs(with_re_metadata(
FindMissingBlobsRequest {
instance_name: self.instance_name.as_str().to_owned(),
blob_digests: digests_to_check.map(|b| tdigest_to(b.0.clone())),
blob_digests: digests_to_check.map(|b| tdigest_to(b.clone())),
..Default::default()
},
metadata.clone(),
self.runtime_opts.use_fbcode_metadata,
))
.await
{
Err(e) => {
// Notify others of the failure
for (_, tx) in &*digests_to_check {
let _ = tx.send(Err(()));
}
digests_to_check.clear();
return Err(e)
.context("Failed to request what blobs are not present on remote");
}
Ok(missing_blobs) => {
let resp: FindMissingBlobsResponse = missing_blobs.into_inner();

// Update the results and the cache
{
let mut find_missing_cache = self.find_missing_cache.lock().unwrap();
for (digest, _) in &*digests_to_check {
remote_results
.insert(digest.clone(), DigestRemoteState::ExistsOnRemote);
find_missing_cache
.put(digest.clone(), DigestRemoteState::ExistsOnRemote);
}
for digest in
&resp.missing_blob_digests.map(|d| tdigest_from(d.clone()))
{
// If it's present in the MissingBlobsResponse, it's expired on the remote and
// needs to be refetched.
remote_results.insert(digest.clone(), DigestRemoteState::Missing);
find_missing_cache.put(digest.clone(), DigestRemoteState::Missing);
}
}
.context("Failed to request what blobs are not present on remote")?;
let resp: FindMissingBlobsResponse = missing_blobs.into_inner();

// Notify others of the success
for (digest, tx) in &*digests_to_check {
let rem_state = remote_results.get(digest).unwrap();
let _ = tx.send(Ok(*rem_state));
}
digests_to_check.clear();
}
};
}
}
// Update the results and the cache
let mut find_missing_cache = self.find_missing_cache.lock().unwrap();
for digest in &digests_to_check {
remote_results.insert(digest.clone(), DigestRemoteState::ExistsOnRemote);
find_missing_cache.put(digest.clone(), DigestRemoteState::ExistsOnRemote);
}

// Wait to receive results from others
for (digest, rx) in &mut digests_waiting {
if let Ok(rs) = rx
.recv()
.await
.context("Failed to receive FindMissingBlobs result from requesting task")?
{
remote_results.insert(digest.clone(), rs);
} else {
return Err(anyhow::anyhow!(
"FindMissingBlobs request failed in other task"
));
for digest in &resp.missing_blob_digests.map(|d| tdigest_from(d.clone())) {
// If it's present in the MissingBlobsResponse, it's expired on the remote and
// needs to be refetched.
remote_results.insert(digest.clone(), DigestRemoteState::Missing);
find_missing_cache.put(digest.clone(), DigestRemoteState::Missing);
}
digests_to_check.clear();
}
}

Expand Down