/build/source/nativelink-store/src/compression_store.rs
Line | Count | Source (jump to first uncovered line) |
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 std::cmp; |
16 | | use std::pin::Pin; |
17 | | use std::sync::Arc; |
18 | | |
19 | | use async_trait::async_trait; |
20 | | use bincode::config::{FixintEncoding, WithOtherIntEncoding}; |
21 | | use bincode::{DefaultOptions, Options}; |
22 | | use byteorder::{ByteOrder, LittleEndian}; |
23 | | use bytes::{Buf, BufMut, BytesMut}; |
24 | | use futures::future::FutureExt; |
25 | | use lz4_flex::block::{compress_into, decompress_into, get_maximum_output_size}; |
26 | | use nativelink_error::{error_if, make_err, Code, Error, ResultExt}; |
27 | | use nativelink_metric::MetricsComponent; |
28 | | use nativelink_util::buf_channel::{ |
29 | | make_buf_channel_pair, DropCloserReadHalf, DropCloserWriteHalf, |
30 | | }; |
31 | | use nativelink_util::health_utils::{default_health_status_indicator, HealthStatusIndicator}; |
32 | | use nativelink_util::spawn; |
33 | | use nativelink_util::store_trait::{Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo}; |
34 | | use serde::{Deserialize, Serialize}; |
35 | | |
36 | | use crate::cas_utils::is_zero_digest; |
37 | | |
38 | | // In the event the bytestream format changes this number should be incremented to prevent |
39 | | // backwards compatibility issues. |
40 | | pub const CURRENT_STREAM_FORMAT_VERSION: u8 = 1; |
41 | | |
42 | | // Default block size that will be used to slice stream into. |
43 | | pub const DEFAULT_BLOCK_SIZE: u32 = 64 * 1024; |
44 | | |
45 | | const U32_SZ: u64 = std::mem::size_of::<u8>() as u64; |
46 | | |
47 | | type BincodeOptions = WithOtherIntEncoding<DefaultOptions, FixintEncoding>; |
48 | | |
49 | | // We use a custom frame format here because I wanted the ability in the future to: |
50 | | // * Read a random part of the data without needing to parse entire file. |
51 | | // * Compress the data on the fly without needing to know the exact input size. |
52 | | // |
53 | | // The frame format that LZ4 uses does not contain an index of where the different |
54 | | // blocks are located. This would mean in the event we only wanted the last byte of |
55 | | // a file, we'd need to seek to the header of each block to find where the next block |
56 | | // offset is until we got to the last block then decompress it. |
57 | | // |
58 | | // By using this custom frame format we keep the ability to have each block reference |
59 | | // the next block (for efficiency), but also after all blocks we have a footer frame |
60 | | // which contains an index of each block in the stream. This means that we can read a |
61 | | // fixed number of bytes of the file, which always contain the size of the index, then |
62 | | // with that size we know the exact size of the footer frame, which contains the entire |
63 | | // index. Once this footer frame is loaded, we could then do some math to figure out |
64 | | // which block each byte is in and the offset of each block within the compressed file. |
65 | | // |
66 | | // The frame format is as follows: |
67 | | // |----------------------------------HEADER-----------------------------------------| |
68 | | // | version(u8) | block_size (u32) | upload_size_type (u32) | upload_size (u32) | |
69 | | // |----------------------------------BLOCK------------------------------------------| |
70 | | // | frame_type(u8) 0x00 | compressed_data_size (u32) | ...DATA... | |
71 | | // | ...DATA... | |
72 | | // | [Possibly repeat block] | |
73 | | // |----------------------------------FOOTER-----------------------------------------| |
74 | | // | frame_type(u8) 0x01 | footer_size (u32) | index_count1 (u64) | |
75 | | // | ...[pos_from_prev_index (u32) - repeat for count {index_count*}]... | |
76 | | // | index_count2 (u32) | uncompressed_data_sz (u64) | block_size (u32) | |
77 | | // | version (u8) |------------------------------------------------------------------| |
78 | | // |---------------------------------------------------------------------------------| |
79 | | // |
80 | | // version - A constant number used to define what version of this format is being |
81 | | // used. Version in header and footer must match. |
82 | | // block_size - Size of each block uncompressed except for last block. This means that |
83 | | // every block uncompressed will be a constant size except last block may |
84 | | // be variable size. Block size in header and footer must match. |
85 | | // upload_size_type - Value of 0 = UploadSizeInfo::ExactSize, 1 = UploadSizeInfo::MaxSize. |
86 | | // This is for debug reasons only. |
87 | | // upload_size - The size of the data. WARNING: Do not rely on this being the uncompressed |
88 | | // payload size. It is a debug field and a "best guess" on how large the data |
89 | | // is. The header does not contain the upload data size. This value is the |
90 | | // value counter part to what the `upload_size_type` field. |
91 | | // frame_type - Type of each frame. 0 = BLOCK frame, 1 = FOOTER frame. Header frame will |
92 | | // always start with the first byte of the stream, so no magic number for it. |
93 | | // compressed_data_size - The size of this block. The bytes after this field should be read |
94 | | // in sequence to get all of the block's data in this block. |
95 | | // footer_size - Size of the footer for bytes after this field. |
96 | | // index_count1 - Number of items in the index. ({index_count1} * 4) represents the number |
97 | | // of bytes that should be read after this field in order to get all index |
98 | | // data. |
99 | | // pos_from_prev_index - Index of an individual block in the stream relative to the previous |
100 | | // block. This field may repeat {index_count} times. |
101 | | // index_count2 - Same as {index_count1} and should always be equal. This field might be |
102 | | // useful if you want to read a random byte from this stream, have random |
103 | | // access to it and know the size of the stream during reading, because |
104 | | // this field is always in the exact same place in the stream relative to |
105 | | // the last byte. |
106 | | // uncompressed_data_sz - Size of the original uncompressed data. |
107 | | // |
108 | | // Note: All fields fields little-endian. |
109 | | |
110 | | /// Number representing a chunk. |
111 | | pub const CHUNK_FRAME_TYPE: u8 = 0; |
112 | | |
113 | | /// Number representing the footer. |
114 | | pub const FOOTER_FRAME_TYPE: u8 = 1; |
115 | | |
116 | | /// This is a partial mirror of nativelink_config::stores::Lz4Config. |
117 | | /// We cannot use that natively here because it could cause our |
118 | | /// serialized format to change if we added more configs. |
119 | 2.45k | #[derive(Serialize, Deserialize, PartialEq, Debug, Default, Copy, Clone)] |
120 | | pub struct Lz4Config { |
121 | | pub block_size: u32, |
122 | | } |
123 | | |
124 | 1.22k | #[derive(Serialize, Deserialize, PartialEq, Debug)] |
125 | | pub struct Header { |
126 | | pub version: u8, |
127 | | pub config: Lz4Config, |
128 | | pub upload_size: UploadSizeInfo, |
129 | | } |
130 | | |
131 | 2.53k | #[derive(Serialize, Deserialize, PartialEq, Debug, Default, Clone, Copy)] |
132 | | pub struct SliceIndex { |
133 | | pub position_from_prev_index: u32, |
134 | | } |
135 | | |
136 | 1.23k | #[derive(Serialize, Deserialize, PartialEq, Debug, Default)] |
137 | | pub struct Footer { |
138 | | pub indexes: Vec<SliceIndex>, |
139 | | pub index_count: u32, |
140 | | pub uncompressed_data_size: u64, |
141 | | pub config: Lz4Config, |
142 | | pub version: u8, |
143 | | } |
144 | | |
145 | | /// lz4_flex::block::get_maximum_output_size() way over estimates, so we use the |
146 | | /// one provided here: https://github.com/torvalds/linux/blob/master/include/linux/lz4.h#L61 |
147 | | /// Local testing shows this gives quite accurate worst case given random input. |
148 | 6 | fn lz4_compress_bound(input_size: u64) -> u64 { |
149 | 6 | input_size + (input_size / 255) + 16 |
150 | 6 | } |
151 | | |
152 | | struct UploadState { |
153 | | header: Header, |
154 | | footer: Footer, |
155 | | max_output_size: u64, |
156 | | input_max_size: u64, |
157 | | } |
158 | | |
159 | | impl UploadState { |
160 | 6 | pub fn new(store: &CompressionStore, upload_size: UploadSizeInfo) -> Result<Self, Error> { |
161 | 6 | let input_max_size = match upload_size { |
162 | 5 | UploadSizeInfo::ExactSize(sz) => sz, |
163 | 1 | UploadSizeInfo::MaxSize(sz) => sz, |
164 | | }; |
165 | | |
166 | 6 | let max_index_count = (input_max_size / store.config.block_size as u64) + 1; |
167 | 6 | |
168 | 6 | let header = Header { |
169 | 6 | version: CURRENT_STREAM_FORMAT_VERSION, |
170 | 6 | config: Lz4Config { |
171 | 6 | block_size: store.config.block_size, |
172 | 6 | }, |
173 | 6 | upload_size, |
174 | 6 | }; |
175 | 6 | let footer = Footer { |
176 | 6 | indexes: vec![ |
177 | 6 | SliceIndex { |
178 | 6 | ..Default::default() |
179 | 6 | }; |
180 | 6 | usize::try_from(max_index_count) |
181 | 6 | .err_tip(|| "Could not convert max_index_count to usize"0 )?0 |
182 | | ], |
183 | 6 | index_count: max_index_count as u32, |
184 | 6 | uncompressed_data_size: 0, // Updated later. |
185 | 6 | config: header.config, |
186 | 6 | version: CURRENT_STREAM_FORMAT_VERSION, |
187 | 6 | }; |
188 | 6 | |
189 | 6 | // This is more accurate of an estimate than what get_maximum_output_size calculates. |
190 | 6 | let max_block_size = lz4_compress_bound(store.config.block_size as u64) + U32_SZ + 1; |
191 | 6 | |
192 | 6 | let max_output_size = { |
193 | 6 | let header_size = store.bincode_options.serialized_size(&header).unwrap(); |
194 | 6 | let max_content_size = max_block_size * max_index_count; |
195 | 6 | let max_footer_size = |
196 | 6 | U32_SZ + 1 + store.bincode_options.serialized_size(&footer).unwrap(); |
197 | 6 | header_size + max_content_size + max_footer_size |
198 | 6 | }; |
199 | 6 | |
200 | 6 | Ok(Self { |
201 | 6 | header, |
202 | 6 | footer, |
203 | 6 | max_output_size, |
204 | 6 | input_max_size, |
205 | 6 | }) |
206 | 6 | } |
207 | | } |
208 | | |
209 | | /// This store will compress data before sending it on to the inner store. |
210 | | /// Note: Currently using get_part() and trying to read part of the data will |
211 | | /// result in the entire contents being read from the inner store but will |
212 | | /// only send the contents requested. |
213 | 0 | #[derive(MetricsComponent)] |
214 | | pub struct CompressionStore { |
215 | | #[metric(group = "inner_store")] |
216 | | inner_store: Store, |
217 | | config: nativelink_config::stores::Lz4Config, |
218 | | bincode_options: BincodeOptions, |
219 | | } |
220 | | |
221 | | impl CompressionStore { |
222 | 7 | pub fn new( |
223 | 7 | compression_config: &nativelink_config::stores::CompressionStore, |
224 | 7 | inner_store: Store, |
225 | 7 | ) -> Result<Arc<Self>, Error> { |
226 | 7 | let lz4_config = match compression_config.compression_algorithm { |
227 | 7 | nativelink_config::stores::CompressionAlgorithm::lz4(mut lz4_config) => { |
228 | 7 | if lz4_config.block_size == 0 { Branch (228:20): [True: 3, False: 4]
Branch (228:20): [Folded - Ignored]
|
229 | 3 | lz4_config.block_size = DEFAULT_BLOCK_SIZE; |
230 | 4 | } |
231 | 7 | if lz4_config.max_decode_block_size == 0 { Branch (231:20): [True: 7, False: 0]
Branch (231:20): [Folded - Ignored]
|
232 | 7 | lz4_config.max_decode_block_size = lz4_config.block_size; |
233 | 7 | }0 |
234 | 7 | lz4_config |
235 | 7 | } |
236 | 7 | }; |
237 | 7 | Ok(Arc::new(CompressionStore { |
238 | 7 | inner_store, |
239 | 7 | config: lz4_config, |
240 | 7 | bincode_options: DefaultOptions::new().with_fixint_encoding(), |
241 | 7 | })) |
242 | 7 | } |
243 | | } |
244 | | |
245 | | #[async_trait] |
246 | | impl StoreDriver for CompressionStore { |
247 | | async fn has_with_results( |
248 | | self: Pin<&Self>, |
249 | | digests: &[StoreKey<'_>], |
250 | | results: &mut [Option<u64>], |
251 | 0 | ) -> Result<(), Error> { |
252 | 0 | self.inner_store.has_with_results(digests, results).await |
253 | 0 | } |
254 | | |
255 | | async fn update( |
256 | | self: Pin<&Self>, |
257 | | key: StoreKey<'_>, |
258 | | mut reader: DropCloserReadHalf, |
259 | | upload_size: UploadSizeInfo, |
260 | 6 | ) -> Result<(), Error> { |
261 | 6 | let mut output_state = UploadState::new(&self, upload_size)?0 ; |
262 | | |
263 | 6 | let (mut tx, rx) = make_buf_channel_pair(); |
264 | 6 | |
265 | 6 | let inner_store = self.inner_store.clone(); |
266 | 6 | let key = key.into_owned(); |
267 | 6 | let update_fut = spawn!("compression_store_update_spawn", async move { |
268 | 6 | inner_store |
269 | 6 | .update( |
270 | 6 | key, |
271 | 6 | rx, |
272 | 6 | UploadSizeInfo::MaxSize(output_state.max_output_size), |
273 | 6 | ) |
274 | 48 | .await |
275 | 6 | .err_tip(|| "Inner store update in compression store failed"0 ) |
276 | 6 | }) |
277 | 6 | .map( |
278 | 6 | |result| match result.err_tip(|| "Failed to run compression update spawn"0 ) { |
279 | 6 | Ok(inner_result) => { |
280 | 6 | inner_result.err_tip(|| "Compression underlying store update failed"0 ) |
281 | | } |
282 | 0 | Err(e) => Err(e), |
283 | 6 | }, |
284 | 6 | ); |
285 | 6 | |
286 | 6 | let write_fut = async move { |
287 | | { |
288 | | // Write Header. |
289 | 6 | let serialized_header = self |
290 | 6 | .bincode_options |
291 | 6 | .serialize(&output_state.header) |
292 | 6 | .map_err(|e| { |
293 | 0 | make_err!(Code::Internal, "Failed to serialize header : {:?}", e) |
294 | 6 | })?0 ; |
295 | 6 | tx.send(serialized_header.into()) |
296 | 0 | .await |
297 | 6 | .err_tip(|| "Failed to write compression header on upload"0 )?0 ; |
298 | | } |
299 | | |
300 | 6 | let mut received_amt = 0; |
301 | 6 | let mut index_count: u32 = 0; |
302 | 99 | for index98 in &mut output_state.footer.indexes { |
303 | 98 | let chunk = reader |
304 | 98 | .consume(Some(self.config.block_size as usize)) |
305 | 0 | .await |
306 | 98 | .err_tip(|| "Failed to read take in update in compression store"0 )?0 ; |
307 | 98 | if chunk.is_empty() { Branch (307:20): [True: 5, False: 93]
Branch (307:20): [Folded - Ignored]
|
308 | 5 | break; // EOF. |
309 | 93 | } |
310 | 93 | |
311 | 93 | received_amt += u64::try_from(chunk.len()) |
312 | 93 | .err_tip(|| "Could not convert chunk.len() to u64"0 )?0 ; |
313 | 0 | error_if!( |
314 | 93 | received_amt > output_state.input_max_size, Branch (314:21): [True: 0, False: 93]
Branch (314:21): [Folded - Ignored]
|
315 | | "Got more data than stated in compression store upload request" |
316 | | ); |
317 | | |
318 | 93 | let max_output_size = get_maximum_output_size(self.config.block_size as usize); |
319 | 93 | let mut compressed_data_buf = BytesMut::with_capacity(max_output_size); |
320 | 93 | compressed_data_buf.put_u8(CHUNK_FRAME_TYPE); |
321 | 93 | compressed_data_buf.put_u32_le(0); // Filled later. |
322 | 93 | |
323 | 93 | // For efficiency reasons we do some raw slice manipulation so we can write directly |
324 | 93 | // into our buffer instead of having to do another allocation. |
325 | 93 | let raw_compressed_data = unsafe { |
326 | 93 | std::slice::from_raw_parts_mut( |
327 | 93 | compressed_data_buf.chunk_mut().as_mut_ptr(), |
328 | 93 | max_output_size, |
329 | 93 | ) |
330 | | }; |
331 | | |
332 | 93 | let compressed_data_sz = compress_into(&chunk, raw_compressed_data) |
333 | 93 | .map_err(|e| make_err!(Code::Internal, "Compression error {:?}", e)0 )?0 ; |
334 | 93 | unsafe { |
335 | 93 | compressed_data_buf.advance_mut(compressed_data_sz); |
336 | 93 | } |
337 | 93 | |
338 | 93 | // Now fill the size in our slice. |
339 | 93 | LittleEndian::write_u32(&mut compressed_data_buf[1..5], compressed_data_sz as u32); |
340 | 93 | |
341 | 93 | // Now send our chunk. |
342 | 93 | tx.send(compressed_data_buf.freeze()) |
343 | 45 | .await |
344 | 93 | .err_tip(|| "Failed to write chunk to inner store in compression store"0 )?0 ; |
345 | | |
346 | 93 | index.position_from_prev_index = compressed_data_sz as u32; |
347 | 93 | |
348 | 93 | index_count += 1; |
349 | | } |
350 | | // Index 0 is actually a pointer to the second chunk. This is because we don't need |
351 | | // an index for the first item, since it starts at position `{header_len}`. |
352 | | // The code above causes us to create 1 more index than we actually need, so we |
353 | | // remove the last index from our vector here, because at this point we are always |
354 | | // one index too many. |
355 | | // Note: We need to be careful that if we don't have any data (zero bytes) it |
356 | | // doesn't go to -1. |
357 | 6 | index_count = index_count.saturating_sub(1); |
358 | 6 | output_state.footer.indexes.resize( |
359 | 6 | index_count as usize, |
360 | 6 | SliceIndex { |
361 | 6 | ..Default::default() |
362 | 6 | }, |
363 | 6 | ); |
364 | 6 | output_state.footer.index_count = output_state.footer.indexes.len() as u32; |
365 | 6 | output_state.footer.uncompressed_data_size = received_amt; |
366 | | { |
367 | | // Write Footer. |
368 | 6 | let serialized_footer = self |
369 | 6 | .bincode_options |
370 | 6 | .serialize(&output_state.footer) |
371 | 6 | .map_err(|e| { |
372 | 0 | make_err!(Code::Internal, "Failed to serialize header : {:?}", e) |
373 | 6 | })?0 ; |
374 | | |
375 | 6 | let mut footer = BytesMut::with_capacity(1 + 4 + serialized_footer.len()); |
376 | 6 | footer.put_u8(FOOTER_FRAME_TYPE); |
377 | 6 | footer.put_u32_le(serialized_footer.len() as u32); |
378 | 6 | footer.extend_from_slice(&serialized_footer); |
379 | 6 | |
380 | 6 | tx.send(footer.freeze()) |
381 | 3 | .await |
382 | 6 | .err_tip(|| "Failed to write footer to inner store in compression store"0 )?0 ; |
383 | 6 | tx.send_eof() |
384 | 6 | .err_tip(|| "Failed writing EOF in compression store update"0 )?0 ; |
385 | | } |
386 | | |
387 | 6 | Result::<(), Error>::Ok(()) |
388 | 6 | }; |
389 | 6 | let (write_result, update_result) = tokio::join!(write_fut, update_fut); |
390 | 6 | write_result.merge(update_result) |
391 | 12 | } |
392 | | |
393 | | async fn get_part( |
394 | | self: Pin<&Self>, |
395 | | key: StoreKey<'_>, |
396 | | writer: &mut DropCloserWriteHalf, |
397 | | offset: u64, |
398 | | length: Option<u64>, |
399 | 1.22k | ) -> Result<(), Error> { |
400 | 1.22k | if is_zero_digest(key.borrow()) { Branch (400:12): [True: 1, False: 1.22k]
Branch (400:12): [Folded - Ignored]
|
401 | 1 | writer |
402 | 1 | .send_eof() |
403 | 1 | .err_tip(|| "Failed to send zero EOF in filesystem store get_part"0 )?0 ; |
404 | 1 | return Ok(()); |
405 | 1.22k | } |
406 | 1.22k | |
407 | 1.22k | let (tx, mut rx) = make_buf_channel_pair(); |
408 | 1.22k | |
409 | 1.22k | let inner_store = self.inner_store.clone(); |
410 | 1.22k | let key = key.into_owned(); |
411 | 1.22k | let get_part_fut = spawn!("compression_store_get_part_spawn", async move { |
412 | 1.22k | inner_store |
413 | 1.22k | .get_part(key, tx, 0, None) |
414 | 0 | .await |
415 | 1.22k | .err_tip(|| "Inner store get in compression store failed"0 ) |
416 | 1.22k | }) |
417 | 1.22k | .map( |
418 | 1.22k | |result| match result.err_tip(|| "Failed to run compression get spawn"0 ) { |
419 | 1.22k | Ok(inner_result) => { |
420 | 1.22k | inner_result.err_tip(|| "Compression underlying store get failed"0 ) |
421 | | } |
422 | 0 | Err(e) => Err(e), |
423 | 1.22k | }, |
424 | 1.22k | ); |
425 | 1.22k | let read_fut = async move { |
426 | 1.22k | let header = { |
427 | | // Read header. |
428 | | static EMPTY_HEADER: Header = Header { |
429 | | version: CURRENT_STREAM_FORMAT_VERSION, |
430 | | config: Lz4Config { block_size: 0 }, |
431 | | upload_size: UploadSizeInfo::ExactSize(0), |
432 | | }; |
433 | 1.22k | let header_size = self.bincode_options.serialized_size(&EMPTY_HEADER).unwrap(); |
434 | 1.22k | let chunk = rx |
435 | 1.22k | .consume(Some(header_size as usize)) |
436 | 1.22k | .await |
437 | 1.22k | .err_tip(|| "Failed to read header in get_part compression store"0 )?0 ; |
438 | 0 | error_if!( |
439 | 1.22k | chunk.len() as u64 != header_size, Branch (439:21): [True: 0, False: 1.22k]
Branch (439:21): [Folded - Ignored]
|
440 | | "Expected inner store to return the proper amount of data in compression store {} != {}", |
441 | 0 | chunk.len(), |
442 | | header_size, |
443 | | ); |
444 | | |
445 | 1.22k | self.bincode_options |
446 | 1.22k | .deserialize::<Header>(&chunk) |
447 | 1.22k | .map_err(|e| { |
448 | 0 | make_err!(Code::Internal, "Failed to deserialize header : {:?}", e) |
449 | 1.22k | })?0 |
450 | | }; |
451 | | |
452 | 0 | error_if!( |
453 | 1.22k | header.version != CURRENT_STREAM_FORMAT_VERSION, Branch (453:17): [True: 0, False: 1.22k]
Branch (453:17): [Folded - Ignored]
|
454 | | "Expected header version to match in get compression, got {}, want {}", |
455 | | header.version, |
456 | | CURRENT_STREAM_FORMAT_VERSION |
457 | | ); |
458 | 0 | error_if!( |
459 | 1.22k | header.config.block_size > self.config.max_decode_block_size, Branch (459:17): [True: 0, False: 1.22k]
Branch (459:17): [Folded - Ignored]
|
460 | | "Block size is too large in compression, got {} > {}", |
461 | | header.config.block_size, |
462 | 0 | self.config.max_decode_block_size |
463 | | ); |
464 | | |
465 | 1.22k | let mut chunk = rx |
466 | 1.22k | .consume(Some(1 + 4)) |
467 | 0 | .await |
468 | 1.22k | .err_tip(|| "Failed to read init frame info in compression store"0 )?0 ; |
469 | 0 | error_if!( |
470 | 1.22k | chunk.len() < 1 + 4, Branch (470:17): [True: 0, False: 1.22k]
Branch (470:17): [Folded - Ignored]
|
471 | | "Received EOF too early while reading init frame info in compression store" |
472 | | ); |
473 | | |
474 | 1.22k | let mut frame_type = chunk.get_u8(); |
475 | 1.22k | let mut frame_sz = chunk.get_u32_le(); |
476 | 1.22k | |
477 | 1.22k | let mut uncompressed_data_sz: u64 = 0; |
478 | 1.22k | let mut remaining_bytes_to_send: u64 = length.unwrap_or(u64::MAX); |
479 | 1.22k | let mut chunks_count: u32 = 0; |
480 | 4.98k | while frame_type != FOOTER_FRAME_TYPE { Branch (480:19): [True: 3.75k, False: 1.22k]
Branch (480:19): [Folded - Ignored]
|
481 | 0 | error_if!( |
482 | 3.75k | frame_type != CHUNK_FRAME_TYPE, Branch (482:21): [True: 0, False: 3.75k]
Branch (482:21): [Folded - Ignored]
|
483 | | "Expected frame to be BODY in compression store, got {} at {}", |
484 | | frame_type, |
485 | | chunks_count |
486 | | ); |
487 | | |
488 | 3.75k | let chunk = rx |
489 | 3.75k | .consume(Some(frame_sz as usize)) |
490 | 0 | .await |
491 | 3.75k | .err_tip(|| "Failed to read chunk in get_part compression store"0 )?0 ; |
492 | 3.75k | if chunk.len() < frame_sz as usize { Branch (492:20): [True: 0, False: 3.75k]
Branch (492:20): [Folded - Ignored]
|
493 | 0 | return Err(make_err!( |
494 | 0 | Code::Internal, |
495 | 0 | "Got EOF earlier than expected. Maybe the data is not compressed or different format?" |
496 | 0 | )); |
497 | 3.75k | } |
498 | 3.75k | { |
499 | 3.75k | let max_output_size = |
500 | 3.75k | get_maximum_output_size(header.config.block_size as usize); |
501 | 3.75k | let mut uncompressed_data = BytesMut::with_capacity(max_output_size); |
502 | 3.75k | |
503 | 3.75k | // For efficiency reasons we do some raw slice manipulation so we can write directly |
504 | 3.75k | // into our buffer instead of having to do another allocation. |
505 | 3.75k | let raw_decompressed_data = unsafe { |
506 | 3.75k | std::slice::from_raw_parts_mut( |
507 | 3.75k | uncompressed_data.chunk_mut().as_mut_ptr(), |
508 | 3.75k | max_output_size, |
509 | 3.75k | ) |
510 | | }; |
511 | | |
512 | 3.75k | let uncompressed_chunk_sz = decompress_into(&chunk, raw_decompressed_data) |
513 | 3.75k | .map_err(|e| make_err!(Code::Internal, "Decompression error {:?}", e)0 )?0 ; |
514 | 3.75k | unsafe { uncompressed_data.advance_mut(uncompressed_chunk_sz) }; |
515 | 3.75k | let new_uncompressed_data_sz = |
516 | 3.75k | uncompressed_data_sz + uncompressed_chunk_sz as u64; |
517 | 3.75k | if new_uncompressed_data_sz >= offset && remaining_bytes_to_send > 02.28k { Branch (517:24): [True: 2.28k, False: 1.47k]
Branch (517:62): [True: 1.95k, False: 328]
Branch (517:24): [Folded - Ignored]
Branch (517:62): [Folded - Ignored]
|
518 | 1.95k | let start_pos = if offset <= uncompressed_data_sz { Branch (518:44): [True: 938, False: 1.02k]
Branch (518:44): [Folded - Ignored]
|
519 | 938 | 0 |
520 | | } else { |
521 | 1.02k | offset - uncompressed_data_sz |
522 | | } as usize; |
523 | 1.95k | let end_pos = cmp::min( |
524 | 1.95k | start_pos + remaining_bytes_to_send as usize, |
525 | 1.95k | uncompressed_chunk_sz, |
526 | 1.95k | ); |
527 | 1.95k | if end_pos != start_pos { Branch (527:28): [True: 1.85k, False: 102]
Branch (527:28): [Folded - Ignored]
|
528 | | // Make sure we don't send an EOF by accident. |
529 | 1.85k | writer |
530 | 1.85k | .send(uncompressed_data.freeze().slice(start_pos..end_pos)) |
531 | 224 | .await |
532 | 1.85k | .err_tip(|| "Failed sending chunk in compression store"0 )?0 ; |
533 | 102 | } |
534 | 1.95k | remaining_bytes_to_send -= (end_pos - start_pos) as u64; |
535 | 1.79k | } |
536 | 3.75k | uncompressed_data_sz = new_uncompressed_data_sz; |
537 | 3.75k | } |
538 | 3.75k | chunks_count += 1; |
539 | | |
540 | 3.75k | let mut chunk = rx |
541 | 3.75k | .consume(Some(1 + 4)) |
542 | 0 | .await |
543 | 3.75k | .err_tip(|| "Failed to read frame info in compression store"0 )?0 ; |
544 | 0 | error_if!( |
545 | 3.75k | chunk.len() < 1 + 4, Branch (545:21): [True: 0, False: 3.75k]
Branch (545:21): [Folded - Ignored]
|
546 | | "Received EOF too early while reading frame info in compression store" |
547 | | ); |
548 | | |
549 | 3.75k | frame_type = chunk.get_u8(); |
550 | 3.75k | frame_sz = chunk.get_u32_le(); |
551 | | } |
552 | | // Index count will always be +1 (unless it is zero bytes long). |
553 | 1.22k | chunks_count = chunks_count.saturating_sub(1); |
554 | | { |
555 | | // Read and validate footer. |
556 | 1.22k | let chunk = rx |
557 | 1.22k | .consume(Some(frame_sz as usize)) |
558 | 0 | .await |
559 | 1.22k | .err_tip(|| "Failed to read chunk in get_part compression store"0 )?0 ; |
560 | 0 | error_if!( |
561 | 1.22k | chunk.len() < frame_sz as usize, Branch (561:21): [True: 0, False: 1.22k]
Branch (561:21): [Folded - Ignored]
|
562 | | "Unexpected EOF when reading footer in compression store get_part" |
563 | | ); |
564 | | |
565 | 1.22k | let footer = self |
566 | 1.22k | .bincode_options |
567 | 1.22k | .deserialize::<Footer>(&chunk) |
568 | 1.22k | .map_err(|e| { |
569 | 0 | make_err!(Code::Internal, "Failed to deserialize footer : {:?}", e) |
570 | 1.22k | })?0 ; |
571 | | |
572 | 0 | error_if!( |
573 | 1.22k | header.version != footer.version, Branch (573:21): [True: 0, False: 1.22k]
Branch (573:21): [Folded - Ignored]
|
574 | | "Expected header and footer versions to match compression store get_part, {} != {}", |
575 | | header.version, |
576 | | footer.version |
577 | | ); |
578 | 0 | error_if!( |
579 | 1.22k | footer.indexes.len() != footer.index_count as usize, Branch (579:21): [True: 0, False: 1.22k]
Branch (579:21): [Folded - Ignored]
|
580 | | "Expected index counts to match in compression store footer in get_part, {} != {}", |
581 | 0 | footer.indexes.len(), |
582 | | footer.index_count |
583 | | ); |
584 | 0 | error_if!( |
585 | 1.22k | footer.index_count != chunks_count, Branch (585:21): [True: 0, False: 1.22k]
Branch (585:21): [Folded - Ignored]
|
586 | | concat!( |
587 | | "Expected index counts to match received chunks count ", |
588 | | "in compression store footer in get_part, {} != {}" |
589 | | ), |
590 | | footer.index_count, |
591 | | chunks_count |
592 | | ); |
593 | 0 | error_if!( |
594 | 1.22k | footer.uncompressed_data_size != uncompressed_data_sz, Branch (594:21): [True: 0, False: 1.22k]
Branch (594:21): [Folded - Ignored]
|
595 | | "Expected uncompressed data sizes to match in compression store footer in get_part, {} != {}", |
596 | | footer.uncompressed_data_size, |
597 | | uncompressed_data_sz |
598 | | ); |
599 | | } |
600 | | |
601 | 1.22k | writer |
602 | 1.22k | .send_eof() |
603 | 1.22k | .err_tip(|| "Failed to send eof in compression store write"0 )?0 ; |
604 | 1.22k | Ok(()) |
605 | 1.22k | }; |
606 | | |
607 | 1.22k | let (read_result, get_part_fut_result) = tokio::join!(read_fut, get_part_fut); |
608 | 1.22k | if let Err(mut e0 ) = read_result { Branch (608:16): [True: 0, False: 1.22k]
Branch (608:16): [Folded - Ignored]
|
609 | | // We may need to propagate the error from reading the data through first. |
610 | 0 | if let Err(err) = get_part_fut_result { Branch (610:20): [True: 0, False: 0]
Branch (610:20): [Folded - Ignored]
|
611 | 0 | e = err.merge(e); |
612 | 0 | } |
613 | 0 | return Err(e); |
614 | 1.22k | } |
615 | 1.22k | Ok(()) |
616 | 2.45k | } |
617 | | |
618 | 0 | fn inner_store(&self, _digest: Option<StoreKey>) -> &dyn StoreDriver { |
619 | 0 | self |
620 | 0 | } |
621 | | |
622 | 0 | fn as_any(&self) -> &(dyn std::any::Any + Sync + Send + 'static) { |
623 | 0 | self |
624 | 0 | } |
625 | | |
626 | 0 | fn as_any_arc(self: Arc<Self>) -> Arc<dyn std::any::Any + Sync + Send + 'static> { |
627 | 0 | self |
628 | 0 | } |
629 | | } |
630 | | |
631 | | default_health_status_indicator!(CompressionStore); |