Skip to content

Commit 75255fa

Browse files
author
tac0turtle
committed
fmt
1 parent 7f54729 commit 75255fa

File tree

4 files changed

+159
-140
lines changed

4 files changed

+159
-140
lines changed

client/crates/client/src/compression.rs

Lines changed: 70 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ const COMPRESSION_HEADER_SIZE: usize = 9;
1313
/// Compression flag for uncompressed data
1414
const FLAG_UNCOMPRESSED: u8 = 0x00;
1515

16-
/// Compression flag for zstd compressed data
16+
/// Compression flag for zstd compressed data
1717
const FLAG_ZSTD: u8 = 0x01;
1818

1919
/// Default zstd compression level
@@ -24,13 +24,13 @@ const DEFAULT_ZSTD_LEVEL: i32 = 3;
2424
pub enum CompressionError {
2525
#[error("invalid compression header")]
2626
InvalidHeader,
27-
27+
2828
#[error("invalid compression flag: {0}")]
2929
InvalidCompressionFlag(u8),
30-
30+
3131
#[error("decompression failed: {0}")]
3232
DecompressionFailed(String),
33-
33+
3434
#[error("zstd error: {0}")]
3535
ZstdError(#[from] io::Error),
3636
}
@@ -66,24 +66,22 @@ impl BlobCompressor {
6666
compression_level: DEFAULT_ZSTD_LEVEL,
6767
}
6868
}
69-
69+
7070
/// Create a new blob compressor with custom compression level
7171
pub fn with_level(compression_level: i32) -> Self {
72-
Self {
73-
compression_level,
74-
}
72+
Self { compression_level }
7573
}
76-
74+
7775
/// Compress a blob
7876
pub fn compress(&self, blob: &[u8]) -> Result<Bytes> {
7977
// For empty blobs, just add uncompressed header
8078
if blob.is_empty() {
8179
return Ok(self.add_compression_header(blob, FLAG_UNCOMPRESSED, 0));
8280
}
83-
81+
8482
// Try to compress with zstd
8583
let compressed = zstd::encode_all(blob, self.compression_level)?;
86-
84+
8785
// Check if compression is beneficial (at least 10% savings)
8886
let compression_ratio = compressed.len() as f64 / blob.len() as f64;
8987
if compression_ratio > 0.9 {
@@ -94,42 +92,42 @@ impl BlobCompressor {
9492
Ok(self.add_compression_header(&compressed, FLAG_ZSTD, blob.len() as u64))
9593
}
9694
}
97-
95+
9896
/// Decompress a blob
9997
pub fn decompress(&self, compressed_blob: &[u8]) -> Result<Bytes> {
10098
// Check if blob is too small to have a header
10199
if compressed_blob.len() < COMPRESSION_HEADER_SIZE {
102100
// Assume legacy uncompressed blob
103101
return Ok(Bytes::copy_from_slice(compressed_blob));
104102
}
105-
103+
106104
// Check the compression flag
107105
let flag = compressed_blob[0];
108-
106+
109107
// Handle invalid flags with legacy blob heuristics
110108
if flag != FLAG_UNCOMPRESSED && flag != FLAG_ZSTD {
111109
// This could be either a legacy blob or a corrupted header
112110
// Use heuristics to determine which
113-
114-
let original_size = u64::from_le_bytes(
115-
compressed_blob[1..9].try_into().unwrap_or([0; 8])
116-
);
117-
111+
112+
let original_size =
113+
u64::from_le_bytes(compressed_blob[1..9].try_into().unwrap_or([0; 8]));
114+
118115
// If flag is in printable ASCII range (32-126) and size is unreasonable,
119116
// it's likely a legacy text blob
120-
if (flag >= 32 && flag <= 126) &&
121-
(original_size == 0 || original_size > (compressed_blob.len() as u64 * 100)) {
117+
if (32..=126).contains(&flag)
118+
&& (original_size == 0 || original_size > (compressed_blob.len() as u64 * 100))
119+
{
122120
// Likely a legacy blob
123121
return Ok(Bytes::copy_from_slice(compressed_blob));
124122
}
125-
123+
126124
// Otherwise, it's likely a corrupted compressed blob
127125
return Err(CompressionError::InvalidCompressionFlag(flag));
128126
}
129-
127+
130128
// Parse the header
131129
let (flag, original_size, payload) = self.parse_compression_header(compressed_blob)?;
132-
130+
133131
match flag {
134132
FLAG_UNCOMPRESSED => {
135133
// Data is uncompressed, just return the payload
@@ -139,15 +137,16 @@ impl BlobCompressor {
139137
// Decompress with zstd
140138
let decompressed = zstd::decode_all(payload)
141139
.map_err(|e| CompressionError::DecompressionFailed(e.to_string()))?;
142-
140+
143141
// Verify the decompressed size matches
144142
if decompressed.len() as u64 != original_size {
145-
return Err(CompressionError::DecompressionFailed(
146-
format!("size mismatch: expected {}, got {}",
147-
original_size, decompressed.len())
148-
));
143+
return Err(CompressionError::DecompressionFailed(format!(
144+
"size mismatch: expected {}, got {}",
145+
original_size,
146+
decompressed.len()
147+
)));
149148
}
150-
149+
151150
Ok(Bytes::from(decompressed))
152151
}
153152
_ => {
@@ -156,7 +155,7 @@ impl BlobCompressor {
156155
}
157156
}
158157
}
159-
158+
160159
/// Get compression information about a blob
161160
pub fn get_compression_info(&self, blob: &[u8]) -> CompressionInfo {
162161
if blob.len() < COMPRESSION_HEADER_SIZE {
@@ -168,7 +167,7 @@ impl BlobCompressor {
168167
compression_ratio: 1.0,
169168
};
170169
}
171-
170+
172171
let flag = blob[0];
173172
if flag != FLAG_UNCOMPRESSED && flag != FLAG_ZSTD {
174173
// Legacy or invalid blob
@@ -180,14 +179,14 @@ impl BlobCompressor {
180179
compression_ratio: 1.0,
181180
};
182181
}
183-
182+
184183
if let Ok((flag, original_size, _)) = self.parse_compression_header(blob) {
185184
let algorithm = match flag {
186185
FLAG_UNCOMPRESSED => "none",
187186
FLAG_ZSTD => "zstd",
188187
_ => "unknown",
189188
};
190-
189+
191190
CompressionInfo {
192191
is_compressed: flag == FLAG_ZSTD,
193192
algorithm: algorithm.to_string(),
@@ -209,40 +208,42 @@ impl BlobCompressor {
209208
}
210209
}
211210
}
212-
211+
213212
/// Add compression header to payload
214213
fn add_compression_header(&self, payload: &[u8], flag: u8, original_size: u64) -> Bytes {
215214
let mut result = BytesMut::with_capacity(COMPRESSION_HEADER_SIZE + payload.len());
216-
215+
217216
// Write flag
218217
result.extend_from_slice(&[flag]);
219-
218+
220219
// Write original size (little-endian)
221220
result.extend_from_slice(&original_size.to_le_bytes());
222-
221+
223222
// Write payload
224223
result.extend_from_slice(payload);
225-
224+
226225
result.freeze()
227226
}
228-
227+
229228
/// Parse compression header from blob
230229
fn parse_compression_header<'a>(&self, blob: &'a [u8]) -> Result<(u8, u64, &'a [u8])> {
231230
if blob.len() < COMPRESSION_HEADER_SIZE {
232231
return Err(CompressionError::InvalidHeader);
233232
}
234-
233+
235234
let flag = blob[0];
236235
let original_size = u64::from_le_bytes(
237-
blob[1..9].try_into().map_err(|_| CompressionError::InvalidHeader)?
236+
blob[1..9]
237+
.try_into()
238+
.map_err(|_| CompressionError::InvalidHeader)?,
238239
);
239240
let payload = &blob[COMPRESSION_HEADER_SIZE..];
240-
241+
241242
// Validate the compression flag
242243
if flag != FLAG_UNCOMPRESSED && flag != FLAG_ZSTD {
243244
return Err(CompressionError::InvalidCompressionFlag(flag));
244245
}
245-
246+
246247
Ok((flag, original_size, payload))
247248
}
248249
}
@@ -271,101 +272,101 @@ pub fn get_compression_info(blob: &[u8]) -> CompressionInfo {
271272
#[cfg(test)]
272273
mod tests {
273274
use super::*;
274-
275+
275276
#[test]
276277
fn test_compress_decompress_roundtrip() {
277278
let compressor = BlobCompressor::new();
278-
279+
279280
// Test with compressible data
280281
let original = b"hello world ".repeat(100);
281282
let compressed = compressor.compress(&original).unwrap();
282283
let decompressed = compressor.decompress(&compressed).unwrap();
283-
284+
284285
assert_eq!(original, decompressed.as_ref());
285-
286+
286287
// Verify it was actually compressed
287288
let info = compressor.get_compression_info(&compressed);
288289
assert!(info.is_compressed);
289290
assert_eq!(info.algorithm, "zstd");
290291
assert!(info.compression_ratio < 0.5); // Should compress well
291292
}
292-
293+
293294
#[test]
294295
fn test_uncompressed_fallback() {
295296
let compressor = BlobCompressor::new();
296-
297+
297298
// Random data that won't compress well
298299
let mut random_data = vec![0u8; 100];
299-
for i in 0..100 {
300-
random_data[i] = (i * 7 + 13) as u8; // Pseudo-random
300+
for (i, item) in random_data.iter_mut().enumerate().take(100) {
301+
*item = (i * 7 + 13) as u8; // Pseudo-random
301302
}
302-
303+
303304
let compressed = compressor.compress(&random_data).unwrap();
304305
let decompressed = compressor.decompress(&compressed).unwrap();
305-
306+
306307
assert_eq!(random_data, decompressed.as_ref());
307-
308+
308309
// Verify it was stored uncompressed
309310
let info = compressor.get_compression_info(&compressed);
310311
assert!(!info.is_compressed);
311312
assert_eq!(info.algorithm, "none");
312313
}
313-
314+
314315
#[test]
315316
fn test_legacy_blob() {
316317
let compressor = BlobCompressor::new();
317-
318+
318319
// Test with legacy blob (no compression header)
319320
let legacy_blob = b"legacy data without header";
320-
321+
321322
// Should return as-is
322323
let decompressed = compressor.decompress(legacy_blob).unwrap();
323324
assert_eq!(legacy_blob, decompressed.as_ref());
324325
}
325-
326+
326327
#[test]
327328
fn test_invalid_compression_flag() {
328329
let compressor = BlobCompressor::new();
329-
330+
330331
// Create blob with invalid flag
331332
let mut invalid_blob = vec![0u8; COMPRESSION_HEADER_SIZE + 10];
332333
invalid_blob[0] = 0xFF; // Invalid flag
333-
334+
334335
// Should return error
335336
let result = compressor.decompress(&invalid_blob);
336337
assert!(result.is_err());
337-
338+
338339
match result.unwrap_err() {
339340
CompressionError::InvalidCompressionFlag(flag) => {
340341
assert_eq!(flag, 0xFF);
341342
}
342343
_ => panic!("Expected InvalidCompressionFlag error"),
343344
}
344345
}
345-
346+
346347
#[test]
347348
fn test_empty_blob() {
348349
let compressor = BlobCompressor::new();
349-
350+
350351
let empty = vec![];
351352
let compressed = compressor.compress(&empty).unwrap();
352353
let decompressed = compressor.decompress(&compressed).unwrap();
353-
354+
354355
assert_eq!(empty, decompressed.as_ref());
355356
}
356-
357+
357358
#[test]
358359
fn test_compression_info() {
359360
let compressor = BlobCompressor::new();
360-
361+
361362
let original = b"compress me ".repeat(100);
362363
let compressed = compressor.compress(&original).unwrap();
363-
364+
364365
let info = compressor.get_compression_info(&compressed);
365366
assert!(info.is_compressed);
366367
assert_eq!(info.algorithm, "zstd");
367368
assert_eq!(info.original_size, original.len() as u64);
368369
assert!(info.compression_ratio < 1.0);
369370
assert!(info.compression_ratio > 0.0);
370371
}
371-
}
372+
}

client/crates/client/src/lib.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,9 @@ pub mod store;
8787

8888
// Re-export main types for convenience
8989
pub use client::{Client, ClientBuilder};
90-
pub use compression::{BlobCompressor, CompressionInfo, compress_blob, decompress_blob, get_compression_info};
90+
pub use compression::{
91+
compress_blob, decompress_blob, get_compression_info, BlobCompressor, CompressionInfo,
92+
};
9193
pub use config::ConfigClient;
9294
pub use error::{ClientError, Result};
9395
pub use health::HealthClient;

0 commit comments

Comments
 (0)