diff --git a/Cargo.lock b/Cargo.lock index a504620..c6f38fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -196,6 +196,21 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + [[package]] name = "crypto-common" version = "0.1.7" @@ -272,6 +287,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hmac" version = "0.12.1" @@ -374,6 +395,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -388,12 +419,15 @@ checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "pg_stage_rs" -version = "0.1.1" +version = "0.1.2" dependencies = [ "chrono", "clap", + "crossbeam-channel", "flate2", "hmac", + "memchr", + "num_cpus", "rand", "regex", "serde", diff --git a/Cargo.toml b/Cargo.toml index 984f38b..27948e6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pg_stage_rs" -version = "0.1.1" +version = "0.1.2" edition = "2021" description = "PostgreSQL dump anonymizer - streaming obfuscation for plain and custom formats" @@ -9,7 +9,10 @@ clap = { version = "4", features = ["derive"] } serde = { version = "1", features = ["derive"] } serde_json = "1" flate2 = "1" -zstd = "0.13" +zstd = { version = "0.13", features = ["zstdmt"] } +memchr = "2" +crossbeam-channel = "0.5" +num_cpus = "1.16" uuid = { version = "1", features = ["v4", "v5"] } regex = "1" rand = "0.8" diff --git a/src/format/custom/blocks.rs b/src/format/custom/blocks.rs index 05ce249..c31df93 100644 --- a/src/format/custom/blocks.rs +++ b/src/format/custom/blocks.rs @@ -1,8 +1,9 @@ -use std::io::{Read, Write}; +use std::io::{self, Read, Write}; use flate2::read::ZlibDecoder; use flate2::write::ZlibEncoder; use flate2::Compression; +use memchr::memrchr; use zstd::stream::read::Decoder as ZstdDecoder; use zstd::stream::write::Encoder as ZstdEncoder; @@ -11,8 +12,71 @@ use crate::format::custom::header::CompressionMethod; use crate::format::custom::io::DumpIO; use crate::processor::DataProcessor; -const OUTPUT_CHUNK_SIZE: usize = 512 * 1024; // 512KB +const OUTPUT_CHUNK_SIZE: usize = 1024 * 1024; // 1MB for better throughput const MAX_CHUNK_SIZE: usize = 50 * 1024 * 1024; // 50MB +const READ_BUF_SIZE: usize = 2 * 1024 * 1024; // 2MB read buffer + +/// Streaming reader that reads chunks on-demand instead of loading entire block into memory. +/// This is critical for large tables (100M+ rows) where compressed blocks can be several GB. +struct ChunkReader<'a, R: Read> { + reader: &'a mut R, + dio: &'a DumpIO, + current_chunk: Vec, + chunk_pos: usize, + done: bool, +} + +impl<'a, R: Read> ChunkReader<'a, R> { + fn new(reader: &'a mut R, dio: &'a DumpIO) -> Self { + Self { + reader, + dio, + current_chunk: Vec::with_capacity(OUTPUT_CHUNK_SIZE), + chunk_pos: 0, + done: false, + } + } +} + +impl Read for ChunkReader<'_, R> { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + // If current chunk exhausted, read next one + if self.chunk_pos >= self.current_chunk.len() { + if self.done { + return Ok(0); + } + + // Read next chunk length + let chunk_len = self.dio.read_int(self.reader) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?; + + if chunk_len == 0 { + self.done = true; + return Ok(0); + } + + let len = chunk_len.unsigned_abs() as usize; + if len > MAX_CHUNK_SIZE { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("Chunk size {} exceeds maximum {}", len, MAX_CHUNK_SIZE), + )); + } + + // Read chunk data + self.current_chunk.resize(len, 0); + self.reader.read_exact(&mut self.current_chunk)?; + self.chunk_pos = 0; + } + + // Copy data from current chunk to output buffer + let available = self.current_chunk.len() - self.chunk_pos; + let to_copy = available.min(buf.len()); + buf[..to_copy].copy_from_slice(&self.current_chunk[self.chunk_pos..self.chunk_pos + to_copy]); + self.chunk_pos += to_copy; + Ok(to_copy) + } +} /// Processes data blocks in a custom format dump. pub struct BlockProcessor<'a> { @@ -116,7 +180,7 @@ impl<'a> BlockProcessor<'a> { }; // Split into complete lines + tail - match data.iter().rposition(|&b| b == b'\n') { + match memrchr(b'\n', &data) { Some(last_nl) => { line_tail = data[last_nl + 1..].to_vec(); self.process_complete_lines(&data[..=last_nl], &mut output_buf); @@ -150,44 +214,20 @@ impl<'a> BlockProcessor<'a> { } /// Streaming processing for zlib-compressed blocks. - /// Reads all compressed chunks (they form one zlib stream), then decompresses - /// and processes incrementally. + /// Uses ChunkReader for on-demand chunk reading to minimize memory usage. fn process_block_zlib( &mut self, reader: &mut R, writer: &mut W, ) -> Result<()> { - // Read all compressed chunks (small compared to decompressed size) - let mut raw = Vec::new(); - loop { - let chunk_len = self.dio.read_int(reader)?; - if chunk_len == 0 { - break; - } - - let len = chunk_len.unsigned_abs() as usize; - if len > MAX_CHUNK_SIZE { - return Err(PgStageError::InvalidFormat(format!( - "Chunk size {} exceeds maximum {}", - len, MAX_CHUNK_SIZE - ))); - } - - let mut buf = vec![0u8; len]; - reader.read_exact(&mut buf)?; - raw.extend_from_slice(&buf); - } - - if raw.is_empty() { - self.dio.write_int(writer, 0)?; - return Ok(()); - } + // Use streaming chunk reader instead of loading entire block into memory + let chunk_reader = ChunkReader::new(reader, self.dio); // Stream: decompress → process lines → compress → write chunks - let mut decoder = ZlibDecoder::new(raw.as_slice()); + let mut decoder = ZlibDecoder::new(chunk_reader); let mut encoder = ZlibEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), Compression::new(6)); - let mut read_buf = vec![0u8; OUTPUT_CHUNK_SIZE]; + let mut read_buf = vec![0u8; READ_BUF_SIZE]; let mut line_tail: Vec = Vec::new(); loop { @@ -204,7 +244,7 @@ impl<'a> BlockProcessor<'a> { line_tail.as_slice() }; - match data.iter().rposition(|&b| b == b'\n') { + match memrchr(b'\n', data) { Some(last_nl) => { let complete = &data[..=last_nl]; let tail = &data[last_nl + 1..]; @@ -250,45 +290,28 @@ impl<'a> BlockProcessor<'a> { } /// Streaming processing for zstd-compressed blocks. - /// Reads all compressed chunks, decompresses, processes, and recompresses with zstd. + /// Uses ChunkReader for on-demand chunk reading to minimize memory usage. fn process_block_zstd( &mut self, reader: &mut R, writer: &mut W, ) -> Result<()> { - // Read all compressed chunks - let mut raw = Vec::new(); - loop { - let chunk_len = self.dio.read_int(reader)?; - if chunk_len == 0 { - break; - } - - let len = chunk_len.unsigned_abs() as usize; - if len > MAX_CHUNK_SIZE { - return Err(PgStageError::InvalidFormat(format!( - "Chunk size {} exceeds maximum {}", - len, MAX_CHUNK_SIZE - ))); - } - - let mut buf = vec![0u8; len]; - reader.read_exact(&mut buf)?; - raw.extend_from_slice(&buf); - } - - if raw.is_empty() { - self.dio.write_int(writer, 0)?; - return Ok(()); - } + // Use streaming chunk reader instead of loading entire block into memory + let chunk_reader = ChunkReader::new(reader, self.dio); // Stream: decompress → process lines → compress → write chunks - let mut decoder = ZstdDecoder::new(raw.as_slice()) + // Use compression level 1 for speed (was 3) + let mut decoder = ZstdDecoder::new(chunk_reader) .map_err(|e| PgStageError::CompressionError(format!("Zstd decoder init failed: {}", e)))?; - let mut encoder = ZstdEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), 3) + + // Use multithread zstd compression for better performance on large data + let mut encoder = ZstdEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), 1) .map_err(|e| PgStageError::CompressionError(format!("Zstd encoder init failed: {}", e)))?; + // Enable multithreaded compression (0 = auto-detect CPU count) + encoder.multithread(0) + .map_err(|e| PgStageError::CompressionError(format!("Zstd multithread init failed: {}", e)))?; - let mut read_buf = vec![0u8; OUTPUT_CHUNK_SIZE]; + let mut read_buf = vec![0u8; READ_BUF_SIZE]; let mut line_tail: Vec = Vec::new(); loop { @@ -305,7 +328,7 @@ impl<'a> BlockProcessor<'a> { line_tail.as_slice() }; - match data.iter().rposition(|&b| b == b'\n') { + match memrchr(b'\n', data) { Some(last_nl) => { let complete = &data[..=last_nl]; let tail = &data[last_nl + 1..]; @@ -352,9 +375,10 @@ impl<'a> BlockProcessor<'a> { /// Process complete lines (each ending with \n) and append results to output buffer. fn process_complete_lines(&mut self, data: &[u8], output: &mut Vec) { + use memchr::memchr; let mut start = 0; while start < data.len() { - let end = data[start..].iter().position(|&b| b == b'\n') + let end = memchr(b'\n', &data[start..]) .map(|p| start + p) .unwrap_or(data.len()); @@ -372,9 +396,10 @@ impl<'a> BlockProcessor<'a> { /// Process complete lines and write results to a Write impl (encoder). fn process_complete_lines_to_writer(&mut self, data: &[u8], writer: &mut W) -> Result<()> { + use memchr::memchr; let mut start = 0; while start < data.len() { - let end = data[start..].iter().position(|&b| b == b'\n') + let end = memchr(b'\n', &data[start..]) .map(|p| start + p) .unwrap_or(data.len()); diff --git a/src/format/custom/io.rs b/src/format/custom/io.rs index 501e77a..bb68073 100644 --- a/src/format/custom/io.rs +++ b/src/format/custom/io.rs @@ -89,16 +89,15 @@ impl DumpIO { (0u8, val) }; - // Sign byte - writer.write_all(&[sign])?; - - // Magnitude bytes (little-endian) + // Write sign + magnitude in one syscall (max 9 bytes: 1 sign + 8 magnitude) + let mut buf = [0u8; 9]; + buf[0] = sign; let mut current = v_abs; - for _ in 0..self.int_size { - let byte = (current & 0xFF) as u8; - writer.write_all(&[byte])?; + for i in 0..self.int_size { + buf[1 + i] = (current & 0xFF) as u8; current >>= 8; } + writer.write_all(&buf[..1 + self.int_size])?; Ok(()) } diff --git a/src/format/custom/mod.rs b/src/format/custom/mod.rs index 2991ff2..959b723 100644 --- a/src/format/custom/mod.rs +++ b/src/format/custom/mod.rs @@ -30,8 +30,9 @@ impl CustomHandler { writer: W, initial_bytes: &[u8], ) -> Result<()> { - let mut reader = BufReader::with_capacity(65536, reader); - let mut writer = BufWriter::with_capacity(65536, writer); + // Large buffers for better throughput on big dumps (2MB each) + let mut reader = BufReader::with_capacity(2 * 1024 * 1024, reader); + let mut writer = BufWriter::with_capacity(2 * 1024 * 1024, writer); // Parse header (bypasses to output) let header = parse_header(&mut reader, &mut writer, initial_bytes)?; diff --git a/src/format/plain.rs b/src/format/plain.rs index 6f26c9f..aa42652 100644 --- a/src/format/plain.rs +++ b/src/format/plain.rs @@ -21,13 +21,14 @@ impl PlainHandler { writer: W, initial_bytes: &[u8], ) -> Result<()> { - let mut writer = BufWriter::with_capacity(65536, writer); + // Large buffers for better throughput (2MB each) + let mut writer = BufWriter::with_capacity(2 * 1024 * 1024, writer); let mut is_data = false; let mut comment_buf: Option = None; // If we have initial bytes, chain them with the reader let combined = std::io::Cursor::new(initial_bytes.to_vec()).chain(reader); - let buf_reader = BufReader::with_capacity(65536, combined); + let buf_reader = BufReader::with_capacity(2 * 1024 * 1024, combined); for line_result in buf_reader.lines() { let line = line_result?; diff --git a/src/mutator/contact.rs b/src/mutator/contact.rs index 114dbcd..21f974e 100644 --- a/src/mutator/contact.rs +++ b/src/mutator/contact.rs @@ -78,7 +78,7 @@ pub fn address(ctx: &mut MutationContext) -> Result { } pub fn deterministic_phone(ctx: &mut MutationContext) -> Result { - let current_value = ctx.current_value.clone(); + let current_value = ctx.current_value; let count = ctx .kwargs .get("obfuscated_numbers_count") diff --git a/src/mutator/mod.rs b/src/mutator/mod.rs index 0553eb9..b9c657a 100644 --- a/src/mutator/mod.rs +++ b/src/mutator/mod.rs @@ -18,7 +18,7 @@ use crate::unique::UniqueTracker; pub struct MutationContext<'a> { pub kwargs: &'a HashMap, - pub current_value: String, + pub current_value: &'a str, // Changed from String to &str to avoid allocation pub rng: &'a mut ThreadRng, pub unique_tracker: &'a mut UniqueTracker, pub locale: Locale, diff --git a/src/processor.rs b/src/processor.rs index 4996481..7a04ec4 100644 --- a/src/processor.rs +++ b/src/processor.rs @@ -179,7 +179,9 @@ impl DataProcessor { // Use Cow to avoid allocating Strings for unmodified columns let mut result_values: Vec> = values.iter().map(|&s| Cow::Borrowed(s)).collect(); - let mut obfuscated_values: HashMap = HashMap::new(); + + // Pre-allocate with expected capacity to reduce reallocations + let mut obfuscated_values: HashMap = HashMap::with_capacity(self.current_mutations.len()); // Iterate by index to avoid cloning sorted_columns for col_sort_idx in 0..self.sorted_columns.len() { @@ -195,8 +197,6 @@ impl DataProcessor { None => continue, }; - let current_value = result_values[col_idx].to_string(); - // Try each mutation spec in order for spec in specs.iter() { // Check conditions — borrow of result_values ends after this call (NLL) @@ -209,11 +209,11 @@ impl DataProcessor { let mut relation_found = false; for relation in &spec.relations { if let Some(&from_idx) = self.column_indices.get(&relation.from_column_name) { - let fk_value = result_values[from_idx].to_string(); + let fk_value = result_values[from_idx].as_ref(); if let Some(existing) = self.relation_tracker.lookup( &relation.table_name, &relation.to_column_name, - &fk_value, + fk_value, ) { let val = existing.clone(); obfuscated_values.insert(col_name.clone(), val.clone()); @@ -228,10 +228,13 @@ impl DataProcessor { } } + // Get current value as &str to avoid allocation + let current_value = result_values[col_idx].as_ref(); + // Dispatch mutation let mut ctx = MutationContext { kwargs: &spec.mutation_kwargs, - current_value: current_value.clone(), + current_value, rng: &mut self.rng, unique_tracker: &mut self.unique_tracker, locale: self.locale, @@ -245,11 +248,10 @@ impl DataProcessor { if !spec.relations.is_empty() { for relation in &spec.relations { if let Some(&from_idx) = self.column_indices.get(&relation.from_column_name) { - let fk_value = result_values[from_idx].to_string(); self.relation_tracker.store( &relation.table_name, &relation.to_column_name, - &fk_value, + result_values[from_idx].as_ref(), &new_val, ); }