Revert "Revert "Performance optimizations for large databases""

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"

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<u8>,
+    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<R: Read> Read for ChunkReader<'_, R> {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        // 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<R: Read, W: Write>(
         &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<u8> = 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<R: Read, W: Write>(
         &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<u8> = 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<u8>) {
+        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<W: Write>(&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());
 

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(())
     }

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)?;

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<String> = 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?;

src/mutator/contact.rs

@@ -78,7 +78,7 @@ pub fn address(ctx: &mut MutationContext) -> Result<String> {
 }
 
 pub fn deterministic_phone(ctx: &mut MutationContext) -> Result<String> {
-    let current_value = ctx.current_value.clone();
+    let current_value = ctx.current_value;
     let count = ctx
         .kwargs
         .get("obfuscated_numbers_count")