From 4ce02628f826b2ae6c0429d317365e7e0bb7a3f8 Mon Sep 17 00:00:00 2001 From: Aleksandr Kharitonov Date: Thu, 29 Jan 2026 11:04:19 +0300 Subject: [PATCH] Revert "Performance optimizations for large databases" --- Cargo.lock | 36 +-------- Cargo.toml | 7 +- src/format/custom/blocks.rs | 155 +++++++++++++++--------------------- src/format/custom/io.rs | 13 +-- src/format/custom/mod.rs | 5 +- src/format/plain.rs | 5 +- src/mutator/contact.rs | 2 +- src/mutator/mod.rs | 2 +- src/processor.rs | 18 ++--- 9 files changed, 89 insertions(+), 154 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c6f38fa..a504620 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -196,21 +196,6 @@ 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" @@ -287,12 +272,6 @@ 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" @@ -395,16 +374,6 @@ 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" @@ -419,15 +388,12 @@ checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "pg_stage_rs" -version = "0.1.2" +version = "0.1.1" dependencies = [ "chrono", "clap", - "crossbeam-channel", "flate2", "hmac", - "memchr", - "num_cpus", "rand", "regex", "serde", diff --git a/Cargo.toml b/Cargo.toml index 27948e6..984f38b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "pg_stage_rs" -version = "0.1.2" +version = "0.1.1" edition = "2021" description = "PostgreSQL dump anonymizer - streaming obfuscation for plain and custom formats" @@ -9,10 +9,7 @@ clap = { version = "4", features = ["derive"] } serde = { version = "1", features = ["derive"] } serde_json = "1" flate2 = "1" -zstd = { version = "0.13", features = ["zstdmt"] } -memchr = "2" -crossbeam-channel = "0.5" -num_cpus = "1.16" +zstd = "0.13" 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 c31df93..05ce249 100644 --- a/src/format/custom/blocks.rs +++ b/src/format/custom/blocks.rs @@ -1,9 +1,8 @@ -use std::io::{self, Read, Write}; +use std::io::{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; @@ -12,71 +11,8 @@ use crate::format::custom::header::CompressionMethod; use crate::format::custom::io::DumpIO; use crate::processor::DataProcessor; -const OUTPUT_CHUNK_SIZE: usize = 1024 * 1024; // 1MB for better throughput +const OUTPUT_CHUNK_SIZE: usize = 512 * 1024; // 512KB 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> { @@ -180,7 +116,7 @@ impl<'a> BlockProcessor<'a> { }; // Split into complete lines + tail - match memrchr(b'\n', &data) { + match data.iter().rposition(|&b| b == b'\n') { Some(last_nl) => { line_tail = data[last_nl + 1..].to_vec(); self.process_complete_lines(&data[..=last_nl], &mut output_buf); @@ -214,20 +150,44 @@ impl<'a> BlockProcessor<'a> { } /// Streaming processing for zlib-compressed blocks. - /// Uses ChunkReader for on-demand chunk reading to minimize memory usage. + /// Reads all compressed chunks (they form one zlib stream), then decompresses + /// and processes incrementally. fn process_block_zlib( &mut self, reader: &mut R, writer: &mut W, ) -> Result<()> { - // Use streaming chunk reader instead of loading entire block into memory - let chunk_reader = ChunkReader::new(reader, self.dio); + // 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(()); + } // Stream: decompress → process lines → compress → write chunks - let mut decoder = ZlibDecoder::new(chunk_reader); + let mut decoder = ZlibDecoder::new(raw.as_slice()); let mut encoder = ZlibEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), Compression::new(6)); - let mut read_buf = vec![0u8; READ_BUF_SIZE]; + let mut read_buf = vec![0u8; OUTPUT_CHUNK_SIZE]; let mut line_tail: Vec = Vec::new(); loop { @@ -244,7 +204,7 @@ impl<'a> BlockProcessor<'a> { line_tail.as_slice() }; - match memrchr(b'\n', data) { + match data.iter().rposition(|&b| b == b'\n') { Some(last_nl) => { let complete = &data[..=last_nl]; let tail = &data[last_nl + 1..]; @@ -290,28 +250,45 @@ impl<'a> BlockProcessor<'a> { } /// Streaming processing for zstd-compressed blocks. - /// Uses ChunkReader for on-demand chunk reading to minimize memory usage. + /// Reads all compressed chunks, decompresses, processes, and recompresses with zstd. fn process_block_zstd( &mut self, reader: &mut R, writer: &mut W, ) -> Result<()> { - // Use streaming chunk reader instead of loading entire block into memory - let chunk_reader = ChunkReader::new(reader, self.dio); + // 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(()); + } // Stream: decompress → process lines → compress → write chunks - // Use compression level 1 for speed (was 3) - let mut decoder = ZstdDecoder::new(chunk_reader) + let mut decoder = ZstdDecoder::new(raw.as_slice()) .map_err(|e| PgStageError::CompressionError(format!("Zstd decoder init failed: {}", e)))?; - - // Use multithread zstd compression for better performance on large data - let mut encoder = ZstdEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), 1) + let mut encoder = ZstdEncoder::new(Vec::with_capacity(OUTPUT_CHUNK_SIZE), 3) .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; READ_BUF_SIZE]; + let mut read_buf = vec![0u8; OUTPUT_CHUNK_SIZE]; let mut line_tail: Vec = Vec::new(); loop { @@ -328,7 +305,7 @@ impl<'a> BlockProcessor<'a> { line_tail.as_slice() }; - match memrchr(b'\n', data) { + match data.iter().rposition(|&b| b == b'\n') { Some(last_nl) => { let complete = &data[..=last_nl]; let tail = &data[last_nl + 1..]; @@ -375,10 +352,9 @@ 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 = memchr(b'\n', &data[start..]) + let end = data[start..].iter().position(|&b| b == b'\n') .map(|p| start + p) .unwrap_or(data.len()); @@ -396,10 +372,9 @@ 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 = memchr(b'\n', &data[start..]) + let end = data[start..].iter().position(|&b| b == b'\n') .map(|p| start + p) .unwrap_or(data.len()); diff --git a/src/format/custom/io.rs b/src/format/custom/io.rs index bb68073..501e77a 100644 --- a/src/format/custom/io.rs +++ b/src/format/custom/io.rs @@ -89,15 +89,16 @@ impl DumpIO { (0u8, val) }; - // Write sign + magnitude in one syscall (max 9 bytes: 1 sign + 8 magnitude) - let mut buf = [0u8; 9]; - buf[0] = sign; + // Sign byte + writer.write_all(&[sign])?; + + // Magnitude bytes (little-endian) let mut current = v_abs; - for i in 0..self.int_size { - buf[1 + i] = (current & 0xFF) as u8; + for _ in 0..self.int_size { + let byte = (current & 0xFF) as u8; + writer.write_all(&[byte])?; 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 959b723..2991ff2 100644 --- a/src/format/custom/mod.rs +++ b/src/format/custom/mod.rs @@ -30,9 +30,8 @@ impl CustomHandler { writer: W, initial_bytes: &[u8], ) -> Result<()> { - // 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); + let mut reader = BufReader::with_capacity(65536, reader); + let mut writer = BufWriter::with_capacity(65536, 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 aa42652..6f26c9f 100644 --- a/src/format/plain.rs +++ b/src/format/plain.rs @@ -21,14 +21,13 @@ impl PlainHandler { writer: W, initial_bytes: &[u8], ) -> Result<()> { - // Large buffers for better throughput (2MB each) - let mut writer = BufWriter::with_capacity(2 * 1024 * 1024, writer); + let mut writer = BufWriter::with_capacity(65536, 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(2 * 1024 * 1024, combined); + let buf_reader = BufReader::with_capacity(65536, 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 21f974e..114dbcd 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; + let current_value = ctx.current_value.clone(); let count = ctx .kwargs .get("obfuscated_numbers_count") diff --git a/src/mutator/mod.rs b/src/mutator/mod.rs index b9c657a..0553eb9 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: &'a str, // Changed from String to &str to avoid allocation + pub current_value: String, 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 7a04ec4..4996481 100644 --- a/src/processor.rs +++ b/src/processor.rs @@ -179,9 +179,7 @@ impl DataProcessor { // Use Cow to avoid allocating Strings for unmodified columns let mut result_values: Vec> = values.iter().map(|&s| Cow::Borrowed(s)).collect(); - - // Pre-allocate with expected capacity to reduce reallocations - let mut obfuscated_values: HashMap = HashMap::with_capacity(self.current_mutations.len()); + let mut obfuscated_values: HashMap = HashMap::new(); // Iterate by index to avoid cloning sorted_columns for col_sort_idx in 0..self.sorted_columns.len() { @@ -197,6 +195,8 @@ 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].as_ref(); + let fk_value = result_values[from_idx].to_string(); 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,13 +228,10 @@ 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: current_value.clone(), rng: &mut self.rng, unique_tracker: &mut self.unique_tracker, locale: self.locale, @@ -248,10 +245,11 @@ 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, - result_values[from_idx].as_ref(), + &fk_value, &new_val, ); }