WIP import_timeline_from_tar

Part of #1467 

This PR adds several performance tests that compare the [PG statistics](https://www.postgresql.org/docs/current/monitoring-stats.html) obtained when running PG benchmarks against Neon and vanilla PG to measure the read/write throughput of the DB.

pageserver/src/import_datadir.rs

@@ -22,6 +22,8 @@ use postgres_ffi::{pg_constants, ControlFileData, DBState_DB_SHUTDOWNED};
 use postgres_ffi::{Oid, TransactionId};
 use utils::lsn::Lsn;
 
+use postgres::CopyOutReader;
+
 ///
 /// Import all relation data pages from local disk into the repository.
 ///
@@ -403,3 +405,126 @@ fn import_wal<R: Repository>(
 
     Ok(())
 }
+
+pub fn import_timeline_from_tar<R: Repository>(
+    tline: &mut DatadirTimeline<R>,
+    copyreader: CopyOutReader,
+    lsn: Lsn,
+) -> Result<()> {
+    let mut ar = tar::Archive::new(copyreader);
+
+    let mut modification = tline.begin_modification(lsn);
+    modification.init_empty()?;
+
+    for e in ar.entries().unwrap() {
+        let mut entry = e.unwrap();
+        let header = entry.header();
+        let file_path = header.path().unwrap().into_owned();
+
+        match header.entry_type() {
+            tar::EntryType::Regular => {
+                let mut buffer = Vec::new();
+                // read the whole entry
+                entry.read_to_end(&mut buffer).unwrap();
+
+                import_file(&mut modification, &file_path.as_ref(), &buffer)?;
+            }
+            tar::EntryType::Directory => {
+                println!("tar::EntryType::Directory {}", file_path.display());
+            }
+            _ => {
+                panic!("tar::EntryType::?? {}", file_path.display());
+            }
+        }
+    }
+    Ok(())
+}
+
+pub fn import_file<R: Repository>(
+    modification: &mut DatadirModification<R>,
+    file_path: &Path,
+    buffer: &[u8],
+) -> Result<()> {
+    if file_path.starts_with("global") {
+        let spcnode = pg_constants::GLOBALTABLESPACE_OID;
+        let dbnode = 0;
+
+        match file_path.file_name().unwrap().to_string_lossy().as_ref() {
+            "pg_control" => {
+                println!("pg_control file {}", file_path.display());
+
+                // Import it as ControlFile
+                modification.put_control_file(Bytes::copy_from_slice(&buffer[..]))?;
+
+                // Extract the checkpoint record and import it separately.
+                let pg_control = ControlFileData::decode(&buffer)?;
+                let checkpoint_bytes = pg_control.checkPointCopy.encode()?;
+                modification.put_checkpoint(checkpoint_bytes)?;
+            }
+            "pg_filenode.map" => {
+                ("pg_filenode.map file {}", file_path.display());
+            }
+            _ => {
+                println!("global relfile {}", file_path.display());
+                //TODO
+            }
+        }
+    } else if file_path.starts_with("base") {
+        let spcnode = pg_constants::DEFAULTTABLESPACE_OID;
+        let dbnode: u32 = file_path
+            .iter()
+            .skip(1)
+            .next()
+            .unwrap()
+            .to_string_lossy()
+            .parse()
+            .unwrap();
+
+        match file_path.file_name().unwrap().to_string_lossy().as_ref() {
+            "pg_filenode.map" => {
+                println!(
+                    "dbnode {} pg_filenode.map file {}",
+                    dbnode,
+                    file_path.display()
+                );
+                modification.put_relmap_file(
+                    spcnode,
+                    dbnode,
+                    Bytes::copy_from_slice(&buffer[..]),
+                )?;
+            }
+            _ => {
+                println!("dbnode {} relfile {}", dbnode, file_path.display());
+                //TODO
+            }
+        }
+    } else if file_path.starts_with("pg_xact") {
+        println!(
+            "pg_xact {} ",
+            file_path.file_name().unwrap().to_string_lossy().as_ref()
+        );
+        // TODO
+    } else if file_path.starts_with("pg_multixact/offset") {
+        println!(
+            "pg_multixact/offset {}",
+            file_path.file_name().unwrap().to_string_lossy().as_ref()
+        );
+        // TODO
+    } else if file_path.starts_with("pg_multixact/members") {
+        println!(
+            "pg_multixact/members {}",
+            file_path.file_name().unwrap().to_string_lossy().as_ref()
+        );
+        // TODO
+    } else if file_path.starts_with("pg_twophase") {
+        let xid = u32::from_str_radix(&file_path.file_name().unwrap().to_string_lossy(), 16)?;
+
+        println!(
+            "xid {} pg_twophase {}",
+            xid,
+            file_path.file_name().unwrap().to_string_lossy().as_ref()
+        );
+        modification.put_twophase_file(xid, Bytes::copy_from_slice(&buffer[..]))?;
+    }
+    Ok(())
+}

test_runner/conftest.py

@@ -1,6 +1,5 @@
-pytest_plugins = (
-    "fixtures.neon_fixtures",
-    "fixtures.benchmark_fixture",
-    "fixtures.compare_fixtures",
-    "fixtures.slow",
-)
+pytest_plugins = ("fixtures.neon_fixtures",
+                  "fixtures.benchmark_fixture",
+                  "fixtures.compare_fixtures",
+                  "fixtures.slow",
+                  "fixtures.pg_stats")

test_runner/fixtures/compare_fixtures.py

@@ -1,12 +1,13 @@
 import pytest
 from contextlib import contextmanager
 from abc import ABC, abstractmethod
+from fixtures.pg_stats import PgStatTable
 
 from fixtures.neon_fixtures import PgBin, PgProtocol, VanillaPostgres, RemotePostgres, NeonEnv
 from fixtures.benchmark_fixture import MetricReport, NeonBenchmarker
 
 # Type-related stuff
-from typing import Iterator
+from typing import Dict, List
 
 
 class PgCompare(ABC):
@@ -51,6 +52,31 @@ class PgCompare(ABC):
     def record_duration(self, out_name):
         pass
 
+    @contextmanager
+    def record_pg_stats(self, pg_stats: List[PgStatTable]):
+        init_data = self._retrieve_pg_stats(pg_stats)
+
+        yield
+
+        data = self._retrieve_pg_stats(pg_stats)
+
+        for k in set(init_data) & set(data):
+            self.zenbenchmark.record(k, data[k] - init_data[k], '', MetricReport.HIGHER_IS_BETTER)
+
+    def _retrieve_pg_stats(self, pg_stats: List[PgStatTable]) -> Dict[str, int]:
+        results: Dict[str, int] = {}
+
+        with self.pg.connect().cursor() as cur:
+            for pg_stat in pg_stats:
+                cur.execute(pg_stat.query)
+                row = cur.fetchone()
+                assert len(row) == len(pg_stat.columns)
+
+                for col, val in zip(pg_stat.columns, row):
+                    results[f"{pg_stat.table}.{col}"] = int(val)
+
+        return results
+
 
 class NeonCompare(PgCompare):
     """PgCompare interface for the neon stack."""

test_runner/fixtures/pg_stats.py

@@ -0,0 +1,52 @@
+from typing import List
+
+import pytest
+
+
+class PgStatTable:
+    table: str
+    columns: List[str]
+    additional_query: str
+
+    def __init__(self, table: str, columns: List[str], filter_query: str = ""):
+        self.table = table
+        self.columns = columns
+        self.additional_query = filter_query
+
+    @property
+    def query(self) -> str:
+        return f"SELECT {','.join(self.columns)} FROM {self.table} {self.additional_query}"
+
+
+@pytest.fixture(scope='function')
+def pg_stats_rw() -> List[PgStatTable]:
+    return [
+        PgStatTable("pg_stat_database",
+                    ["tup_returned", "tup_fetched", "tup_inserted", "tup_updated", "tup_deleted"],
+                    "WHERE datname='postgres'"),
+    ]
+
+
+@pytest.fixture(scope='function')
+def pg_stats_ro() -> List[PgStatTable]:
+    return [
+        PgStatTable("pg_stat_database", ["tup_returned", "tup_fetched"],
+                    "WHERE datname='postgres'"),
+    ]
+
+
+@pytest.fixture(scope='function')
+def pg_stats_wo() -> List[PgStatTable]:
+    return [
+        PgStatTable("pg_stat_database", ["tup_inserted", "tup_updated", "tup_deleted"],
+                    "WHERE datname='postgres'"),
+    ]
+
+
+@pytest.fixture(scope='function')
+def pg_stats_wal() -> List[PgStatTable]:
+    return [
+        PgStatTable("pg_stat_wal",
+                    ["wal_records", "wal_fpi", "wal_bytes", "wal_buffers_full", "wal_write"],
+                    "")
+    ]

test_runner/performance/test_compare_pg_stats.py

@@ -0,0 +1,101 @@
+import os
+from typing import List
+
+import pytest
+from fixtures.compare_fixtures import PgCompare
+from fixtures.pg_stats import PgStatTable
+
+from performance.test_perf_pgbench import get_durations_matrix, get_scales_matrix
+
+
+def get_seeds_matrix(default: int = 100):
+    seeds = os.getenv("TEST_PG_BENCH_SEEDS_MATRIX", default=str(default))
+    return list(map(int, seeds.split(",")))
+
+
+@pytest.mark.parametrize("seed", get_seeds_matrix())
+@pytest.mark.parametrize("scale", get_scales_matrix())
+@pytest.mark.parametrize("duration", get_durations_matrix(5))
+def test_compare_pg_stats_rw_with_pgbench_default(neon_with_baseline: PgCompare,
+                                                  seed: int,
+                                                  scale: int,
+                                                  duration: int,
+                                                  pg_stats_rw: List[PgStatTable]):
+    env = neon_with_baseline
+    # initialize pgbench
+    env.pg_bin.run_capture(['pgbench', f'-s{scale}', '-i', env.pg.connstr()])
+    env.flush()
+
+    with env.record_pg_stats(pg_stats_rw):
+        env.pg_bin.run_capture(
+            ['pgbench', f'-T{duration}', f'--random-seed={seed}', '-Mprepared', env.pg.connstr()])
+        env.flush()
+
+
+@pytest.mark.parametrize("seed", get_seeds_matrix())
+@pytest.mark.parametrize("scale", get_scales_matrix())
+@pytest.mark.parametrize("duration", get_durations_matrix(5))
+def test_compare_pg_stats_wo_with_pgbench_simple_update(neon_with_baseline: PgCompare,
+                                                        seed: int,
+                                                        scale: int,
+                                                        duration: int,
+                                                        pg_stats_wo: List[PgStatTable]):
+    env = neon_with_baseline
+    # initialize pgbench
+    env.pg_bin.run_capture(['pgbench', f'-s{scale}', '-i', env.pg.connstr()])
+    env.flush()
+
+    with env.record_pg_stats(pg_stats_wo):
+        env.pg_bin.run_capture([
+            'pgbench',
+            '-N',
+            f'-T{duration}',
+            f'--random-seed={seed}',
+            '-Mprepared',
+            env.pg.connstr()
+        ])
+        env.flush()
+
+
+@pytest.mark.parametrize("seed", get_seeds_matrix())
+@pytest.mark.parametrize("scale", get_scales_matrix())
+@pytest.mark.parametrize("duration", get_durations_matrix(5))
+def test_compare_pg_stats_ro_with_pgbench_select_only(neon_with_baseline: PgCompare,
+                                                      seed: int,
+                                                      scale: int,
+                                                      duration: int,
+                                                      pg_stats_ro: List[PgStatTable]):
+    env = neon_with_baseline
+    # initialize pgbench
+    env.pg_bin.run_capture(['pgbench', f'-s{scale}', '-i', env.pg.connstr()])
+    env.flush()
+
+    with env.record_pg_stats(pg_stats_ro):
+        env.pg_bin.run_capture([
+            'pgbench',
+            '-S',
+            f'-T{duration}',
+            f'--random-seed={seed}',
+            '-Mprepared',
+            env.pg.connstr()
+        ])
+        env.flush()
+
+
+@pytest.mark.parametrize("seed", get_seeds_matrix())
+@pytest.mark.parametrize("scale", get_scales_matrix())
+@pytest.mark.parametrize("duration", get_durations_matrix(5))
+def test_compare_pg_stats_wal_with_pgbench_default(neon_with_baseline: PgCompare,
+                                                   seed: int,
+                                                   scale: int,
+                                                   duration: int,
+                                                   pg_stats_wal: List[PgStatTable]):
+    env = neon_with_baseline
+    # initialize pgbench
+    env.pg_bin.run_capture(['pgbench', f'-s{scale}', '-i', env.pg.connstr()])
+    env.flush()
+
+    with env.record_pg_stats(pg_stats_wal):
+        env.pg_bin.run_capture(
+            ['pgbench', f'-T{duration}', f'--random-seed={seed}', '-Mprepared', env.pg.connstr()])
+        env.flush()

test_runner/performance/test_perf_pgbench.py

@@ -79,7 +79,7 @@ def run_test_pgbench(env: PgCompare, scale: int, duration: int):
     # Run simple-update workload
     run_pgbench(env,
                 "simple-update",
-                ['pgbench', '-n', '-c4', f'-T{duration}', '-P2', '-Mprepared', env.pg.connstr()])
+                ['pgbench', '-N', '-c4', f'-T{duration}', '-P2', '-Mprepared', env.pg.connstr()])
 
     # Run SELECT workload
     run_pgbench(env,
@@ -89,13 +89,13 @@ def run_test_pgbench(env: PgCompare, scale: int, duration: int):
     env.report_size()
 
 
-def get_durations_matrix():
-    durations = os.getenv("TEST_PG_BENCH_DURATIONS_MATRIX", default="45")
+def get_durations_matrix(default: int = 45):
+    durations = os.getenv("TEST_PG_BENCH_DURATIONS_MATRIX", default=str(default))
     return list(map(int, durations.split(",")))
 
 
-def get_scales_matrix():
-    scales = os.getenv("TEST_PG_BENCH_SCALES_MATRIX", default="10")
+def get_scales_matrix(default: int = 10):
+    scales = os.getenv("TEST_PG_BENCH_SCALES_MATRIX", default=str(default))
     return list(map(int, scales.split(",")))