Merge branch 'main' into completeness-checking-store-restructuring

Author: blakehatch

Cargo.toml

@@ -42,6 +42,8 @@ members = [
   "gencargo/evicting_map",
   "gencargo/evicting_map_test",
   "gencargo/execution_server",
+  "gencargo/existence_store",
+  "gencargo/existence_store_test",
   "gencargo/fast_slow_store",
   "gencargo/fast_slow_store_test",
   "gencargo/fastcdc",
@@ -190,6 +192,8 @@ error = { path = "gencargo/error" }
 evicting_map = { path = "gencargo/evicting_map" }
 evicting_map_test = { path = "gencargo/evicting_map_test" }
 execution_server = { path = "gencargo/execution_server" }
+existence_store = { path = "gencargo/existence_store" }
+existence_store_test = { path = "gencargo/existence_store_test" }
 fast_slow_store = { path = "gencargo/fast_slow_store" }
 fast_slow_store_test = { path = "gencargo/fast_slow_store_test" }
 fastcdc = { path = "gencargo/fastcdc" }

cas/grpc_service/tests/worker_api_server_test.rs

@@ -28,7 +28,7 @@ use proto::build::bazel::remote::execution::v2::{
     ActionResult as ProtoActionResult, ExecuteResponse, ExecutedActionMetadata, LogFile, OutputDirectory, OutputFile,
     OutputSymlink,
 };
-use proto::com::github::allada::turbo_cache::remote_execution::{
+use proto::com::github::trace_machina::turbo_cache::remote_execution::{
     execute_result, update_for_worker, worker_api_server::WorkerApi, ExecuteResult, KeepAliveRequest,
     SupportedProperties,
 };

cas/grpc_service/worker_api_server.rs

@@ -29,7 +29,7 @@ use common::DigestInfo;
 use config::cas_server::WorkerApiConfig;
 use error::{make_err, Code, Error, ResultExt};
 use platform_property_manager::PlatformProperties;
-use proto::com::github::allada::turbo_cache::remote_execution::{
+use proto::com::github::trace_machina::turbo_cache::remote_execution::{
     execute_result, worker_api_server::WorkerApi, worker_api_server::WorkerApiServer as Server, ExecuteResult,
     GoingAwayRequest, KeepAliveRequest, SupportedProperties, UpdateForWorker,
 };

cas/scheduler/tests/simple_scheduler_test.rs

@@ -27,7 +27,7 @@ use common::DigestInfo;
 use error::{make_err, Code, Error, ResultExt};
 use platform_property_manager::{PlatformProperties, PlatformPropertyValue};
 use proto::build::bazel::remote::execution::v2::{digest_function, ExecuteRequest};
-use proto::com::github::allada::turbo_cache::remote_execution::{
+use proto::com::github::trace_machina::turbo_cache::remote_execution::{
     update_for_worker, ConnectionResult, StartExecute, UpdateForWorker,
 };
 use scheduler::{ActionScheduler, WorkerScheduler};

cas/scheduler/worker.rs

@@ -23,7 +23,7 @@ use action_messages::ActionInfo;
 use error::{make_err, make_input_err, Code, Error, ResultExt};
 use metrics_utils::{CollectorState, CounterWithTime, FuncCounterWrapper, MetricsComponent};
 use platform_property_manager::{PlatformProperties, PlatformPropertyValue};
-use proto::com::github::allada::turbo_cache::remote_execution::{
+use proto::com::github::trace_machina::turbo_cache::remote_execution::{
     update_for_worker, ConnectionResult, StartExecute, UpdateForWorker,
 };
 use tokio::sync::mpsc::UnboundedSender;

cas/store/BUILD

@@ -171,6 +171,27 @@ rust_library(
     ],
 )
 
+rust_library(
+    name = "existence_store",
+    srcs = ["existence_store.rs"],
+    proc_macro_deps = ["@crate_index//:async-trait"],
+    visibility = ["//cas:__pkg__"],
+    deps = [
+        ":traits",
+        ":ac_utils",
+        "//proto",
+        "//config",
+        "//util:buf_channel",
+        "//util:common",
+        "//util:error",
+        "//util:metrics_utils",
+        "@crate_index//:hashbrown",
+        "@crate_index//:hex",
+        "@crate_index//:sha2",
+        "@crate_index//:tokio",
+    ],
+)
+
 rust_library(
     name = "fast_slow_store",
     srcs = ["fast_slow_store.rs"],
@@ -454,6 +475,23 @@ rust_test(
     ],
 )
 
+rust_test(
+    name = "existence_store_test",
+    srcs = ["tests/existence_store_test.rs"],
+    deps = [
+        ":memory_store",
+        ":traits",
+        ":existence_store",
+        "//config",
+        "//util:buf_channel",
+        "//util:common",
+        "//util:error",
+        "@crate_index//:futures",
+        "@crate_index//:pretty_assertions",
+        "@crate_index//:tokio",
+    ],
+)
+
 rust_test(
     name = "s3_store_test",
     srcs = ["tests/s3_store_test.rs"],

cas/store/default_store_factory.rs

@@ -23,6 +23,7 @@ use compression_store::CompressionStore;
 use config::{self, stores::StoreConfig};
 use dedup_store::DedupStore;
 use error::Error;
+use existence_store::ExistenceStore;
 use fast_slow_store::FastSlowStore;
 use filesystem_store::FilesystemStore;
 use grpc_store::GrpcStore;
@@ -63,6 +64,9 @@ pub fn store_factory<'a>(
                 store_factory(&config.index_store, store_manager, None).await?,
                 store_factory(&config.content_store, store_manager, None).await?,
             )),
+            StoreConfig::existence_store(config) => Arc::new(ExistenceStore::new(
+                store_factory(&config.inner, store_manager, None).await?,
+            )),
             StoreConfig::fast_slow(config) => Arc::new(FastSlowStore::new(
                 config,
                 store_factory(&config.fast, store_manager, None).await?,

cas/store/existence_store.rs

@@ -0,0 +1,105 @@
+// Copyright 2022 The Turbo Cache Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use hashbrown::HashSet;
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+
+use async_trait::async_trait;
+
+use buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
+use common::DigestInfo;
+use error::Error;
+use traits::{StoreTrait, UploadSizeInfo};
+
+pub struct ExistenceStore {
+    inner_store: Arc<dyn StoreTrait>,
+    pub existence_cache: Mutex<HashSet<DigestInfo>>,
+}
+
+impl ExistenceStore {
+    pub fn new(inner_store: Arc<dyn StoreTrait>) -> Self {
+        Self {
+            inner_store,
+            // TODO (BlakeHatch):
+            // Consider using RwLock in a future commit.
+            // Since HashSet implements Send and Sync this should
+            // be a drop-in replacement in theory.
+            // Make sure benchmark is done to justify.
+            existence_cache: Mutex::new(HashSet::new()),
+        }
+    }
+
+    fn pin_inner(&self) -> Pin<&dyn StoreTrait> {
+        Pin::new(self.inner_store.as_ref())
+    }
+}
+
+#[async_trait]
+impl StoreTrait for ExistenceStore {
+    async fn has_with_results(
+        self: Pin<&Self>,
+        digests: &[DigestInfo],
+        results: &mut [Option<usize>],
+    ) -> Result<(), Error> {
+        let mut pruned_digests = Vec::new();
+
+        for (i, digest) in digests.iter().enumerate() {
+            if self.existence_cache.lock().unwrap().contains(digest) {
+                results[i] = Some(1);
+            } else {
+                pruned_digests.push(*digest);
+            }
+        }
+
+        if !pruned_digests.is_empty() {
+            let mut inner_results = vec![None; pruned_digests.len()];
+            self.pin_inner()
+                .has_with_results(&pruned_digests, &mut inner_results)
+                .await?;
+
+            for (i, result) in inner_results.iter().enumerate() {
+                if result.is_some() {
+                    self.existence_cache.lock().unwrap().insert(pruned_digests[i]);
+                    results[i] = Some(1);
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn update(
+        self: Pin<&Self>,
+        digest: DigestInfo,
+        reader: DropCloserReadHalf,
+        size_info: UploadSizeInfo,
+    ) -> Result<(), Error> {
+        self.pin_inner().update(digest, reader, size_info).await
+    }
+
+    async fn get_part_ref(
+        self: Pin<&Self>,
+        digest: DigestInfo,
+        writer: &mut DropCloserWriteHalf,
+        offset: usize,
+        length: Option<usize>,
+    ) -> Result<(), Error> {
+        self.pin_inner().get_part_ref(digest, writer, offset, length).await
+    }
+
+    fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> {
+        Box::new(self)
+    }
+}

cas/store/grpc_store.rs

@@ -25,6 +25,7 @@ use prost::Message;
 use proto::build::bazel::remote::execution::v2::digest_function;
 use rand::{rngs::OsRng, Rng};
 use tokio::time::sleep;
+use tonic::transport::Channel;
 use tonic::{transport, IntoRequest, Request, Response, Streaming};
 use uuid::Uuid;
 
@@ -111,16 +112,20 @@ impl GrpcStore {
         })
     }
 
+    fn get_retry_config(&self) -> impl Iterator<Item = Duration> + '_ {
+        ExponentialBackoff::new(Duration::from_millis(self.retry.delay as u64))
+            .map(|d| (self.jitter_fn)(d))
+            .take(self.retry.max_retries) // Remember this is number of retries, so will run max_retries + 1.
+    }
+
     async fn perform_request<F, Fut, R, I>(&self, input: I, mut request: F) -> Result<R, Error>
     where
         F: FnMut(I) -> Fut + Send + Copy,
         Fut: Future<Output = Result<R, Error>> + Send,
         R: Send,
         I: Send + Clone,
     {
-        let retry_config = ExponentialBackoff::new(Duration::from_millis(self.retry.delay as u64))
-            .map(|d| (self.jitter_fn)(d))
-            .take(self.retry.max_retries); // Remember this is number of retries, so will run max_retries + 1.
+        let retry_config = self.get_retry_config();
         self.retrier
             .retry(
                 retry_config,
@@ -263,50 +268,86 @@ impl GrpcStore {
             "CAS operation on AC store"
         );
 
-        let mut client = self.bytestream_client.clone();
-
-        let error = Arc::new(Mutex::new(None));
-        struct LocalState {
+        struct LocalState<T, E>
+        where
+            T: Stream<Item = Result<WriteRequest, E>> + Unpin + Send + 'static,
+            E: Into<Error> + 'static,
+        {
             instance_name: String,
-            error: Arc<Mutex<Option<Error>>>,
+            error: Mutex<Option<Error>>,
+            read_stream: Mutex<Option<WriteRequestStreamWrapper<T, E>>>,
+            client: ByteStreamClient<Channel>,
         }
 
-        let local_state = LocalState {
+        let local_state = Arc::new(LocalState {
             instance_name: self.instance_name.clone(),
-            error: error.clone(),
-        };
+            error: Mutex::new(None),
+            read_stream: Mutex::new(Some(stream)),
+            client: self.bytestream_client.clone(),
+        });
 
-        let stream = unfold((stream, local_state), move |(mut stream, local_state)| async {
-            let maybe_message = stream.next().await;
-            if let Ok(maybe_message) = maybe_message {
-                if let Some(mut message) = maybe_message {
-                    // `resource_name` pattern is: "{instance_name}/uploads/{uuid}/blobs/{hash}/{size}".
-                    let first_slash_pos = match message.resource_name.find('/') {
-                        Some(pos) => pos,
-                        None => {
-                            log::error!("{}", "Resource name should follow pattern {instance_name}/uploads/{uuid}/blobs/{hash}/{size}");
+        let retry_config = self.get_retry_config();
+        let result = self
+            .retrier
+            .retry(
+                retry_config,
+                unfold(local_state, move |local_state| async move {
+                    let stream = unfold((None, local_state.clone()), move |(stream, local_state)| async {
+                        // Only consume the stream on the first request to read,
+                        // then pass it for future requests in the unfold.
+                        let Some(mut stream) = stream.or_else(|| local_state.read_stream.lock().take()) else {
+                            return None;
+                        };
+                        let maybe_message = stream.next().await;
+                        if let Ok(maybe_message) = maybe_message {
+                            if let Some(mut message) = maybe_message {
+                                // `resource_name` pattern is: "{instance_name}/uploads/{uuid}/blobs/{hash}/{size}".
+                                let first_slash_pos = match message.resource_name.find('/') {
+                                    Some(pos) => pos,
+                                    None => {
+                                        log::error!("{}", "Resource name should follow pattern {instance_name}/uploads/{uuid}/blobs/{hash}/{size}");
+                                        return None;
+                                    }
+                                };
+                                message.resource_name = format!(
+                                    "{}/{}",
+                                    &local_state.instance_name,
+                                    message.resource_name.get((first_slash_pos + 1)..).unwrap()
+                                );
+                                return Some((message, (Some(stream), local_state)));
+                            }
                             return None;
                         }
+                        // TODO(allada) I'm sure there's a way to do this without a mutex, but rust can be super
+                        // picky with borrowing through a stream await.
+                        *local_state.error.lock() = Some(maybe_message.unwrap_err());
+                        None
+                    });
+
+                    let result = local_state.client.clone()
+                            .write(stream)
+                            .await
+                            .err_tip(|| "in GrpcStore::write");
+
+                    // If the stream has been consumed, don't retry, but
+                    // otherwise it's ok to try again.
+                    let result = if local_state.read_stream.lock().is_some() {
+                        result.map_or_else(RetryResult::Retry, RetryResult::Ok)
+                    } else {
+                        result.map_or_else(RetryResult::Err, RetryResult::Ok)
                     };
-                    message.resource_name = format!(
-                        "{}/{}",
-                        &local_state.instance_name,
-                        message.resource_name.get((first_slash_pos + 1)..).unwrap()
-                    );
-                    return Some((message, (stream, local_state)));
-                }
-                return None;
-            }
-            // TODO(allada) I'm sure there's a way to do this without a mutex, but rust can be super
-            // picky with borrowing through a stream await.
-            *local_state.error.lock() = Some(maybe_message.unwrap_err());
-            None
-        });
 
-        let result = client.write(stream).await.err_tip(|| "in GrpcStore::write")?;
-        if let Some(err) = error.lock().take() {
-            return Err(err);
-        }
+                    // If there was an error with the stream, then don't retry.
+                    let result = if let Some(err) = local_state.error.lock().take() {
+                        RetryResult::Err(err)
+                    } else {
+                        result
+                    };
+
+                    Some((result, local_state))
+                }),
+            )
+            .await?;
         Ok(result)
     }