/build/source/nativelink-store/src/compression_store.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (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 | | // See LICENSE file for details |
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 byteorder::{ByteOrder, LittleEndian}; |
21 | | use bytes::{Buf, BufMut, BytesMut}; |
22 | | use futures::future::FutureExt; |
23 | | use lz4_flex::block::{compress_into, decompress_into, get_maximum_output_size}; |
24 | | use nativelink_config::stores::CompressionSpec; |
25 | | use nativelink_error::{Code, Error, ResultExt, error_if, make_err}; |
26 | | use nativelink_metric::MetricsComponent; |
27 | | use nativelink_util::buf_channel::{ |
28 | | DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair, |
29 | | }; |
30 | | use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; |
31 | | use nativelink_util::spawn; |
32 | | use nativelink_util::store_trait::{ |
33 | | RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, |
34 | | }; |
35 | | use serde::{Deserialize, Serialize}; |
36 | | use wincode::{SchemaRead, SchemaWrite}; |
37 | | |
38 | | use crate::cas_utils::is_zero_digest; |
39 | | |
40 | | // In the event the bytestream format changes this number should be incremented to prevent |
41 | | // backwards compatibility issues. |
42 | | pub const CURRENT_STREAM_FORMAT_VERSION: u8 = 1; |
43 | | |
44 | | // Default block size that will be used to slice stream into. |
45 | | pub const DEFAULT_BLOCK_SIZE: u32 = 64 * 1024; |
46 | | |
47 | | const U32_SZ: u64 = size_of::<u8>() as u64; |
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 | | #[derive( |
120 | 0 | Serialize, Deserialize, PartialEq, Eq, Debug, Default, Copy, Clone, SchemaWrite, SchemaRead, |
121 | | )] |
122 | | pub struct Lz4Config { |
123 | | pub block_size: u32, |
124 | | } |
125 | | |
126 | 0 | #[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone, Copy, SchemaWrite, SchemaRead)] |
127 | | pub struct Header { |
128 | | pub version: u8, |
129 | | pub config: Lz4Config, |
130 | | pub upload_size: UploadSizeInfo, |
131 | | } |
132 | | |
133 | | #[derive( |
134 | 0 | Serialize, Deserialize, PartialEq, Eq, Debug, Default, Clone, Copy, SchemaWrite, SchemaRead, |
135 | | )] |
136 | | pub struct SliceIndex { |
137 | | pub position_from_prev_index: u32, |
138 | | } |
139 | | |
140 | 0 | #[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Default, SchemaWrite, SchemaRead)] |
141 | | pub struct Footer { |
142 | | pub indexes: Vec<SliceIndex>, |
143 | | pub index_count: u32, |
144 | | pub uncompressed_data_size: u64, |
145 | | pub config: Lz4Config, |
146 | | pub version: u8, |
147 | | } |
148 | | |
149 | | /// `lz4_flex::block::get_maximum_output_size()` way over estimates, so we use the |
150 | | /// one provided here: <https://github.com/torvalds/linux/blob/master/include/linux/lz4.h#L61> |
151 | | /// Local testing shows this gives quite accurate worst case given random input. |
152 | 7 | const fn lz4_compress_bound(input_size: u64) -> u64 { |
153 | 7 | input_size + (input_size / 255) + 16 |
154 | 7 | } |
155 | | |
156 | 1.26k | fn serialized_size<C, V>(value: &V, config: C) -> Result<u64, Error> |
157 | 1.26k | where |
158 | 1.26k | C: wincode::config::Config, |
159 | 1.26k | V: SchemaWrite<C, Src = V>, |
160 | | { |
161 | 1.26k | wincode::config::serialized_size(value, config).map_err(|e| {0 |
162 | 0 | make_err!( |
163 | 0 | Code::Internal, |
164 | | "Failed to calculate serialized size: {:?}", |
165 | | e |
166 | | ) |
167 | 0 | }) |
168 | 1.26k | } |
169 | | |
170 | | struct UploadState { |
171 | | header: Header, |
172 | | footer: Footer, |
173 | | max_output_size: u64, |
174 | | input_max_size: u64, |
175 | | } |
176 | | |
177 | | impl UploadState { |
178 | 7 | pub(crate) fn new( |
179 | 7 | store: &CompressionStore, |
180 | 7 | upload_size: UploadSizeInfo, |
181 | 7 | ) -> Result<Self, Error> { |
182 | 7 | let input_max_size = match upload_size { |
183 | 7 | UploadSizeInfo::MaxSize(sz1 ) | UploadSizeInfo::ExactSize(sz6 ) => sz, |
184 | | }; |
185 | | |
186 | 7 | let max_index_count = (input_max_size / u64::from(store.config.block_size)) + 1; |
187 | | |
188 | 7 | let header = Header { |
189 | 7 | version: CURRENT_STREAM_FORMAT_VERSION, |
190 | 7 | config: Lz4Config { |
191 | 7 | block_size: store.config.block_size, |
192 | 7 | }, |
193 | 7 | upload_size, |
194 | 7 | }; |
195 | 7 | let footer = Footer { |
196 | 7 | indexes: vec![ |
197 | 7 | SliceIndex::default(); |
198 | 7 | usize::try_from(max_index_count) |
199 | 7 | .err_tip(|| "Could not convert max_index_count to usize")?0 |
200 | | ], |
201 | 7 | index_count: u32::try_from(max_index_count) |
202 | 7 | .err_tip(|| "Could not convert max_index_count to u32")?0 , |
203 | | uncompressed_data_size: 0, // Updated later. |
204 | 7 | config: header.config, |
205 | | version: CURRENT_STREAM_FORMAT_VERSION, |
206 | | }; |
207 | | |
208 | | // This is more accurate of an estimate than what get_maximum_output_size calculates. |
209 | 7 | let max_block_size = lz4_compress_bound(u64::from(store.config.block_size)) + U32_SZ + 1; |
210 | | |
211 | 7 | let max_output_size = { |
212 | 7 | let header_size = serialized_size(&header, store.wincode_config)?0 ; |
213 | 7 | let max_content_size = max_block_size * max_index_count; |
214 | 7 | let max_footer_size = U32_SZ + 1 + serialized_size(&footer, store.wincode_config)?0 ; |
215 | 7 | header_size + max_content_size + max_footer_size |
216 | | }; |
217 | | |
218 | 7 | Ok(Self { |
219 | 7 | header, |
220 | 7 | footer, |
221 | 7 | max_output_size, |
222 | 7 | input_max_size, |
223 | 7 | }) |
224 | 7 | } |
225 | | } |
226 | | |
227 | | // Based off of an older Bincode config |
228 | | pub type WincodeConfig = wincode::config::Configuration< |
229 | | true, |
230 | | { wincode::config::PREALLOCATION_SIZE_LIMIT_DISABLED }, |
231 | | wincode::len::BincodeLen, |
232 | | wincode::int_encoding::LittleEndian, |
233 | | wincode::int_encoding::FixInt, |
234 | | u32, |
235 | | >; |
236 | | |
237 | | /// This store will compress data before sending it on to the inner store. |
238 | | /// Note: Currently using `get_part()` and trying to read part of the data will |
239 | | /// result in the entire contents being read from the inner store but will |
240 | | /// only send the contents requested. |
241 | | #[derive(MetricsComponent)] |
242 | | pub struct CompressionStore { |
243 | | #[metric(group = "inner_store")] |
244 | | inner_store: Store, |
245 | | config: nativelink_config::stores::Lz4Config, |
246 | | wincode_config: WincodeConfig, |
247 | | } |
248 | | |
249 | | impl core::fmt::Debug for CompressionStore { |
250 | 0 | fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { |
251 | 0 | f.debug_struct("CompressionStore") |
252 | 0 | .field("inner_store", &self.inner_store) |
253 | 0 | .field("config", &self.config) |
254 | 0 | .finish_non_exhaustive() |
255 | 0 | } |
256 | | } |
257 | | |
258 | | impl CompressionStore { |
259 | 13 | pub fn new(spec: &CompressionSpec, inner_store: Store) -> Result<Arc<Self>, Error> { |
260 | 13 | let lz4_config = match spec.compression_algorithm { |
261 | 13 | nativelink_config::stores::CompressionAlgorithm::Lz4(mut lz4_config) => { |
262 | 13 | if lz4_config.block_size == 0 { |
263 | 8 | lz4_config.block_size = DEFAULT_BLOCK_SIZE; |
264 | 8 | }5 |
265 | 13 | if lz4_config.max_decode_block_size == 0 { |
266 | 13 | lz4_config.max_decode_block_size = lz4_config.block_size; |
267 | 13 | }0 |
268 | 13 | lz4_config |
269 | | } |
270 | | }; |
271 | 13 | Ok(Arc::new(Self { |
272 | 13 | inner_store, |
273 | 13 | config: lz4_config, |
274 | 13 | wincode_config: WincodeConfig::new(), |
275 | 13 | })) |
276 | 13 | } |
277 | | } |
278 | | |
279 | | #[async_trait] |
280 | | impl StoreDriver for CompressionStore { |
281 | 0 | async fn post_init(self: Arc<Self>) -> Result<(), Error> { |
282 | | self.inner_store.clone().into_inner().post_init().await?; |
283 | | Ok(()) |
284 | 0 | } |
285 | | |
286 | | async fn has_with_results( |
287 | | self: Pin<&Self>, |
288 | | digests: &[StoreKey<'_>], |
289 | | results: &mut [Option<u64>], |
290 | 3 | ) -> Result<(), Error> { |
291 | 4 | if !digests.iter().any(|digest| is_zero_digest(digest.borrow())) { |
292 | | return self.inner_store.has_with_results(digests, results).await; |
293 | | } |
294 | | |
295 | | for (digest, result) in digests.iter().zip(results.iter_mut()) { |
296 | | if is_zero_digest(digest.borrow()) { |
297 | | *result = Some(0); |
298 | | } |
299 | | } |
300 | | |
301 | | let nonzero_digests = digests |
302 | | .iter() |
303 | 6 | .filter(|digest| !is_zero_digest(digest.borrow())) |
304 | | .map(StoreKey::borrow) |
305 | | .collect::<Vec<_>>(); |
306 | | if nonzero_digests.is_empty() { |
307 | | return Ok(()); |
308 | | } |
309 | | |
310 | | let mut nonzero_results = vec![None; nonzero_digests.len()]; |
311 | | self.inner_store |
312 | | .has_with_results(&nonzero_digests, &mut nonzero_results) |
313 | | .await?; |
314 | | for (result, inner_result) in digests |
315 | | .iter() |
316 | | .zip(results.iter_mut()) |
317 | 4 | .filter_map(|(digest, result)| (!is_zero_digest(digest.borrow())).then_some(result)) |
318 | | .zip(nonzero_results) |
319 | | { |
320 | | *result = inner_result; |
321 | | } |
322 | | Ok(()) |
323 | 3 | } |
324 | | |
325 | | async fn update( |
326 | | self: Pin<&Self>, |
327 | | key: StoreKey<'_>, |
328 | | mut reader: DropCloserReadHalf, |
329 | | upload_size: UploadSizeInfo, |
330 | 11 | ) -> Result<u64, Error> { |
331 | | if is_zero_digest(key.borrow()) { |
332 | 4 | return reader.recv().await.and_then(|should_be_empty| { |
333 | 4 | if should_be_empty.is_empty() { |
334 | 2 | Ok(0) |
335 | | } else { |
336 | 2 | Err(make_err!(Code::Internal, "Zero byte hash not empty")) |
337 | | } |
338 | 4 | }); |
339 | | } |
340 | | |
341 | | let mut output_state = UploadState::new(&self, upload_size)?; |
342 | | |
343 | | let (mut tx, rx) = make_buf_channel_pair(); |
344 | | |
345 | | let inner_store = self.inner_store.clone(); |
346 | | let key = key.into_owned(); |
347 | 7 | let update_fut = spawn!("compression_store_update_spawn", async move { |
348 | 7 | inner_store |
349 | 7 | .update( |
350 | 7 | key, |
351 | 7 | rx, |
352 | 7 | UploadSizeInfo::MaxSize(output_state.max_output_size), |
353 | 7 | ) |
354 | 7 | .await |
355 | 7 | .err_tip(|| "Inner store update in compression store failed") |
356 | 7 | }) |
357 | | .map( |
358 | 7 | |result| match result.err_tip(|| "Failed to run compression update spawn") { |
359 | 7 | Ok(inner_result) => { |
360 | 7 | inner_result.err_tip(|| "Compression underlying store update failed") |
361 | | } |
362 | 0 | Err(e) => Err(e), |
363 | 7 | }, |
364 | | ); |
365 | | |
366 | 7 | let write_fut = async move { |
367 | | { |
368 | | // Write Header. |
369 | 7 | let serialized_header = |
370 | 7 | wincode::config::serialize(&output_state.header, self.wincode_config).map_err( |
371 | 0 | |e| make_err!(Code::Internal, "Failed to serialize header : {:?}", e), |
372 | 0 | )?; |
373 | 7 | tx.send(serialized_header.into()) |
374 | 7 | .await |
375 | 7 | .err_tip(|| "Failed to write compression header on upload")?0 ; |
376 | | } |
377 | | |
378 | 7 | let mut received_amt = 0; |
379 | 7 | let mut index_count: u32 = 0; |
380 | 102 | for index in &mut output_state.footer.indexes7 { |
381 | 102 | let chunk = reader |
382 | 102 | .consume(Some(self.config.block_size as usize)) |
383 | 102 | .await |
384 | 102 | .err_tip(|| "Failed to read take in update in compression store")?0 ; |
385 | 102 | if chunk.is_empty() { |
386 | 6 | break; // EOF. |
387 | 96 | } |
388 | | |
389 | 96 | received_amt += u64::try_from(chunk.len()) |
390 | 96 | .err_tip(|| "Could not convert chunk.len() to u64")?0 ; |
391 | 0 | error_if!( |
392 | 96 | received_amt > output_state.input_max_size, |
393 | | "Got more data than stated in compression store upload request" |
394 | | ); |
395 | | |
396 | 96 | let max_output_size = get_maximum_output_size(self.config.block_size as usize); |
397 | 96 | let mut compressed_data_buf = BytesMut::with_capacity(max_output_size); |
398 | 96 | compressed_data_buf.put_u8(CHUNK_FRAME_TYPE); |
399 | 96 | compressed_data_buf.put_u32_le(0); // Filled later. |
400 | | |
401 | | // For efficiency reasons we do some raw slice manipulation so we can write directly |
402 | | // into our buffer instead of having to do another allocation. |
403 | 96 | let raw_compressed_data = unsafe { |
404 | 96 | core::slice::from_raw_parts_mut( |
405 | 96 | compressed_data_buf.chunk_mut().as_mut_ptr(), |
406 | 96 | max_output_size, |
407 | 96 | ) |
408 | | }; |
409 | | |
410 | 96 | let compressed_data_sz = compress_into(&chunk, raw_compressed_data) |
411 | 96 | .map_err(|e| make_err!0 (Code::Internal0 , "Compression error {:?}", e))?0 ; |
412 | 96 | unsafe { |
413 | 96 | compressed_data_buf.advance_mut(compressed_data_sz); |
414 | 96 | } |
415 | | |
416 | | // Now fill the size in our slice. |
417 | 96 | LittleEndian::write_u32( |
418 | 96 | &mut compressed_data_buf[1..5], |
419 | 96 | u32::try_from(compressed_data_sz).unwrap_or(u32::MAX), |
420 | | ); |
421 | | |
422 | | // Now send our chunk. |
423 | 96 | tx.send(compressed_data_buf.freeze()) |
424 | 96 | .await |
425 | 96 | .err_tip(|| "Failed to write chunk to inner store in compression store")?0 ; |
426 | | |
427 | 96 | index.position_from_prev_index = |
428 | 96 | u32::try_from(compressed_data_sz).unwrap_or(u32::MAX); |
429 | | |
430 | 96 | index_count += 1; |
431 | | } |
432 | | // Index 0 is actually a pointer to the second chunk. This is because we don't need |
433 | | // an index for the first item, since it starts at position `{header_len}`. |
434 | | // The code above causes us to create 1 more index than we actually need, so we |
435 | | // remove the last index from our vector here, because at this point we are always |
436 | | // one index too many. |
437 | | // Note: We need to be careful that if we don't have any data (zero bytes) it |
438 | | // doesn't go to -1. |
439 | 7 | index_count = index_count.saturating_sub(1); |
440 | 7 | output_state |
441 | 7 | .footer |
442 | 7 | .indexes |
443 | 7 | .resize(index_count as usize, SliceIndex::default()); |
444 | 7 | output_state.footer.index_count = |
445 | 7 | u32::try_from(output_state.footer.indexes.len()).unwrap_or(u32::MAX); |
446 | 7 | output_state.footer.uncompressed_data_size = received_amt; |
447 | | { |
448 | | // Write Footer. |
449 | 7 | let serialized_footer = |
450 | 7 | wincode::config::serialize(&output_state.footer, self.wincode_config).map_err( |
451 | 0 | |e| make_err!(Code::Internal, "Failed to serialize header : {:?}", e), |
452 | 0 | )?; |
453 | | |
454 | 7 | let mut footer = BytesMut::with_capacity(1 + 4 + serialized_footer.len()); |
455 | 7 | footer.put_u8(FOOTER_FRAME_TYPE); |
456 | 7 | footer.put_u32_le(u32::try_from(serialized_footer.len()).unwrap_or(u32::MAX)); |
457 | 7 | footer.extend_from_slice(&serialized_footer); |
458 | | |
459 | 7 | tx.send(footer.freeze()) |
460 | 7 | .await |
461 | 7 | .err_tip(|| "Failed to write footer to inner store in compression store")?0 ; |
462 | 7 | tx.send_eof() |
463 | 7 | .err_tip(|| "Failed writing EOF in compression store update")?0 ; |
464 | | } |
465 | | |
466 | 7 | Result::<u64, Error>::Ok(received_amt) |
467 | 7 | }; |
468 | | let (write_result, update_result) = tokio::join!(write_fut, update_fut); |
469 | | match (write_result, update_result) { |
470 | | (Ok(size), Ok(_)) => Ok(size), |
471 | | (Err(e), _) | (_, Err(e)) => Err(e), |
472 | | } |
473 | 11 | } |
474 | | |
475 | | async fn get_part( |
476 | | self: Pin<&Self>, |
477 | | key: StoreKey<'_>, |
478 | | writer: &mut DropCloserWriteHalf, |
479 | | offset: u64, |
480 | | length: Option<u64>, |
481 | 1.24k | ) -> Result<(), Error> { |
482 | | if is_zero_digest(key.borrow()) { |
483 | | writer |
484 | | .send_eof() |
485 | | .err_tip(|| "Failed to send zero EOF in filesystem store get_part")?; |
486 | | return Ok(()); |
487 | | } |
488 | | |
489 | | let (tx, mut rx) = make_buf_channel_pair(); |
490 | | |
491 | | let inner_store = self.inner_store.clone(); |
492 | | let key = key.into_owned(); |
493 | 1.24k | let get_part_fut = spawn!("compression_store_get_part_spawn", async move { |
494 | 1.24k | inner_store |
495 | 1.24k | .get_part(key, tx, 0, None) |
496 | 1.24k | .await |
497 | 1.24k | .err_tip(|| "Inner store get in compression store failed") |
498 | 1.24k | }) |
499 | | .map( |
500 | 1.24k | |result| match result.err_tip(|| "Failed to run compression get spawn") { |
501 | 1.24k | Ok(inner_result) => { |
502 | 1.24k | inner_result.err_tip(|| "Compression underlying store get failed") |
503 | | } |
504 | 0 | Err(e) => Err(e), |
505 | 1.24k | }, |
506 | | ); |
507 | 1.24k | let read_fut = async move { |
508 | 1.24k | let header = { |
509 | | // Read header. |
510 | | const EMPTY_HEADER: Header = Header { |
511 | | version: CURRENT_STREAM_FORMAT_VERSION, |
512 | | config: Lz4Config { block_size: 0 }, |
513 | | upload_size: UploadSizeInfo::ExactSize(0), |
514 | | }; |
515 | 1.24k | let header_size = serialized_size(&EMPTY_HEADER, self.wincode_config)?0 ; |
516 | 1.24k | let chunk = rx |
517 | 1.24k | .consume(Some(usize::try_from(header_size).unwrap_or(usize::MAX))) |
518 | 1.24k | .await |
519 | 1.24k | .err_tip(|| "Failed to read header in get_part compression store")?0 ; |
520 | 0 | error_if!( |
521 | 1.24k | chunk.len() as u64 != header_size, |
522 | | "Expected inner store to return the proper amount of data in compression store {} != {}", |
523 | 0 | chunk.len(), |
524 | | header_size, |
525 | | ); |
526 | | |
527 | 1.24k | wincode::config::deserialize::<Header, WincodeConfig>( |
528 | 1.24k | &chunk, |
529 | 1.24k | self.wincode_config, |
530 | | ) |
531 | 1.24k | .map_err(|e| make_err!0 (Code::Internal0 , "Failed to deserialize header : {:?}", e))?0 |
532 | | }; |
533 | | |
534 | 0 | error_if!( |
535 | 1.24k | header.version != CURRENT_STREAM_FORMAT_VERSION, |
536 | | "Expected header version to match in get compression, got {}, want {}", |
537 | | header.version, |
538 | | CURRENT_STREAM_FORMAT_VERSION |
539 | | ); |
540 | 0 | error_if!( |
541 | 1.24k | header.config.block_size > self.config.max_decode_block_size, |
542 | | "Block size is too large in compression, got {} > {}", |
543 | | header.config.block_size, |
544 | 0 | self.config.max_decode_block_size |
545 | | ); |
546 | | |
547 | 1.24k | let mut chunk = rx |
548 | 1.24k | .consume(Some(1 + 4)) |
549 | 1.24k | .await |
550 | 1.24k | .err_tip(|| "Failed to read init frame info in compression store")?0 ; |
551 | 0 | error_if!( |
552 | 1.24k | chunk.len() < 1 + 4, |
553 | | "Received EOF too early while reading init frame info in compression store" |
554 | | ); |
555 | | |
556 | 1.24k | let mut frame_type = chunk.get_u8(); |
557 | 1.24k | let mut frame_sz = chunk.get_u32_le(); |
558 | | |
559 | 1.24k | let mut uncompressed_data_sz: u64 = 0; |
560 | 1.24k | let mut remaining_bytes_to_send: u64 = length.unwrap_or(u64::MAX); |
561 | 1.24k | let mut chunks_count: u32 = 0; |
562 | 5.05k | while frame_type != FOOTER_FRAME_TYPE { |
563 | 0 | error_if!( |
564 | 3.81k | frame_type != CHUNK_FRAME_TYPE, |
565 | | "Expected frame to be BODY in compression store, got {} at {}", |
566 | | frame_type, |
567 | | chunks_count |
568 | | ); |
569 | | |
570 | 3.81k | let chunk = rx |
571 | 3.81k | .consume(Some(frame_sz as usize)) |
572 | 3.81k | .await |
573 | 3.81k | .err_tip(|| "Failed to read chunk in get_part compression store")?0 ; |
574 | 3.81k | if chunk.len() < frame_sz as usize { |
575 | 0 | return Err(make_err!( |
576 | 0 | Code::Internal, |
577 | 0 | "Got EOF earlier than expected. Maybe the data is not compressed or different format?" |
578 | 0 | )); |
579 | 3.81k | } |
580 | | { |
581 | 3.81k | let max_output_size = |
582 | 3.81k | get_maximum_output_size(header.config.block_size as usize); |
583 | 3.81k | let mut uncompressed_data = BytesMut::with_capacity(max_output_size); |
584 | | |
585 | | // For efficiency reasons we do some raw slice manipulation so we can write directly |
586 | | // into our buffer instead of having to do another allocation. |
587 | 3.81k | let raw_decompressed_data = unsafe { |
588 | 3.81k | core::slice::from_raw_parts_mut( |
589 | 3.81k | uncompressed_data.chunk_mut().as_mut_ptr(), |
590 | 3.81k | max_output_size, |
591 | 3.81k | ) |
592 | | }; |
593 | | |
594 | 3.81k | let uncompressed_chunk_sz = decompress_into(&chunk, raw_decompressed_data) |
595 | 3.81k | .map_err(|e| make_err!0 (Code::Internal0 , "Decompression error {:?}", e))?0 ; |
596 | 3.81k | unsafe { uncompressed_data.advance_mut(uncompressed_chunk_sz) }; |
597 | 3.81k | let new_uncompressed_data_sz = |
598 | 3.81k | uncompressed_data_sz + uncompressed_chunk_sz as u64; |
599 | 3.81k | if new_uncompressed_data_sz >= offset && remaining_bytes_to_send > 02.33k { |
600 | 1.98k | let start_pos = |
601 | 1.98k | usize::try_from(offset.saturating_sub(uncompressed_data_sz)) |
602 | 1.98k | .unwrap_or(usize::MAX); |
603 | 1.98k | let end_pos = cmp::min( |
604 | 1.98k | start_pos.saturating_add( |
605 | 1.98k | usize::try_from(remaining_bytes_to_send).unwrap_or(usize::MAX), |
606 | | ), |
607 | 1.98k | uncompressed_chunk_sz, |
608 | | ); |
609 | 1.98k | if end_pos != start_pos { |
610 | | // Make sure we don't send an EOF by accident. |
611 | 1.87k | writer |
612 | 1.87k | .send(uncompressed_data.freeze().slice(start_pos..end_pos)) |
613 | 1.87k | .await |
614 | 1.87k | .err_tip(|| "Failed sending chunk in compression store")?0 ; |
615 | 112 | } |
616 | 1.98k | remaining_bytes_to_send -= (end_pos - start_pos) as u64; |
617 | 1.82k | } |
618 | 3.81k | uncompressed_data_sz = new_uncompressed_data_sz; |
619 | | } |
620 | 3.81k | chunks_count += 1; |
621 | | |
622 | 3.81k | let mut chunk = rx |
623 | 3.81k | .consume(Some(1 + 4)) |
624 | 3.81k | .await |
625 | 3.81k | .err_tip(|| "Failed to read frame info in compression store")?0 ; |
626 | 0 | error_if!( |
627 | 3.81k | chunk.len() < 1 + 4, |
628 | | "Received EOF too early while reading frame info in compression store" |
629 | | ); |
630 | | |
631 | 3.81k | frame_type = chunk.get_u8(); |
632 | 3.81k | frame_sz = chunk.get_u32_le(); |
633 | | } |
634 | | // Index count will always be +1 (unless it is zero bytes long). |
635 | 1.24k | chunks_count = chunks_count.saturating_sub(1); |
636 | | { |
637 | | // Read and validate footer. |
638 | 1.24k | let chunk = rx |
639 | 1.24k | .consume(Some(frame_sz as usize)) |
640 | 1.24k | .await |
641 | 1.24k | .err_tip(|| "Failed to read chunk in get_part compression store")?0 ; |
642 | 0 | error_if!( |
643 | 1.24k | chunk.len() < frame_sz as usize, |
644 | | "Unexpected EOF when reading footer in compression store get_part" |
645 | | ); |
646 | | |
647 | 1.24k | let footer = wincode::config::deserialize::<Footer, WincodeConfig>( |
648 | 1.24k | &chunk, |
649 | 1.24k | self.wincode_config, |
650 | | ) |
651 | 1.24k | .map_err(|e| make_err!0 (Code::Internal0 , "Failed to deserialize footer : {:?}", e))?0 ; |
652 | | |
653 | 0 | error_if!( |
654 | 1.24k | header.version != footer.version, |
655 | | "Expected header and footer versions to match compression store get_part, {} != {}", |
656 | | header.version, |
657 | | footer.version |
658 | | ); |
659 | 0 | error_if!( |
660 | 1.24k | footer.indexes.len() != footer.index_count as usize, |
661 | | "Expected index counts to match in compression store footer in get_part, {} != {}", |
662 | 0 | footer.indexes.len(), |
663 | | footer.index_count |
664 | | ); |
665 | 0 | error_if!( |
666 | 1.24k | footer.index_count != chunks_count, |
667 | | concat!( |
668 | | "Expected index counts to match received chunks count ", |
669 | | "in compression store footer in get_part, {} != {}" |
670 | | ), |
671 | | footer.index_count, |
672 | | chunks_count |
673 | | ); |
674 | 0 | error_if!( |
675 | 1.24k | footer.uncompressed_data_size != uncompressed_data_sz, |
676 | | "Expected uncompressed data sizes to match in compression store footer in get_part, {} != {}", |
677 | | footer.uncompressed_data_size, |
678 | | uncompressed_data_sz |
679 | | ); |
680 | | } |
681 | | |
682 | 1.24k | writer |
683 | 1.24k | .send_eof() |
684 | 1.24k | .err_tip(|| "Failed to send eof in compression store write")?0 ; |
685 | 1.24k | Ok(()) |
686 | 1.24k | }; |
687 | | |
688 | | let (read_result, get_part_fut_result) = tokio::join!(read_fut, get_part_fut); |
689 | | if let Err(mut e) = read_result { |
690 | | // We may need to propagate the error from reading the data through first. |
691 | | if let Err(err) = get_part_fut_result { |
692 | | e = err.merge(e); |
693 | | } |
694 | | return Err(e); |
695 | | } |
696 | | Ok(()) |
697 | 1.24k | } |
698 | | |
699 | 0 | fn inner_store(&self, _digest: Option<StoreKey>) -> &dyn StoreDriver { |
700 | 0 | self |
701 | 0 | } |
702 | | |
703 | 0 | fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { |
704 | 0 | self |
705 | 0 | } |
706 | | |
707 | 0 | fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> { |
708 | 0 | self |
709 | 0 | } |
710 | | |
711 | 0 | fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> { |
712 | 0 | self.inner_store.register_remove_callback(callback) |
713 | 0 | } |
714 | | } |
715 | | |
716 | | default_health_status_indicator!(CompressionStore); |