/build/source/nativelink-store/src/dedup_store.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink 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 core::cmp; |
16 | | use core::pin::Pin; |
17 | | use std::sync::Arc; |
18 | | |
19 | | use async_trait::async_trait; |
20 | | use bincode::serde::{decode_from_slice, encode_to_vec}; |
21 | | use futures::stream::{self, FuturesOrdered, StreamExt, TryStreamExt}; |
22 | | use nativelink_config::stores::DedupSpec; |
23 | | use nativelink_error::{Code, Error, ResultExt, make_err}; |
24 | | use nativelink_metric::MetricsComponent; |
25 | | use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; |
26 | | use nativelink_util::common::DigestInfo; |
27 | | use nativelink_util::fastcdc::FastCDC; |
28 | | use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; |
29 | | use nativelink_util::store_trait::{Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo}; |
30 | | use serde::{Deserialize, Serialize}; |
31 | | use tokio_util::codec::FramedRead; |
32 | | use tokio_util::io::StreamReader; |
33 | | use tracing::warn; |
34 | | |
35 | | use crate::cas_utils::is_zero_digest; |
36 | | |
37 | | // NOTE: If these change update the comments in `stores.rs` to reflect |
38 | | // the new defaults. |
39 | | const DEFAULT_MIN_SIZE: u64 = 64 * 1024; |
40 | | const DEFAULT_NORM_SIZE: u64 = 256 * 1024; |
41 | | const DEFAULT_MAX_SIZE: u64 = 512 * 1024; |
42 | | const DEFAULT_MAX_CONCURRENT_FETCH_PER_GET: usize = 10; |
43 | | |
44 | | #[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Default, Clone)] |
45 | | pub struct DedupIndex { |
46 | | pub entries: Vec<DigestInfo>, |
47 | | } |
48 | | |
49 | | type LegacyBincodeConfig = bincode::config::Configuration< |
50 | | bincode::config::LittleEndian, |
51 | | bincode::config::Fixint, |
52 | | bincode::config::NoLimit, |
53 | | >; |
54 | | |
55 | | #[derive(MetricsComponent)] |
56 | | pub struct DedupStore { |
57 | | #[metric(group = "index_store")] |
58 | | index_store: Store, |
59 | | #[metric(group = "content_store")] |
60 | | content_store: Store, |
61 | | fast_cdc_decoder: FastCDC, |
62 | | #[metric(help = "Maximum number of concurrent fetches per get")] |
63 | | max_concurrent_fetch_per_get: usize, |
64 | | bincode_config: LegacyBincodeConfig, |
65 | | } |
66 | | |
67 | | impl core::fmt::Debug for DedupStore { |
68 | 0 | fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { |
69 | 0 | f.debug_struct("DedupStore") |
70 | 0 | .field("index_store", &self.index_store) |
71 | 0 | .field("content_store", &self.content_store) |
72 | 0 | .field("fast_cdc_decoder", &self.fast_cdc_decoder) |
73 | 0 | .field( |
74 | 0 | "max_concurrent_fetch_per_get", |
75 | 0 | &self.max_concurrent_fetch_per_get, |
76 | 0 | ) |
77 | 0 | .finish_non_exhaustive() |
78 | 0 | } |
79 | | } |
80 | | |
81 | | impl DedupStore { |
82 | 8 | pub fn new( |
83 | 8 | spec: &DedupSpec, |
84 | 8 | index_store: Store, |
85 | 8 | content_store: Store, |
86 | 8 | ) -> Result<Arc<Self>, Error> { |
87 | 8 | let min_size = if spec.min_size == 0 { Branch (87:27): [True: 0, False: 8]
Branch (87:27): [Folded - Ignored]
|
88 | 0 | DEFAULT_MIN_SIZE |
89 | | } else { |
90 | 8 | u64::from(spec.min_size) |
91 | | }; |
92 | 8 | let normal_size = if spec.normal_size == 0 { Branch (92:30): [True: 0, False: 8]
Branch (92:30): [Folded - Ignored]
|
93 | 0 | DEFAULT_NORM_SIZE |
94 | | } else { |
95 | 8 | u64::from(spec.normal_size) |
96 | | }; |
97 | 8 | let max_size = if spec.max_size == 0 { Branch (97:27): [True: 0, False: 8]
Branch (97:27): [Folded - Ignored]
|
98 | 0 | DEFAULT_MAX_SIZE |
99 | | } else { |
100 | 8 | u64::from(spec.max_size) |
101 | | }; |
102 | 8 | let max_concurrent_fetch_per_get = if spec.max_concurrent_fetch_per_get == 0 { Branch (102:47): [True: 0, False: 8]
Branch (102:47): [Folded - Ignored]
|
103 | 0 | DEFAULT_MAX_CONCURRENT_FETCH_PER_GET |
104 | | } else { |
105 | 8 | spec.max_concurrent_fetch_per_get as usize |
106 | | }; |
107 | 8 | Ok(Arc::new(Self { |
108 | 8 | index_store, |
109 | 8 | content_store, |
110 | 8 | fast_cdc_decoder: FastCDC::new( |
111 | 8 | usize::try_from(min_size).err_tip(|| "Could not convert min_size to usize")?0 , |
112 | 8 | usize::try_from(normal_size) |
113 | 8 | .err_tip(|| "Could not convert normal_size to usize")?0 , |
114 | 8 | usize::try_from(max_size).err_tip(|| "Could not convert max_size to usize")?0 , |
115 | | ), |
116 | 8 | max_concurrent_fetch_per_get, |
117 | 8 | bincode_config: bincode::config::legacy(), |
118 | | })) |
119 | 8 | } |
120 | | |
121 | 4 | async fn has(self: Pin<&Self>, key: StoreKey<'_>) -> Result<Option<u64>, Error> { |
122 | | // First we need to load the index that contains where the individual parts actually |
123 | | // can be fetched from. |
124 | 3 | let index_entries = { |
125 | 4 | let maybe_data = self |
126 | 4 | .index_store |
127 | 4 | .get_part_unchunked(key.borrow(), 0, None) |
128 | 4 | .await |
129 | 4 | .err_tip(|| "Failed to read index store in dedup store"); |
130 | 4 | let data3 = match maybe_data { |
131 | 1 | Err(e) => { |
132 | 1 | if e.code == Code::NotFound { Branch (132:24): [True: 1, False: 0]
Branch (132:24): [Folded - Ignored]
|
133 | 1 | return Ok(None); |
134 | 0 | } |
135 | 0 | return Err(e); |
136 | | } |
137 | 3 | Ok(data) => data, |
138 | | }; |
139 | | |
140 | 3 | match decode_from_slice::<DedupIndex, _>(&data, self.bincode_config) { |
141 | 3 | Ok((dedup_index, _)) => dedup_index, |
142 | 0 | Err(err) => { |
143 | 0 | warn!(?key, ?err, "Failed to deserialize index in dedup store",); |
144 | | // We return the equivalent of NotFound here so the client is happy. |
145 | 0 | return Ok(None); |
146 | | } |
147 | | } |
148 | | }; |
149 | | |
150 | 3 | let digests: Vec<_> = index_entries |
151 | 3 | .entries |
152 | 3 | .into_iter() |
153 | 3 | .map(StoreKey::Digest) |
154 | 3 | .collect(); |
155 | 3 | let mut sum = 0; |
156 | 8 | for size in self.content_store3 .has_many3 (&digests).await?0 { |
157 | 8 | let Some(size7 ) = size else { Branch (157:17): [True: 7, False: 1]
Branch (157:17): [Folded - Ignored]
|
158 | | // A part is missing so return None meaning not-found. |
159 | | // This will abort all in-flight queries related to this request. |
160 | 1 | return Ok(None); |
161 | | }; |
162 | 7 | sum += size; |
163 | | } |
164 | 2 | Ok(Some(sum)) |
165 | 4 | } |
166 | | } |
167 | | |
168 | | #[async_trait] |
169 | | impl StoreDriver for DedupStore { |
170 | | async fn has_with_results( |
171 | | self: Pin<&Self>, |
172 | | digests: &[StoreKey<'_>], |
173 | | results: &mut [Option<u64>], |
174 | 10 | ) -> Result<(), Error> { |
175 | 5 | digests |
176 | 5 | .iter() |
177 | 5 | .zip(results.iter_mut()) |
178 | 5 | .map(|(key, result)| async move { |
179 | 5 | if is_zero_digest(key.borrow()) { Branch (179:20): [True: 1, False: 4]
Branch (179:20): [Folded - Ignored]
|
180 | 1 | *result = Some(0); |
181 | 1 | return Ok(()); |
182 | 4 | } |
183 | | |
184 | 4 | match self.has(key.borrow()).await { |
185 | 4 | Ok(maybe_size) => { |
186 | 4 | *result = maybe_size; |
187 | 4 | Ok(()) |
188 | | } |
189 | 0 | Err(err) => Err(err), |
190 | | } |
191 | 10 | }) |
192 | 5 | .collect::<FuturesOrdered<_>>() |
193 | 5 | .try_collect() |
194 | 5 | .await |
195 | 10 | } |
196 | | |
197 | | async fn update( |
198 | | self: Pin<&Self>, |
199 | | key: StoreKey<'_>, |
200 | | reader: DropCloserReadHalf, |
201 | | _size_info: UploadSizeInfo, |
202 | 14 | ) -> Result<(), Error> { |
203 | 7 | let mut bytes_reader = StreamReader::new(reader); |
204 | 7 | let frame_reader = FramedRead::new(&mut bytes_reader, self.fast_cdc_decoder.clone()); |
205 | 7 | let index_entries = frame_reader |
206 | 83 | .map7 (|r| r.err_tip(|| "Failed to decode frame from fast_cdc")) |
207 | 83 | .map_ok7 (|frame| async move { |
208 | 83 | let hash = blake3::hash(&frame[..]).into(); |
209 | 83 | let index_entry = DigestInfo::new(hash, frame.len() as u64); |
210 | 83 | if self Branch (210:20): [True: 0, False: 83]
Branch (210:20): [Folded - Ignored]
|
211 | 83 | .content_store |
212 | 83 | .has(index_entry) |
213 | 83 | .await |
214 | 83 | .err_tip(|| "Failed to call .has() in DedupStore::update()")?0 |
215 | 83 | .is_some() |
216 | | { |
217 | | // If our store has this digest, we don't need to upload it. |
218 | 0 | return Result::<_, Error>::Ok(index_entry); |
219 | 83 | } |
220 | 83 | self.content_store |
221 | 83 | .update_oneshot(index_entry, frame) |
222 | 83 | .await |
223 | 83 | .err_tip(|| "Failed to update content store in dedup_store")?0 ; |
224 | 83 | Ok(index_entry) |
225 | 166 | }) |
226 | 7 | .try_buffered(self.max_concurrent_fetch_per_get) |
227 | 7 | .try_collect() |
228 | 7 | .await?0 ; |
229 | | |
230 | 7 | let serialized_index = encode_to_vec( |
231 | 7 | &DedupIndex { |
232 | 7 | entries: index_entries, |
233 | 7 | }, |
234 | 7 | self.bincode_config, |
235 | | ) |
236 | 7 | .map_err(|e| {0 |
237 | 0 | make_err!( |
238 | 0 | Code::Internal, |
239 | | "Failed to serialize index in dedup_store : {:?}", |
240 | | e |
241 | | ) |
242 | 0 | })?; |
243 | | |
244 | 7 | self.index_store |
245 | 7 | .update_oneshot(key, serialized_index.into()) |
246 | 7 | .await |
247 | 7 | .err_tip(|| "Failed to insert our index entry to index_store in dedup_store")?0 ; |
248 | | |
249 | 7 | Ok(()) |
250 | 14 | } |
251 | | |
252 | | async fn get_part( |
253 | | self: Pin<&Self>, |
254 | | key: StoreKey<'_>, |
255 | | writer: &mut DropCloserWriteHalf, |
256 | | offset: u64, |
257 | | length: Option<u64>, |
258 | 1.87k | ) -> Result<(), Error> { |
259 | | // Special case for if a client tries to read zero bytes. |
260 | 935 | if length == Some(0) { Branch (260:12): [True: 30, False: 905]
Branch (260:12): [Folded - Ignored]
|
261 | 30 | writer |
262 | 30 | .send_eof() |
263 | 30 | .err_tip(|| "Failed to write EOF out from get_part dedup")?0 ; |
264 | 30 | return Ok(()); |
265 | 905 | } |
266 | | // First we need to download the index that contains where the individual parts actually |
267 | | // can be fetched from. |
268 | 905 | let index_entries = { |
269 | 905 | let data = self |
270 | 905 | .index_store |
271 | 905 | .get_part_unchunked(key, 0, None) |
272 | 905 | .await |
273 | 905 | .err_tip(|| "Failed to read index store in dedup store")?0 ; |
274 | 905 | let (dedup_index, _) = decode_from_slice::<DedupIndex, _>(&data, self.bincode_config) |
275 | 905 | .map_err(|e| {0 |
276 | 0 | make_err!( |
277 | 0 | Code::Internal, |
278 | | "Failed to deserialize index in dedup_store::get_part : {:?}", |
279 | | e |
280 | | ) |
281 | 0 | })?; |
282 | 905 | dedup_index |
283 | | }; |
284 | | |
285 | 905 | let mut start_byte_in_stream: u64 = 0; |
286 | 905 | let entries = { |
287 | 905 | if offset == 0 && length31 .is_none31 () { Branch (287:16): [True: 31, False: 874]
Branch (287:31): [True: 2, False: 29]
Branch (287:16): [Folded - Ignored]
Branch (287:31): [Folded - Ignored]
|
288 | 2 | index_entries.entries |
289 | | } else { |
290 | 903 | let mut current_entries_sum = 0; |
291 | 903 | let mut entries = Vec::with_capacity(index_entries.entries.len()); |
292 | 4.82k | for entry4.24k in index_entries.entries { |
293 | 4.24k | let first_byte = current_entries_sum; |
294 | 4.24k | let entry_size = entry.size_bytes(); |
295 | 4.24k | current_entries_sum += entry_size; |
296 | | // Filter any items who's end byte is before the first requested byte. |
297 | 4.24k | if current_entries_sum <= offset { Branch (297:24): [True: 1.86k, False: 2.38k]
Branch (297:24): [Folded - Ignored]
|
298 | 1.86k | start_byte_in_stream = current_entries_sum; |
299 | 1.86k | continue; |
300 | 2.38k | } |
301 | | // If we are not going to read any bytes past the length we are done. |
302 | 2.38k | if let Some(length2.37k ) = length { Branch (302:28): [True: 2.37k, False: 8]
Branch (302:28): [Folded - Ignored]
|
303 | 2.37k | if first_byte >= offset + length { Branch (303:28): [True: 326, False: 2.04k]
Branch (303:28): [Folded - Ignored]
|
304 | 326 | break; |
305 | 2.04k | } |
306 | 8 | } |
307 | 2.05k | entries.push(entry); |
308 | | } |
309 | 903 | entries |
310 | | } |
311 | | }; |
312 | | |
313 | | // Second we we create a stream of futures for each chunk, but buffer/limit them so only |
314 | | // `max_concurrent_fetch_per_get` will be executed at a time. |
315 | | // The results will be streamed out in the same order they are in the entries table. |
316 | | // The results will execute in a "window-like" fashion, meaning that if we limit to |
317 | | // 5 requests at a time, and request 3 is stalled, request 1 & 2 can be output and |
318 | | // request 4 & 5 can be executing (or finished) while waiting for 3 to finish. |
319 | | // Note: We will buffer our data here up to: |
320 | | // `spec.max_size * spec.max_concurrent_fetch_per_get` per `get_part()` request. |
321 | 905 | let mut entries_stream = stream::iter(entries) |
322 | 2.11k | .map905 (move |index_entry| async move { |
323 | 2.11k | let data2.11k = self |
324 | 2.11k | .content_store |
325 | 2.11k | .get_part_unchunked(index_entry, 0, None) |
326 | 2.11k | .await |
327 | 2.11k | .err_tip(|| "Failed to get_part in content_store in dedup_store")?1 ; |
328 | | |
329 | 2.11k | Result::<_, Error>::Ok(data) |
330 | 4.23k | }) |
331 | 905 | .buffered(self.max_concurrent_fetch_per_get); |
332 | | |
333 | | // Stream out the buffered data one at a time and write the data to our writer stream. |
334 | | // In the event any of these error, we will abort early and abandon all the rest of the |
335 | | // streamed data. |
336 | | // Note: Need to take special care to ensure we send the proper slice of data requested. |
337 | 905 | let mut bytes_to_skip = usize::try_from(offset - start_byte_in_stream) |
338 | 905 | .err_tip(|| "Could not convert (offset - start_byte_in_stream) to usize")?0 ; |
339 | 905 | let mut bytes_to_send = usize::try_from(length.unwrap_or(u64::MAX - offset)) |
340 | 905 | .err_tip(|| "Could not convert length to usize")?0 ; |
341 | 3.01k | while let Some(result2.11k ) = entries_stream.next().await { Branch (341:19): [True: 2.11k, False: 904]
Branch (341:19): [Folded - Ignored]
|
342 | 2.11k | let mut data2.11k = result.err_tip(|| "Inner store iterator closed early in DedupStore")?1 ; |
343 | 2.11k | assert!( |
344 | 2.11k | bytes_to_skip <= data.len(), |
345 | 0 | "Formula above must be wrong, {} > {}", |
346 | | bytes_to_skip, |
347 | 0 | data.len() |
348 | | ); |
349 | 2.11k | let end_pos = cmp::min(data.len(), bytes_to_send + bytes_to_skip); |
350 | 2.11k | if bytes_to_skip != 0 || data1.38k .len() > bytes_to_send { Branch (350:16): [True: 728, False: 1.38k]
Branch (350:38): [True: 324, False: 1.06k]
Branch (350:16): [Folded - Ignored]
Branch (350:38): [Folded - Ignored]
|
351 | 1.05k | data = data.slice(bytes_to_skip..end_pos); |
352 | 1.06k | } |
353 | 2.11k | writer |
354 | 2.11k | .send(data) |
355 | 2.11k | .await |
356 | 2.11k | .err_tip(|| "Failed to write data to get_part dedup")?0 ; |
357 | 2.11k | bytes_to_send -= end_pos - bytes_to_skip; |
358 | 2.11k | bytes_to_skip = 0; |
359 | | } |
360 | | |
361 | | // Finish our stream by writing our EOF and shutdown the stream. |
362 | 904 | writer |
363 | 904 | .send_eof() |
364 | 904 | .err_tip(|| "Failed to write EOF out from get_part dedup")?0 ; |
365 | 904 | Ok(()) |
366 | 1.87k | } |
367 | | |
368 | 0 | fn inner_store(&self, _digest: Option<StoreKey>) -> &dyn StoreDriver { |
369 | 0 | self |
370 | 0 | } |
371 | | |
372 | 0 | fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { |
373 | 0 | self |
374 | 0 | } |
375 | | |
376 | 0 | fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> { |
377 | 0 | self |
378 | 0 | } |
379 | | } |
380 | | |
381 | | default_health_status_indicator!(DedupStore); |