|
| 1 | +// Copyright 2023 The Turbo Cache Authors. All rights reserved. |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +use std::pin::Pin; |
| 16 | +use std::sync::Arc; |
| 17 | + |
| 18 | +use ac_utils::get_and_decode_digest; |
| 19 | +use action_messages::ActionResult; |
| 20 | +use async_trait::async_trait; |
| 21 | +use futures::stream::FuturesUnordered; |
| 22 | +use futures::FutureExt; |
| 23 | +use futures::{ |
| 24 | + future::BoxFuture, |
| 25 | + stream::{StreamExt, TryStreamExt}, |
| 26 | +}; |
| 27 | +use proto::build::bazel::remote::execution::v2::{ActionResult as ProtoActionResult, Directory as ProtoDirectory}; |
| 28 | + |
| 29 | +use buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; |
| 30 | +use common::DigestInfo; |
| 31 | +use error::{ |
| 32 | + //make_err, Code, |
| 33 | + Error, |
| 34 | + ResultExt, |
| 35 | +}; |
| 36 | +use traits::{StoreTrait, UploadSizeInfo}; |
| 37 | + |
| 38 | +/// Aggressively check if the digests of files exist in the cas. This function |
| 39 | +/// will spawn unbounded number of futures check all of the files. The store itself |
| 40 | +/// should be rate limited if spawning too many requests at once is an issue. |
| 41 | +// Sadly we cannot use `async fn` here because the rust compiler cannot determine the auto traits |
| 42 | +// of the future. So we need to force this function to return a dynamic future instead. |
| 43 | +// see: https://github.com/rust-lang/rust/issues/78649 |
| 44 | +pub fn check_directory_files_in_cas<'a>( |
| 45 | + cas_store: Pin<&'a dyn StoreTrait>, |
| 46 | + digest: &'a DigestInfo, |
| 47 | + current_directory: &'a str, |
| 48 | +) -> BoxFuture<'a, Result<(), Error>> { |
| 49 | + async move { |
| 50 | + let directory = get_and_decode_digest::<ProtoDirectory>(cas_store, digest) |
| 51 | + .await |
| 52 | + .err_tip(|| "Converting digest to Directory")?; |
| 53 | + let mut futures = FuturesUnordered::new(); |
| 54 | + |
| 55 | + for file in directory.files { |
| 56 | + let digest: DigestInfo = file |
| 57 | + .digest |
| 58 | + .err_tip(|| "Expected Digest to exist in Directory::file::digest")? |
| 59 | + .try_into() |
| 60 | + .err_tip(|| "In Directory::file::digest")?; |
| 61 | + // Maybe could be made more efficient |
| 62 | + futures.push(cas_store.has(digest).boxed()); |
| 63 | + } |
| 64 | + |
| 65 | + for directory in directory.directories { |
| 66 | + let digest: DigestInfo = directory |
| 67 | + .digest |
| 68 | + .err_tip(|| "Expected Digest to exist in Directory::directories::digest")? |
| 69 | + .try_into() |
| 70 | + .err_tip(|| "In Directory::file::digest")?; |
| 71 | + let new_directory_path = format!("{}/{}", current_directory, directory.name); |
| 72 | + futures.push( |
| 73 | + async move { |
| 74 | + check_directory_files_in_cas(cas_store, &digest, &new_directory_path) |
| 75 | + .await |
| 76 | + .err_tip(|| format!("in traverse_ : {new_directory_path}"))?; |
| 77 | + Ok(Some(1)) |
| 78 | + } |
| 79 | + .boxed(), |
| 80 | + ); |
| 81 | + } |
| 82 | + |
| 83 | + while futures.try_next().await?.is_some() {} |
| 84 | + Ok(()) |
| 85 | + } |
| 86 | + .boxed() |
| 87 | +} |
| 88 | + |
| 89 | +pub struct CompletenessCheckingStore { |
| 90 | + cas_store: Arc<dyn StoreTrait>, |
| 91 | +} |
| 92 | + |
| 93 | +impl CompletenessCheckingStore { |
| 94 | + pub fn new(_ac_store: Arc<dyn StoreTrait>, cas_store: Arc<dyn StoreTrait>) -> Self { |
| 95 | + CompletenessCheckingStore { cas_store } |
| 96 | + } |
| 97 | + |
| 98 | + fn pin_cas(&self) -> Pin<&dyn StoreTrait> { |
| 99 | + Pin::new(self.cas_store.as_ref()) |
| 100 | + } |
| 101 | +} |
| 102 | + |
| 103 | +#[async_trait] |
| 104 | +impl StoreTrait for CompletenessCheckingStore { |
| 105 | + async fn has_with_results( |
| 106 | + self: Pin<&Self>, |
| 107 | + digests: &[DigestInfo], |
| 108 | + _results: &mut [Option<usize>], |
| 109 | + ) -> Result<(), Error> { |
| 110 | + // The proto promises that all results will exist in the CAS when |
| 111 | + // requested using the associated actions. However we currently allow |
| 112 | + // stores to prune themselves which violates this condition, because the |
| 113 | + // action cache and CAS are different. Therefore we need this completeness checking |
| 114 | + // store to check that all results exist in the CAS before we allow the the action result |
| 115 | + // to be returned to the client. |
| 116 | + // * Take the root digest which is the serialzied action proto hash |
| 117 | + // * Deserialize the action proto |
| 118 | + // * Check files in the root tree exist in the CAS but there's directories that needs |
| 119 | + // to be traversed as the directories can have directories and files in them and esure all |
| 120 | + // of them exist in the CAS. |
| 121 | + |
| 122 | + let pin_cas = self.pin_cas(); |
| 123 | + let mut futures = FuturesUnordered::new(); |
| 124 | + |
| 125 | + for digest in digests { |
| 126 | + let action_result: ActionResult = get_and_decode_digest::<ProtoActionResult>(pin_cas, digest) |
| 127 | + .await? |
| 128 | + .try_into() |
| 129 | + .err_tip(|| "Action result could not be converted in completeness checking store")?; |
| 130 | + |
| 131 | + for output_file in &action_result.output_files { |
| 132 | + let file = output_file.digest; |
| 133 | + let file_digest = DigestInfo::try_new(&file.hash_str(), file.size_bytes as usize)?; |
| 134 | + futures.push( |
| 135 | + async move { |
| 136 | + if let Err(e) = check_directory_files_in_cas(pin_cas, &file_digest, "").await { |
| 137 | + eprintln!("Error: {:?}", e); |
| 138 | + } |
| 139 | + } |
| 140 | + .boxed(), |
| 141 | + ); |
| 142 | + } |
| 143 | + |
| 144 | + for output_directory in action_result.output_folders { |
| 145 | + let path = output_directory.path; |
| 146 | + let tree_digest = output_directory.tree_digest; |
| 147 | + futures.push( |
| 148 | + async move { |
| 149 | + if let Err(e) = check_directory_files_in_cas(pin_cas, &tree_digest, &path).await { |
| 150 | + eprintln!("Error: {:?}", e); |
| 151 | + } |
| 152 | + } |
| 153 | + .boxed(), |
| 154 | + ); |
| 155 | + } |
| 156 | + } |
| 157 | + |
| 158 | + while (futures.next().await).is_some() {} |
| 159 | + |
| 160 | + Ok(()) |
| 161 | + } |
| 162 | + |
| 163 | + async fn update( |
| 164 | + self: Pin<&Self>, |
| 165 | + digest: DigestInfo, |
| 166 | + reader: DropCloserReadHalf, |
| 167 | + size_info: UploadSizeInfo, |
| 168 | + ) -> Result<(), Error> { |
| 169 | + self.pin_cas().update(digest, reader, size_info).await |
| 170 | + } |
| 171 | + |
| 172 | + async fn get_part_ref( |
| 173 | + self: Pin<&Self>, |
| 174 | + digest: DigestInfo, |
| 175 | + writer: &mut DropCloserWriteHalf, |
| 176 | + offset: usize, |
| 177 | + length: Option<usize>, |
| 178 | + ) -> Result<(), Error> { |
| 179 | + self.pin_cas().get_part_ref(digest, writer, offset, length).await |
| 180 | + } |
| 181 | + |
| 182 | + fn as_any(self: Arc<Self>) -> Box<dyn std::any::Any + Send> { |
| 183 | + Box::new(self) |
| 184 | + } |
| 185 | +} |
0 commit comments