cleanup logging

- move extension download to a separate thread;
- add timer around shared preload libraries downloading

Cargo.lock

@@ -593,7 +593,7 @@ dependencies = [
  "cc",
  "cfg-if",
  "libc",
- "miniz_oxide",
+ "miniz_oxide 0.6.2",
  "object",
  "rustc-demangle",
 ]
@@ -885,6 +885,7 @@ dependencies = [
  "chrono",
  "clap",
  "compute_api",
+ "flate2",
  "futures",
  "hyper",
  "notify",
@@ -892,12 +893,14 @@ dependencies = [
  "opentelemetry",
  "postgres",
  "regex",
+ "remote_storage",
  "reqwest",
  "serde",
  "serde_json",
  "tar",
  "tokio",
  "tokio-postgres",
+ "toml_edit",
  "tracing",
  "tracing-opentelemetry",
  "tracing-subscriber",
@@ -965,6 +968,7 @@ dependencies = [
  "tar",
  "thiserror",
  "toml",
+ "tracing",
  "url",
  "utils",
  "workspace_hack",
@@ -1367,6 +1371,16 @@ version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
 
+[[package]]
+name = "flate2"
+version = "1.0.26"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3b9429470923de8e8cbd4d2dc513535400b4b3fef0319fb5c4e1f520a7bef743"
+dependencies = [
+ "crc32fast",
+ "miniz_oxide 0.7.1",
+]
+
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -2151,6 +2165,15 @@ dependencies = [
  "adler",
 ]
 
+[[package]]
+name = "miniz_oxide"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7"
+dependencies = [
+ "adler",
+]
+
 [[package]]
 name = "mio"
 version = "0.8.6"
@@ -2617,16 +2640,6 @@ dependencies = [
  "windows-sys 0.45.0",
 ]
 
-[[package]]
-name = "pbkdf2"
-version = "0.12.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f0ca0b5a68607598bf3bad68f32227a8164f6254833f84eafaac409cd6746c31"
-dependencies = [
- "digest",
- "hmac",
-]
-
 [[package]]
 name = "peeking_take_while"
 version = "0.1.2"
@@ -3020,7 +3033,6 @@ dependencies = [
  "once_cell",
  "opentelemetry",
  "parking_lot 0.12.1",
- "pbkdf2",
  "pin-project-lite",
  "postgres-native-tls",
  "postgres_backend",

Cargo.toml

@@ -86,7 +86,6 @@ opentelemetry = "0.18.0"
 opentelemetry-otlp = { version = "0.11.0", default_features=false, features = ["http-proto", "trace", "http", "reqwest-client"] }
 opentelemetry-semantic-conventions = "0.10.0"
 parking_lot = "0.12"
-pbkdf2 = "0.12.1"
 pin-project-lite = "0.2"
 prometheus = {version = "0.13", default_features=false, features = ["process"]} # removes protobuf dependency
 prost = "0.11"

Dockerfile.compute-node

@@ -189,8 +189,8 @@ RUN wget https://github.com/df7cb/postgresql-unit/archive/refs/tags/7.7.tar.gz -
 FROM build-deps AS vector-pg-build
 COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
 
-RUN wget https://github.com/pgvector/pgvector/archive/refs/tags/v0.4.4.tar.gz -O pgvector.tar.gz && \
-    echo "1cb70a63f8928e396474796c22a20be9f7285a8a013009deb8152445b61b72e6 pgvector.tar.gz" | sha256sum --check && \
+RUN wget https://github.com/pgvector/pgvector/archive/refs/tags/v0.4.0.tar.gz -O pgvector.tar.gz && \
+    echo "b76cf84ddad452cc880a6c8c661d137ddd8679c000a16332f4f03ecf6e10bcc8 pgvector.tar.gz" | sha256sum --check && \
     mkdir pgvector-src && cd pgvector-src && tar xvzf ../pgvector.tar.gz --strip-components=1 -C . && \
     make -j $(getconf _NPROCESSORS_ONLN) PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
     make -j $(getconf _NPROCESSORS_ONLN) install PG_CONFIG=/usr/local/pgsql/bin/pg_config && \
@@ -515,25 +515,6 @@ RUN wget https://github.com/ChenHuajun/pg_roaringbitmap/archive/refs/tags/v0.5.4
     make -j $(getconf _NPROCESSORS_ONLN) install && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/roaringbitmap.control
 
-#########################################################################################
-#
-# Layer "pg-embedding-pg-build"
-# compile pg_embedding extension
-#
-#########################################################################################
-FROM build-deps AS pg-embedding-pg-build
-COPY --from=pg-build /usr/local/pgsql/ /usr/local/pgsql/
-
-ENV PATH "/usr/local/pgsql/bin/:$PATH"
-# 2465f831ea1f8d49c1d74f8959adb7fc277d70cd made on 05/07/2023
-# There is no release tag yet
-RUN wget https://github.com/neondatabase/pg_embedding/archive/2465f831ea1f8d49c1d74f8959adb7fc277d70cd.tar.gz -O pg_embedding.tar.gz && \
-    echo "047af2b1f664a1e6e37867bd4eeaf5934fa27d6ba3d6c4461efa388ddf7cd1d5 pg_embedding.tar.gz" | sha256sum --check && \
-    mkdir pg_embedding-src && cd pg_embedding-src && tar xvzf ../pg_embedding.tar.gz --strip-components=1 -C . && \
-    make -j $(getconf _NPROCESSORS_ONLN) && \
-    make -j $(getconf _NPROCESSORS_ONLN) install && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/embedding.control
-
 #########################################################################################
 #
 # Layer "pg-anon-pg-build"
@@ -690,7 +671,6 @@ COPY --from=pg-pgx-ulid-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=rdkit-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pg-uuidv7-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY --from=pg-roaringbitmap-pg-build /usr/local/pgsql/ /usr/local/pgsql/
-COPY --from=pg-embedding-pg-build /usr/local/pgsql/ /usr/local/pgsql/
 COPY pgxn/ pgxn/
 
 RUN make -j $(getconf _NPROCESSORS_ONLN) \

compute_tools/Cargo.toml

@@ -30,3 +30,6 @@ url.workspace = true
 compute_api.workspace = true
 utils.workspace = true
 workspace_hack.workspace = true
+toml_edit.workspace = true
+remote_storage = { version = "0.1", path = "../libs/remote_storage/" }
+flate2 = "1.0.26"

compute_tools/src/bin/compute_ctl.rs

@@ -5,6 +5,8 @@
 //! - `compute_ctl` accepts cluster (compute node) specification as a JSON file.
 //! - Every start is a fresh start, so the data directory is removed and
 //!   initialized again on each run.
+//! - If remote_extension_config is provided, it will be used to fetch extensions list
+//!  and download `shared_preload_libraries` from the remote storage.
 //! - Next it will put configuration files into the `PGDATA` directory.
 //! - Sync safekeepers and get commit LSN.
 //! - Get `basebackup` from pageserver using the returned on the previous step LSN.
@@ -27,7 +29,8 @@
 //! compute_ctl -D /var/db/postgres/compute \
 //!             -C 'postgresql://cloud_admin@localhost/postgres' \
 //!             -S /var/db/postgres/specs/current.json \
-//!             -b /usr/local/bin/postgres
+//!             -b /usr/local/bin/postgres \
+//!             -r {"bucket": "my-bucket", "region": "eu-central-1", "endpoint": "http:://localhost:9000"} \
 //! ```
 //!
 use std::collections::HashMap;
@@ -35,7 +38,7 @@ use std::fs::File;
 use std::panic;
 use std::path::Path;
 use std::process::exit;
-use std::sync::{mpsc, Arc, Condvar, Mutex};
+use std::sync::{mpsc, Arc, Condvar, Mutex, OnceLock};
 use std::{thread, time::Duration};
 
 use anyhow::{Context, Result};
@@ -48,6 +51,8 @@ use compute_api::responses::ComputeStatus;
 
 use compute_tools::compute::{ComputeNode, ComputeState, ParsedSpec};
 use compute_tools::configurator::launch_configurator;
+use compute_tools::extension_server::launch_download_extensions;
+use compute_tools::extension_server::{get_pg_version, init_remote_storage};
 use compute_tools::http::api::launch_http_server;
 use compute_tools::logger::*;
 use compute_tools::monitor::launch_monitor;
@@ -60,10 +65,21 @@ fn main() -> Result<()> {
     init_tracing_and_logging(DEFAULT_LOG_LEVEL)?;
 
     let build_tag = option_env!("BUILD_TAG").unwrap_or(BUILD_TAG_DEFAULT);
-
     info!("build_tag: {build_tag}");
 
     let matches = cli().get_matches();
+    let pgbin_default = String::from("postgres");
+    let pgbin = matches.get_one::<String>("pgbin").unwrap_or(&pgbin_default);
+
+    let remote_ext_config = matches.get_one::<String>("remote-ext-config");
+    // NOTE TODO: until control-plane changes, we can use the following line to forcibly enable remote extensions
+    // let remote_ext_config = Some(
+    //     r#"{"bucket": "neon-dev-extensions", "region": "eu-central-1", "endpoint": null, "prefix": "5555"}"#.to_string(),
+    // );
+    let ext_remote_storage = remote_ext_config.map(|x| {
+        init_remote_storage(x, build_tag)
+            .expect("cannot initialize remote extension storage from config")
+    });
 
     let http_port = *matches
         .get_one::<u16>("http-port")
@@ -128,9 +144,6 @@ fn main() -> Result<()> {
     let compute_id = matches.get_one::<String>("compute-id");
     let control_plane_uri = matches.get_one::<String>("control-plane-uri");
 
-    // Try to use just 'postgres' if no path is provided
-    let pgbin = matches.get_one::<String>("pgbin").unwrap();
-
     let spec;
     let mut live_config_allowed = false;
     match spec_json {
@@ -168,6 +181,7 @@ fn main() -> Result<()> {
 
     let mut new_state = ComputeState::new();
     let spec_set;
+
     if let Some(spec) = spec {
         let pspec = ParsedSpec::try_from(spec).map_err(|msg| anyhow::anyhow!(msg))?;
         new_state.pspec = Some(pspec);
@@ -179,9 +193,12 @@ fn main() -> Result<()> {
         connstr: Url::parse(connstr).context("cannot parse connstr as a URL")?,
         pgdata: pgdata.to_string(),
         pgbin: pgbin.to_string(),
+        pgversion: get_pg_version(pgbin),
         live_config_allowed,
         state: Mutex::new(new_state),
         state_changed: Condvar::new(),
+        ext_remote_storage,
+        available_extensions: OnceLock::new(),
     };
     let compute = Arc::new(compute_node);
 
@@ -190,6 +207,8 @@ fn main() -> Result<()> {
     let _http_handle =
         launch_http_server(http_port, &compute).expect("cannot launch http endpoint thread");
 
+    let extension_server_port: u16 = http_port;
+
     if !spec_set {
         // No spec provided, hang waiting for it.
         info!("no compute spec provided, waiting");
@@ -227,10 +246,13 @@ fn main() -> Result<()> {
     let _configurator_handle =
         launch_configurator(&compute).expect("cannot launch configurator thread");
 
+    let _download_extensions_handle =
+        launch_download_extensions(&compute).expect("cannot launch download extensions thread");
+
     // Start Postgres
     let mut delay_exit = false;
     let mut exit_code = None;
-    let pg = match compute.start_compute() {
+    let pg = match compute.start_compute(extension_server_port) {
         Ok(pg) => Some(pg),
         Err(err) => {
             error!("could not start the compute node: {:?}", err);
@@ -359,6 +381,12 @@ fn cli() -> clap::Command {
                 .long("control-plane-uri")
                 .value_name("CONTROL_PLANE_API_BASE_URI"),
         )
+        .arg(
+            Arg::new("remote-ext-config")
+                .short('r')
+                .long("remote-ext-config")
+                .value_name("REMOTE_EXT_CONFIG"),
+        )
 }
 
 #[test]

compute_tools/src/compute.rs

@@ -1,13 +1,16 @@
+use std::collections::HashSet;
 use std::fs;
 use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
 use std::process::{Command, Stdio};
 use std::str::FromStr;
-use std::sync::{Condvar, Mutex};
+use std::sync::{Condvar, Mutex, OnceLock};
 
 use anyhow::{Context, Result};
 use chrono::{DateTime, Utc};
+use futures::future::join_all;
 use postgres::{Client, NoTls};
+use tokio;
 use tokio_postgres;
 use tracing::{info, instrument, warn};
 use utils::id::{TenantId, TimelineId};
@@ -16,9 +19,11 @@ use utils::lsn::Lsn;
 use compute_api::responses::{ComputeMetrics, ComputeStatus};
 use compute_api::spec::{ComputeMode, ComputeSpec};
 
-use crate::config;
+use remote_storage::GenericRemoteStorage;
+
 use crate::pg_helpers::*;
 use crate::spec::*;
+use crate::{config, extension_server};
 
 /// Compute node info shared across several `compute_ctl` threads.
 pub struct ComputeNode {
@@ -26,6 +31,7 @@ pub struct ComputeNode {
     pub connstr: url::Url,
     pub pgdata: String,
     pub pgbin: String,
+    pub pgversion: String,
     /// We should only allow live re- / configuration of the compute node if
     /// it uses 'pull model', i.e. it can go to control-plane and fetch
     /// the latest configuration. Otherwise, there could be a case:
@@ -45,6 +51,10 @@ pub struct ComputeNode {
     pub state: Mutex<ComputeState>,
     /// `Condvar` to allow notifying waiters about state changes.
     pub state_changed: Condvar,
+    ///  the S3 bucket that we search for extensions in
+    pub ext_remote_storage: Option<GenericRemoteStorage>,
+    // cached lists of available extensions and libraries
+    pub available_extensions: OnceLock<HashSet<String>>,
 }
 
 #[derive(Clone, Debug)]
@@ -323,14 +333,22 @@ impl ComputeNode {
     /// Do all the preparations like PGDATA directory creation, configuration,
     /// safekeepers sync, basebackup, etc.
     #[instrument(skip_all)]
-    pub fn prepare_pgdata(&self, compute_state: &ComputeState) -> Result<()> {
+    pub fn prepare_pgdata(
+        &self,
+        compute_state: &ComputeState,
+        extension_server_port: u16,
+    ) -> Result<()> {
         let pspec = compute_state.pspec.as_ref().expect("spec must be set");
         let spec = &pspec.spec;
         let pgdata_path = Path::new(&self.pgdata);
 
         // Remove/create an empty pgdata directory and put configuration there.
         self.create_pgdata()?;
-        config::write_postgres_conf(&pgdata_path.join("postgresql.conf"), &pspec.spec)?;
+        config::write_postgres_conf(
+            &pgdata_path.join("postgresql.conf"),
+            &pspec.spec,
+            Some(extension_server_port),
+        )?;
 
         // Syncing safekeepers is only safe with primary nodes: if a primary
         // is already connected it will be kicked out, so a secondary (standby)
@@ -472,7 +490,7 @@ impl ComputeNode {
 
         // Write new config
         let pgdata_path = Path::new(&self.pgdata);
-        config::write_postgres_conf(&pgdata_path.join("postgresql.conf"), &spec)?;
+        config::write_postgres_conf(&pgdata_path.join("postgresql.conf"), &spec, None)?;
 
         let mut client = Client::connect(self.connstr.as_str(), NoTls)?;
         self.pg_reload_conf(&mut client)?;
@@ -502,7 +520,7 @@ impl ComputeNode {
     }
 
     #[instrument(skip_all)]
-    pub fn start_compute(&self) -> Result<std::process::Child> {
+    pub fn start_compute(&self, extension_server_port: u16) -> Result<std::process::Child> {
         let compute_state = self.state.lock().unwrap().clone();
         let pspec = compute_state.pspec.as_ref().expect("spec must be set");
         info!(
@@ -513,12 +531,31 @@ impl ComputeNode {
             pspec.timeline_id,
         );
 
-        self.prepare_pgdata(&compute_state)?;
+        // This part is sync, because we need to download
+        // remote shared_preload_libraries before postgres start (if any)
+        {
+            let library_load_start_time = Utc::now();
+            self.prepare_preload_libraries(&compute_state)?;
+
+            let library_load_time = Utc::now()
+                .signed_duration_since(library_load_start_time)
+                .to_std()
+                .unwrap()
+                .as_millis() as u64;
+            let mut state = self.state.lock().unwrap();
+            state.metrics.load_libraries_ms = library_load_time;
+            info!(
+                "Loading shared_preload_libraries took {:?}ms",
+                library_load_time
+            );
+        }
+
+        self.prepare_pgdata(&compute_state, extension_server_port)?;
 
         let start_time = Utc::now();
+
         let pg = self.start_postgres(pspec.storage_auth_token.clone())?;
 
-        let config_time = Utc::now();
         if pspec.spec.mode == ComputeMode::Primary && !pspec.spec.skip_pg_catalog_updates {
             self.apply_config(&compute_state)?;
         }
@@ -526,13 +563,8 @@ impl ComputeNode {
         let startup_end_time = Utc::now();
         {
             let mut state = self.state.lock().unwrap();
-            state.metrics.start_postgres_ms = config_time
-                .signed_duration_since(start_time)
-                .to_std()
-                .unwrap()
-                .as_millis() as u64;
             state.metrics.config_ms = startup_end_time
-                .signed_duration_since(config_time)
+                .signed_duration_since(start_time)
                 .to_std()
                 .unwrap()
                 .as_millis() as u64;
@@ -654,4 +686,92 @@ LIMIT 100",
             "{{\"pg_stat_statements\": []}}".to_string()
         }
     }
+
+    // If remote extension storage is configured,
+    // download extension control files
+    #[tokio::main]
+    pub async fn prepare_external_extensions(&self, compute_state: &ComputeState) -> Result<()> {
+        if let Some(ref ext_remote_storage) = self.ext_remote_storage {
+            let pspec = compute_state.pspec.as_ref().expect("spec must be set");
+            let spec = &pspec.spec;
+            let custom_ext_prefixes = spec.custom_extensions.clone().unwrap_or(Vec::new());
+            info!("custom_ext_prefixes: {:?}", &custom_ext_prefixes);
+            let available_extensions = extension_server::get_available_extensions(
+                ext_remote_storage,
+                &self.pgbin,
+                &self.pgversion,
+                &custom_ext_prefixes,
+            )
+            .await?;
+            self.available_extensions
+                .set(available_extensions)
+                .expect("available_extensions.set error");
+        }
+        Ok(())
+    }
+
+    pub async fn download_extension(&self, ext_name: &str) -> Result<()> {
+        match &self.ext_remote_storage {
+            None => anyhow::bail!("No remote extension storage"),
+            Some(remote_storage) => {
+                extension_server::download_extension(
+                    ext_name,
+                    remote_storage,
+                    &self.pgbin,
+                    &self.pgversion,
+                )
+                .await
+            }
+        }
+    }
+
+    #[tokio::main]
+    pub async fn prepare_preload_libraries(&self, compute_state: &ComputeState) -> Result<()> {
+        if self.ext_remote_storage.is_none() {
+            return Ok(());
+        }
+        let pspec = compute_state.pspec.as_ref().expect("spec must be set");
+        let spec = &pspec.spec;
+
+        info!("parse shared_preload_libraries from spec.cluster.settings");
+        let mut libs_vec = Vec::new();
+        if let Some(libs) = spec.cluster.settings.find("shared_preload_libraries") {
+            libs_vec = libs
+                .split(&[',', '\'', ' '])
+                .filter(|s| *s != "neon" && !s.is_empty())
+                .map(str::to_string)
+                .collect();
+        }
+        info!("parse shared_preload_libraries from provided postgresql.conf");
+        // that is used in neon_local and python tests
+        if let Some(conf) = &spec.cluster.postgresql_conf {
+            let conf_lines = conf.split('\n').collect::<Vec<&str>>();
+            let mut shared_preload_libraries_line = "";
+            for line in conf_lines {
+                if line.starts_with("shared_preload_libraries") {
+                    shared_preload_libraries_line = line;
+                }
+            }
+            let mut preload_libs_vec = Vec::new();
+            if let Some(libs) = shared_preload_libraries_line.split("='").nth(1) {
+                preload_libs_vec = libs
+                    .split(&[',', '\'', ' '])
+                    .filter(|s| *s != "neon" && !s.is_empty())
+                    .map(str::to_string)
+                    .collect();
+            }
+            libs_vec.extend(preload_libs_vec);
+        }
+
+        info!("Downloading to shared preload libraries: {:?}", &libs_vec);
+        let mut download_tasks = Vec::new();
+        for library in &libs_vec {
+            download_tasks.push(self.download_extension(library));
+        }
+        let results = join_all(download_tasks).await;
+        for result in results {
+            result?; // propogate any errors
+        }
+        Ok(())
+    }
 }

compute_tools/src/config.rs

@@ -33,7 +33,11 @@ pub fn line_in_file(path: &Path, line: &str) -> Result<bool> {
 }
 
 /// Create or completely rewrite configuration file specified by `path`
-pub fn write_postgres_conf(path: &Path, spec: &ComputeSpec) -> Result<()> {
+pub fn write_postgres_conf(
+    path: &Path,
+    spec: &ComputeSpec,
+    extension_server_port: Option<u16>,
+) -> Result<()> {
     // File::create() destroys the file content if it exists.
     let mut file = File::create(path)?;
 
@@ -95,5 +99,9 @@ pub fn write_postgres_conf(path: &Path, spec: &ComputeSpec) -> Result<()> {
         writeln!(file, "# Managed by compute_ctl: end")?;
     }
 
+    if let Some(port) = extension_server_port {
+        writeln!(file, "neon.extension_server_port={}", port)?;
+    }
+
     Ok(())
 }

compute_tools/src/configurator.rs

@@ -42,13 +42,15 @@ fn configurator_main_loop(compute: &Arc<ComputeNode>) {
     }
 }
 
-pub fn launch_configurator(compute: &Arc<ComputeNode>) -> Result<thread::JoinHandle<()>> {
+pub fn launch_configurator(
+    compute: &Arc<ComputeNode>,
+) -> Result<thread::JoinHandle<()>, std::io::Error> {
     let compute = Arc::clone(compute);
 
-    Ok(thread::Builder::new()
+    thread::Builder::new()
         .name("compute-configurator".into())
         .spawn(move || {
             configurator_main_loop(&compute);
             info!("configurator thread is exited");
-        })?)
+        })
 }

compute_tools/src/extension_server.rs

@@ -0,0 +1,237 @@
+// Download extension files from the extension store
+// and put them in the right place in the postgres directory
+/*
+The layout of the S3 bucket is as follows:
+
+v14/ext_index.json
+    -- this contains information necessary to create control files
+v14/extensions/test_ext1.tar.gz
+    -- this contains the library files and sql files necessary to create this extension
+v14/extensions/custom_ext1.tar.gz
+
+The difference between a private and public extensions is determined by who can
+load the extension this is specified in ext_index.json
+
+Speicially, ext_index.json has a list of public extensions, and a list of
+extensions enabled for specific tenant-ids.
+*/
+use crate::compute::ComputeNode;
+use anyhow::Context;
+use anyhow::{self, Result};
+use flate2::read::GzDecoder;
+use remote_storage::*;
+use serde_json::{self, Value};
+use std::collections::HashSet;
+use std::num::{NonZeroU32, NonZeroUsize};
+use std::path::Path;
+use std::str;
+use std::sync::Arc;
+use std::thread;
+use tar::Archive;
+use tokio::io::AsyncReadExt;
+use tracing::info;
+
+fn get_pg_config(argument: &str, pgbin: &str) -> String {
+    // gives the result of `pg_config [argument]`
+    // where argument is a flag like `--version` or `--sharedir`
+    let pgconfig = pgbin
+        .strip_suffix("postgres")
+        .expect("bad pgbin")
+        .to_owned()
+        + "/pg_config";
+    let config_output = std::process::Command::new(pgconfig)
+        .arg(argument)
+        .output()
+        .expect("pg_config error");
+    std::str::from_utf8(&config_output.stdout)
+        .expect("pg_config error")
+        .trim()
+        .to_string()
+}
+
+pub fn get_pg_version(pgbin: &str) -> String {
+    // pg_config --version returns a (platform specific) human readable string
+    // such as "PostgreSQL 15.4". We parse this to v14/v15
+    let human_version = get_pg_config("--version", pgbin);
+    if human_version.contains("15") {
+        return "v15".to_string();
+    } else if human_version.contains("14") {
+        return "v14".to_string();
+    }
+    panic!("Unsuported postgres version {human_version}");
+}
+
+// download extension control files
+// if custom_ext_prefixes is provided - search also in custom extension paths
+pub async fn get_available_extensions(
+    remote_storage: &GenericRemoteStorage,
+    pgbin: &str,
+    pg_version: &str,
+    custom_ext_prefixes: &[String],
+) -> Result<HashSet<String>> {
+    let local_sharedir = Path::new(&get_pg_config("--sharedir", pgbin)).join("extension");
+    let index_path = pg_version.to_owned() + "/ext_index.json";
+    let index_path = RemotePath::new(Path::new(&index_path)).context("error forming path")?;
+    info!("download ext_index.json: {:?}", &index_path);
+
+    // TODO: potential optimization: cache ext_index.json
+    let mut download = remote_storage.download(&index_path).await?;
+    let mut write_data_buffer = Vec::new();
+    download
+        .download_stream
+        .read_to_end(&mut write_data_buffer)
+        .await?;
+    let ext_index_str = match str::from_utf8(&write_data_buffer) {
+        Ok(v) => v,
+        Err(e) => panic!("Invalid UTF-8 sequence: {}", e),
+    };
+
+    let ext_index_full: Value = serde_json::from_str(ext_index_str)?;
+    let ext_index_full = ext_index_full.as_object().context("error parsing json")?;
+    let control_data = ext_index_full["control_data"]
+        .as_object()
+        .context("json parse error")?;
+    let enabled_extensions = ext_index_full["enabled_extensions"]
+        .as_object()
+        .context("json parse error")?;
+    info!("{:?}", control_data.clone());
+    info!("{:?}", enabled_extensions.clone());
+
+    let mut prefixes = vec!["public".to_string()];
+    prefixes.extend(custom_ext_prefixes.to_owned());
+    info!("{:?}", &prefixes);
+    let mut all_extensions = HashSet::new();
+    for prefix in prefixes {
+        let prefix_extensions = match enabled_extensions.get(&prefix) {
+            Some(Value::Array(ext_name)) => ext_name,
+            _ => {
+                info!("prefix {} has no extensions", prefix);
+                continue;
+            }
+        };
+        info!("{:?}", prefix_extensions);
+        for ext_name in prefix_extensions {
+            all_extensions.insert(ext_name.as_str().context("json parse error")?.to_string());
+        }
+    }
+
+    for prefix in &all_extensions {
+        let control_contents = control_data[prefix].as_str().context("json parse error")?;
+        let control_path = local_sharedir.join(prefix.to_owned() + ".control");
+
+        info!("WRITING FILE {:?}{:?}", control_path, control_contents);
+        std::fs::write(control_path, control_contents)?;
+    }
+
+    Ok(all_extensions.into_iter().collect())
+}
+
+// download all sqlfiles (and possibly data files) for a given extension name
+pub async fn download_extension(
+    ext_name: &str,
+    remote_storage: &GenericRemoteStorage,
+    pgbin: &str,
+    pg_version: &str,
+) -> Result<()> {
+    // TODO: potential optimization: only download the extension if it doesn't exist
+    // problem: how would we tell if it exists?
+    let ext_name = ext_name.replace(".so", "");
+    let ext_name_targz = ext_name.to_owned() + ".tar.gz";
+    if Path::new(&ext_name_targz).exists() {
+        info!("extension {:?} already exists", ext_name_targz);
+        return Ok(());
+    }
+    let ext_path = RemotePath::new(
+        &Path::new(pg_version)
+            .join("extensions")
+            .join(ext_name_targz.clone()),
+    )?;
+    info!(
+        "Start downloading extension {:?} from {:?}",
+        ext_name, ext_path
+    );
+    let mut download = remote_storage.download(&ext_path).await?;
+    let mut write_data_buffer = Vec::new();
+    download
+        .download_stream
+        .read_to_end(&mut write_data_buffer)
+        .await?;
+    let unzip_dest = pgbin.strip_suffix("/bin/postgres").expect("bad pgbin");
+    let tar = GzDecoder::new(std::io::Cursor::new(write_data_buffer));
+    let mut archive = Archive::new(tar);
+    archive.unpack(unzip_dest)?;
+    info!("Download + unzip {:?} completed successfully", &ext_path);
+
+    let local_sharedir = Path::new(&get_pg_config("--sharedir", pgbin)).join("extension");
+    let zip_sharedir = format!("{unzip_dest}/extensions/{ext_name}/share/extension");
+    info!("mv {zip_sharedir:?}/* {local_sharedir:?}");
+    for file in std::fs::read_dir(zip_sharedir)? {
+        let old_file = file?.path();
+        let new_file =
+            Path::new(&local_sharedir).join(old_file.file_name().context("error parsing file")?);
+        std::fs::rename(old_file, new_file)?;
+    }
+    let local_libdir = Path::new(&get_pg_config("--libdir", pgbin)).join("postgresql");
+    let zip_libdir = format!("{unzip_dest}/extensions/{ext_name}/lib");
+    info!("mv {zip_libdir:?}/* {local_libdir:?}");
+    for file in std::fs::read_dir(zip_libdir)? {
+        let old_file = file?.path();
+        let new_file =
+            Path::new(&local_libdir).join(old_file.file_name().context("error parsing file")?);
+        std::fs::rename(old_file, new_file)?;
+    }
+    Ok(())
+}
+
+// This function initializes the necessary structs to use remmote storage (should be fairly cheap)
+pub fn init_remote_storage(
+    remote_ext_config: &str,
+    default_prefix: &str,
+) -> anyhow::Result<GenericRemoteStorage> {
+    let remote_ext_config: serde_json::Value = serde_json::from_str(remote_ext_config)?;
+
+    let remote_ext_bucket = remote_ext_config["bucket"]
+        .as_str()
+        .context("config parse error")?;
+    let remote_ext_region = remote_ext_config["region"]
+        .as_str()
+        .context("config parse error")?;
+    let remote_ext_endpoint = remote_ext_config["endpoint"].as_str();
+    let remote_ext_prefix = remote_ext_config["prefix"]
+        .as_str()
+        .unwrap_or(default_prefix)
+        .to_string();
+
+    // TODO: potentially allow modification of other parameters
+    // however, default values should be fine for now
+    let config = S3Config {
+        bucket_name: remote_ext_bucket.to_string(),
+        bucket_region: remote_ext_region.to_string(),
+        prefix_in_bucket: Some(remote_ext_prefix),
+        endpoint: remote_ext_endpoint.map(|x| x.to_string()),
+        concurrency_limit: NonZeroUsize::new(100).expect("100 != 0"),
+        max_keys_per_list_response: None,
+    };
+    let config = RemoteStorageConfig {
+        max_concurrent_syncs: NonZeroUsize::new(100).expect("100 != 0"),
+        max_sync_errors: NonZeroU32::new(100).expect("100 != 0"),
+        storage: RemoteStorageKind::AwsS3(config),
+    };
+    GenericRemoteStorage::from_config(&config)
+}
+
+pub fn launch_download_extensions(
+    compute: &Arc<ComputeNode>,
+) -> Result<thread::JoinHandle<()>, std::io::Error> {
+    let compute = Arc::clone(compute);
+    thread::Builder::new()
+        .name("download-extensions".into())
+        .spawn(move || {
+            info!("start download_extension_files");
+            let compute_state = compute.state.lock().expect("error unlocking compute.state");
+            compute
+                .prepare_external_extensions(&compute_state)
+                .expect("error preparing extensions");
+            info!("download_extension_files done, exiting thread");
+        })
+}

compute_tools/src/http/api.rs

@@ -121,6 +121,27 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
             }
         }
 
+        // download extension files from S3 on demand
+        (&Method::POST, route) if route.starts_with("/extension_server/") => {
+            info!("serving {:?} POST request", route);
+            info!("req.uri {:?}", req.uri());
+            let filename = route.split('/').last().unwrap().to_string();
+            info!(
+                "serving /extension_server POST request, filename: {:?}",
+                &filename
+            );
+
+            match compute.download_extension(&filename).await {
+                Ok(_) => Response::new(Body::from("OK")),
+                Err(e) => {
+                    error!("extension download failed: {}", e);
+                    let mut resp = Response::new(Body::from(e.to_string()));
+                    *resp.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
+                    resp
+                }
+            }
+        }
+
         // Return the `404 Not Found` for any other routes.
         _ => {
             let mut not_found = Response::new(Body::from("404 Not Found"));

compute_tools/src/http/openapi_spec.yaml

@@ -139,6 +139,34 @@ paths:
             application/json:
               schema:
                 $ref: "#/components/schemas/GenericError"
+  /extension_server:
+    post:
+      tags:
+      - Extension
+      summary: Download extension from S3 to local folder.
+      description: ""
+      operationId: downloadExtension
+      responses:
+        200:
+          description: Extension downloaded
+          content:
+            text/plain:
+              schema:
+                type: string
+                description: Error text or 'OK' if download succeeded.
+                example: "OK"
+        400:
+        description: Request is invalid.
+        content:
+          application/json:
+            schema:
+              $ref: "#/components/schemas/GenericError"
+        500:
+        description: Extension download request failed.
+        content:
+          application/json:
+            schema:
+              $ref: "#/components/schemas/GenericError"
 
 components:
   securitySchemes:

compute_tools/src/lib.rs

@@ -9,6 +9,7 @@ pub mod http;
 #[macro_use]
 pub mod logger;
 pub mod compute;
+pub mod extension_server;
 pub mod monitor;
 pub mod params;
 pub mod pg_helpers;

compute_tools/src/monitor.rs

@@ -105,10 +105,10 @@ fn watch_compute_activity(compute: &ComputeNode) {
 }
 
 /// Launch a separate compute monitor thread and return its `JoinHandle`.
-pub fn launch_monitor(state: &Arc<ComputeNode>) -> Result<thread::JoinHandle<()>> {
+pub fn launch_monitor(state: &Arc<ComputeNode>) -> Result<thread::JoinHandle<()>, std::io::Error> {
     let state = Arc::clone(state);
 
-    Ok(thread::Builder::new()
+    thread::Builder::new()
         .name("compute-monitor".into())
-        .spawn(move || watch_compute_activity(&state))?)
+        .spawn(move || watch_compute_activity(&state))
 }

compute_tools/src/spec.rs

@@ -124,7 +124,7 @@ pub fn get_spec_from_control_plane(
 pub fn handle_configuration(spec: &ComputeSpec, pgdata_path: &Path) -> Result<()> {
     // File `postgresql.conf` is no longer included into `basebackup`, so just
     // always write all config into it creating new file.
-    config::write_postgres_conf(&pgdata_path.join("postgresql.conf"), spec)?;
+    config::write_postgres_conf(&pgdata_path.join("postgresql.conf"), spec, None)?;
 
     update_pg_hba(pgdata_path)?;
 

control_plane/Cargo.toml

@@ -32,3 +32,4 @@ utils.workspace = true
 
 compute_api.workspace = true
 workspace_hack.workspace = true
+tracing.workspace = true

control_plane/src/bin/neon_local.rs

@@ -658,6 +658,8 @@ fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<(
                 .get_one::<String>("endpoint_id")
                 .ok_or_else(|| anyhow!("No endpoint ID was provided to start"))?;
 
+            let remote_ext_config = sub_args.get_one::<String>("remote-ext-config");
+
             // If --safekeepers argument is given, use only the listed safekeeper nodes.
             let safekeepers =
                 if let Some(safekeepers_str) = sub_args.get_one::<String>("safekeepers") {
@@ -699,7 +701,7 @@ fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<(
                     _ => {}
                 }
                 println!("Starting existing endpoint {endpoint_id}...");
-                endpoint.start(&auth_token, safekeepers)?;
+                endpoint.start(&auth_token, safekeepers, remote_ext_config)?;
             } else {
                 let branch_name = sub_args
                     .get_one::<String>("branch-name")
@@ -743,7 +745,7 @@ fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Result<(
                     pg_version,
                     mode,
                 )?;
-                ep.start(&auth_token, safekeepers)?;
+                ep.start(&auth_token, safekeepers, remote_ext_config)?;
             }
         }
         "stop" => {
@@ -1003,6 +1005,12 @@ fn cli() -> Command {
         .help("Additional pageserver's configuration options or overrides, refer to pageserver's 'config-override' CLI parameter docs for more")
         .required(false);
 
+    let remote_ext_config_args = Arg::new("remote-ext-config")
+        .long("remote-ext-config")
+        .num_args(1)
+        .help("Configure the S3 bucket that we search for extensions in.")
+        .required(false);
+
     let lsn_arg = Arg::new("lsn")
         .long("lsn")
         .help("Specify Lsn on the timeline to start from. By default, end of the timeline would be used.")
@@ -1161,6 +1169,7 @@ fn cli() -> Command {
                     .arg(pg_version_arg)
                     .arg(hot_standby_arg)
                     .arg(safekeepers_arg)
+                    .arg(remote_ext_config_args)
                 )
                 .subcommand(
                     Command::new("stop")

control_plane/src/endpoint.rs

@@ -311,7 +311,7 @@ impl Endpoint {
 
                 // TODO: use future host field from safekeeper spec
                 // Pass the list of safekeepers to the replica so that it can connect to any of them,
-                // whichever is availiable.
+                // whichever is available.
                 let sk_ports = self
                     .env
                     .safekeepers
@@ -418,7 +418,12 @@ impl Endpoint {
         Ok(())
     }
 
-    pub fn start(&self, auth_token: &Option<String>, safekeepers: Vec<NodeId>) -> Result<()> {
+    pub fn start(
+        &self,
+        auth_token: &Option<String>,
+        safekeepers: Vec<NodeId>,
+        remote_ext_config: Option<&String>,
+    ) -> Result<()> {
         if self.status() == "running" {
             anyhow::bail!("The endpoint is already running");
         }
@@ -486,6 +491,15 @@ impl Endpoint {
             pageserver_connstring: Some(pageserver_connstring),
             safekeeper_connstrings,
             storage_auth_token: auth_token.clone(),
+            // TODO FIXME: This is a hack to test custom extensions locally.
+            // In test_download_extensions, we assume that the custom extension
+            // prefix is the tenant ID. So we set it here.
+            //
+            // The proper way to implement this is to pass the custom extension
+            // in spec, but we don't have a way to do that yet in the python tests.
+            // NEW HACK: we enable the anon custom extension for everyone! this is of course just for testing
+            // how will we do it for real?
+            custom_extensions: Some(vec!["123454321".to_string(), self.tenant_id.to_string()]),
         };
         let spec_path = self.endpoint_path().join("spec.json");
         std::fs::write(spec_path, serde_json::to_string_pretty(&spec)?)?;
@@ -517,6 +531,11 @@ impl Endpoint {
             .stdin(std::process::Stdio::null())
             .stderr(logfile.try_clone()?)
             .stdout(logfile);
+
+        if let Some(remote_ext_config) = remote_ext_config {
+            cmd.args(["--remote-ext-config", remote_ext_config]);
+        }
+
         let child = cmd.spawn()?;
 
         // Write down the pid so we can wait for it when we want to stop

docker-compose/docker-compose.yml

@@ -189,7 +189,7 @@ services:
       - "/bin/bash"
       - "-c"
     command:
-      - "until pg_isready -h compute -p 55433 -U cloud_admin ; do
+      - "until pg_isready -h compute -p 55433 ; do
             echo 'Waiting to start compute...' && sleep 1;
          done"
     depends_on:

docs/docker.md

@@ -48,7 +48,6 @@ Creating docker-compose_storage_broker_1       ... done
 2. connect compute node
 ```
 $ echo "localhost:55433:postgres:cloud_admin:cloud_admin" >> ~/.pgpass
-$ chmod 600 ~/.pgpass
 $ psql -h localhost -p 55433 -U cloud_admin
 postgres=# CREATE TABLE t(key int primary key, value text);
 CREATE TABLE

docs/rfcs/024-extension-loading.md

@@ -0,0 +1,183 @@
+# Supporting custom user Extensions (Dynamic Extension Loading)
+Created 2023-05-03
+
+## Motivation
+
+There are many extensions in the PostgreSQL ecosystem, and not all extensions
+are of a quality that we can confidently support them. Additionally, our
+current extension inclusion mechanism has several problems because we build all
+extensions into the primary Compute image: We build the extensions every time
+we build the compute image regardless of whether we actually need to rebuild
+the image, and the inclusion of these extensions in the image adds a hard
+dependency on all supported extensions - thus increasing the image size, and
+with it the time it takes to download that image - increasing first start
+latency.
+
+This RFC proposes a dynamic loading mechanism that solves most of these
+problems.
+
+## Summary
+
+`compute_ctl` is made responsible for loading extensions on-demand into
+the container's file system for dynamically loaded extensions, and will also
+make sure that the extensions in `shared_preload_libraries` are downloaded
+before the compute node starts.
+
+## Components
+
+compute_ctl, PostgreSQL, neon (extension), Compute Host Node, Extension Store
+
+## Requirements
+
+Compute nodes with no extra extensions should not be negatively impacted by
+the existence of support for many extensions.
+
+Installing an extension into PostgreSQL should be easy.
+
+Non-preloaded extensions shouldn't impact startup latency.
+
+Uninstalled extensions shouldn't impact query latency.
+
+A small latency penalty for dynamically loaded extensions is acceptable in
+the first seconds of compute startup, but not in steady-state operations.
+
+## Proposed implementation
+
+### On-demand, JIT-loading of extensions
+
+Before postgres starts we download 
+- control files for all extensions available to that compute node;
+- all `shared_preload_libraries`;
+
+After postgres is running, `compute_ctl` listens for requests to load files.
+When PostgreSQL requests a file, `compute_ctl` downloads it.
+
+PostgreSQL requests files in the following cases:
+- When loading a preload library set in `local_preload_libraries`
+- When explicitly loading a library with `LOAD`
+- Wnen creating extension with `CREATE EXTENSION` (download sql scripts, (optional) extension data files and (optional) library files)))
+
+
+#### Summary
+
+Pros:
+ - Startup is only as slow as it takes to load all (shared_)preload_libraries
+ - Supports BYO Extension
+
+Cons:
+ - O(sizeof(extensions)) IO requirement for loading all extensions.
+
+### Alternative solutions
+
+1. Allow users to add their extensions to the base image
+   
+   Pros:
+    - Easy to deploy
+
+   Cons:
+    - Doesn't scale - first start size is dependent on image size;
+    - All extensions are shared across all users: It doesn't allow users to
+      bring their own restrictive-licensed extensions
+
+2. Bring Your Own compute image
+   
+   Pros:
+    - Still easy to deploy
+    - User can bring own patched version of PostgreSQL
+
+   Cons:
+    - First start latency is O(sizeof(extensions image))
+    - Warm instance pool for skipping pod schedule latency is not feasible with
+      O(n) custom images
+    - Support channels are difficult to manage
+
+3. Download all user extensions in bulk on compute start
+   
+   Pros:
+    - Easy to deploy
+    - No startup latency issues for "clean" users.
+    - Warm instance pool for skipping pod schedule latency is possible
+
+   Cons:
+    - Downloading all extensions in advance takes a lot of time, thus startup
+      latency issues
+
+4. Store user's extensions in persistent storage
+   
+   Pros:
+    - Easy to deploy
+    - No startup latency issues
+    - Warm instance pool for skipping pod schedule latency is possible
+
+   Cons:
+    - EC2 instances have only limited number of attachments shared between EBS
+      volumes, direct-attached NVMe drives, and ENIs.
+    - Compute instance migration isn't trivially solved for EBS mounts (e.g.
+      the device is unavailable whilst moving the mount between instances).
+    - EBS can only mount on one instance at a time (except the expensive IO2
+      device type).
+
+5. Store user's extensions in network drive
+   
+   Pros:
+    - Easy to deploy
+    - Few startup latency issues
+    - Warm instance pool for skipping pod schedule latency is possible
+
+   Cons:
+    - We'd need networked drives, and a lot of them, which would store many
+      duplicate extensions.
+    - **UNCHECKED:** Compute instance migration may not work nicely with
+      networked IOs
+
+
+### Idea extensions
+
+The extension store does not have to be S3 directly, but could be a Node-local
+caching service on top of S3. This would reduce the load on the network for
+popular extensions.
+
+## Extension Storage implementation
+
+Extension Storage in our case is an S3 bucket with a "directory" per build and postgres version,
+where extension files are stored as plain files in the bucket following the same directory structure as in the postgres.
+
+i.e.
+
+`s3://<the-bucket>/<build-version>/<postgres-version>/lib/postgis-3.1.so`
+`s3://<the-bucket>/<build-version>/<postgres-version>/share/extension/postgis.control`
+`s3://<the-bucket>/<build-version>/<postgres-version>/share/extension/postgis--3.1.sql`
+
+To handle custom extensions, that available only to specific users, we use per-extension subdirectories:
+
+i.e.
+`s3://<the-bucket>/<build-version>/<postgres-version>/<custom-ext-prefix>/lib/ext-name.so`, etc.
+`s3://<the-bucket>/<build-version>/<postgres-version>/<custom-ext-prefix>/share/extension/ext-name.control`, etc.
+
+On compute start, `compute_ctl` accepts a list of custom_ext_prefixes.
+
+To get the list of available extensions,`compute_ctl` downloads control files from all prefixes:
+
+`s3://<the-bucket>/<build-version>/<postgres-version>/share/extension/`
+`s3://<the-bucket>/<build-version>/<postgres-version>/<custom-ext-prefix1>/share/extension/`
+`s3://<the-bucket>/<build-version>/<postgres-version>/<custom-ext-prefix2>/share/extension/`
+
+
+
+### How to add new extension to the Extension Storage?
+
+Simply upload build artifacts to the S3 bucket.
+Implement a CI step for that. Splitting it from ompute-node-image build.
+
+### How do we deal with extension versions and updates?
+
+Currently, we rebuild extensions on every compute-node-image build and store them in the <build-version> prefix.
+This is needed to ensure that `/share` and `/lib` files are in sync.
+
+For extension updates, we rely on the PostgreSQL extension versioning mechanism (sql update scripts) and extension authors to not break backwards compatibility within one major version of PostgreSQL.
+
+### Alternatives
+
+For extensions written on trusted languages we can also adopt
+`dbdev` PostgreSQL Package Manager based on `pg_tle` by Supabase.
+This will increase the amount supported extensions and decrease the amount of work required to support them.

libs/compute_api/src/responses.rs

@@ -71,9 +71,9 @@ pub struct ComputeMetrics {
     pub wait_for_spec_ms: u64,
     pub sync_safekeepers_ms: u64,
     pub basebackup_ms: u64,
-    pub start_postgres_ms: u64,
     pub config_ms: u64,
     pub total_startup_ms: u64,
+    pub load_libraries_ms: u64,
 }
 
 /// Response of the `/computes/{compute_id}/spec` control-plane API.

libs/compute_api/src/spec.rs

@@ -60,6 +60,9 @@ pub struct ComputeSpec {
     /// If set, 'storage_auth_token' is used as the password to authenticate to
     /// the pageserver and safekeepers.
     pub storage_auth_token: Option<String>,
+
+    // list of prefixes to search for custom extensions in remote extension storage
+    pub custom_extensions: Option<Vec<String>>,
 }
 
 #[serde_as]

libs/postgres_ffi/src/nonrelfile_utils.rs

@@ -57,9 +57,9 @@ pub fn slru_may_delete_clogsegment(segpage: u32, cutoff_page: u32) -> bool {
 // Multixact utils
 
 pub fn mx_offset_to_flags_offset(xid: MultiXactId) -> usize {
-    ((xid / pg_constants::MULTIXACT_MEMBERS_PER_MEMBERGROUP as u32)
-        % pg_constants::MULTIXACT_MEMBERGROUPS_PER_PAGE as u32
-        * pg_constants::MULTIXACT_MEMBERGROUP_SIZE as u32) as usize
+    ((xid / pg_constants::MULTIXACT_MEMBERS_PER_MEMBERGROUP as u32) as u16
+        % pg_constants::MULTIXACT_MEMBERGROUPS_PER_PAGE
+        * pg_constants::MULTIXACT_MEMBERGROUP_SIZE) as usize
 }
 
 pub fn mx_offset_to_flags_bitshift(xid: MultiXactId) -> u16 {

libs/remote_storage/src/lib.rs

@@ -184,6 +184,20 @@ pub enum GenericRemoteStorage {
 }
 
 impl GenericRemoteStorage {
+    // A function for listing all the files in a "directory"
+    // Example:
+    // list_files("foo/bar") = ["foo/bar/a.txt", "foo/bar/b.txt"]
+    pub async fn list_files(&self, folder: Option<&RemotePath>) -> anyhow::Result<Vec<RemotePath>> {
+        match self {
+            Self::LocalFs(s) => s.list_files(folder).await,
+            Self::AwsS3(s) => s.list_files(folder).await,
+            Self::Unreliable(s) => s.list_files(folder).await,
+        }
+    }
+
+    // lists common *prefixes*, if any of files
+    // Example:
+    // list_prefixes("foo123","foo567","bar123","bar432") = ["foo", "bar"]
     pub async fn list_prefixes(
         &self,
         prefix: Option<&RemotePath>,
@@ -195,14 +209,6 @@ impl GenericRemoteStorage {
         }
     }
 
-    pub async fn list_files(&self, folder: Option<&RemotePath>) -> anyhow::Result<Vec<RemotePath>> {
-        match self {
-            Self::LocalFs(s) => s.list_files(folder).await,
-            Self::AwsS3(s) => s.list_files(folder).await,
-            Self::Unreliable(s) => s.list_files(folder).await,
-        }
-    }
-
     pub async fn upload(
         &self,
         from: impl io::AsyncRead + Unpin + Send + Sync + 'static,

libs/remote_storage/src/s3_bucket.rs

@@ -349,10 +349,17 @@ impl RemoteStorage for S3Bucket {
 
     /// See the doc for `RemoteStorage::list_files`
     async fn list_files(&self, folder: Option<&RemotePath>) -> anyhow::Result<Vec<RemotePath>> {
-        let folder_name = folder
+        let mut folder_name = folder
             .map(|p| self.relative_path_to_s3_object(p))
             .or_else(|| self.prefix_in_bucket.clone());
 
+        // remove leading "/" if one exists
+        if let Some(folder_name_slash) = folder_name.clone() {
+            if folder_name_slash.starts_with(REMOTE_STORAGE_PREFIX_SEPARATOR) {
+                folder_name = Some(folder_name_slash[1..].to_string());
+            }
+        }
+
         // AWS may need to break the response into several parts
         let mut continuation_token = None;
         let mut all_files = vec![];

pageserver/src/consumption_metrics.rs

@@ -234,18 +234,14 @@ pub async fn collect_metrics_iteration(
         // Note that this metric is calculated in a separate bgworker
         // Here we only use cached value, which may lag behind the real latest one
         let tenant_synthetic_size = tenant.get_cached_synthetic_size();
-
-        if tenant_synthetic_size != 0 {
-            // only send non-zeroes because otherwise these show up as errors in logs
-            current_metrics.push((
-                PageserverConsumptionMetricsKey {
-                    tenant_id,
-                    timeline_id: None,
-                    metric: SYNTHETIC_STORAGE_SIZE,
-                },
-                tenant_synthetic_size,
-            ));
-        }
+        current_metrics.push((
+            PageserverConsumptionMetricsKey {
+                tenant_id,
+                timeline_id: None,
+                metric: SYNTHETIC_STORAGE_SIZE,
+            },
+            tenant_synthetic_size,
+        ));
     }
 
     // Filter metrics, unless we want to send all metrics, including cached ones.

pageserver/src/disk_usage_eviction_task.rs

@@ -110,6 +110,7 @@ pub fn launch_disk_usage_global_eviction_task(
 
             disk_usage_eviction_task(&state, task_config, storage, &conf.tenants_path(), cancel)
                 .await;
+            info!("disk usage based eviction task finishing");
             Ok(())
         },
     );
@@ -125,16 +126,13 @@ async fn disk_usage_eviction_task(
     tenants_dir: &Path,
     cancel: CancellationToken,
 ) {
-    scopeguard::defer! {
-        info!("disk usage based eviction task finishing");
-    };
-
     use crate::tenant::tasks::random_init_delay;
     {
         if random_init_delay(task_config.period, &cancel)
             .await
             .is_err()
         {
+            info!("shutting down");
             return;
         }
     }
@@ -169,6 +167,7 @@ async fn disk_usage_eviction_task(
         tokio::select! {
             _ = tokio::time::sleep_until(sleep_until) => {},
             _ = cancel.cancelled() => {
+                info!("shutting down");
                 break
             }
         }
@@ -315,7 +314,7 @@ pub async fn disk_usage_eviction_task_iteration_impl<U: Usage>(
             partition,
             candidate.layer.get_tenant_id(),
             candidate.layer.get_timeline_id(),
-            candidate.layer,
+            candidate.layer.filename().file_name(),
         );
     }
 

pageserver/src/metrics.rs

@@ -1,9 +1,9 @@
 use metrics::metric_vec_duration::DurationResultObserver;
 use metrics::{
     register_counter_vec, register_histogram, register_histogram_vec, register_int_counter,
-    register_int_counter_vec, register_int_gauge, register_int_gauge_vec, register_uint_gauge,
-    register_uint_gauge_vec, Counter, CounterVec, Histogram, HistogramVec, IntCounter,
-    IntCounterVec, IntGauge, IntGaugeVec, UIntGauge, UIntGaugeVec,
+    register_int_counter_vec, register_int_gauge, register_int_gauge_vec, register_uint_gauge_vec,
+    Counter, CounterVec, Histogram, HistogramVec, IntCounter, IntCounterVec, IntGauge, IntGaugeVec,
+    UIntGauge, UIntGaugeVec,
 };
 use once_cell::sync::Lazy;
 use pageserver_api::models::TenantState;
@@ -130,122 +130,6 @@ pub static MATERIALIZED_PAGE_CACHE_HIT: Lazy<IntCounter> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
-pub struct PageCacheMetrics {
-    pub read_accesses_materialized_page: IntCounter,
-    pub read_accesses_ephemeral: IntCounter,
-    pub read_accesses_immutable: IntCounter,
-
-    pub read_hits_ephemeral: IntCounter,
-    pub read_hits_immutable: IntCounter,
-    pub read_hits_materialized_page_exact: IntCounter,
-    pub read_hits_materialized_page_older_lsn: IntCounter,
-}
-
-static PAGE_CACHE_READ_HITS: Lazy<IntCounterVec> = Lazy::new(|| {
-    register_int_counter_vec!(
-        "pageserver_page_cache_read_hits_total",
-        "Number of read accesses to the page cache that hit",
-        &["key_kind", "hit_kind"]
-    )
-    .expect("failed to define a metric")
-});
-
-static PAGE_CACHE_READ_ACCESSES: Lazy<IntCounterVec> = Lazy::new(|| {
-    register_int_counter_vec!(
-        "pageserver_page_cache_read_accesses_total",
-        "Number of read accesses to the page cache",
-        &["key_kind"]
-    )
-    .expect("failed to define a metric")
-});
-
-pub static PAGE_CACHE: Lazy<PageCacheMetrics> = Lazy::new(|| PageCacheMetrics {
-    read_accesses_materialized_page: {
-        PAGE_CACHE_READ_ACCESSES
-            .get_metric_with_label_values(&["materialized_page"])
-            .unwrap()
-    },
-
-    read_accesses_ephemeral: {
-        PAGE_CACHE_READ_ACCESSES
-            .get_metric_with_label_values(&["ephemeral"])
-            .unwrap()
-    },
-
-    read_accesses_immutable: {
-        PAGE_CACHE_READ_ACCESSES
-            .get_metric_with_label_values(&["immutable"])
-            .unwrap()
-    },
-
-    read_hits_ephemeral: {
-        PAGE_CACHE_READ_HITS
-            .get_metric_with_label_values(&["ephemeral", "-"])
-            .unwrap()
-    },
-
-    read_hits_immutable: {
-        PAGE_CACHE_READ_HITS
-            .get_metric_with_label_values(&["immutable", "-"])
-            .unwrap()
-    },
-
-    read_hits_materialized_page_exact: {
-        PAGE_CACHE_READ_HITS
-            .get_metric_with_label_values(&["materialized_page", "exact"])
-            .unwrap()
-    },
-
-    read_hits_materialized_page_older_lsn: {
-        PAGE_CACHE_READ_HITS
-            .get_metric_with_label_values(&["materialized_page", "older_lsn"])
-            .unwrap()
-    },
-});
-
-pub struct PageCacheSizeMetrics {
-    pub max_bytes: UIntGauge,
-
-    pub current_bytes_ephemeral: UIntGauge,
-    pub current_bytes_immutable: UIntGauge,
-    pub current_bytes_materialized_page: UIntGauge,
-}
-
-static PAGE_CACHE_SIZE_CURRENT_BYTES: Lazy<UIntGaugeVec> = Lazy::new(|| {
-    register_uint_gauge_vec!(
-        "pageserver_page_cache_size_current_bytes",
-        "Current size of the page cache in bytes, by key kind",
-        &["key_kind"]
-    )
-    .expect("failed to define a metric")
-});
-
-pub static PAGE_CACHE_SIZE: Lazy<PageCacheSizeMetrics> = Lazy::new(|| PageCacheSizeMetrics {
-    max_bytes: {
-        register_uint_gauge!(
-            "pageserver_page_cache_size_max_bytes",
-            "Maximum size of the page cache in bytes"
-        )
-        .expect("failed to define a metric")
-    },
-
-    current_bytes_ephemeral: {
-        PAGE_CACHE_SIZE_CURRENT_BYTES
-            .get_metric_with_label_values(&["ephemeral"])
-            .unwrap()
-    },
-    current_bytes_immutable: {
-        PAGE_CACHE_SIZE_CURRENT_BYTES
-            .get_metric_with_label_values(&["immutable"])
-            .unwrap()
-    },
-    current_bytes_materialized_page: {
-        PAGE_CACHE_SIZE_CURRENT_BYTES
-            .get_metric_with_label_values(&["materialized_page"])
-            .unwrap()
-    },
-});
-
 static WAIT_LSN_TIME: Lazy<HistogramVec> = Lazy::new(|| {
     register_histogram_vec!(
         "pageserver_wait_lsn_seconds",
@@ -1084,6 +968,7 @@ impl RemoteTimelineClientMetrics {
         op_kind: &RemoteOpKind,
         status: &'static str,
     ) -> Histogram {
+        // XXX would be nice to have an upgradable RwLock
         let mut guard = self.remote_operation_time.lock().unwrap();
         let key = (file_kind.as_str(), op_kind.as_str(), status);
         let metric = guard.entry(key).or_insert_with(move || {
@@ -1105,6 +990,7 @@ impl RemoteTimelineClientMetrics {
         file_kind: &RemoteOpFileKind,
         op_kind: &RemoteOpKind,
     ) -> IntGauge {
+        // XXX would be nice to have an upgradable RwLock
         let mut guard = self.calls_unfinished_gauge.lock().unwrap();
         let key = (file_kind.as_str(), op_kind.as_str());
         let metric = guard.entry(key).or_insert_with(move || {
@@ -1125,6 +1011,7 @@ impl RemoteTimelineClientMetrics {
         file_kind: &RemoteOpFileKind,
         op_kind: &RemoteOpKind,
     ) -> Histogram {
+        // XXX would be nice to have an upgradable RwLock
         let mut guard = self.calls_started_hist.lock().unwrap();
         let key = (file_kind.as_str(), op_kind.as_str());
         let metric = guard.entry(key).or_insert_with(move || {
@@ -1145,6 +1032,7 @@ impl RemoteTimelineClientMetrics {
         file_kind: &RemoteOpFileKind,
         op_kind: &RemoteOpKind,
     ) -> IntCounter {
+        // XXX would be nice to have an upgradable RwLock
         let mut guard = self.bytes_started_counter.lock().unwrap();
         let key = (file_kind.as_str(), op_kind.as_str());
         let metric = guard.entry(key).or_insert_with(move || {
@@ -1165,6 +1053,7 @@ impl RemoteTimelineClientMetrics {
         file_kind: &RemoteOpFileKind,
         op_kind: &RemoteOpKind,
     ) -> IntCounter {
+        // XXX would be nice to have an upgradable RwLock
         let mut guard = self.bytes_finished_counter.lock().unwrap();
         let key = (file_kind.as_str(), op_kind.as_str());
         let metric = guard.entry(key).or_insert_with(move || {

pageserver/src/page_cache.rs

@@ -53,8 +53,8 @@ use utils::{
     lsn::Lsn,
 };
 
+use crate::repository::Key;
 use crate::tenant::writeback_ephemeral_file;
-use crate::{metrics::PageCacheSizeMetrics, repository::Key};
 
 static PAGE_CACHE: OnceCell<PageCache> = OnceCell::new();
 const TEST_PAGE_CACHE_SIZE: usize = 50;
@@ -187,8 +187,6 @@ pub struct PageCache {
     /// Index of the next candidate to evict, for the Clock replacement algorithm.
     /// This is interpreted modulo the page cache size.
     next_evict_slot: AtomicUsize,
-
-    size_metrics: &'static PageCacheSizeMetrics,
 }
 
 ///
@@ -315,10 +313,6 @@ impl PageCache {
         key: &Key,
         lsn: Lsn,
     ) -> Option<(Lsn, PageReadGuard)> {
-        crate::metrics::PAGE_CACHE
-            .read_accesses_materialized_page
-            .inc();
-
         let mut cache_key = CacheKey::MaterializedPage {
             hash_key: MaterializedPageHashKey {
                 tenant_id,
@@ -329,21 +323,8 @@ impl PageCache {
         };
 
         if let Some(guard) = self.try_lock_for_read(&mut cache_key) {
-            if let CacheKey::MaterializedPage {
-                hash_key: _,
-                lsn: available_lsn,
-            } = cache_key
-            {
-                if available_lsn == lsn {
-                    crate::metrics::PAGE_CACHE
-                        .read_hits_materialized_page_exact
-                        .inc();
-                } else {
-                    crate::metrics::PAGE_CACHE
-                        .read_hits_materialized_page_older_lsn
-                        .inc();
-                }
-                Some((available_lsn, guard))
+            if let CacheKey::MaterializedPage { hash_key: _, lsn } = cache_key {
+                Some((lsn, guard))
             } else {
                 panic!("unexpected key type in slot");
             }
@@ -518,31 +499,11 @@ impl PageCache {
     /// ```
     ///
     fn lock_for_read(&self, cache_key: &mut CacheKey) -> anyhow::Result<ReadBufResult> {
-        let (read_access, hit) = match cache_key {
-            CacheKey::MaterializedPage { .. } => {
-                unreachable!("Materialized pages use lookup_materialized_page")
-            }
-            CacheKey::EphemeralPage { .. } => (
-                &crate::metrics::PAGE_CACHE.read_accesses_ephemeral,
-                &crate::metrics::PAGE_CACHE.read_hits_ephemeral,
-            ),
-            CacheKey::ImmutableFilePage { .. } => (
-                &crate::metrics::PAGE_CACHE.read_accesses_immutable,
-                &crate::metrics::PAGE_CACHE.read_hits_immutable,
-            ),
-        };
-        read_access.inc();
-
-        let mut is_first_iteration = true;
         loop {
             // First check if the key already exists in the cache.
             if let Some(read_guard) = self.try_lock_for_read(cache_key) {
-                if is_first_iteration {
-                    hit.inc();
-                }
                 return Ok(ReadBufResult::Found(read_guard));
             }
-            is_first_iteration = false;
 
             // Not found. Find a victim buffer
             let (slot_idx, mut inner) =
@@ -720,9 +681,6 @@ impl PageCache {
 
                     if let Ok(version_idx) = versions.binary_search_by_key(old_lsn, |v| v.lsn) {
                         versions.remove(version_idx);
-                        self.size_metrics
-                            .current_bytes_materialized_page
-                            .sub_page_sz(1);
                         if versions.is_empty() {
                             old_entry.remove_entry();
                         }
@@ -735,13 +693,11 @@ impl PageCache {
                 let mut map = self.ephemeral_page_map.write().unwrap();
                 map.remove(&(*file_id, *blkno))
                     .expect("could not find old key in mapping");
-                self.size_metrics.current_bytes_ephemeral.sub_page_sz(1);
             }
             CacheKey::ImmutableFilePage { file_id, blkno } => {
                 let mut map = self.immutable_page_map.write().unwrap();
                 map.remove(&(*file_id, *blkno))
                     .expect("could not find old key in mapping");
-                self.size_metrics.current_bytes_immutable.sub_page_sz(1);
             }
         }
     }
@@ -769,9 +725,6 @@ impl PageCache {
                                 slot_idx,
                             },
                         );
-                        self.size_metrics
-                            .current_bytes_materialized_page
-                            .add_page_sz(1);
                         None
                     }
                 }
@@ -782,7 +735,6 @@ impl PageCache {
                     Entry::Occupied(entry) => Some(*entry.get()),
                     Entry::Vacant(entry) => {
                         entry.insert(slot_idx);
-                        self.size_metrics.current_bytes_ephemeral.add_page_sz(1);
                         None
                     }
                 }
@@ -793,7 +745,6 @@ impl PageCache {
                     Entry::Occupied(entry) => Some(*entry.get()),
                     Entry::Vacant(entry) => {
                         entry.insert(slot_idx);
-                        self.size_metrics.current_bytes_immutable.add_page_sz(1);
                         None
                     }
                 }
@@ -893,12 +844,6 @@ impl PageCache {
 
         let page_buffer = Box::leak(vec![0u8; num_pages * PAGE_SZ].into_boxed_slice());
 
-        let size_metrics = &crate::metrics::PAGE_CACHE_SIZE;
-        size_metrics.max_bytes.set_page_sz(num_pages);
-        size_metrics.current_bytes_ephemeral.set_page_sz(0);
-        size_metrics.current_bytes_immutable.set_page_sz(0);
-        size_metrics.current_bytes_materialized_page.set_page_sz(0);
-
         let slots = page_buffer
             .chunks_exact_mut(PAGE_SZ)
             .map(|chunk| {
@@ -921,30 +866,6 @@ impl PageCache {
             immutable_page_map: Default::default(),
             slots,
             next_evict_slot: AtomicUsize::new(0),
-            size_metrics,
         }
     }
 }
-
-trait PageSzBytesMetric {
-    fn set_page_sz(&self, count: usize);
-    fn add_page_sz(&self, count: usize);
-    fn sub_page_sz(&self, count: usize);
-}
-
-#[inline(always)]
-fn count_times_page_sz(count: usize) -> u64 {
-    u64::try_from(count).unwrap() * u64::try_from(PAGE_SZ).unwrap()
-}
-
-impl PageSzBytesMetric for metrics::UIntGauge {
-    fn set_page_sz(&self, count: usize) {
-        self.set(count_times_page_sz(count));
-    }
-    fn add_page_sz(&self, count: usize) {
-        self.add(count_times_page_sz(count));
-    }
-    fn sub_page_sz(&self, count: usize) {
-        self.sub(count_times_page_sz(count));
-    }
-}

pageserver/src/tenant.rs

@@ -11,7 +11,7 @@
 //! parent timeline, and the last LSN that has been written to disk.
 //!
 
-use anyhow::{bail, Context};
+use anyhow::{bail, ensure, Context};
 use futures::FutureExt;
 use pageserver_api::models::TimelineState;
 use remote_storage::DownloadError;
@@ -49,8 +49,6 @@ use std::time::{Duration, Instant};
 use self::config::TenantConf;
 use self::metadata::TimelineMetadata;
 use self::remote_timeline_client::RemoteTimelineClient;
-use self::timeline::uninit::TimelineUninitMark;
-use self::timeline::uninit::UninitializedTimeline;
 use self::timeline::EvictionTaskTenantState;
 use crate::config::PageServerConf;
 use crate::context::{DownloadBehavior, RequestContext};
@@ -70,7 +68,6 @@ use crate::tenant::storage_layer::ImageLayer;
 use crate::tenant::storage_layer::Layer;
 use crate::InitializationOrder;
 
-use crate::tenant::timeline::uninit::cleanup_timeline_directory;
 use crate::virtual_file::VirtualFile;
 use crate::walredo::PostgresRedoManager;
 use crate::walredo::WalRedoManager;
@@ -90,7 +87,6 @@ pub mod disk_btree;
 pub(crate) mod ephemeral_file;
 pub mod layer_map;
 pub mod manifest;
-mod span;
 
 pub mod metadata;
 mod par_fsync;
@@ -106,7 +102,7 @@ mod timeline;
 
 pub mod size;
 
-pub(crate) use timeline::span::debug_assert_current_span_has_tenant_and_timeline_id;
+pub(crate) use timeline::debug_assert_current_span_has_tenant_and_timeline_id;
 pub use timeline::{
     LocalLayerInfoForDiskUsageEviction, LogicalSizeCalculationCause, PageReconstructError, Timeline,
 };
@@ -165,6 +161,200 @@ pub struct Tenant {
     eviction_task_tenant_state: tokio::sync::Mutex<EvictionTaskTenantState>,
 }
 
+/// A timeline with some of its files on disk, being initialized.
+/// This struct ensures the atomicity of the timeline init: it's either properly created and inserted into pageserver's memory, or
+/// its local files are removed. In the worst case of a crash, an uninit mark file is left behind, which causes the directory
+/// to be removed on next restart.
+///
+/// The caller is responsible for proper timeline data filling before the final init.
+#[must_use]
+pub struct UninitializedTimeline<'t> {
+    owning_tenant: &'t Tenant,
+    timeline_id: TimelineId,
+    raw_timeline: Option<(Arc<Timeline>, TimelineUninitMark)>,
+}
+
+/// An uninit mark file, created along the timeline dir to ensure the timeline either gets fully initialized and loaded into pageserver's memory,
+/// or gets removed eventually.
+///
+/// XXX: it's important to create it near the timeline dir, not inside it to ensure timeline dir gets removed first.
+#[must_use]
+struct TimelineUninitMark {
+    uninit_mark_deleted: bool,
+    uninit_mark_path: PathBuf,
+    timeline_path: PathBuf,
+}
+
+impl UninitializedTimeline<'_> {
+    /// Finish timeline creation: insert it into the Tenant's timelines map and remove the
+    /// uninit mark file.
+    ///
+    /// This function launches the flush loop if not already done.
+    ///
+    /// The caller is responsible for activating the timeline (function `.activate()`).
+    fn finish_creation(mut self) -> anyhow::Result<Arc<Timeline>> {
+        let timeline_id = self.timeline_id;
+        let tenant_id = self.owning_tenant.tenant_id;
+
+        let (new_timeline, uninit_mark) = self.raw_timeline.take().with_context(|| {
+            format!("No timeline for initalization found for {tenant_id}/{timeline_id}")
+        })?;
+
+        // Check that the caller initialized disk_consistent_lsn
+        let new_disk_consistent_lsn = new_timeline.get_disk_consistent_lsn();
+        ensure!(
+            new_disk_consistent_lsn.is_valid(),
+            "new timeline {tenant_id}/{timeline_id} has invalid disk_consistent_lsn"
+        );
+
+        let mut timelines = self.owning_tenant.timelines.lock().unwrap();
+        match timelines.entry(timeline_id) {
+            Entry::Occupied(_) => anyhow::bail!(
+                "Found freshly initialized timeline {tenant_id}/{timeline_id} in the tenant map"
+            ),
+            Entry::Vacant(v) => {
+                uninit_mark.remove_uninit_mark().with_context(|| {
+                    format!(
+                        "Failed to remove uninit mark file for timeline {tenant_id}/{timeline_id}"
+                    )
+                })?;
+                v.insert(Arc::clone(&new_timeline));
+
+                new_timeline.maybe_spawn_flush_loop();
+            }
+        }
+
+        Ok(new_timeline)
+    }
+
+    /// Prepares timeline data by loading it from the basebackup archive.
+    pub async fn import_basebackup_from_tar(
+        self,
+        copyin_read: &mut (impl tokio::io::AsyncRead + Send + Sync + Unpin),
+        base_lsn: Lsn,
+        broker_client: storage_broker::BrokerClientChannel,
+        ctx: &RequestContext,
+    ) -> anyhow::Result<Arc<Timeline>> {
+        let raw_timeline = self.raw_timeline()?;
+
+        import_datadir::import_basebackup_from_tar(raw_timeline, copyin_read, base_lsn, ctx)
+            .await
+            .context("Failed to import basebackup")?;
+
+        // Flush the new layer files to disk, before we make the timeline as available to
+        // the outside world.
+        //
+        // Flush loop needs to be spawned in order to be able to flush.
+        raw_timeline.maybe_spawn_flush_loop();
+
+        fail::fail_point!("before-checkpoint-new-timeline", |_| {
+            bail!("failpoint before-checkpoint-new-timeline");
+        });
+
+        raw_timeline
+            .freeze_and_flush()
+            .await
+            .context("Failed to flush after basebackup import")?;
+
+        // All the data has been imported. Insert the Timeline into the tenant's timelines
+        // map and remove the uninit mark file.
+        let tl = self.finish_creation()?;
+        tl.activate(broker_client, None, ctx);
+        Ok(tl)
+    }
+
+    fn raw_timeline(&self) -> anyhow::Result<&Arc<Timeline>> {
+        Ok(&self
+            .raw_timeline
+            .as_ref()
+            .with_context(|| {
+                format!(
+                    "No raw timeline {}/{} found",
+                    self.owning_tenant.tenant_id, self.timeline_id
+                )
+            })?
+            .0)
+    }
+}
+
+impl Drop for UninitializedTimeline<'_> {
+    fn drop(&mut self) {
+        if let Some((_, uninit_mark)) = self.raw_timeline.take() {
+            let _entered = info_span!("drop_uninitialized_timeline", tenant = %self.owning_tenant.tenant_id, timeline = %self.timeline_id).entered();
+            error!("Timeline got dropped without initializing, cleaning its files");
+            cleanup_timeline_directory(uninit_mark);
+        }
+    }
+}
+
+fn cleanup_timeline_directory(uninit_mark: TimelineUninitMark) {
+    let timeline_path = &uninit_mark.timeline_path;
+    match ignore_absent_files(|| fs::remove_dir_all(timeline_path)) {
+        Ok(()) => {
+            info!("Timeline dir {timeline_path:?} removed successfully, removing the uninit mark")
+        }
+        Err(e) => {
+            error!("Failed to clean up uninitialized timeline directory {timeline_path:?}: {e:?}")
+        }
+    }
+    drop(uninit_mark); // mark handles its deletion on drop, gets retained if timeline dir exists
+}
+
+impl TimelineUninitMark {
+    fn new(uninit_mark_path: PathBuf, timeline_path: PathBuf) -> Self {
+        Self {
+            uninit_mark_deleted: false,
+            uninit_mark_path,
+            timeline_path,
+        }
+    }
+
+    fn remove_uninit_mark(mut self) -> anyhow::Result<()> {
+        if !self.uninit_mark_deleted {
+            self.delete_mark_file_if_present()?;
+        }
+
+        Ok(())
+    }
+
+    fn delete_mark_file_if_present(&mut self) -> anyhow::Result<()> {
+        let uninit_mark_file = &self.uninit_mark_path;
+        let uninit_mark_parent = uninit_mark_file
+            .parent()
+            .with_context(|| format!("Uninit mark file {uninit_mark_file:?} has no parent"))?;
+        ignore_absent_files(|| fs::remove_file(uninit_mark_file)).with_context(|| {
+            format!("Failed to remove uninit mark file at path {uninit_mark_file:?}")
+        })?;
+        crashsafe::fsync(uninit_mark_parent).context("Failed to fsync uninit mark parent")?;
+        self.uninit_mark_deleted = true;
+
+        Ok(())
+    }
+}
+
+impl Drop for TimelineUninitMark {
+    fn drop(&mut self) {
+        if !self.uninit_mark_deleted {
+            if self.timeline_path.exists() {
+                error!(
+                    "Uninit mark {} is not removed, timeline {} stays uninitialized",
+                    self.uninit_mark_path.display(),
+                    self.timeline_path.display()
+                )
+            } else {
+                // unblock later timeline creation attempts
+                warn!(
+                    "Removing intermediate uninit mark file {}",
+                    self.uninit_mark_path.display()
+                );
+                if let Err(e) = self.delete_mark_file_if_present() {
+                    error!("Failed to remove the uninit mark file: {e}")
+                }
+            }
+        }
+    }
+}
+
 // We should not blindly overwrite local metadata with remote one.
 // For example, consider the following case:
 //     Image layer is flushed to disk as a new delta layer, we update local metadata and start upload task but after that
@@ -505,7 +695,7 @@ impl Tenant {
     /// No background tasks are started as part of this routine.
     ///
     async fn attach(self: &Arc<Tenant>, ctx: &RequestContext) -> anyhow::Result<()> {
-        span::debug_assert_current_span_has_tenant_id();
+        debug_assert_current_span_has_tenant_id();
 
         let marker_file = self.conf.tenant_attaching_mark_file_path(&self.tenant_id);
         if !tokio::fs::try_exists(&marker_file)
@@ -643,7 +833,7 @@ impl Tenant {
         remote_client: RemoteTimelineClient,
         ctx: &RequestContext,
     ) -> anyhow::Result<()> {
-        span::debug_assert_current_span_has_tenant_id();
+        debug_assert_current_span_has_tenant_id();
 
         info!("downloading index file for timeline {}", timeline_id);
         tokio::fs::create_dir_all(self.conf.timeline_path(&timeline_id, &self.tenant_id))
@@ -722,7 +912,7 @@ impl Tenant {
         init_order: Option<InitializationOrder>,
         ctx: &RequestContext,
     ) -> Arc<Tenant> {
-        span::debug_assert_current_span_has_tenant_id();
+        debug_assert_current_span_has_tenant_id();
 
         let tenant_conf = match Self::load_tenant_config(conf, tenant_id) {
             Ok(conf) => conf,
@@ -908,7 +1098,7 @@ impl Tenant {
         init_order: Option<&InitializationOrder>,
         ctx: &RequestContext,
     ) -> anyhow::Result<()> {
-        span::debug_assert_current_span_has_tenant_id();
+        debug_assert_current_span_has_tenant_id();
 
         debug!("loading tenant task");
 
@@ -954,7 +1144,7 @@ impl Tenant {
         init_order: Option<&InitializationOrder>,
         ctx: &RequestContext,
     ) -> anyhow::Result<()> {
-        span::debug_assert_current_span_has_tenant_id();
+        debug_assert_current_span_has_tenant_id();
 
         let remote_client = self.remote_storage.as_ref().map(|remote_storage| {
             RemoteTimelineClient::new(
@@ -1545,7 +1735,7 @@ impl Tenant {
         timeline_id: TimelineId,
         _ctx: &RequestContext,
     ) -> Result<(), DeleteTimelineError> {
-        debug_assert_current_span_has_tenant_and_timeline_id();
+        timeline::debug_assert_current_span_has_tenant_and_timeline_id();
 
         // Transition the timeline into TimelineState::Stopping.
         // This should prevent new operations from starting.
@@ -1709,7 +1899,7 @@ impl Tenant {
         background_jobs_can_start: Option<&completion::Barrier>,
         ctx: &RequestContext,
     ) {
-        span::debug_assert_current_span_has_tenant_id();
+        debug_assert_current_span_has_tenant_id();
 
         let mut activating = false;
         self.state.send_modify(|current_state| {
@@ -1780,7 +1970,7 @@ impl Tenant {
     ///
     /// This will attempt to shutdown even if tenant is broken.
     pub(crate) async fn shutdown(&self, freeze_and_flush: bool) -> Result<(), ShutdownError> {
-        span::debug_assert_current_span_has_tenant_id();
+        debug_assert_current_span_has_tenant_id();
         // Set tenant (and its timlines) to Stoppping state.
         //
         // Since we can only transition into Stopping state after activation is complete,
@@ -2822,11 +3012,11 @@ impl Tenant {
 
         debug!("Successfully created initial files for timeline {tenant_id}/{new_timeline_id}");
 
-        Ok(UninitializedTimeline::new(
-            self,
-            new_timeline_id,
-            Some((timeline_struct, uninit_mark)),
-        ))
+        Ok(UninitializedTimeline {
+            owning_tenant: self,
+            timeline_id: new_timeline_id,
+            raw_timeline: Some((timeline_struct, uninit_mark)),
+        })
     }
 
     fn create_timeline_files(
@@ -4381,3 +4571,28 @@ mod tests {
         Ok(())
     }
 }
+
+#[cfg(not(debug_assertions))]
+#[inline]
+pub(crate) fn debug_assert_current_span_has_tenant_id() {}
+
+#[cfg(debug_assertions)]
+pub static TENANT_ID_EXTRACTOR: once_cell::sync::Lazy<
+    utils::tracing_span_assert::MultiNameExtractor<2>,
+> = once_cell::sync::Lazy::new(|| {
+    utils::tracing_span_assert::MultiNameExtractor::new("TenantId", ["tenant_id", "tenant"])
+});
+
+#[cfg(debug_assertions)]
+#[inline]
+pub(crate) fn debug_assert_current_span_has_tenant_id() {
+    use utils::tracing_span_assert;
+
+    match tracing_span_assert::check_fields_present([&*TENANT_ID_EXTRACTOR]) {
+        Ok(()) => (),
+        Err(missing) => panic!(
+            "missing extractors: {:?}",
+            missing.into_iter().map(|e| e.name()).collect::<Vec<_>>()
+        ),
+    }
+}

pageserver/src/tenant/remote_timeline_client.rs

@@ -608,7 +608,10 @@ impl RemoteTimelineClient {
         self.calls_unfinished_metric_begin(&op);
         upload_queue.queued_operations.push_back(op);
 
-        info!("scheduled layer file upload {layer_file_name}");
+        info!(
+            "scheduled layer file upload {}",
+            layer_file_name.file_name()
+        );
 
         // Launch the task immediately, if possible
         self.launch_queued_tasks(upload_queue);
@@ -661,7 +664,7 @@ impl RemoteTimelineClient {
                 });
                 self.calls_unfinished_metric_begin(&op);
                 upload_queue.queued_operations.push_back(op);
-                info!("scheduled layer file deletion {name}");
+                info!("scheduled layer file deletion {}", name.file_name());
             }
 
             // Launch the tasks immediately, if possible
@@ -825,7 +828,7 @@ impl RemoteTimelineClient {
                     .queued_operations
                     .push_back(op);
 
-                info!("scheduled layer file deletion {name}");
+                info!("scheduled layer file deletion {}", name.file_name());
                 deletions_queued += 1;
             }
 

pageserver/src/tenant/remote_timeline_client/download.rs

@@ -16,7 +16,7 @@ use tracing::{info, warn};
 
 use crate::config::PageServerConf;
 use crate::tenant::storage_layer::LayerFileName;
-use crate::tenant::timeline::span::debug_assert_current_span_has_tenant_and_timeline_id;
+use crate::tenant::timeline::debug_assert_current_span_has_tenant_and_timeline_id;
 use crate::{exponential_backoff, DEFAULT_BASE_BACKOFF_SECONDS, DEFAULT_MAX_BACKOFF_SECONDS};
 use remote_storage::{DownloadError, GenericRemoteStorage};
 use utils::crashsafe::path_with_suffix_extension;

pageserver/src/tenant/span.rs

@@ -1,20 +0,0 @@
-#[cfg(debug_assertions)]
-use utils::tracing_span_assert::{check_fields_present, MultiNameExtractor};
-
-#[cfg(not(debug_assertions))]
-pub(crate) fn debug_assert_current_span_has_tenant_id() {}
-
-#[cfg(debug_assertions)]
-pub(crate) static TENANT_ID_EXTRACTOR: once_cell::sync::Lazy<MultiNameExtractor<2>> =
-    once_cell::sync::Lazy::new(|| MultiNameExtractor::new("TenantId", ["tenant_id", "tenant"]));
-
-#[cfg(debug_assertions)]
-#[track_caller]
-pub(crate) fn debug_assert_current_span_has_tenant_id() {
-    if let Err(missing) = check_fields_present([&*TENANT_ID_EXTRACTOR]) {
-        panic!(
-            "missing extractors: {:?}",
-            missing.into_iter().map(|e| e.name()).collect::<Vec<_>>()
-        )
-    }
-}

pageserver/src/tenant/storage_layer.rs

@@ -335,7 +335,7 @@ impl LayerAccessStats {
 /// All layers should implement a minimal `std::fmt::Debug` without tenant or
 /// timeline names, because those are known in the context of which the layers
 /// are used in (timeline).
-pub trait Layer: std::fmt::Debug + std::fmt::Display + Send + Sync {
+pub trait Layer: std::fmt::Debug + Send + Sync {
     /// Range of keys that this layer covers
     fn get_key_range(&self) -> Range<Key>;
 
@@ -373,6 +373,9 @@ pub trait Layer: std::fmt::Debug + std::fmt::Display + Send + Sync {
         ctx: &RequestContext,
     ) -> Result<ValueReconstructResult>;
 
+    /// A short ID string that uniquely identifies the given layer within a [`LayerMap`].
+    fn short_id(&self) -> String;
+
     /// Dump summary of the contents of the layer to stdout
     fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()>;
 }
@@ -509,12 +512,10 @@ pub mod tests {
         fn is_incremental(&self) -> bool {
             self.layer_desc().is_incremental
         }
-    }
 
-    /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
-    impl std::fmt::Display for LayerDescriptor {
-        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-            write!(f, "{}", self.layer_desc().short_id())
+        /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
+        fn short_id(&self) -> String {
+            self.layer_desc().short_id()
         }
     }
 

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -222,14 +222,13 @@ impl Layer for DeltaLayer {
     /// debugging function to print out the contents of the layer
     fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
         println!(
-            "----- delta layer for ten {} tli {} keys {}-{} lsn {}-{} size {} ----",
+            "----- delta layer for ten {} tli {} keys {}-{} lsn {}-{} ----",
             self.desc.tenant_id,
             self.desc.timeline_id,
             self.desc.key_range.start,
             self.desc.key_range.end,
             self.desc.lsn_range.start,
-            self.desc.lsn_range.end,
-            self.desc.file_size,
+            self.desc.lsn_range.end
         );
 
         if !verbose {
@@ -395,11 +394,10 @@ impl Layer for DeltaLayer {
     fn is_incremental(&self) -> bool {
         self.layer_desc().is_incremental
     }
-}
-/// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
-impl std::fmt::Display for DeltaLayer {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.layer_desc().short_id())
+
+    /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
+    fn short_id(&self) -> String {
+        self.layer_desc().short_id()
     }
 }
 

pageserver/src/tenant/storage_layer/filename.rs

@@ -210,15 +210,9 @@ pub enum LayerFileName {
 
 impl LayerFileName {
     pub fn file_name(&self) -> String {
-        self.to_string()
-    }
-}
-
-impl fmt::Display for LayerFileName {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
-            Self::Image(fname) => write!(f, "{fname}"),
-            Self::Delta(fname) => write!(f, "{fname}"),
+            Self::Image(fname) => fname.to_string(),
+            Self::Delta(fname) => fname.to_string(),
         }
     }
 }

pageserver/src/tenant/storage_layer/image_layer.rs

@@ -153,14 +153,12 @@ impl Layer for ImageLayer {
     /// debugging function to print out the contents of the layer
     fn dump(&self, verbose: bool, ctx: &RequestContext) -> Result<()> {
         println!(
-            "----- image layer for ten {} tli {} key {}-{} at {} is_incremental {} size {} ----",
+            "----- image layer for ten {} tli {} key {}-{} at {} ----",
             self.desc.tenant_id,
             self.desc.timeline_id,
             self.desc.key_range.start,
             self.desc.key_range.end,
-            self.lsn,
-            self.desc.is_incremental,
-            self.desc.file_size
+            self.lsn
         );
 
         if !verbose {
@@ -232,12 +230,10 @@ impl Layer for ImageLayer {
     fn is_incremental(&self) -> bool {
         self.layer_desc().is_incremental
     }
-}
 
-/// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
-impl std::fmt::Display for ImageLayer {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.layer_desc().short_id())
+    /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
+    fn short_id(&self) -> String {
+        self.layer_desc().short_id()
     }
 }
 

pageserver/src/tenant/storage_layer/inmemory_layer.rs

@@ -131,6 +131,13 @@ impl Layer for InMemoryLayer {
         true
     }
 
+    fn short_id(&self) -> String {
+        let inner = self.inner.read().unwrap();
+
+        let end_lsn = inner.end_lsn.unwrap_or(Lsn(u64::MAX));
+        format!("inmem-{:016X}-{:016X}", self.start_lsn.0, end_lsn.0)
+    }
+
     /// debugging function to print out the contents of the layer
     fn dump(&self, verbose: bool, _ctx: &RequestContext) -> Result<()> {
         let inner = self.inner.read().unwrap();
@@ -233,15 +240,6 @@ impl Layer for InMemoryLayer {
     }
 }
 
-impl std::fmt::Display for InMemoryLayer {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let inner = self.inner.read().unwrap();
-
-        let end_lsn = inner.end_lsn.unwrap_or(Lsn(u64::MAX));
-        write!(f, "inmem-{:016X}-{:016X}", self.start_lsn.0, end_lsn.0)
-    }
-}
-
 impl InMemoryLayer {
     ///
     /// Get layer size on the disk

pageserver/src/tenant/storage_layer/layer_desc.rs

@@ -1,5 +1,4 @@
 use anyhow::Result;
-use core::fmt::Display;
 use std::ops::Range;
 use utils::{
     id::{TenantId, TimelineId},
@@ -49,8 +48,8 @@ impl PersistentLayerDesc {
         }
     }
 
-    pub fn short_id(&self) -> impl Display {
-        self.filename()
+    pub fn short_id(&self) -> String {
+        self.filename().file_name()
     }
 
     #[cfg(test)]
@@ -174,16 +173,13 @@ impl PersistentLayerDesc {
 
     pub fn dump(&self, _verbose: bool, _ctx: &RequestContext) -> Result<()> {
         println!(
-            "----- layer for ten {} tli {} keys {}-{} lsn {}-{} is_delta {} is_incremental {} size {} ----",
+            "----- layer for ten {} tli {} keys {}-{} lsn {}-{} ----",
             self.tenant_id,
             self.timeline_id,
             self.key_range.start,
             self.key_range.end,
             self.lsn_range.start,
-            self.lsn_range.end,
-            self.is_delta,
-            self.is_incremental,
-            self.file_size,
+            self.lsn_range.end
         );
 
         Ok(())

pageserver/src/tenant/storage_layer/remote_layer.rs

@@ -71,22 +71,22 @@ impl Layer for RemoteLayer {
         _reconstruct_state: &mut ValueReconstructState,
         _ctx: &RequestContext,
     ) -> Result<ValueReconstructResult> {
-        bail!("layer {self} needs to be downloaded");
+        bail!(
+            "layer {} needs to be downloaded",
+            self.filename().file_name()
+        );
     }
 
     /// debugging function to print out the contents of the layer
     fn dump(&self, _verbose: bool, _ctx: &RequestContext) -> Result<()> {
         println!(
-            "----- remote layer for ten {} tli {} keys {}-{} lsn {}-{} is_delta {} is_incremental {} size {} ----",
+            "----- remote layer for ten {} tli {} keys {}-{} lsn {}-{} ----",
             self.desc.tenant_id,
             self.desc.timeline_id,
             self.desc.key_range.start,
             self.desc.key_range.end,
             self.desc.lsn_range.start,
-            self.desc.lsn_range.end,
-            self.desc.is_delta,
-            self.desc.is_incremental,
-            self.desc.file_size,
+            self.desc.lsn_range.end
         );
 
         Ok(())
@@ -106,12 +106,10 @@ impl Layer for RemoteLayer {
     fn is_incremental(&self) -> bool {
         self.layer_desc().is_incremental
     }
-}
 
-/// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
-impl std::fmt::Display for RemoteLayer {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.layer_desc().short_id())
+    /// Boilerplate to implement the Layer trait, always use layer_desc for persistent layers.
+    fn short_id(&self) -> String {
+        self.layer_desc().short_id()
     }
 }
 

pageserver/src/tenant/timeline.rs

@@ -1,9 +1,6 @@
 //!
 
 mod eviction_task;
-mod logical_size;
-pub mod span;
-pub mod uninit;
 mod walreceiver;
 
 use anyhow::{anyhow, bail, ensure, Context, Result};
@@ -11,6 +8,7 @@ use bytes::Bytes;
 use fail::fail_point;
 use futures::StreamExt;
 use itertools::Itertools;
+use once_cell::sync::OnceCell;
 use pageserver_api::models::{
     DownloadRemoteLayersTaskInfo, DownloadRemoteLayersTaskSpawnRequest,
     DownloadRemoteLayersTaskState, LayerMapInfo, LayerResidenceEventReason, LayerResidenceStatus,
@@ -19,7 +17,7 @@ use pageserver_api::models::{
 use remote_storage::GenericRemoteStorage;
 use serde_with::serde_as;
 use storage_broker::BrokerClientChannel;
-use tokio::sync::{oneshot, watch, TryAcquireError};
+use tokio::sync::{oneshot, watch, Semaphore, TryAcquireError};
 use tokio_util::sync::CancellationToken;
 use tracing::*;
 use utils::id::TenantTimelineId;
@@ -30,7 +28,7 @@ use std::fs;
 use std::ops::{Deref, Range};
 use std::path::{Path, PathBuf};
 use std::pin::pin;
-use std::sync::atomic::Ordering as AtomicOrdering;
+use std::sync::atomic::{AtomicI64, Ordering as AtomicOrdering};
 use std::sync::{Arc, Mutex, RwLock, Weak};
 use std::time::{Duration, Instant, SystemTime};
 
@@ -40,7 +38,6 @@ use crate::tenant::storage_layer::{
     DeltaFileName, DeltaLayerWriter, ImageFileName, ImageLayerWriter, InMemoryLayer,
     LayerAccessStats, LayerFileName, RemoteLayer,
 };
-use crate::tenant::timeline::logical_size::CurrentLogicalSize;
 use crate::tenant::{
     ephemeral_file::is_ephemeral_file,
     layer_map::{LayerMap, SearchResult},
@@ -82,7 +79,6 @@ use crate::{is_temporary, task_mgr};
 
 pub(super) use self::eviction_task::EvictionTaskTenantState;
 use self::eviction_task::EvictionTaskTimelineState;
-use self::logical_size::LogicalSize;
 use self::walreceiver::{WalReceiver, WalReceiverConf};
 
 use super::config::TenantConf;
@@ -132,7 +128,7 @@ impl LayerFileManager {
         // A layer's descriptor is present in the LayerMap => the LayerFileManager contains a layer for the descriptor.
         self.0
             .get(&desc.key())
-            .with_context(|| format!("get layer from desc: {}", desc.filename()))
+            .with_context(|| format!("get layer from desc: {}", desc.filename().file_name()))
             .expect("not found")
             .clone()
     }
@@ -369,6 +365,126 @@ pub struct Timeline {
     initial_logical_size_attempt: Mutex<Option<completion::Completion>>,
 }
 
+/// Internal structure to hold all data needed for logical size calculation.
+///
+/// Calculation consists of two stages:
+///
+/// 1. Initial size calculation. That might take a long time, because it requires
+/// reading all layers containing relation sizes at `initial_part_end`.
+///
+/// 2. Collecting an incremental part and adding that to the initial size.
+/// Increments are appended on walreceiver writing new timeline data,
+/// which result in increase or decrease of the logical size.
+struct LogicalSize {
+    /// Size, potentially slow to compute. Calculating this might require reading multiple
+    /// layers, and even ancestor's layers.
+    ///
+    /// NOTE: size at a given LSN is constant, but after a restart we will calculate
+    /// the initial size at a different LSN.
+    initial_logical_size: OnceCell<u64>,
+
+    /// Semaphore to track ongoing calculation of `initial_logical_size`.
+    initial_size_computation: Arc<tokio::sync::Semaphore>,
+
+    /// Latest Lsn that has its size uncalculated, could be absent for freshly created timelines.
+    initial_part_end: Option<Lsn>,
+
+    /// All other size changes after startup, combined together.
+    ///
+    /// Size shouldn't ever be negative, but this is signed for two reasons:
+    ///
+    /// 1. If we initialized the "baseline" size lazily, while we already
+    /// process incoming WAL, the incoming WAL records could decrement the
+    /// variable and temporarily make it negative. (This is just future-proofing;
+    /// the initialization is currently not done lazily.)
+    ///
+    /// 2. If there is a bug and we e.g. forget to increment it in some cases
+    /// when size grows, but remember to decrement it when it shrinks again, the
+    /// variable could go negative. In that case, it seems better to at least
+    /// try to keep tracking it, rather than clamp or overflow it. Note that
+    /// get_current_logical_size() will clamp the returned value to zero if it's
+    /// negative, and log an error. Could set it permanently to zero or some
+    /// special value to indicate "broken" instead, but this will do for now.
+    ///
+    /// Note that we also expose a copy of this value as a prometheus metric,
+    /// see `current_logical_size_gauge`. Use the `update_current_logical_size`
+    /// to modify this, it will also keep the prometheus metric in sync.
+    size_added_after_initial: AtomicI64,
+}
+
+/// Normalized current size, that the data in pageserver occupies.
+#[derive(Debug, Clone, Copy)]
+enum CurrentLogicalSize {
+    /// The size is not yet calculated to the end, this is an intermediate result,
+    /// constructed from walreceiver increments and normalized: logical data could delete some objects, hence be negative,
+    /// yet total logical size cannot be below 0.
+    Approximate(u64),
+    // Fully calculated logical size, only other future walreceiver increments are changing it, and those changes are
+    // available for observation without any calculations.
+    Exact(u64),
+}
+
+impl CurrentLogicalSize {
+    fn size(&self) -> u64 {
+        *match self {
+            Self::Approximate(size) => size,
+            Self::Exact(size) => size,
+        }
+    }
+}
+
+impl LogicalSize {
+    fn empty_initial() -> Self {
+        Self {
+            initial_logical_size: OnceCell::with_value(0),
+            //  initial_logical_size already computed, so, don't admit any calculations
+            initial_size_computation: Arc::new(Semaphore::new(0)),
+            initial_part_end: None,
+            size_added_after_initial: AtomicI64::new(0),
+        }
+    }
+
+    fn deferred_initial(compute_to: Lsn) -> Self {
+        Self {
+            initial_logical_size: OnceCell::new(),
+            initial_size_computation: Arc::new(Semaphore::new(1)),
+            initial_part_end: Some(compute_to),
+            size_added_after_initial: AtomicI64::new(0),
+        }
+    }
+
+    fn current_size(&self) -> anyhow::Result<CurrentLogicalSize> {
+        let size_increment: i64 = self.size_added_after_initial.load(AtomicOrdering::Acquire);
+        //                  ^^^ keep this type explicit so that the casts in this function break if
+        //                  we change the type.
+        match self.initial_logical_size.get() {
+            Some(initial_size) => {
+                initial_size.checked_add_signed(size_increment)
+                    .with_context(|| format!("Overflow during logical size calculation, initial_size: {initial_size}, size_increment: {size_increment}"))
+                    .map(CurrentLogicalSize::Exact)
+            }
+            None => {
+                let non_negative_size_increment = u64::try_from(size_increment).unwrap_or(0);
+                Ok(CurrentLogicalSize::Approximate(non_negative_size_increment))
+            }
+        }
+    }
+
+    fn increment_size(&self, delta: i64) {
+        self.size_added_after_initial
+            .fetch_add(delta, AtomicOrdering::SeqCst);
+    }
+
+    /// Make the value computed by initial logical size computation
+    /// available for re-use. This doesn't contain the incremental part.
+    fn initialized_size(&self, lsn: Lsn) -> Option<u64> {
+        match self.initial_part_end {
+            Some(v) if v == lsn => self.initial_logical_size.get().copied(),
+            _ => None,
+        }
+    }
+}
+
 pub struct WalReceiverInfo {
     pub wal_source_connconf: PgConnectionConfig,
     pub last_received_msg_lsn: Lsn,
@@ -1265,9 +1381,9 @@ impl Timeline {
                         .read()
                         .unwrap()
                         .observe(delta);
-                    info!(layer=%local_layer, residence_millis=delta.as_millis(), "evicted layer after known residence period");
+                    info!(layer=%local_layer.short_id(), residence_millis=delta.as_millis(), "evicted layer after known residence period");
                 } else {
-                    info!(layer=%local_layer, "evicted layer after unknown residence period");
+                    info!(layer=%local_layer.short_id(), "evicted layer after unknown residence period");
                 }
 
                 true
@@ -2123,7 +2239,7 @@ impl Timeline {
         ctx: &RequestContext,
         cancel: CancellationToken,
     ) -> Result<u64, CalculateLogicalSizeError> {
-        span::debug_assert_current_span_has_tenant_and_timeline_id();
+        debug_assert_current_span_has_tenant_and_timeline_id();
 
         let mut timeline_state_updates = self.subscribe_for_state_updates();
         let self_calculation = Arc::clone(self);
@@ -2346,7 +2462,11 @@ impl TraversalLayerExt for Arc<dyn PersistentLayer> {
                 format!("{}", local_path.display())
             }
             None => {
-                format!("remote {}/{self}", self.get_timeline_id())
+                format!(
+                    "remote {}/{}",
+                    self.get_timeline_id(),
+                    self.filename().file_name()
+                )
             }
         }
     }
@@ -2354,7 +2474,11 @@ impl TraversalLayerExt for Arc<dyn PersistentLayer> {
 
 impl TraversalLayerExt for Arc<InMemoryLayer> {
     fn traversal_id(&self) -> TraversalId {
-        format!("timeline {} in-memory {self}", self.get_timeline_id())
+        format!(
+            "timeline {} in-memory {}",
+            self.get_timeline_id(),
+            self.short_id()
+        )
     }
 }
 
@@ -2874,7 +2998,7 @@ impl Timeline {
     }
 
     /// Flush one frozen in-memory layer to disk, as a new delta layer.
-    #[instrument(skip_all, fields(tenant_id=%self.tenant_id, timeline_id=%self.timeline_id, layer=%frozen_layer))]
+    #[instrument(skip_all, fields(tenant_id=%self.tenant_id, timeline_id=%self.timeline_id, layer=%frozen_layer.short_id()))]
     async fn flush_frozen_layer(
         self: &Arc<Self>,
         frozen_layer: Arc<InMemoryLayer>,
@@ -3553,7 +3677,7 @@ impl Timeline {
         let remotes = deltas_to_compact
             .iter()
             .filter(|l| l.is_remote_layer())
-            .inspect(|l| info!("compact requires download of {l}"))
+            .inspect(|l| info!("compact requires download of {}", l.filename().file_name()))
             .map(|l| {
                 l.clone()
                     .downcast_remote_layer()
@@ -3577,7 +3701,7 @@ impl Timeline {
         );
 
         for l in deltas_to_compact.iter() {
-            info!("compact includes {l}");
+            info!("compact includes {}", l.filename().file_name());
         }
 
         // We don't need the original list of layers anymore. Drop it so that
@@ -4192,8 +4316,8 @@ impl Timeline {
             if l.get_lsn_range().end > horizon_cutoff {
                 debug!(
                     "keeping {} because it's newer than horizon_cutoff {}",
-                    l.filename(),
-                    horizon_cutoff,
+                    l.filename().file_name(),
+                    horizon_cutoff
                 );
                 result.layers_needed_by_cutoff += 1;
                 continue 'outer;
@@ -4203,8 +4327,8 @@ impl Timeline {
             if l.get_lsn_range().end > pitr_cutoff {
                 debug!(
                     "keeping {} because it's newer than pitr_cutoff {}",
-                    l.filename(),
-                    pitr_cutoff,
+                    l.filename().file_name(),
+                    pitr_cutoff
                 );
                 result.layers_needed_by_pitr += 1;
                 continue 'outer;
@@ -4222,7 +4346,7 @@ impl Timeline {
                 if &l.get_lsn_range().start <= retain_lsn {
                     debug!(
                         "keeping {} because it's still might be referenced by child branch forked at {} is_dropped: xx is_incremental: {}",
-                        l.filename(),
+                        l.filename().file_name(),
                         retain_lsn,
                         l.is_incremental(),
                     );
@@ -4253,7 +4377,10 @@ impl Timeline {
             if !layers
                 .image_layer_exists(&l.get_key_range(), &(l.get_lsn_range().end..new_gc_cutoff))?
             {
-                debug!("keeping {} because it is the latest layer", l.filename());
+                debug!(
+                    "keeping {} because it is the latest layer",
+                    l.filename().file_name()
+                );
                 // Collect delta key ranges that need image layers to allow garbage
                 // collecting the layers.
                 // It is not so obvious whether we need to propagate information only about
@@ -4270,7 +4397,7 @@ impl Timeline {
             // We didn't find any reason to keep this file, so remove it.
             debug!(
                 "garbage collecting {} is_dropped: xx is_incremental: {}",
-                l.filename(),
+                l.filename().file_name(),
                 l.is_incremental(),
             );
             layers_to_remove.push(Arc::clone(&l));
@@ -4424,12 +4551,12 @@ impl Timeline {
     /// If the caller has a deadline or needs a timeout, they can simply stop polling:
     /// we're **cancellation-safe** because the download happens in a separate task_mgr task.
     /// So, the current download attempt will run to completion even if we stop polling.
-    #[instrument(skip_all, fields(layer=%remote_layer))]
+    #[instrument(skip_all, fields(layer=%remote_layer.short_id()))]
     pub async fn download_remote_layer(
         &self,
         remote_layer: Arc<RemoteLayer>,
     ) -> anyhow::Result<()> {
-        span::debug_assert_current_span_has_tenant_and_timeline_id();
+        debug_assert_current_span_has_tenant_and_timeline_id();
 
         use std::sync::atomic::Ordering::Relaxed;
 
@@ -4462,7 +4589,7 @@ impl Timeline {
             TaskKind::RemoteDownloadTask,
             Some(self.tenant_id),
             Some(self.timeline_id),
-            &format!("download layer {}", remote_layer),
+            &format!("download layer {}", remote_layer.short_id()),
             false,
             async move {
                 let remote_client = self_clone.remote_client.as_ref().unwrap();
@@ -4738,12 +4865,15 @@ impl Timeline {
                 continue;
             }
 
-            let last_activity_ts = l.access_stats().latest_activity().unwrap_or_else(|| {
-                // We only use this fallback if there's an implementation error.
-                // `latest_activity` already does rate-limited warn!() log.
-                debug!(layer=%l, "last_activity returns None, using SystemTime::now");
-                SystemTime::now()
-            });
+            let last_activity_ts = l
+                .access_stats()
+                .latest_activity()
+                .unwrap_or_else(|| {
+                    // We only use this fallback if there's an implementation error.
+                    // `latest_activity` already does rate-limited warn!() log.
+                    debug!(layer=%l.filename().file_name(), "last_activity returns None, using SystemTime::now");
+                    SystemTime::now()
+                });
 
             resident_layers.push(LocalLayerInfoForDiskUsageEviction {
                 layer: l,
@@ -4863,6 +4993,33 @@ fn rename_to_backup(path: &Path) -> anyhow::Result<()> {
     bail!("couldn't find an unused backup number for {:?}", path)
 }
 
+#[cfg(not(debug_assertions))]
+#[inline]
+pub(crate) fn debug_assert_current_span_has_tenant_and_timeline_id() {}
+
+#[cfg(debug_assertions)]
+#[inline]
+pub(crate) fn debug_assert_current_span_has_tenant_and_timeline_id() {
+    use utils::tracing_span_assert;
+
+    pub static TIMELINE_ID_EXTRACTOR: once_cell::sync::Lazy<
+        tracing_span_assert::MultiNameExtractor<2>,
+    > = once_cell::sync::Lazy::new(|| {
+        tracing_span_assert::MultiNameExtractor::new("TimelineId", ["timeline_id", "timeline"])
+    });
+
+    match tracing_span_assert::check_fields_present([
+        &*super::TENANT_ID_EXTRACTOR,
+        &*TIMELINE_ID_EXTRACTOR,
+    ]) {
+        Ok(()) => (),
+        Err(missing) => panic!(
+            "missing extractors: {:?}",
+            missing.into_iter().map(|e| e.name()).collect::<Vec<_>>()
+        ),
+    }
+}
+
 /// Similar to `Arc::ptr_eq`, but only compares the object pointers, not vtables.
 ///
 /// Returns `true` if the two `Arc` point to the same layer, false otherwise.

pageserver/src/tenant/timeline/eviction_task.rs

@@ -70,6 +70,7 @@ impl Timeline {
                 };
 
                 self_clone.eviction_task(cancel).await;
+                info!("eviction task finishing");
                 Ok(())
             },
         );
@@ -77,9 +78,6 @@ impl Timeline {
 
     #[instrument(skip_all, fields(tenant_id = %self.tenant_id, timeline_id = %self.timeline_id))]
     async fn eviction_task(self: Arc<Self>, cancel: CancellationToken) {
-        scopeguard::defer! {
-            info!("eviction task finishing");
-        }
         use crate::tenant::tasks::random_init_delay;
         {
             let policy = self.get_eviction_policy();
@@ -88,6 +86,7 @@ impl Timeline {
                 EvictionPolicy::NoEviction => Duration::from_secs(10),
             };
             if random_init_delay(period, &cancel).await.is_err() {
+                info!("shutting down");
                 return;
             }
         }
@@ -102,6 +101,7 @@ impl Timeline {
                 ControlFlow::Continue(sleep_until) => {
                     tokio::select! {
                         _ = cancel.cancelled() => {
+                            info!("shutting down");
                             break;
                         }
                         _ = tokio::time::sleep_until(sleep_until) => { }
@@ -209,7 +209,7 @@ impl Timeline {
                 let last_activity_ts = hist_layer.access_stats().latest_activity().unwrap_or_else(|| {
                     // We only use this fallback if there's an implementation error.
                     // `latest_activity` already does rate-limited warn!() log.
-                    debug!(layer=%hist_layer, "last_activity returns None, using SystemTime::now");
+                    debug!(layer=%hist_layer.filename().file_name(), "last_activity returns None, using SystemTime::now");
                     SystemTime::now()
                 });
 

pageserver/src/tenant/timeline/logical_size.rs

@@ -1,128 +0,0 @@
-use anyhow::Context;
-use once_cell::sync::OnceCell;
-
-use tokio::sync::Semaphore;
-use utils::lsn::Lsn;
-
-use std::sync::atomic::{AtomicI64, Ordering as AtomicOrdering};
-use std::sync::Arc;
-
-/// Internal structure to hold all data needed for logical size calculation.
-///
-/// Calculation consists of two stages:
-///
-/// 1. Initial size calculation. That might take a long time, because it requires
-/// reading all layers containing relation sizes at `initial_part_end`.
-///
-/// 2. Collecting an incremental part and adding that to the initial size.
-/// Increments are appended on walreceiver writing new timeline data,
-/// which result in increase or decrease of the logical size.
-pub(super) struct LogicalSize {
-    /// Size, potentially slow to compute. Calculating this might require reading multiple
-    /// layers, and even ancestor's layers.
-    ///
-    /// NOTE: size at a given LSN is constant, but after a restart we will calculate
-    /// the initial size at a different LSN.
-    pub initial_logical_size: OnceCell<u64>,
-
-    /// Semaphore to track ongoing calculation of `initial_logical_size`.
-    pub initial_size_computation: Arc<tokio::sync::Semaphore>,
-
-    /// Latest Lsn that has its size uncalculated, could be absent for freshly created timelines.
-    pub initial_part_end: Option<Lsn>,
-
-    /// All other size changes after startup, combined together.
-    ///
-    /// Size shouldn't ever be negative, but this is signed for two reasons:
-    ///
-    /// 1. If we initialized the "baseline" size lazily, while we already
-    /// process incoming WAL, the incoming WAL records could decrement the
-    /// variable and temporarily make it negative. (This is just future-proofing;
-    /// the initialization is currently not done lazily.)
-    ///
-    /// 2. If there is a bug and we e.g. forget to increment it in some cases
-    /// when size grows, but remember to decrement it when it shrinks again, the
-    /// variable could go negative. In that case, it seems better to at least
-    /// try to keep tracking it, rather than clamp or overflow it. Note that
-    /// get_current_logical_size() will clamp the returned value to zero if it's
-    /// negative, and log an error. Could set it permanently to zero or some
-    /// special value to indicate "broken" instead, but this will do for now.
-    ///
-    /// Note that we also expose a copy of this value as a prometheus metric,
-    /// see `current_logical_size_gauge`. Use the `update_current_logical_size`
-    /// to modify this, it will also keep the prometheus metric in sync.
-    pub size_added_after_initial: AtomicI64,
-}
-
-/// Normalized current size, that the data in pageserver occupies.
-#[derive(Debug, Clone, Copy)]
-pub(super) enum CurrentLogicalSize {
-    /// The size is not yet calculated to the end, this is an intermediate result,
-    /// constructed from walreceiver increments and normalized: logical data could delete some objects, hence be negative,
-    /// yet total logical size cannot be below 0.
-    Approximate(u64),
-    // Fully calculated logical size, only other future walreceiver increments are changing it, and those changes are
-    // available for observation without any calculations.
-    Exact(u64),
-}
-
-impl CurrentLogicalSize {
-    pub(super) fn size(&self) -> u64 {
-        *match self {
-            Self::Approximate(size) => size,
-            Self::Exact(size) => size,
-        }
-    }
-}
-
-impl LogicalSize {
-    pub(super) fn empty_initial() -> Self {
-        Self {
-            initial_logical_size: OnceCell::with_value(0),
-            //  initial_logical_size already computed, so, don't admit any calculations
-            initial_size_computation: Arc::new(Semaphore::new(0)),
-            initial_part_end: None,
-            size_added_after_initial: AtomicI64::new(0),
-        }
-    }
-
-    pub(super) fn deferred_initial(compute_to: Lsn) -> Self {
-        Self {
-            initial_logical_size: OnceCell::new(),
-            initial_size_computation: Arc::new(Semaphore::new(1)),
-            initial_part_end: Some(compute_to),
-            size_added_after_initial: AtomicI64::new(0),
-        }
-    }
-
-    pub(super) fn current_size(&self) -> anyhow::Result<CurrentLogicalSize> {
-        let size_increment: i64 = self.size_added_after_initial.load(AtomicOrdering::Acquire);
-        //                  ^^^ keep this type explicit so that the casts in this function break if
-        //                  we change the type.
-        match self.initial_logical_size.get() {
-            Some(initial_size) => {
-                initial_size.checked_add_signed(size_increment)
-                    .with_context(|| format!("Overflow during logical size calculation, initial_size: {initial_size}, size_increment: {size_increment}"))
-                    .map(CurrentLogicalSize::Exact)
-            }
-            None => {
-                let non_negative_size_increment = u64::try_from(size_increment).unwrap_or(0);
-                Ok(CurrentLogicalSize::Approximate(non_negative_size_increment))
-            }
-        }
-    }
-
-    pub(super) fn increment_size(&self, delta: i64) {
-        self.size_added_after_initial
-            .fetch_add(delta, AtomicOrdering::SeqCst);
-    }
-
-    /// Make the value computed by initial logical size computation
-    /// available for re-use. This doesn't contain the incremental part.
-    pub(super) fn initialized_size(&self, lsn: Lsn) -> Option<u64> {
-        match self.initial_part_end {
-            Some(v) if v == lsn => self.initial_logical_size.get().copied(),
-            _ => None,
-        }
-    }
-}

pageserver/src/tenant/timeline/span.rs

@@ -1,25 +0,0 @@
-#[cfg(debug_assertions)]
-use utils::tracing_span_assert::{check_fields_present, Extractor, MultiNameExtractor};
-
-#[cfg(not(debug_assertions))]
-pub(crate) fn debug_assert_current_span_has_tenant_and_timeline_id() {}
-
-#[cfg(debug_assertions)]
-#[track_caller]
-pub(crate) fn debug_assert_current_span_has_tenant_and_timeline_id() {
-    static TIMELINE_ID_EXTRACTOR: once_cell::sync::Lazy<MultiNameExtractor<2>> =
-        once_cell::sync::Lazy::new(|| {
-            MultiNameExtractor::new("TimelineId", ["timeline_id", "timeline"])
-        });
-
-    let fields: [&dyn Extractor; 2] = [
-        &*crate::tenant::span::TENANT_ID_EXTRACTOR,
-        &*TIMELINE_ID_EXTRACTOR,
-    ];
-    if let Err(missing) = check_fields_present(fields) {
-        panic!(
-            "missing extractors: {:?}",
-            missing.into_iter().map(|e| e.name()).collect::<Vec<_>>()
-        )
-    }
-}

pageserver/src/tenant/timeline/uninit.rs

@@ -1,219 +0,0 @@
-use std::{collections::hash_map::Entry, fs, path::PathBuf, sync::Arc};
-
-use anyhow::Context;
-use tracing::{error, info, info_span, warn};
-use utils::{crashsafe, id::TimelineId, lsn::Lsn};
-
-use crate::{
-    context::RequestContext,
-    import_datadir,
-    tenant::{ignore_absent_files, Tenant},
-};
-
-use super::Timeline;
-
-/// A timeline with some of its files on disk, being initialized.
-/// This struct ensures the atomicity of the timeline init: it's either properly created and inserted into pageserver's memory, or
-/// its local files are removed. In the worst case of a crash, an uninit mark file is left behind, which causes the directory
-/// to be removed on next restart.
-///
-/// The caller is responsible for proper timeline data filling before the final init.
-#[must_use]
-pub struct UninitializedTimeline<'t> {
-    pub(crate) owning_tenant: &'t Tenant,
-    timeline_id: TimelineId,
-    raw_timeline: Option<(Arc<Timeline>, TimelineUninitMark)>,
-}
-
-impl<'t> UninitializedTimeline<'t> {
-    pub(crate) fn new(
-        owning_tenant: &'t Tenant,
-        timeline_id: TimelineId,
-        raw_timeline: Option<(Arc<Timeline>, TimelineUninitMark)>,
-    ) -> Self {
-        Self {
-            owning_tenant,
-            timeline_id,
-            raw_timeline,
-        }
-    }
-
-    /// Finish timeline creation: insert it into the Tenant's timelines map and remove the
-    /// uninit mark file.
-    ///
-    /// This function launches the flush loop if not already done.
-    ///
-    /// The caller is responsible for activating the timeline (function `.activate()`).
-    pub(crate) fn finish_creation(mut self) -> anyhow::Result<Arc<Timeline>> {
-        let timeline_id = self.timeline_id;
-        let tenant_id = self.owning_tenant.tenant_id;
-
-        let (new_timeline, uninit_mark) = self.raw_timeline.take().with_context(|| {
-            format!("No timeline for initalization found for {tenant_id}/{timeline_id}")
-        })?;
-
-        // Check that the caller initialized disk_consistent_lsn
-        let new_disk_consistent_lsn = new_timeline.get_disk_consistent_lsn();
-        anyhow::ensure!(
-            new_disk_consistent_lsn.is_valid(),
-            "new timeline {tenant_id}/{timeline_id} has invalid disk_consistent_lsn"
-        );
-
-        let mut timelines = self.owning_tenant.timelines.lock().unwrap();
-        match timelines.entry(timeline_id) {
-            Entry::Occupied(_) => anyhow::bail!(
-                "Found freshly initialized timeline {tenant_id}/{timeline_id} in the tenant map"
-            ),
-            Entry::Vacant(v) => {
-                uninit_mark.remove_uninit_mark().with_context(|| {
-                    format!(
-                        "Failed to remove uninit mark file for timeline {tenant_id}/{timeline_id}"
-                    )
-                })?;
-                v.insert(Arc::clone(&new_timeline));
-
-                new_timeline.maybe_spawn_flush_loop();
-            }
-        }
-
-        Ok(new_timeline)
-    }
-
-    /// Prepares timeline data by loading it from the basebackup archive.
-    pub(crate) async fn import_basebackup_from_tar(
-        self,
-        copyin_read: &mut (impl tokio::io::AsyncRead + Send + Sync + Unpin),
-        base_lsn: Lsn,
-        broker_client: storage_broker::BrokerClientChannel,
-        ctx: &RequestContext,
-    ) -> anyhow::Result<Arc<Timeline>> {
-        let raw_timeline = self.raw_timeline()?;
-
-        import_datadir::import_basebackup_from_tar(raw_timeline, copyin_read, base_lsn, ctx)
-            .await
-            .context("Failed to import basebackup")?;
-
-        // Flush the new layer files to disk, before we make the timeline as available to
-        // the outside world.
-        //
-        // Flush loop needs to be spawned in order to be able to flush.
-        raw_timeline.maybe_spawn_flush_loop();
-
-        fail::fail_point!("before-checkpoint-new-timeline", |_| {
-            anyhow::bail!("failpoint before-checkpoint-new-timeline");
-        });
-
-        raw_timeline
-            .freeze_and_flush()
-            .await
-            .context("Failed to flush after basebackup import")?;
-
-        // All the data has been imported. Insert the Timeline into the tenant's timelines
-        // map and remove the uninit mark file.
-        let tl = self.finish_creation()?;
-        tl.activate(broker_client, None, ctx);
-        Ok(tl)
-    }
-
-    pub(crate) fn raw_timeline(&self) -> anyhow::Result<&Arc<Timeline>> {
-        Ok(&self
-            .raw_timeline
-            .as_ref()
-            .with_context(|| {
-                format!(
-                    "No raw timeline {}/{} found",
-                    self.owning_tenant.tenant_id, self.timeline_id
-                )
-            })?
-            .0)
-    }
-}
-
-impl Drop for UninitializedTimeline<'_> {
-    fn drop(&mut self) {
-        if let Some((_, uninit_mark)) = self.raw_timeline.take() {
-            let _entered = info_span!("drop_uninitialized_timeline", tenant = %self.owning_tenant.tenant_id, timeline = %self.timeline_id).entered();
-            error!("Timeline got dropped without initializing, cleaning its files");
-            cleanup_timeline_directory(uninit_mark);
-        }
-    }
-}
-
-pub(crate) fn cleanup_timeline_directory(uninit_mark: TimelineUninitMark) {
-    let timeline_path = &uninit_mark.timeline_path;
-    match ignore_absent_files(|| fs::remove_dir_all(timeline_path)) {
-        Ok(()) => {
-            info!("Timeline dir {timeline_path:?} removed successfully, removing the uninit mark")
-        }
-        Err(e) => {
-            error!("Failed to clean up uninitialized timeline directory {timeline_path:?}: {e:?}")
-        }
-    }
-    drop(uninit_mark); // mark handles its deletion on drop, gets retained if timeline dir exists
-}
-
-/// An uninit mark file, created along the timeline dir to ensure the timeline either gets fully initialized and loaded into pageserver's memory,
-/// or gets removed eventually.
-///
-/// XXX: it's important to create it near the timeline dir, not inside it to ensure timeline dir gets removed first.
-#[must_use]
-pub(crate) struct TimelineUninitMark {
-    uninit_mark_deleted: bool,
-    uninit_mark_path: PathBuf,
-    pub(crate) timeline_path: PathBuf,
-}
-
-impl TimelineUninitMark {
-    pub(crate) fn new(uninit_mark_path: PathBuf, timeline_path: PathBuf) -> Self {
-        Self {
-            uninit_mark_deleted: false,
-            uninit_mark_path,
-            timeline_path,
-        }
-    }
-
-    fn remove_uninit_mark(mut self) -> anyhow::Result<()> {
-        if !self.uninit_mark_deleted {
-            self.delete_mark_file_if_present()?;
-        }
-
-        Ok(())
-    }
-
-    fn delete_mark_file_if_present(&mut self) -> anyhow::Result<()> {
-        let uninit_mark_file = &self.uninit_mark_path;
-        let uninit_mark_parent = uninit_mark_file
-            .parent()
-            .with_context(|| format!("Uninit mark file {uninit_mark_file:?} has no parent"))?;
-        ignore_absent_files(|| fs::remove_file(uninit_mark_file)).with_context(|| {
-            format!("Failed to remove uninit mark file at path {uninit_mark_file:?}")
-        })?;
-        crashsafe::fsync(uninit_mark_parent).context("Failed to fsync uninit mark parent")?;
-        self.uninit_mark_deleted = true;
-
-        Ok(())
-    }
-}
-
-impl Drop for TimelineUninitMark {
-    fn drop(&mut self) {
-        if !self.uninit_mark_deleted {
-            if self.timeline_path.exists() {
-                error!(
-                    "Uninit mark {} is not removed, timeline {} stays uninitialized",
-                    self.uninit_mark_path.display(),
-                    self.timeline_path.display()
-                )
-            } else {
-                // unblock later timeline creation attempts
-                warn!(
-                    "Removing intermediate uninit mark file {}",
-                    self.uninit_mark_path.display()
-                );
-                if let Err(e) = self.delete_mark_file_if_present() {
-                    error!("Failed to remove the uninit mark file: {e}")
-                }
-            }
-        }
-    }
-}

pageserver/src/tenant/timeline/walreceiver/walreceiver_connection.rs

@@ -71,8 +71,6 @@ pub(super) async fn handle_walreceiver_connection(
     ctx: RequestContext,
     node: NodeId,
 ) -> anyhow::Result<()> {
-    debug_assert_current_span_has_tenant_and_timeline_id();
-
     WALRECEIVER_STARTED_CONNECTIONS.inc();
 
     // Connect to the database in replication mode.
@@ -142,9 +140,6 @@ pub(super) async fn handle_walreceiver_connection(
             }
             Ok(())
         }
-        // Enrich the log lines emitted by this closure with meaningful context.
-        // TODO: technically, this task outlives the surrounding function, so, the
-        // spans won't be properly nested.
         .instrument(tracing::info_span!("poller")),
     );
 

pageserver/src/virtual_file.rs

@@ -302,6 +302,15 @@ impl VirtualFile {
             .observe_closure_duration(|| self.open_options.open(&self.path))?;
 
         // Perform the requested operation on it
+        //
+        // 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.
+        //
+        // XXX: `parking_lot::RwLock` can enable such downgrades, yet its implementation 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.tenant_id, &self.timeline_id])
             .observe_closure_duration(|| func(&file));

pgxn/hnsw/hnsw.c

@@ -122,43 +122,6 @@ hnsw_populate(HierarchicalNSW* hnsw, Relation indexRel, Relation heapRel)
 						   true, true, hnsw_build_callback, (void *) hnsw, NULL);
 }
 
-#ifdef __APPLE__
-
-#include <sys/types.h>
-#include <sys/sysctl.h>
-
-static void
-hnsw_check_available_memory(Size requested)
-{
-	size_t total;
-	if (sysctlbyname("hw.memsize", NULL, &total, NULL, 0) < 0)
-		elog(ERROR, "Failed to get amount of RAM: %m");
-
-	if ((Size)NBuffers*BLCKSZ + requested >= total)
-		elog(ERROR, "HNSW index requeries %ld bytes while only %ld are available",
-			requested, total - (Size)NBuffers*BLCKSZ);
-}
-
-#else
-
-#include <sys/sysinfo.h>
-
-static void
-hnsw_check_available_memory(Size requested)
-{
-	struct sysinfo si;
-	Size total;
-	if (sysinfo(&si) < 0)
-		elog(ERROR, "Failed to get amount of RAM: %m");
-
-	total = si.totalram*si.mem_unit;
-	if ((Size)NBuffers*BLCKSZ + requested >= total)
-		elog(ERROR, "HNSW index requeries %ld bytes while only %ld are available",
-			requested, total - (Size)NBuffers*BLCKSZ);
-}
-
-#endif
-
 static HierarchicalNSW*
 hnsw_get_index(Relation indexRel, Relation heapRel)
 {
@@ -193,8 +156,6 @@ hnsw_get_index(Relation indexRel, Relation heapRel)
 		size_data_per_element = size_links_level0 + data_size + sizeof(label_t);
 		shmem_size =  hnsw_sizeof() + maxelements * size_data_per_element;
 
-		hnsw_check_available_memory(shmem_size);
-
 		/* first try to attach to existed index */
 		if (!dsm_impl_op(DSM_OP_ATTACH, handle, 0, &impl_private,
 						 &mapped_address, &mapped_size, DEBUG1))
@@ -580,7 +541,6 @@ l2_distance(PG_FUNCTION_ARGS)
 				 errmsg("different array dimensions %d and %d", a_dim, b_dim)));
 	}
 
-	#pragma clang loop vectorize(enable)
 	for (int i = 0; i < a_dim; i++)
 	{
 		diff = ax[i] - bx[i];

pgxn/hnsw/hnsw.control

@@ -1,4 +1,4 @@
-comment = 'hnsw index'
+comment = 'hNsw index'
 default_version = '0.1.0'
 module_pathname = '$libdir/hnsw'
 relocatable = true

pgxn/hnsw/hnswalg.cpp

@@ -223,7 +223,6 @@ dist_t fstdistfunc_scalar(const coord_t *x, const coord_t *y, size_t n)
 {
     dist_t 	distance = 0.0;
 
-    #pragma clang loop vectorize(enable)
     for (size_t i = 0; i < n; i++)
     {
         dist_t diff = x[i] - y[i];

pgxn/neon/Makefile

@@ -4,6 +4,7 @@
 MODULE_big = neon
 OBJS = \
 	$(WIN32RES) \
+	extension_server.o \
 	file_cache.o \
 	libpagestore.o \
 	libpqwalproposer.o \

pgxn/neon/extension_server.c

@@ -0,0 +1,104 @@
+
+/*-------------------------------------------------------------------------
+ *
+ * extension_server.c
+ *	  Request compute_ctl to download extension files.
+ *
+ * IDENTIFICATION
+ *	 contrib/neon/extension_server.c
+ *
+ *-------------------------------------------------------------------------
+ */
+#include "postgres.h"
+#include "tcop/pquery.h"
+#include "tcop/utility.h"
+#include "access/xact.h"
+#include "utils/hsearch.h"
+#include "utils/memutils.h"
+#include "commands/defrem.h"
+#include "miscadmin.h"
+#include "utils/acl.h"
+#include "fmgr.h"
+#include "utils/guc.h"
+#include "port.h"
+#include "fmgr.h"
+
+#include <curl/curl.h>
+
+static int extension_server_port = 0;
+
+static download_extension_file_hook_type prev_download_extension_file_hook = NULL;
+
+// to download all SQL (and data) files for an extension:
+// curl -X POST http://localhost:8080/extension_server/postgis
+// it covers two possible extension files layouts:
+// 1. extension_name--version--platform.sql
+// 2. extension_name/extension_name--version.sql
+//    extension_name/extra_files.csv
+//
+// to download specific library file:
+// curl -X POST http://localhost:8080/extension_server/postgis-3.so?is_library=true
+static bool
+neon_download_extension_file_http(const char *filename, bool is_library)
+{
+    CURL *curl;
+    CURLcode res;
+    char *compute_ctl_url;
+    char *postdata;
+    bool ret = false;
+
+    if ((curl = curl_easy_init()) == NULL)
+    {
+        elog(ERROR, "Failed to initialize curl handle");
+    }
+
+    compute_ctl_url = psprintf("http://localhost:%d/extension_server/%s%s",
+                               extension_server_port, filename, is_library ? "?is_library=true" : "");
+
+    elog(LOG, "Sending request to compute_ctl: %s", compute_ctl_url);
+
+    curl_easy_setopt(curl, CURLOPT_CUSTOMREQUEST, "POST");
+    curl_easy_setopt(curl, CURLOPT_URL, compute_ctl_url);
+    // NOTE: 15L may be insufficient time for large extensions like postgis
+    curl_easy_setopt(curl, CURLOPT_TIMEOUT, 15L /* seconds */);
+
+    if (curl)
+    {
+        /* Perform the request, res will get the return code */
+        res = curl_easy_perform(curl);
+        /* Check for errors */
+        if (res == CURLE_OK)
+        {
+            ret = true;
+        }
+        else
+        {
+            // Don't error here because postgres will try to find the file
+            // and will fail with some proper error message if it's not found.
+            elog(WARNING, "neon_download_extension_file_http failed: %s\n", curl_easy_strerror(res));
+        }
+
+        /* always cleanup */
+        curl_easy_cleanup(curl);
+    }
+
+    return ret;
+}
+
+void pg_init_extension_server()
+{
+    // Port to connect to compute_ctl on localhost
+    // to request extension files.
+    DefineCustomIntVariable("neon.extension_server_port",
+                            "connection string to the compute_ctl",
+                            NULL,
+                            &extension_server_port,
+                            0, 0, INT_MAX,
+                            PGC_POSTMASTER,
+                            0, /* no flags required */
+                            NULL, NULL, NULL);
+
+    // set download_extension_file_hook
+    prev_download_extension_file_hook = download_extension_file_hook;
+    download_extension_file_hook = neon_download_extension_file_http;
+}

pgxn/neon/extension_server.h

@@ -0,0 +1 @@
+

pgxn/neon/libpagestore.c

@@ -34,6 +34,7 @@
 
 #define PageStoreTrace DEBUG5
 
+#define MAX_RECONNECT_ATTEMPTS 5
 #define RECONNECT_INTERVAL_USEC 1000000
 
 bool		connected = false;
@@ -54,15 +55,13 @@ int32		max_cluster_size;
 char	   *page_server_connstring;
 char	   *neon_auth_token;
 
-int			readahead_buffer_size = 128;
+int			n_unflushed_requests = 0;
 int			flush_every_n_requests = 8;
-
-int			n_reconnect_attempts = 0;
-int			max_reconnect_attempts = 60;
+int			readahead_buffer_size = 128;
 
 bool	(*old_redo_read_buffer_filter) (XLogReaderState *record, uint8 block_id) = NULL;
 
-static bool pageserver_flush(void);
+static void pageserver_flush(void);
 
 static bool
 pageserver_connect(int elevel)
@@ -233,17 +232,16 @@ pageserver_disconnect(void)
 	}
 }
 
-static bool
+static void
 pageserver_send(NeonRequest * request)
 {
 	StringInfoData req_buff;
+	int n_reconnect_attempts = 0;
 
 	/* If the connection was lost for some reason, reconnect */
 	if (connected && PQstatus(pageserver_conn) == CONNECTION_BAD)
-	{
-		neon_log(LOG, "pageserver_send disconnect bad connection");
 		pageserver_disconnect();
-	}
+
 
 	req_buff = nm_pack_request(request);
 
@@ -254,36 +252,53 @@ pageserver_send(NeonRequest * request)
 	 * See https://github.com/neondatabase/neon/issues/1138
 	 * So try to reestablish connection in case of failure.
 	 */
-	if (!connected)
+	while (true)
 	{
-		while (!pageserver_connect(n_reconnect_attempts < max_reconnect_attempts ? LOG : ERROR))
+		if (!connected)
 		{
-			n_reconnect_attempts += 1;
-			pg_usleep(RECONNECT_INTERVAL_USEC);
+			if (!pageserver_connect(n_reconnect_attempts < MAX_RECONNECT_ATTEMPTS ? LOG : ERROR))
+			{
+				n_reconnect_attempts += 1;
+				pg_usleep(RECONNECT_INTERVAL_USEC);
+				continue;
+			}
 		}
-		n_reconnect_attempts = 0;
-	}
 
-	/*
-	 * Send request.
-	 *
-	 * In principle, this could block if the output buffer is full, and we
-	 * should use async mode and check for interrupts while waiting. In
-	 * practice, our requests are small enough to always fit in the output and
-	 * TCP buffer.
-	 */
-	if (PQputCopyData(pageserver_conn, req_buff.data, req_buff.len) <= 0)
-	{
-		char	   *msg = pchomp(PQerrorMessage(pageserver_conn));
-		pageserver_disconnect();
-		neon_log(LOG, "pageserver_send disconnect because failed to send page request (try to reconnect): %s", msg);
-		pfree(msg);
-		pfree(req_buff.data);
-		return false;
+		/*
+		 * Send request.
+		 *
+		 * In principle, this could block if the output buffer is full, and we
+		 * should use async mode and check for interrupts while waiting. In
+		 * practice, our requests are small enough to always fit in the output and
+		 * TCP buffer.
+		 */
+		if (PQputCopyData(pageserver_conn, req_buff.data, req_buff.len) <= 0)
+		{
+			char	   *msg = pchomp(PQerrorMessage(pageserver_conn));
+			if (n_reconnect_attempts < MAX_RECONNECT_ATTEMPTS)
+			{
+				neon_log(LOG, "failed to send page request (try to reconnect): %s", msg);
+				if (n_reconnect_attempts != 0) /* do not sleep before first reconnect attempt, assuming that pageserver is already restarted */
+					pg_usleep(RECONNECT_INTERVAL_USEC);
+				n_reconnect_attempts += 1;
+				continue;
+			}
+			else
+			{
+				pageserver_disconnect();
+				neon_log(ERROR, "failed to send page request: %s", msg);
+			}
+		}
+		break;
 	}
 
 	pfree(req_buff.data);
 
+	n_unflushed_requests++;
+
+	if (flush_every_n_requests > 0 && n_unflushed_requests >= flush_every_n_requests)
+		pageserver_flush();
+
 	if (message_level_is_interesting(PageStoreTrace))
 	{
 		char	   *msg = nm_to_string((NeonMessage *) request);
@@ -291,7 +306,6 @@ pageserver_send(NeonRequest * request)
 		neon_log(PageStoreTrace, "sent request: %s", msg);
 		pfree(msg);
 	}
-	return true;
 }
 
 static NeonResponse *
@@ -326,25 +340,16 @@ pageserver_receive(void)
 		}
 		else if (rc == -1)
 		{
-			neon_log(LOG, "pageserver_receive disconnect because call_PQgetCopyData returns -1: %s", pchomp(PQerrorMessage(pageserver_conn)));
 			pageserver_disconnect();
 			resp = NULL;
 		}
 		else if (rc == -2)
-		{
-			char* msg = pchomp(PQerrorMessage(pageserver_conn));
-			pageserver_disconnect();
-			neon_log(ERROR, "pageserver_receive disconnect because could not read COPY data: %s", msg);
-		}
+			neon_log(ERROR, "could not read COPY data: %s", pchomp(PQerrorMessage(pageserver_conn)));
 		else
-		{
-			pageserver_disconnect();
-			neon_log(ERROR, "pageserver_receive disconnect because unexpected PQgetCopyData return value: %d", rc);
-		}
+			neon_log(ERROR, "unexpected PQgetCopyData return value: %d", rc);
 	}
 	PG_CATCH();
 	{
-		neon_log(LOG, "pageserver_receive disconnect due to caught exception");
 		pageserver_disconnect();
 		PG_RE_THROW();
 	}
@@ -354,25 +359,21 @@ pageserver_receive(void)
 }
 
 
-static bool
+static void
 pageserver_flush(void)
 {
 	if (!connected)
 	{
 		neon_log(WARNING, "Tried to flush while disconnected");
 	}
-	else
+	else if (PQflush(pageserver_conn))
 	{
-		if (PQflush(pageserver_conn))
-		{
-			char	   *msg = pchomp(PQerrorMessage(pageserver_conn));
-			pageserver_disconnect();
-			neon_log(LOG, "pageserver_flush disconnect because failed to flush page requests: %s", msg);
-			pfree(msg);
-			return false;
-		}
+		char	   *msg = pchomp(PQerrorMessage(pageserver_conn));
+
+		pageserver_disconnect();
+		neon_log(ERROR, "failed to flush page requests: %s", msg);
 	}
-	return true;
+	n_unflushed_requests = 0;
 }
 
 page_server_api api = {
@@ -438,14 +439,6 @@ pg_init_libpagestore(void)
 							PGC_USERSET,
 							0,	/* no flags required */
 							NULL, NULL, NULL);
-	DefineCustomIntVariable("neon.max_reconnect_attempts",
-							"Maximal attempts to reconnect to pages server (with 1 second timeout)",
-							NULL,
-							&max_reconnect_attempts,
-							10, 0, INT_MAX,
-							PGC_USERSET,
-							0,
-							NULL, NULL, NULL);
 	DefineCustomIntVariable("neon.readahead_buffer_size",
 							"number of prefetches to buffer",
 							"This buffer is used to hold and manage prefetched "

pgxn/neon/neon.c

@@ -35,8 +35,11 @@ _PG_init(void)
 {
 	pg_init_libpagestore();
 	pg_init_walproposer();
+
 	InitControlPlaneConnector();
 
+	pg_init_extension_server();
+
         // Important: This must happen after other parts of the extension
         // are loaded, otherwise any settings to GUCs that were set before
         // the extension was loaded will be removed.

pgxn/neon/neon.h

@@ -21,6 +21,8 @@ extern char *neon_tenant;
 extern void pg_init_libpagestore(void);
 extern void pg_init_walproposer(void);
 
+extern void pg_init_extension_server(void);
+
 /*
  * Returns true if we shouldn't do REDO on that block in record indicated by
  * block_id; false otherwise.

pgxn/neon/pagestore_client.h

@@ -145,9 +145,9 @@ extern char *nm_to_string(NeonMessage * msg);
 
 typedef struct
 {
-	bool		(*send) (NeonRequest * request);
+	void		(*send) (NeonRequest * request);
 	NeonResponse *(*receive) (void);
-	bool		(*flush) (void);
+	void		(*flush) (void);
 }			page_server_api;
 
 extern void prefetch_on_ps_disconnect(void);

pgxn/neon/pagestore_smgr.c

@@ -489,8 +489,7 @@ prefetch_wait_for(uint64 ring_index)
 	if (MyPState->ring_flush <= ring_index &&
 		MyPState->ring_unused > MyPState->ring_flush)
 	{
-		if (!page_server->flush())
-			return false;
+		page_server->flush();
 		MyPState->ring_flush = MyPState->ring_unused;
 	}
 
@@ -667,7 +666,7 @@ prefetch_do_request(PrefetchRequest *slot, bool *force_latest, XLogRecPtr *force
 		 * smaller than the current WAL insert/redo pointer, which is already
 		 * larger than this prefetch_lsn. So in any case, that would
 		 * invalidate this cache.
-		 *
+		 * 
 		 * The best LSN to use for effective_request_lsn would be
 		 * XLogCtl->Insert.RedoRecPtr, but that's expensive to access.
 		 */
@@ -678,8 +677,7 @@ prefetch_do_request(PrefetchRequest *slot, bool *force_latest, XLogRecPtr *force
 
 	Assert(slot->response == NULL);
 	Assert(slot->my_ring_index == MyPState->ring_unused);
-
-	while (!page_server->send((NeonRequest *) &request));
+	page_server->send((NeonRequest *) &request);
 
 	/* update prefetch state */
 	MyPState->n_requests_inflight += 1;
@@ -689,7 +687,6 @@ prefetch_do_request(PrefetchRequest *slot, bool *force_latest, XLogRecPtr *force
 	/* update slot state */
 	slot->status = PRFS_REQUESTED;
 
-
 	prfh_insert(MyPState->prf_hash, slot, &found);
 	Assert(!found);
 }
@@ -746,7 +743,6 @@ prefetch_register_buffer(BufferTag tag, bool *force_latest, XLogRecPtr *force_ls
 					prefetch_set_unused(ring_index);
 					entry = NULL;
 				}
-
 			}
 			/* if we don't want the latest version, only accept requests with the exact same LSN */
 			else
@@ -760,23 +756,20 @@ prefetch_register_buffer(BufferTag tag, bool *force_latest, XLogRecPtr *force_ls
 			}
 		}
 
-		if (entry != NULL)
+		/*
+		 * We received a prefetch for a page that was recently read and
+		 * removed from the buffers. Remove that request from the buffers.
+		 */
+		else if (slot->status == PRFS_TAG_REMAINS)
 		{
-			/*
-			 * We received a prefetch for a page that was recently read and
-			 * removed from the buffers. Remove that request from the buffers.
-			 */
-			if (slot->status == PRFS_TAG_REMAINS)
-			{
-				prefetch_set_unused(ring_index);
-				entry = NULL;
-			}
-			else
-			{
-				/* The buffered request is good enough, return that index */
-				pgBufferUsage.prefetch.duplicates++;
-				return ring_index;
-			}
+			prefetch_set_unused(ring_index);
+			entry = NULL;
+		}
+		else
+		{
+			/* The buffered request is good enough, return that index */
+			pgBufferUsage.prefetch.duplicates++;
+			return ring_index;
 		}
 	}
 
@@ -866,7 +859,8 @@ page_server_request(void const *req)
 {
 	NeonResponse* resp;
 	do {
-		while (!page_server->send((NeonRequest *) req) || !page_server->flush());
+		page_server->send((NeonRequest *) req);
+		page_server->flush();
 		MyPState->ring_flush = MyPState->ring_unused;
 		consume_prefetch_responses();
 		resp = page_server->receive();

proxy/Cargo.toml

@@ -29,7 +29,6 @@ metrics.workspace = true
 once_cell.workspace = true
 opentelemetry.workspace = true
 parking_lot.workspace = true
-pbkdf2.workspace = true
 pin-project-lite.workspace = true
 postgres_backend.workspace = true
 pq_proto.workspace = true

proxy/src/compute.rs

@@ -136,17 +136,18 @@ impl Default for ConnCfg {
 
 impl ConnCfg {
     /// Establish a raw TCP connection to the compute node.
-    async fn connect_raw(&self, timeout: Duration) -> io::Result<(SocketAddr, TcpStream, &str)> {
+    async fn connect_raw(&self) -> io::Result<(SocketAddr, TcpStream, &str)> {
         use tokio_postgres::config::Host;
 
         // wrap TcpStream::connect with timeout
         let connect_with_timeout = |host, port| {
-            tokio::time::timeout(timeout, TcpStream::connect((host, port))).map(
+            let connection_timeout = Duration::from_millis(10000);
+            tokio::time::timeout(connection_timeout, TcpStream::connect((host, port))).map(
                 move |res| match res {
                     Ok(tcpstream_connect_res) => tcpstream_connect_res,
                     Err(_) => Err(io::Error::new(
                         io::ErrorKind::TimedOut,
-                        format!("exceeded connection timeout {timeout:?}"),
+                        format!("exceeded connection timeout {connection_timeout:?}"),
                     )),
                 },
             )
@@ -222,9 +223,8 @@ impl ConnCfg {
     async fn do_connect(
         &self,
         allow_self_signed_compute: bool,
-        timeout: Duration,
     ) -> Result<PostgresConnection, ConnectionError> {
-        let (socket_addr, stream, host) = self.connect_raw(timeout).await?;
+        let (socket_addr, stream, host) = self.connect_raw().await?;
 
         let tls_connector = native_tls::TlsConnector::builder()
             .danger_accept_invalid_certs(allow_self_signed_compute)
@@ -264,9 +264,8 @@ impl ConnCfg {
     pub async fn connect(
         &self,
         allow_self_signed_compute: bool,
-        timeout: Duration,
     ) -> Result<PostgresConnection, ConnectionError> {
-        self.do_connect(allow_self_signed_compute, timeout)
+        self.do_connect(allow_self_signed_compute)
             .inspect_err(|err| {
                 // Immediately log the error we have at our disposal.
                 error!("couldn't connect to compute node: {err}");

proxy/src/config.rs

@@ -212,7 +212,7 @@ pub struct CacheOptions {
 
 impl CacheOptions {
     /// Default options for [`crate::auth::caches::NodeInfoCache`].
-    pub const DEFAULT_OPTIONS_NODE_INFO: &str = "size=4000,ttl=4m";
+    pub const DEFAULT_OPTIONS_NODE_INFO: &str = "size=4000,ttl=5m";
 
     /// Parse cache options passed via cmdline.
     /// Example: [`Self::DEFAULT_OPTIONS_NODE_INFO`].

proxy/src/http.rs

@@ -2,7 +2,6 @@
 //! Other modules should use stuff from this module instead of
 //! directly relying on deps like `reqwest` (think loose coupling).
 
-pub mod conn_pool;
 pub mod server;
 pub mod sql_over_http;
 pub mod websocket;

proxy/src/http/conn_pool.rs

@@ -1,278 +0,0 @@
-use parking_lot::Mutex;
-use pq_proto::StartupMessageParams;
-use std::fmt;
-use std::{collections::HashMap, sync::Arc};
-
-use futures::TryFutureExt;
-
-use crate::config;
-use crate::{auth, console};
-
-use super::sql_over_http::MAX_RESPONSE_SIZE;
-
-use crate::proxy::invalidate_cache;
-use crate::proxy::NUM_RETRIES_WAKE_COMPUTE;
-
-use tracing::error;
-use tracing::info;
-
-pub const APP_NAME: &str = "sql_over_http";
-const MAX_CONNS_PER_ENDPOINT: usize = 20;
-
-#[derive(Debug)]
-pub struct ConnInfo {
-    pub username: String,
-    pub dbname: String,
-    pub hostname: String,
-    pub password: String,
-}
-
-impl ConnInfo {
-    // hm, change to hasher to avoid cloning?
-    pub fn db_and_user(&self) -> (String, String) {
-        (self.dbname.clone(), self.username.clone())
-    }
-}
-
-impl fmt::Display for ConnInfo {
-    // use custom display to avoid logging password
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}@{}/{}", self.username, self.hostname, self.dbname)
-    }
-}
-
-struct ConnPoolEntry {
-    conn: tokio_postgres::Client,
-    _last_access: std::time::Instant,
-}
-
-// Per-endpoint connection pool, (dbname, username) -> Vec<ConnPoolEntry>
-// Number of open connections is limited by the `max_conns_per_endpoint`.
-pub struct EndpointConnPool {
-    pools: HashMap<(String, String), Vec<ConnPoolEntry>>,
-    total_conns: usize,
-}
-
-pub struct GlobalConnPool {
-    // endpoint -> per-endpoint connection pool
-    //
-    // That should be a fairly conteded map, so return reference to the per-endpoint
-    // pool as early as possible and release the lock.
-    global_pool: Mutex<HashMap<String, Arc<Mutex<EndpointConnPool>>>>,
-
-    // Maximum number of connections per one endpoint.
-    // Can mix different (dbname, username) connections.
-    // When running out of free slots for a particular endpoint,
-    // falls back to opening a new connection for each request.
-    max_conns_per_endpoint: usize,
-
-    proxy_config: &'static crate::config::ProxyConfig,
-}
-
-impl GlobalConnPool {
-    pub fn new(config: &'static crate::config::ProxyConfig) -> Arc<Self> {
-        Arc::new(Self {
-            global_pool: Mutex::new(HashMap::new()),
-            max_conns_per_endpoint: MAX_CONNS_PER_ENDPOINT,
-            proxy_config: config,
-        })
-    }
-
-    pub async fn get(
-        &self,
-        conn_info: &ConnInfo,
-        force_new: bool,
-    ) -> anyhow::Result<tokio_postgres::Client> {
-        let mut client: Option<tokio_postgres::Client> = None;
-
-        if !force_new {
-            let pool = self.get_endpoint_pool(&conn_info.hostname).await;
-
-            // find a pool entry by (dbname, username) if exists
-            let mut pool = pool.lock();
-            let pool_entries = pool.pools.get_mut(&conn_info.db_and_user());
-            if let Some(pool_entries) = pool_entries {
-                if let Some(entry) = pool_entries.pop() {
-                    client = Some(entry.conn);
-                    pool.total_conns -= 1;
-                }
-            }
-        }
-
-        // ok return cached connection if found and establish a new one otherwise
-        if let Some(client) = client {
-            if client.is_closed() {
-                info!("pool: cached connection '{conn_info}' is closed, opening a new one");
-                connect_to_compute(self.proxy_config, conn_info).await
-            } else {
-                info!("pool: reusing connection '{conn_info}'");
-                Ok(client)
-            }
-        } else {
-            info!("pool: opening a new connection '{conn_info}'");
-            connect_to_compute(self.proxy_config, conn_info).await
-        }
-    }
-
-    pub async fn put(
-        &self,
-        conn_info: &ConnInfo,
-        client: tokio_postgres::Client,
-    ) -> anyhow::Result<()> {
-        let pool = self.get_endpoint_pool(&conn_info.hostname).await;
-
-        // return connection to the pool
-        let mut total_conns;
-        let mut returned = false;
-        let mut per_db_size = 0;
-        {
-            let mut pool = pool.lock();
-            total_conns = pool.total_conns;
-
-            let pool_entries: &mut Vec<ConnPoolEntry> = pool
-                .pools
-                .entry(conn_info.db_and_user())
-                .or_insert_with(|| Vec::with_capacity(1));
-            if total_conns < self.max_conns_per_endpoint {
-                pool_entries.push(ConnPoolEntry {
-                    conn: client,
-                    _last_access: std::time::Instant::now(),
-                });
-
-                total_conns += 1;
-                returned = true;
-                per_db_size = pool_entries.len();
-
-                pool.total_conns += 1;
-            }
-        }
-
-        // do logging outside of the mutex
-        if returned {
-            info!("pool: returning connection '{conn_info}' back to the pool, total_conns={total_conns}, for this (db, user)={per_db_size}");
-        } else {
-            info!("pool: throwing away connection '{conn_info}' because pool is full, total_conns={total_conns}");
-        }
-
-        Ok(())
-    }
-
-    async fn get_endpoint_pool(&self, endpoint: &String) -> Arc<Mutex<EndpointConnPool>> {
-        // find or create a pool for this endpoint
-        let mut created = false;
-        let mut global_pool = self.global_pool.lock();
-        let pool = global_pool
-            .entry(endpoint.clone())
-            .or_insert_with(|| {
-                created = true;
-                Arc::new(Mutex::new(EndpointConnPool {
-                    pools: HashMap::new(),
-                    total_conns: 0,
-                }))
-            })
-            .clone();
-        let global_pool_size = global_pool.len();
-        drop(global_pool);
-
-        // log new global pool size
-        if created {
-            info!(
-                "pool: created new pool for '{endpoint}', global pool size now {global_pool_size}"
-            );
-        }
-
-        pool
-    }
-}
-
-//
-// Wake up the destination if needed. Code here is a bit involved because
-// we reuse the code from the usual proxy and we need to prepare few structures
-// that this code expects.
-//
-async fn connect_to_compute(
-    config: &config::ProxyConfig,
-    conn_info: &ConnInfo,
-) -> anyhow::Result<tokio_postgres::Client> {
-    let tls = config.tls_config.as_ref();
-    let common_names = tls.and_then(|tls| tls.common_names.clone());
-
-    let credential_params = StartupMessageParams::new([
-        ("user", &conn_info.username),
-        ("database", &conn_info.dbname),
-        ("application_name", APP_NAME),
-    ]);
-
-    let creds = config
-        .auth_backend
-        .as_ref()
-        .map(|_| {
-            auth::ClientCredentials::parse(
-                &credential_params,
-                Some(&conn_info.hostname),
-                common_names,
-            )
-        })
-        .transpose()?;
-    let extra = console::ConsoleReqExtra {
-        session_id: uuid::Uuid::new_v4(),
-        application_name: Some(APP_NAME),
-    };
-
-    let node_info = &mut creds.wake_compute(&extra).await?.expect("msg");
-
-    // This code is a copy of `connect_to_compute` from `src/proxy.rs` with
-    // the difference that it uses `tokio_postgres` for the connection.
-    let mut num_retries: usize = NUM_RETRIES_WAKE_COMPUTE;
-    loop {
-        match connect_to_compute_once(node_info, conn_info).await {
-            Err(e) if num_retries > 0 => {
-                info!("compute node's state has changed; requesting a wake-up");
-                match creds.wake_compute(&extra).await? {
-                    // Update `node_info` and try one more time.
-                    Some(new) => {
-                        *node_info = new;
-                    }
-                    // Link auth doesn't work that way, so we just exit.
-                    None => return Err(e),
-                }
-            }
-            other => return other,
-        }
-
-        num_retries -= 1;
-        info!("retrying after wake-up ({num_retries} attempts left)");
-    }
-}
-
-async fn connect_to_compute_once(
-    node_info: &console::CachedNodeInfo,
-    conn_info: &ConnInfo,
-) -> anyhow::Result<tokio_postgres::Client> {
-    let mut config = (*node_info.config).clone();
-
-    let (client, connection) = config
-        .user(&conn_info.username)
-        .password(&conn_info.password)
-        .dbname(&conn_info.dbname)
-        .max_backend_message_size(MAX_RESPONSE_SIZE)
-        .connect(tokio_postgres::NoTls)
-        .inspect_err(|e: &tokio_postgres::Error| {
-            error!(
-                "failed to connect to compute node hosts={:?} ports={:?}: {}",
-                node_info.config.get_hosts(),
-                node_info.config.get_ports(),
-                e
-            );
-            invalidate_cache(node_info)
-        })
-        .await?;
-
-    tokio::spawn(async move {
-        if let Err(e) = connection.await {
-            error!("connection error: {}", e);
-        }
-    });
-
-    Ok(client)
-}

proxy/src/http/sql_over_http.rs

@@ -1,21 +1,25 @@
-use std::sync::Arc;
-
 use futures::pin_mut;
 use futures::StreamExt;
+use futures::TryFutureExt;
 use hyper::body::HttpBody;
 use hyper::http::HeaderName;
 use hyper::http::HeaderValue;
 use hyper::{Body, HeaderMap, Request};
+use pq_proto::StartupMessageParams;
 use serde_json::json;
 use serde_json::Map;
 use serde_json::Value;
 use tokio_postgres::types::Kind;
 use tokio_postgres::types::Type;
 use tokio_postgres::Row;
+use tracing::error;
+use tracing::info;
+use tracing::instrument;
 use url::Url;
 
-use super::conn_pool::ConnInfo;
-use super::conn_pool::GlobalConnPool;
+use crate::proxy::invalidate_cache;
+use crate::proxy::NUM_RETRIES_WAKE_COMPUTE;
+use crate::{auth, config::ProxyConfig, console};
 
 #[derive(serde::Deserialize)]
 struct QueryData {
@@ -23,13 +27,12 @@ struct QueryData {
     params: Vec<serde_json::Value>,
 }
 
-pub const MAX_RESPONSE_SIZE: usize = 1024 * 1024; // 1 MB
+const APP_NAME: &str = "sql_over_http";
+const MAX_RESPONSE_SIZE: usize = 1024 * 1024; // 1 MB
 const MAX_REQUEST_SIZE: u64 = 1024 * 1024; // 1 MB
 
 static RAW_TEXT_OUTPUT: HeaderName = HeaderName::from_static("neon-raw-text-output");
 static ARRAY_MODE: HeaderName = HeaderName::from_static("neon-array-mode");
-static ALLOW_POOL: HeaderName = HeaderName::from_static("neon-pool-opt-in");
-
 static HEADER_VALUE_TRUE: HeaderValue = HeaderValue::from_static("true");
 
 //
@@ -93,6 +96,13 @@ fn json_array_to_pg_array(value: &Value) -> Result<Option<String>, serde_json::E
     }
 }
 
+struct ConnInfo {
+    username: String,
+    dbname: String,
+    hostname: String,
+    password: String,
+}
+
 fn get_conn_info(
     headers: &HeaderMap,
     sni_hostname: Option<String>,
@@ -159,23 +169,50 @@ fn get_conn_info(
 
 // TODO: return different http error codes
 pub async fn handle(
+    config: &'static ProxyConfig,
     request: Request<Body>,
     sni_hostname: Option<String>,
-    conn_pool: Arc<GlobalConnPool>,
 ) -> anyhow::Result<Value> {
     //
     // Determine the destination and connection params
     //
     let headers = request.headers();
     let conn_info = get_conn_info(headers, sni_hostname)?;
+    let credential_params = StartupMessageParams::new([
+        ("user", &conn_info.username),
+        ("database", &conn_info.dbname),
+        ("application_name", APP_NAME),
+    ]);
 
     // Determine the output options. Default behaviour is 'false'. Anything that is not
     // strictly 'true' assumed to be false.
     let raw_output = headers.get(&RAW_TEXT_OUTPUT) == Some(&HEADER_VALUE_TRUE);
     let array_mode = headers.get(&ARRAY_MODE) == Some(&HEADER_VALUE_TRUE);
 
-    // Allow connection pooling only if explicitly requested
-    let allow_pool = headers.get(&ALLOW_POOL) == Some(&HEADER_VALUE_TRUE);
+    //
+    // Wake up the destination if needed. Code here is a bit involved because
+    // we reuse the code from the usual proxy and we need to prepare few structures
+    // that this code expects.
+    //
+    let tls = config.tls_config.as_ref();
+    let common_names = tls.and_then(|tls| tls.common_names.clone());
+    let creds = config
+        .auth_backend
+        .as_ref()
+        .map(|_| {
+            auth::ClientCredentials::parse(
+                &credential_params,
+                Some(&conn_info.hostname),
+                common_names,
+            )
+        })
+        .transpose()?;
+    let extra = console::ConsoleReqExtra {
+        session_id: uuid::Uuid::new_v4(),
+        application_name: Some(APP_NAME),
+    };
+
+    let mut node_info = creds.wake_compute(&extra).await?.expect("msg");
 
     let request_content_length = match request.body().size_hint().upper() {
         Some(v) => v,
@@ -198,8 +235,7 @@ pub async fn handle(
     //
     // Now execute the query and return the result
     //
-    let client = conn_pool.get(&conn_info, !allow_pool).await?;
-
+    let client = connect_to_compute(&mut node_info, &extra, &creds, &conn_info).await?;
     let row_stream = client.query_raw_txt(query, query_params).await?;
 
     // Manually drain the stream into a vector to leave row_stream hanging
@@ -256,13 +292,6 @@ pub async fn handle(
         .map(|row| pg_text_row_to_json(row, raw_output, array_mode))
         .collect::<Result<Vec<_>, _>>()?;
 
-    if allow_pool {
-        // return connection to the pool
-        tokio::task::spawn(async move {
-            let _ = conn_pool.put(&conn_info, client).await;
-        });
-    }
-
     // resulting JSON format is based on the format of node-postgres result
     Ok(json!({
         "command": command_tag_name,
@@ -273,6 +302,70 @@ pub async fn handle(
     }))
 }
 
+/// This function is a copy of `connect_to_compute` from `src/proxy.rs` with
+/// the difference that it uses `tokio_postgres` for the connection.
+#[instrument(skip_all)]
+async fn connect_to_compute(
+    node_info: &mut console::CachedNodeInfo,
+    extra: &console::ConsoleReqExtra<'_>,
+    creds: &auth::BackendType<'_, auth::ClientCredentials<'_>>,
+    conn_info: &ConnInfo,
+) -> anyhow::Result<tokio_postgres::Client> {
+    let mut num_retries: usize = NUM_RETRIES_WAKE_COMPUTE;
+
+    loop {
+        match connect_to_compute_once(node_info, conn_info).await {
+            Err(e) if num_retries > 0 => {
+                info!("compute node's state has changed; requesting a wake-up");
+                match creds.wake_compute(extra).await? {
+                    // Update `node_info` and try one more time.
+                    Some(new) => {
+                        *node_info = new;
+                    }
+                    // Link auth doesn't work that way, so we just exit.
+                    None => return Err(e),
+                }
+            }
+            other => return other,
+        }
+
+        num_retries -= 1;
+        info!("retrying after wake-up ({num_retries} attempts left)");
+    }
+}
+
+async fn connect_to_compute_once(
+    node_info: &console::CachedNodeInfo,
+    conn_info: &ConnInfo,
+) -> anyhow::Result<tokio_postgres::Client> {
+    let mut config = (*node_info.config).clone();
+
+    let (client, connection) = config
+        .user(&conn_info.username)
+        .password(&conn_info.password)
+        .dbname(&conn_info.dbname)
+        .max_backend_message_size(MAX_RESPONSE_SIZE)
+        .connect(tokio_postgres::NoTls)
+        .inspect_err(|e: &tokio_postgres::Error| {
+            error!(
+                "failed to connect to compute node hosts={:?} ports={:?}: {}",
+                node_info.config.get_hosts(),
+                node_info.config.get_ports(),
+                e
+            );
+            invalidate_cache(node_info)
+        })
+        .await?;
+
+    tokio::spawn(async move {
+        if let Err(e) = connection.await {
+            error!("connection error: {}", e);
+        }
+    });
+
+    Ok(client)
+}
+
 //
 // Convert postgres row with text-encoded values to JSON object
 //

proxy/src/http/websocket.rs

@@ -35,7 +35,7 @@ use utils::http::{error::ApiError, json::json_response};
 // Tracking issue: https://github.com/rust-lang/rust/issues/98407.
 use sync_wrapper::SyncWrapper;
 
-use super::{conn_pool::GlobalConnPool, sql_over_http};
+use super::sql_over_http;
 
 pin_project! {
     /// This is a wrapper around a [`WebSocketStream`] that
@@ -164,7 +164,6 @@ async fn serve_websocket(
 async fn ws_handler(
     mut request: Request<Body>,
     config: &'static ProxyConfig,
-    conn_pool: Arc<GlobalConnPool>,
     cancel_map: Arc<CancelMap>,
     session_id: uuid::Uuid,
     sni_hostname: Option<String>,
@@ -193,7 +192,7 @@ async fn ws_handler(
     // TODO: that deserves a refactor as now this function also handles http json client besides websockets.
     // Right now I don't want to blow up sql-over-http patch with file renames and do that as a follow up instead.
     } else if request.uri().path() == "/sql" && request.method() == Method::POST {
-        let result = sql_over_http::handle(request, sni_hostname, conn_pool)
+        let result = sql_over_http::handle(config, request, sni_hostname)
             .instrument(info_span!("sql-over-http"))
             .await;
         let status_code = match result {
@@ -235,8 +234,6 @@ pub async fn task_main(
         info!("websocket server has shut down");
     }
 
-    let conn_pool: Arc<GlobalConnPool> = GlobalConnPool::new(config);
-
     let tls_config = config.tls_config.as_ref().map(|cfg| cfg.to_server_config());
     let tls_acceptor: tokio_rustls::TlsAcceptor = match tls_config {
         Some(config) => config.into(),
@@ -261,18 +258,15 @@ pub async fn task_main(
     let make_svc =
         hyper::service::make_service_fn(|stream: &tokio_rustls::server::TlsStream<AddrStream>| {
             let sni_name = stream.get_ref().1.sni_hostname().map(|s| s.to_string());
-            let conn_pool = conn_pool.clone();
 
             async move {
                 Ok::<_, Infallible>(hyper::service::service_fn(move |req: Request<Body>| {
                     let sni_name = sni_name.clone();
-                    let conn_pool = conn_pool.clone();
-
                     async move {
                         let cancel_map = Arc::new(CancelMap::default());
                         let session_id = uuid::Uuid::new_v4();
 
-                        ws_handler(req, config, conn_pool, cancel_map, session_id, sni_name)
+                        ws_handler(req, config, cancel_map, session_id, sni_name)
                             .instrument(info_span!(
                                 "ws-client",
                                 session = format_args!("{session_id}")

proxy/src/proxy.rs

@@ -16,10 +16,7 @@ use metrics::{register_int_counter, register_int_counter_vec, IntCounter, IntCou
 use once_cell::sync::Lazy;
 use pq_proto::{BeMessage as Be, FeStartupPacket, StartupMessageParams};
 use std::sync::Arc;
-use tokio::{
-    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
-    time,
-};
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio_util::sync::CancellationToken;
 use tracing::{error, info, warn};
 use utils::measured_stream::MeasuredStream;
@@ -308,13 +305,12 @@ pub fn invalidate_cache(node_info: &console::CachedNodeInfo) {
 #[tracing::instrument(name = "connect_once", skip_all)]
 async fn connect_to_compute_once(
     node_info: &console::CachedNodeInfo,
-    timeout: time::Duration,
 ) -> Result<PostgresConnection, compute::ConnectionError> {
     let allow_self_signed_compute = node_info.allow_self_signed_compute;
 
     node_info
         .config
-        .connect(allow_self_signed_compute, timeout)
+        .connect(allow_self_signed_compute)
         .inspect_err(|_: &compute::ConnectionError| invalidate_cache(node_info))
         .await
 }
@@ -332,27 +328,7 @@ async fn connect_to_compute(
     loop {
         // Apply startup params to the (possibly, cached) compute node info.
         node_info.config.set_startup_params(params);
-
-        // Set a shorter timeout for the initial connection attempt.
-        //
-        // In case we try to connect to an outdated address that is no longer valid, the
-        // default behavior of Kubernetes is to drop the packets, causing us to wait for
-        // the entire timeout period. We want to fail fast in such cases.
-        //
-        // A specific case to consider is when we have cached compute node information
-        // with a 4-minute TTL (Time To Live), but the user has executed a `/suspend` API
-        // call, resulting in the nonexistence of the compute node.
-        //
-        // We only use caching in case of scram proxy backed by the console, so reduce
-        // the timeout only in that case.
-        let is_scram_proxy = matches!(creds, auth::BackendType::Console(_, _));
-        let timeout = if is_scram_proxy && num_retries == NUM_RETRIES_WAKE_COMPUTE {
-            time::Duration::from_secs(2)
-        } else {
-            time::Duration::from_secs(10)
-        };
-
-        match connect_to_compute_once(node_info, timeout).await {
+        match connect_to_compute_once(node_info).await {
             Err(e) if num_retries > 0 => {
                 info!("compute node's state has changed; requesting a wake-up");
                 match creds.wake_compute(extra).map_err(io_error).await? {

proxy/src/scram.rs

@@ -45,74 +45,17 @@ fn hmac_sha256<'a>(key: &[u8], parts: impl IntoIterator<Item = &'a [u8]>) -> [u8
     let mut mac = Hmac::<Sha256>::new_from_slice(key).expect("bad key size");
     parts.into_iter().for_each(|s| mac.update(s));
 
-    mac.finalize().into_bytes().into()
+    // TODO: maybe newer `hmac` et al already migrated to regular arrays?
+    let mut result = [0u8; 32];
+    result.copy_from_slice(mac.finalize().into_bytes().as_slice());
+    result
 }
 
 fn sha256<'a>(parts: impl IntoIterator<Item = &'a [u8]>) -> [u8; 32] {
     let mut hasher = Sha256::new();
     parts.into_iter().for_each(|s| hasher.update(s));
 
-    hasher.finalize().into()
-}
-
-#[cfg(test)]
-mod tests {
-    use crate::sasl::{Mechanism, Step};
-
-    use super::{password::SaltedPassword, Exchange, ServerSecret};
-
-    #[test]
-    fn happy_path() {
-        let iterations = 4096;
-        let salt_base64 = "QSXCR+Q6sek8bf92";
-        let pw = SaltedPassword::new(
-            b"pencil",
-            base64::decode(salt_base64).unwrap().as_slice(),
-            iterations,
-        );
-
-        let secret = ServerSecret {
-            iterations,
-            salt_base64: salt_base64.to_owned(),
-            stored_key: pw.client_key().sha256(),
-            server_key: pw.server_key(),
-            doomed: false,
-        };
-        const NONCE: [u8; 18] = [
-            1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
-        ];
-        let mut exchange = Exchange::new(&secret, || NONCE, None);
-
-        let client_first = "n,,n=user,r=rOprNGfwEbeRWgbNEkqO";
-        let client_final = "c=biws,r=rOprNGfwEbeRWgbNEkqOAQIDBAUGBwgJCgsMDQ4PEBES,p=rw1r5Kph5ThxmaUBC2GAQ6MfXbPnNkFiTIvdb/Rear0=";
-        let server_first =
-            "r=rOprNGfwEbeRWgbNEkqOAQIDBAUGBwgJCgsMDQ4PEBES,s=QSXCR+Q6sek8bf92,i=4096";
-        let server_final = "v=qtUDIofVnIhM7tKn93EQUUt5vgMOldcDVu1HC+OH0o0=";
-
-        exchange = match exchange.exchange(client_first).unwrap() {
-            Step::Continue(exchange, message) => {
-                assert_eq!(message, server_first);
-                exchange
-            }
-            Step::Success(_, _) => panic!("expected continue, got success"),
-            Step::Failure(f) => panic!("{f}"),
-        };
-
-        let key = match exchange.exchange(client_final).unwrap() {
-            Step::Success(key, message) => {
-                assert_eq!(message, server_final);
-                key
-            }
-            Step::Continue(_, _) => panic!("expected success, got continue"),
-            Step::Failure(f) => panic!("{f}"),
-        };
-
-        assert_eq!(
-            key.as_bytes(),
-            [
-                74, 103, 1, 132, 12, 31, 200, 48, 28, 54, 82, 232, 207, 12, 138, 189, 40, 32, 134,
-                27, 125, 170, 232, 35, 171, 167, 166, 41, 70, 228, 182, 112,
-            ]
-        );
-    }
+    let mut result = [0u8; 32];
+    result.copy_from_slice(hasher.finalize().as_slice());
+    result
 }

proxy/src/scram/password.rs

@@ -14,7 +14,19 @@ impl SaltedPassword {
     /// See `scram-common.c : scram_SaltedPassword` for details.
     /// Further reading: <https://datatracker.ietf.org/doc/html/rfc2898> (see `PBKDF2`).
     pub fn new(password: &[u8], salt: &[u8], iterations: u32) -> SaltedPassword {
-        pbkdf2::pbkdf2_hmac_array::<sha2::Sha256, 32>(password, salt, iterations).into()
+        let one = 1_u32.to_be_bytes(); // magic
+
+        let mut current = super::hmac_sha256(password, [salt, &one]);
+        let mut result = current;
+        for _ in 1..iterations {
+            current = super::hmac_sha256(password, [current.as_ref()]);
+            // TODO: result = current.zip(result).map(|(x, y)| x ^ y), issue #80094
+            for (i, x) in current.iter().enumerate() {
+                result[i] ^= x;
+            }
+        }
+
+        result.into()
     }
 
     /// Derive `ClientKey` from a salted hashed password.
@@ -34,41 +46,3 @@ impl From<[u8; SALTED_PASSWORD_LEN]> for SaltedPassword {
         Self { bytes }
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::SaltedPassword;
-
-    fn legacy_pbkdf2_impl(password: &[u8], salt: &[u8], iterations: u32) -> SaltedPassword {
-        let one = 1_u32.to_be_bytes(); // magic
-
-        let mut current = super::super::hmac_sha256(password, [salt, &one]);
-        let mut result = current;
-        for _ in 1..iterations {
-            current = super::super::hmac_sha256(password, [current.as_ref()]);
-            // TODO: result = current.zip(result).map(|(x, y)| x ^ y), issue #80094
-            for (i, x) in current.iter().enumerate() {
-                result[i] ^= x;
-            }
-        }
-
-        result.into()
-    }
-
-    #[test]
-    fn pbkdf2() {
-        let password = "a-very-secure-password";
-        let salt = "such-a-random-salt";
-        let iterations = 4096;
-        let output = [
-            203, 18, 206, 81, 4, 154, 193, 100, 147, 41, 211, 217, 177, 203, 69, 210, 194, 211,
-            101, 1, 248, 156, 96, 0, 8, 223, 30, 87, 158, 41, 20, 42,
-        ];
-
-        let actual = SaltedPassword::new(password.as_bytes(), salt.as_bytes(), iterations);
-        let expected = legacy_pbkdf2_impl(password.as_bytes(), salt.as_bytes(), iterations);
-
-        assert_eq!(actual.bytes, output);
-        assert_eq!(actual.bytes, expected.bytes);
-    }
-}

safekeeper/src/control_file.rs

@@ -191,12 +191,6 @@ impl Storage for FileStorage {
                 control_partial_path.display()
             )
         })?;
-        control_partial.flush().await.with_context(|| {
-            format!(
-                "failed to flush safekeeper state into control file at: {}",
-                control_partial_path.display()
-            )
-        })?;
 
         // fsync the file
         if !self.conf.no_sync {

safekeeper/src/pull_timeline.rs

@@ -188,7 +188,6 @@ async fn pull_timeline(status: TimelineStatus, host: String) -> Result<Response>
         let mut response = client.get(&http_url).send().await?;
         while let Some(chunk) = response.chunk().await? {
             file.write_all(&chunk).await?;
-            file.flush().await?;
         }
     }
 

safekeeper/src/send_wal.rs

@@ -403,18 +403,16 @@ impl SafekeeperPostgresHandler {
         };
 
         // take the latest commit_lsn if don't have stop_pos
-        let end_pos = stop_pos.unwrap_or(*commit_lsn_watch_rx.borrow());
+        let mut end_pos = stop_pos.unwrap_or(*commit_lsn_watch_rx.borrow());
 
         if end_pos < start_pos {
-            warn!(
-                "requested start_pos {} is ahead of available WAL end_pos {}",
-                start_pos, end_pos
-            );
+            warn!("start_pos {} is ahead of end_pos {}", start_pos, end_pos);
+            end_pos = start_pos;
         }
 
         info!(
-            "starting streaming from {:?} till {:?}, available WAL ends at {}",
-            start_pos, stop_pos, end_pos
+            "starting streaming from {:?} till {:?}",
+            start_pos, stop_pos
         );
 
         // switch to copy
@@ -549,14 +547,12 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> WalSender<'_, IO> {
             self.end_pos = *self.commit_lsn_watch_rx.borrow();
             if self.end_pos > self.start_pos {
                 // We have something to send.
-                trace!("got end_pos {:?}, streaming", self.end_pos);
                 return Ok(());
             }
 
             // Wait for WAL to appear, now self.end_pos == self.start_pos.
             if let Some(lsn) = wait_for_lsn(&mut self.commit_lsn_watch_rx, self.start_pos).await? {
                 self.end_pos = lsn;
-                trace!("got end_pos {:?}, streaming", self.end_pos);
                 return Ok(());
             }
 

safekeeper/src/wal_storage.rs

@@ -248,10 +248,6 @@ impl PhysicalStorage {
         };
 
         file.write_all(buf).await?;
-        // Note: flush just ensures write above reaches the OS (this is not
-        // needed in case of sync IO as Write::write there calls directly write
-        // syscall, but needed in case of async). It does *not* fsyncs the file.
-        file.flush().await?;
 
         if xlogoff + buf.len() == self.wal_seg_size {
             // If we reached the end of a WAL segment, flush and close it.
@@ -720,7 +716,6 @@ async fn write_zeroes(file: &mut File, mut count: usize) -> Result<()> {
         count -= XLOG_BLCKSZ;
     }
     file.write_all(&ZERO_BLOCK[0..count]).await?;
-    file.flush().await?;
     Ok(())
 }
 

test_runner/fixtures/metrics.py

@@ -59,10 +59,6 @@ PAGESERVER_GLOBAL_METRICS: Tuple[str, ...] = (
     "libmetrics_tracing_event_count_total",
     "pageserver_materialized_cache_hits_total",
     "pageserver_materialized_cache_hits_direct_total",
-    "pageserver_page_cache_read_hits_total",
-    "pageserver_page_cache_read_accesses_total",
-    "pageserver_page_cache_size_current_bytes",
-    "pageserver_page_cache_size_max_bytes",
     "pageserver_getpage_reconstruct_seconds_bucket",
     "pageserver_getpage_reconstruct_seconds_count",
     "pageserver_getpage_reconstruct_seconds_sum",

test_runner/fixtures/neon_fixtures.py

@@ -529,6 +529,16 @@ def available_remote_storages() -> List[RemoteStorageKind]:
     return remote_storages
 
 
+def available_s3_storages() -> List[RemoteStorageKind]:
+    remote_storages = [RemoteStorageKind.MOCK_S3]
+    if os.getenv("ENABLE_REAL_S3_REMOTE_STORAGE") is not None:
+        remote_storages.append(RemoteStorageKind.REAL_S3)
+        log.info("Enabling real s3 storage for tests")
+    else:
+        log.info("Using mock implementations to test remote storage")
+    return remote_storages
+
+
 @dataclass
 class LocalFsStorage:
     root: Path
@@ -549,6 +559,16 @@ class S3Storage:
             "AWS_SECRET_ACCESS_KEY": self.secret_key,
         }
 
+    def to_string(self) -> str:
+        return json.dumps(
+            {
+                "bucket": self.bucket_name,
+                "region": self.bucket_region,
+                "endpoint": self.endpoint,
+                "prefix": self.prefix_in_bucket,
+            }
+        )
+
 
 RemoteStorage = Union[LocalFsStorage, S3Storage]
 
@@ -615,10 +635,12 @@ class NeonEnvBuilder:
         self.rust_log_override = rust_log_override
         self.port_distributor = port_distributor
         self.remote_storage = remote_storage
+        self.ext_remote_storage: Optional[S3Storage] = None
+        self.remote_storage_client: Optional[Any] = None
         self.remote_storage_users = remote_storage_users
         self.broker = broker
         self.run_id = run_id
-        self.mock_s3_server = mock_s3_server
+        self.mock_s3_server: MockS3Server = mock_s3_server
         self.pageserver_config_override = pageserver_config_override
         self.num_safekeepers = num_safekeepers
         self.safekeepers_id_start = safekeepers_id_start
@@ -666,15 +688,24 @@ class NeonEnvBuilder:
         remote_storage_kind: RemoteStorageKind,
         test_name: str,
         force_enable: bool = True,
+        enable_remote_extensions: bool = False,
     ):
         if remote_storage_kind == RemoteStorageKind.NOOP:
             return
         elif remote_storage_kind == RemoteStorageKind.LOCAL_FS:
             self.enable_local_fs_remote_storage(force_enable=force_enable)
         elif remote_storage_kind == RemoteStorageKind.MOCK_S3:
-            self.enable_mock_s3_remote_storage(bucket_name=test_name, force_enable=force_enable)
+            self.enable_mock_s3_remote_storage(
+                bucket_name=test_name,
+                force_enable=force_enable,
+                enable_remote_extensions=enable_remote_extensions,
+            )
         elif remote_storage_kind == RemoteStorageKind.REAL_S3:
-            self.enable_real_s3_remote_storage(test_name=test_name, force_enable=force_enable)
+            self.enable_real_s3_remote_storage(
+                test_name=test_name,
+                force_enable=force_enable,
+                enable_remote_extensions=enable_remote_extensions,
+            )
         else:
             raise RuntimeError(f"Unknown storage type: {remote_storage_kind}")
 
@@ -688,11 +719,15 @@ class NeonEnvBuilder:
         assert force_enable or self.remote_storage is None, "remote storage is enabled already"
         self.remote_storage = LocalFsStorage(Path(self.repo_dir / "local_fs_remote_storage"))
 
-    def enable_mock_s3_remote_storage(self, bucket_name: str, force_enable: bool = True):
+    def enable_mock_s3_remote_storage(
+        self, bucket_name: str, force_enable: bool = True, enable_remote_extensions: bool = False
+    ):
         """
         Sets up the pageserver to use the S3 mock server, creates the bucket, if it's not present already.
         Starts up the mock server, if that does not run yet.
         Errors, if the pageserver has some remote storage configuration already, unless `force_enable` is not set to `True`.
+
+        Also creates the bucket for extensions, self.ext_remote_storage bucket
         """
         assert force_enable or self.remote_storage is None, "remote storage is enabled already"
         mock_endpoint = self.mock_s3_server.endpoint()
@@ -713,9 +748,22 @@ class NeonEnvBuilder:
             bucket_region=mock_region,
             access_key=self.mock_s3_server.access_key(),
             secret_key=self.mock_s3_server.secret_key(),
+            prefix_in_bucket="pageserver",
         )
 
-    def enable_real_s3_remote_storage(self, test_name: str, force_enable: bool = True):
+        if enable_remote_extensions:
+            self.ext_remote_storage = S3Storage(
+                bucket_name=bucket_name,
+                endpoint=mock_endpoint,
+                bucket_region=mock_region,
+                access_key=self.mock_s3_server.access_key(),
+                secret_key=self.mock_s3_server.secret_key(),
+                prefix_in_bucket="ext",
+            )
+
+    def enable_real_s3_remote_storage(
+        self, test_name: str, force_enable: bool = True, enable_remote_extensions: bool = False
+    ):
         """
         Sets up configuration to use real s3 endpoint without mock server
         """
@@ -755,6 +803,15 @@ class NeonEnvBuilder:
             prefix_in_bucket=self.remote_storage_prefix,
         )
 
+        if enable_remote_extensions:
+            self.ext_remote_storage = S3Storage(
+                bucket_name="neon-dev-extensions",
+                bucket_region="eu-central-1",
+                access_key=access_key,
+                secret_key=secret_key,
+                prefix_in_bucket="5555",
+            )
+
     def cleanup_local_storage(self):
         if self.preserve_database_files:
             return
@@ -788,6 +845,7 @@ class NeonEnvBuilder:
         # `self.remote_storage_prefix` is coupled with `S3Storage` storage type,
         # so this line effectively a no-op
         assert isinstance(self.remote_storage, S3Storage)
+        assert self.remote_storage_client is not None
 
         if self.keep_remote_storage_contents:
             log.info("keep_remote_storage_contents skipping remote storage cleanup")
@@ -917,6 +975,8 @@ class NeonEnv:
         self.neon_binpath = config.neon_binpath
         self.pg_distrib_dir = config.pg_distrib_dir
         self.endpoint_counter = 0
+        self.remote_storage_client = config.remote_storage_client
+        self.ext_remote_storage = config.ext_remote_storage
 
         # generate initial tenant ID here instead of letting 'neon init' generate it,
         # so that we don't need to dig it out of the config file afterwards.
@@ -1503,6 +1563,7 @@ class NeonCli(AbstractNeonCli):
         safekeepers: Optional[List[int]] = None,
         tenant_id: Optional[TenantId] = None,
         lsn: Optional[Lsn] = None,
+        remote_ext_config: Optional[str] = None,
     ) -> "subprocess.CompletedProcess[str]":
         args = [
             "endpoint",
@@ -1512,6 +1573,8 @@ class NeonCli(AbstractNeonCli):
             "--pg-version",
             self.env.pg_version,
         ]
+        if remote_ext_config is not None:
+            args.extend(["--remote-ext-config", remote_ext_config])
         if lsn is not None:
             args.append(f"--lsn={lsn}")
         args.extend(["--pg-port", str(pg_port)])
@@ -2373,7 +2436,7 @@ class Endpoint(PgProtocol):
 
         return self
 
-    def start(self) -> "Endpoint":
+    def start(self, remote_ext_config: Optional[str] = None) -> "Endpoint":
         """
         Start the Postgres instance.
         Returns self.
@@ -2389,6 +2452,7 @@ class Endpoint(PgProtocol):
             http_port=self.http_port,
             tenant_id=self.tenant_id,
             safekeepers=self.active_safekeepers,
+            remote_ext_config=remote_ext_config,
         )
         self.running = True
 
@@ -2478,6 +2542,7 @@ class Endpoint(PgProtocol):
         hot_standby: bool = False,
         lsn: Optional[Lsn] = None,
         config_lines: Optional[List[str]] = None,
+        remote_ext_config: Optional[str] = None,
     ) -> "Endpoint":
         """
         Create an endpoint, apply config, and start Postgres.
@@ -2492,7 +2557,7 @@ class Endpoint(PgProtocol):
             config_lines=config_lines,
             hot_standby=hot_standby,
             lsn=lsn,
-        ).start()
+        ).start(remote_ext_config=remote_ext_config)
 
         log.info(f"Postgres startup took {time.time() - started_at} seconds")
 
@@ -2526,6 +2591,7 @@ class EndpointFactory:
         lsn: Optional[Lsn] = None,
         hot_standby: bool = False,
         config_lines: Optional[List[str]] = None,
+        remote_ext_config: Optional[str] = None,
     ) -> Endpoint:
         ep = Endpoint(
             self.env,
@@ -2542,6 +2608,7 @@ class EndpointFactory:
             hot_standby=hot_standby,
             config_lines=config_lines,
             lsn=lsn,
+            remote_ext_config=remote_ext_config,
         )
 
     def create(

test_runner/fixtures/types.py

@@ -89,6 +89,9 @@ class TenantId(Id):
     def __repr__(self) -> str:
         return f'`TenantId("{self.id.hex()}")'
 
+    def __str__(self) -> str:
+        return self.id.hex()
+
 
 class TimelineId(Id):
     def __repr__(self) -> str:

test_runner/performance/test_startup.py

@@ -52,7 +52,6 @@ def test_startup_simple(neon_env_builder: NeonEnvBuilder, zenbenchmark: NeonBenc
             "wait_for_spec_ms": f"{i}_wait_for_spec",
             "sync_safekeepers_ms": f"{i}_sync_safekeepers",
             "basebackup_ms": f"{i}_basebackup",
-            "start_postgres_ms": f"{i}_start_postgres",
             "config_ms": f"{i}_config",
             "total_startup_ms": f"{i}_total_startup",
         }

test_runner/regress/data/extension_test/v14/ext_index.json

@@ -0,0 +1,14 @@
+{
+  "enabled_extensions": {
+    "123454321": [
+      "anon"
+    ],
+    "public": [
+      "embedding"
+    ]
+  },
+  "control_data": {
+    "embedding": "comment = 'hnsw index' \ndefault_version = '0.1.0' \nmodule_pathname = '$libdir/embedding' \nrelocatable = true \ntrusted = true",
+    "anon": "# PostgreSQL Anonymizer (anon) extension \ncomment = 'Data anonymization tools' \ndefault_version = '1.1.0' \ndirectory='extension/anon' \nrelocatable = false \nrequires = 'pgcrypto' \nsuperuser = false \nmodule_pathname = '$libdir/anon' \ntrusted = true \n"
+  }
+}

test_runner/regress/data/extension_test/v15/ext_index.json

@@ -0,0 +1,14 @@
+{
+  "enabled_extensions": {
+    "123454321": [
+      "anon"
+    ],
+    "public": [
+      "embedding"
+    ]
+  },
+  "control_data": {
+    "embedding": "comment = 'hnsw index' \ndefault_version = '0.1.0' \nmodule_pathname = '$libdir/embedding' \nrelocatable = true \ntrusted = true",
+    "anon": "# PostgreSQL Anonymizer (anon) extension \ncomment = 'Data anonymization tools' \ndefault_version = '1.1.0' \ndirectory='extension/anon' \nrelocatable = false \nrequires = 'pgcrypto' \nsuperuser = false \nmodule_pathname = '$libdir/anon' \ntrusted = true \n"
+  }
+}

test_runner/regress/test_download_extensions.py

@@ -0,0 +1,122 @@
+import os
+import shutil
+from contextlib import closing
+
+import pytest
+from fixtures.log_helper import log
+from fixtures.neon_fixtures import (
+    NeonEnvBuilder,
+    RemoteStorageKind,
+    available_s3_storages,
+)
+from fixtures.pg_version import PgVersion
+
+# Generate mock extension files and upload them to the mock bucket.
+#
+# NOTE: You must have appropriate AWS credentials to run REAL_S3 test.
+# It may also be necessary to set the following environment variables for MOCK_S3 test:
+#   export AWS_ACCESS_KEY_ID='test' #   export AWS_SECRET_ACCESS_KEY='test'
+#   export AWS_SECURITY_TOKEN='test' #   export AWS_SESSION_TOKEN='test'
+#   export AWS_DEFAULT_REGION='us-east-1'
+
+
+@pytest.mark.parametrize("remote_storage_kind", available_s3_storages())
+def test_remote_extensions(
+    neon_env_builder: NeonEnvBuilder,
+    remote_storage_kind: RemoteStorageKind,
+    pg_version: PgVersion,
+):
+    neon_env_builder.enable_remote_storage(
+        remote_storage_kind=remote_storage_kind,
+        test_name="test_remote_extensions",
+        enable_remote_extensions=True,
+    )
+    neon_env_builder.num_safekeepers = 3
+    env = neon_env_builder.init_start()
+    tenant_id, _ = env.neon_cli.create_tenant()
+    env.neon_cli.create_timeline("test_remote_extensions", tenant_id=tenant_id)
+
+    # For MOCK_S3 we upload test files.
+    # For REAL_S3 we use the files already in the bucket
+    if remote_storage_kind == RemoteStorageKind.MOCK_S3:
+        log.info("Uploading test files to mock bucket")
+
+        def upload_test_file(from_path, to_path):
+            assert env.ext_remote_storage is not None  # satisfy mypy
+            assert env.remote_storage_client is not None  # satisfy mypy
+            with open(
+                f"test_runner/regress/data/extension_test/v{pg_version}/{from_path}", "rb"
+            ) as f:
+                env.remote_storage_client.upload_fileobj(
+                    f,
+                    env.ext_remote_storage.bucket_name,
+                    f"ext/v{pg_version}/{to_path}",
+                )
+
+        upload_test_file("ext_index.json", "ext_index.json")
+        upload_test_file("anon.tar.gz", "extensions/anon.tar.gz")
+        upload_test_file("embedding.tar.gz", "extensions/embedding.tar.gz")
+
+    assert env.ext_remote_storage is not None  # satisfy mypy
+    assert env.remote_storage_client is not None  # satisfy mypy
+    try:
+        # Start a compute node and check that it can download the extensions
+        # and use them to CREATE EXTENSION and LOAD
+        endpoint = env.endpoints.create_start(
+            "test_remote_extensions",
+            tenant_id=tenant_id,
+            remote_ext_config=env.ext_remote_storage.to_string(),
+            # config_lines=["log_min_messages=debug3"],
+        )
+        with closing(endpoint.connect()) as conn:
+            with conn.cursor() as cur:
+                # Check that appropriate control files were downloaded
+                cur.execute("SELECT * FROM pg_available_extensions")
+                all_extensions = [x[0] for x in cur.fetchall()]
+                log.info(all_extensions)
+                assert "anon" in all_extensions
+                assert "embedding" in all_extensions
+                # TODO: check that we cant't download custom extensions for other tenant ids
+
+                # check that we can download public extension
+                cur.execute("CREATE EXTENSION embedding")
+                cur.execute("SELECT extname FROM pg_extension")
+                assert "embedding" in [x[0] for x in cur.fetchall()]
+
+                # check that we can download private extension
+                try:
+                    cur.execute("CREATE EXTENSION anon")
+                except Exception as err:
+                    log.info("error creating anon extension")
+                    assert "pgcrypto" in str(err), "unexpected error creating anon extension"
+
+                # TODO: try to load libraries as well
+
+    finally:
+        # Cleaning up downloaded files is important for local tests
+        # or else one test could reuse the files from another test or another test run
+        cleanup_files = [
+            "lib/postgresql/anon.so",
+            "lib/postgresql/embedding.so",
+            "share/postgresql/extension/anon.control",
+            "share/postgresql/extension/embedding--0.1.0.sql",
+            "share/postgresql/extension/embedding.control",
+        ]
+        cleanup_files = [f"pg_install/v{pg_version}/" + x for x in cleanup_files]
+        cleanup_folders = [
+            "extensions",
+            f"pg_install/v{pg_version}/share/postgresql/extension/anon",
+            f"pg_install/v{pg_version}/extensions",
+        ]
+        for file in cleanup_files:
+            try:
+                os.remove(file)
+                log.info(f"removed file {file}")
+            except Exception as err:
+                log.info(f"error removing file {file}: {err}")
+        for folder in cleanup_folders:
+            try:
+                shutil.rmtree(folder)
+                log.info(f"removed folder {folder}")
+            except Exception as err:
+                log.info(f"error removing file {file}: {err}")

test_runner/regress/test_multixact.py

@@ -1,7 +1,3 @@
-import random
-import threading
-from threading import Thread
-
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import NeonEnv, check_restored_datadir_content
 from fixtures.utils import query_scalar
@@ -19,17 +15,11 @@ def test_multixact(neon_simple_env: NeonEnv, test_output_dir):
     endpoint = env.endpoints.create_start("test_multixact")
 
     log.info("postgres is running on 'test_multixact' branch")
-
-    n_records = 100
-    n_threads =   5
-    n_iters =  1000
-    n_restarts = 10
-
     cur = endpoint.connect().cursor()
     cur.execute(
-        f"""
-        CREATE TABLE t1(pk int primary key, val integer);
-        INSERT INTO t1 values (generate_series(1, {n_records}), 0);
+        """
+        CREATE TABLE t1(i int primary key);
+        INSERT INTO t1 select * from generate_series(1, 100);
     """
     )
 
@@ -38,32 +28,26 @@ def test_multixact(neon_simple_env: NeonEnv, test_output_dir):
     )
 
     # Lock entries using parallel connections in a round-robin fashion.
-    def do_updates():
+    nclients = 20
+    connections = []
+    for i in range(nclients):
+        # Do not turn on autocommit. We want to hold the key-share locks.
         conn = endpoint.connect(autocommit=False)
-        for i in range(n_iters):
-            pk = random.randrange(1, n_records)
-            conn.cursor().execute(f"update t1 set val=val+1 where pk={pk}")
-            conn.cursor().execute("select * from t1 for key share")
-            conn.commit()
-        conn.close()
+        connections.append(conn)
 
-    for iter in range(n_restarts):
-        threads: List[threading.Thread] = []
-        for i in range(n_threads):
-            threads.append(threading.Thread(target=do_updates, args=(), daemon=False))
-            threads[-1].start()
+    # On each iteration, we commit the previous transaction on a connection,
+    # and issue antoher select. Each SELECT generates a new multixact that
+    # includes the new XID, and the XIDs of all the other parallel transactions.
+    # This generates enough traffic on both multixact offsets and members SLRUs
+    # to cross page boundaries.
+    for i in range(5000):
+        conn = connections[i % nclients]
+        conn.commit()
+        conn.cursor().execute("select * from t1 for key share")
 
-        for thread in threads:
-            thread.join()
-
-        # Restart endpoint
-        endpoint.stop()
-        endpoint.start()
-
-        conn = endpoint.connect()
-        cur = conn.cursor()
-        cur.execute("select count(*) from t1")
-        assert cur.fetchone() == (n_records,)
+    # We have multixacts now. We can close the connections.
+    for c in connections:
+        c.close()
 
     # force wal flush
     cur.execute("checkpoint")
@@ -90,3 +74,6 @@ def test_multixact(neon_simple_env: NeonEnv, test_output_dir):
 
     # Check that we restored pg_controlfile correctly
     assert next_multixact_id_new == next_multixact_id
+
+    # Check that we can restore the content of the datadir correctly
+    check_restored_datadir_content(test_output_dir, env, endpoint)

test_runner/regress/test_timeline_delete.py

@@ -275,6 +275,7 @@ def assert_prefix_empty(neon_env_builder: NeonEnvBuilder, prefix: Optional[str]
     assert isinstance(neon_env_builder.remote_storage, S3Storage)
 
     # Note that this doesnt use pagination, so list is not guaranteed to be exhaustive.
+    assert neon_env_builder.remote_storage_client is not None
     response = neon_env_builder.remote_storage_client.list_objects_v2(
         Bucket=neon_env_builder.remote_storage.bucket_name,
         Prefix=prefix or neon_env_builder.remote_storage.prefix_in_bucket or "",
@@ -628,7 +629,7 @@ def test_timeline_delete_works_for_remote_smoke(
         )
 
     # for some reason the check above doesnt immediately take effect for the below.
-    # Assume it is mock server incosistency and check twice.
+    # Assume it is mock server inconsistency and check twice.
     wait_until(
         2,
         0.5,