postgres_ffi: add WalStreamDecoder::complete_record() benchmark (#10158)

Touches #10097.

libs/postgres_ffi/Cargo.toml

@@ -9,9 +9,11 @@ regex.workspace = true
 bytes.workspace = true
 anyhow.workspace = true
 crc32c.workspace = true
+criterion.workspace = true
 once_cell.workspace = true
 log.workspace = true
 memoffset.workspace = true
+pprof.workspace = true
 thiserror.workspace = true
 serde.workspace = true
 utils.workspace = true
@@ -24,3 +26,7 @@ postgres.workspace = true
 [build-dependencies]
 anyhow.workspace = true
 bindgen.workspace = true
+
+[[bench]]
+name = "waldecoder"
+harness = false

libs/postgres_ffi/benches/README.md

@@ -0,0 +1,26 @@
+## Benchmarks
+
+To run benchmarks:
+
+```sh
+# All benchmarks.
+cargo bench --package postgres_ffi
+
+# Specific file.
+cargo bench --package postgres_ffi --bench waldecoder
+
+# Specific benchmark.
+cargo bench --package postgres_ffi --bench waldecoder complete_record/size=1024
+
+# List available benchmarks.
+cargo bench --package postgres_ffi --benches -- --list
+
+# Generate flamegraph profiles using pprof-rs, profiling for 10 seconds.
+# Output in target/criterion/*/profile/flamegraph.svg.
+cargo bench --package postgres_ffi --bench waldecoder complete_record/size=1024 -- --profile-time 10
+```
+
+Additional charts and statistics are available in `target/criterion/report/index.html`.
+
+Benchmarks are automatically compared against the previous run. To compare against other runs, see
+`--baseline` and `--save-baseline`.

libs/postgres_ffi/benches/waldecoder.rs

@@ -0,0 +1,49 @@
+use std::ffi::CStr;
+
+use criterion::{criterion_group, criterion_main, Bencher, Criterion};
+use postgres_ffi::v17::wal_generator::LogicalMessageGenerator;
+use postgres_ffi::v17::waldecoder_handler::WalStreamDecoderHandler;
+use postgres_ffi::waldecoder::WalStreamDecoder;
+use pprof::criterion::{Output, PProfProfiler};
+use utils::lsn::Lsn;
+
+const KB: usize = 1024;
+
+// Register benchmarks with Criterion.
+criterion_group!(
+    name = benches;
+    config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
+    targets = bench_complete_record,
+);
+criterion_main!(benches);
+
+/// Benchmarks WalStreamDecoder::complete_record() for a logical message of varying size.
+fn bench_complete_record(c: &mut Criterion) {
+    let mut g = c.benchmark_group("complete_record");
+    for size in [64, KB, 8 * KB, 128 * KB] {
+        // Kind of weird to change the group throughput per benchmark, but it's the only way
+        // to vary it per benchmark. It works.
+        g.throughput(criterion::Throughput::Bytes(size as u64));
+        g.bench_function(format!("size={size}"), |b| run_bench(b, size).unwrap());
+    }
+
+    fn run_bench(b: &mut Bencher, size: usize) -> anyhow::Result<()> {
+        const PREFIX: &CStr = c"";
+        let value_size = LogicalMessageGenerator::make_value_size(size, PREFIX);
+        let value = vec![1; value_size];
+
+        let mut decoder = WalStreamDecoder::new(Lsn(0), 170000);
+        let msg = LogicalMessageGenerator::new(PREFIX, &value)
+            .next()
+            .unwrap()
+            .encode(Lsn(0));
+        assert_eq!(msg.len(), size);
+
+        b.iter(|| {
+            let msg = msg.clone(); // Bytes::clone() is cheap
+            decoder.complete_record(msg).unwrap();
+        });
+
+        Ok(())
+    }
+}

libs/postgres_ffi/src/wal_generator.rs

@@ -231,6 +231,22 @@ impl LogicalMessageGenerator {
         };
         [&header.encode(), prefix, message].concat().into()
     }
+
+    /// Computes how large a value must be to get a record of the given size. Convenience method to
+    /// construct records of pre-determined size. Panics if the record size is too small.
+    pub fn make_value_size(record_size: usize, prefix: &CStr) -> usize {
+        let xlog_header_size = XLOG_SIZE_OF_XLOG_RECORD;
+        let lm_header_size = size_of::<XlLogicalMessage>();
+        let prefix_size = prefix.to_bytes_with_nul().len();
+        let data_header_size = match record_size - xlog_header_size - 2 {
+            0..=255 => 2,
+            256..=258 => panic!("impossible record_size {record_size}"),
+            259.. => 5,
+        };
+        record_size
+            .checked_sub(xlog_header_size + lm_header_size + prefix_size + data_header_size)
+            .expect("record_size too small")
+    }
 }
 
 impl Iterator for LogicalMessageGenerator {