set pageserver id in dockerfile

This adds node id parameter to pageserver configuration. Also I use a
simple builder to construct pageserver config struct to avoid setting
node id to some temporary invalid value. Some of the changes in test
fixtures are needed to split init and start operations for envrionment.

.circleci/config.yml

@@ -54,7 +54,8 @@ jobs:
             if [ ! -e tmp_install/bin/postgres ]; then
               # "depth 1" saves some time by not cloning the whole repo
               git submodule update --init --depth 1
-              mold -run make postgres -j$(nproc)
+              # bail out on any warnings
+              COPT='-Werror' mold -run make postgres -j$(nproc)
             fi
 
       - save_cache:
@@ -297,6 +298,7 @@ jobs:
             - PLATFORM: zenith-local-ci
           command: |
             PERF_REPORT_DIR="$(realpath test_runner/perf-report-local)"
+            rm -rf $PERF_REPORT_DIR
 
             TEST_SELECTION="test_runner/<< parameters.test_selection >>"
             EXTRA_PARAMS="<< parameters.extra_params >>"
@@ -341,7 +343,6 @@ jobs:
 
             if << parameters.save_perf_report >>; then
               if [[ $CIRCLE_BRANCH == "main" ]]; then
-                # TODO: reuse scripts/git-upload
                 export REPORT_FROM="$PERF_REPORT_DIR"
                 export REPORT_TO=local
                 scripts/generate_and_push_perf_report.sh
@@ -597,6 +598,7 @@ workflows:
             - build-postgres-<< matrix.build_type >>
       - run-pytest:
           name: pg_regress-tests-<< matrix.build_type >>
+          context: PERF_TEST_RESULT_CONNSTR
           matrix:
             parameters:
               build_type: ["debug", "release"]
@@ -614,6 +616,7 @@ workflows:
             - build-zenith-<< matrix.build_type >>
       - run-pytest:
           name: benchmarks
+          context: PERF_TEST_RESULT_CONNSTR
           build_type: release
           test_selection: performance
           run_in_parallel: false

.github/workflows/benchmarking.yml

@@ -3,7 +3,7 @@ name: benchmarking
 on:
   # uncomment to run on push for debugging your PR
   # push:
-  #   branches: [ mybranch ]
+  #   branches: [ your branch ]
   schedule:
     # * is a special character in YAML so you have to quote this string
     #          ┌───────────── minute (0 - 59)
@@ -41,7 +41,7 @@ jobs:
       run: |
         python3 -m pip install --upgrade poetry wheel
         # since pip/poetry caches are reused there shouldn't be any troubles with install every time
-        poetry install
+        ./scripts/pysync
 
     - name: Show versions
       run: |
@@ -89,11 +89,15 @@ jobs:
         BENCHMARK_CONNSTR: "${{ secrets.BENCHMARK_STAGING_CONNSTR }}"
         REMOTE_ENV: "1" # indicate to test harness that we do not have zenith binaries locally
       run: |
+        # just to be sure that no data was cached on self hosted runner
+        # since it might generate duplicates when calling ingest_perf_test_result.py
+        rm -rf perf-report-staging
         mkdir -p perf-report-staging
         ./scripts/pytest test_runner/performance/ -v -m "remote_cluster" --skip-interfering-proc-check --out-dir perf-report-staging
 
     - name: Submit result
       env:
         VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
+        PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
       run: |
         REPORT_FROM=$(realpath perf-report-staging) REPORT_TO=staging scripts/generate_and_push_perf_report.sh

compute_tools/Cargo.toml

@@ -1,17 +1,14 @@
 [package]
 name = "compute_tools"
 version = "0.1.0"
-authors = ["Alexey Kondratov <kondratov.aleksey@gmail.com>"]
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 libc = "0.2"
 anyhow = "1.0"
 chrono = "0.4"
-clap = "2.33"
-env_logger = "0.8"
+clap = "3.0"
+env_logger = "0.9"
 hyper = { version = "0.14", features = ["full"] }
 log = { version = "0.4", features = ["std", "serde"] }
 postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="9eb0dbfbeb6a6c1b79099b9f7ae4a8c021877858" }

compute_tools/src/bin/zenith_ctl.rs

@@ -34,6 +34,7 @@ use std::sync::{Arc, RwLock};
 
 use anyhow::{Context, Result};
 use chrono::Utc;
+use clap::Arg;
 use log::info;
 use postgres::{Client, NoTls};
 
@@ -162,34 +163,34 @@ fn main() -> Result<()> {
     let matches = clap::App::new("zenith_ctl")
         .version(version.unwrap_or("unknown"))
         .arg(
-            clap::Arg::with_name("connstr")
-                .short("C")
+            Arg::new("connstr")
+                .short('C')
                 .long("connstr")
                 .value_name("DATABASE_URL")
                 .required(true),
         )
         .arg(
-            clap::Arg::with_name("pgdata")
-                .short("D")
+            Arg::new("pgdata")
+                .short('D')
                 .long("pgdata")
                 .value_name("DATADIR")
                 .required(true),
         )
         .arg(
-            clap::Arg::with_name("pgbin")
-                .short("b")
+            Arg::new("pgbin")
+                .short('b')
                 .long("pgbin")
                 .value_name("POSTGRES_PATH"),
         )
         .arg(
-            clap::Arg::with_name("spec")
-                .short("s")
+            Arg::new("spec")
+                .short('s')
                 .long("spec")
                 .value_name("SPEC_JSON"),
         )
         .arg(
-            clap::Arg::with_name("spec-path")
-                .short("S")
+            Arg::new("spec-path")
+                .short('S')
                 .long("spec-path")
                 .value_name("SPEC_PATH"),
         )

control_plane/Cargo.toml

@@ -1,11 +1,8 @@
 [package]
 name = "control_plane"
 version = "0.1.0"
-authors = ["Stas Kelvich <stas@zenith.tech>"]
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 tar = "0.4.33"
 postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }

control_plane/safekeepers.conf

@@ -1,20 +1,20 @@
 # Page server and three safekeepers.
 [pageserver]
-listen_pg_addr = 'localhost:64000'
-listen_http_addr = 'localhost:9898'
+listen_pg_addr = '127.0.0.1:64000'
+listen_http_addr = '127.0.0.1:9898'
 auth_type = 'Trust'
 
 [[safekeepers]]
-name = 'sk1'
+id = 1
 pg_port = 5454
 http_port = 7676
 
 [[safekeepers]]
-name = 'sk2'
+id = 2
 pg_port = 5455
 http_port = 7677
 
 [[safekeepers]]
-name = 'sk3'
+id = 3
 pg_port = 5456
 http_port = 7678

control_plane/simple.conf

@@ -1,11 +1,11 @@
 # Minimal zenith environment with one safekeeper. This is equivalent to the built-in
 # defaults that you get with no --config
 [pageserver]
-listen_pg_addr = 'localhost:64000'
-listen_http_addr = 'localhost:9898'
+listen_pg_addr = '127.0.0.1:64000'
+listen_http_addr = '127.0.0.1:9898'
 auth_type = 'Trust'
 
 [[safekeepers]]
-name = 'single'
+id = 1
 pg_port = 5454
 http_port = 7676

control_plane/src/compute.rs

