Added completeness checking store

Author: blakehatch

Cargo.lock

@@ -29,6 +29,8 @@ members = [
   "gencargo/cas_server",
   "gencargo/cas_server_test",
   "gencargo/common",
+  "gencargo/completeness_checking_store",
+  "gencargo/completeness_checking_store_test",
   "gencargo/compression_store",
   "gencargo/compression_store_test",
   "gencargo/config",
@@ -175,6 +177,8 @@ cas = { path = "gencargo/cas" }
 cas_server = { path = "gencargo/cas_server" }
 cas_server_test = { path = "gencargo/cas_server_test" }
 common = { path = "gencargo/common" }
+completeness_checking_store = { path = "gencargo/completeness_checking_store" }
+completeness_checking_store_test = { path = "gencargo/completeness_checking_store_test" }
 compression_store = { path = "gencargo/compression_store" }
 compression_store_test = { path = "gencargo/compression_store_test" }
 config = { path = "gencargo/config" }

Cargo.toml

@@ -748,6 +748,53 @@ impl From<ActionResult> for ProtoActionResult {
     }
 }
 
+impl From<ProtoActionResult> for ActionResult {
+    fn from(val: ProtoActionResult) -> Self {
+        let output_file_symlinks = val
+            .output_file_symlinks
+            .into_iter()
+            .map(|output_symlink| SymlinkInfo::try_from(output_symlink).unwrap())
+            .collect::<Vec<_>>();
+
+        let output_directory_symlinks = val
+            .output_directory_symlinks
+            .into_iter()
+            .map(|output_symlink| SymlinkInfo::try_from(output_symlink).unwrap())
+            .collect::<Vec<_>>();
+
+        Self {
+            output_files: val
+                .output_files
+                .into_iter()
+                .map(|output_file| output_file.try_into().unwrap())
+                .collect(),
+            output_folders: val
+                .output_directories
+                .into_iter()
+                .map(|output_directory| output_directory.try_into().unwrap())
+                .collect(),
+            output_file_symlinks,
+            output_directory_symlinks,
+            exit_code: val.exit_code,
+            stdout_digest: val
+                .stdout_digest
+                .map(|digest| digest.try_into().unwrap_or(DigestInfo::empty_digest()))
+                .unwrap_or(DigestInfo::empty_digest()),
+            stderr_digest: val
+                .stderr_digest
+                .map(|digest| digest.try_into().unwrap_or(DigestInfo::empty_digest()))
+                .unwrap_or_else(DigestInfo::empty_digest),
+            execution_metadata: val
+                .execution_metadata
+                .map(|metadata| metadata.try_into().unwrap_or(ExecutionMetadata::default()))
+                .unwrap_or_default(),
+            server_logs: Default::default(),
+            error: None,
+            message: String::new(),
+        }
+    }
+}
+
 impl TryFrom<ExecuteResponse> for ActionStage {
     type Error = Error;
 

cas/scheduler/action_messages.rs

@@ -35,6 +35,7 @@ rust_library(
         ":store",
         ":traits",
         ":verify_store",
+        ":completeness_checking_store",
         "//config",
         "//util:error",
         "//util:metrics_utils",
@@ -125,18 +126,44 @@ rust_library(
     ],
 )
 
+rust_library(
+    name = "completeness_checking_store",
+    srcs = ["completeness_checking_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",   
+        "//cas/scheduler:action_messages",
+        "@crate_index//:futures",
+        "@crate_index//:hashbrown",
+        "@crate_index//:hex",
+        "@crate_index//:sha2",
+        "@crate_index//:tokio",
+    ],
+)
+
 rust_library(
     name = "verify_store",
     srcs = ["verify_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",
@@ -388,6 +415,27 @@ rust_test(
     ],
 )
 
+rust_test(
+    name = "completeness_checking_store_test",
+    srcs = ["tests/completeness_checking_store_test.rs"],
+    deps = [
+        ":memory_store",
+        ":traits",
+        ":completeness_checking_store",
+        ":store",
+        "//proto",
+        "//config",
+        "//util:buf_channel",
+        "//util:common",
+        "//util:error",
+        "//cas/scheduler:action_messages",
+        "@crate_index//:prost",
+        "@crate_index//:futures",
+        "@crate_index//:pretty_assertions",
+        "@crate_index//:tokio",
+    ],
+)
+
 rust_test(
     name = "verify_store_test",
     srcs = ["tests/verify_store_test.rs"],
@@ -498,4 +546,4 @@ rust_test(
         "@crate_index//:rand",
         "@crate_index//:tokio",
     ],
-)
+)

cas/store/BUILD

@@ -0,0 +1,122 @@
+// Copyright 2023 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 std::pin::Pin;
+use std::sync::Arc;
+
+use ac_utils::get_and_decode_digest;
+use action_messages::ActionResult;
+use async_trait::async_trait;
+use futures::stream::{self, StreamExt, TryStreamExt};
+use proto::build::bazel::remote::execution::v2::{ActionResult as ProtoActionResult, Tree};
+
+use buf_channel::{DropCloserReadHalf, DropCloserWriteHalf};
+use common::DigestInfo;
+use error::{Error, ResultExt};
+use traits::{StoreTrait, UploadSizeInfo};
+
+pub struct CompletenessCheckingStore {
+    cas_store: Arc<dyn StoreTrait>,
+}
+
+impl CompletenessCheckingStore {
+    pub fn new(_ac_store: Arc<dyn StoreTrait>, cas_store: Arc<dyn StoreTrait>) -> Self {
+        CompletenessCheckingStore { cas_store }
+    }
+
+    fn pin_cas(&self) -> Pin<&dyn StoreTrait> {
+        Pin::new(self.cas_store.as_ref())
+    }
+}
+
+#[async_trait]
+impl StoreTrait for CompletenessCheckingStore {
+    async fn has_with_results(
+        self: Pin<&Self>,
+        digests: &[DigestInfo],
+        _results: &mut [Option<usize>],
+    ) -> Result<(), Error> {
+        // The proto promises that all results will exist in the CAS when
+        // requested using the associated actions. However we currently allow
+        // stores to prune themselves which violates this condition, because the
+        // action cache and CAS are different. Therefore we need this completeness checking
+        // store to check that all results exist in the CAS before we allow the the action result
+        // to be returned to the client.
+        // * Take the root digest which is the serialzied action proto hash
+        // * Deserialize the action proto
+        // * Check files in the root tree exist in the CAS but there's directories that needs
+        // to be traversed as the directories can have directories and files in them and esure all
+        // of them exist in the CAS.
+
+        for digest in digests {
+            let action_result: ActionResult = get_and_decode_digest::<ProtoActionResult>(self.pin_cas(), digest)
+                .await?
+                .try_into()
+                .err_tip(|| "Action result could not be converted in completeness checking store has")?;
+
+            // Something that takes a stream of iterators and flattens it into a stream of items
+            // 'temp' is a stream of digests obtained from the output folders of the action result
+            let _temp =
+                stream::iter(action_result.output_folders)
+                    .then(|directory| async move {
+                        get_and_decode_digest::<Tree>(self.pin_cas(), &directory.tree_digest).await
+                    })
+                    .map_ok(|tree| async {
+                        let trees =
+                            tree.children
+                                .into_iter()
+                                .chain(tree.root.into_iter())
+                                .flat_map(|child_directory| {
+                                    child_directory.files.into_iter().filter_map(|file| file.digest)
+                                });
+
+                        let item_stream = stream::iter(trees);
+                        item_stream
+                            .map(|file_digest| async move {
+                                let digest_info =
+                                    DigestInfo::try_new(&file_digest.hash, file_digest.size_bytes as usize)?;
+                                self.pin_cas().has(digest_info).await
+                            })
+                            .buffer_unordered(10)
+                            .try_collect::<Vec<_>>()
+                            .await
+                    });
+        }
+
+        Ok(())
+    }
+
+    async fn update(
+        self: Pin<&Self>,
+        digest: DigestInfo,
+        reader: DropCloserReadHalf,
+        size_info: UploadSizeInfo,
+    ) -> Result<(), Error> {
+        self.pin_cas().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_cas().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/completeness_checking_store.rs

@@ -18,6 +18,7 @@ use std::sync::Arc;
 use futures::stream::FuturesOrdered;
 use futures::{Future, TryStreamExt};
 
+use completeness_checking_store::CompletenessCheckingStore;
 use compression_store::CompressionStore;
 use config::{self, stores::StoreConfig};
 use dedup_store::DedupStore;
@@ -45,6 +46,10 @@ pub fn store_factory<'a>(
         let store: Arc<dyn Store> = match backend {
             StoreConfig::memory(config) => Arc::new(MemoryStore::new(config)),
             StoreConfig::s3_store(config) => Arc::new(S3Store::new(config)?),
+            StoreConfig::completeness_checking_store(config) => Arc::new(CompletenessCheckingStore::new(
+                store_factory(&config.backend, store_manager, None).await?,
+                store_factory(&config.backend, store_manager, None).await?,
+            )),
             StoreConfig::verify(config) => Arc::new(VerifyStore::new(
                 config,
                 store_factory(&config.backend, store_manager, None).await?,