@@ -334,14 +334,26 @@ impl PostgresNode {
         if let Some(lsn) = self.lsn {
             conf.append("recovery_target_lsn", &lsn.to_string());
         }
+
         conf.append_line("");
+        // Configure backpressure
+        // - Replication write lag depends on how fast the walreceiver can process incoming WAL.
+        //   This lag determines latency of get_page_at_lsn. Speed of applying WAL is about 10MB/sec,
+        //   so to avoid expiration of 1 minute timeout, this lag should not be larger than 600MB.
+        //   Actually latency should be much smaller (better if < 1sec). But we assume that recently
+        //   updates pages are not requested from pageserver.
+        // - Replication flush lag depends on speed of persisting data by checkpointer (creation of
+        //   delta/image layers) and advancing disk_consistent_lsn. Safekeepers are able to
+        //   remove/archive WAL only beyond disk_consistent_lsn. Too large a lag can cause long
+        //   recovery time (in case of pageserver crash) and disk space overflow at safekeepers.
+        // - Replication apply lag depends on speed of uploading changes to S3 by uploader thread.
+        //   To be able to restore database in case of pageserver node crash, safekeeper should not
+        //   remove WAL beyond this point. Too large lag can cause space exhaustion in safekeepers
+        //   (if they are not able to upload WAL to S3).
+        conf.append("max_replication_write_lag", "500MB");
+        conf.append("max_replication_flush_lag", "10GB");
 
         if !self.env.safekeepers.is_empty() {
-            // Configure backpressure
-            // In setup with safekeepers apply_lag depends on
-            // speed of data checkpointing on pageserver (see disk_consistent_lsn).
-            conf.append("max_replication_apply_lag", "1500MB");
-
             // Configure the node to connect to the safekeepers
             conf.append("synchronous_standby_names", "walproposer");
 
@@ -354,11 +366,6 @@ impl PostgresNode {
                 .join(",");
             conf.append("wal_acceptors", &wal_acceptors);
         } else {
-            // Configure backpressure
-            // In setup without safekeepers, flush_lag depends on
-            // speed of of data checkpointing on pageserver (see disk_consistent_lsn)
-            conf.append("max_replication_flush_lag", "1500MB");
-
             // We only use setup without safekeepers for tests,
             // and don't care about data durability on pageserver,
             // so set more relaxed synchronous_commit.

control_plane/src/local_env.rs

@@ -12,7 +12,9 @@ use std::path::{Path, PathBuf};
 use std::process::{Command, Stdio};
 use zenith_utils::auth::{encode_from_key_file, Claims, Scope};
 use zenith_utils::postgres_backend::AuthType;
-use zenith_utils::zid::{opt_display_serde, ZTenantId};
+use zenith_utils::zid::{opt_display_serde, ZNodeId, ZTenantId};
+
+use crate::safekeeper::SafekeeperNode;
 
 //
 // This data structures represents zenith CLI config
@@ -62,6 +64,8 @@ pub struct LocalEnv {
 #[derive(Serialize, Deserialize, Clone, Debug)]
 #[serde(default)]
 pub struct PageServerConf {
+    // node id
+    pub id: ZNodeId,
     // Pageserver connection settings
     pub listen_pg_addr: String,
     pub listen_http_addr: String,
@@ -76,6 +80,7 @@ pub struct PageServerConf {
 impl Default for PageServerConf {
     fn default() -> Self {
         Self {
+            id: ZNodeId(0),
             listen_pg_addr: String::new(),
             listen_http_addr: String::new(),
             auth_type: AuthType::Trust,
@@ -87,7 +92,7 @@ impl Default for PageServerConf {
 #[derive(Serialize, Deserialize, Clone, Debug)]
 #[serde(default)]
 pub struct SafekeeperConf {
-    pub name: String,
+    pub id: ZNodeId,
     pub pg_port: u16,
     pub http_port: u16,
     pub sync: bool,
@@ -96,7 +101,7 @@ pub struct SafekeeperConf {
 impl Default for SafekeeperConf {
     fn default() -> Self {
         Self {
-            name: String::new(),
+            id: ZNodeId(0),
             pg_port: 0,
             http_port: 0,
             sync: true,
@@ -136,8 +141,8 @@ impl LocalEnv {
         self.base_data_dir.clone()
     }
 
-    pub fn safekeeper_data_dir(&self, node_name: &str) -> PathBuf {
-        self.base_data_dir.join("safekeepers").join(node_name)
+    pub fn safekeeper_data_dir(&self, data_dir_name: &str) -> PathBuf {
+        self.base_data_dir.join("safekeepers").join(data_dir_name)
     }
 
     /// Create a LocalEnv from a config file.
@@ -285,7 +290,7 @@ impl LocalEnv {
         fs::create_dir_all(self.pg_data_dirs_path())?;
 
         for safekeeper in &self.safekeepers {
-            fs::create_dir_all(self.safekeeper_data_dir(&safekeeper.name))?;
+            fs::create_dir_all(SafekeeperNode::datadir_path_by_id(self, safekeeper.id))?;
         }
 
         let mut conf_content = String::new();

control_plane/src/safekeeper.rs

@@ -15,6 +15,7 @@ use reqwest::blocking::{Client, RequestBuilder, Response};
 use reqwest::{IntoUrl, Method};
 use thiserror::Error;
 use zenith_utils::http::error::HttpErrorBody;
+use zenith_utils::zid::ZNodeId;
 
 use crate::local_env::{LocalEnv, SafekeeperConf};
 use crate::storage::PageServerNode;
@@ -61,7 +62,7 @@ impl ResponseErrorMessageExt for Response {
 //
 #[derive(Debug)]
 pub struct SafekeeperNode {
-    pub name: String,
+    pub id: ZNodeId,
 
     pub conf: SafekeeperConf,
 
@@ -77,15 +78,15 @@ impl SafekeeperNode {
     pub fn from_env(env: &LocalEnv, conf: &SafekeeperConf) -> SafekeeperNode {
         let pageserver = Arc::new(PageServerNode::from_env(env));
 
-        println!("initializing for {} for {}", conf.name, conf.http_port);
+        println!("initializing for sk {} for {}", conf.id, conf.http_port);
 
         SafekeeperNode {
-            name: conf.name.clone(),
+            id: conf.id,
             conf: conf.clone(),
             pg_connection_config: Self::safekeeper_connection_config(conf.pg_port),
             env: env.clone(),
             http_client: Client::new(),
-            http_base_url: format!("http://localhost:{}/v1", conf.http_port),
+            http_base_url: format!("http://127.0.0.1:{}/v1", conf.http_port),
             pageserver,
         }
     }
@@ -93,13 +94,17 @@ impl SafekeeperNode {
     /// Construct libpq connection string for connecting to this safekeeper.
     fn safekeeper_connection_config(port: u16) -> Config {
         // TODO safekeeper authentication not implemented yet
-        format!("postgresql://no_user@localhost:{}/no_db", port)
+        format!("postgresql://no_user@127.0.0.1:{}/no_db", port)
             .parse()
             .unwrap()
     }
 
+    pub fn datadir_path_by_id(env: &LocalEnv, sk_id: ZNodeId) -> PathBuf {
+        env.safekeeper_data_dir(format!("sk{}", sk_id).as_ref())
+    }
+
     pub fn datadir_path(&self) -> PathBuf {
-        self.env.safekeeper_data_dir(&self.name)
+        SafekeeperNode::datadir_path_by_id(&self.env, self.id)
     }
 
     pub fn pid_file(&self) -> PathBuf {
@@ -114,12 +119,13 @@ impl SafekeeperNode {
         );
         io::stdout().flush().unwrap();
 
-        let listen_pg = format!("localhost:{}", self.conf.pg_port);
-        let listen_http = format!("localhost:{}", self.conf.http_port);
+        let listen_pg = format!("127.0.0.1:{}", self.conf.pg_port);
+        let listen_http = format!("127.0.0.1:{}", self.conf.http_port);
 
         let mut cmd = Command::new(self.env.safekeeper_bin()?);
         fill_rust_env_vars(
             cmd.args(&["-D", self.datadir_path().to_str().unwrap()])
+                .args(&["--id", self.id.to_string().as_ref()])
                 .args(&["--listen-pg", &listen_pg])
                 .args(&["--listen-http", &listen_http])
                 .args(&["--recall", "1 second"])
@@ -183,7 +189,7 @@ impl SafekeeperNode {
     pub fn stop(&self, immediate: bool) -> anyhow::Result<()> {
         let pid_file = self.pid_file();
         if !pid_file.exists() {
-            println!("Safekeeper {} is already stopped", self.name);
+            println!("Safekeeper {} is already stopped", self.id);
             return Ok(());
         }
         let pid = read_pidfile(&pid_file)?;

control_plane/src/storage.rs

@@ -103,6 +103,8 @@ impl PageServerNode {
     ) -> anyhow::Result<()> {
         let mut cmd = Command::new(self.env.pageserver_bin()?);
 
+        let id = format!("id={}", self.env.pageserver.id);
+
         // FIXME: the paths should be shell-escaped to handle paths with spaces, quotas etc.
         let base_data_dir_param = self.env.base_data_dir.display().to_string();
         let pg_distrib_dir_param =
@@ -122,6 +124,7 @@ impl PageServerNode {
         args.extend(["-c", &authg_type_param]);
         args.extend(["-c", &listen_http_addr_param]);
         args.extend(["-c", &listen_pg_addr_param]);
+        args.extend(["-c", &id]);
 
         for config_override in config_overrides {
             args.extend(["-c", config_override]);

docker-entrypoint.sh

@@ -4,7 +4,7 @@ set -eux
 if [ "$1" = 'pageserver' ]; then
     if [ ! -d "/data/tenants" ]; then
         echo "Initializing pageserver data directory"
-        pageserver --init -D /data -c "pg_distrib_dir='/usr/local'"
+        pageserver --init -D /data -c "pg_distrib_dir='/usr/local'" -c "id=10"
     fi
     echo "Staring pageserver at 0.0.0.0:6400"
     pageserver -c "listen_pg_addr='0.0.0.0:6400'" -c "listen_http_addr='0.0.0.0:9898'" -D /data

pageserver/Cargo.toml

@@ -1,7 +1,6 @@
 [package]
 name = "pageserver"
 version = "0.1.0"
-authors = ["Stas Kelvich <stas@zenith.tech>"]
 edition = "2021"
 
 [dependencies]
@@ -15,7 +14,7 @@ futures = "0.3.13"
 hyper = "0.14"
 lazy_static = "1.4.0"
 log = "0.4.14"
-clap = "2.33.0"
+clap = "3.0"
 daemonize = "0.4.1"
 tokio = { version = "1.11", features = ["process", "sync", "macros", "fs", "rt", "io-util", "time"] }
 postgres-types = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
@@ -23,7 +22,6 @@ postgres-protocol = { git = "https://github.com/zenithdb/rust-postgres.git", rev
 postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
 tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
 tokio-stream = "0.1.8"
-routerify = "2"
 anyhow = { version = "1.0", features = ["backtrace"] }
 crc32c = "0.6.0"
 thiserror = "1.0"
@@ -32,7 +30,7 @@ tar = "0.4.33"
 humantime = "2.1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
-toml_edit = { version = "0.12", features = ["easy"] }
+toml_edit = { version = "0.13", features = ["easy"] }
 scopeguard = "1.1.0"
 async-trait = "0.1"
 const_format = "0.2.21"
@@ -42,8 +40,8 @@ signal-hook = "0.3.10"
 url = "2"
 nix = "0.23"
 once_cell = "1.8.0"
-parking_lot = "0.11.2"
 crossbeam-utils = "0.8.5"
+fail = "0.5.0"
 
 rust-s3 = { version = "0.28", default-features = false, features = ["no-verify-ssl", "tokio-rustls-tls"] }
 async-compression = {version = "0.3", features = ["zstd", "tokio"]}

pageserver/src/bin/dump_layerfile.rs

@@ -13,7 +13,7 @@ fn main() -> Result<()> {
         .about("Dump contents of one layer file, for debugging")
         .version(GIT_VERSION)
         .arg(
-            Arg::with_name("path")
+            Arg::new("path")
                 .help("Path to file to dump")
                 .required(true)
                 .index(1),

pageserver/src/bin/pageserver.rs

@@ -27,27 +27,27 @@ fn main() -> Result<()> {
         .about("Materializes WAL stream to pages and serves them to the postgres")
         .version(GIT_VERSION)
         .arg(
-            Arg::with_name("daemonize")
-                .short("d")
+            Arg::new("daemonize")
+                .short('d')
                 .long("daemonize")
                 .takes_value(false)
                 .help("Run in the background"),
         )
         .arg(
-            Arg::with_name("init")
+            Arg::new("init")
                 .long("init")
                 .takes_value(false)
                 .help("Initialize pageserver repo"),
         )
         .arg(
-            Arg::with_name("workdir")
-                .short("D")
+            Arg::new("workdir")
+                .short('D')
                 .long("workdir")
                 .takes_value(true)
                 .help("Working directory for the pageserver"),
         )
         .arg(
-            Arg::with_name("create-tenant")
+            Arg::new("create-tenant")
                 .long("create-tenant")
                 .takes_value(true)
                 .help("Create tenant during init")
@@ -55,13 +55,13 @@ fn main() -> Result<()> {
         )
         // See `settings.md` for more details on the extra configuration patameters pageserver can process
         .arg(
-            Arg::with_name("config-override")
-                .short("c")
+            Arg::new("config-override")
+                .short('c')
                 .takes_value(true)
                 .number_of_values(1)
-                .multiple(true)
+                .multiple_occurrences(true)
                 .help("Additional configuration overrides of the ones from the toml config file (or new ones to add there).
-                Any option has to be a valid toml document, example: `-c \"foo='hey'\"` `-c \"foo={value=1}\"`"),
+                Any option has to be a valid toml document, example: `-c=\"foo='hey'\"` `-c=\"foo={value=1}\"`"),
         )
         .get_matches();
 
@@ -115,7 +115,14 @@ fn main() -> Result<()> {
                     option_line
                 )
             })?;
+
             for (key, item) in doc.iter() {
+                if key == "id" {
+                    anyhow::ensure!(
+                        init,
+                        "node id can only be set during pageserver init and cannot be overridden"
+                    );
+                }
                 toml.insert(key, item.clone());
             }
         }

pageserver/src/bin/pageserver_zst.rs

@@ -0,0 +1,334 @@
+//! A CLI helper to deal with remote storage (S3, usually) blobs as archives.
+//! See [`compression`] for more details about the archives.
+
+use std::{collections::BTreeSet, path::Path};
+
+use anyhow::{bail, ensure, Context};
+use clap::{App, Arg};
+use pageserver::{
+    layered_repository::metadata::{TimelineMetadata, METADATA_FILE_NAME},
+    remote_storage::compression,
+};
+use tokio::{fs, io};
+use zenith_utils::GIT_VERSION;
+
+const LIST_SUBCOMMAND: &str = "list";
+const ARCHIVE_ARG_NAME: &str = "archive";
+
+const EXTRACT_SUBCOMMAND: &str = "extract";
+const TARGET_DIRECTORY_ARG_NAME: &str = "target_directory";
+
+const CREATE_SUBCOMMAND: &str = "create";
+const SOURCE_DIRECTORY_ARG_NAME: &str = "source_directory";
+
+#[tokio::main(flavor = "current_thread")]
+async fn main() -> anyhow::Result<()> {
+    let arg_matches = App::new("pageserver zst blob [un]compressor utility")
+        .version(GIT_VERSION)
+        .subcommands(vec![
+            App::new(LIST_SUBCOMMAND)
+                .about("List the archive contents")
+                .arg(
+                    Arg::new(ARCHIVE_ARG_NAME)
+                        .required(true)
+                        .takes_value(true)
+                        .help("An archive to list the contents of"),
+                ),
+            App::new(EXTRACT_SUBCOMMAND)
+                .about("Extracts the archive into the directory")
+                .arg(
+                    Arg::new(ARCHIVE_ARG_NAME)
+                        .required(true)
+                        .takes_value(true)
+                        .help("An archive to extract"),
+                )
+                .arg(
+                    Arg::new(TARGET_DIRECTORY_ARG_NAME)
+                        .required(false)
+                        .takes_value(true)
+                        .help("A directory to extract the archive into. Optional, will use the current directory if not specified"),
+                ),
+            App::new(CREATE_SUBCOMMAND)
+                .about("Creates an archive with the contents of a directory (only the first level files are taken, metadata file has to be present in the same directory)")
+                .arg(
+                    Arg::new(SOURCE_DIRECTORY_ARG_NAME)
+                        .required(true)
+                        .takes_value(true)
+                        .help("A directory to use for creating the archive"),
+                )
+                .arg(
+                    Arg::new(TARGET_DIRECTORY_ARG_NAME)
+                        .required(false)
+                        .takes_value(true)
+                        .help("A directory to create the archive in. Optional, will use the current directory if not specified"),
+                ),
+        ])
+        .get_matches();
+
+    let subcommand_name = match arg_matches.subcommand_name() {
+        Some(name) => name,
+        None => bail!("No subcommand specified"),
+    };
+
+    let subcommand_matches = match arg_matches.subcommand_matches(subcommand_name) {
+        Some(matches) => matches,
+        None => bail!(
+            "No subcommand arguments were recognized for subcommand '{}'",
+            subcommand_name
+        ),
+    };
+
+    let target_dir = Path::new(
+        subcommand_matches
+            .value_of(TARGET_DIRECTORY_ARG_NAME)
+            .unwrap_or("./"),
+    );
+
+    match subcommand_name {
+        LIST_SUBCOMMAND => {
+            let archive = match subcommand_matches.value_of(ARCHIVE_ARG_NAME) {
+                Some(archive) => Path::new(archive),
+                None => bail!("No '{}' argument is specified", ARCHIVE_ARG_NAME),
+            };
+            list_archive(archive).await
+        }
+        EXTRACT_SUBCOMMAND => {
+            let archive = match subcommand_matches.value_of(ARCHIVE_ARG_NAME) {
+                Some(archive) => Path::new(archive),
+                None => bail!("No '{}' argument is specified", ARCHIVE_ARG_NAME),
+            };
+            extract_archive(archive, target_dir).await
+        }
+        CREATE_SUBCOMMAND => {
+            let source_dir = match subcommand_matches.value_of(SOURCE_DIRECTORY_ARG_NAME) {
+                Some(source) => Path::new(source),
+                None => bail!("No '{}' argument is specified", SOURCE_DIRECTORY_ARG_NAME),
+            };
+            create_archive(source_dir, target_dir).await
+        }
+        unknown => bail!("Unknown subcommand {}", unknown),
+    }
+}
+
+async fn list_archive(archive: &Path) -> anyhow::Result<()> {
+    let archive = archive.canonicalize().with_context(|| {
+        format!(
+            "Failed to get the absolute path for the archive path '{}'",
+            archive.display()
+        )
+    })?;
+    ensure!(
+        archive.is_file(),
+        "Path '{}' is not an archive file",
+        archive.display()
+    );
+    println!("Listing an archive at path '{}'", archive.display());
+    let archive_name = match archive.file_name().and_then(|name| name.to_str()) {
+        Some(name) => name,
+        None => bail!(
+            "Failed to get the archive name from the path '{}'",
+            archive.display()
+        ),
+    };
+
+    let archive_bytes = fs::read(&archive)
+        .await
+        .context("Failed to read the archive bytes")?;
+
+    let header = compression::read_archive_header(archive_name, &mut archive_bytes.as_slice())
+        .await
+        .context("Failed to read the archive header")?;
+
+    let empty_path = Path::new("");
+    println!("-------------------------------");
+
+    let longest_path_in_archive = header
+        .files
+        .iter()
+        .filter_map(|file| Some(file.subpath.as_path(empty_path).to_str()?.len()))
+        .max()
+        .unwrap_or_default()
+        .max(METADATA_FILE_NAME.len());
+
+    for regular_file in &header.files {
+        println!(
+            "File: {:width$} uncompressed size: {} bytes",
+            regular_file.subpath.as_path(empty_path).display(),
+            regular_file.size,
+            width = longest_path_in_archive,
+        )
+    }
+    println!(
+        "File: {:width$} uncompressed size: {} bytes",
+        METADATA_FILE_NAME,
+        header.metadata_file_size,
+        width = longest_path_in_archive,
+    );
+    println!("-------------------------------");
+
+    Ok(())
+}
+
+async fn extract_archive(archive: &Path, target_dir: &Path) -> anyhow::Result<()> {
+    let archive = archive.canonicalize().with_context(|| {
+        format!(
+            "Failed to get the absolute path for the archive path '{}'",
+            archive.display()
+        )
+    })?;
+    ensure!(
+        archive.is_file(),
+        "Path '{}' is not an archive file",
+        archive.display()
+    );
+    let archive_name = match archive.file_name().and_then(|name| name.to_str()) {
+        Some(name) => name,
+        None => bail!(
+            "Failed to get the archive name from the path '{}'",
+            archive.display()
+        ),
+    };
+
+    if !target_dir.exists() {
+        fs::create_dir_all(target_dir).await.with_context(|| {
+            format!(
+                "Failed to create the target dir at path '{}'",
+                target_dir.display()
+            )
+        })?;
+    }
+    let target_dir = target_dir.canonicalize().with_context(|| {
+        format!(
+            "Failed to get the absolute path for the target dir path '{}'",
+            target_dir.display()
+        )
+    })?;
+    ensure!(
+        target_dir.is_dir(),
+        "Path '{}' is not a directory",
+        target_dir.display()
+    );
+    let mut dir_contents = fs::read_dir(&target_dir)
+        .await
+        .context("Failed to list the target directory contents")?;
+    let dir_entry = dir_contents
+        .next_entry()
+        .await
+        .context("Failed to list the target directory contents")?;
+    ensure!(
+        dir_entry.is_none(),
+        "Target directory '{}' is not empty",
+        target_dir.display()
+    );
+
+    println!(
+        "Extracting an archive at path '{}' into directory '{}'",
+        archive.display(),
+        target_dir.display()
+    );
+
+    let mut archive_file = fs::File::open(&archive).await.with_context(|| {
+        format!(
+            "Failed to get the archive name from the path '{}'",
+            archive.display()
+        )
+    })?;
+    let header = compression::read_archive_header(archive_name, &mut archive_file)
+        .await
+        .context("Failed to read the archive header")?;
+    compression::uncompress_with_header(&BTreeSet::new(), &target_dir, header, &mut archive_file)
+        .await
+        .context("Failed to extract the archive")
+}
+
+async fn create_archive(source_dir: &Path, target_dir: &Path) -> anyhow::Result<()> {
+    let source_dir = source_dir.canonicalize().with_context(|| {
+        format!(
+            "Failed to get the absolute path for the source dir path '{}'",
+            source_dir.display()
+        )
+    })?;
+    ensure!(
+        source_dir.is_dir(),
+        "Path '{}' is not a directory",
+        source_dir.display()
+    );
+
+    if !target_dir.exists() {
+        fs::create_dir_all(target_dir).await.with_context(|| {
+            format!(
+                "Failed to create the target dir at path '{}'",
+                target_dir.display()
+            )
+        })?;
+    }
+    let target_dir = target_dir.canonicalize().with_context(|| {
+        format!(
+            "Failed to get the absolute path for the target dir path '{}'",
+            target_dir.display()
+        )
+    })?;
+    ensure!(
+        target_dir.is_dir(),
+        "Path '{}' is not a directory",
+        target_dir.display()
+    );
+
+    println!(
+        "Compressing directory '{}' and creating resulting archive in directory '{}'",
+        source_dir.display(),
+        target_dir.display()
+    );
+
+    let mut metadata_file_contents = None;
+    let mut files_co_archive = Vec::new();
+
+    let mut source_dir_contents = fs::read_dir(&source_dir)
+        .await
+        .context("Failed to read the source directory contents")?;
+
+    while let Some(source_dir_entry) = source_dir_contents
+        .next_entry()
+        .await
+        .context("Failed to read a source dir entry")?
+    {
+        let entry_path = source_dir_entry.path();
+        if entry_path.is_file() {
+            if entry_path.file_name().and_then(|name| name.to_str()) == Some(METADATA_FILE_NAME) {
+                let metadata_bytes = fs::read(entry_path)
+                    .await
+                    .context("Failed to read metata file bytes in the source dir")?;
+                metadata_file_contents = Some(
+                    TimelineMetadata::from_bytes(&metadata_bytes)
+                        .context("Failed to parse metata file contents in the source dir")?,
+                );
+            } else {
+                files_co_archive.push(entry_path);
+            }
+        }
+    }
+
+    let metadata = match metadata_file_contents {
+        Some(metadata) => metadata,
+        None => bail!(
+            "No metadata file found in the source dir '{}', cannot create the archive",
+            source_dir.display()
+        ),
+    };
+
+    let _ = compression::archive_files_as_stream(
+        &source_dir,
+        files_co_archive.iter(),
+        &metadata,
+        move |mut archive_streamer, archive_name| async move {
+            let archive_target = target_dir.join(&archive_name);
+            let mut archive_file = fs::File::create(&archive_target).await?;
+            io::copy(&mut archive_streamer, &mut archive_file).await?;
+            Ok(archive_target)
+        },
+    )
+    .await
+    .context("Failed to create an archive")?;
+
+    Ok(())
+}

pageserver/src/bin/update_metadata.rs

@@ -14,20 +14,20 @@ fn main() -> Result<()> {
         .about("Dump or update metadata file")
         .version(GIT_VERSION)
         .arg(
-            Arg::with_name("path")
+            Arg::new("path")
                 .help("Path to metadata file")
                 .required(true),
         )
         .arg(
-            Arg::with_name("disk_lsn")
-                .short("d")
+            Arg::new("disk_lsn")
+                .short('d')
                 .long("disk_lsn")
                 .takes_value(true)
                 .help("Replace disk constistent lsn"),
         )
         .arg(
-            Arg::with_name("prev_lsn")
-                .short("p")
+            Arg::new("prev_lsn")
+                .short('p')
                 .long("prev_lsn")
                 .takes_value(true)
                 .help("Previous record LSN"),

pageserver/src/config.rs

@@ -8,7 +8,7 @@ use anyhow::{bail, ensure, Context, Result};
 use toml_edit;
 use toml_edit::{Document, Item};
 use zenith_utils::postgres_backend::AuthType;
-use zenith_utils::zid::{ZTenantId, ZTimelineId};
+use zenith_utils::zid::{ZNodeId, ZTenantId, ZTimelineId};
 
 use std::convert::TryInto;
 use std::env;
@@ -72,6 +72,10 @@ pub mod defaults {
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct PageServerConf {
+    // Identifier of that particular pageserver so e g safekeepers
+    // can safely distinguish different pageservers
+    pub id: ZNodeId,
+
     /// Example (default): 127.0.0.1:64000
     pub listen_pg_addr: String,
     /// Example (default): 127.0.0.1:9898
@@ -106,6 +110,184 @@ pub struct PageServerConf {
     pub remote_storage_config: Option<RemoteStorageConfig>,
 }
 
+// use dedicated enum for builder to better indicate the intention
+// and avoid possible confusion with nested options
+pub enum BuilderValue<T> {
+    Set(T),
+    NotSet,
+}
+
+impl<T> BuilderValue<T> {
+    pub fn ok_or<E>(self, err: E) -> Result<T, E> {
+        match self {
+            Self::Set(v) => Ok(v),
+            Self::NotSet => Err(err),
+        }
+    }
+}
+
+// needed to simplify config construction
+struct PageServerConfigBuilder {
+    listen_pg_addr: BuilderValue<String>,
+
+    listen_http_addr: BuilderValue<String>,
+
+    checkpoint_distance: BuilderValue<u64>,
+    checkpoint_period: BuilderValue<Duration>,
+
+    gc_horizon: BuilderValue<u64>,
+    gc_period: BuilderValue<Duration>,
+    superuser: BuilderValue<String>,
+
+    page_cache_size: BuilderValue<usize>,
+    max_file_descriptors: BuilderValue<usize>,
+
+    workdir: BuilderValue<PathBuf>,
+
+    pg_distrib_dir: BuilderValue<PathBuf>,
+
+    auth_type: BuilderValue<AuthType>,
+
+    //
+    auth_validation_public_key_path: BuilderValue<Option<PathBuf>>,
+    remote_storage_config: BuilderValue<Option<RemoteStorageConfig>>,
+
+    id: BuilderValue<ZNodeId>,
+}
+
+impl Default for PageServerConfigBuilder {
+    fn default() -> Self {
+        use self::BuilderValue::*;
+        use defaults::*;
+        Self {
+            listen_pg_addr: Set(DEFAULT_PG_LISTEN_ADDR.to_string()),
+            listen_http_addr: Set(DEFAULT_HTTP_LISTEN_ADDR.to_string()),
+            checkpoint_distance: Set(DEFAULT_CHECKPOINT_DISTANCE),
+            checkpoint_period: Set(humantime::parse_duration(DEFAULT_CHECKPOINT_PERIOD)
+                .expect("cannot parse default checkpoint period")),
+            gc_horizon: Set(DEFAULT_GC_HORIZON),
+            gc_period: Set(humantime::parse_duration(DEFAULT_GC_PERIOD)
+                .expect("cannot parse default gc period")),
+            superuser: Set(DEFAULT_SUPERUSER.to_string()),
+            page_cache_size: Set(DEFAULT_PAGE_CACHE_SIZE),
+            max_file_descriptors: Set(DEFAULT_MAX_FILE_DESCRIPTORS),
+            workdir: Set(PathBuf::new()),
+            pg_distrib_dir: Set(env::current_dir()
+                .expect("cannot access current directory")
+                .join("tmp_install")),
+            auth_type: Set(AuthType::Trust),
+            auth_validation_public_key_path: Set(None),
+            remote_storage_config: Set(None),
+            id: NotSet,
+        }
+    }
+}
+
+impl PageServerConfigBuilder {
+    pub fn listen_pg_addr(&mut self, listen_pg_addr: String) {
+        self.listen_pg_addr = BuilderValue::Set(listen_pg_addr)
+    }
+
+    pub fn listen_http_addr(&mut self, listen_http_addr: String) {
+        self.listen_http_addr = BuilderValue::Set(listen_http_addr)
+    }
+
+    pub fn checkpoint_distance(&mut self, checkpoint_distance: u64) {
+        self.checkpoint_distance = BuilderValue::Set(checkpoint_distance)
+    }
+
+    pub fn checkpoint_period(&mut self, checkpoint_period: Duration) {
+        self.checkpoint_period = BuilderValue::Set(checkpoint_period)
+    }
+
+    pub fn gc_horizon(&mut self, gc_horizon: u64) {
+        self.gc_horizon = BuilderValue::Set(gc_horizon)
+    }
+
+    pub fn gc_period(&mut self, gc_period: Duration) {
+        self.gc_period = BuilderValue::Set(gc_period)
+    }
+
+    pub fn superuser(&mut self, superuser: String) {
+        self.superuser = BuilderValue::Set(superuser)
+    }
+
+    pub fn page_cache_size(&mut self, page_cache_size: usize) {
+        self.page_cache_size = BuilderValue::Set(page_cache_size)
+    }
+
+    pub fn max_file_descriptors(&mut self, max_file_descriptors: usize) {
+        self.max_file_descriptors = BuilderValue::Set(max_file_descriptors)
+    }
+
+    pub fn workdir(&mut self, workdir: PathBuf) {
+        self.workdir = BuilderValue::Set(workdir)
+    }
+
+    pub fn pg_distrib_dir(&mut self, pg_distrib_dir: PathBuf) {
+        self.pg_distrib_dir = BuilderValue::Set(pg_distrib_dir)
+    }
+
+    pub fn auth_type(&mut self, auth_type: AuthType) {
+        self.auth_type = BuilderValue::Set(auth_type)
+    }
+
+    pub fn auth_validation_public_key_path(
+        &mut self,
+        auth_validation_public_key_path: Option<PathBuf>,
+    ) {
+        self.auth_validation_public_key_path = BuilderValue::Set(auth_validation_public_key_path)
+    }
+
+    pub fn remote_storage_config(&mut self, remote_storage_config: Option<RemoteStorageConfig>) {
+        self.remote_storage_config = BuilderValue::Set(remote_storage_config)
+    }
+
+    pub fn id(&mut self, node_id: ZNodeId) {
+        self.id = BuilderValue::Set(node_id)
+    }
+
+    pub fn build(self) -> Result<PageServerConf> {
+        Ok(PageServerConf {
+            listen_pg_addr: self
+                .listen_pg_addr
+                .ok_or(anyhow::anyhow!("missing listen_pg_addr"))?,
+            listen_http_addr: self
+                .listen_http_addr
+                .ok_or(anyhow::anyhow!("missing listen_http_addr"))?,
+            checkpoint_distance: self
+                .checkpoint_distance
+                .ok_or(anyhow::anyhow!("missing checkpoint_distance"))?,
+            checkpoint_period: self
+                .checkpoint_period
+                .ok_or(anyhow::anyhow!("missing checkpoint_period"))?,
+            gc_horizon: self
+                .gc_horizon
+                .ok_or(anyhow::anyhow!("missing gc_horizon"))?,
+            gc_period: self.gc_period.ok_or(anyhow::anyhow!("missing gc_period"))?,
+            superuser: self.superuser.ok_or(anyhow::anyhow!("missing superuser"))?,
+            page_cache_size: self
+                .page_cache_size
+                .ok_or(anyhow::anyhow!("missing page_cache_size"))?,
+            max_file_descriptors: self
+                .max_file_descriptors
+                .ok_or(anyhow::anyhow!("missing max_file_descriptors"))?,
+            workdir: self.workdir.ok_or(anyhow::anyhow!("missing workdir"))?,
+            pg_distrib_dir: self
+                .pg_distrib_dir
+                .ok_or(anyhow::anyhow!("missing pg_distrib_dir"))?,
+            auth_type: self.auth_type.ok_or(anyhow::anyhow!("missing auth_type"))?,
+            auth_validation_public_key_path: self
+                .auth_validation_public_key_path
+                .ok_or(anyhow::anyhow!("missing auth_validation_public_key_path"))?,
+            remote_storage_config: self
+                .remote_storage_config
+                .ok_or(anyhow::anyhow!("missing remote_storage_config"))?,
+            id: self.id.ok_or(anyhow::anyhow!("missing id"))?,
+        })
+    }
+}
+
 /// External backup storage configuration, enough for creating a client for that storage.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct RemoteStorageConfig {
@@ -221,57 +403,39 @@ impl PageServerConf {
     ///
     /// This leaves any options not present in the file in the built-in defaults.
     pub fn parse_and_validate(toml: &Document, workdir: &Path) -> Result<Self> {
-        use defaults::*;
-
-        let mut conf = PageServerConf {
-            workdir: workdir.to_path_buf(),
-
-            listen_pg_addr: DEFAULT_PG_LISTEN_ADDR.to_string(),
-            listen_http_addr: DEFAULT_HTTP_LISTEN_ADDR.to_string(),
-            checkpoint_distance: DEFAULT_CHECKPOINT_DISTANCE,
-            checkpoint_period: humantime::parse_duration(DEFAULT_CHECKPOINT_PERIOD)?,
-            gc_horizon: DEFAULT_GC_HORIZON,
-            gc_period: humantime::parse_duration(DEFAULT_GC_PERIOD)?,
-            page_cache_size: DEFAULT_PAGE_CACHE_SIZE,
-            max_file_descriptors: DEFAULT_MAX_FILE_DESCRIPTORS,
-
-            pg_distrib_dir: PathBuf::new(),
-            auth_validation_public_key_path: None,
-            auth_type: AuthType::Trust,
-
-            remote_storage_config: None,
-
-            superuser: DEFAULT_SUPERUSER.to_string(),
-        };
+        let mut builder = PageServerConfigBuilder::default();
+        builder.workdir(workdir.to_owned());
 
         for (key, item) in toml.iter() {
             match key {
-                "listen_pg_addr" => conf.listen_pg_addr = parse_toml_string(key, item)?,
-                "listen_http_addr" => conf.listen_http_addr = parse_toml_string(key, item)?,
-                "checkpoint_distance" => conf.checkpoint_distance = parse_toml_u64(key, item)?,
-                "checkpoint_period" => conf.checkpoint_period = parse_toml_duration(key, item)?,
-                "gc_horizon" => conf.gc_horizon = parse_toml_u64(key, item)?,
-                "gc_period" => conf.gc_period = parse_toml_duration(key, item)?,
-                "initial_superuser_name" => conf.superuser = parse_toml_string(key, item)?,
-                "page_cache_size" => conf.page_cache_size = parse_toml_u64(key, item)? as usize,
+                "listen_pg_addr" => builder.listen_pg_addr(parse_toml_string(key, item)?),
+                "listen_http_addr" => builder.listen_http_addr(parse_toml_string(key, item)?),
+                "checkpoint_distance" => builder.checkpoint_distance(parse_toml_u64(key, item)?),
+                "checkpoint_period" => builder.checkpoint_period(parse_toml_duration(key, item)?),
+                "gc_horizon" => builder.gc_horizon(parse_toml_u64(key, item)?),
+                "gc_period" => builder.gc_period(parse_toml_duration(key, item)?),
+                "initial_superuser_name" => builder.superuser(parse_toml_string(key, item)?),
+                "page_cache_size" => builder.page_cache_size(parse_toml_u64(key, item)? as usize),
                 "max_file_descriptors" => {
-                    conf.max_file_descriptors = parse_toml_u64(key, item)? as usize
+                    builder.max_file_descriptors(parse_toml_u64(key, item)? as usize)
                 }
                 "pg_distrib_dir" => {
-                    conf.pg_distrib_dir = PathBuf::from(parse_toml_string(key, item)?)
+                    builder.pg_distrib_dir(PathBuf::from(parse_toml_string(key, item)?))
                 }
-                "auth_validation_public_key_path" => {
-                    conf.auth_validation_public_key_path =
-                        Some(PathBuf::from(parse_toml_string(key, item)?))
-                }
-                "auth_type" => conf.auth_type = parse_toml_auth_type(key, item)?,
+                "auth_validation_public_key_path" => builder.auth_validation_public_key_path(Some(
+                    PathBuf::from(parse_toml_string(key, item)?),
+                )),
+                "auth_type" => builder.auth_type(parse_toml_auth_type(key, item)?),
                 "remote_storage" => {
-                    conf.remote_storage_config = Some(Self::parse_remote_storage_config(item)?)
+                    builder.remote_storage_config(Some(Self::parse_remote_storage_config(item)?))
                 }
+                "id" => builder.id(ZNodeId(parse_toml_u64(key, item)?)),
                 _ => bail!("unrecognized pageserver option '{}'", key),
             }
         }
 
+        let mut conf = builder.build().context("invalid config")?;
+
         if conf.auth_type == AuthType::ZenithJWT {
             let auth_validation_public_key_path = conf
                 .auth_validation_public_key_path
@@ -285,9 +449,6 @@ impl PageServerConf {
             );
         }
 
-        if conf.pg_distrib_dir == PathBuf::new() {
-            conf.pg_distrib_dir = env::current_dir()?.join("tmp_install")
-        };
         if !conf.pg_distrib_dir.join("bin/postgres").exists() {
             bail!(
                 "Can't find postgres binary at {}",
@@ -382,6 +543,7 @@ impl PageServerConf {
     #[cfg(test)]
     pub fn dummy_conf(repo_dir: PathBuf) -> Self {
         PageServerConf {
+            id: ZNodeId(0),
             checkpoint_distance: defaults::DEFAULT_CHECKPOINT_DISTANCE,
             checkpoint_period: Duration::from_secs(10),
             gc_horizon: defaults::DEFAULT_GC_HORIZON,
@@ -461,15 +623,16 @@ max_file_descriptors = 333
 
 # initial superuser role name to use when creating a new tenant
 initial_superuser_name = 'zzzz'
+id = 10
 
-    "#;
+"#;
 
     #[test]
     fn parse_defaults() -> anyhow::Result<()> {
         let tempdir = tempdir()?;
         let (workdir, pg_distrib_dir) = prepare_fs(&tempdir)?;
         // we have to create dummy pathes to overcome the validation errors
-        let config_string = format!("pg_distrib_dir='{}'", pg_distrib_dir.display());
+        let config_string = format!("pg_distrib_dir='{}'\nid=10", pg_distrib_dir.display());
         let toml = config_string.parse()?;
 
         let parsed_config =
@@ -480,6 +643,7 @@ initial_superuser_name = 'zzzz'
         assert_eq!(
             parsed_config,
             PageServerConf {
+                id: ZNodeId(10),
                 listen_pg_addr: defaults::DEFAULT_PG_LISTEN_ADDR.to_string(),
                 listen_http_addr: defaults::DEFAULT_HTTP_LISTEN_ADDR.to_string(),
                 checkpoint_distance: defaults::DEFAULT_CHECKPOINT_DISTANCE,
@@ -521,6 +685,7 @@ initial_superuser_name = 'zzzz'
         assert_eq!(
             parsed_config,
             PageServerConf {
+                id: ZNodeId(10),
                 listen_pg_addr: "127.0.0.1:64000".to_string(),
                 listen_http_addr: "127.0.0.1:9898".to_string(),
                 checkpoint_distance: 111,

pageserver/src/http/models.rs

@@ -1,6 +1,7 @@
 use serde::{Deserialize, Serialize};
 
 use crate::ZTenantId;
+use zenith_utils::zid::ZNodeId;
 
 #[derive(Serialize, Deserialize)]
 pub struct BranchCreateRequest {
@@ -15,3 +16,8 @@ pub struct TenantCreateRequest {
     #[serde(with = "hex")]
     pub tenant_id: ZTenantId,
 }
+
+#[derive(Serialize)]
+pub struct StatusResponse {
+    pub id: ZNodeId,
+}

pageserver/src/http/openapi_spec.yml

@@ -17,6 +17,11 @@ paths:
             application/json:
               schema:
                 type: object
+                required:
+                - id
+                properties:
+                  id:
+                    type: integer
   /v1/timeline/{tenant_id}:
     parameters:
       - name: tenant_id
@@ -234,9 +239,7 @@ paths:
           content:
             application/json:
               schema:
-                type: array
-                items:
-                  $ref: "#/components/schemas/BranchInfo"
+                $ref: "#/components/schemas/BranchInfo"
         "400":
           description: Malformed branch create request
           content:
@@ -370,12 +373,15 @@ components:
           format: hex
         ancestor_id:
           type: string
+          format: hex
         ancestor_lsn:
           type: string
         current_logical_size:
           type: integer
         current_logical_size_non_incremental:
           type: integer
+        latest_valid_lsn:
+          type: integer
     TimelineInfo:
       type: object
       required:

pageserver/src/http/routes.rs

@@ -1,10 +1,8 @@
 use std::sync::Arc;
 
 use anyhow::{Context, Result};
-use hyper::header;
 use hyper::StatusCode;
 use hyper::{Body, Request, Response, Uri};
-use routerify::{ext::RequestExt, RouterBuilder};
 use serde::Serialize;
 use tracing::*;
 use zenith_utils::auth::JwtAuth;
@@ -19,10 +17,12 @@ use zenith_utils::http::{
     request::get_request_param,
     request::parse_request_param,
 };
+use zenith_utils::http::{RequestExt, RouterBuilder};
 use zenith_utils::lsn::Lsn;
 use zenith_utils::zid::{opt_display_serde, ZTimelineId};
 
 use super::models::BranchCreateRequest;
+use super::models::StatusResponse;
 use super::models::TenantCreateRequest;
 use crate::branches::BranchInfo;
 use crate::repository::RepositoryTimeline;
@@ -64,12 +64,12 @@ fn get_config(request: &Request<Body>) -> &'static PageServerConf {
 }
 
 // healthcheck handler
-async fn status_handler(_: Request<Body>) -> Result<Response<Body>, ApiError> {
-    Ok(Response::builder()
-        .status(StatusCode::OK)
-        .header(header::CONTENT_TYPE, "application/json")
-        .body(Body::from("{}"))
-        .map_err(ApiError::from_err)?)
+async fn status_handler(request: Request<Body>) -> Result<Response<Body>, ApiError> {
+    let config = get_config(&request);
+    Ok(json_response(
+        StatusCode::OK,
+        StatusResponse { id: config.id },
+    )?)
 }
 
 async fn branch_create_handler(mut request: Request<Body>) -> Result<Response<Body>, ApiError> {

pageserver/src/layered_repository/ephemeral_file.rs

@@ -175,7 +175,10 @@ impl Write for EphemeralFile {
     }
 
     fn flush(&mut self) -> Result<(), std::io::Error> {
-        todo!()
+        // we don't need to flush data:
+        // * we either write input bytes or not, not keeping any intermediate data buffered
+        // * rust unix file `flush` impl does not flush things either, returning `Ok(())`
+        Ok(())
     }
 }
 

pageserver/src/page_service.rs

@@ -27,13 +27,10 @@ use zenith_utils::lsn::Lsn;
 use zenith_utils::postgres_backend::is_socket_read_timed_out;
 use zenith_utils::postgres_backend::PostgresBackend;
 use zenith_utils::postgres_backend::{self, AuthType};
-use zenith_utils::pq_proto::{
-    BeMessage, FeMessage, RowDescriptor, HELLO_WORLD_ROW, SINGLE_COL_ROWDESC,
-};
+use zenith_utils::pq_proto::{BeMessage, FeMessage, RowDescriptor, SINGLE_COL_ROWDESC};
 use zenith_utils::zid::{ZTenantId, ZTimelineId};
 
 use crate::basebackup;
-use crate::branches;
 use crate::config::PageServerConf;
 use crate::relish::*;
 use crate::repository::Timeline;
@@ -662,79 +659,21 @@ impl postgres_backend::Handler for PageServerHandler {
             walreceiver::launch_wal_receiver(self.conf, tenantid, timelineid, &connstr)?;
 
             pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
-        } else if query_string.starts_with("branch_create ") {
-            let err = || format!("invalid branch_create: '{}'", query_string);
-
-            // branch_create <tenantid> <branchname> <startpoint>
-            // TODO lazy static
-            // TODO: escaping, to allow branch names with spaces
-            let re = Regex::new(r"^branch_create ([[:xdigit:]]+) (\S+) ([^\r\n\s;]+)[\r\n\s;]*;?$")
-                .unwrap();
-            let caps = re.captures(query_string).with_context(err)?;
-
-            let tenantid = ZTenantId::from_str(caps.get(1).unwrap().as_str())?;
-            let branchname = caps.get(2).with_context(err)?.as_str().to_owned();
-            let startpoint_str = caps.get(3).with_context(err)?.as_str().to_owned();
-
-            self.check_permission(Some(tenantid))?;
-
-            let _enter =
-                info_span!("branch_create", name = %branchname, tenant = %tenantid).entered();
-
-            let branch =
-                branches::create_branch(self.conf, &branchname, &startpoint_str, &tenantid)?;
-            let branch = serde_json::to_vec(&branch)?;
-
-            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
-                .write_message_noflush(&BeMessage::DataRow(&[Some(&branch)]))?
-                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
-        } else if query_string.starts_with("branch_list ") {
-            // branch_list <zenith tenantid as hex string>
-            let re = Regex::new(r"^branch_list ([[:xdigit:]]+)$").unwrap();
-            let caps = re
-                .captures(query_string)
-                .with_context(|| format!("invalid branch_list: '{}'", query_string))?;
-
-            let tenantid = ZTenantId::from_str(caps.get(1).unwrap().as_str())?;
-
-            // since these handlers for tenant/branch commands are deprecated (in favor of http based ones)
-            // just use false in place of include non incremental logical size
-            let branches = crate::branches::get_branches(self.conf, &tenantid, false)?;
-            let branches_buf = serde_json::to_vec(&branches)?;
-
-            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
-                .write_message_noflush(&BeMessage::DataRow(&[Some(&branches_buf)]))?
-                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
-        } else if query_string.starts_with("tenant_list") {
-            let tenants = crate::tenant_mgr::list_tenants()?;
-            let tenants_buf = serde_json::to_vec(&tenants)?;
-
-            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
-                .write_message_noflush(&BeMessage::DataRow(&[Some(&tenants_buf)]))?
-                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
-        } else if query_string.starts_with("tenant_create") {
-            let err = || format!("invalid tenant_create: '{}'", query_string);
-
-            // tenant_create <tenantid>
-            let re = Regex::new(r"^tenant_create ([[:xdigit:]]+)$").unwrap();
-            let caps = re.captures(query_string).with_context(err)?;
-
-            self.check_permission(None)?;
-
-            let tenantid = ZTenantId::from_str(caps.get(1).unwrap().as_str())?;
-
-            tenant_mgr::create_repository_for_tenant(self.conf, tenantid)?;
-
-            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
-                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
-        } else if query_string.starts_with("status") {
-            pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
-                .write_message_noflush(&HELLO_WORLD_ROW)?
-                .write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
         } else if query_string.to_ascii_lowercase().starts_with("set ") {
             // important because psycopg2 executes "SET datestyle TO 'ISO'"
             // on connect
             pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
+        } else if query_string.starts_with("failpoints ") {
+            let (_, failpoints) = query_string.split_at("failpoints ".len());
+            for failpoint in failpoints.split(';') {
+                if let Some((name, actions)) = failpoint.split_once('=') {
+                    info!("cfg failpoint: {} {}", name, actions);
+                    fail::cfg(name, actions).unwrap();
+                } else {
+                    bail!("Invalid failpoints format");
+                }
+            }
+            pgb.write_message_noflush(&BeMessage::CommandComplete(b"SELECT 1"))?;
         } else if query_string.starts_with("do_gc ") {
             // Run GC immediately on given timeline.
             // FIXME: This is just for tests. See test_runner/batch_others/test_gc.py.

pageserver/src/remote_storage.rs

@@ -94,7 +94,7 @@ use std::{
 use anyhow::{bail, Context};
 use tokio::io;
 use tracing::{error, info};
-use zenith_utils::zid::{ZTenantId, ZTimelineId};
+use zenith_utils::zid::{ZTenantId, ZTenantTimelineId, ZTimelineId};
 
 pub use self::storage_sync::{schedule_timeline_checkpoint_upload, schedule_timeline_download};
 use self::{local_fs::LocalFs, rust_s3::S3};
@@ -104,16 +104,7 @@ use crate::{
     repository::TimelineSyncState,
 };
 
-/// Any timeline has its own id and its own tenant it belongs to,
-/// the sync processes group timelines by both for simplicity.
-#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
-pub struct TimelineSyncId(ZTenantId, ZTimelineId);
-
-impl std::fmt::Display for TimelineSyncId {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "(tenant: {}, timeline: {})", self.0, self.1)
-    }
-}
+pub use storage_sync::compression;
 
 /// A structure to combine all synchronization data to share with pageserver after a successful sync loop initialization.
 /// Successful initialization includes a case when sync loop is not started, in which case the startup data is returned still,
@@ -167,7 +158,7 @@ pub fn start_local_timeline_sync(
                 ZTenantId,
                 HashMap<ZTimelineId, TimelineSyncState>,
             > = HashMap::new();
-            for (TimelineSyncId(tenant_id, timeline_id), (timeline_metadata, _)) in
+            for (ZTenantTimelineId{tenant_id, timeline_id}, (timeline_metadata, _)) in
                 local_timeline_files
             {
                 initial_timeline_states
@@ -187,7 +178,7 @@ pub fn start_local_timeline_sync(
 
 fn local_tenant_timeline_files(
     config: &'static PageServerConf,
-) -> anyhow::Result<HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)>> {
+) -> anyhow::Result<HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)>> {
     let mut local_tenant_timeline_files = HashMap::new();
     let tenants_dir = config.tenants_path();
     for tenants_dir_entry in fs::read_dir(&tenants_dir)
@@ -222,8 +213,9 @@ fn local_tenant_timeline_files(
 fn collect_timelines_for_tenant(
     config: &'static PageServerConf,
     tenant_path: &Path,
-) -> anyhow::Result<HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)>> {
-    let mut timelines: HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)> = HashMap::new();
+) -> anyhow::Result<HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)>> {
+    let mut timelines: HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)> =
+        HashMap::new();
     let tenant_id = tenant_path
         .file_name()
         .and_then(ffi::OsStr::to_str)
@@ -244,7 +236,10 @@ fn collect_timelines_for_tenant(
                 match collect_timeline_files(&timeline_path) {
                     Ok((timeline_id, metadata, timeline_files)) => {
                         timelines.insert(
-                            TimelineSyncId(tenant_id, timeline_id),
+                            ZTenantTimelineId {
+                                tenant_id,
+                                timeline_id,
+                            },
                             (metadata, timeline_files),
                         );
                     }

pageserver/src/remote_storage/storage_sync.rs

@@ -70,7 +70,8 @@
 //!
 //! When pageserver signals shutdown, current sync task gets finished and the loop exists.
 
-mod compression;
+/// Expose the module for a binary CLI tool that deals with the corresponding blobs.
+pub mod compression;
 mod download;
 pub mod index;
 mod upload;
@@ -105,7 +106,7 @@ use self::{
     },
     upload::upload_timeline_checkpoint,
 };
-use super::{RemoteStorage, SyncStartupData, TimelineSyncId};
+use super::{RemoteStorage, SyncStartupData, ZTenantTimelineId};
 use crate::{
     config::PageServerConf, layered_repository::metadata::TimelineMetadata,
     remote_storage::storage_sync::compression::read_archive_header, repository::TimelineSyncState,
@@ -242,13 +243,13 @@ mod sync_queue {
 /// Limited by the number of retries, after certain threshold the failing task gets evicted and the timeline disabled.
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
 pub struct SyncTask {
-    sync_id: TimelineSyncId,
+    sync_id: ZTenantTimelineId,
     retries: u32,
     kind: SyncKind,
 }
 
 impl SyncTask {
-    fn new(sync_id: TimelineSyncId, retries: u32, kind: SyncKind) -> Self {
+    fn new(sync_id: ZTenantTimelineId, retries: u32, kind: SyncKind) -> Self {
         Self {
             sync_id,
             retries,
@@ -307,7 +308,10 @@ pub fn schedule_timeline_checkpoint_upload(
     }
 
     if !sync_queue::push(SyncTask::new(
-        TimelineSyncId(tenant_id, timeline_id),
+        ZTenantTimelineId {
+            tenant_id,
+            timeline_id,
+        },
         0,
         SyncKind::Upload(NewCheckpoint { layers, metadata }),
     )) {
@@ -338,7 +342,10 @@ pub fn schedule_timeline_download(tenant_id: ZTenantId, timeline_id: ZTimelineId
         tenant_id, timeline_id
     );
     sync_queue::push(SyncTask::new(
-        TimelineSyncId(tenant_id, timeline_id),
+        ZTenantTimelineId {
+            tenant_id,
+            timeline_id,
+        },
         0,
         SyncKind::Download(TimelineDownload {
             files_to_skip: Arc::new(BTreeSet::new()),
@@ -354,7 +361,7 @@ pub(super) fn spawn_storage_sync_thread<
     S: RemoteStorage<StoragePath = P> + Send + Sync + 'static,
 >(
     conf: &'static PageServerConf,
-    local_timeline_files: HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)>,
+    local_timeline_files: HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)>,
     storage: S,
     max_concurrent_sync: NonZeroUsize,
     max_sync_errors: NonZeroU32,
@@ -510,7 +517,7 @@ async fn loop_step<
                 Err(e) => {
                     error!(
                         "Failed to process storage sync task for tenant {}, timeline {}: {:?}",
-                        sync_id.0, sync_id.1, e
+                        sync_id.tenant_id, sync_id.timeline_id, e
                     );
                     None
                 }
@@ -524,7 +531,10 @@ async fn loop_step<
     while let Some((sync_id, state_update)) = task_batch.next().await {
         debug!("Finished storage sync task for sync id {}", sync_id);
         if let Some(state_update) = state_update {
-            let TimelineSyncId(tenant_id, timeline_id) = sync_id;
+            let ZTenantTimelineId {
+                tenant_id,
+                timeline_id,
+            } = sync_id;
             new_timeline_states
                 .entry(tenant_id)
                 .or_default()
@@ -618,7 +628,7 @@ async fn process_task<
 
 fn schedule_first_sync_tasks(
     index: &RemoteTimelineIndex,
-    local_timeline_files: HashMap<TimelineSyncId, (TimelineMetadata, Vec<PathBuf>)>,
+    local_timeline_files: HashMap<ZTenantTimelineId, (TimelineMetadata, Vec<PathBuf>)>,
 ) -> HashMap<ZTenantId, HashMap<ZTimelineId, TimelineSyncState>> {
     let mut initial_timeline_statuses: HashMap<ZTenantId, HashMap<ZTimelineId, TimelineSyncState>> =
         HashMap::new();
@@ -629,7 +639,10 @@ fn schedule_first_sync_tasks(
     for (sync_id, (local_metadata, local_files)) in local_timeline_files {
         let local_disk_consistent_lsn = local_metadata.disk_consistent_lsn();
 
-        let TimelineSyncId(tenant_id, timeline_id) = sync_id;
+        let ZTenantTimelineId {
+            tenant_id,
+            timeline_id,
+        } = sync_id;
         match index.timeline_entry(&sync_id) {
             Some(index_entry) => {
                 let timeline_status = compare_local_and_remote_timeline(
@@ -672,10 +685,10 @@ fn schedule_first_sync_tasks(
         }
     }
 
-    let unprocessed_remote_ids = |remote_id: &TimelineSyncId| {
+    let unprocessed_remote_ids = |remote_id: &ZTenantTimelineId| {
         initial_timeline_statuses
-            .get(&remote_id.0)
-            .and_then(|timelines| timelines.get(&remote_id.1))
+            .get(&remote_id.tenant_id)
+            .and_then(|timelines| timelines.get(&remote_id.timeline_id))
             .is_none()
     };
     for unprocessed_remote_id in index
@@ -683,7 +696,10 @@ fn schedule_first_sync_tasks(
         .filter(unprocessed_remote_ids)
         .collect::<Vec<_>>()
     {
-        let TimelineSyncId(cloud_only_tenant_id, cloud_only_timeline_id) = unprocessed_remote_id;
+        let ZTenantTimelineId {
+            tenant_id: cloud_only_tenant_id,
+            timeline_id: cloud_only_timeline_id,
+        } = unprocessed_remote_id;
         match index
             .timeline_entry(&unprocessed_remote_id)
             .and_then(TimelineIndexEntry::disk_consistent_lsn)
@@ -712,7 +728,7 @@ fn schedule_first_sync_tasks(
 
 fn compare_local_and_remote_timeline(
     new_sync_tasks: &mut VecDeque<SyncTask>,
-    sync_id: TimelineSyncId,
+    sync_id: ZTenantTimelineId,
     local_metadata: TimelineMetadata,
     local_files: Vec<PathBuf>,
     remote_entry: &TimelineIndexEntry,
@@ -769,7 +785,7 @@ async fn update_index_description<
 >(
     (storage, index): &(S, RwLock<RemoteTimelineIndex>),
     timeline_dir: &Path,
-    id: TimelineSyncId,
+    id: ZTenantTimelineId,
 ) -> anyhow::Result<RemoteTimeline> {
     let mut index_write = index.write().await;
     let full_index = match index_write.timeline_entry(&id) {
@@ -792,7 +808,7 @@ async fn update_index_description<
                         Ok((archive_id, header_size, header)) => full_index.update_archive_contents(archive_id.0, header, header_size),
                         Err((e, archive_id)) => bail!(
                             "Failed to download archive header for tenant {}, timeline {}, archive for Lsn {}: {}",
-                            id.0, id.1, archive_id.0,
+                            id.tenant_id, id.timeline_id, archive_id.0,
                             e
                         ),
                     }
@@ -870,7 +886,7 @@ mod test_utils {
         timeline_id: ZTimelineId,
         new_upload: NewCheckpoint,
     ) {
-        let sync_id = TimelineSyncId(harness.tenant_id, timeline_id);
+        let sync_id = ZTenantTimelineId::new(harness.tenant_id, timeline_id);
         upload_timeline_checkpoint(
             harness.conf,
             Arc::clone(&remote_assets),
@@ -926,7 +942,7 @@ mod test_utils {
 
     pub async fn expect_timeline(
         index: &RwLock<RemoteTimelineIndex>,
-        sync_id: TimelineSyncId,
+        sync_id: ZTenantTimelineId,
     ) -> RemoteTimeline {
         if let Some(TimelineIndexEntry::Full(remote_timeline)) =
             index.read().await.timeline_entry(&sync_id)
@@ -961,18 +977,18 @@ mod test_utils {
         let mut expected_timeline_entries = BTreeMap::new();
         for sync_id in actual_sync_ids {
             actual_branches.insert(
-                sync_id.1,
+                sync_id.tenant_id,
                 index_read
-                    .branch_files(sync_id.0)
+                    .branch_files(sync_id.tenant_id)
                     .into_iter()
                     .flat_map(|branch_paths| branch_paths.iter())
                     .cloned()
                     .collect::<BTreeSet<_>>(),
             );
             expected_branches.insert(
-                sync_id.1,
+                sync_id.tenant_id,
                 expected_index_with_descriptions
-                    .branch_files(sync_id.0)
+                    .branch_files(sync_id.tenant_id)
                     .into_iter()
                     .flat_map(|branch_paths| branch_paths.iter())
                     .cloned()

pageserver/src/remote_storage/storage_sync/compression.rs

@@ -248,7 +248,7 @@ fn archive_name(disk_consistent_lsn: Lsn, header_size: u64) -> String {
     archive_name
 }
 
-async fn uncompress_with_header(
+pub async fn uncompress_with_header(
     files_to_skip: &BTreeSet<PathBuf>,
     destination_dir: &Path,
     header: ArchiveHeader,

pageserver/src/remote_storage/storage_sync/download.rs

@@ -17,7 +17,7 @@ use crate::{
             compression, index::TimelineIndexEntry, sync_queue, tenant_branch_files,
             update_index_description, SyncKind, SyncTask,
         },
-        RemoteStorage, TimelineSyncId,
+        RemoteStorage, ZTenantTimelineId,
     },
 };
 
@@ -52,13 +52,16 @@ pub(super) async fn download_timeline<
 >(
     conf: &'static PageServerConf,
     remote_assets: Arc<(S, RwLock<RemoteTimelineIndex>)>,
-    sync_id: TimelineSyncId,
+    sync_id: ZTenantTimelineId,
     mut download: TimelineDownload,
     retries: u32,
 ) -> DownloadedTimeline {
     debug!("Downloading layers for sync id {}", sync_id);
 
-    let TimelineSyncId(tenant_id, timeline_id) = sync_id;
+    let ZTenantTimelineId {
+        tenant_id,
+        timeline_id,
+    } = sync_id;
     let index_read = remote_assets.1.read().await;
     let remote_timeline = match index_read.timeline_entry(&sync_id) {
         None => {
@@ -110,7 +113,8 @@ pub(super) async fn download_timeline<
         }
     };
 
-    if let Err(e) = download_missing_branches(conf, remote_assets.as_ref(), sync_id.0).await {
+    if let Err(e) = download_missing_branches(conf, remote_assets.as_ref(), sync_id.tenant_id).await
+    {
         error!(
             "Failed to download missing branches for sync id {}: {:?}",
             sync_id, e
@@ -180,7 +184,10 @@ async fn try_download_archive<
     S: RemoteStorage<StoragePath = P> + Send + Sync + 'static,
 >(
     conf: &'static PageServerConf,
-    TimelineSyncId(tenant_id, timeline_id): TimelineSyncId,
+    ZTenantTimelineId {
+        tenant_id,
+        timeline_id,
+    }: ZTenantTimelineId,
     remote_assets: Arc<(S, RwLock<RemoteTimelineIndex>)>,
     remote_timeline: &RemoteTimeline,
     archive_id: ArchiveId,
@@ -343,7 +350,7 @@ mod tests {
     #[tokio::test]
     async fn test_download_timeline() -> anyhow::Result<()> {
         let repo_harness = RepoHarness::create("test_download_timeline")?;
-        let sync_id = TimelineSyncId(repo_harness.tenant_id, TIMELINE_ID);
+        let sync_id = ZTenantTimelineId::new(repo_harness.tenant_id, TIMELINE_ID);
         let storage = LocalFs::new(tempdir()?.path().to_owned(), &repo_harness.conf.workdir)?;
         let index = RwLock::new(RemoteTimelineIndex::try_parse_descriptions_from_paths(
             repo_harness.conf,

pageserver/src/remote_storage/storage_sync/index.rs

@@ -22,7 +22,7 @@ use crate::{
     layered_repository::TIMELINES_SEGMENT_NAME,
     remote_storage::{
         storage_sync::compression::{parse_archive_name, FileEntry},
-        TimelineSyncId,
+        ZTenantTimelineId,
     },
 };
 
@@ -53,7 +53,7 @@ impl RelativePath {
 #[derive(Debug, Clone)]
 pub struct RemoteTimelineIndex {
     branch_files: HashMap<ZTenantId, HashSet<RelativePath>>,
-    timeline_files: HashMap<TimelineSyncId, TimelineIndexEntry>,
+    timeline_files: HashMap<ZTenantTimelineId, TimelineIndexEntry>,
 }
 
 impl RemoteTimelineIndex {
@@ -80,19 +80,22 @@ impl RemoteTimelineIndex {
         index
     }
 
-    pub fn timeline_entry(&self, id: &TimelineSyncId) -> Option<&TimelineIndexEntry> {
+    pub fn timeline_entry(&self, id: &ZTenantTimelineId) -> Option<&TimelineIndexEntry> {
         self.timeline_files.get(id)
     }
 
-    pub fn timeline_entry_mut(&mut self, id: &TimelineSyncId) -> Option<&mut TimelineIndexEntry> {
+    pub fn timeline_entry_mut(
+        &mut self,
+        id: &ZTenantTimelineId,
+    ) -> Option<&mut TimelineIndexEntry> {
         self.timeline_files.get_mut(id)
     }
 
-    pub fn add_timeline_entry(&mut self, id: TimelineSyncId, entry: TimelineIndexEntry) {
+    pub fn add_timeline_entry(&mut self, id: ZTenantTimelineId, entry: TimelineIndexEntry) {
         self.timeline_files.insert(id, entry);
     }
 
-    pub fn all_sync_ids(&self) -> impl Iterator<Item = TimelineSyncId> + '_ {
+    pub fn all_sync_ids(&self) -> impl Iterator<Item = ZTenantTimelineId> + '_ {
         self.timeline_files.keys().copied()
     }
 
@@ -348,7 +351,10 @@ fn try_parse_index_entry(
                 .to_string_lossy()
                 .to_string();
 
-            let sync_id = TimelineSyncId(tenant_id, timeline_id);
+            let sync_id = ZTenantTimelineId {
+                tenant_id,
+                timeline_id,
+            };
             let timeline_index_entry = index
                 .timeline_files
                 .entry(sync_id)

pageserver/src/remote_storage/storage_sync/upload.rs

@@ -17,7 +17,7 @@ use crate::{
             index::{RemoteTimeline, TimelineIndexEntry},
             sync_queue, tenant_branch_files, update_index_description, SyncKind, SyncTask,
         },
-        RemoteStorage, TimelineSyncId,
+        RemoteStorage, ZTenantTimelineId,
     },
 };
 
@@ -36,12 +36,13 @@ pub(super) async fn upload_timeline_checkpoint<
 >(
     config: &'static PageServerConf,
     remote_assets: Arc<(S, RwLock<RemoteTimelineIndex>)>,
-    sync_id: TimelineSyncId,
+    sync_id: ZTenantTimelineId,
     new_checkpoint: NewCheckpoint,
     retries: u32,
 ) -> Option<bool> {
     debug!("Uploading checkpoint for sync id {}", sync_id);
-    if let Err(e) = upload_missing_branches(config, remote_assets.as_ref(), sync_id.0).await {
+    if let Err(e) = upload_missing_branches(config, remote_assets.as_ref(), sync_id.tenant_id).await
+    {
         error!(
             "Failed to upload missing branches for sync id {}: {:?}",
             sync_id, e
@@ -57,7 +58,10 @@ pub(super) async fn upload_timeline_checkpoint<
 
     let index = &remote_assets.1;
 
-    let TimelineSyncId(tenant_id, timeline_id) = sync_id;
+    let ZTenantTimelineId {
+        tenant_id,
+        timeline_id,
+    } = sync_id;
     let timeline_dir = config.timeline_path(&timeline_id, &tenant_id);
 
     let index_read = index.read().await;
@@ -151,11 +155,14 @@ async fn try_upload_checkpoint<
 >(
     config: &'static PageServerConf,
     remote_assets: Arc<(S, RwLock<RemoteTimelineIndex>)>,
-    sync_id: TimelineSyncId,
+    sync_id: ZTenantTimelineId,
     new_checkpoint: &NewCheckpoint,
     files_to_skip: BTreeSet<PathBuf>,
 ) -> anyhow::Result<(ArchiveHeader, u64)> {
-    let TimelineSyncId(tenant_id, timeline_id) = sync_id;
+    let ZTenantTimelineId {
+        tenant_id,
+        timeline_id,
+    } = sync_id;
     let timeline_dir = config.timeline_path(&timeline_id, &tenant_id);
 
     let files_to_upload = new_checkpoint
@@ -288,7 +295,7 @@ mod tests {
     #[tokio::test]
     async fn reupload_timeline() -> anyhow::Result<()> {
         let repo_harness = RepoHarness::create("reupload_timeline")?;
-        let sync_id = TimelineSyncId(repo_harness.tenant_id, TIMELINE_ID);
+        let sync_id = ZTenantTimelineId::new(repo_harness.tenant_id, TIMELINE_ID);
         let storage = LocalFs::new(tempdir()?.path().to_owned(), &repo_harness.conf.workdir)?;
         let index = RwLock::new(RemoteTimelineIndex::try_parse_descriptions_from_paths(
             repo_harness.conf,
@@ -484,7 +491,7 @@ mod tests {
     #[tokio::test]
     async fn reupload_timeline_rejected() -> anyhow::Result<()> {
         let repo_harness = RepoHarness::create("reupload_timeline_rejected")?;
-        let sync_id = TimelineSyncId(repo_harness.tenant_id, TIMELINE_ID);
+        let sync_id = ZTenantTimelineId::new(repo_harness.tenant_id, TIMELINE_ID);
         let storage = LocalFs::new(tempdir()?.path().to_owned(), &repo_harness.conf.workdir)?;
         let index = RwLock::new(RemoteTimelineIndex::try_parse_descriptions_from_paths(
             repo_harness.conf,

pageserver/src/repository.rs

@@ -306,8 +306,12 @@ pub enum ZenithWalRecord {
     /// Native PostgreSQL WAL record
     Postgres { will_init: bool, rec: Bytes },
 
-    /// Set bits in heap visibility map. (heap blkno, flag bits to clear)
-    ClearVisibilityMapFlags { heap_blkno: u32, flags: u8 },
+    /// Clear bits in heap visibility map. ('flags' is bitmap of bits to clear)
+    ClearVisibilityMapFlags {
+        new_heap_blkno: Option<u32>,
+        old_heap_blkno: Option<u32>,
+        flags: u8,
+    },
     /// Mark transaction IDs as committed on a CLOG page
     ClogSetCommitted { xids: Vec<TransactionId> },
     /// Mark transaction IDs as aborted on a CLOG page

pageserver/src/virtual_file.rs

@@ -332,8 +332,11 @@ impl VirtualFile {
         // TODO: We could downgrade the locks to read mode before calling
         // 'func', to allow a little bit more concurrency, but the standard
         // library RwLock doesn't allow downgrading without releasing the lock,
-        // and that doesn't seem worth the trouble. (parking_lot RwLock would
-        // allow it)
+        // and that doesn't seem worth the trouble.
+        //
+        // XXX: `parking_lot::RwLock` can enable such downgrades, yet its implemenation is fair and
+        // may deadlock on subsequent read calls.
+        // Simply replacing all `RwLock` in project causes deadlocks, so use it sparingly.
         let result = STORAGE_IO_TIME
             .with_label_values(&[op, &self.tenantid, &self.timelineid])
             .observe_closure_duration(|| func(&file));

pageserver/src/walingest.rs

@@ -349,49 +349,25 @@ impl WalIngest {
         decoded: &mut DecodedWALRecord,
     ) -> Result<()> {
         // Handle VM bit updates that are implicitly part of heap records.
+
+        // First, look at the record to determine which VM bits need
+        // to be cleared. If either of these variables is set, we
+        // need to clear the corresponding bits in the visibility map.
+        let mut new_heap_blkno: Option<u32> = None;
+        let mut old_heap_blkno: Option<u32> = None;
         if decoded.xl_rmid == pg_constants::RM_HEAP_ID {
             let info = decoded.xl_info & pg_constants::XLOG_HEAP_OPMASK;
             if info == pg_constants::XLOG_HEAP_INSERT {
                 let xlrec = XlHeapInsert::decode(buf);
                 assert_eq!(0, buf.remaining());
-                if (xlrec.flags
-                    & (pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED
-                        | pg_constants::XLH_INSERT_ALL_FROZEN_SET))
-                    != 0
-                {
-                    timeline.put_wal_record(
-                        lsn,
-                        RelishTag::Relation(RelTag {
-                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
-                            spcnode: decoded.blocks[0].rnode_spcnode,
-                            dbnode: decoded.blocks[0].rnode_dbnode,
-                            relnode: decoded.blocks[0].rnode_relnode,
-                        }),
-                        decoded.blocks[0].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
-                        ZenithWalRecord::ClearVisibilityMapFlags {
-                            heap_blkno: decoded.blocks[0].blkno,
-                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
-                        },
-                    )?;
+                if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                    new_heap_blkno = Some(decoded.blocks[0].blkno);
                 }
             } else if info == pg_constants::XLOG_HEAP_DELETE {
                 let xlrec = XlHeapDelete::decode(buf);
                 assert_eq!(0, buf.remaining());
                 if (xlrec.flags & pg_constants::XLH_DELETE_ALL_VISIBLE_CLEARED) != 0 {
-                    timeline.put_wal_record(
-                        lsn,
-                        RelishTag::Relation(RelTag {
-                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
-                            spcnode: decoded.blocks[0].rnode_spcnode,
-                            dbnode: decoded.blocks[0].rnode_dbnode,
-                            relnode: decoded.blocks[0].rnode_relnode,
-                        }),
-                        decoded.blocks[0].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
-                        ZenithWalRecord::ClearVisibilityMapFlags {
-                            heap_blkno: decoded.blocks[0].blkno,
-                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
-                        },
-                    )?;
+                    new_heap_blkno = Some(decoded.blocks[0].blkno);
                 }
             } else if info == pg_constants::XLOG_HEAP_UPDATE
                 || info == pg_constants::XLOG_HEAP_HOT_UPDATE
@@ -400,39 +376,15 @@ impl WalIngest {
                 // the size of tuple data is inferred from the size of the record.
                 // we can't validate the remaining number of bytes without parsing
                 // the tuple data.
-                if (xlrec.flags & pg_constants::XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED) != 0 {
-                    timeline.put_wal_record(
-                        lsn,
-                        RelishTag::Relation(RelTag {
-                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
-                            spcnode: decoded.blocks[0].rnode_spcnode,
-                            dbnode: decoded.blocks[0].rnode_dbnode,
-                            relnode: decoded.blocks[0].rnode_relnode,
-                        }),
-                        decoded.blocks[0].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
-                        ZenithWalRecord::ClearVisibilityMapFlags {
-                            heap_blkno: decoded.blocks[0].blkno,
-                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
-                        },
-                    )?;
+                if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0 {
+                    old_heap_blkno = Some(decoded.blocks[0].blkno);
                 }
-                if (xlrec.flags & pg_constants::XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED) != 0
-                    && decoded.blocks.len() > 1
-                {
-                    timeline.put_wal_record(
-                        lsn,
-                        RelishTag::Relation(RelTag {
-                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
-                            spcnode: decoded.blocks[1].rnode_spcnode,
-                            dbnode: decoded.blocks[1].rnode_dbnode,
-                            relnode: decoded.blocks[1].rnode_relnode,
-                        }),
-                        decoded.blocks[1].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
-                        ZenithWalRecord::ClearVisibilityMapFlags {
-                            heap_blkno: decoded.blocks[1].blkno,
-                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
-                        },
-                    )?;
+                if (xlrec.flags & pg_constants::XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED) != 0 {
+                    // PostgreSQL only uses XLH_UPDATE_NEW_ALL_VISIBLE_CLEARED on a
+                    // non-HOT update where the new tuple goes to different page than
+                    // the old one. Otherwise, only XLH_UPDATE_OLD_ALL_VISIBLE_CLEARED is
+                    // set.
+                    new_heap_blkno = Some(decoded.blocks[1].blkno);
                 }
             }
         } else if decoded.xl_rmid == pg_constants::RM_HEAP2_ID {
@@ -448,23 +400,60 @@ impl WalIngest {
                 };
                 assert_eq!(offset_array_len, buf.remaining());
 
-                // FIXME: why also ALL_FROZEN_SET?
-                if (xlrec.flags
-                    & (pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED
-                        | pg_constants::XLH_INSERT_ALL_FROZEN_SET))
-                    != 0
-                {
+                if (xlrec.flags & pg_constants::XLH_INSERT_ALL_VISIBLE_CLEARED) != 0 {
+                    new_heap_blkno = Some(decoded.blocks[0].blkno);
+                }
+            }
+        }
+        // FIXME: What about XLOG_HEAP_LOCK and XLOG_HEAP2_LOCK_UPDATED?
+
+        // Clear the VM bits if required.
+        if new_heap_blkno.is_some() || old_heap_blkno.is_some() {
+            let vm_relish = RelishTag::Relation(RelTag {
+                forknum: pg_constants::VISIBILITYMAP_FORKNUM,
+                spcnode: decoded.blocks[0].rnode_spcnode,
+                dbnode: decoded.blocks[0].rnode_dbnode,
+                relnode: decoded.blocks[0].rnode_relnode,
+            });
+
+            let new_vm_blk = new_heap_blkno.map(pg_constants::HEAPBLK_TO_MAPBLOCK);
+            let old_vm_blk = old_heap_blkno.map(pg_constants::HEAPBLK_TO_MAPBLOCK);
+            if new_vm_blk == old_vm_blk {
+                // An UPDATE record that needs to clear the bits for both old and the
+                // new page, both of which reside on the same VM page.
+                timeline.put_wal_record(
+                    lsn,
+                    vm_relish,
+                    new_vm_blk.unwrap(),
+                    ZenithWalRecord::ClearVisibilityMapFlags {
+                        new_heap_blkno,
+                        old_heap_blkno,
+                        flags: pg_constants::VISIBILITYMAP_VALID_BITS,
+                    },
+                )?;
+            } else {
+                // Clear VM bits for one heap page, or for two pages that reside on
+                // different VM pages.
+                if let Some(new_vm_blk) = new_vm_blk {
                     timeline.put_wal_record(
                         lsn,
-                        RelishTag::Relation(RelTag {
-                            forknum: pg_constants::VISIBILITYMAP_FORKNUM,
-                            spcnode: decoded.blocks[0].rnode_spcnode,
-                            dbnode: decoded.blocks[0].rnode_dbnode,
-                            relnode: decoded.blocks[0].rnode_relnode,
-                        }),
-                        decoded.blocks[0].blkno / pg_constants::HEAPBLOCKS_PER_PAGE as u32,
+                        vm_relish,
+                        new_vm_blk,
                         ZenithWalRecord::ClearVisibilityMapFlags {
-                            heap_blkno: decoded.blocks[0].blkno,
+                            new_heap_blkno,
+                            old_heap_blkno: None,
+                            flags: pg_constants::VISIBILITYMAP_VALID_BITS,
+                        },
+                    )?;
+                }
+                if let Some(old_vm_blk) = old_vm_blk {
+                    timeline.put_wal_record(
+                        lsn,
+                        vm_relish,
+                        old_vm_blk,
+                        ZenithWalRecord::ClearVisibilityMapFlags {
+                            new_heap_blkno: None,
+                            old_heap_blkno,
                             flags: pg_constants::VISIBILITYMAP_VALID_BITS,
                         },
                     )?;
@@ -472,8 +461,6 @@ impl WalIngest {
             }
         }
 
-        // FIXME: What about XLOG_HEAP_LOCK and XLOG_HEAP2_LOCK_UPDATED?
-
         Ok(())
     }
 

pageserver/src/walreceiver.rs

@@ -12,14 +12,15 @@ use crate::thread_mgr::ThreadKind;
 use crate::walingest::WalIngest;
 use anyhow::{bail, Context, Error, Result};
 use bytes::BytesMut;
+use fail::fail_point;
 use lazy_static::lazy_static;
-use parking_lot::Mutex;
 use postgres_ffi::waldecoder::*;
 use postgres_protocol::message::backend::ReplicationMessage;
 use postgres_types::PgLsn;
 use std::cell::Cell;
 use std::collections::HashMap;
 use std::str::FromStr;
+use std::sync::Mutex;
 use std::thread_local;
 use std::time::SystemTime;
 use tokio::pin;
@@ -31,6 +32,7 @@ use zenith_utils::lsn::Lsn;
 use zenith_utils::pq_proto::ZenithFeedback;
 use zenith_utils::zid::ZTenantId;
 use zenith_utils::zid::ZTimelineId;
+
 //
 // We keep one WAL Receiver active per timeline.
 //
@@ -51,7 +53,7 @@ thread_local! {
 }
 
 fn drop_wal_receiver(tenantid: ZTenantId, timelineid: ZTimelineId) {
-    let mut receivers = WAL_RECEIVERS.lock();
+    let mut receivers = WAL_RECEIVERS.lock().unwrap();
     receivers.remove(&(tenantid, timelineid));
 }
 
@@ -62,7 +64,7 @@ pub fn launch_wal_receiver(
     timelineid: ZTimelineId,
     wal_producer_connstr: &str,
 ) -> Result<()> {
-    let mut receivers = WAL_RECEIVERS.lock();
+    let mut receivers = WAL_RECEIVERS.lock().unwrap();
 
     match receivers.get_mut(&(tenantid, timelineid)) {
         Some(receiver) => {
@@ -95,7 +97,7 @@ pub fn launch_wal_receiver(
 
 // Look up current WAL producer connection string in the hash table
 fn get_wal_producer_connstr(tenantid: ZTenantId, timelineid: ZTimelineId) -> String {
-    let receivers = WAL_RECEIVERS.lock();
+    let receivers = WAL_RECEIVERS.lock().unwrap();
 
     receivers
         .get(&(tenantid, timelineid))
@@ -160,7 +162,7 @@ fn walreceiver_main(
     // This is from tokio-postgres docs, but it is a bit weird in our case because we extensively use block_on
     runtime.spawn(async move {
         if let Err(e) = connection.await {
-            eprintln!("connection error: {}", e);
+            error!("connection error: {}", e);
         }
     });
 
@@ -254,6 +256,8 @@ fn walreceiver_main(
                     let writer = timeline.writer();
                     walingest.ingest_record(writer.as_ref(), recdata, lsn)?;
 
+                    fail_point!("walreceiver-after-ingest");
+
                     last_rec_lsn = lsn;
                 }
 

pageserver/src/walredo.rs

@@ -363,25 +363,44 @@ impl PostgresRedoManager {
                 will_init: _,
                 rec: _,
             } => panic!("tried to pass postgres wal record to zenith WAL redo"),
-            ZenithWalRecord::ClearVisibilityMapFlags { heap_blkno, flags } => {
-                // Calculate the VM block and offset that corresponds to the heap block.
-                let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(*heap_blkno);
-                let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(*heap_blkno);
-                let map_offset = pg_constants::HEAPBLK_TO_OFFSET(*heap_blkno);
-
-                // Check that we're modifying the correct VM block.
+            ZenithWalRecord::ClearVisibilityMapFlags {
+                new_heap_blkno,
+                old_heap_blkno,
+                flags,
+            } => {
+                // sanity check that this is modifying the correct relish
                 assert!(
                     check_forknum(&rel, pg_constants::VISIBILITYMAP_FORKNUM),
                     "ClearVisibilityMapFlags record on unexpected rel {:?}",
                     rel
                 );
-                assert!(map_block == blknum);
+                if let Some(heap_blkno) = *new_heap_blkno {
+                    // Calculate the VM block and offset that corresponds to the heap block.
+                    let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(heap_blkno);
+                    let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(heap_blkno);
+                    let map_offset = pg_constants::HEAPBLK_TO_OFFSET(heap_blkno);
 
-                // equivalent to PageGetContents(page)
-                let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..];
+                    // Check that we're modifying the correct VM block.
+                    assert!(map_block == blknum);
 
-                let mask: u8 = flags << map_offset;
-                map[map_byte as usize] &= !mask;
+                    // equivalent to PageGetContents(page)
+                    let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..];
+
+                    map[map_byte as usize] &= !(flags << map_offset);
+                }
+
+                // Repeat for 'old_heap_blkno', if any
+                if let Some(heap_blkno) = *old_heap_blkno {
+                    let map_block = pg_constants::HEAPBLK_TO_MAPBLOCK(heap_blkno);
+                    let map_byte = pg_constants::HEAPBLK_TO_MAPBYTE(heap_blkno);
+                    let map_offset = pg_constants::HEAPBLK_TO_OFFSET(heap_blkno);
+
+                    assert!(map_block == blknum);
+
+                    let map = &mut page[pg_constants::MAXALIGN_SIZE_OF_PAGE_HEADER_DATA..];
+
+                    map[map_byte as usize] &= !(flags << map_offset);
+                }
             }
             // Non-relational WAL records are handled here, with custom code that has the
             // same effects as the corresponding Postgres WAL redo function.

postgres_ffi/Cargo.toml

@@ -1,11 +1,8 @@
 [package]
 name = "postgres_ffi"
 version = "0.1.0"
-authors = ["Heikki Linnakangas <heikki@zenith.tech>"]
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 chrono = "0.4.19"
 rand = "0.8.3"

proxy/Cargo.toml

@@ -1,38 +1,32 @@
 [package]
 name = "proxy"
 version = "0.1.0"
-authors = ["Stas Kelvich <stas.kelvich@gmail.com>"]
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 anyhow = "1.0"
 bytes = { version = "1.0.1", features = ['serde'] }
-lazy_static = "1.4.0"
-md5 = "0.7.0"
-rand = "0.8.3"
+clap = "3.0"
+futures = "0.3.13"
+hashbrown = "0.11.2"
 hex = "0.4.3"
 hyper = "0.14"
-routerify = "2"
+lazy_static = "1.4.0"
+md5 = "0.7.0"
 parking_lot = "0.11.2"
-hashbrown = "0.11.2"
+pin-project-lite = "0.2.7"
+rand = "0.8.3"
+reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
+rustls = "0.19.1"
+scopeguard = "1.1.0"
 serde = "1"
 serde_json = "1"
 tokio = { version = "1.11", features = ["macros"] }
 tokio-postgres = { git = "https://github.com/zenithdb/rust-postgres.git", rev="2949d98df52587d562986aad155dd4e889e408b7" }
 tokio-rustls = "0.22.0"
-clap = "2.33.0"
-rustls = "0.19.1"
-reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
-pin-project-lite = "0.2.7"
-futures = "0.3.13"
-scopeguard = "1.1.0"
 
 zenith_utils = { path = "../zenith_utils" }
 zenith_metrics = { path = "../zenith_metrics" }
-base64 = "0.13.0"
-async-trait = "0.1.52"
 
 [dev-dependencies]
 tokio-postgres-rustls = "0.8.0"

proxy/src/auth.rs

@@ -1,41 +1,169 @@
-use crate::db::AuthSecret;
+use crate::compute::DatabaseInfo;
+use crate::config::ProxyConfig;
+use crate::cplane_api::{self, CPlaneApi};
 use crate::stream::PqStream;
-use bytes::Bytes;
+use anyhow::{anyhow, bail, Context};
+use std::collections::HashMap;
 use tokio::io::{AsyncRead, AsyncWrite};
-use zenith_utils::pq_proto::BeMessage as Be;
+use zenith_utils::pq_proto::{BeMessage as Be, BeParameterStatusMessage, FeMessage as Fe};
 
-
-/// Stored secret for authenticating the user via md5 but authenticating
-/// to the compute database with a (possibly different) plaintext password.
-pub struct PlaintextStoredSecret {
-    pub salt: [u8; 4],
-    pub hashed_salted_password: Bytes,
-    pub compute_db_password: String,
+/// Various client credentials which we use for authentication.
+#[derive(Debug, PartialEq, Eq)]
+pub struct ClientCredentials {
+    pub user: String,
+    pub dbname: String,
 }
 
-/// Sufficient information to auth user and create AuthSecret
-#[non_exhaustive]
-pub enum StoredSecret {
-    PlaintextPassword(PlaintextStoredSecret),
-    // TODO add md5 option?
-    // TODO add SCRAM option
-}
+impl TryFrom<HashMap<String, String>> for ClientCredentials {
+    type Error = anyhow::Error;
 
-pub async fn authenticate(
-    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    stored_secret: StoredSecret
-) -> anyhow::Result<AuthSecret> {
-    match stored_secret {
-        StoredSecret::PlaintextPassword(stored) => {
-            client.write_message(&Be::AuthenticationMD5Password(&stored.salt)).await?;
-            let provided = client.read_password_message().await?;
-            anyhow::ensure!(provided == stored.hashed_salted_password);
-            Ok(AuthSecret::Password(stored.compute_db_password))
-        },
+    fn try_from(mut value: HashMap<String, String>) -> Result<Self, Self::Error> {
+        let mut get_param = |key| {
+            value
+                .remove(key)
+                .with_context(|| format!("{} is missing in startup packet", key))
+        };
+
+        let user = get_param("user")?;
+        let db = get_param("database")?;
+
+        Ok(Self { user, dbname: db })
     }
 }
 
-#[async_trait::async_trait]
-pub trait SecretStore {
-    async fn get_stored_secret(&self, creds: &crate::cplane_api::ClientCredentials) -> anyhow::Result<StoredSecret>;
+impl ClientCredentials {
+    /// Use credentials to authenticate the user.
+    pub async fn authenticate(
+        self,
+        config: &ProxyConfig,
+        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    ) -> anyhow::Result<DatabaseInfo> {
+        use crate::config::ClientAuthMethod::*;
+        use crate::config::RouterConfig::*;
+        let db_info = match &config.router_config {
+            Static { host, port } => handle_static(host.clone(), *port, client, self).await,
+            Dynamic(Mixed) => {
+                if self.user.ends_with("@zenith") {
+                    handle_existing_user(config, client, self).await
+                } else {
+                    handle_new_user(config, client).await
+                }
+            }
+            Dynamic(Password) => handle_existing_user(config, client, self).await,
+            Dynamic(Link) => handle_new_user(config, client).await,
+        };
+
+        db_info.context("failed to authenticate client")
+    }
+}
+
+fn new_psql_session_id() -> String {
+    hex::encode(rand::random::<[u8; 8]>())
+}
+
+async fn handle_static(
+    host: String,
+    port: u16,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    creds: ClientCredentials,
+) -> anyhow::Result<DatabaseInfo> {
+    client
+        .write_message(&Be::AuthenticationCleartextPassword)
+        .await?;
+
+    // Read client's password bytes
+    let msg = match client.read_message().await? {
+        Fe::PasswordMessage(msg) => msg,
+        bad => bail!("unexpected message type: {:?}", bad),
+    };
+
+    let cleartext_password = std::str::from_utf8(&msg)?.split('\0').next().unwrap();
+
+    let db_info = DatabaseInfo {
+        host,
+        port,
+        dbname: creds.dbname.clone(),
+        user: creds.user.clone(),
+        password: Some(cleartext_password.into()),
+    };
+
+    client
+        .write_message_noflush(&Be::AuthenticationOk)?
+        .write_message_noflush(&BeParameterStatusMessage::encoding())?;
+
+    Ok(db_info)
+}
+
+async fn handle_existing_user(
+    config: &ProxyConfig,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    creds: ClientCredentials,
+) -> anyhow::Result<DatabaseInfo> {
+    let psql_session_id = new_psql_session_id();
+    let md5_salt = rand::random();
+
+    client
+        .write_message(&Be::AuthenticationMD5Password(&md5_salt))
+        .await?;
+
+    // Read client's password hash
+    let msg = match client.read_message().await? {
+        Fe::PasswordMessage(msg) => msg,
+        bad => bail!("unexpected message type: {:?}", bad),
+    };
+
+    let (_trailing_null, md5_response) = msg
+        .split_last()
+        .ok_or_else(|| anyhow!("unexpected password message"))?;
+
+    let cplane = CPlaneApi::new(&config.auth_endpoint);
+    let db_info = cplane
+        .authenticate_proxy_request(creds, md5_response, &md5_salt, &psql_session_id)
+        .await?;
+
+    client
+        .write_message_noflush(&Be::AuthenticationOk)?
+        .write_message_noflush(&BeParameterStatusMessage::encoding())?;
+
+    Ok(db_info)
+}
+
+async fn handle_new_user(
+    config: &ProxyConfig,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+) -> anyhow::Result<DatabaseInfo> {
+    let psql_session_id = new_psql_session_id();
+    let greeting = hello_message(&config.redirect_uri, &psql_session_id);
+
+    let db_info = cplane_api::with_waiter(psql_session_id, |waiter| async {
+        // Give user a URL to spawn a new database
+        client
+            .write_message_noflush(&Be::AuthenticationOk)?
+            .write_message_noflush(&BeParameterStatusMessage::encoding())?
+            .write_message(&Be::NoticeResponse(greeting))
+            .await?;
+
+        // Wait for web console response
+        waiter.await?.map_err(|e| anyhow!(e))
+    })
+    .await?;
+
+    client.write_message_noflush(&Be::NoticeResponse("Connecting to database.".into()))?;
+
+    Ok(db_info)
+}
+
+fn hello_message(redirect_uri: &str, session_id: &str) -> String {
+    format!(
+        concat![
+            "☀️  Welcome to Zenith!\n",
+            "To proceed with database creation, open the following link:\n\n",
+            "    {redirect_uri}{session_id}\n\n",
+            "It needs to be done once and we will send you '.pgpass' file,\n",
+            "which will allow you to access or create ",
+            "databases without opening your web browser."
+        ],
+        redirect_uri = redirect_uri,
+        session_id = session_id,
+    )
 }

proxy/src/cancellation.rs

@@ -1,15 +1,57 @@
 use anyhow::{anyhow, Context};
 use hashbrown::HashMap;
-use lazy_static::lazy_static;
 use parking_lot::Mutex;
 use std::net::SocketAddr;
 use tokio::net::TcpStream;
 use tokio_postgres::{CancelToken, NoTls};
 use zenith_utils::pq_proto::CancelKeyData;
 
-lazy_static! {
-    /// Enables serving CancelRequests.
-    static ref CANCEL_MAP: Mutex<HashMap<CancelKeyData, Option<CancelClosure>>> = Default::default();
+/// Enables serving CancelRequests.
+#[derive(Default)]
+pub struct CancelMap(Mutex<HashMap<CancelKeyData, Option<CancelClosure>>>);
+
+impl CancelMap {
+    /// Cancel a running query for the corresponding connection.
+    pub async fn cancel_session(&self, key: CancelKeyData) -> anyhow::Result<()> {
+        let cancel_closure = self
+            .0
+            .lock()
+            .get(&key)
+            .and_then(|x| x.clone())
+            .with_context(|| format!("unknown session: {:?}", key))?;
+
+        cancel_closure.try_cancel_query().await
+    }
+
+    /// Run async action within an ephemeral session identified by [`CancelKeyData`].
+    pub async fn with_session<'a, F, R, V>(&'a self, f: F) -> anyhow::Result<V>
+    where
+        F: FnOnce(Session<'a>) -> R,
+        R: std::future::Future<Output = anyhow::Result<V>>,
+    {
+        // HACK: We'd rather get the real backend_pid but tokio_postgres doesn't
+        // expose it and we don't want to do another roundtrip to query
+        // for it. The client will be able to notice that this is not the
+        // actual backend_pid, but backend_pid is not used for anything
+        // so it doesn't matter.
+        let key = rand::random();
+
+        // Random key collisions are unlikely to happen here, but they're still possible,
+        // which is why we have to take care not to rewrite an existing key.
+        self.0
+            .lock()
+            .try_insert(key, None)
+            .map_err(|_| anyhow!("session already exists: {:?}", key))?;
+
+        // This will guarantee that the session gets dropped
+        // as soon as the future is finished.
+        scopeguard::defer! {
+            self.0.lock().remove(&key);
+        }
+
+        let session = Session::new(key, self);
+        f(session).await
+    }
 }
 
 /// This should've been a [`std::future::Future`], but
@@ -38,53 +80,27 @@ impl CancelClosure {
     }
 }
 
-/// Cancel a running query for the corresponding connection.
-pub async fn cancel_session(key: CancelKeyData) -> anyhow::Result<()> {
-    let cancel_closure = CANCEL_MAP
-        .lock()
-        .get(&key)
-        .and_then(|x| x.clone())
-        .with_context(|| format!("unknown session: {:?}", key))?;
-
-    cancel_closure.try_cancel_query().await
-}
-
 /// Helper for registering query cancellation tokens.
-pub struct Session(CancelKeyData);
+pub struct Session<'a> {
+    /// The user-facing key identifying this session.
+    key: CancelKeyData,
+    /// The [`CancelMap`] this session belongs to.
+    cancel_map: &'a CancelMap,
+}
+
+impl<'a> Session<'a> {
+    fn new(key: CancelKeyData, cancel_map: &'a CancelMap) -> Self {
+        Self { key, cancel_map }
+    }
 
-impl Session {
     /// Store the cancel token for the given session.
+    /// This enables query cancellation in [`crate::proxy::handshake`].
     pub fn enable_cancellation(self, cancel_closure: CancelClosure) -> CancelKeyData {
-        CANCEL_MAP.lock().insert(self.0, Some(cancel_closure));
-        self.0
+        self.cancel_map
+            .0
+            .lock()
+            .insert(self.key, Some(cancel_closure));
+
+        self.key
     }
 }
-
-/// Run async action within an ephemeral session identified by [`CancelKeyData`].
-pub async fn with_session<F, R, V>(f: F) -> anyhow::Result<V>
-where
-    F: FnOnce(Session) -> R,
-    R: std::future::Future<Output = anyhow::Result<V>>,
-{
-    // HACK: We'd rather get the real backend_pid but tokio_postgres doesn't
-    // expose it and we don't want to do another roundtrip to query
-    // for it. The client will be able to notice that this is not the
-    // actual backend_pid, but backend_pid is not used for anything
-    // so it doesn't matter.
-    let key = rand::random();
-
-    // The birthday problem is unlikely to happen here, but it's still possible
-    CANCEL_MAP
-        .lock()
-        .try_insert(key, None)
-        .map_err(|_| anyhow!("session already exists: {:?}", key))?;
-
-    // This will guarantee that the session gets dropped
-    // as soon as the future is finished.
-    scopeguard::defer! {
-        CANCEL_MAP.lock().remove(&key);
-    }
-
-    let session = Session(key);
-    f(session).await
-}

proxy/src/compute.rs

@@ -1,7 +1,42 @@
-use crate::{cplane_api::ClientCredentials, db::DatabaseConnInfo};
+use anyhow::Context;
+use serde::{Deserialize, Serialize};
+use std::net::{SocketAddr, ToSocketAddrs};
 
-
-#[async_trait::async_trait]
-pub trait ComputeProvider {
-    async fn get_compute_node(&self, creds: &ClientCredentials) -> anyhow::Result<DatabaseConnInfo>;
+/// Compute node connection params.
+#[derive(Serialize, Deserialize, Debug, Default)]
+pub struct DatabaseInfo {
+    pub host: String,
+    pub port: u16,
+    pub dbname: String,
+    pub user: String,
+    pub password: Option<String>,
+}
+
+impl DatabaseInfo {
+    pub fn socket_addr(&self) -> anyhow::Result<SocketAddr> {
+        let host_port = format!("{}:{}", self.host, self.port);
+        host_port
+            .to_socket_addrs()
+            .with_context(|| format!("cannot resolve {} to SocketAddr", host_port))?
+            .next()
+            .context("cannot resolve at least one SocketAddr")
+    }
+}
+
+impl From<DatabaseInfo> for tokio_postgres::Config {
+    fn from(db_info: DatabaseInfo) -> Self {
+        let mut config = tokio_postgres::Config::new();
+
+        config
+            .host(&db_info.host)
+            .port(db_info.port)
+            .dbname(&db_info.dbname)
+            .user(&db_info.user);
+
+        if let Some(password) = db_info.password {
+            config.password(password);
+        }
+
+        config
+    }
 }

proxy/src/config.rs

@@ -1,15 +1,46 @@
-use crate::cplane_api::DatabaseInfo;
 use anyhow::{anyhow, ensure, Context};
 use rustls::{internal::pemfile, NoClientAuth, ProtocolVersion, ServerConfig};
 use std::net::SocketAddr;
+use std::str::FromStr;
 use std::sync::Arc;
 
-pub type SslConfig = Arc<ServerConfig>;
+pub type TlsConfig = Arc<ServerConfig>;
+
+#[non_exhaustive]
+pub enum ClientAuthMethod {
+    Password,
+    Link,
+
+    /// Use password auth only if username ends with "@zenith"
+    Mixed,
+}
+
+pub enum RouterConfig {
+    Static { host: String, port: u16 },
+    Dynamic(ClientAuthMethod),
+}
+
+impl FromStr for ClientAuthMethod {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> anyhow::Result<Self> {
+        use ClientAuthMethod::*;
+        match s {
+            "password" => Ok(Password),
+            "link" => Ok(Link),
+            "mixed" => Ok(Mixed),
+            _ => Err(anyhow::anyhow!("Invlid option for router")),
+        }
+    }
+}
 
 pub struct ProxyConfig {
     /// main entrypoint for users to connect to
     pub proxy_address: SocketAddr,
 
+    /// method of assigning compute nodes
+    pub router_config: RouterConfig,
+
     /// internally used for status and prometheus metrics
     pub http_address: SocketAddr,
 
@@ -24,26 +55,10 @@ pub struct ProxyConfig {
     /// control plane address where we would check auth.
     pub auth_endpoint: String,
 
-    pub ssl_config: Option<SslConfig>,
+    pub tls_config: Option<TlsConfig>,
 }
 
-pub type ProxyWaiters = crate::waiters::Waiters<Result<DatabaseInfo, String>>;
-
-pub struct ProxyState {
-    pub conf: ProxyConfig,
-    pub waiters: ProxyWaiters,
-}
-
-impl ProxyState {
-    pub fn new(conf: ProxyConfig) -> Self {
-        Self {
-            conf,
-            waiters: ProxyWaiters::default(),
-        }
-    }
-}
-
-pub fn configure_ssl(key_path: &str, cert_path: &str) -> anyhow::Result<SslConfig> {
+pub fn configure_ssl(key_path: &str, cert_path: &str) -> anyhow::Result<TlsConfig> {
     let key = {
         let key_bytes = std::fs::read(key_path).context("SSL key file")?;
         let mut keys = pemfile::pkcs8_private_keys(&mut &key_bytes[..])

proxy/src/cplane_api.rs

@@ -1,101 +1,87 @@
-use anyhow::{anyhow, bail, Context};
+use crate::auth::ClientCredentials;
+use crate::compute::DatabaseInfo;
+use crate::waiters::{Waiter, Waiters};
+use anyhow::{anyhow, bail};
+use lazy_static::lazy_static;
 use serde::{Deserialize, Serialize};
-use std::net::{SocketAddr, ToSocketAddrs};
-use std::collections::HashMap;
 
-use crate::state::ProxyWaiters;
-
-#[derive(Debug, PartialEq, Eq)]
-pub struct ClientCredentials {
-    pub user: String,
-    pub dbname: String,
+lazy_static! {
+    static ref CPLANE_WAITERS: Waiters<Result<DatabaseInfo, String>> = Default::default();
 }
 
-impl TryFrom<HashMap<String, String>> for ClientCredentials {
-    type Error = anyhow::Error;
-
-    fn try_from(mut value: HashMap<String, String>) -> Result<Self, Self::Error> {
-        let mut get_param = |key| {
-            value
-                .remove(key)
-                .with_context(|| format!("{} is missing in startup packet", key))
-        };
-
-        let user = get_param("user")?;
-        let db = get_param("database")?;
-
-        Ok(Self { user, dbname: db })
-    }
+/// Give caller an opportunity to wait for cplane's reply.
+pub async fn with_waiter<F, R, T>(psql_session_id: impl Into<String>, f: F) -> anyhow::Result<T>
+where
+    F: FnOnce(Waiter<'static, Result<DatabaseInfo, String>>) -> R,
+    R: std::future::Future<Output = anyhow::Result<T>>,
+{
+    let waiter = CPLANE_WAITERS.register(psql_session_id.into())?;
+    f(waiter).await
 }
 
-#[derive(Serialize, Deserialize, Debug, Default)]
-pub struct DatabaseInfo {
-    pub host: String,
-    pub port: u16,
-    pub dbname: String,
-    pub user: String,
-    pub password: Option<String>,
-}
-
-#[derive(Serialize, Deserialize, Debug)]
-#[serde(untagged)]
-enum ProxyAuthResponse {
-    Ready { conn_info: DatabaseInfo },
-    Error { error: String },
-    NotReady { ready: bool }, // TODO: get rid of `ready`
+pub fn notify(psql_session_id: &str, msg: Result<DatabaseInfo, String>) -> anyhow::Result<()> {
+    CPLANE_WAITERS.notify(psql_session_id, msg)
 }
 
+/// Zenith console API wrapper.
 pub struct CPlaneApi<'a> {
     auth_endpoint: &'a str,
-    waiters: &'a ProxyWaiters,
 }
 
 impl<'a> CPlaneApi<'a> {
-    pub fn new(auth_endpoint: &'a str, waiters: &'a ProxyWaiters) -> Self {
-        Self {
-            auth_endpoint,
-            waiters,
-        }
+    pub fn new(auth_endpoint: &'a str) -> Self {
+        Self { auth_endpoint }
     }
 }
 
 impl CPlaneApi<'_> {
-    pub fn authenticate_proxy_request(
+    pub async fn authenticate_proxy_request(
         &self,
-        user: &str,
-        database: &str,
+        creds: ClientCredentials,
         md5_response: &[u8],
         salt: &[u8; 4],
         psql_session_id: &str,
     ) -> anyhow::Result<DatabaseInfo> {
         let mut url = reqwest::Url::parse(self.auth_endpoint)?;
         url.query_pairs_mut()
-            .append_pair("login", user)
-            .append_pair("database", database)
+            .append_pair("login", &creds.user)
+            .append_pair("database", &creds.dbname)
             .append_pair("md5response", std::str::from_utf8(md5_response)?)
             .append_pair("salt", &hex::encode(salt))
             .append_pair("psql_session_id", psql_session_id);
 
-        let waiter = self.waiters.register(psql_session_id.to_owned());
+        with_waiter(psql_session_id, |waiter| async {
+            println!("cplane request: {}", url);
+            // TODO: leverage `reqwest::Client` to reuse connections
+            let resp = reqwest::get(url).await?;
+            if !resp.status().is_success() {
+                bail!("Auth failed: {}", resp.status())
+            }
 
-        println!("cplane request: {}", url);
-        let resp = reqwest::blocking::get(url)?;
-        if !resp.status().is_success() {
-            bail!("Auth failed: {}", resp.status())
-        }
+            let auth_info: ProxyAuthResponse = serde_json::from_str(resp.text().await?.as_str())?;
+            println!("got auth info: #{:?}", auth_info);
 
-        let auth_info: ProxyAuthResponse = serde_json::from_str(resp.text()?.as_str())?;
-        println!("got auth info: #{:?}", auth_info);
-
-        use ProxyAuthResponse::*;
-        match auth_info {
-            Ready { conn_info } => Ok(conn_info),
-            Error { error } => bail!(error),
-            NotReady { .. } => waiter.wait()?.map_err(|e| anyhow!(e)),
-        }
+            use ProxyAuthResponse::*;
+            match auth_info {
+                Ready { conn_info } => Ok(conn_info),
+                Error { error } => bail!(error),
+                NotReady { .. } => waiter.await?.map_err(|e| anyhow!(e)),
+            }
+        })
+        .await
     }
 }
 
+// NOTE: the order of constructors is important.
+// https://serde.rs/enum-representations.html#untagged
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(untagged)]
+enum ProxyAuthResponse {
+    Ready { conn_info: DatabaseInfo },
+    Error { error: String },
+    NotReady { ready: bool }, // TODO: get rid of `ready`
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

proxy/src/db.rs

@@ -1,58 +0,0 @@
-///
-/// Utils for connecting with the postgres dataabase.
-///
-
-use std::net::{SocketAddr, ToSocketAddrs};
-use anyhow::{Context, anyhow};
-
-use crate::cplane_api::ClientCredentials;
-
-pub struct DatabaseConnInfo {
-    pub host: String,
-    pub port: u16,
-}
-
-pub struct DatabaseAuthInfo {
-    pub conn_info: DatabaseConnInfo,
-    pub creds: ClientCredentials,
-    pub auth_secret: AuthSecret,
-}
-
-/// Sufficient information to auth with database
-#[non_exhaustive]
-#[derive(Debug)]
-pub enum AuthSecret {
-    Password(String),
-    // TODO add SCRAM option
-}
-
-impl From<DatabaseAuthInfo> for tokio_postgres::Config {
-    fn from(auth_info: DatabaseAuthInfo) -> Self {
-        let mut config = tokio_postgres::Config::new();
-
-        config
-            .host(&auth_info.conn_info.host)
-            .port(auth_info.conn_info.port)
-            .dbname(&auth_info.creds.dbname)
-            .user(&auth_info.creds.user);
-
-        match auth_info.auth_secret {
-            AuthSecret::Password(password) => {
-                config.password(password);
-            }
-        }
-
-        config
-    }
-}
-
-impl DatabaseConnInfo {
-    pub fn socket_addr(&self) -> anyhow::Result<SocketAddr> {
-        let host_port = format!("{}:{}", self.host, self.port);
-        host_port
-            .to_socket_addrs()
-            .with_context(|| format!("cannot resolve {} to SocketAddr", host_port))?
-            .next()
-            .ok_or_else(|| anyhow!("cannot resolve at least one SocketAddr"))
-    }
-}

proxy/src/http.rs

@@ -1,10 +1,10 @@
 use anyhow::anyhow;
 use hyper::{Body, Request, Response, StatusCode};
-use routerify::RouterBuilder;
 use std::net::TcpListener;
 use zenith_utils::http::endpoint;
 use zenith_utils::http::error::ApiError;
 use zenith_utils::http::json::json_response;
+use zenith_utils::http::{RouterBuilder, RouterService};
 
 async fn status_handler(_: Request<Body>) -> Result<Response<Body>, ApiError> {
     Ok(json_response(StatusCode::OK, "")?)
@@ -16,7 +16,11 @@ fn make_router() -> RouterBuilder<hyper::Body, ApiError> {
 }
 
 pub async fn thread_main(http_listener: TcpListener) -> anyhow::Result<()> {
-    let service = || routerify::RouterService::new(make_router().build()?);
+    scopeguard::defer! {
+        println!("http has shut down");
+    }
+
+    let service = || RouterService::new(make_router().build()?);
 
     hyper::Server::from_tcp(http_listener)?
         .serve(service().map_err(|e| anyhow!(e))?)

proxy/src/main.rs

@@ -5,126 +5,162 @@
 /// (control plane API in our case) and can create new databases and accounts
 /// in somewhat transparent manner (again via communication with control plane API).
 ///
-use anyhow::bail;
+use anyhow::{bail, Context};
 use clap::{App, Arg};
-use state::{ProxyConfig, ProxyState};
-use zenith_utils::{tcp_listener, GIT_VERSION};
+use config::ProxyConfig;
+use futures::FutureExt;
+use std::future::Future;
+use tokio::{net::TcpListener, task::JoinError};
+use zenith_utils::GIT_VERSION;
+
+use crate::config::{ClientAuthMethod, RouterConfig};
 
-mod compute;
-mod mock;
 mod auth;
-mod db;
 mod cancellation;
+mod compute;
+mod config;
 mod cplane_api;
 mod http;
 mod mgmt;
 mod proxy;
-mod state;
 mod stream;
 mod waiters;
 
+/// Flattens Result<Result<T>> into Result<T>.
+async fn flatten_err(
+    f: impl Future<Output = Result<anyhow::Result<()>, JoinError>>,
+) -> anyhow::Result<()> {
+    f.map(|r| r.context("join error").and_then(|x| x)).await
+}
+
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     zenith_metrics::set_common_metrics_prefix("zenith_proxy");
     let arg_matches = App::new("Zenith proxy/router")
         .version(GIT_VERSION)
         .arg(
-            Arg::with_name("proxy")
-                .short("p")
+            Arg::new("proxy")
+                .short('p')
                 .long("proxy")
                 .takes_value(true)
                 .help("listen for incoming client connections on ip:port")
                 .default_value("127.0.0.1:4432"),
         )
         .arg(
-            Arg::with_name("mgmt")
-                .short("m")
+            Arg::new("auth-method")
+                .long("auth-method")
+                .takes_value(true)
+                .help("Possible values: password | link | mixed")
+                .default_value("mixed"),
+        )
+        .arg(
+            Arg::new("static-router")
+                .short('s')
+                .long("static-router")
+                .takes_value(true)
+                .help("Route all clients to host:port"),
+        )
+        .arg(
+            Arg::new("mgmt")
+                .short('m')
                 .long("mgmt")
                 .takes_value(true)
                 .help("listen for management callback connection on ip:port")
                 .default_value("127.0.0.1:7000"),
         )
         .arg(
-            Arg::with_name("http")
-                .short("h")
+            Arg::new("http")
+                .short('h')
                 .long("http")
                 .takes_value(true)
                 .help("listen for incoming http connections (metrics, etc) on ip:port")
                 .default_value("127.0.0.1:7001"),
         )
         .arg(
-            Arg::with_name("uri")
-                .short("u")
+            Arg::new("uri")
+                .short('u')
                 .long("uri")
                 .takes_value(true)
                 .help("redirect unauthenticated users to given uri")
                 .default_value("http://localhost:3000/psql_session/"),
         )
         .arg(
-            Arg::with_name("auth-endpoint")
-                .short("a")
+            Arg::new("auth-endpoint")
+                .short('a')
                 .long("auth-endpoint")
                 .takes_value(true)
                 .help("API endpoint for authenticating users")
                 .default_value("http://localhost:3000/authenticate_proxy_request/"),
         )
         .arg(
-            Arg::with_name("ssl-key")
-                .short("k")
+            Arg::new("ssl-key")
+                .short('k')
                 .long("ssl-key")
                 .takes_value(true)
                 .help("path to SSL key for client postgres connections"),
         )
         .arg(
-            Arg::with_name("ssl-cert")
-                .short("c")
+            Arg::new("ssl-cert")
+                .short('c')
                 .long("ssl-cert")
                 .takes_value(true)
                 .help("path to SSL cert for client postgres connections"),
         )
         .get_matches();
 
-    let ssl_config = match (
+    let tls_config = match (
         arg_matches.value_of("ssl-key"),
         arg_matches.value_of("ssl-cert"),
     ) {
-        (Some(key_path), Some(cert_path)) => {
-            Some(crate::state::configure_ssl(key_path, cert_path)?)
-        }
+        (Some(key_path), Some(cert_path)) => Some(config::configure_ssl(key_path, cert_path)?),
         (None, None) => None,
         _ => bail!("either both or neither ssl-key and ssl-cert must be specified"),
     };
 
-    let config = ProxyConfig {
+    let auth_method = arg_matches.value_of("auth-method").unwrap().parse()?;
+    let router_config = match arg_matches.value_of("static-router") {
+        None => RouterConfig::Dynamic(auth_method),
+        Some(addr) => {
+            if let ClientAuthMethod::Password = auth_method {
+                let (host, port) = addr.split_once(":").unwrap();
+                RouterConfig::Static {
+                    host: host.to_string(),
+                    port: port.parse().unwrap(),
+                }
+            } else {
+                bail!("static-router requires --auth-method password")
+            }
+        }
+    };
+
+    let config: &ProxyConfig = Box::leak(Box::new(ProxyConfig {
+        router_config,
         proxy_address: arg_matches.value_of("proxy").unwrap().parse()?,
         mgmt_address: arg_matches.value_of("mgmt").unwrap().parse()?,
         http_address: arg_matches.value_of("http").unwrap().parse()?,
         redirect_uri: arg_matches.value_of("uri").unwrap().parse()?,
         auth_endpoint: arg_matches.value_of("auth-endpoint").unwrap().parse()?,
-        ssl_config,
-    };
-    let state: &ProxyState = Box::leak(Box::new(ProxyState::new(config)));
+        tls_config,
+    }));
 
     println!("Version: {}", GIT_VERSION);
 
     // Check that we can bind to address before further initialization
-    println!("Starting http on {}", state.conf.http_address);
-    let http_listener = tcp_listener::bind(state.conf.http_address)?;
+    println!("Starting http on {}", config.http_address);
+    let http_listener = TcpListener::bind(config.http_address).await?.into_std()?;
 
-    println!("Starting proxy on {}", state.conf.proxy_address);
-    let proxy_listener = tokio::net::TcpListener::bind(state.conf.proxy_address).await?;
+    println!("Starting mgmt on {}", config.mgmt_address);
+    let mgmt_listener = TcpListener::bind(config.mgmt_address).await?.into_std()?;
 
-    println!("Starting mgmt on {}", state.conf.mgmt_address);
-    let mgmt_listener = tcp_listener::bind(state.conf.mgmt_address)?;
+    println!("Starting proxy on {}", config.proxy_address);
+    let proxy_listener = TcpListener::bind(config.proxy_address).await?;
 
     let http = tokio::spawn(http::thread_main(http_listener));
-    let proxy = tokio::spawn(proxy::thread_main(state, proxy_listener));
-    let mgmt = tokio::task::spawn_blocking(move || mgmt::thread_main(state, mgmt_listener));
+    let proxy = tokio::spawn(proxy::thread_main(config, proxy_listener));
+    let mgmt = tokio::task::spawn_blocking(move || mgmt::thread_main(mgmt_listener));
 
-    let _ = futures::future::try_join_all([http, proxy, mgmt])
-        .await?
-        .into_iter()
-        .collect::<Result<Vec<()>, _>>()?;
+    let tasks = [flatten_err(http), flatten_err(proxy), flatten_err(mgmt)];
+    let _: Vec<()> = futures::future::try_join_all(tasks).await?;
 
     Ok(())
 }

proxy/src/mgmt.rs

@@ -1,44 +1,49 @@
+use crate::{compute::DatabaseInfo, cplane_api};
+use anyhow::Context;
+use serde::Deserialize;
 use std::{
     net::{TcpListener, TcpStream},
     thread,
 };
-
-use serde::Deserialize;
 use zenith_utils::{
     postgres_backend::{self, AuthType, PostgresBackend},
     pq_proto::{BeMessage, SINGLE_COL_ROWDESC},
 };
 
-use crate::{cplane_api::DatabaseInfo, ProxyState};
-
 ///
 /// Main proxy listener loop.
 ///
 /// Listens for connections, and launches a new handler thread for each.
 ///
-pub fn thread_main(state: &'static ProxyState, listener: TcpListener) -> anyhow::Result<()> {
+pub fn thread_main(listener: TcpListener) -> anyhow::Result<()> {
+    scopeguard::defer! {
+        println!("mgmt has shut down");
+    }
+
+    listener
+        .set_nonblocking(false)
+        .context("failed to set listener to blocking")?;
     loop {
-        let (socket, peer_addr) = listener.accept()?;
+        let (socket, peer_addr) = listener.accept().context("failed to accept a new client")?;
         println!("accepted connection from {}", peer_addr);
-        socket.set_nodelay(true).unwrap();
+        socket
+            .set_nodelay(true)
+            .context("failed to set client socket option")?;
 
         thread::spawn(move || {
-            if let Err(err) = handle_connection(state, socket) {
+            if let Err(err) = handle_connection(socket) {
                 println!("error: {}", err);
             }
         });
     }
 }
 
-fn handle_connection(state: &ProxyState, socket: TcpStream) -> anyhow::Result<()> {
-    let mut conn_handler = MgmtHandler { state };
+fn handle_connection(socket: TcpStream) -> anyhow::Result<()> {
     let pgbackend = PostgresBackend::new(socket, AuthType::Trust, None, true)?;
-    pgbackend.run(&mut conn_handler)
+    pgbackend.run(&mut MgmtHandler)
 }
 
-struct MgmtHandler<'a> {
-    state: &'a ProxyState,
-}
+struct MgmtHandler;
 
 /// Serialized examples:
 // {
@@ -74,13 +79,13 @@ enum PsqlSessionResult {
     Failure(String),
 }
 
-impl postgres_backend::Handler for MgmtHandler<'_> {
+impl postgres_backend::Handler for MgmtHandler {
     fn process_query(
         &mut self,
         pgb: &mut PostgresBackend,
         query_string: &str,
     ) -> anyhow::Result<()> {
-        let res = try_process_query(self, pgb, query_string);
+        let res = try_process_query(pgb, query_string);
         // intercept and log error message
         if res.is_err() {
             println!("Mgmt query failed: #{:?}", res);
@@ -89,11 +94,7 @@ impl postgres_backend::Handler for MgmtHandler<'_> {
     }
 }
 
-fn try_process_query(
-    mgmt: &mut MgmtHandler,
-    pgb: &mut PostgresBackend,
-    query_string: &str,
-) -> anyhow::Result<()> {
+fn try_process_query(pgb: &mut PostgresBackend, query_string: &str) -> anyhow::Result<()> {
     println!("Got mgmt query: '{}'", query_string);
 
     let resp: PsqlSessionResponse = serde_json::from_str(query_string)?;
@@ -104,7 +105,7 @@ fn try_process_query(
         Failure(message) => Err(message),
     };
 
-    match mgmt.state.waiters.notify(&resp.session_id, msg) {
+    match cplane_api::notify(&resp.session_id, msg) {
         Ok(()) => {
             pgb.write_message_noflush(&SINGLE_COL_ROWDESC)?
                 .write_message_noflush(&BeMessage::DataRow(&[Some(b"ok")]))?

proxy/src/mock.rs

@@ -1,32 +0,0 @@
-use bytes::Bytes;
-
-use crate::{auth::{PlaintextStoredSecret, SecretStore, StoredSecret}, compute::ComputeProvider, cplane_api::ClientCredentials, db::DatabaseConnInfo};
-
-
-pub struct MockConsole {
-}
-
-#[async_trait::async_trait]
-impl SecretStore for MockConsole {
-    async fn get_stored_secret(&self, creds: &ClientCredentials) -> anyhow::Result<StoredSecret> {
-        let salt = [0; 4];
-        match (&creds.user[..], &creds.dbname[..]) {
-            ("postgres", "postgres") => Ok(StoredSecret::PlaintextPassword(PlaintextStoredSecret {
-                salt,
-                hashed_salted_password: "md52fff09cd9def51601fc5445943b3a11f\0".into(),
-                compute_db_password: "postgres".into(),
-            })),
-            _ => unimplemented!()
-        }
-    }
-}
-
-#[async_trait::async_trait]
-impl ComputeProvider for MockConsole{
-    async fn get_compute_node(&self, creds: &ClientCredentials) -> anyhow::Result<DatabaseConnInfo> {
-        return Ok(DatabaseConnInfo {
-            host: "127.0.0.1".into(),
-            port: 5432,
-        })
-    }
-}

proxy/src/proxy.rs

@@ -1,14 +1,11 @@
-use crate::auth::{self, StoredSecret, SecretStore};
-use crate::cancellation::{self, CancelClosure};
-use crate::compute::ComputeProvider;
-use crate::cplane_api as cplane;
-use crate::db::{AuthSecret, DatabaseAuthInfo};
-use crate::mock::MockConsole;
-use crate::state::SslConfig;
-use crate::stream::{PqStream, Stream};
-use crate::ProxyState;
+use crate::auth;
+use crate::cancellation::{self, CancelClosure, CancelMap};
+use crate::compute::DatabaseInfo;
+use crate::config::{ProxyConfig, TlsConfig};
+use crate::stream::{MetricsStream, PqStream, Stream};
 use anyhow::{bail, Context};
 use lazy_static::lazy_static;
+use std::sync::Arc;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
 use tokio_postgres::NoTls;
@@ -33,38 +30,44 @@ lazy_static! {
     .unwrap();
 }
 
-pub async fn thread_main(
-    state: &'static ProxyState,
-    listener: tokio::net::TcpListener,
-) -> anyhow::Result<()> {
-    loop {
-        let (socket, peer_addr) = listener.accept().await?;
-        println!("accepted connection from {}", peer_addr);
-
-        tokio::spawn(log_error(async {
-            socket
-                .set_nodelay(true)
-                .context("failed to set socket option")?;
-
-            let tls = state.conf.ssl_config.clone();
-            handle_client(socket, tls).await
-        }));
-    }
-}
-
 async fn log_error<R, F>(future: F) -> F::Output
 where
     F: std::future::Future<Output = anyhow::Result<R>>,
 {
     future.await.map_err(|err| {
-        println!("error: {}", err.to_string());
+        println!("error: {}", err);
         err
     })
 }
 
+pub async fn thread_main(
+    config: &'static ProxyConfig,
+    listener: tokio::net::TcpListener,
+) -> anyhow::Result<()> {
+    scopeguard::defer! {
+        println!("proxy has shut down");
+    }
+
+    let cancel_map = Arc::new(CancelMap::default());
+    loop {
+        let (socket, peer_addr) = listener.accept().await?;
+        println!("accepted connection from {}", peer_addr);
+
+        let cancel_map = Arc::clone(&cancel_map);
+        tokio::spawn(log_error(async move {
+            socket
+                .set_nodelay(true)
+                .context("failed to set socket option")?;
+
+            handle_client(config, &cancel_map, socket).await
+        }));
+    }
+}
+
 async fn handle_client(
+    config: &ProxyConfig,
+    cancel_map: &CancelMap,
     stream: impl AsyncRead + AsyncWrite + Unpin,
-    tls: Option<SslConfig>,
 ) -> anyhow::Result<()> {
     // The `closed` counter will increase when this future is destroyed.
     NUM_CONNECTIONS_ACCEPTED_COUNTER.inc();
@@ -72,11 +75,13 @@ async fn handle_client(
         NUM_CONNECTIONS_CLOSED_COUNTER.inc();
     }
 
-    if let Some((stream, creds)) = handshake(stream, tls).await? {
-        cancellation::with_session(|session| async {
-            connect_client_to_db(stream, creds, session).await
-        })
-        .await?;
+    let tls = config.tls_config.clone();
+    if let Some((client, creds)) = handshake(stream, tls, cancel_map).await? {
+        cancel_map
+            .with_session(|session| async {
+                connect_client_to_db(config, session, client, creds).await
+            })
+            .await?;
     }
 
     Ok(())
@@ -87,8 +92,9 @@ async fn handle_client(
 /// any object satisfying the traits.
 async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
     stream: S,
-    mut tls: Option<SslConfig>,
-) -> anyhow::Result<Option<(PqStream<Stream<S>>, cplane::ClientCredentials)>> {
+    mut tls: Option<TlsConfig>,
+    cancel_map: &CancelMap,
+) -> anyhow::Result<Option<(PqStream<Stream<S>>, auth::ClientCredentials)>> {
     // Client may try upgrading to each protocol only once
     let (mut tried_ssl, mut tried_gss) = (false, false);
 
@@ -136,7 +142,7 @@ async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
                 break Ok(Some((stream, params.try_into()?)));
             }
             CancelRequest(cancel_key_data) => {
-                cancellation::cancel_session(cancel_key_data).await?;
+                cancel_map.cancel_session(cancel_key_data).await?;
 
                 break Ok(None);
             }
@@ -145,30 +151,16 @@ async fn handshake<S: AsyncRead + AsyncWrite + Unpin>(
 }
 
 async fn connect_client_to_db(
+    config: &ProxyConfig,
+    session: cancellation::Session<'_>,
     mut client: PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    creds: cplane::ClientCredentials,
-    session: cancellation::Session,
+    creds: auth::ClientCredentials,
 ) -> anyhow::Result<()> {
-    // Authenticate
-    // TODO use real console
-    let console = MockConsole {};
-    let stored_secret = console.get_stored_secret(&creds).await?;
-    let auth_secret = auth::authenticate(&mut client, stored_secret).await?;
-    let conn_info = console.get_compute_node(&creds).await?;
-    let db_auth_info = DatabaseAuthInfo {
-        conn_info,
-        creds,
-        auth_secret,
-    };
-
-    // Connect to db
-    let (mut db, version, cancel_closure) = connect_to_db(db_auth_info).await?;
+    let db_info = creds.authenticate(config, &mut client).await?;
+    let (db, version, cancel_closure) = connect_to_db(db_info).await?;
     let cancel_key_data = session.enable_cancellation(cancel_closure);
 
-    // Report success to client
     client
-        .write_message_noflush(&Be::AuthenticationOk)?
-        .write_message_noflush(&BeParameterStatusMessage::encoding())?
         .write_message_noflush(&BeMessage::ParameterStatus(
             BeParameterStatusMessage::ServerVersion(&version),
         ))?
@@ -176,33 +168,26 @@ async fn connect_client_to_db(
         .write_message(&BeMessage::ReadyForQuery)
         .await?;
 
-    let mut client = client.into_inner();
+    // This function will be called for writes to either direction.
+    fn inc_proxied(cnt: usize) {
+        // Consider inventing something more sophisticated
+        // if this ever becomes a bottleneck (cacheline bouncing).
+        NUM_BYTES_PROXIED_COUNTER.inc_by(cnt as u64);
+    }
+
+    let mut db = MetricsStream::new(db, inc_proxied);
+    let mut client = MetricsStream::new(client.into_inner(), inc_proxied);
     let _ = tokio::io::copy_bidirectional(&mut client, &mut db).await?;
 
     Ok(())
 }
 
-fn hello_message(redirect_uri: &str, session_id: &str) -> String {
-    format!(
-        concat![
-            "☀️  Welcome to Zenith!\n",
-            "To proceed with database creation, open the following link:\n\n",
-            "    {redirect_uri}{session_id}\n\n",
-            "It needs to be done once and we will send you '.pgpass' file,\n",
-            "which will allow you to access or create ",
-            "databases without opening your web browser."
-        ],
-        redirect_uri = redirect_uri,
-        session_id = session_id,
-    )
-}
-
 /// Connect to a corresponding compute node.
 async fn connect_to_db(
-    db_info: DatabaseAuthInfo,
+    db_info: DatabaseInfo,
 ) -> anyhow::Result<(TcpStream, String, CancelClosure)> {
     // TODO: establish a secure connection to the DB
-    let socket_addr = db_info.conn_info.socket_addr()?;
+    let socket_addr = db_info.socket_addr()?;
     let mut socket = TcpStream::connect(socket_addr).await?;
 
     let (client, conn) = tokio_postgres::Config::from(db_info)
@@ -230,10 +215,14 @@ mod tests {
 
     async fn dummy_proxy(
         client: impl AsyncRead + AsyncWrite + Unpin,
-        tls: Option<SslConfig>,
+        tls: Option<TlsConfig>,
     ) -> anyhow::Result<()> {
+        let cancel_map = CancelMap::default();
+
         // TODO: add some infra + tests for credentials
-        let (mut stream, _creds) = handshake(client, tls).await?.context("no stream")?;
+        let (mut stream, _creds) = handshake(client, tls, &cancel_map)
+            .await?
+            .context("no stream")?;
 
         stream
             .write_message_noflush(&Be::AuthenticationOk)?

proxy/src/stream.rs

@@ -1,3 +1,4 @@
+use anyhow::Context;
 use bytes::BytesMut;
 use pin_project_lite::pin_project;
 use rustls::ServerConfig;
@@ -50,12 +51,10 @@ impl<S: AsyncRead + Unpin> PqStream<S> {
         }
     }
 
-    pub async fn read_password_message(&mut self) -> anyhow::Result<bytes::Bytes> {
-        match FeMessage::read_fut(&mut self.stream).await? {
-            Some(FeMessage::PasswordMessage(msg)) => Ok(msg),
-            None => anyhow::bail!("connection is lost"),
-            other => anyhow::bail!("bad message type: {:?}", other),
-        }
+    pub async fn read_message(&mut self) -> anyhow::Result<FeMessage> {
+        FeMessage::read_fut(&mut self.stream)
+            .await?
+            .context("connection is lost")
     }
 }
 
@@ -164,3 +163,68 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for Stream<S> {
         }
     }
 }
+
+pin_project! {
+    /// This stream tracks all writes and calls user provided
+    /// callback when the underlying stream is flushed.
+    pub struct MetricsStream<S, W> {
+        #[pin]
+        stream: S,
+        write_count: usize,
+        inc_write_count: W,
+    }
+}
+
+impl<S, W> MetricsStream<S, W> {
+    pub fn new(stream: S, inc_write_count: W) -> Self {
+        Self {
+            stream,
+            write_count: 0,
+            inc_write_count,
+        }
+    }
+}
+
+impl<S: AsyncRead + Unpin, W> AsyncRead for MetricsStream<S, W> {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        self.project().stream.poll_read(context, buf)
+    }
+}
+
+impl<S: AsyncWrite + Unpin, W: FnMut(usize)> AsyncWrite for MetricsStream<S, W> {
+    fn poll_write(
+        self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+        buf: &[u8],
+    ) -> task::Poll<io::Result<usize>> {
+        let this = self.project();
+        this.stream.poll_write(context, buf).map_ok(|cnt| {
+            // Increment the write count.
+            *this.write_count += cnt;
+            cnt
+        })
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        let this = self.project();
+        this.stream.poll_flush(context).map_ok(|()| {
+            // Call the user provided callback and reset the write count.
+            (this.inc_write_count)(*this.write_count);
+            *this.write_count = 0;
+        })
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        self.project().stream.poll_shutdown(context)
+    }
+}

proxy/src/waiters.rs

@@ -1,8 +1,12 @@
-use anyhow::Context;
-use std::collections::HashMap;
-use std::sync::{mpsc, Mutex};
+use anyhow::{anyhow, Context};
+use hashbrown::HashMap;
+use parking_lot::Mutex;
+use pin_project_lite::pin_project;
+use std::pin::Pin;
+use std::task;
+use tokio::sync::oneshot;
 
-pub struct Waiters<T>(pub(self) Mutex<HashMap<String, mpsc::Sender<T>>>);
+pub struct Waiters<T>(pub(self) Mutex<HashMap<String, oneshot::Sender<T>>>);
 
 impl<T> Default for Waiters<T> {
     fn default() -> Self {
@@ -11,48 +15,86 @@ impl<T> Default for Waiters<T> {
 }
 
 impl<T> Waiters<T> {
-    pub fn register(&self, key: String) -> Waiter<T> {
-        let (tx, rx) = mpsc::channel();
+    pub fn register(&self, key: String) -> anyhow::Result<Waiter<T>> {
+        let (tx, rx) = oneshot::channel();
 
-        // TODO: use `try_insert` (unstable)
-        let prev = self.0.lock().unwrap().insert(key.clone(), tx);
-        assert!(matches!(prev, None)); // assert_matches! is nightly-only
+        self.0
+            .lock()
+            .try_insert(key.clone(), tx)
+            .map_err(|_| anyhow!("waiter already registered"))?;
 
-        Waiter {
+        Ok(Waiter {
             receiver: rx,
-            registry: self,
-            key,
-        }
+            guard: DropKey {
+                registry: self,
+                key,
+            },
+        })
     }
 
     pub fn notify(&self, key: &str, value: T) -> anyhow::Result<()>
     where
-        T: Send + Sync + 'static,
+        T: Send + Sync,
     {
         let tx = self
             .0
             .lock()
-            .unwrap()
             .remove(key)
             .with_context(|| format!("key {} not found", key))?;
-        tx.send(value).context("channel hangup")
+
+        tx.send(value).map_err(|_| anyhow!("waiter channel hangup"))
     }
 }
 
-pub struct Waiter<'a, T> {
-    receiver: mpsc::Receiver<T>,
-    registry: &'a Waiters<T>,
+struct DropKey<'a, T> {
     key: String,
+    registry: &'a Waiters<T>,
 }
 
-impl<T> Waiter<'_, T> {
-    pub fn wait(self) -> anyhow::Result<T> {
-        self.receiver.recv().context("channel hangup")
-    }
-}
-
-impl<T> Drop for Waiter<'_, T> {
+impl<'a, T> Drop for DropKey<'a, T> {
     fn drop(&mut self) {
-        self.registry.0.lock().unwrap().remove(&self.key);
+        self.registry.0.lock().remove(&self.key);
+    }
+}
+
+pin_project! {
+    pub struct Waiter<'a, T> {
+        #[pin]
+        receiver: oneshot::Receiver<T>,
+        guard: DropKey<'a, T>,
+    }
+}
+
+impl<T> std::future::Future for Waiter<'_, T> {
+    type Output = anyhow::Result<T>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
+        self.project()
+            .receiver
+            .poll(cx)
+            .map_err(|_| anyhow!("channel hangup"))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Arc;
+
+    #[tokio::test]
+    async fn test_waiter() -> anyhow::Result<()> {
+        let waiters = Arc::new(Waiters::default());
+
+        let key = "Key";
+        let waiter = waiters.register(key.to_owned())?;
+
+        let waiters = Arc::clone(&waiters);
+        let notifier = tokio::spawn(async move {
+            waiters.notify(key, Default::default())?;
+            Ok(())
+        });
+
+        let () = waiter.await?;
+        notifier.await?
     }
 }

pyproject.toml

@@ -2,7 +2,7 @@
 name = "zenith"
 version = "0.1.0"
 description = ""
-authors = ["Dmitry Rodionov <dmitry@zenith.tech>"]
+authors = []
 
 [tool.poetry.dependencies]
 python = "^3.7"

scripts/generate_and_push_perf_report.sh

@@ -1,27 +1,24 @@
 #!/bin/bash
 
 # this is a shortcut script to avoid duplication in CI
-
 set -eux -o pipefail
 
 SCRIPT_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
 
-git clone https://$VIP_VAP_ACCESS_TOKEN@github.com/zenithdb/zenith-perf-data.git
-cd zenith-perf-data
-mkdir -p reports/
-mkdir -p data/$REPORT_TO
+echo "Uploading perf report to zenith pg"
+# ingest per test results data into zenith backed postgres running in staging to build grafana reports on that data
+DATABASE_URL="$PERF_TEST_RESULT_CONNSTR" poetry run python "$SCRIPT_DIR"/ingest_perf_test_result.py --ingest "$REPORT_FROM"
 
-cp $REPORT_FROM/* data/$REPORT_TO
+# Activate poetry's venv. Needed because git upload does not run in a project dir (it uses tmp to store the repository)
+# so the problem occurs because poetry cannot find pyproject.toml in temp dir created by git upload
+# shellcheck source=/dev/null
+. "$(poetry env info --path)"/bin/activate
 
-echo "Generating report"
-poetry run python $SCRIPT_DIR/generate_perf_report_page.py --input-dir data/$REPORT_TO --out reports/$REPORT_TO.html
-echo "Uploading perf result"
-git add data reports
-git \
-    -c "user.name=vipvap" \
-    -c "user.email=vipvap@zenith.tech" \
-    commit \
-    --author="vipvap <vipvap@zenith.tech>" \
-    -m "add performance test result for $GITHUB_SHA zenith revision"
-
-git push https://$VIP_VAP_ACCESS_TOKEN@github.com/zenithdb/zenith-perf-data.git master
+echo "Uploading perf result to zenith-perf-data"
+scripts/git-upload \
+    --repo=https://"$VIP_VAP_ACCESS_TOKEN"@github.com/zenithdb/zenith-perf-data.git \
+    --message="add performance test result for $GITHUB_SHA zenith revision" \
+    --branch=master \
+    copy "$REPORT_FROM" "data/$REPORT_TO" `# COPY FROM TO_RELATIVE`\
+    --merge \
+    --run-cmd "python $SCRIPT_DIR/generate_perf_report_page.py --input-dir data/$REPORT_TO --out reports/$REPORT_TO.html"

scripts/git-upload

@@ -1,7 +1,9 @@
 #!/usr/bin/env python3
 
 from contextlib import contextmanager
+import shlex
 from tempfile import TemporaryDirectory
+from distutils.dir_util import copy_tree
 from pathlib import Path
 
 import argparse
@@ -9,6 +11,8 @@ import os
 import shutil
 import subprocess
 import sys
+import textwrap
+from typing import Optional
 
 
 def absolute_path(path):
@@ -38,13 +42,21 @@ def run(cmd, *args, **kwargs):
 
 
 class GitRepo:
-    def __init__(self, url):
+    def __init__(self, url, branch: Optional[str] = None):
         self.url = url
         self.cwd = TemporaryDirectory()
+        self.branch = branch
 
-        subprocess.check_call([
+        args = [
             'git',
             'clone',
+            '--single-branch',
+        ]
+        if self.branch:
+            args.extend(['--branch', self.branch])
+
+        subprocess.check_call([
+            *args,
             str(url),
             self.cwd.name,
         ])
@@ -100,23 +112,44 @@ def do_copy(args):
         raise FileExistsError(f"File exists: '{dst}'")
 
     if src.is_dir():
-        shutil.rmtree(dst, ignore_errors=True)
-        shutil.copytree(src, dst)
+        if not args.merge:
+            shutil.rmtree(dst, ignore_errors=True)
+        # distutils is deprecated, but this is a temporary workaround before python version bump
+        # here we need dir_exists_ok=True from shutil.copytree which is available in python 3.8+
+        copy_tree(str(src), str(dst))
     else:
         shutil.copy(src, dst)
 
+    if args.run_cmd:
+        run(shlex.split(args.run_cmd))
+
 
 def main():
     parser = argparse.ArgumentParser(description='Git upload tool')
     parser.add_argument('--repo', type=str, metavar='URL', required=True, help='git repo url')
     parser.add_argument('--message', type=str, metavar='TEXT', help='commit message')
+    parser.add_argument('--branch', type=str, metavar='TEXT', help='target git repo branch')
 
     commands = parser.add_subparsers(title='commands', dest='subparser_name')
 
-    p_copy = commands.add_parser('copy', help='copy file into the repo')
+    p_copy = commands.add_parser(
+        'copy',
+        help='copy file into the repo',
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
     p_copy.add_argument('src', type=absolute_path, help='source path')
     p_copy.add_argument('dst', type=relative_path, help='relative dest path')
     p_copy.add_argument('--forbid-overwrite', action='store_true', help='do not allow overwrites')
+    p_copy.add_argument(
+        '--merge',
+        action='store_true',
+        help='when copying a directory do not delete existing data, but add new files')
+    p_copy.add_argument('--run-cmd',
+                        help=textwrap.dedent('''\
+                run arbitrary cmd on top of copied files,
+                example usage is static content generation
+                based on current repository state\
+            '''))
 
     args = parser.parse_args()
 
@@ -127,7 +160,7 @@ def main():
     action = commands.get(args.subparser_name)
     if action:
         message = args.message or 'update'
-        GitRepo(args.repo).update(message, lambda: action(args))
+        GitRepo(args.repo, args.branch).update(message, lambda: action(args))
     else:
         parser.print_usage()
 

scripts/ingest_perf_test_result.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+import argparse
+from contextlib import contextmanager
+import json
+import os
+import psycopg2
+import psycopg2.extras
+from pathlib import Path
+from datetime import datetime
+
+CREATE_TABLE = """
+CREATE TABLE IF NOT EXISTS perf_test_results (
+    id SERIAL PRIMARY KEY,
+    suit TEXT,
+    revision CHAR(40),
+    platform TEXT,
+    metric_name TEXT,
+    metric_value NUMERIC,
+    metric_unit VARCHAR(10),
+    metric_report_type TEXT,
+    recorded_at_timestamp TIMESTAMP WITH TIME ZONE DEFAULT NOW()
+)
+"""
+
+
+def err(msg):
+    print(f'error: {msg}')
+    exit(1)
+
+
+@contextmanager
+def get_connection_cursor():
+    connstr = os.getenv('DATABASE_URL')
+    if not connstr:
+        err('DATABASE_URL environment variable is not set')
+    with psycopg2.connect(connstr) as conn:
+        with conn.cursor() as cur:
+            yield cur
+
+
+def create_table(cur):
+    cur.execute(CREATE_TABLE)
+
+
+def ingest_perf_test_result(cursor, data_dile: Path, recorded_at_timestamp: int) -> int:
+    run_data = json.loads(data_dile.read_text())
+    revision = run_data['revision']
+    platform = run_data['platform']
+
+    run_result = run_data['result']
+    args_list = []
+
+    for suit_result in run_result:
+        suit = suit_result['suit']
+        total_duration = suit_result['total_duration']
+
+        suit_result['data'].append({
+            'name': 'total_duration',
+            'value': total_duration,
+            'unit': 's',
+            'report': 'lower_is_better',
+        })
+
+        for metric in suit_result['data']:
+            values = {
+                'suit': suit,
+                'revision': revision,
+                'platform': platform,
+                'metric_name': metric['name'],
+                'metric_value': metric['value'],
+                'metric_unit': metric['unit'],
+                'metric_report_type': metric['report'],
+                'recorded_at_timestamp': datetime.utcfromtimestamp(recorded_at_timestamp),
+            }
+            args_list.append(values)
+
+    psycopg2.extras.execute_values(
+        cursor,
+        """
+        INSERT INTO perf_test_results (
+            suit,
+            revision,
+            platform,
+            metric_name,
+            metric_value,
+            metric_unit,
+            metric_report_type,
+            recorded_at_timestamp
+        ) VALUES %s
+        """,
+        args_list,
+        template="""(
+            %(suit)s,
+            %(revision)s,
+            %(platform)s,
+            %(metric_name)s,
+            %(metric_value)s,
+            %(metric_unit)s,
+            %(metric_report_type)s,
+            %(recorded_at_timestamp)s
+        )""",
+    )
+    return len(args_list)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Perf test result uploader. \
+            Database connection string should be provided via DATABASE_URL environment variable', )
+    parser.add_argument(
+        '--ingest',
+        type=Path,
+        help='Path to perf test result file, or directory with perf test result files')
+    parser.add_argument('--initdb', action='store_true', help='Initialuze database')
+
+    args = parser.parse_args()
+    with get_connection_cursor() as cur:
+        if args.initdb:
+            create_table(cur)
+
+        if not args.ingest.exists():
+            err(f'ingest path {args.ingest} does not exist')
+
+        if args.ingest:
+            if args.ingest.is_dir():
+                for item in sorted(args.ingest.iterdir(), key=lambda x: int(x.name.split('_')[0])):
+                    recorded_at_timestamp = int(item.name.split('_')[0])
+                    ingested = ingest_perf_test_result(cur, item, recorded_at_timestamp)
+                    print(f'Ingested {ingested} metric values from {item}')
+            else:
+                recorded_at_timestamp = int(args.ingest.name.split('_')[0])
+                ingested = ingest_perf_test_result(cur, args.ingest, recorded_at_timestamp)
+                print(f'Ingested {ingested} metric values from {args.ingest}')
+
+
+if __name__ == '__main__':
+    main()

test_runner/README.md

@@ -89,7 +89,7 @@ def test_foobar(zenith_env_builder: ZenithEnvBuilder):
 
     # Now create the environment. This initializes the repository, and starts
     # up the page server and the safekeepers
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     # Run the test
     ...

test_runner/batch_others/test_auth.py

@@ -1,45 +1,49 @@
 from contextlib import closing
 from typing import Iterator
-from uuid import uuid4
+from uuid import UUID, uuid4
 import psycopg2
-from fixtures.zenith_fixtures import ZenithEnvBuilder
+from fixtures.zenith_fixtures import ZenithEnvBuilder, ZenithPageserverApiException
 import pytest
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 def test_pageserver_auth(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.pageserver_auth_enabled = True
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     ps = env.pageserver
 
-    tenant_token = env.auth_keys.generate_tenant_token(env.initial_tenant)
+    tenant_token = env.auth_keys.generate_tenant_token(env.initial_tenant.hex)
+    tenant_http_client = env.pageserver.http_client(tenant_token)
     invalid_tenant_token = env.auth_keys.generate_tenant_token(uuid4().hex)
+    invalid_tenant_http_client = env.pageserver.http_client(invalid_tenant_token)
+
     management_token = env.auth_keys.generate_management_token()
+    management_http_client = env.pageserver.http_client(management_token)
 
     # this does not invoke auth check and only decodes jwt and checks it for validity
     # check both tokens
-    ps.safe_psql("status", password=tenant_token)
-    ps.safe_psql("status", password=management_token)
+    ps.safe_psql("set FOO", password=tenant_token)
+    ps.safe_psql("set FOO", password=management_token)
 
     # tenant can create branches
-    ps.safe_psql(f"branch_create {env.initial_tenant} new1 main", password=tenant_token)
+    tenant_http_client.branch_create(env.initial_tenant, 'new1', 'main')
     # console can create branches for tenant
-    ps.safe_psql(f"branch_create {env.initial_tenant} new2 main", password=management_token)
+    management_http_client.branch_create(env.initial_tenant, 'new2', 'main')
 
-    # fail to create branch using token with different tenantid
-    with pytest.raises(psycopg2.DatabaseError, match='Tenant id mismatch. Permission denied'):
-        ps.safe_psql(f"branch_create {env.initial_tenant} new2 main", password=invalid_tenant_token)
+    # fail to create branch using token with different tenant_id
+    with pytest.raises(ZenithPageserverApiException,
+                       match='Forbidden: Tenant id mismatch. Permission denied'):
+        invalid_tenant_http_client.branch_create(env.initial_tenant, "new3", "main")
 
     # create tenant using management token
-    ps.safe_psql(f"tenant_create {uuid4().hex}", password=management_token)
+    management_http_client.tenant_create(uuid4())
 
     # fail to create tenant using tenant token
     with pytest.raises(
-            psycopg2.DatabaseError,
-            match='Attempt to access management api with tenant scope. Permission denied'):
-        ps.safe_psql(f"tenant_create {uuid4().hex}", password=tenant_token)
+            ZenithPageserverApiException,
+            match='Forbidden: Attempt to access management api with tenant scope. Permission denied'
+    ):
+        tenant_http_client.tenant_create(uuid4())
 
 
 @pytest.mark.parametrize('with_wal_acceptors', [False, True])
@@ -47,10 +51,10 @@ def test_compute_auth_to_pageserver(zenith_env_builder: ZenithEnvBuilder, with_w
     zenith_env_builder.pageserver_auth_enabled = True
     if with_wal_acceptors:
         zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     branch = f"test_compute_auth_to_pageserver{with_wal_acceptors}"
-    env.zenith_cli(["branch", branch, "main"])
+    env.zenith_cli.create_branch(branch, "main")
 
     pg = env.postgres.create_start(branch)
 

test_runner/batch_others/test_backpressure.py

@@ -0,0 +1,154 @@
+from contextlib import closing, contextmanager
+import psycopg2.extras
+from fixtures.zenith_fixtures import ZenithEnvBuilder
+from fixtures.log_helper import log
+import os
+import time
+import asyncpg
+from fixtures.zenith_fixtures import Postgres
+import threading
+
+pytest_plugins = ("fixtures.zenith_fixtures")
+
+
+@contextmanager
+def pg_cur(pg):
+    with closing(pg.connect()) as conn:
+        with conn.cursor() as cur:
+            yield cur
+
+
+# Periodically check that all backpressure lags are below the configured threshold,
+# assert if they are not.
+# If the check query fails, stop the thread. Main thread should notice that and stop the test.
+def check_backpressure(pg: Postgres, stop_event: threading.Event, polling_interval=5):
+    log.info("checks started")
+
+    with pg_cur(pg) as cur:
+        cur.execute("CREATE EXTENSION zenith")  # TODO move it to zenith_fixtures?
+
+        cur.execute("select pg_size_bytes(current_setting('max_replication_write_lag'))")
+        res = cur.fetchone()
+        max_replication_write_lag_bytes = res[0]
+        log.info(f"max_replication_write_lag: {max_replication_write_lag_bytes} bytes")
+
+        cur.execute("select pg_size_bytes(current_setting('max_replication_flush_lag'))")
+        res = cur.fetchone()
+        max_replication_flush_lag_bytes = res[0]
+        log.info(f"max_replication_flush_lag: {max_replication_flush_lag_bytes} bytes")
+
+        cur.execute("select pg_size_bytes(current_setting('max_replication_apply_lag'))")
+        res = cur.fetchone()
+        max_replication_apply_lag_bytes = res[0]
+        log.info(f"max_replication_apply_lag: {max_replication_apply_lag_bytes} bytes")
+
+    with pg_cur(pg) as cur:
+        while not stop_event.is_set():
+            try:
+                cur.execute('''
+                select pg_wal_lsn_diff(pg_current_wal_flush_lsn(),received_lsn) as received_lsn_lag,
+                pg_wal_lsn_diff(pg_current_wal_flush_lsn(),disk_consistent_lsn) as disk_consistent_lsn_lag,
+                pg_wal_lsn_diff(pg_current_wal_flush_lsn(),remote_consistent_lsn) as remote_consistent_lsn_lag,
+                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_flush_lsn(),received_lsn)),
+                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_flush_lsn(),disk_consistent_lsn)),
+                pg_size_pretty(pg_wal_lsn_diff(pg_current_wal_flush_lsn(),remote_consistent_lsn))
+                from backpressure_lsns();
+                ''')
+
+                res = cur.fetchone()
+                received_lsn_lag = res[0]
+                disk_consistent_lsn_lag = res[1]
+                remote_consistent_lsn_lag = res[2]
+
+                log.info(f"received_lsn_lag = {received_lsn_lag} ({res[3]}), "
+                         f"disk_consistent_lsn_lag = {disk_consistent_lsn_lag} ({res[4]}), "
+                         f"remote_consistent_lsn_lag = {remote_consistent_lsn_lag} ({res[5]})")
+
+                # Since feedback from pageserver is not immediate, we should allow some lag overflow
+                lag_overflow = 5 * 1024 * 1024  # 5MB
+
+                if max_replication_write_lag_bytes > 0:
+                    assert received_lsn_lag < max_replication_write_lag_bytes + lag_overflow
+                if max_replication_flush_lag_bytes > 0:
+                    assert disk_consistent_lsn_lag < max_replication_flush_lag_bytes + lag_overflow
+                if max_replication_apply_lag_bytes > 0:
+                    assert remote_consistent_lsn_lag < max_replication_apply_lag_bytes + lag_overflow
+
+                time.sleep(polling_interval)
+
+            except Exception as e:
+                log.info(f"backpressure check query failed: {e}")
+                stop_event.set()
+
+    log.info('check thread stopped')
+
+
+# This test illustrates how to tune backpressure to control the lag
+# between the WAL flushed on compute node and WAL digested by pageserver.
+#
+# To test it, throttle walreceiver ingest using failpoint and run heavy write load.
+# If backpressure is disabled or not tuned properly, the query will timeout, because the walreceiver cannot keep up.
+# If backpressure is enabled and tuned properly, insertion will be throttled, but the query will not timeout.
+
+
+def test_backpressure_received_lsn_lag(zenith_env_builder: ZenithEnvBuilder):
+    zenith_env_builder.num_safekeepers = 1
+    env = zenith_env_builder.init()
+    # Create a branch for us
+    env.zenith_cli.create_branch("test_backpressure", "main")
+
+    pg = env.postgres.create_start('test_backpressure',
+                                   config_lines=['max_replication_write_lag=30MB'])
+    log.info("postgres is running on 'test_backpressure' branch")
+
+    # setup check thread
+    check_stop_event = threading.Event()
+    check_thread = threading.Thread(target=check_backpressure, args=(pg, check_stop_event))
+    check_thread.start()
+
+    # Configure failpoint to slow down walreceiver ingest
+    with closing(env.pageserver.connect()) as psconn:
+        with psconn.cursor(cursor_factory=psycopg2.extras.DictCursor) as pscur:
+            pscur.execute("failpoints walreceiver-after-ingest=sleep(20)")
+
+    # FIXME
+    # Wait for the check thread to start
+    #
+    # Now if load starts too soon,
+    # check thread cannot auth, because it is not able to connect to the database
+    # because of the lag and waiting for lsn to replay to arrive.
+    time.sleep(2)
+
+    with pg_cur(pg) as cur:
+        # Create and initialize test table
+        cur.execute("CREATE TABLE foo(x bigint)")
+
+        inserts_to_do = 2000000
+        rows_inserted = 0
+
+        while check_thread.is_alive() and rows_inserted < inserts_to_do:
+            try:
+                cur.execute("INSERT INTO foo select from generate_series(1, 100000)")
+                rows_inserted += 100000
+            except Exception as e:
+                if check_thread.is_alive():
+                    log.info('stopping check thread')
+                    check_stop_event.set()
+                    check_thread.join()
+                    assert False, f"Exception {e} while inserting rows, but WAL lag is within configured threshold. That means backpressure is not tuned properly"
+                else:
+                    assert False, f"Exception {e} while inserting rows and WAL lag overflowed configured threshold. That means backpressure doesn't work."
+
+        log.info(f"inserted {rows_inserted} rows")
+
+    if check_thread.is_alive():
+        log.info('stopping check thread')
+        check_stop_event.set()
+        check_thread.join()
+        log.info('check thread stopped')
+    else:
+        assert False, "WAL lag overflowed configured threshold. That means backpressure doesn't work."
+
+
+#TODO test_backpressure_disk_consistent_lsn_lag. Play with pageserver's checkpoint settings
+#TODO test_backpressure_remote_consistent_lsn_lag

test_runner/batch_others/test_branch_behind.py

@@ -7,8 +7,6 @@ from fixtures.log_helper import log
 from fixtures.utils import print_gc_result
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Create a couple of branches off the main branch, at a historical point in time.
@@ -21,10 +19,10 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
     #
     # See https://github.com/zenithdb/zenith/issues/1068
     zenith_env_builder.num_safekeepers = 1
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     # Branch at the point where only 100 rows were inserted
-    env.zenith_cli(["branch", "test_branch_behind", "main"])
+    env.zenith_cli.create_branch("test_branch_behind", "main")
 
     pgmain = env.postgres.create_start('test_branch_behind')
     log.info("postgres is running on 'test_branch_behind' branch")
@@ -62,7 +60,7 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
     log.info(f'LSN after 200100 rows: {lsn_b}')
 
     # Branch at the point where only 100 rows were inserted
-    env.zenith_cli(["branch", "test_branch_behind_hundred", "test_branch_behind@" + lsn_a])
+    env.zenith_cli.create_branch("test_branch_behind_hundred", "test_branch_behind@" + lsn_a)
 
     # Insert many more rows. This generates enough WAL to fill a few segments.
     main_cur.execute('''
@@ -77,7 +75,7 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
     log.info(f'LSN after 400100 rows: {lsn_c}')
 
     # Branch at the point where only 200100 rows were inserted
-    env.zenith_cli(["branch", "test_branch_behind_more", "test_branch_behind@" + lsn_b])
+    env.zenith_cli.create_branch("test_branch_behind_more", "test_branch_behind@" + lsn_b)
 
     pg_hundred = env.postgres.create_start("test_branch_behind_hundred")
     pg_more = env.postgres.create_start("test_branch_behind_more")
@@ -101,7 +99,7 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
     # Check bad lsn's for branching
 
     # branch at segment boundary
-    env.zenith_cli(["branch", "test_branch_segment_boundary", "test_branch_behind@0/3000000"])
+    env.zenith_cli.create_branch("test_branch_segment_boundary", "test_branch_behind@0/3000000")
     pg = env.postgres.create_start("test_branch_segment_boundary")
     cur = pg.connect().cursor()
     cur.execute('SELECT 1')
@@ -109,23 +107,23 @@ def test_branch_behind(zenith_env_builder: ZenithEnvBuilder):
 
     # branch at pre-initdb lsn
     with pytest.raises(Exception, match="invalid branch start lsn"):
-        env.zenith_cli(["branch", "test_branch_preinitdb", "main@0/42"])
+        env.zenith_cli.create_branch("test_branch_preinitdb", "main@0/42")
 
     # branch at pre-ancestor lsn
     with pytest.raises(Exception, match="less than timeline ancestor lsn"):
-        env.zenith_cli(["branch", "test_branch_preinitdb", "test_branch_behind@0/42"])
+        env.zenith_cli.create_branch("test_branch_preinitdb", "test_branch_behind@0/42")
 
     # check that we cannot create branch based on garbage collected data
     with closing(env.pageserver.connect()) as psconn:
         with psconn.cursor(cursor_factory=psycopg2.extras.DictCursor) as pscur:
             # call gc to advace latest_gc_cutoff_lsn
-            pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+            pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
             row = pscur.fetchone()
             print_gc_result(row)
 
     with pytest.raises(Exception, match="invalid branch start lsn"):
         # this gced_lsn is pretty random, so if gc is disabled this woudln't fail
-        env.zenith_cli(["branch", "test_branch_create_fail", f"test_branch_behind@{gced_lsn}"])
+        env.zenith_cli.create_branch("test_branch_create_fail", f"test_branch_behind@{gced_lsn}")
 
     # check that after gc everything is still there
     hundred_cur.execute('SELECT count(*) FROM foo')

test_runner/batch_others/test_clog_truncate.py

@@ -6,16 +6,13 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test compute node start after clog truncation
 #
 def test_clog_truncate(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_clog_truncate", "empty"])
+    env.zenith_cli.create_branch("test_clog_truncate", "empty")
 
     # set agressive autovacuum to make sure that truncation will happen
     config = [
@@ -65,8 +62,8 @@ def test_clog_truncate(zenith_simple_env: ZenithEnv):
 
     # create new branch after clog truncation and start a compute node on it
     log.info(f'create branch at lsn_after_truncation {lsn_after_truncation}')
-    env.zenith_cli(
-        ["branch", "test_clog_truncate_new", "test_clog_truncate@" + lsn_after_truncation])
+    env.zenith_cli.create_branch("test_clog_truncate_new",
+                                 "test_clog_truncate@" + lsn_after_truncation)
 
     pg2 = env.postgres.create_start('test_clog_truncate_new')
     log.info('postgres is running on test_clog_truncate_new branch')

test_runner/batch_others/test_config.py

@@ -3,16 +3,13 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test starting Postgres with custom options
 #
 def test_config(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_config", "empty"])
+    env.zenith_cli.create_branch("test_config", "empty")
 
     # change config
     pg = env.postgres.create_start('test_config', config_lines=['log_min_messages=debug1'])

test_runner/batch_others/test_createdropdb.py

@@ -5,15 +5,13 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv, check_restored_datadir_content
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test CREATE DATABASE when there have been relmapper changes
 #
 def test_createdb(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli(["branch", "test_createdb", "empty"])
+    env.zenith_cli.create_branch("test_createdb", "empty")
 
     pg = env.postgres.create_start('test_createdb')
     log.info("postgres is running on 'test_createdb' branch")
@@ -29,7 +27,7 @@ def test_createdb(zenith_simple_env: ZenithEnv):
             lsn = cur.fetchone()[0]
 
     # Create a branch
-    env.zenith_cli(["branch", "test_createdb2", "test_createdb@" + lsn])
+    env.zenith_cli.create_branch("test_createdb2", "test_createdb@" + lsn)
 
     pg2 = env.postgres.create_start('test_createdb2')
 
@@ -43,7 +41,7 @@ def test_createdb(zenith_simple_env: ZenithEnv):
 #
 def test_dropdb(zenith_simple_env: ZenithEnv, test_output_dir):
     env = zenith_simple_env
-    env.zenith_cli(["branch", "test_dropdb", "empty"])
+    env.zenith_cli.create_branch("test_dropdb", "empty")
 
     pg = env.postgres.create_start('test_dropdb')
     log.info("postgres is running on 'test_dropdb' branch")
@@ -68,10 +66,10 @@ def test_dropdb(zenith_simple_env: ZenithEnv, test_output_dir):
             lsn_after_drop = cur.fetchone()[0]
 
     # Create two branches before and after database drop.
-    env.zenith_cli(["branch", "test_before_dropdb", "test_dropdb@" + lsn_before_drop])
+    env.zenith_cli.create_branch("test_before_dropdb", "test_dropdb@" + lsn_before_drop)
     pg_before = env.postgres.create_start('test_before_dropdb')
 
-    env.zenith_cli(["branch", "test_after_dropdb", "test_dropdb@" + lsn_after_drop])
+    env.zenith_cli.create_branch("test_after_dropdb", "test_dropdb@" + lsn_after_drop)
     pg_after = env.postgres.create_start('test_after_dropdb')
 
     # Test that database exists on the branch before drop

test_runner/batch_others/test_createuser.py

@@ -3,15 +3,13 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test CREATE USER to check shared catalog restore
 #
 def test_createuser(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli(["branch", "test_createuser", "empty"])
+    env.zenith_cli.create_branch("test_createuser", "empty")
 
     pg = env.postgres.create_start('test_createuser')
     log.info("postgres is running on 'test_createuser' branch")
@@ -27,7 +25,7 @@ def test_createuser(zenith_simple_env: ZenithEnv):
             lsn = cur.fetchone()[0]
 
     # Create a branch
-    env.zenith_cli(["branch", "test_createuser2", "test_createuser@" + lsn])
+    env.zenith_cli.create_branch("test_createuser2", "test_createuser@" + lsn)
 
     pg2 = env.postgres.create_start('test_createuser2')
 

test_runner/batch_others/test_gc_aggressive.py

@@ -7,8 +7,6 @@ import random
 from fixtures.zenith_fixtures import ZenithEnv, Postgres, Safekeeper
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 # Test configuration
 #
 # Create a table with {num_rows} rows, and perform {updates_to_perform} random
@@ -36,7 +34,7 @@ async def gc(env: ZenithEnv, timeline: str):
     psconn = await env.pageserver.connect_async()
 
     while updates_performed < updates_to_perform:
-        await psconn.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+        await psconn.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
 
 
 # At the same time, run UPDATEs and GC
@@ -57,9 +55,7 @@ async def update_and_gc(env: ZenithEnv, pg: Postgres, timeline: str):
 #
 def test_gc_aggressive(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_gc_aggressive", "empty"])
-
+    env.zenith_cli.create_branch("test_gc_aggressive", "empty")
     pg = env.postgres.create_start('test_gc_aggressive')
     log.info('postgres is running on test_gc_aggressive branch')
 

test_runner/batch_others/test_multixact.py

@@ -1,8 +1,6 @@
 from fixtures.zenith_fixtures import ZenithEnv, check_restored_datadir_content
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test multixact state after branching
@@ -12,8 +10,7 @@ pytest_plugins = ("fixtures.zenith_fixtures")
 #
 def test_multixact(zenith_simple_env: ZenithEnv, test_output_dir):
     env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_multixact", "empty"])
+    env.zenith_cli.create_branch("test_multixact", "empty")
     pg = env.postgres.create_start('test_multixact')
 
     log.info("postgres is running on 'test_multixact' branch")
@@ -63,7 +60,7 @@ def test_multixact(zenith_simple_env: ZenithEnv, test_output_dir):
     assert int(next_multixact_id) > int(next_multixact_id_old)
 
     # Branch at this point
-    env.zenith_cli(["branch", "test_multixact_new", "test_multixact@" + lsn])
+    env.zenith_cli.create_branch("test_multixact_new", "test_multixact@" + lsn)
     pg_new = env.postgres.create_start('test_multixact_new')
 
     log.info("postgres is running on 'test_multixact_new' branch")

test_runner/batch_others/test_next_xid.py

@@ -5,15 +5,13 @@ import time
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 # Test restarting page server, while safekeeper and compute node keep
 # running.
 def test_next_xid(zenith_env_builder: ZenithEnvBuilder):
     # One safekeeper is enough for this test.
     zenith_env_builder.num_safekeepers = 1
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     pg = env.postgres.create_start('main')
 

test_runner/batch_others/test_old_request_lsn.py

@@ -3,8 +3,6 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test where Postgres generates a lot of WAL, and it's garbage collected away, but
@@ -18,8 +16,7 @@ pytest_plugins = ("fixtures.zenith_fixtures")
 #
 def test_old_request_lsn(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_old_request_lsn", "empty"])
+    env.zenith_cli.create_branch("test_old_request_lsn", "empty")
     pg = env.postgres.create_start('test_old_request_lsn')
     log.info('postgres is running on test_old_request_lsn branch')
 
@@ -57,7 +54,7 @@ def test_old_request_lsn(zenith_simple_env: ZenithEnv):
     # Make a lot of updates on a single row, generating a lot of WAL. Trigger
     # garbage collections so that the page server will remove old page versions.
     for i in range(10):
-        pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+        pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
         for j in range(100):
             cur.execute('UPDATE foo SET val = val + 1 WHERE id = 1;')
 

test_runner/batch_others/test_pageserver_api.py

@@ -1,95 +1,22 @@
-import json
 from uuid import uuid4, UUID
 import pytest
-import psycopg2
-import requests
-from fixtures.zenith_fixtures import ZenithEnv, ZenithEnvBuilder, ZenithPageserverHttpClient
-from typing import cast
-
-pytest_plugins = ("fixtures.zenith_fixtures")
+from fixtures.zenith_fixtures import ZenithEnv, ZenithEnvBuilder, ZenithPageserverHttpClient, zenith_binpath
 
 
-def test_status_psql(zenith_simple_env: ZenithEnv):
-    env = zenith_simple_env
-    assert env.pageserver.safe_psql('status') == [
-        ('hello world', ),
-    ]
-
-
-def test_branch_list_psql(zenith_simple_env: ZenithEnv):
-    env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_branch_list_main", "empty"])
-
-    conn = env.pageserver.connect()
-    cur = conn.cursor()
-
-    cur.execute(f'branch_list {env.initial_tenant}')
-    branches = json.loads(cur.fetchone()[0])
-    # Filter out branches created by other tests
-    branches = [x for x in branches if x['name'].startswith('test_branch_list')]
-
-    assert len(branches) == 1
-    assert branches[0]['name'] == 'test_branch_list_main'
-    assert 'timeline_id' in branches[0]
-    assert 'latest_valid_lsn' in branches[0]
-    assert 'ancestor_id' in branches[0]
-    assert 'ancestor_lsn' in branches[0]
-
-    # Create another branch, and start Postgres on it
-    env.zenith_cli(['branch', 'test_branch_list_experimental', 'test_branch_list_main'])
-    env.zenith_cli(['pg', 'create', 'test_branch_list_experimental'])
-
-    cur.execute(f'branch_list {env.initial_tenant}')
-    new_branches = json.loads(cur.fetchone()[0])
-    # Filter out branches created by other tests
-    new_branches = [x for x in new_branches if x['name'].startswith('test_branch_list')]
-    assert len(new_branches) == 2
-    new_branches.sort(key=lambda k: k['name'])
-
-    assert new_branches[0]['name'] == 'test_branch_list_experimental'
-    assert new_branches[0]['timeline_id'] != branches[0]['timeline_id']
-
-    # TODO: do the LSNs have to match here?
-    assert new_branches[1] == branches[0]
-
-    conn.close()
-
-
-def test_tenant_list_psql(zenith_env_builder: ZenithEnvBuilder):
-    # don't use zenith_simple_env, because there might be other tenants there,
-    # left over from other tests.
+# test that we cannot override node id
+def test_pageserver_init_node_id(zenith_env_builder: ZenithEnvBuilder):
     env = zenith_env_builder.init()
-
-    res = env.zenith_cli(["tenant", "list"])
-    res.check_returncode()
-    tenants = sorted(map(lambda t: t.split()[0], res.stdout.splitlines()))
-    assert tenants == [env.initial_tenant]
-
-    conn = env.pageserver.connect()
-    cur = conn.cursor()
-
-    # check same tenant cannot be created twice
-    with pytest.raises(psycopg2.DatabaseError,
-                       match=f'repo for {env.initial_tenant} already exists'):
-        cur.execute(f'tenant_create {env.initial_tenant}')
-
-    # create one more tenant
-    tenant1 = uuid4().hex
-    cur.execute(f'tenant_create {tenant1}')
-
-    cur.execute('tenant_list')
-
-    # compare tenants list
-    new_tenants = sorted(map(lambda t: cast(str, t['id']), json.loads(cur.fetchone()[0])))
-    assert sorted([env.initial_tenant, tenant1]) == new_tenants
+    with pytest.raises(
+            Exception,
+            match="node id can only be set during pageserver init and cannot be overridden"):
+        env.pageserver.start(overrides=['--pageserver-config-override=id=10'])
 
 
-def check_client(client: ZenithPageserverHttpClient, initial_tenant: str):
+def check_client(client: ZenithPageserverHttpClient, initial_tenant: UUID):
     client.check_status()
 
     # check initial tenant is there
-    assert initial_tenant in {t['id'] for t in client.tenant_list()}
+    assert initial_tenant.hex in {t['id'] for t in client.tenant_list()}
 
     # create new tenant and check it is also there
     tenant_id = uuid4()
@@ -121,7 +48,7 @@ def test_pageserver_http_api_client(zenith_simple_env: ZenithEnv):
 
 def test_pageserver_http_api_client_auth_enabled(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.pageserver_auth_enabled = True
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     management_token = env.auth_keys.generate_management_token()
 

test_runner/batch_others/test_pageserver_catchup.py

@@ -7,8 +7,6 @@ from multiprocessing import Process, Value
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 # Test safekeeper sync and pageserver catch up
 # while initial compute node is down and pageserver is lagging behind safekeepers.
@@ -16,9 +14,9 @@ pytest_plugins = ("fixtures.zenith_fixtures")
 # and new compute node contains all data.
 def test_pageserver_catchup_while_compute_down(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_pageserver_catchup_while_compute_down", "main"])
+    env.zenith_cli.create_branch("test_pageserver_catchup_while_compute_down", "main")
     pg = env.postgres.create_start('test_pageserver_catchup_while_compute_down')
 
     pg_conn = pg.connect()

test_runner/batch_others/test_pageserver_restart.py

@@ -7,17 +7,15 @@ from multiprocessing import Process, Value
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 # Test restarting page server, while safekeeper and compute node keep
 # running.
 def test_pageserver_restart(zenith_env_builder: ZenithEnvBuilder):
     # One safekeeper is enough for this test.
     zenith_env_builder.num_safekeepers = 1
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_pageserver_restart", "main"])
+    env.zenith_cli.create_branch("test_pageserver_restart", "main")
     pg = env.postgres.create_start('test_pageserver_restart')
 
     pg_conn = pg.connect()

test_runner/batch_others/test_parallel_copy.py

@@ -5,8 +5,6 @@ import subprocess
 from fixtures.zenith_fixtures import ZenithEnv, Postgres
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 async def repeat_bytes(buf, repetitions: int):
     for i in range(repetitions):
@@ -39,9 +37,7 @@ async def parallel_load_same_table(pg: Postgres, n_parallel: int):
 # Load data into one table with COPY TO from 5 parallel connections
 def test_parallel_copy(zenith_simple_env: ZenithEnv, n_parallel=5):
     env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_parallel_copy", "empty"])
-
+    env.zenith_cli.create_branch("test_parallel_copy", "empty")
     pg = env.postgres.create_start('test_parallel_copy')
     log.info("postgres is running on 'test_parallel_copy' branch")
 

test_runner/batch_others/test_pgbench.py

@@ -1,14 +1,10 @@
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 def test_pgbench(zenith_simple_env: ZenithEnv, pg_bin):
     env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_pgbench", "empty"])
-
+    env.zenith_cli.create_branch("test_pgbench", "empty")
     pg = env.postgres.create_start('test_pgbench')
     log.info("postgres is running on 'test_pgbench' branch")
 

test_runner/batch_others/test_readonly_node.py

@@ -2,8 +2,6 @@ import pytest
 from fixtures.log_helper import log
 from fixtures.zenith_fixtures import ZenithEnv
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Create read-only compute nodes, anchored at historical points in time.
@@ -13,7 +11,7 @@ pytest_plugins = ("fixtures.zenith_fixtures")
 #
 def test_readonly_node(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli(["branch", "test_readonly_node", "empty"])
+    env.zenith_cli.create_branch("test_readonly_node", "empty")
 
     pgmain = env.postgres.create_start('test_readonly_node')
     log.info("postgres is running on 'test_readonly_node' branch")
@@ -88,4 +86,5 @@ def test_readonly_node(zenith_simple_env: ZenithEnv):
     # Create node at pre-initdb lsn
     with pytest.raises(Exception, match="invalid basebackup lsn"):
         # compute node startup with invalid LSN should fail
-        env.zenith_cli(["pg", "start", "test_readonly_node_preinitdb", "test_readonly_node@0/42"])
+        env.zenith_cli.pg_start("test_readonly_node_preinitdb",
+                                timeline_spec="test_readonly_node@0/42")

test_runner/batch_others/test_remote_storage.py

@@ -9,8 +9,6 @@ from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 import pytest
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Tests that a piece of data is backed up and restored correctly:
@@ -44,7 +42,7 @@ def test_remote_storage_backup_and_restore(zenith_env_builder: ZenithEnvBuilder,
     data_secret = 'very secret secret'
 
     ##### First start, insert secret data and upload it to the remote storage
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
     pg = env.postgres.create_start()
 
     tenant_id = pg.safe_psql("show zenith.zenith_tenant")[0][0]

test_runner/batch_others/test_restart_compute.py

@@ -4,8 +4,6 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test restarting and recreating a postgres instance
@@ -15,9 +13,9 @@ def test_restart_compute(zenith_env_builder: ZenithEnvBuilder, with_wal_acceptor
     zenith_env_builder.pageserver_auth_enabled = True
     if with_wal_acceptors:
         zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_restart_compute", "main"])
+    env.zenith_cli.create_branch("test_restart_compute", "main")
 
     pg = env.postgres.create_start('test_restart_compute')
     log.info("postgres is running on 'test_restart_compute' branch")

test_runner/batch_others/test_snapfiles_gc.py

@@ -5,8 +5,6 @@ from fixtures.utils import print_gc_result
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test Garbage Collection of old layer files
@@ -16,7 +14,7 @@ pytest_plugins = ("fixtures.zenith_fixtures")
 #
 def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli(["branch", "test_layerfiles_gc", "empty"])
+    env.zenith_cli.create_branch("test_layerfiles_gc", "empty")
     pg = env.postgres.create_start('test_layerfiles_gc')
 
     with closing(pg.connect()) as conn:
@@ -50,7 +48,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     cur.execute("DELETE FROM foo")
 
                     log.info("Running GC before test")
-                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     # remember the number of files
@@ -63,7 +61,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     # removing the old image and delta layer.
                     log.info("Inserting one row and running GC")
                     cur.execute("INSERT INTO foo VALUES (1)")
-                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     assert row['layer_relfiles_total'] == layer_relfiles_remain + 2
@@ -77,7 +75,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     cur.execute("INSERT INTO foo VALUES (2)")
                     cur.execute("INSERT INTO foo VALUES (3)")
 
-                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     assert row['layer_relfiles_total'] == layer_relfiles_remain + 2
@@ -89,7 +87,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     cur.execute("INSERT INTO foo VALUES (2)")
                     cur.execute("INSERT INTO foo VALUES (3)")
 
-                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     assert row['layer_relfiles_total'] == layer_relfiles_remain + 2
@@ -98,7 +96,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
 
                     # Run GC again, with no changes in the database. Should not remove anything.
                     log.info("Run GC again, with nothing to do")
-                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
                     assert row['layer_relfiles_total'] == layer_relfiles_remain
@@ -111,7 +109,7 @@ def test_layerfiles_gc(zenith_simple_env: ZenithEnv):
                     log.info("Drop table and run GC again")
                     cur.execute("DROP TABLE foo")
 
-                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
+                    pscur.execute(f"do_gc {env.initial_tenant.hex} {timeline} 0")
                     row = pscur.fetchone()
                     print_gc_result(row)
 

test_runner/batch_others/test_subxacts.py

@@ -1,8 +1,6 @@
 from fixtures.zenith_fixtures import ZenithEnv, check_restored_datadir_content
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 # Test subtransactions
 #
@@ -12,8 +10,7 @@ pytest_plugins = ("fixtures.zenith_fixtures")
 # CLOG.
 def test_subxacts(zenith_simple_env: ZenithEnv, test_output_dir):
     env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_subxacts", "empty"])
+    env.zenith_cli.create_branch("test_subxacts", "empty")
     pg = env.postgres.create_start('test_subxacts')
 
     log.info("postgres is running on 'test_subxacts' branch")

test_runner/batch_others/test_tenant_relocation.py

@@ -108,8 +108,8 @@ def load(pg: Postgres, stop_event: threading.Event, load_ok_event: threading.Eve
     log.info('load thread stopped')
 
 
-def assert_local(pageserver_http_client: ZenithPageserverHttpClient, tenant: str, timeline: str):
-    timeline_detail = pageserver_http_client.timeline_detail(UUID(tenant), UUID(timeline))
+def assert_local(pageserver_http_client: ZenithPageserverHttpClient, tenant: UUID, timeline: str):
+    timeline_detail = pageserver_http_client.timeline_detail(tenant, UUID(timeline))
     assert timeline_detail.get('type') == "Local", timeline_detail
     return timeline_detail
 
@@ -122,15 +122,15 @@ def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
     zenith_env_builder.num_safekeepers = 1
     zenith_env_builder.enable_local_fs_remote_storage()
 
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     # create folder for remote storage mock
     remote_storage_mock_path = env.repo_dir / 'local_fs_remote_storage'
 
-    tenant = env.create_tenant("74ee8b079a0e437eb0afea7d26a07209")
+    tenant = env.create_tenant(UUID("74ee8b079a0e437eb0afea7d26a07209"))
     log.info("tenant to relocate %s", tenant)
 
-    env.zenith_cli(["branch", "test_tenant_relocation", "main", f"--tenantid={tenant}"])
+    env.zenith_cli.create_branch("test_tenant_relocation", "main", tenant_id=tenant)
 
     tenant_pg = env.postgres.create_start(
         "test_tenant_relocation",
@@ -167,11 +167,11 @@ def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
     # run checkpoint manually to be sure that data landed in remote storage
     with closing(env.pageserver.connect()) as psconn:
         with psconn.cursor() as pscur:
-            pscur.execute(f"do_gc {tenant} {timeline}")
+            pscur.execute(f"do_gc {tenant.hex} {timeline}")
 
     # ensure upload is completed
     pageserver_http_client = env.pageserver.http_client()
-    timeline_detail = pageserver_http_client.timeline_detail(UUID(tenant), UUID(timeline))
+    timeline_detail = pageserver_http_client.timeline_detail(tenant, UUID(timeline))
     assert timeline_detail['disk_consistent_lsn'] == timeline_detail['timeline_state']['Ready']
 
     log.info("inititalizing new pageserver")
@@ -194,7 +194,7 @@ def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
                                new_pageserver_http_port):
 
         # call to attach timeline to new pageserver
-        new_pageserver_http_client.timeline_attach(UUID(tenant), UUID(timeline))
+        new_pageserver_http_client.timeline_attach(tenant, UUID(timeline))
         # FIXME cannot handle duplicate download requests, subject to fix in https://github.com/zenithdb/zenith/issues/997
         time.sleep(5)
         # new pageserver should in sync (modulo wal tail or vacuum activity) with the old one because there was no new writes since checkpoint
@@ -241,7 +241,7 @@ def test_tenant_relocation(zenith_env_builder: ZenithEnvBuilder,
         # detach tenant from old pageserver before we check
         # that all the data is there to be sure that old pageserver
         # is no longer involved, and if it is, we will see the errors
-        pageserver_http_client.timeline_detach(UUID(tenant), UUID(timeline))
+        pageserver_http_client.timeline_detach(tenant, UUID(timeline))
 
         with pg_cur(tenant_pg) as cur:
             # check that data is still there

test_runner/batch_others/test_tenants.py

@@ -10,23 +10,17 @@ def test_tenants_normal_work(zenith_env_builder: ZenithEnvBuilder, with_wal_acce
     if with_wal_acceptors:
         zenith_env_builder.num_safekeepers = 3
 
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
     """Tests tenants with and without wal acceptors"""
     tenant_1 = env.create_tenant()
     tenant_2 = env.create_tenant()
 
-    env.zenith_cli([
-        "branch",
-        f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",
-        "main",
-        f"--tenantid={tenant_1}"
-    ])
-    env.zenith_cli([
-        "branch",
-        f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",
-        "main",
-        f"--tenantid={tenant_2}"
-    ])
+    env.zenith_cli.create_branch(f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",
+                                 "main",
+                                 tenant_id=tenant_1)
+    env.zenith_cli.create_branch(f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",
+                                 "main",
+                                 tenant_id=tenant_2)
 
     pg_tenant1 = env.postgres.create_start(
         f"test_tenants_normal_work_with_wal_acceptors{with_wal_acceptors}",

test_runner/batch_others/test_timeline_size.py

@@ -1,17 +1,19 @@
 from contextlib import closing
 from uuid import UUID
 import psycopg2.extras
-from fixtures.zenith_fixtures import ZenithEnv
+import psycopg2.errors
+from fixtures.zenith_fixtures import ZenithEnv, ZenithEnvBuilder, Postgres
 from fixtures.log_helper import log
+import time
 
 
 def test_timeline_size(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
     # Branch at the point where only 100 rows were inserted
-    env.zenith_cli(["branch", "test_timeline_size", "empty"])
+    env.zenith_cli.create_branch("test_timeline_size", "empty")
 
     client = env.pageserver.http_client()
-    res = client.branch_detail(UUID(env.initial_tenant), "test_timeline_size")
+    res = client.branch_detail(env.initial_tenant, "test_timeline_size")
     assert res["current_logical_size"] == res["current_logical_size_non_incremental"]
 
     pgmain = env.postgres.create_start("test_timeline_size")
@@ -29,9 +31,102 @@ def test_timeline_size(zenith_simple_env: ZenithEnv):
                     FROM generate_series(1, 10) g
             """)
 
-            res = client.branch_detail(UUID(env.initial_tenant), "test_timeline_size")
+            res = client.branch_detail(env.initial_tenant, "test_timeline_size")
             assert res["current_logical_size"] == res["current_logical_size_non_incremental"]
             cur.execute("TRUNCATE foo")
 
-            res = client.branch_detail(UUID(env.initial_tenant), "test_timeline_size")
+            res = client.branch_detail(env.initial_tenant, "test_timeline_size")
             assert res["current_logical_size"] == res["current_logical_size_non_incremental"]
+
+
+# wait until received_lsn_lag is 0
+def wait_for_pageserver_catchup(pgmain: Postgres, polling_interval=1, timeout=60):
+    started_at = time.time()
+
+    received_lsn_lag = 1
+    while received_lsn_lag > 0:
+        elapsed = time.time() - started_at
+        if elapsed > timeout:
+            raise RuntimeError(
+                f"timed out waiting for pageserver to reach pg_current_wal_flush_lsn()")
+
+        with closing(pgmain.connect()) as conn:
+            with conn.cursor() as cur:
+
+                cur.execute('''
+                    select  pg_size_pretty(pg_cluster_size()),
+                    pg_wal_lsn_diff(pg_current_wal_flush_lsn(),received_lsn) as received_lsn_lag
+                    FROM backpressure_lsns();
+                ''')
+                res = cur.fetchone()
+                log.info(f"pg_cluster_size = {res[0]}, received_lsn_lag = {res[1]}")
+                received_lsn_lag = res[1]
+
+        time.sleep(polling_interval)
+
+
+def test_timeline_size_quota(zenith_env_builder: ZenithEnvBuilder):
+    zenith_env_builder.num_safekeepers = 1
+    env = zenith_env_builder.init_start()
+    env.zenith_cli.create_branch("test_timeline_size_quota", "main")
+
+    client = env.pageserver.http_client()
+    res = client.branch_detail(env.initial_tenant, "test_timeline_size_quota")
+    assert res["current_logical_size"] == res["current_logical_size_non_incremental"]
+
+    pgmain = env.postgres.create_start(
+        "test_timeline_size_quota",
+        # Set small limit for the test
+        config_lines=['zenith.max_cluster_size=30MB'],
+    )
+    log.info("postgres is running on 'test_timeline_size_quota' branch")
+
+    with closing(pgmain.connect()) as conn:
+        with conn.cursor() as cur:
+            cur.execute("CREATE EXTENSION zenith")  # TODO move it to zenith_fixtures?
+
+            cur.execute("CREATE TABLE foo (t text)")
+
+            wait_for_pageserver_catchup(pgmain)
+
+            # Insert many rows. This query must fail because of space limit
+            try:
+                cur.execute('''
+                    INSERT INTO foo
+                        SELECT 'long string to consume some space' || g
+                        FROM generate_series(1, 100000) g
+                ''')
+
+                wait_for_pageserver_catchup(pgmain)
+
+                cur.execute('''
+                    INSERT INTO foo
+                        SELECT 'long string to consume some space' || g
+                        FROM generate_series(1, 500000) g
+                ''')
+
+                # If we get here, the timeline size limit failed
+                log.error("Query unexpectedly succeeded")
+                assert False
+
+            except psycopg2.errors.DiskFull as err:
+                log.info(f"Query expectedly failed with: {err}")
+
+            # drop table to free space
+            cur.execute('DROP TABLE foo')
+
+            wait_for_pageserver_catchup(pgmain)
+
+            # create it again and insert some rows. This query must succeed
+            cur.execute("CREATE TABLE foo (t text)")
+            cur.execute('''
+                INSERT INTO foo
+                    SELECT 'long string to consume some space' || g
+                    FROM generate_series(1, 10000) g
+            ''')
+
+            wait_for_pageserver_catchup(pgmain)
+
+            cur.execute("SELECT * from pg_size_pretty(pg_cluster_size())")
+            pg_cluster_size = cur.fetchone()
+            log.info(f"pg_cluster_size = {pg_cluster_size}")

test_runner/batch_others/test_twophase.py

@@ -3,15 +3,13 @@ import os
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test branching, when a transaction is in prepared state
 #
 def test_twophase(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    env.zenith_cli(["branch", "test_twophase", "empty"])
+    env.zenith_cli.create_branch("test_twophase", "empty")
 
     pg = env.postgres.create_start('test_twophase', config_lines=['max_prepared_transactions=5'])
     log.info("postgres is running on 'test_twophase' branch")
@@ -58,7 +56,7 @@ def test_twophase(zenith_simple_env: ZenithEnv):
     assert len(twophase_files) == 2
 
     # Create a branch with the transaction in prepared state
-    env.zenith_cli(["branch", "test_twophase_prepared", "test_twophase"])
+    env.zenith_cli.create_branch("test_twophase_prepared", "test_twophase")
 
     # Start compute on the new branch
     pg2 = env.postgres.create_start(

test_runner/batch_others/test_vm_bits.py

@@ -1,8 +1,6 @@
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 #
 # Test that the VM bit is cleared correctly at a HEAP_DELETE and
@@ -11,8 +9,7 @@ pytest_plugins = ("fixtures.zenith_fixtures")
 def test_vm_bit_clear(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
 
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_vm_bit_clear", "empty"])
+    env.zenith_cli.create_branch("test_vm_bit_clear", "empty")
     pg = env.postgres.create_start('test_vm_bit_clear')
 
     log.info("postgres is running on 'test_vm_bit_clear' branch")
@@ -36,7 +33,7 @@ def test_vm_bit_clear(zenith_simple_env: ZenithEnv):
     cur.execute('UPDATE vmtest_update SET id = 5000 WHERE id = 1')
 
     # Branch at this point, to test that later
-    env.zenith_cli(["branch", "test_vm_bit_clear_new", "test_vm_bit_clear"])
+    env.zenith_cli.create_branch("test_vm_bit_clear_new", "test_vm_bit_clear")
 
     # Clear the buffer cache, to force the VM page to be re-fetched from
     # the page server

test_runner/batch_others/test_wal_acceptor.py

@@ -17,16 +17,14 @@ from fixtures.utils import lsn_to_hex, mkdir_if_needed
 from fixtures.log_helper import log
 from typing import List, Optional, Any
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 # basic test, write something in setup with wal acceptors, ensure that commits
 # succeed and data is written
 def test_normal_work(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_wal_acceptors_normal_work", "main"])
+    env.zenith_cli.create_branch("test_wal_acceptors_normal_work", "main")
 
     pg = env.postgres.create_start('test_wal_acceptors_normal_work')
 
@@ -53,7 +51,7 @@ class BranchMetrics:
 # against different timelines.
 def test_many_timelines(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     n_timelines = 3
 
@@ -62,10 +60,10 @@ def test_many_timelines(zenith_env_builder: ZenithEnvBuilder):
     # start postgres on each timeline
     pgs = []
     for branch in branches:
-        env.zenith_cli(["branch", branch, "main"])
+        env.zenith_cli.create_branch(branch, "main")
         pgs.append(env.postgres.create_start(branch))
 
-    tenant_id = uuid.UUID(env.initial_tenant)
+    tenant_id = env.initial_tenant
 
     def collect_metrics(message: str) -> List[BranchMetrics]:
         with env.pageserver.http_client() as pageserver_http:
@@ -92,8 +90,8 @@ def test_many_timelines(zenith_env_builder: ZenithEnvBuilder):
                     latest_valid_lsn=branch_detail["latest_valid_lsn"],
                 )
                 for sk_m in sk_metrics:
-                    m.flush_lsns.append(sk_m.flush_lsn_inexact[timeline_id])
-                    m.commit_lsns.append(sk_m.commit_lsn_inexact[timeline_id])
+                    m.flush_lsns.append(sk_m.flush_lsn_inexact[(tenant_id.hex, timeline_id)])
+                    m.commit_lsns.append(sk_m.commit_lsn_inexact[(tenant_id.hex, timeline_id)])
 
                 for flush_lsn, commit_lsn in zip(m.flush_lsns, m.commit_lsns):
                     # Invariant. May be < when transaction is in progress.
@@ -183,9 +181,9 @@ def test_restarts(zenith_env_builder: ZenithEnvBuilder):
     n_acceptors = 3
 
     zenith_env_builder.num_safekeepers = n_acceptors
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_wal_acceptors_restarts", "main"])
+    env.zenith_cli.create_branch("test_wal_acceptors_restarts", "main")
     pg = env.postgres.create_start('test_wal_acceptors_restarts')
 
     # we rely upon autocommit after each statement
@@ -220,9 +218,9 @@ def delayed_wal_acceptor_start(wa):
 # When majority of acceptors is offline, commits are expected to be frozen
 def test_unavailability(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 2
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_wal_acceptors_unavailability", "main"])
+    env.zenith_cli.create_branch("test_wal_acceptors_unavailability", "main")
     pg = env.postgres.create_start('test_wal_acceptors_unavailability')
 
     # we rely upon autocommit after each statement
@@ -291,9 +289,9 @@ def stop_value():
 def test_race_conditions(zenith_env_builder: ZenithEnvBuilder, stop_value):
 
     zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_wal_acceptors_race_conditions", "main"])
+    env.zenith_cli.create_branch("test_wal_acceptors_race_conditions", "main")
     pg = env.postgres.create_start('test_wal_acceptors_race_conditions')
 
     # we rely upon autocommit after each statement
@@ -321,16 +319,16 @@ class ProposerPostgres(PgProtocol):
     def __init__(self,
                  pgdata_dir: str,
                  pg_bin,
-                 timeline_id: str,
-                 tenant_id: str,
+                 timeline_id: uuid.UUID,
+                 tenant_id: uuid.UUID,
                  listen_addr: str,
                  port: int):
-        super().__init__(host=listen_addr, port=port)
+        super().__init__(host=listen_addr, port=port, username='zenith_admin')
 
         self.pgdata_dir: str = pgdata_dir
         self.pg_bin: PgBin = pg_bin
-        self.timeline_id: str = timeline_id
-        self.tenant_id: str = tenant_id
+        self.timeline_id: uuid.UUID = timeline_id
+        self.tenant_id: uuid.UUID = tenant_id
         self.listen_addr: str = listen_addr
         self.port: int = port
 
@@ -350,8 +348,8 @@ class ProposerPostgres(PgProtocol):
             cfg = [
                 "synchronous_standby_names = 'walproposer'\n",
                 "shared_preload_libraries = 'zenith'\n",
-                f"zenith.zenith_timeline = '{self.timeline_id}'\n",
-                f"zenith.zenith_tenant = '{self.tenant_id}'\n",
+                f"zenith.zenith_timeline = '{self.timeline_id.hex}'\n",
+                f"zenith.zenith_tenant = '{self.tenant_id.hex}'\n",
                 f"zenith.page_server_connstring = ''\n",
                 f"wal_acceptors = '{wal_acceptors}'\n",
                 f"listen_addresses = '{self.listen_addr}'\n",
@@ -406,10 +404,10 @@ def test_sync_safekeepers(zenith_env_builder: ZenithEnvBuilder,
     # We don't really need the full environment for this test, just the
     # safekeepers would be enough.
     zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    timeline_id = uuid.uuid4().hex
-    tenant_id = uuid.uuid4().hex
+    timeline_id = uuid.uuid4()
+    tenant_id = uuid.uuid4()
 
     # write config for proposer
     pgdata_dir = os.path.join(env.repo_dir, "proposer_pgdata")
@@ -456,9 +454,9 @@ def test_sync_safekeepers(zenith_env_builder: ZenithEnvBuilder,
 def test_timeline_status(zenith_env_builder: ZenithEnvBuilder):
 
     zenith_env_builder.num_safekeepers = 1
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_timeline_status", "main"])
+    env.zenith_cli.create_branch("test_timeline_status", "main")
     pg = env.postgres.create_start('test_timeline_status')
 
     wa = env.safekeepers[0]
@@ -495,15 +493,15 @@ class SafekeeperEnv:
         self.bin_safekeeper = os.path.join(str(zenith_binpath), 'safekeeper')
         self.safekeepers: Optional[List[subprocess.CompletedProcess[Any]]] = None
         self.postgres: Optional[ProposerPostgres] = None
-        self.tenant_id: Optional[str] = None
-        self.timeline_id: Optional[str] = None
+        self.tenant_id: Optional[uuid.UUID] = None
+        self.timeline_id: Optional[uuid.UUID] = None
 
     def init(self) -> "SafekeeperEnv":
         assert self.postgres is None, "postgres is already initialized"
         assert self.safekeepers is None, "safekeepers are already initialized"
 
-        self.timeline_id = uuid.uuid4().hex
-        self.tenant_id = uuid.uuid4().hex
+        self.timeline_id = uuid.uuid4()
+        self.tenant_id = uuid.uuid4()
         mkdir_if_needed(str(self.repo_dir))
 
         # Create config and a Safekeeper object for each safekeeper
@@ -523,12 +521,7 @@ class SafekeeperEnv:
             http=self.port_distributor.get_port(),
         )
 
-        if self.num_safekeepers == 1:
-            name = "single"
-        else:
-            name = f"sk{i}"
-
-        safekeeper_dir = os.path.join(self.repo_dir, name)
+        safekeeper_dir = os.path.join(self.repo_dir, f"sk{i}")
         mkdir_if_needed(safekeeper_dir)
 
         args = [
@@ -539,6 +532,8 @@ class SafekeeperEnv:
             f"127.0.0.1:{port.http}",
             "-D",
             safekeeper_dir,
+            "--id",
+            str(i),
             "--daemonize"
         ]
 
@@ -606,9 +601,8 @@ def test_safekeeper_without_pageserver(test_output_dir: str,
 
 
 def test_replace_safekeeper(zenith_env_builder: ZenithEnvBuilder):
-    def safekeepers_guc(env: ZenithEnv, sk_names: List[str]) -> str:
-        return ','.join(
-            [f'localhost:{sk.port.pg}' for sk in env.safekeepers if sk.name in sk_names])
+    def safekeepers_guc(env: ZenithEnv, sk_names: List[int]) -> str:
+        return ','.join([f'localhost:{sk.port.pg}' for sk in env.safekeepers if sk.id in sk_names])
 
     def execute_payload(pg: Postgres):
         with closing(pg.connect()) as conn:
@@ -630,17 +624,17 @@ def test_replace_safekeeper(zenith_env_builder: ZenithEnvBuilder):
             http_cli = sk.http_client()
             try:
                 status = http_cli.timeline_status(tenant_id, timeline_id)
-                log.info(f"Safekeeper {sk.name} status: {status}")
+                log.info(f"Safekeeper {sk.id} status: {status}")
             except Exception as e:
-                log.info(f"Safekeeper {sk.name} status error: {e}")
+                log.info(f"Safekeeper {sk.id} status error: {e}")
 
     zenith_env_builder.num_safekeepers = 4
-    env = zenith_env_builder.init()
-    env.zenith_cli(["branch", "test_replace_safekeeper", "main"])
+    env = zenith_env_builder.init_start()
+    env.zenith_cli.create_branch("test_replace_safekeeper", "main")
 
     log.info("Use only first 3 safekeepers")
     env.safekeepers[3].stop()
-    active_safekeepers = ['sk1', 'sk2', 'sk3']
+    active_safekeepers = [1, 2, 3]
     pg = env.postgres.create('test_replace_safekeeper')
     pg.adjust_for_wal_acceptors(safekeepers_guc(env, active_safekeepers))
     pg.start()
@@ -680,7 +674,7 @@ def test_replace_safekeeper(zenith_env_builder: ZenithEnvBuilder):
 
     log.info("Recreate postgres to replace failed sk1 with new sk4")
     pg.stop_and_destroy().create('test_replace_safekeeper')
-    active_safekeepers = ['sk2', 'sk3', 'sk4']
+    active_safekeepers = [2, 3, 4]
     env.safekeepers[3].start()
     pg.adjust_for_wal_acceptors(safekeepers_guc(env, active_safekeepers))
     pg.start()

test_runner/batch_others/test_wal_acceptor_async.py

@@ -9,7 +9,6 @@ from fixtures.utils import lsn_from_hex, lsn_to_hex
 from typing import List
 
 log = getLogger('root.wal_acceptor_async')
-pytest_plugins = ("fixtures.zenith_fixtures")
 
 
 class BankClient(object):
@@ -201,9 +200,9 @@ async def run_restarts_under_load(pg: Postgres, acceptors: List[Safekeeper], n_w
 # restart acceptors one by one, while executing and validating bank transactions
 def test_restarts_under_load(zenith_env_builder: ZenithEnvBuilder):
     zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
-    env.zenith_cli(["branch", "test_wal_acceptors_restarts_under_load", "main"])
+    env.zenith_cli.create_branch("test_wal_acceptors_restarts_under_load", "main")
     pg = env.postgres.create_start('test_wal_acceptors_restarts_under_load')
 
     asyncio.run(run_restarts_under_load(pg, env.safekeepers))

test_runner/batch_others/test_zenith_cli.py

@@ -1,31 +1,28 @@
 import json
 import uuid
+import requests
 
 from psycopg2.extensions import cursor as PgCursor
-from fixtures.zenith_fixtures import ZenithEnv
+from fixtures.zenith_fixtures import ZenithEnv, ZenithEnvBuilder, ZenithPageserverHttpClient
 from typing import cast
 
-pytest_plugins = ("fixtures.zenith_fixtures")
 
-
-def helper_compare_branch_list(page_server_cur: PgCursor, env: ZenithEnv, initial_tenant: str):
+def helper_compare_branch_list(pageserver_http_client: ZenithPageserverHttpClient,
+                               env: ZenithEnv,
+                               initial_tenant: uuid.UUID):
     """
     Compare branches list returned by CLI and directly via API.
     Filters out branches created by other tests.
     """
-
-    page_server_cur.execute(f'branch_list {initial_tenant}')
-    branches_api = sorted(
-        map(lambda b: cast(str, b['name']), json.loads(page_server_cur.fetchone()[0])))
+    branches = pageserver_http_client.branch_list(initial_tenant)
+    branches_api = sorted(map(lambda b: cast(str, b['name']), branches))
     branches_api = [b for b in branches_api if b.startswith('test_cli_') or b in ('empty', 'main')]
 
-    res = env.zenith_cli(["branch"])
-    res.check_returncode()
+    res = env.zenith_cli.list_branches()
     branches_cli = sorted(map(lambda b: b.split(':')[-1].strip(), res.stdout.strip().split("\n")))
     branches_cli = [b for b in branches_cli if b.startswith('test_cli_') or b in ('empty', 'main')]
 
-    res = env.zenith_cli(["branch", f"--tenantid={initial_tenant}"])
-    res.check_returncode()
+    res = env.zenith_cli.list_branches(tenant_id=initial_tenant)
     branches_cli_with_tenant_arg = sorted(
         map(lambda b: b.split(':')[-1].strip(), res.stdout.strip().split("\n")))
     branches_cli_with_tenant_arg = [
@@ -37,24 +34,20 @@ def helper_compare_branch_list(page_server_cur: PgCursor, env: ZenithEnv, initia
 
 def test_cli_branch_list(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    page_server_conn = env.pageserver.connect()
-    page_server_cur = page_server_conn.cursor()
+    pageserver_http_client = env.pageserver.http_client()
 
     # Initial sanity check
-    helper_compare_branch_list(page_server_cur, env, env.initial_tenant)
-
-    # Create a branch for us
-    res = env.zenith_cli(["branch", "test_cli_branch_list_main", "empty"])
-    assert res.stderr == ''
-    helper_compare_branch_list(page_server_cur, env, env.initial_tenant)
+    helper_compare_branch_list(pageserver_http_client, env, env.initial_tenant)
+    env.zenith_cli.create_branch("test_cli_branch_list_main", "empty")
+    helper_compare_branch_list(pageserver_http_client, env, env.initial_tenant)
 
     # Create a nested branch
-    res = env.zenith_cli(["branch", "test_cli_branch_list_nested", "test_cli_branch_list_main"])
+    res = env.zenith_cli.create_branch("test_cli_branch_list_nested", "test_cli_branch_list_main")
     assert res.stderr == ''
-    helper_compare_branch_list(page_server_cur, env, env.initial_tenant)
+    helper_compare_branch_list(pageserver_http_client, env, env.initial_tenant)
 
     # Check that all new branches are visible via CLI
-    res = env.zenith_cli(["branch"])
+    res = env.zenith_cli.list_branches()
     assert res.stderr == ''
     branches_cli = sorted(map(lambda b: b.split(':')[-1].strip(), res.stdout.strip().split("\n")))
 
@@ -62,12 +55,11 @@ def test_cli_branch_list(zenith_simple_env: ZenithEnv):
     assert 'test_cli_branch_list_nested' in branches_cli
 
 
-def helper_compare_tenant_list(page_server_cur: PgCursor, env: ZenithEnv):
-    page_server_cur.execute(f'tenant_list')
-    tenants_api = sorted(
-        map(lambda t: cast(str, t['id']), json.loads(page_server_cur.fetchone()[0])))
+def helper_compare_tenant_list(pageserver_http_client: ZenithPageserverHttpClient, env: ZenithEnv):
+    tenants = pageserver_http_client.tenant_list()
+    tenants_api = sorted(map(lambda t: cast(str, t['id']), tenants))
 
-    res = env.zenith_cli(["tenant", "list"])
+    res = env.zenith_cli.list_tenants()
     assert res.stderr == ''
     tenants_cli = sorted(map(lambda t: t.split()[0], res.stdout.splitlines()))
 
@@ -76,32 +68,62 @@ def helper_compare_tenant_list(page_server_cur: PgCursor, env: ZenithEnv):
 
 def test_cli_tenant_list(zenith_simple_env: ZenithEnv):
     env = zenith_simple_env
-    page_server_conn = env.pageserver.connect()
-    page_server_cur = page_server_conn.cursor()
-
+    pageserver_http_client = env.pageserver.http_client()
     # Initial sanity check
-    helper_compare_tenant_list(page_server_cur, env)
+    helper_compare_tenant_list(pageserver_http_client, env)
 
     # Create new tenant
-    tenant1 = uuid.uuid4().hex
-    res = env.zenith_cli(["tenant", "create", tenant1])
-    res.check_returncode()
+    tenant1 = uuid.uuid4()
+    env.zenith_cli.create_tenant(tenant1)
 
     # check tenant1 appeared
-    helper_compare_tenant_list(page_server_cur, env)
+    helper_compare_tenant_list(pageserver_http_client, env)
 
     # Create new tenant
-    tenant2 = uuid.uuid4().hex
-    res = env.zenith_cli(["tenant", "create", tenant2])
-    res.check_returncode()
+    tenant2 = uuid.uuid4()
+    env.zenith_cli.create_tenant(tenant2)
 
     # check tenant2 appeared
-    helper_compare_tenant_list(page_server_cur, env)
+    helper_compare_tenant_list(pageserver_http_client, env)
 
-    res = env.zenith_cli(["tenant", "list"])
-    res.check_returncode()
+    res = env.zenith_cli.list_tenants()
     tenants = sorted(map(lambda t: t.split()[0], res.stdout.splitlines()))
 
-    assert env.initial_tenant in tenants
-    assert tenant1 in tenants
-    assert tenant2 in tenants
+    assert env.initial_tenant.hex in tenants
+    assert tenant1.hex in tenants
+    assert tenant2.hex in tenants
+
+
+def test_cli_ipv4_listeners(zenith_env_builder: ZenithEnvBuilder):
+    # Start with single sk
+    zenith_env_builder.num_safekeepers = 1
+    env = zenith_env_builder.init_start()
+
+    # Connect to sk port on v4 loopback
+    res = requests.get(f'http://127.0.0.1:{env.safekeepers[0].port.http}/v1/status')
+    assert res.ok
+
+    # FIXME Test setup is using localhost:xx in ps config.
+    # Perhaps consider switching test suite to v4 loopback.
+
+    # Connect to ps port on v4 loopback
+    # res = requests.get(f'http://127.0.0.1:{env.pageserver.service_port.http}/v1/status')
+    # assert res.ok
+
+
+def test_cli_start_stop(zenith_env_builder: ZenithEnvBuilder):
+    # Start with single sk
+    zenith_env_builder.num_safekeepers = 1
+    env = zenith_env_builder.init_start()
+
+    # Stop default ps/sk
+    env.zenith_cli.pageserver_stop()
+    env.zenith_cli.safekeeper_stop()
+
+    # Default start
+    res = env.zenith_cli.raw_cli(["start"])
+    res.check_returncode()
+
+    # Default stop
+    res = env.zenith_cli.raw_cli(["stop"])
+    res.check_returncode()

test_runner/batch_pg_regress/test_isolation.py

@@ -3,15 +3,11 @@ import os
 from fixtures.utils import mkdir_if_needed
 from fixtures.zenith_fixtures import ZenithEnv, base_dir, pg_distrib_dir
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 def test_isolation(zenith_simple_env: ZenithEnv, test_output_dir, pg_bin, capsys):
     env = zenith_simple_env
 
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_isolation", "empty"])
-
+    env.zenith_cli.create_branch("test_isolation", "empty")
     # Connect to postgres and create a database called "regression".
     # isolation tests use prepared transactions, so enable them
     pg = env.postgres.create_start('test_isolation', config_lines=['max_prepared_transactions=100'])

test_runner/batch_pg_regress/test_pg_regress.py

@@ -3,15 +3,11 @@ import os
 from fixtures.utils import mkdir_if_needed
 from fixtures.zenith_fixtures import ZenithEnv, check_restored_datadir_content, base_dir, pg_distrib_dir
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 def test_pg_regress(zenith_simple_env: ZenithEnv, test_output_dir: str, pg_bin, capsys):
     env = zenith_simple_env
 
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_pg_regress", "empty"])
-
+    env.zenith_cli.create_branch("test_pg_regress", "empty")
     # Connect to postgres and create a database called "regression".
     pg = env.postgres.create_start('test_pg_regress')
     pg.safe_psql('CREATE DATABASE regression')

test_runner/batch_pg_regress/test_zenith_regress.py

@@ -7,15 +7,11 @@ from fixtures.zenith_fixtures import (ZenithEnv,
                                       pg_distrib_dir)
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures")
-
 
 def test_zenith_regress(zenith_simple_env: ZenithEnv, test_output_dir, pg_bin, capsys):
     env = zenith_simple_env
 
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_zenith_regress", "empty"])
-
+    env.zenith_cli.create_branch("test_zenith_regress", "empty")
     # Connect to postgres and create a database called "regression".
     pg = env.postgres.create_start('test_zenith_regress')
     pg.safe_psql('CREATE DATABASE regression')

test_runner/conftest.py

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

test_runner/fixtures/benchmark_fixture.py

@@ -8,6 +8,7 @@ import timeit
 import calendar
 import enum
 from datetime import datetime
+import uuid
 import pytest
 from _pytest.config import Config
 from _pytest.terminal import TerminalReporter
@@ -26,8 +27,6 @@ bencmark, and then record the result by calling zenbenchmark.record. For example
 import timeit
 from fixtures.zenith_fixtures import ZenithEnv
 
-pytest_plugins = ("fixtures.zenith_fixtures", "fixtures.benchmark_fixture")
-
 def test_mybench(zenith_simple_env: env, zenbenchmark):
 
     # Initialize the test
@@ -40,6 +39,8 @@ def test_mybench(zenith_simple_env: env, zenbenchmark):
     # Record another measurement
     zenbenchmark.record('speed_of_light', 300000, 'km/s')
 
+There's no need to import this file to use it. It should be declared as a plugin
+inside conftest.py, and that makes it available to all tests.
 
 You can measure multiple things in one test, and record each one with a separate
 call to zenbenchmark. For example, you could time the bulk loading that happens
@@ -276,11 +277,11 @@ class ZenithBenchmarker:
         assert matches
         return int(round(float(matches.group(1))))
 
-    def get_timeline_size(self, repo_dir: Path, tenantid: str, timelineid: str):
+    def get_timeline_size(self, repo_dir: Path, tenantid: uuid.UUID, timelineid: str):
         """
         Calculate the on-disk size of a timeline
         """
-        path = "{}/tenants/{}/timelines/{}".format(repo_dir, tenantid, timelineid)
+        path = "{}/tenants/{}/timelines/{}".format(repo_dir, tenantid.hex, timelineid)
 
         totalbytes = 0
         for root, dirs, files in os.walk(path):

test_runner/fixtures/compare_fixtures.py

@@ -0,0 +1,200 @@
+import pytest
+from contextlib import contextmanager
+from abc import ABC, abstractmethod
+
+from fixtures.zenith_fixtures import PgBin, PgProtocol, VanillaPostgres, ZenithEnv
+from fixtures.benchmark_fixture import MetricReport, ZenithBenchmarker
+
+# Type-related stuff
+from typing import Iterator
+
+
+class PgCompare(ABC):
+    """Common interface of all postgres implementations, useful for benchmarks.
+
+    This class is a helper class for the zenith_with_baseline fixture. See its documentation
+    for more details.
+    """
+    @property
+    @abstractmethod
+    def pg(self) -> PgProtocol:
+        pass
+
+    @property
+    @abstractmethod
+    def pg_bin(self) -> PgBin:
+        pass
+
+    @property
+    def zenbenchmark(self) -> ZenithBenchmarker:
+        pass
+
+    @abstractmethod
+    def flush(self) -> None:
+        pass
+
+    @abstractmethod
+    def report_peak_memory_use(self) -> None:
+        pass
+
+    @abstractmethod
+    def report_size(self) -> None:
+        pass
+
+    @contextmanager
+    @abstractmethod
+    def record_pageserver_writes(self, out_name):
+        pass
+
+    @contextmanager
+    @abstractmethod
+    def record_duration(self, out_name):
+        pass
+
+
+class ZenithCompare(PgCompare):
+    """PgCompare interface for the zenith stack."""
+    def __init__(self,
+                 zenbenchmark: ZenithBenchmarker,
+                 zenith_simple_env: ZenithEnv,
+                 pg_bin: PgBin,
+                 branch_name):
+        self.env = zenith_simple_env
+        self._zenbenchmark = zenbenchmark
+        self._pg_bin = pg_bin
+
+        # We only use one branch and one timeline
+        self.branch = branch_name
+        self.env.zenith_cli.create_branch(self.branch, "empty")
+        self._pg = self.env.postgres.create_start(self.branch)
+        self.timeline = self.pg.safe_psql("SHOW zenith.zenith_timeline")[0][0]
+
+        # Long-lived cursor, useful for flushing
+        self.psconn = self.env.pageserver.connect()
+        self.pscur = self.psconn.cursor()
+
+    @property
+    def pg(self):
+        return self._pg
+
+    @property
+    def zenbenchmark(self):
+        return self._zenbenchmark
+
+    @property
+    def pg_bin(self):
+        return self._pg_bin
+
+    def flush(self):
+        self.pscur.execute(f"do_gc {self.env.initial_tenant.hex} {self.timeline} 0")
+
+    def report_peak_memory_use(self) -> None:
+        self.zenbenchmark.record("peak_mem",
+                                 self.zenbenchmark.get_peak_mem(self.env.pageserver) / 1024,
+                                 'MB',
+                                 report=MetricReport.LOWER_IS_BETTER)
+
+    def report_size(self) -> None:
+        timeline_size = self.zenbenchmark.get_timeline_size(self.env.repo_dir,
+                                                            self.env.initial_tenant,
+                                                            self.timeline)
+        self.zenbenchmark.record('size',
+                                 timeline_size / (1024 * 1024),
+                                 'MB',
+                                 report=MetricReport.LOWER_IS_BETTER)
+
+    def record_pageserver_writes(self, out_name):
+        return self.zenbenchmark.record_pageserver_writes(self.env.pageserver, out_name)
+
+    def record_duration(self, out_name):
+        return self.zenbenchmark.record_duration(out_name)
+
+
+class VanillaCompare(PgCompare):
+    """PgCompare interface for vanilla postgres."""
+    def __init__(self, zenbenchmark, vanilla_pg: VanillaPostgres):
+        self._pg = vanilla_pg
+        self._zenbenchmark = zenbenchmark
+        vanilla_pg.configure(['shared_buffers=1MB'])
+        vanilla_pg.start()
+
+        # Long-lived cursor, useful for flushing
+        self.conn = self.pg.connect()
+        self.cur = self.conn.cursor()
+
+    @property
+    def pg(self):
+        return self._pg
+
+    @property
+    def zenbenchmark(self):
+        return self._zenbenchmark
+
+    @property
+    def pg_bin(self):
+        return self._pg.pg_bin
+
+    def flush(self):
+        self.cur.execute("checkpoint")
+
+    def report_peak_memory_use(self) -> None:
+        pass  # TODO find something
+
+    def report_size(self) -> None:
+        data_size = self.pg.get_subdir_size('base')
+        self.zenbenchmark.record('data_size',
+                                 data_size / (1024 * 1024),
+                                 'MB',
+                                 report=MetricReport.LOWER_IS_BETTER)
+        wal_size = self.pg.get_subdir_size('pg_wal')
+        self.zenbenchmark.record('wal_size',
+                                 wal_size / (1024 * 1024),
+                                 'MB',
+                                 report=MetricReport.LOWER_IS_BETTER)
+
+    @contextmanager
+    def record_pageserver_writes(self, out_name):
+        yield  # Do nothing
+
+    def record_duration(self, out_name):
+        return self.zenbenchmark.record_duration(out_name)
+
+
+@pytest.fixture(scope='function')
+def zenith_compare(request, zenbenchmark, pg_bin, zenith_simple_env) -> ZenithCompare:
+    branch_name = request.node.name
+    return ZenithCompare(zenbenchmark, zenith_simple_env, pg_bin, branch_name)
+
+
+@pytest.fixture(scope='function')
+def vanilla_compare(zenbenchmark, vanilla_pg) -> VanillaCompare:
+    return VanillaCompare(zenbenchmark, vanilla_pg)
+
+
+@pytest.fixture(params=["vanilla_compare", "zenith_compare"], ids=["vanilla", "zenith"])
+def zenith_with_baseline(request) -> PgCompare:
+    """Parameterized fixture that helps compare zenith against vanilla postgres.
+
+    A test that uses this fixture turns into a parameterized test that runs against:
+    1. A vanilla postgres instance
+    2. A simple zenith env (see zenith_simple_env)
+    3. Possibly other postgres protocol implementations.
+
+    The main goal of this fixture is to make it easier for people to read and write
+    performance tests. Easy test writing leads to more tests.
+
+    Perfect encapsulation of the postgres implementations is **not** a goal because
+    it's impossible. Operational and configuration differences in the different
+    implementations sometimes matter, and the writer of the test should be mindful
+    of that.
+
+    If a test requires some one-off special implementation-specific logic, use of
+    isinstance(zenith_with_baseline, ZenithCompare) is encouraged. Though if that
+    implementation-specific logic is widely useful across multiple tests, it might
+    make sense to add methods to the PgCompare class.
+    """
+    fixture = request.getfixturevalue(request.param)
+    if isinstance(fixture, PgCompare):
+        return fixture
+    else:
+        raise AssertionError(f"test error: fixture {request.param} is not PgCompare")

test_runner/fixtures/slow.py

@@ -0,0 +1,26 @@
+import pytest
+"""
+This plugin allows tests to be marked as slow using pytest.mark.slow. By default slow
+tests are excluded. They need to be specifically requested with the --runslow flag in
+order to run.
+
+Copied from here: https://docs.pytest.org/en/latest/example/simple.html
+"""
+
+
+def pytest_addoption(parser):
+    parser.addoption("--runslow", action="store_true", default=False, help="run slow tests")
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "slow: mark test as slow to run")
+
+
+def pytest_collection_modifyitems(config, items):
+    if config.getoption("--runslow"):
+        # --runslow given in cli: do not skip slow tests
+        return
+    skip_slow = pytest.mark.skip(reason="need --runslow option to run")
+    for item in items:
+        if "slow" in item.keywords:
+            item.add_marker(skip_slow)

test_runner/fixtures/zenith_fixtures.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+import textwrap
 from cached_property import cached_property
 import asyncpg
 import os
@@ -26,7 +27,7 @@ from dataclasses import dataclass
 
 # Type-related stuff
 from psycopg2.extensions import connection as PgConnection
-from typing import Any, Callable, Dict, Iterator, List, Optional, TypeVar, cast, Union
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, TypeVar, cast, Union, Tuple
 from typing_extensions import Literal
 import pytest
 
@@ -44,9 +45,8 @@ the standard pytest.fixture with some extra behavior.
 There are several environment variables that can control the running of tests:
 ZENITH_BIN, POSTGRES_DISTRIB_DIR, etc. See README.md for more information.
 
-To use fixtures in a test file, add this line of code:
-
->>> pytest_plugins = ("fixtures.zenith_fixtures")
+There's no need to import this file to use it. It should be declared as a plugin
+inside conftest.py, and that makes it available to all tests.
 
 Don't import functions from this file, or pytest will emit warnings. Instead
 put directly-importable functions into utils.py or another separate file.
@@ -184,6 +184,16 @@ def worker_base_port(worker_seq_no: int):
     return BASE_PORT + worker_seq_no * WORKER_PORT_NUM
 
 
+def get_dir_size(path: str) -> int:
+    """Return size in bytes."""
+    totalbytes = 0
+    for root, dirs, files in os.walk(path):
+        for name in files:
+            totalbytes += os.path.getsize(os.path.join(root, name))
+
+    return totalbytes
+
+
 def can_bind(host: str, port: int) -> bool:
     """
     Check whether a host:port is available to bind for listening
@@ -230,7 +240,7 @@ class PgProtocol:
     def __init__(self, host: str, port: int, username: Optional[str] = None):
         self.host = host
         self.port = port
-        self.username = username or "zenith_admin"
+        self.username = username
 
     def connstr(self,
                 *,
@@ -242,10 +252,15 @@ class PgProtocol:
         """
 
         username = username or self.username
-        res = f'host={self.host} port={self.port} user={username} dbname={dbname}'
-        if not password:
-            return res
-        return f'{res} password={password}'
+        res = f'host={self.host} port={self.port} dbname={dbname}'
+
+        if username:
+            res = f'{res} user={username}'
+
+        if password:
+            res = f'{res} password={password}'
+
+        return res
 
     # autocommit=True here by default because that's what we need most of the time
     def connect(self,
@@ -410,6 +425,14 @@ class ZenithEnvBuilder:
         self.env = ZenithEnv(self)
         return self.env
 
+    def start(self):
+        self.env.start()
+
+    def init_start(self) -> ZenithEnv:
+        env = self.init()
+        self.start()
+        return env
+
     """
     Sets up the pageserver to use the local fs at the `test_dir/local_fs_remote_storage` path.
     Errors, if the pageserver has some remote storage configuration already, unless `force_enable` is not set to `True`.
@@ -501,6 +524,7 @@ class ZenithEnv:
         self.rust_log_override = config.rust_log_override
         self.port_distributor = config.port_distributor
         self.s3_mock_server = config.s3_mock_server
+        self.zenith_cli = ZenithCli(env=self)
 
         self.postgres = PostgresFactory(self)
 
@@ -508,12 +532,12 @@ class ZenithEnv:
 
         # generate initial tenant ID here instead of letting 'zenith init' generate it,
         # so that we don't need to dig it out of the config file afterwards.
-        self.initial_tenant = uuid.uuid4().hex
+        self.initial_tenant = uuid.uuid4()
 
         # Create a config file corresponding to the options
-        toml = f"""
-default_tenantid = '{self.initial_tenant}'
-        """
+        toml = textwrap.dedent(f"""
+            default_tenantid = '{self.initial_tenant.hex}'
+        """)
 
         # Create config for pageserver
         pageserver_port = PageserverPort(
@@ -522,12 +546,13 @@ default_tenantid = '{self.initial_tenant}'
         )
         pageserver_auth_type = "ZenithJWT" if config.pageserver_auth_enabled else "Trust"
 
-        toml += f"""
-[pageserver]
-listen_pg_addr = 'localhost:{pageserver_port.pg}'
-listen_http_addr = 'localhost:{pageserver_port.http}'
-auth_type = '{pageserver_auth_type}'
-        """
+        toml += textwrap.dedent(f"""
+            [pageserver]
+            id=1
+            listen_pg_addr = 'localhost:{pageserver_port.pg}'
+            listen_http_addr = 'localhost:{pageserver_port.http}'
+            auth_type = '{pageserver_auth_type}'
+        """)
 
         # Create a corresponding ZenithPageserver object
         self.pageserver = ZenithPageserver(self,
@@ -540,33 +565,22 @@ auth_type = '{pageserver_auth_type}'
                 pg=self.port_distributor.get_port(),
                 http=self.port_distributor.get_port(),
             )
-
-            if config.num_safekeepers == 1:
-                name = "single"
-            else:
-                name = f"sk{i}"
-            toml += f"""
-[[safekeepers]]
-name = '{name}'
-pg_port = {port.pg}
-http_port = {port.http}
-sync = false # Disable fsyncs to make the tests go faster
-            """
-            safekeeper = Safekeeper(env=self, name=name, port=port)
+            id = i  # assign ids sequentially
+            toml += textwrap.dedent(f"""
+                [[safekeepers]]
+                id = {id}
+                pg_port = {port.pg}
+                http_port = {port.http}
+                sync = false # Disable fsyncs to make the tests go faster
+            """)
+            safekeeper = Safekeeper(env=self, id=id, port=port)
             self.safekeepers.append(safekeeper)
 
         log.info(f"Config: {toml}")
 
-        # Run 'zenith init' using the config file we constructed
-        with tempfile.NamedTemporaryFile(mode='w+') as tmp:
-            tmp.write(toml)
-            tmp.flush()
-
-            cmd = ['init', f'--config={tmp.name}']
-            append_pageserver_param_overrides(cmd, config.pageserver_remote_storage)
-
-            self.zenith_cli(cmd)
+        self.zenith_cli.init(toml)
 
+    def start(self):
         # Start up the page server and all the safekeepers
         self.pageserver.start()
 
@@ -577,69 +591,12 @@ sync = false # Disable fsyncs to make the tests go faster
         """ Get list of safekeeper endpoints suitable for wal_acceptors GUC  """
         return ','.join([f'localhost:{wa.port.pg}' for wa in self.safekeepers])
 
-    def create_tenant(self, tenant_id: Optional[str] = None):
+    def create_tenant(self, tenant_id: Optional[uuid.UUID] = None) -> uuid.UUID:
         if tenant_id is None:
-            tenant_id = uuid.uuid4().hex
-        res = self.zenith_cli(['tenant', 'create', tenant_id])
-        res.check_returncode()
+            tenant_id = uuid.uuid4()
+        self.zenith_cli.create_tenant(tenant_id)
         return tenant_id
 
-    def zenith_cli(self, arguments: List[str]) -> 'subprocess.CompletedProcess[str]':
-        """
-        Run "zenith" with the specified arguments.
-
-        Arguments must be in list form, e.g. ['pg', 'create']
-
-        Return both stdout and stderr, which can be accessed as
-
-        >>> result = env.zenith_cli(...)
-        >>> assert result.stderr == ""
-        >>> log.info(result.stdout)
-        """
-
-        assert type(arguments) == list
-
-        bin_zenith = os.path.join(str(zenith_binpath), 'zenith')
-
-        args = [bin_zenith] + arguments
-        log.info('Running command "{}"'.format(' '.join(args)))
-        log.info(f'Running in "{self.repo_dir}"')
-
-        env_vars = os.environ.copy()
-        env_vars['ZENITH_REPO_DIR'] = str(self.repo_dir)
-        env_vars['POSTGRES_DISTRIB_DIR'] = str(pg_distrib_dir)
-
-        if self.rust_log_override is not None:
-            env_vars['RUST_LOG'] = self.rust_log_override
-
-        # Pass coverage settings
-        var = 'LLVM_PROFILE_FILE'
-        val = os.environ.get(var)
-        if val:
-            env_vars[var] = val
-
-        # Intercept CalledProcessError and print more info
-        try:
-            res = subprocess.run(args,
-                                 env=env_vars,
-                                 check=True,
-                                 universal_newlines=True,
-                                 stdout=subprocess.PIPE,
-                                 stderr=subprocess.PIPE)
-            log.info(f"Run success: {res.stdout}")
-        except subprocess.CalledProcessError as exc:
-            # this way command output will be in recorded and shown in CI in failure message
-            msg = f"""\
-            Run failed: {exc}
-              stdout: {exc.stdout}
-              stderr: {exc.stderr}
-            """
-            log.info(msg)
-
-            raise Exception(msg) from exc
-
-        return res
-
     @cached_property
     def auth_keys(self) -> AuthKeys:
         pub = (Path(self.repo_dir) / 'auth_public_key.pem').read_bytes()
@@ -664,10 +621,10 @@ def _shared_simple_env(request: Any, port_distributor) -> Iterator[ZenithEnv]:
 
     with ZenithEnvBuilder(Path(repo_dir), port_distributor) as builder:
 
-        env = builder.init()
+        env = builder.init_start()
 
         # For convenience in tests, create a branch from the freshly-initialized cluster.
-        env.zenith_cli(["branch", "empty", "main"])
+        env.zenith_cli.create_branch("empty", "main")
 
         # Return the builder to the caller
         yield env
@@ -698,7 +655,7 @@ def zenith_env_builder(test_output_dir, port_distributor) -> Iterator[ZenithEnvB
     To use, define 'zenith_env_builder' fixture in your test to get access to the
     builder object. Set properties on it to describe the environment.
     Finally, initialize and start up the environment by calling
-    zenith_env_builder.init().
+    zenith_env_builder.init_start().
 
     After the initialization, you can launch compute nodes by calling
     the functions in the 'env.postgres' factory object, stop/start the
@@ -713,6 +670,10 @@ def zenith_env_builder(test_output_dir, port_distributor) -> Iterator[ZenithEnvB
         yield builder
 
 
+class ZenithPageserverApiException(Exception):
+    pass
+
+
 class ZenithPageserverHttpClient(requests.Session):
     def __init__(self, port: int, auth_token: Optional[str] = None) -> None:
         super().__init__()
@@ -722,22 +683,32 @@ class ZenithPageserverHttpClient(requests.Session):
         if auth_token is not None:
             self.headers['Authorization'] = f'Bearer {auth_token}'
 
+    def verbose_error(self, res: requests.Response):
+        try:
+            res.raise_for_status()
+        except requests.RequestException as e:
+            try:
+                msg = res.json()['msg']
+            except:
+                msg = ''
+            raise ZenithPageserverApiException(msg) from e
+
     def check_status(self):
         self.get(f"http://localhost:{self.port}/v1/status").raise_for_status()
 
     def timeline_attach(self, tenant_id: uuid.UUID, timeline_id: uuid.UUID):
         res = self.post(
             f"http://localhost:{self.port}/v1/timeline/{tenant_id.hex}/{timeline_id.hex}/attach", )
-        res.raise_for_status()
+        self.verbose_error(res)
 
     def timeline_detach(self, tenant_id: uuid.UUID, timeline_id: uuid.UUID):
         res = self.post(
             f"http://localhost:{self.port}/v1/timeline/{tenant_id.hex}/{timeline_id.hex}/detach", )
-        res.raise_for_status()
+        self.verbose_error(res)
 
     def branch_list(self, tenant_id: uuid.UUID) -> List[Dict[Any, Any]]:
         res = self.get(f"http://localhost:{self.port}/v1/branch/{tenant_id.hex}")
-        res.raise_for_status()
+        self.verbose_error(res)
         res_json = res.json()
         assert isinstance(res_json, list)
         return res_json
@@ -749,7 +720,7 @@ class ZenithPageserverHttpClient(requests.Session):
                             'name': name,
                             'start_point': start_point,
                         })
-        res.raise_for_status()
+        self.verbose_error(res)
         res_json = res.json()
         assert isinstance(res_json, dict)
         return res_json
@@ -758,14 +729,14 @@ class ZenithPageserverHttpClient(requests.Session):
         res = self.get(
             f"http://localhost:{self.port}/v1/branch/{tenant_id.hex}/{name}?include-non-incremental-logical-size=1",
         )
-        res.raise_for_status()
+        self.verbose_error(res)
         res_json = res.json()
         assert isinstance(res_json, dict)
         return res_json
 
     def tenant_list(self) -> List[Dict[Any, Any]]:
         res = self.get(f"http://localhost:{self.port}/v1/tenant")
-        res.raise_for_status()
+        self.verbose_error(res)
         res_json = res.json()
         assert isinstance(res_json, list)
         return res_json
@@ -777,27 +748,27 @@ class ZenithPageserverHttpClient(requests.Session):
                 'tenant_id': tenant_id.hex,
             },
         )
-        res.raise_for_status()
+        self.verbose_error(res)
         return res.json()
 
     def timeline_list(self, tenant_id: uuid.UUID) -> List[str]:
         res = self.get(f"http://localhost:{self.port}/v1/timeline/{tenant_id.hex}")
-        res.raise_for_status()
+        self.verbose_error(res)
         res_json = res.json()
         assert isinstance(res_json, list)
         return res_json
 
-    def timeline_detail(self, tenant_id: uuid.UUID, timeline_id: uuid.UUID):
+    def timeline_detail(self, tenant_id: uuid.UUID, timeline_id: uuid.UUID) -> Dict[Any, Any]:
         res = self.get(
             f"http://localhost:{self.port}/v1/timeline/{tenant_id.hex}/{timeline_id.hex}")
-        res.raise_for_status()
+        self.verbose_error(res)
         res_json = res.json()
         assert isinstance(res_json, dict)
         return res_json
 
     def get_metrics(self) -> str:
         res = self.get(f"http://localhost:{self.port}/metrics")
-        res.raise_for_status()
+        self.verbose_error(res)
         return res.text
 
 
@@ -824,6 +795,190 @@ class S3Storage:
 RemoteStorage = Union[LocalFsStorage, S3Storage]
 
 
+class ZenithCli:
+    """
+    A typed wrapper around the `zenith` CLI tool.
+    Supports main commands via typed methods and a way to run arbitrary command directly via CLI.
+    """
+    def __init__(self, env: ZenithEnv) -> None:
+        self.env = env
+        pass
+
+    def create_tenant(self, tenant_id: Optional[uuid.UUID] = None) -> uuid.UUID:
+        if tenant_id is None:
+            tenant_id = uuid.uuid4()
+        self.raw_cli(['tenant', 'create', tenant_id.hex])
+        return tenant_id
+
+    def list_tenants(self) -> 'subprocess.CompletedProcess[str]':
+        return self.raw_cli(['tenant', 'list'])
+
+    def create_branch(self,
+                      branch_name: str,
+                      starting_point: str,
+                      tenant_id: Optional[uuid.UUID] = None) -> 'subprocess.CompletedProcess[str]':
+        args = ['branch']
+        if tenant_id is not None:
+            args.extend(['--tenantid', tenant_id.hex])
+        args.extend([branch_name, starting_point])
+
+        return self.raw_cli(args)
+
+    def list_branches(self,
+                      tenant_id: Optional[uuid.UUID] = None) -> 'subprocess.CompletedProcess[str]':
+        args = ['branch']
+        if tenant_id is not None:
+            args.extend(['--tenantid', tenant_id.hex])
+        return self.raw_cli(args)
+
+    def init(self, config_toml: str) -> 'subprocess.CompletedProcess[str]':
+        with tempfile.NamedTemporaryFile(mode='w+') as tmp:
+            tmp.write(config_toml)
+            tmp.flush()
+
+            cmd = ['init', f'--config={tmp.name}']
+            append_pageserver_param_overrides(cmd, self.env.pageserver.remote_storage)
+
+            return self.raw_cli(cmd)
+
+    def pageserver_start(self, overrides=()) -> 'subprocess.CompletedProcess[str]':
+        start_args = ['pageserver', 'start', *overrides]
+
+        append_pageserver_param_overrides(start_args, self.env.pageserver.remote_storage)
+        return self.raw_cli(start_args)
+
+    def pageserver_stop(self, immediate=False) -> 'subprocess.CompletedProcess[str]':
+        cmd = ['pageserver', 'stop']
+        if immediate:
+            cmd.extend(['-m', 'immediate'])
+
+        log.info(f"Stopping pageserver with {cmd}")
+        return self.raw_cli(cmd)
+
+    def safekeeper_start(self, id: int) -> 'subprocess.CompletedProcess[str]':
+        return self.raw_cli(['safekeeper', 'start', str(id)])
+
+    def safekeeper_stop(self,
+                        id: Optional[int] = None,
+                        immediate=False) -> 'subprocess.CompletedProcess[str]':
+        args = ['safekeeper', 'stop']
+        if id is not None:
+            args.extend(str(id))
+        if immediate:
+            args.extend(['-m', 'immediate'])
+        return self.raw_cli(args)
+
+    def pg_create(
+        self,
+        node_name: str,
+        tenant_id: Optional[uuid.UUID] = None,
+        timeline_spec: Optional[str] = None,
+        port: Optional[int] = None,
+    ) -> 'subprocess.CompletedProcess[str]':
+        args = ['pg', 'create']
+        if tenant_id is not None:
+            args.extend(['--tenantid', tenant_id.hex])
+        if port is not None:
+            args.append(f'--port={port}')
+        args.append(node_name)
+        if timeline_spec is not None:
+            args.append(timeline_spec)
+        return self.raw_cli(args)
+
+    def pg_start(
+        self,
+        node_name: str,
+        tenant_id: Optional[uuid.UUID] = None,
+        timeline_spec: Optional[str] = None,
+        port: Optional[int] = None,
+    ) -> 'subprocess.CompletedProcess[str]':
+        args = ['pg', 'start']
+        if tenant_id is not None:
+            args.extend(['--tenantid', tenant_id.hex])
+        if port is not None:
+            args.append(f'--port={port}')
+        args.append(node_name)
+        if timeline_spec is not None:
+            args.append(timeline_spec)
+
+        return self.raw_cli(args)
+
+    def pg_stop(
+        self,
+        node_name: str,
+        tenant_id: Optional[uuid.UUID] = None,
+        destroy=False,
+    ) -> 'subprocess.CompletedProcess[str]':
+        args = ['pg', 'stop']
+        if tenant_id is not None:
+            args.extend(['--tenantid', tenant_id.hex])
+        if destroy:
+            args.append('--destroy')
+        args.append(node_name)
+
+        return self.raw_cli(args)
+
+    def raw_cli(self,
+                arguments: List[str],
+                check_return_code=True) -> 'subprocess.CompletedProcess[str]':
+        """
+        Run "zenith" with the specified arguments.
+
+        Arguments must be in list form, e.g. ['pg', 'create']
+
+        Return both stdout and stderr, which can be accessed as
+
+        >>> result = env.zenith_cli.raw_cli(...)
+        >>> assert result.stderr == ""
+        >>> log.info(result.stdout)
+        """
+
+        assert type(arguments) == list
+
+        bin_zenith = os.path.join(str(zenith_binpath), 'zenith')
+
+        args = [bin_zenith] + arguments
+        log.info('Running command "{}"'.format(' '.join(args)))
+        log.info(f'Running in "{self.env.repo_dir}"')
+
+        env_vars = os.environ.copy()
+        env_vars['ZENITH_REPO_DIR'] = str(self.env.repo_dir)
+        env_vars['POSTGRES_DISTRIB_DIR'] = str(pg_distrib_dir)
+
+        if self.env.rust_log_override is not None:
+            env_vars['RUST_LOG'] = self.env.rust_log_override
+
+        # Pass coverage settings
+        var = 'LLVM_PROFILE_FILE'
+        val = os.environ.get(var)
+        if val:
+            env_vars[var] = val
+
+        # Intercept CalledProcessError and print more info
+        try:
+            res = subprocess.run(args,
+                                 env=env_vars,
+                                 check=True,
+                                 universal_newlines=True,
+                                 stdout=subprocess.PIPE,
+                                 stderr=subprocess.PIPE)
+            log.info(f"Run success: {res.stdout}")
+        except subprocess.CalledProcessError as exc:
+            # this way command output will be in recorded and shown in CI in failure message
+            msg = f"""\
+            Run failed: {exc}
+              stdout: {exc.stdout}
+              stderr: {exc.stderr}
+            """
+            log.info(msg)
+
+            raise Exception(msg) from exc
+
+        if check_return_code:
+            res.check_returncode()
+        return res
+
+
 class ZenithPageserver(PgProtocol):
     """
     An object representing a running pageserver.
@@ -835,23 +990,20 @@ class ZenithPageserver(PgProtocol):
                  port: PageserverPort,
                  remote_storage: Optional[RemoteStorage] = None,
                  enable_auth=False):
-        super().__init__(host='localhost', port=port.pg)
+        super().__init__(host='localhost', port=port.pg, username='zenith_admin')
         self.env = env
         self.running = False
         self.service_port = port  # do not shadow PgProtocol.port which is just int
         self.remote_storage = remote_storage
 
-    def start(self) -> 'ZenithPageserver':
+    def start(self, overrides=()) -> 'ZenithPageserver':
         """
         Start the page server.
         Returns self.
         """
         assert self.running == False
 
-        start_args = ['pageserver', 'start']
-        append_pageserver_param_overrides(start_args, self.remote_storage)
-
-        self.env.zenith_cli(start_args)
+        self.env.zenith_cli.pageserver_start(overrides=overrides)
         self.running = True
         return self
 
@@ -860,13 +1012,8 @@ class ZenithPageserver(PgProtocol):
         Stop the page server.
         Returns self.
         """
-        cmd = ['pageserver', 'stop']
-        if immediate:
-            cmd.extend(['-m', 'immediate'])
-
-        log.info(f"Stopping pageserver with {cmd}")
         if self.running:
-            self.env.zenith_cli(cmd)
+            self.env.zenith_cli.pageserver_stop(immediate)
             self.running = False
 
         return self
@@ -973,10 +1120,54 @@ def pg_bin(test_output_dir: str) -> PgBin:
     return PgBin(test_output_dir)
 
 
+class VanillaPostgres(PgProtocol):
+    def __init__(self, pgdatadir: str, pg_bin: PgBin, port: int):
+        super().__init__(host='localhost', port=port)
+        self.pgdatadir = pgdatadir
+        self.pg_bin = pg_bin
+        self.running = False
+        self.pg_bin.run_capture(['initdb', '-D', pgdatadir])
+
+    def configure(self, options: List[str]) -> None:
+        """Append lines into postgresql.conf file."""
+        assert not self.running
+        with open(os.path.join(self.pgdatadir, 'postgresql.conf'), 'a') as conf_file:
+            conf_file.writelines(options)
+
+    def start(self) -> None:
+        assert not self.running
+        self.running = True
+        self.pg_bin.run_capture(['pg_ctl', '-D', self.pgdatadir, 'start'])
+
+    def stop(self) -> None:
+        assert self.running
+        self.running = False
+        self.pg_bin.run_capture(['pg_ctl', '-D', self.pgdatadir, 'stop'])
+
+    def get_subdir_size(self, subdir) -> int:
+        """Return size of pgdatadir subdirectory in bytes."""
+        return get_dir_size(os.path.join(self.pgdatadir, subdir))
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        if self.running:
+            self.stop()
+
+
+@pytest.fixture(scope='function')
+def vanilla_pg(test_output_dir: str) -> Iterator[VanillaPostgres]:
+    pgdatadir = os.path.join(test_output_dir, "pgdata-vanilla")
+    pg_bin = PgBin(test_output_dir)
+    with VanillaPostgres(pgdatadir, pg_bin, 5432) as vanilla_pg:
+        yield vanilla_pg
+
+
 class Postgres(PgProtocol):
     """ An object representing a running postgres daemon. """
-    def __init__(self, env: ZenithEnv, tenant_id: str, port: int):
-        super().__init__(host='localhost', port=port)
+    def __init__(self, env: ZenithEnv, tenant_id: uuid.UUID, port: int):
+        super().__init__(host='localhost', port=port, username='zenith_admin')
 
         self.env = env
         self.running = False
@@ -1002,16 +1193,12 @@ class Postgres(PgProtocol):
         if branch is None:
             branch = node_name
 
-        self.env.zenith_cli([
-            'pg',
-            'create',
-            f'--tenantid={self.tenant_id}',
-            f'--port={self.port}',
-            node_name,
-            branch
-        ])
+        self.env.zenith_cli.pg_create(node_name,
+                                      tenant_id=self.tenant_id,
+                                      port=self.port,
+                                      timeline_spec=branch)
         self.node_name = node_name
-        path = pathlib.Path('pgdatadirs') / 'tenants' / self.tenant_id / self.node_name
+        path = pathlib.Path('pgdatadirs') / 'tenants' / self.tenant_id.hex / self.node_name
         self.pgdata_dir = os.path.join(self.env.repo_dir, path)
 
         if config_lines is None:
@@ -1030,8 +1217,9 @@ class Postgres(PgProtocol):
 
         log.info(f"Starting postgres node {self.node_name}")
 
-        run_result = self.env.zenith_cli(
-            ['pg', 'start', f'--tenantid={self.tenant_id}', f'--port={self.port}', self.node_name])
+        run_result = self.env.zenith_cli.pg_start(self.node_name,
+                                                  tenant_id=self.tenant_id,
+                                                  port=self.port)
         self.running = True
 
         log.info(f"stdout: {run_result.stdout}")
@@ -1041,7 +1229,7 @@ class Postgres(PgProtocol):
     def pg_data_dir_path(self) -> str:
         """ Path to data directory """
         assert self.node_name
-        path = pathlib.Path('pgdatadirs') / 'tenants' / self.tenant_id / self.node_name
+        path = pathlib.Path('pgdatadirs') / 'tenants' / self.tenant_id.hex / self.node_name
         return os.path.join(self.env.repo_dir, path)
 
     def pg_xact_dir_path(self) -> str:
@@ -1101,7 +1289,7 @@ class Postgres(PgProtocol):
 
         if self.running:
             assert self.node_name is not None
-            self.env.zenith_cli(['pg', 'stop', self.node_name, f'--tenantid={self.tenant_id}'])
+            self.env.zenith_cli.pg_stop(self.node_name, tenant_id=self.tenant_id)
             self.running = False
 
         return self
@@ -1113,8 +1301,7 @@ class Postgres(PgProtocol):
         """
 
         assert self.node_name is not None
-        self.env.zenith_cli(
-            ['pg', 'stop', '--destroy', self.node_name, f'--tenantid={self.tenant_id}'])
+        self.env.zenith_cli.pg_stop(self.node_name, self.tenant_id, destroy=True)
         self.node_name = None
 
         return self
@@ -1156,7 +1343,7 @@ class PostgresFactory:
     def create_start(self,
                      node_name: str = "main",
                      branch: Optional[str] = None,
-                     tenant_id: Optional[str] = None,
+                     tenant_id: Optional[uuid.UUID] = None,
                      config_lines: Optional[List[str]] = None) -> Postgres:
 
         pg = Postgres(
@@ -1176,7 +1363,7 @@ class PostgresFactory:
     def create(self,
                node_name: str = "main",
                branch: Optional[str] = None,
-               tenant_id: Optional[str] = None,
+               tenant_id: Optional[uuid.UUID] = None,
                config_lines: Optional[List[str]] = None) -> Postgres:
 
         pg = Postgres(
@@ -1217,12 +1404,14 @@ class Safekeeper:
     """ An object representing a running safekeeper daemon. """
     env: ZenithEnv
     port: SafekeeperPort
-    name: str  # identifier for logging
+    id: int
     auth_token: Optional[str] = None
+    running: bool = False
 
     def start(self) -> 'Safekeeper':
-        self.env.zenith_cli(['safekeeper', 'start', self.name])
-
+        assert self.running == False
+        self.env.zenith_cli.safekeeper_start(self.id)
+        self.running = True
         # wait for wal acceptor start by checking its status
         started_at = time.time()
         while True:
@@ -1240,16 +1429,14 @@ class Safekeeper:
         return self
 
     def stop(self, immediate=False) -> 'Safekeeper':
-        cmd = ['safekeeper', 'stop']
-        if immediate:
-            cmd.extend(['-m', 'immediate'])
-        cmd.append(self.name)
-
-        log.info('Stopping safekeeper {}'.format(self.name))
-        self.env.zenith_cli(cmd)
+        log.info('Stopping safekeeper {}'.format(self.id))
+        self.env.zenith_cli.safekeeper_stop(self.id, immediate)
+        self.running = False
         return self
 
-    def append_logical_message(self, tenant_id: str, timeline_id: str,
+    def append_logical_message(self,
+                               tenant_id: uuid.UUID,
+                               timeline_id: uuid.UUID,
                                request: Dict[str, Any]) -> Dict[str, Any]:
         """
         Send JSON_CTRL query to append LogicalMessage to WAL and modify
@@ -1259,7 +1446,7 @@ class Safekeeper:
 
         # "replication=0" hacks psycopg not to send additional queries
         # on startup, see https://github.com/psycopg/psycopg2/pull/482
-        connstr = f"host=localhost port={self.port.pg} replication=0 options='-c ztimelineid={timeline_id} ztenantid={tenant_id}'"
+        connstr = f"host=localhost port={self.port.pg} replication=0 options='-c ztimelineid={timeline_id.hex} ztenantid={tenant_id.hex}'"
 
         with closing(psycopg2.connect(connstr)) as conn:
             # server doesn't support transactions
@@ -1288,8 +1475,8 @@ class SafekeeperTimelineStatus:
 class SafekeeperMetrics:
     # These are metrics from Prometheus which uses float64 internally.
     # As a consequence, values may differ from real original int64s.
-    flush_lsn_inexact: Dict[str, int] = field(default_factory=dict)
-    commit_lsn_inexact: Dict[str, int] = field(default_factory=dict)
+    flush_lsn_inexact: Dict[Tuple[str, str], int] = field(default_factory=dict)
+    commit_lsn_inexact: Dict[Tuple[str, str], int] = field(default_factory=dict)
 
 
 class SafekeeperHttpClient(requests.Session):
@@ -1313,14 +1500,16 @@ class SafekeeperHttpClient(requests.Session):
         all_metrics_text = request_result.text
 
         metrics = SafekeeperMetrics()
-        for match in re.finditer(r'^safekeeper_flush_lsn{ztli="([0-9a-f]+)"} (\S+)$',
-                                 all_metrics_text,
-                                 re.MULTILINE):
-            metrics.flush_lsn_inexact[match.group(1)] = int(match.group(2))
-        for match in re.finditer(r'^safekeeper_commit_lsn{ztli="([0-9a-f]+)"} (\S+)$',
-                                 all_metrics_text,
-                                 re.MULTILINE):
-            metrics.commit_lsn_inexact[match.group(1)] = int(match.group(2))
+        for match in re.finditer(
+                r'^safekeeper_flush_lsn{tenant_id="([0-9a-f]+)",timeline_id="([0-9a-f]+)"} (\S+)$',
+                all_metrics_text,
+                re.MULTILINE):
+            metrics.flush_lsn_inexact[(match.group(1), match.group(2))] = int(match.group(3))
+        for match in re.finditer(
+                r'^safekeeper_commit_lsn{tenant_id="([0-9a-f]+)",timeline_id="([0-9a-f]+)"} (\S+)$',
+                all_metrics_text,
+                re.MULTILINE):
+            metrics.commit_lsn_inexact[(match.group(1), match.group(2))] = int(match.group(3))
         return metrics
 
 
@@ -1429,7 +1618,7 @@ def check_restored_datadir_content(test_output_dir: str, env: ZenithEnv, pg: Pos
         {psql_path}                                    \
             --no-psqlrc                                \
             postgres://localhost:{env.pageserver.service_port.pg}  \
-            -c 'basebackup {pg.tenant_id} {timeline}'  \
+            -c 'basebackup {pg.tenant_id.hex} {timeline}'  \
          | tar -x -C {restored_dir_path}
     """
 

test_runner/performance/test_bulk_insert.py

@@ -2,8 +2,7 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 from fixtures.benchmark_fixture import MetricReport, ZenithBenchmarker
-
-pytest_plugins = ("fixtures.zenith_fixtures", "fixtures.benchmark_fixture")
+from fixtures.compare_fixtures import PgCompare, VanillaCompare, ZenithCompare
 
 
 #
@@ -16,47 +15,19 @@ pytest_plugins = ("fixtures.zenith_fixtures", "fixtures.benchmark_fixture")
 # 3. Disk space used
 # 4. Peak memory usage
 #
-def test_bulk_insert(zenith_simple_env: ZenithEnv, zenbenchmark: ZenithBenchmarker):
-    env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_bulk_insert", "empty"])
-
-    pg = env.postgres.create_start('test_bulk_insert')
-    log.info("postgres is running on 'test_bulk_insert' branch")
-
-    # Open a connection directly to the page server that we'll use to force
-    # flushing the layers to disk
-    psconn = env.pageserver.connect()
-    pscur = psconn.cursor()
+def test_bulk_insert(zenith_with_baseline: PgCompare):
+    env = zenith_with_baseline
 
     # Get the timeline ID of our branch. We need it for the 'do_gc' command
-    with closing(pg.connect()) as conn:
+    with closing(env.pg.connect()) as conn:
         with conn.cursor() as cur:
-            cur.execute("SHOW zenith.zenith_timeline")
-            timeline = cur.fetchone()[0]
-
             cur.execute("create table huge (i int, j int);")
 
             # Run INSERT, recording the time and I/O it takes
-            with zenbenchmark.record_pageserver_writes(env.pageserver, 'pageserver_writes'):
-                with zenbenchmark.record_duration('insert'):
+            with env.record_pageserver_writes('pageserver_writes'):
+                with env.record_duration('insert'):
                     cur.execute("insert into huge values (generate_series(1, 5000000), 0);")
+                    env.flush()
 
-                    # Flush the layers from memory to disk. This is included in the reported
-                    # time and I/O
-                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
-
-            # Record peak memory usage
-            zenbenchmark.record("peak_mem",
-                                zenbenchmark.get_peak_mem(env.pageserver) / 1024,
-                                'MB',
-                                report=MetricReport.LOWER_IS_BETTER)
-
-            # Report disk space used by the repository
-            timeline_size = zenbenchmark.get_timeline_size(env.repo_dir,
-                                                           env.initial_tenant,
-                                                           timeline)
-            zenbenchmark.record('size',
-                                timeline_size / (1024 * 1024),
-                                'MB',
-                                report=MetricReport.LOWER_IS_BETTER)
+            env.report_peak_memory_use()
+            env.report_size()

test_runner/performance/test_bulk_tenant_create.py

@@ -4,8 +4,6 @@ import pytest
 
 from fixtures.zenith_fixtures import ZenithEnvBuilder
 
-pytest_plugins = ("fixtures.benchmark_fixture")
-
 # Run bulk tenant creation test.
 #
 # Collects metrics:
@@ -25,7 +23,7 @@ def test_bulk_tenant_create(
     """Measure tenant creation time (with and without wal acceptors)"""
     if use_wal_acceptors == 'with_wa':
         zenith_env_builder.num_safekeepers = 3
-    env = zenith_env_builder.init()
+    env = zenith_env_builder.init_start()
 
     time_slices = []
 
@@ -33,12 +31,10 @@ def test_bulk_tenant_create(
         start = timeit.default_timer()
 
         tenant = env.create_tenant()
-        env.zenith_cli([
-            "branch",
+        env.zenith_cli.create_branch(
             f"test_bulk_tenant_create_{tenants_count}_{i}_{use_wal_acceptors}",
             "main",
-            f"--tenantid={tenant}"
-        ])
+            tenant_id=tenant)
 
         # FIXME: We used to start new safekeepers here. Did that make sense? Should we do it now?
         #if use_wal_acceptors == 'with_wa':

test_runner/performance/test_copy.py

@@ -2,11 +2,10 @@ from contextlib import closing
 from fixtures.zenith_fixtures import ZenithEnv
 from fixtures.log_helper import log
 from fixtures.benchmark_fixture import MetricReport, ZenithBenchmarker
+from fixtures.compare_fixtures import PgCompare, VanillaCompare, ZenithCompare
 from io import BufferedReader, RawIOBase
 from itertools import repeat
 
-pytest_plugins = ("fixtures.zenith_fixtures", "fixtures.benchmark_fixture")
-
 
 class CopyTestData(RawIOBase):
     def __init__(self, rows: int):
@@ -42,77 +41,41 @@ def copy_test_data(rows: int):
 #
 # COPY performance tests.
 #
-def test_copy(zenith_simple_env: ZenithEnv, zenbenchmark: ZenithBenchmarker):
-    env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_copy", "empty"])
-
-    pg = env.postgres.create_start('test_copy')
-    log.info("postgres is running on 'test_copy' branch")
-
-    # Open a connection directly to the page server that we'll use to force
-    # flushing the layers to disk
-    psconn = env.pageserver.connect()
-    pscur = psconn.cursor()
+def test_copy(zenith_with_baseline: PgCompare):
+    env = zenith_with_baseline
 
     # Get the timeline ID of our branch. We need it for the pageserver 'checkpoint' command
-    with closing(pg.connect()) as conn:
+    with closing(env.pg.connect()) as conn:
         with conn.cursor() as cur:
-            cur.execute("SHOW zenith.zenith_timeline")
-            timeline = cur.fetchone()[0]
-
             cur.execute("create table copytest (i int, t text);")
 
             # Load data with COPY, recording the time and I/O it takes.
             #
             # Since there's no data in the table previously, this extends it.
-            with zenbenchmark.record_pageserver_writes(env.pageserver,
-                                                       'copy_extend_pageserver_writes'):
-                with zenbenchmark.record_duration('copy_extend'):
+            with env.record_pageserver_writes('copy_extend_pageserver_writes'):
+                with env.record_duration('copy_extend'):
                     cur.copy_from(copy_test_data(1000000), 'copytest')
-                    # Flush the layers from memory to disk. This is included in the reported
-                    # time and I/O
-                    pscur.execute(f"checkpoint {env.initial_tenant} {timeline}")
+                    env.flush()
 
             # Delete most rows, and VACUUM to make the space available for reuse.
-            with zenbenchmark.record_pageserver_writes(env.pageserver, 'delete_pageserver_writes'):
-                with zenbenchmark.record_duration('delete'):
+            with env.record_pageserver_writes('delete_pageserver_writes'):
+                with env.record_duration('delete'):
                     cur.execute("delete from copytest where i % 100 <> 0;")
-                    # Flush the layers from memory to disk. This is included in the reported
-                    # time and I/O
-                    pscur.execute(f"checkpoint {env.initial_tenant} {timeline}")
+                    env.flush()
 
-            with zenbenchmark.record_pageserver_writes(env.pageserver, 'vacuum_pageserver_writes'):
-                with zenbenchmark.record_duration('vacuum'):
+            with env.record_pageserver_writes('vacuum_pageserver_writes'):
+                with env.record_duration('vacuum'):
                     cur.execute("vacuum copytest")
-                    # Flush the layers from memory to disk. This is included in the reported
-                    # time and I/O
-                    pscur.execute(f"checkpoint {env.initial_tenant} {timeline}")
+                    env.flush()
 
             # Load data into the table again. This time, this will use the space free'd
             # by the VACUUM.
             #
             # This will also clear all the VM bits.
-            with zenbenchmark.record_pageserver_writes(env.pageserver,
-                                                       'copy_reuse_pageserver_writes'):
-                with zenbenchmark.record_duration('copy_reuse'):
+            with env.record_pageserver_writes('copy_reuse_pageserver_writes'):
+                with env.record_duration('copy_reuse'):
                     cur.copy_from(copy_test_data(1000000), 'copytest')
+                    env.flush()
 
-                    # Flush the layers from memory to disk. This is included in the reported
-                    # time and I/O
-                    pscur.execute(f"checkpoint {env.initial_tenant} {timeline}")
-
-            # Record peak memory usage
-            zenbenchmark.record("peak_mem",
-                                zenbenchmark.get_peak_mem(env.pageserver) / 1024,
-                                'MB',
-                                report=MetricReport.LOWER_IS_BETTER)
-
-            # Report disk space used by the repository
-            timeline_size = zenbenchmark.get_timeline_size(env.repo_dir,
-                                                           env.initial_tenant,
-                                                           timeline)
-            zenbenchmark.record('size',
-                                timeline_size / (1024 * 1024),
-                                'MB',
-                                report=MetricReport.LOWER_IS_BETTER)
+            env.report_peak_memory_use()
+            env.report_size()

test_runner/performance/test_gist_build.py

@@ -2,34 +2,20 @@ import os
 from contextlib import closing
 from fixtures.benchmark_fixture import MetricReport
 from fixtures.zenith_fixtures import ZenithEnv
+from fixtures.compare_fixtures import PgCompare, VanillaCompare, ZenithCompare
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures", "fixtures.benchmark_fixture")
-
 
 #
 # Test buffering GisT build. It WAL-logs the whole relation, in 32-page chunks.
 # As of this writing, we're duplicate those giant WAL records for each page,
 # which makes the delta layer about 32x larger than it needs to be.
 #
-def test_gist_buffering_build(zenith_simple_env: ZenithEnv, zenbenchmark):
-    env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_gist_buffering_build", "empty"])
+def test_gist_buffering_build(zenith_with_baseline: PgCompare):
+    env = zenith_with_baseline
 
-    pg = env.postgres.create_start('test_gist_buffering_build')
-    log.info("postgres is running on 'test_gist_buffering_build' branch")
-
-    # Open a connection directly to the page server that we'll use to force
-    # flushing the layers to disk
-    psconn = env.pageserver.connect()
-    pscur = psconn.cursor()
-
-    # Get the timeline ID of our branch. We need it for the 'do_gc' command
-    with closing(pg.connect()) as conn:
+    with closing(env.pg.connect()) as conn:
         with conn.cursor() as cur:
-            cur.execute("SHOW zenith.zenith_timeline")
-            timeline = cur.fetchone()[0]
 
             # Create test table.
             cur.execute("create table gist_point_tbl(id int4, p point)")
@@ -38,27 +24,12 @@ def test_gist_buffering_build(zenith_simple_env: ZenithEnv, zenbenchmark):
             )
 
             # Build the index.
-            with zenbenchmark.record_pageserver_writes(env.pageserver, 'pageserver_writes'):
-                with zenbenchmark.record_duration('build'):
+            with env.record_pageserver_writes('pageserver_writes'):
+                with env.record_duration('build'):
                     cur.execute(
                         "create index gist_pointidx2 on gist_point_tbl using gist(p) with (buffering = on)"
                     )
+                    env.flush()
 
-                    # Flush the layers from memory to disk. This is included in the reported
-                    # time and I/O
-                    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 1000000")
-
-            # Record peak memory usage
-            zenbenchmark.record("peak_mem",
-                                zenbenchmark.get_peak_mem(env.pageserver) / 1024,
-                                'MB',
-                                report=MetricReport.LOWER_IS_BETTER)
-
-            # Report disk space used by the repository
-            timeline_size = zenbenchmark.get_timeline_size(env.repo_dir,
-                                                           env.initial_tenant,
-                                                           timeline)
-            zenbenchmark.record('size',
-                                timeline_size / (1024 * 1024),
-                                'MB',
-                                report=MetricReport.LOWER_IS_BETTER)
+            env.report_peak_memory_use()
+            env.report_size()

test_runner/performance/test_parallel_copy_to.py

@@ -1,11 +1,10 @@
 from io import BytesIO
 import asyncio
 import asyncpg
-from fixtures.zenith_fixtures import ZenithEnv, Postgres
+from fixtures.zenith_fixtures import ZenithEnv, Postgres, PgProtocol
 from fixtures.log_helper import log
 from fixtures.benchmark_fixture import MetricReport, ZenithBenchmarker
-
-pytest_plugins = ("fixtures.zenith_fixtures", "fixtures.benchmark_fixture")
+from fixtures.compare_fixtures import PgCompare, VanillaCompare, ZenithCompare
 
 
 async def repeat_bytes(buf, repetitions: int):
@@ -13,7 +12,7 @@ async def repeat_bytes(buf, repetitions: int):
         yield buf
 
 
-async def copy_test_data_to_table(pg: Postgres, worker_id: int, table_name: str):
+async def copy_test_data_to_table(pg: PgProtocol, worker_id: int, table_name: str):
     buf = BytesIO()
     for i in range(1000):
         buf.write(
@@ -26,7 +25,7 @@ async def copy_test_data_to_table(pg: Postgres, worker_id: int, table_name: str)
     await pg_conn.copy_to_table(table_name, source=copy_input)
 
 
-async def parallel_load_different_tables(pg: Postgres, n_parallel: int):
+async def parallel_load_different_tables(pg: PgProtocol, n_parallel: int):
     workers = []
     for worker_id in range(n_parallel):
         worker = copy_test_data_to_table(pg, worker_id, f'copytest_{worker_id}')
@@ -37,54 +36,25 @@ async def parallel_load_different_tables(pg: Postgres, n_parallel: int):
 
 
 # Load 5 different tables in parallel with COPY TO
-def test_parallel_copy_different_tables(zenith_simple_env: ZenithEnv,
-                                        zenbenchmark: ZenithBenchmarker,
-                                        n_parallel=5):
+def test_parallel_copy_different_tables(zenith_with_baseline: PgCompare, n_parallel=5):
 
-    env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_parallel_copy_different_tables", "empty"])
-
-    pg = env.postgres.create_start('test_parallel_copy_different_tables')
-    log.info("postgres is running on 'test_parallel_copy_different_tables' branch")
-
-    # Open a connection directly to the page server that we'll use to force
-    # flushing the layers to disk
-    psconn = env.pageserver.connect()
-    pscur = psconn.cursor()
-
-    # Get the timeline ID of our branch. We need it for the 'do_gc' command
-    conn = pg.connect()
+    env = zenith_with_baseline
+    conn = env.pg.connect()
     cur = conn.cursor()
-    cur.execute("SHOW zenith.zenith_timeline")
-    timeline = cur.fetchone()[0]
 
     for worker_id in range(n_parallel):
         cur.execute(f'CREATE TABLE copytest_{worker_id} (i int, t text)')
 
-    with zenbenchmark.record_pageserver_writes(env.pageserver, 'pageserver_writes'):
-        with zenbenchmark.record_duration('load'):
-            asyncio.run(parallel_load_different_tables(pg, n_parallel))
+    with env.record_pageserver_writes('pageserver_writes'):
+        with env.record_duration('load'):
+            asyncio.run(parallel_load_different_tables(env.pg, n_parallel))
+            env.flush()
 
-            # Flush the layers from memory to disk. This is included in the reported
-            # time and I/O
-            pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
-
-    # Record peak memory usage
-    zenbenchmark.record("peak_mem",
-                        zenbenchmark.get_peak_mem(env.pageserver) / 1024,
-                        'MB',
-                        report=MetricReport.LOWER_IS_BETTER)
-
-    # Report disk space used by the repository
-    timeline_size = zenbenchmark.get_timeline_size(env.repo_dir, env.initial_tenant, timeline)
-    zenbenchmark.record('size',
-                        timeline_size / (1024 * 1024),
-                        'MB',
-                        report=MetricReport.LOWER_IS_BETTER)
+    env.report_peak_memory_use()
+    env.report_size()
 
 
-async def parallel_load_same_table(pg: Postgres, n_parallel: int):
+async def parallel_load_same_table(pg: PgProtocol, n_parallel: int):
     workers = []
     for worker_id in range(n_parallel):
         worker = copy_test_data_to_table(pg, worker_id, f'copytest')
@@ -95,46 +65,17 @@ async def parallel_load_same_table(pg: Postgres, n_parallel: int):
 
 
 # Load data into one table with COPY TO from 5 parallel connections
-def test_parallel_copy_same_table(zenith_simple_env: ZenithEnv,
-                                  zenbenchmark: ZenithBenchmarker,
-                                  n_parallel=5):
-    env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_parallel_copy_same_table", "empty"])
-
-    pg = env.postgres.create_start('test_parallel_copy_same_table')
-    log.info("postgres is running on 'test_parallel_copy_same_table' branch")
-
-    # Open a connection directly to the page server that we'll use to force
-    # flushing the layers to disk
-    psconn = env.pageserver.connect()
-    pscur = psconn.cursor()
-
-    # Get the timeline ID of our branch. We need it for the 'do_gc' command
-    conn = pg.connect()
+def test_parallel_copy_same_table(zenith_with_baseline: PgCompare, n_parallel=5):
+    env = zenith_with_baseline
+    conn = env.pg.connect()
     cur = conn.cursor()
-    cur.execute("SHOW zenith.zenith_timeline")
-    timeline = cur.fetchone()[0]
 
     cur.execute(f'CREATE TABLE copytest (i int, t text)')
 
-    with zenbenchmark.record_pageserver_writes(env.pageserver, 'pageserver_writes'):
-        with zenbenchmark.record_duration('load'):
-            asyncio.run(parallel_load_same_table(pg, n_parallel))
+    with env.record_pageserver_writes('pageserver_writes'):
+        with env.record_duration('load'):
+            asyncio.run(parallel_load_same_table(env.pg, n_parallel))
+            env.flush()
 
-            # Flush the layers from memory to disk. This is included in the reported
-            # time and I/O
-            pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
-
-    # Record peak memory usage
-    zenbenchmark.record("peak_mem",
-                        zenbenchmark.get_peak_mem(env.pageserver) / 1024,
-                        'MB',
-                        report=MetricReport.LOWER_IS_BETTER)
-
-    # Report disk space used by the repository
-    timeline_size = zenbenchmark.get_timeline_size(env.repo_dir, env.initial_tenant, timeline)
-    zenbenchmark.record('size',
-                        timeline_size / (1024 * 1024),
-                        'MB',
-                        report=MetricReport.LOWER_IS_BETTER)
+    env.report_peak_memory_use()
+    env.report_size()

test_runner/performance/test_perf_pgbench.py

@@ -1,11 +1,10 @@
 from contextlib import closing
-from fixtures.zenith_fixtures import PgBin, ZenithEnv
+from fixtures.zenith_fixtures import PgBin, VanillaPostgres, ZenithEnv
+from fixtures.compare_fixtures import PgCompare, VanillaCompare, ZenithCompare
 
 from fixtures.benchmark_fixture import MetricReport, ZenithBenchmarker
 from fixtures.log_helper import log
 
-pytest_plugins = ("fixtures.zenith_fixtures", "fixtures.benchmark_fixture")
-
 
 #
 # Run a very short pgbench test.
@@ -16,47 +15,16 @@ pytest_plugins = ("fixtures.zenith_fixtures", "fixtures.benchmark_fixture")
 # 2. Time to run 5000 pgbench transactions
 # 3. Disk space used
 #
-def test_pgbench(zenith_simple_env: ZenithEnv, pg_bin: PgBin, zenbenchmark: ZenithBenchmarker):
-    env = zenith_simple_env
-    # Create a branch for us
-    env.zenith_cli(["branch", "test_pgbench_perf", "empty"])
+def test_pgbench(zenith_with_baseline: PgCompare):
+    env = zenith_with_baseline
 
-    pg = env.postgres.create_start('test_pgbench_perf')
-    log.info("postgres is running on 'test_pgbench_perf' branch")
+    with env.record_pageserver_writes('pageserver_writes'):
+        with env.record_duration('init'):
+            env.pg_bin.run_capture(['pgbench', '-s5', '-i', env.pg.connstr()])
+            env.flush()
 
-    # Open a connection directly to the page server that we'll use to force
-    # flushing the layers to disk
-    psconn = env.pageserver.connect()
-    pscur = psconn.cursor()
+    with env.record_duration('5000_xacts'):
+        env.pg_bin.run_capture(['pgbench', '-c1', '-t5000', env.pg.connstr()])
+    env.flush()
 
-    # Get the timeline ID of our branch. We need it for the 'do_gc' command
-    with closing(pg.connect()) as conn:
-        with conn.cursor() as cur:
-            cur.execute("SHOW zenith.zenith_timeline")
-            timeline = cur.fetchone()[0]
-
-    connstr = pg.connstr()
-
-    # Initialize pgbench database, recording the time and I/O it takes
-    with zenbenchmark.record_pageserver_writes(env.pageserver, 'pageserver_writes'):
-        with zenbenchmark.record_duration('init'):
-            pg_bin.run_capture(['pgbench', '-s5', '-i', connstr])
-
-            # Flush the layers from memory to disk. This is included in the reported
-            # time and I/O
-            pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
-
-    # Run pgbench for 5000 transactions
-    with zenbenchmark.record_duration('5000_xacts'):
-        pg_bin.run_capture(['pgbench', '-c1', '-t5000', connstr])
-
-    # Flush the layers to disk again. This is *not' included in the reported time,
-    # though.
-    pscur.execute(f"do_gc {env.initial_tenant} {timeline} 0")
-
-    # Report disk space used by the repository
-    timeline_size = zenbenchmark.get_timeline_size(env.repo_dir, env.initial_tenant, timeline)
-    zenbenchmark.record('size',
-                        timeline_size / (1024 * 1024),
-                        'MB',
-                        report=MetricReport.LOWER_IS_BETTER)
+    env.report_size()

test_runner/performance/test_perf_pgbench_remote.py

@@ -9,8 +9,6 @@ import calendar
 import timeit
 import os
 
-pytest_plugins = ("fixtures.benchmark_fixture", )
-
 
 def utc_now_timestamp() -> int:
     return calendar.timegm(datetime.utcnow().utctimetuple())

test_runner/performance/test_random_writes.py

@@ -0,0 +1,79 @@
+import os
+from contextlib import closing
+from fixtures.benchmark_fixture import MetricReport
+from fixtures.zenith_fixtures import ZenithEnv
+from fixtures.compare_fixtures import PgCompare, VanillaCompare, ZenithCompare
+from fixtures.log_helper import log
+
+import psycopg2.extras
+import random
+import time
+from fixtures.utils import print_gc_result
+
+
+# This is a clear-box test that demonstrates the worst case scenario for the
+# "1 segment per layer" implementation of the pageserver. It writes to random
+# rows, while almost never writing to the same segment twice before flushing.
+# A naive pageserver implementation would create a full image layer for each
+# dirty segment, leading to write_amplification = segment_size / page_size,
+# when compared to vanilla postgres. With segment_size = 10MB, that's 1250.
+def test_random_writes(zenith_with_baseline: PgCompare):
+    env = zenith_with_baseline
+
+    # Number of rows in the test database. 1M rows runs quickly, but implies
+    # a small effective_checkpoint_distance, which makes the test less realistic.
+    # Using a 300 TB database would imply a 250 MB effective_checkpoint_distance,
+    # but it will take a very long time to run. From what I've seen so far,
+    # increasing n_rows doesn't have impact on the (zenith_runtime / vanilla_runtime)
+    # performance ratio.
+    n_rows = 1 * 1000 * 1000  # around 36 MB table
+
+    # Number of writes per 3 segments. A value of 1 should produce a random
+    # workload where we almost never write to the same segment twice. Larger
+    # values of load_factor produce a larger effective_checkpoint_distance,
+    # making the test more realistic, but less effective. If you want a realistic
+    # worst case scenario and you have time to wait you should increase n_rows instead.
+    load_factor = 1
+
+    # Not sure why but this matters in a weird way (up to 2x difference in perf).
+    # TODO look into it
+    n_iterations = 1
+
+    with closing(env.pg.connect()) as conn:
+        with conn.cursor() as cur:
+            # Create the test table
+            with env.record_duration('init'):
+                cur.execute("""
+                    CREATE TABLE Big(
+                        pk integer primary key,
+                        count integer default 0
+                    );
+                """)
+                cur.execute(f"INSERT INTO Big (pk) values (generate_series(1,{n_rows}))")
+
+            # Get table size (can't be predicted because padding and alignment)
+            cur.execute("SELECT pg_relation_size('Big');")
+            row = cur.fetchone()
+            table_size = row[0]
+            env.zenbenchmark.record("table_size", table_size, 'bytes', MetricReport.TEST_PARAM)
+
+            # Decide how much to write, based on knowledge of pageserver implementation.
+            # Avoiding segment collisions maximizes (zenith_runtime / vanilla_runtime).
+            segment_size = 10 * 1024 * 1024
+            n_segments = table_size // segment_size
+            n_writes = load_factor * n_segments // 3
+
+            # The closer this is to 250 MB, the more realistic the test is.
+            effective_checkpoint_distance = table_size * n_writes // n_rows
+            env.zenbenchmark.record("effective_checkpoint_distance",
+                                    effective_checkpoint_distance,
+                                    'bytes',
+                                    MetricReport.TEST_PARAM)
+
+            # Update random keys
+            with env.record_duration('run'):
+                for it in range(n_iterations):
+                    for i in range(n_writes):
+                        key = random.randint(1, n_rows)
+                        cur.execute(f"update Big set count=count+1 where pk={key}")
+                    env.flush()