drain instead of mark offline

Signed-off-by: Alex Chi Z <chi@neon.tech>

.dockerignore

@@ -19,7 +19,7 @@
 !pageserver/
 !pgxn/
 !proxy/
-!object_storage/
+!endpoint_storage/
 !storage_scrubber/
 !safekeeper/
 !storage_broker/

.github/actions/run-python-test-set/action.yml

@@ -133,6 +133,7 @@ runs:
         fi
 
         PERF_REPORT_DIR="$(realpath test_runner/perf-report-local)"
+        echo "PERF_REPORT_DIR=${PERF_REPORT_DIR}" >> ${GITHUB_ENV}
         rm -rf $PERF_REPORT_DIR
 
         TEST_SELECTION="test_runner/${{ inputs.test_selection }}"
@@ -209,11 +210,12 @@ runs:
           --verbose \
           -rA $TEST_SELECTION $EXTRA_PARAMS
 
-        if [[ "${{ inputs.save_perf_report }}" == "true" ]]; then
-          export REPORT_FROM="$PERF_REPORT_DIR"
-          export REPORT_TO="$PLATFORM"
-          scripts/generate_and_push_perf_report.sh
-        fi
+    - name: Upload performance report
+      if: ${{ !cancelled() && inputs.save_perf_report == 'true' }}
+      shell: bash -euxo pipefail {0}
+      run: |
+        export REPORT_FROM="${PERF_REPORT_DIR}"
+        scripts/generate_and_push_perf_report.sh
 
     - name: Upload compatibility snapshot
       # Note, that we use `github.base_ref` which is a target branch for a PR

Cargo.lock

@@ -2037,6 +2037,33 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "endpoint_storage"
+version = "0.0.1"
+dependencies = [
+ "anyhow",
+ "axum",
+ "axum-extra",
+ "camino",
+ "camino-tempfile",
+ "futures",
+ "http-body-util",
+ "itertools 0.10.5",
+ "jsonwebtoken",
+ "prometheus",
+ "rand 0.8.5",
+ "remote_storage",
+ "serde",
+ "serde_json",
+ "test-log",
+ "tokio",
+ "tokio-util",
+ "tower 0.5.2",
+ "tracing",
+ "utils",
+ "workspace_hack",
+]
+
 [[package]]
 name = "enum-map"
 version = "2.5.0"
@@ -3998,33 +4025,6 @@ dependencies = [
  "memchr",
 ]
 
-[[package]]
-name = "object_storage"
-version = "0.0.1"
-dependencies = [
- "anyhow",
- "axum",
- "axum-extra",
- "camino",
- "camino-tempfile",
- "futures",
- "http-body-util",
- "itertools 0.10.5",
- "jsonwebtoken",
- "prometheus",
- "rand 0.8.5",
- "remote_storage",
- "serde",
- "serde_json",
- "test-log",
- "tokio",
- "tokio-util",
- "tower 0.5.2",
- "tracing",
- "utils",
- "workspace_hack",
-]
-
 [[package]]
 name = "once_cell"
 version = "1.20.2"

Cargo.toml

@@ -40,7 +40,7 @@ members = [
     "libs/proxy/postgres-protocol2",
     "libs/proxy/postgres-types2",
     "libs/proxy/tokio-postgres2",
-    "object_storage",
+    "endpoint_storage",
 ]
 
 [workspace.package]

Dockerfile

@@ -89,7 +89,7 @@ RUN set -e \
       --bin storage_broker  \
       --bin storage_controller  \
       --bin proxy  \
-      --bin object_storage \
+      --bin endpoint_storage \
       --bin neon_local \
       --bin storage_scrubber \
       --locked --release
@@ -122,7 +122,7 @@ COPY --from=build --chown=neon:neon /home/nonroot/target/release/safekeeper
 COPY --from=build --chown=neon:neon /home/nonroot/target/release/storage_broker      /usr/local/bin
 COPY --from=build --chown=neon:neon /home/nonroot/target/release/storage_controller  /usr/local/bin
 COPY --from=build --chown=neon:neon /home/nonroot/target/release/proxy               /usr/local/bin
-COPY --from=build --chown=neon:neon /home/nonroot/target/release/object_storage      /usr/local/bin
+COPY --from=build --chown=neon:neon /home/nonroot/target/release/endpoint_storage    /usr/local/bin
 COPY --from=build --chown=neon:neon /home/nonroot/target/release/neon_local          /usr/local/bin
 COPY --from=build --chown=neon:neon /home/nonroot/target/release/storage_scrubber    /usr/local/bin
 

compute/compute-node.Dockerfile

@@ -1677,7 +1677,7 @@ RUN set -e \
     && apt clean && rm -rf /var/lib/apt/lists/*
 
 # Use `dist_man_MANS=` to skip manpage generation (which requires python3/pandoc)
-ENV PGBOUNCER_TAG=pgbouncer_1_22_1
+ENV PGBOUNCER_TAG=pgbouncer_1_24_1
 RUN set -e \
     && git clone --recurse-submodules --depth 1 --branch ${PGBOUNCER_TAG} https://github.com/pgbouncer/pgbouncer.git pgbouncer \
     && cd pgbouncer \

compute/patches/pg_repack.patch

@@ -11,6 +11,14 @@ index bf6edcb..89b4c7f 100644
  
  USE_PGXS = 1	# use pgxs if not in contrib directory
  PGXS := $(shell $(PG_CONFIG) --pgxs)
+diff --git a/regress/expected/init-extension.out b/regress/expected/init-extension.out
+index 9f2e171..f6e4f8d 100644
+--- a/regress/expected/init-extension.out
++++ b/regress/expected/init-extension.out
+@@ -1,3 +1,2 @@
+ SET client_min_messages = warning;
+ CREATE EXTENSION pg_repack;
+-RESET client_min_messages;
 diff --git a/regress/expected/nosuper.out b/regress/expected/nosuper.out
 index 8d0a94e..63b68bf 100644
 --- a/regress/expected/nosuper.out
@@ -42,6 +50,14 @@ index 8d0a94e..63b68bf 100644
  INFO: repacking table "public.tbl_cluster"
  ERROR: query failed: ERROR:  current transaction is aborted, commands ignored until end of transaction block
  DETAIL: query was: RESET lock_timeout
+diff --git a/regress/sql/init-extension.sql b/regress/sql/init-extension.sql
+index 9f2e171..f6e4f8d 100644
+--- a/regress/sql/init-extension.sql
++++ b/regress/sql/init-extension.sql
+@@ -1,3 +1,2 @@
+ SET client_min_messages = warning;
+ CREATE EXTENSION pg_repack;
+-RESET client_min_messages;
 diff --git a/regress/sql/nosuper.sql b/regress/sql/nosuper.sql
 index 072f0fa..dbe60f8 100644
 --- a/regress/sql/nosuper.sql

compute_tools/src/compute.rs

@@ -1500,9 +1500,24 @@ impl ComputeNode {
             let mut conf = conf.as_ref().clone();
             conf.application_name("compute_ctl:migrations");
 
-            if let Err(e) = handle_migrations(conf).await {
-                error!("Failed to run migrations: {}", e);
-            }
+            match conf.connect(NoTls).await {
+                Ok((mut client, connection)) => {
+                    tokio::spawn(async move {
+                        if let Err(e) = connection.await {
+                            eprintln!("connection error: {}", e);
+                        }
+                    });
+                    if let Err(e) = handle_migrations(&mut client).await {
+                        error!("Failed to run migrations: {}", e);
+                    }
+                }
+                Err(e) => {
+                    error!(
+                        "Failed to connect to the compute for running migrations: {}",
+                        e
+                    );
+                }
+            };
         });
 
         Ok::<(), anyhow::Error>(())

compute_tools/src/migration.rs

@@ -1,60 +1,29 @@
 use anyhow::{Context, Result};
 use fail::fail_point;
-use postgres::NoTls;
-use tokio_postgres::{Client, Config, Transaction};
-use tracing::{error, info, warn};
+use tokio_postgres::{Client, Transaction};
+use tracing::{error, info};
 
 use crate::metrics::DB_MIGRATION_FAILED;
 
 /// Runs a series of migrations on a target database
-use compute_api::spec::{Database, PgIdent};
-
-use crate::pg_helpers::{Escaping, get_existing_dbs_async};
-
-pub(crate) enum Migration<'m> {
-    /// Cluster migrations are things like catalog updates, where they can be
-    /// run in the default Postgres database, but affect every database in the
-    /// cluster.
-    Cluster(&'m str),
-
-    /// Per-database migrations will be run in every database of the cluster.
-    /// The migration will not be marked as completed until after it has been
-    /// run in every database. We will save the `postgres` database for last so
-    /// that we can commit the transaction as applied in the
-    /// neon_migration.migration_id table.
-    ///
-    /// Please be aware of the race condition that exists for this type of
-    /// migration. At the beginning of running the series of migrations, we get
-    /// the current list of databases. However, we run migrations in a separate
-    /// thread in order to not block connections to the compute. If after the
-    /// time we have gotten the list of databases in the cluster, a user creates
-    /// a new database, that database will not receive the migration, but we
-    /// will have marked the migration as completed successfully, assuming all
-    /// previous databases ran the migration to completion.
-    PerDatabase(&'m str),
-}
-
 pub(crate) struct MigrationRunner<'m> {
-    /// Postgres client configuration.
-    config: Config,
-
-    /// List of migrations to run.
-    migrations: &'m [Migration<'m>],
+    client: &'m mut Client,
+    migrations: &'m [&'m str],
 }
 
 impl<'m> MigrationRunner<'m> {
     /// Create a new migration runner
-    pub fn new(config: Config, migrations: &'m [Migration<'m>]) -> Result<Self> {
-        // The neon_migration.migration_id::id column is a bigint, which is
-        // equivalent to an i64
-        debug_assert!(migrations.len() + 1 < i64::MAX as usize);
+    pub fn new(client: &'m mut Client, migrations: &'m [&'m str]) -> Self {
+        // The neon_migration.migration_id::id column is a bigint, which is equivalent to an i64
+        assert!(migrations.len() + 1 < i64::MAX as usize);
 
-        Ok(Self { config, migrations })
+        Self { client, migrations }
     }
 
     /// Get the current value neon_migration.migration_id
-    async fn get_migration_id(client: &mut Client) -> Result<i64> {
-        let row = client
+    async fn get_migration_id(&mut self) -> Result<i64> {
+        let row = self
+            .client
             .query_one("SELECT id FROM neon_migration.migration_id", &[])
             .await?;
 
@@ -67,8 +36,9 @@ impl<'m> MigrationRunner<'m> {
     /// used if you would like to fail the application of a series of migrations
     /// at some point.
     async fn update_migration_id(txn: &mut Transaction<'_>, migration_id: i64) -> Result<()> {
-        // We use this fail point in order to check that failing in the middle
-        // of applying a series of migrations fails in an expected manner
+        // We use this fail point in order to check that failing in the
+        // middle of applying a series of migrations fails in an expected
+        // manner
         if cfg!(feature = "testing") {
             let fail = (|| {
                 fail_point!("compute-migration", |fail_migration_id| {
@@ -97,95 +67,48 @@ impl<'m> MigrationRunner<'m> {
     }
 
     /// Prepare the migrations the target database for handling migrations
-    async fn prepare_database(client: &mut Client) -> Result<()> {
-        client
+    async fn prepare_database(&mut self) -> Result<()> {
+        self.client
             .simple_query("CREATE SCHEMA IF NOT EXISTS neon_migration")
             .await?;
-        client.simple_query("CREATE TABLE IF NOT EXISTS neon_migration.migration_id (key INT NOT NULL PRIMARY KEY, id bigint NOT NULL DEFAULT 0)").await?;
-        client
+        self.client.simple_query("CREATE TABLE IF NOT EXISTS neon_migration.migration_id (key INT NOT NULL PRIMARY KEY, id bigint NOT NULL DEFAULT 0)").await?;
+        self.client
             .simple_query(
                 "INSERT INTO neon_migration.migration_id VALUES (0, 0) ON CONFLICT DO NOTHING",
             )
             .await?;
-        client
+        self.client
             .simple_query("ALTER SCHEMA neon_migration OWNER TO cloud_admin")
             .await?;
-        client
+        self.client
             .simple_query("REVOKE ALL ON SCHEMA neon_migration FROM PUBLIC")
             .await?;
 
         Ok(())
     }
 
-    /// Helper function for allowing/disallowing connections to a Postgres
-    /// database.
-    async fn allow_connections_to_db(
-        client: &mut Client,
-        dbname: &PgIdent,
-        allow: bool,
-    ) -> Result<()> {
-        client
-            .simple_query(
-                format!(
-                    "ALTER DATABASE {} WITH ALLOW_CONNECTIONS {}",
-                    dbname.pg_quote(),
-                    allow
-                )
-                .as_str(),
-            )
-            .await?;
-
-        Ok(())
-    }
-
-    /// Connect to the configured Postgres database. Spawns a tokio task to
-    /// handle the connection.
-    async fn connect(config: &Config) -> Result<Client> {
-        let (client, connection) = config.connect(NoTls).await?;
-
-        tokio::spawn(async move {
-            if let Err(e) = connection.await {
-                error!("connection error: {}", e);
-            }
-        });
-
-        Ok(client)
-    }
-
-    async fn run_migration(
-        client: &mut Client,
-        db: &str,
-        migration_id: i64,
-        migration: &str,
-        update_migration_id: bool,
-    ) -> Result<()> {
+    /// Run an individual migration in a separate transaction block.
+    async fn run_migration(client: &mut Client, migration_id: i64, migration: &str) -> Result<()> {
         let mut txn = client
             .transaction()
             .await
             .with_context(|| format!("begin transaction for migration {migration_id}"))?;
 
         if migration.starts_with("-- SKIP") {
-            info!("Skipping migration id={} db=\"{}\"", migration_id, db);
+            info!("Skipping migration id={}", migration_id);
+
+            // Even though we are skipping the migration, updating the
+            // migration ID should help keep logic easy to understand when
+            // trying to understand the state of a cluster.
+            Self::update_migration_id(&mut txn, migration_id).await?;
         } else {
-            info!(
-                "Running migration id={} db=\"{}\":\n{}\n",
-                migration_id, db, migration
-            );
+            info!("Running migration id={}:\n{}\n", migration_id, migration);
 
-            if let Err(e) = txn.simple_query(migration).await {
-                error!("Failed to run the migration: {}", e);
-                return Err(anyhow::anyhow!(e));
-            }
-        }
+            txn.simple_query(migration)
+                .await
+                .with_context(|| format!("apply migration {migration_id}"))?;
 
-        if update_migration_id {
-            if let Err(e) = Self::update_migration_id(&mut txn, migration_id).await {
-                error!(
-                    "Failed to update the migration id to {}: {}",
-                    migration_id, e
-                );
-                return Err(e);
-            }
+            Self::update_migration_id(&mut txn, migration_id).await?;
         }
 
         txn.commit()
@@ -195,201 +118,32 @@ impl<'m> MigrationRunner<'m> {
         Ok(())
     }
 
-    /// Run the migration for the entire cluster. See [`Migration::Cluster`] for
-    /// more information.
-    async fn run_cluster_migration(
-        client: &mut Client,
-        db: &str,
-        migration_id: i64,
-        migration: &str,
-    ) -> Result<()> {
-        Self::run_migration(client, db, migration_id, migration, true).await
-    }
-
-    /// Run the migration in the specified database. See
-    /// [`Migration::PerDatabase`] for more information.
-    async fn run_database_migration(
-        cluster_client: &mut Client,
-        config: Config,
-        db: &Database,
-        migration_id: i64,
-        migration: &str,
-    ) -> Result<()> {
-        // There are 2 race conditions here. Migrations get ran in a separate
-        // thread to not block the ability to connect to the compute. The race
-        // conditions are as follow:
-        //
-        //   1. If between the time we have retrieved the list of databases in
-        //      the cluster and before we set ALLOW_CONNECTIONS back to false,
-        //      the user has changed allowed connections to the database, we
-        //      will have overwritten their change.
-        //
-        //      This is not the end of the world, but an inconvenience,
-        //      nonetheless.
-        //
-        //   2. If between the time we have allowed connections to the database
-        //      and the time the migration is performed, the user disallows
-        //      connections to the database, we will fail to connect to the
-        //      database.
-        //
-        //      This is not much of a problem since we will re-run the migration
-        //      the next time we run migrations.
-        if db.restrict_conn {
-            info!("Allowing connections to \"{}\" for migrations", db.name);
-
-            Self::allow_connections_to_db(cluster_client, &db.name, true)
-                .await
-                .context("Failed to allow connections to the database")?;
-        }
-
-        let mut db_client = Self::connect(&config)
+    /// Run the configured set of migrations
+    pub async fn run_migrations(mut self) -> Result<()> {
+        self.prepare_database()
             .await
-            .context("Failed to connect to the database")?;
-
-        let result = Self::run_migration(&mut db_client, &db.name, migration_id, migration, false)
-            .await
-            .context("Failed to run the migration");
-
-        // Reset the connection restriction
-        if db.restrict_conn {
-            info!(
-                "Disallowing connections to \"{}\" because migration {} is done",
-                db.name, migration_id
-            );
-
-            // Failing here is not the end of the world
-            if let Err(e) = Self::allow_connections_to_db(cluster_client, &db.name, false).await {
-                warn!(
-                    "failed to reset ALLOW_CONNECTIONS on \"{}\": {}",
-                    db.name, e
-                )
-            }
-        }
-
-        result
-    }
-
-    /// Run the configured set of migrations.
-    pub async fn run_migrations(self) -> Result<()> {
-        // Owns the connection to the database containing the
-        // neon_migration.migration_id table. In addition, all Cluster
-        // migrations will be run on this connection.
-        let mut cluster_client = Self::connect(&self.config)
-            .await
-            .context("failed to connect to cluster")?;
-
-        Self::prepare_database(&mut cluster_client)
-            .await
-            .context("failed to prepare database to handle migrations")?;
-
-        let mut current_migration = Self::get_migration_id(&mut cluster_client)
-            .await
-            .context("failed to get the current migration ID")?
-            as usize;
-
-        // All databases within the cluster
-        let dbs: Option<Vec<Database>> = {
-            // Then check if we actually need to run any, and if so, check if
-            // any need to run in each individual database
-            if current_migration < self.migrations.len()
-                && self.migrations[current_migration..]
-                    .iter()
-                    .any(|m| matches!(m, Migration::PerDatabase(_)))
-            {
-                match get_existing_dbs_async(&cluster_client).await {
-                    Ok(dbs) => Some(
-                        // Filter out invalid database (datconnectivity = -2)
-                        dbs.into_values().filter(|d| !d.invalid).collect::<Vec<_>>(),
-                    ),
-                    Err(e) => {
-                        error!("Failed to collect the existing databases: {}", e);
-                        return Err(e);
-                    }
-                }
-            } else {
-                None::<Vec<_>>
-            }
-        };
-
-        let admin_db = self.config.get_dbname().unwrap();
+            .context("prepare database to handle migrations")?;
 
+        let mut current_migration = self.get_migration_id().await? as usize;
         while current_migration < self.migrations.len() {
+            // The index lags the migration ID by 1, so the current migration
+            // ID is also the next index
             let migration_id = (current_migration + 1) as i64;
+            let migration = self.migrations[current_migration];
 
-            let result: Result<()> = match &self.migrations[current_migration] {
-                Migration::Cluster(migration) => {
-                    Self::run_cluster_migration(
-                        &mut cluster_client,
-                        admin_db,
-                        migration_id,
-                        migration,
-                    )
-                    .await
+            match Self::run_migration(self.client, migration_id, migration).await {
+                Ok(_) => {
+                    info!("Finished migration id={}", migration_id);
                 }
-                Migration::PerDatabase(migration) => {
-                    let mut result: Result<()> = Ok(());
-                    for db in dbs.as_ref().unwrap() {
-                        // Once all the databases have run the migration, then we can run it in the
-                        // admin database to mark the migration as complete. See the run for the
-                        // admin database outside this loop.
-                        if db.name == admin_db {
-                            continue;
-                        }
-
-                        let mut config = self.config.clone();
-                        config.dbname(&db.name);
-
-                        // If we failed to run the migration in the current
-                        // database, stop trying to run this migration
-                        if let Err(e) = Self::run_database_migration(
-                            &mut cluster_client,
-                            config,
-                            db,
-                            migration_id,
-                            migration,
-                        )
-                        .await
-                        {
-                            result = Err(e);
-                            break;
-                        }
-                    }
-
-                    match result {
-                        Ok(_) => {
-                            // Finally, run the migration for the admin database,
-                            // and update the migration ID
-                            Self::run_migration(
-                                &mut cluster_client,
-                                admin_db,
-                                migration_id,
-                                migration,
-                                true,
-                            )
-                            .await
-                            .map_err(|e| {
-                                error!("failed to commit the per-database migration: {}", e);
-                                e
-                            })
-                        }
-                        Err(e) => Err(e),
-                    }
+                Err(e) => {
+                    error!("Failed to run migration id={}: {:?}", migration_id, e);
+                    DB_MIGRATION_FAILED
+                        .with_label_values(&[migration_id.to_string().as_str()])
+                        .inc();
+                    return Err(e);
                 }
-            };
-
-            // If failed, mark the metric and return
-            if let Err(e) = result {
-                DB_MIGRATION_FAILED
-                    .with_label_values(&[migration_id.to_string().as_str()])
-                    .inc();
-
-                return Err(anyhow::anyhow!(format!(
-                    "failed at migration {migration_id}: {e}"
-                )));
             }
 
-            info!("Finished migration id={}", migration_id);
-
             current_migration += 1;
         }
 

compute_tools/src/migrations/0012-fix-CVE-2024-4317.sql

@@ -1,235 +0,0 @@
-/*
- * fix-CVE-2024-4317.sql
- *
- * Copyright (c) 2024, PostgreSQL Global Development Group
- *
- * src/backend/catalog/fix-CVE-2024-4317.sql
- *
- * This file should be run in every database in the cluster to address
- * CVE-2024-4317.
- */
-
-DO $$
-DECLARE
-    server_version_num numeric;
-BEGIN
-    SET search_path = pg_catalog;
-
-    SELECT setting::numeric FROM pg_settings INTO server_version_num WHERE name = 'server_version_num';
-
-    -- Everything after Postgres 17 will have the fix
-    IF server_version_num >= 170000 THEN
-        RETURN;
-    END IF;
-
-    -- pg_statistic_ext_data doesn't have the stxdinherit column in 14 and below
-    IF server_version_num < 150000 THEN
-        CREATE OR REPLACE VIEW pg_stats_ext WITH (security_barrier) AS
-            SELECT cn.nspname AS schemaname,
-                   c.relname AS tablename,
-                   sn.nspname AS statistics_schemaname,
-                   s.stxname AS statistics_name,
-                   pg_get_userbyid(s.stxowner) AS statistics_owner,
-                   ( SELECT array_agg(a.attname ORDER BY a.attnum)
-                     FROM unnest(s.stxkeys) k
-                          JOIN pg_attribute a
-                               ON (a.attrelid = s.stxrelid AND a.attnum = k)
-                   ) AS attnames,
-                   pg_get_statisticsobjdef_expressions(s.oid) as exprs,
-                   s.stxkind AS kinds,
-                   sd.stxdndistinct AS n_distinct,
-                   sd.stxddependencies AS dependencies,
-                   m.most_common_vals,
-                   m.most_common_val_nulls,
-                   m.most_common_freqs,
-                   m.most_common_base_freqs
-            FROM pg_statistic_ext s JOIN pg_class c ON (c.oid = s.stxrelid)
-                 JOIN pg_statistic_ext_data sd ON (s.oid = sd.stxoid)
-                 LEFT JOIN pg_namespace cn ON (cn.oid = c.relnamespace)
-                 LEFT JOIN pg_namespace sn ON (sn.oid = s.stxnamespace)
-                 LEFT JOIN LATERAL
-                           ( SELECT array_agg(values) AS most_common_vals,
-                                    array_agg(nulls) AS most_common_val_nulls,
-                                    array_agg(frequency) AS most_common_freqs,
-                                    array_agg(base_frequency) AS most_common_base_freqs
-                             FROM pg_mcv_list_items(sd.stxdmcv)
-                           ) m ON sd.stxdmcv IS NOT NULL
-            WHERE pg_has_role(c.relowner, 'USAGE')
-            AND (c.relrowsecurity = false OR NOT row_security_active(c.oid));
-
-        CREATE OR REPLACE VIEW pg_stats_ext_exprs WITH (security_barrier) AS
-            SELECT cn.nspname AS schemaname,
-                   c.relname AS tablename,
-                   sn.nspname AS statistics_schemaname,
-                   s.stxname AS statistics_name,
-                   pg_get_userbyid(s.stxowner) AS statistics_owner,
-                   stat.expr,
-                   (stat.a).stanullfrac AS null_frac,
-                   (stat.a).stawidth AS avg_width,
-                   (stat.a).stadistinct AS n_distinct,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 1 THEN (stat.a).stavalues1
-                       WHEN (stat.a).stakind2 = 1 THEN (stat.a).stavalues2
-                       WHEN (stat.a).stakind3 = 1 THEN (stat.a).stavalues3
-                       WHEN (stat.a).stakind4 = 1 THEN (stat.a).stavalues4
-                       WHEN (stat.a).stakind5 = 1 THEN (stat.a).stavalues5
-                   END) AS most_common_vals,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 1 THEN (stat.a).stanumbers1
-                       WHEN (stat.a).stakind2 = 1 THEN (stat.a).stanumbers2
-                       WHEN (stat.a).stakind3 = 1 THEN (stat.a).stanumbers3
-                       WHEN (stat.a).stakind4 = 1 THEN (stat.a).stanumbers4
-                       WHEN (stat.a).stakind5 = 1 THEN (stat.a).stanumbers5
-                   END) AS most_common_freqs,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 2 THEN (stat.a).stavalues1
-                       WHEN (stat.a).stakind2 = 2 THEN (stat.a).stavalues2
-                       WHEN (stat.a).stakind3 = 2 THEN (stat.a).stavalues3
-                       WHEN (stat.a).stakind4 = 2 THEN (stat.a).stavalues4
-                       WHEN (stat.a).stakind5 = 2 THEN (stat.a).stavalues5
-                   END) AS histogram_bounds,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 3 THEN (stat.a).stanumbers1[1]
-                       WHEN (stat.a).stakind2 = 3 THEN (stat.a).stanumbers2[1]
-                       WHEN (stat.a).stakind3 = 3 THEN (stat.a).stanumbers3[1]
-                       WHEN (stat.a).stakind4 = 3 THEN (stat.a).stanumbers4[1]
-                       WHEN (stat.a).stakind5 = 3 THEN (stat.a).stanumbers5[1]
-                   END) correlation,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 4 THEN (stat.a).stavalues1
-                       WHEN (stat.a).stakind2 = 4 THEN (stat.a).stavalues2
-                       WHEN (stat.a).stakind3 = 4 THEN (stat.a).stavalues3
-                       WHEN (stat.a).stakind4 = 4 THEN (stat.a).stavalues4
-                       WHEN (stat.a).stakind5 = 4 THEN (stat.a).stavalues5
-                   END) AS most_common_elems,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 4 THEN (stat.a).stanumbers1
-                       WHEN (stat.a).stakind2 = 4 THEN (stat.a).stanumbers2
-                       WHEN (stat.a).stakind3 = 4 THEN (stat.a).stanumbers3
-                       WHEN (stat.a).stakind4 = 4 THEN (stat.a).stanumbers4
-                       WHEN (stat.a).stakind5 = 4 THEN (stat.a).stanumbers5
-                   END) AS most_common_elem_freqs,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 5 THEN (stat.a).stanumbers1
-                       WHEN (stat.a).stakind2 = 5 THEN (stat.a).stanumbers2
-                       WHEN (stat.a).stakind3 = 5 THEN (stat.a).stanumbers3
-                       WHEN (stat.a).stakind4 = 5 THEN (stat.a).stanumbers4
-                       WHEN (stat.a).stakind5 = 5 THEN (stat.a).stanumbers5
-                   END) AS elem_count_histogram
-            FROM pg_statistic_ext s JOIN pg_class c ON (c.oid = s.stxrelid)
-                 LEFT JOIN pg_statistic_ext_data sd ON (s.oid = sd.stxoid)
-                 LEFT JOIN pg_namespace cn ON (cn.oid = c.relnamespace)
-                 LEFT JOIN pg_namespace sn ON (sn.oid = s.stxnamespace)
-                 JOIN LATERAL (
-                     SELECT unnest(pg_get_statisticsobjdef_expressions(s.oid)) AS expr,
-                            unnest(sd.stxdexpr)::pg_statistic AS a
-                 ) stat ON (stat.expr IS NOT NULL)
-            WHERE pg_has_role(c.relowner, 'USAGE')
-            AND (c.relrowsecurity = false OR NOT row_security_active(c.oid));
-    ELSE
-        CREATE OR REPLACE VIEW pg_stats_ext WITH (security_barrier) AS
-            SELECT cn.nspname AS schemaname,
-                   c.relname AS tablename,
-                   sn.nspname AS statistics_schemaname,
-                   s.stxname AS statistics_name,
-                   pg_get_userbyid(s.stxowner) AS statistics_owner,
-                   ( SELECT array_agg(a.attname ORDER BY a.attnum)
-                     FROM unnest(s.stxkeys) k
-                          JOIN pg_attribute a
-                               ON (a.attrelid = s.stxrelid AND a.attnum = k)
-                   ) AS attnames,
-                   pg_get_statisticsobjdef_expressions(s.oid) as exprs,
-                   s.stxkind AS kinds,
-                   sd.stxdinherit AS inherited,
-                   sd.stxdndistinct AS n_distinct,
-                   sd.stxddependencies AS dependencies,
-                   m.most_common_vals,
-                   m.most_common_val_nulls,
-                   m.most_common_freqs,
-                   m.most_common_base_freqs
-            FROM pg_statistic_ext s JOIN pg_class c ON (c.oid = s.stxrelid)
-                 JOIN pg_statistic_ext_data sd ON (s.oid = sd.stxoid)
-                 LEFT JOIN pg_namespace cn ON (cn.oid = c.relnamespace)
-                 LEFT JOIN pg_namespace sn ON (sn.oid = s.stxnamespace)
-                 LEFT JOIN LATERAL
-                           ( SELECT array_agg(values) AS most_common_vals,
-                                    array_agg(nulls) AS most_common_val_nulls,
-                                    array_agg(frequency) AS most_common_freqs,
-                                    array_agg(base_frequency) AS most_common_base_freqs
-                             FROM pg_mcv_list_items(sd.stxdmcv)
-                           ) m ON sd.stxdmcv IS NOT NULL
-            WHERE pg_has_role(c.relowner, 'USAGE')
-            AND (c.relrowsecurity = false OR NOT row_security_active(c.oid));
-
-        CREATE OR REPLACE VIEW pg_stats_ext_exprs WITH (security_barrier) AS
-            SELECT cn.nspname AS schemaname,
-                   c.relname AS tablename,
-                   sn.nspname AS statistics_schemaname,
-                   s.stxname AS statistics_name,
-                   pg_get_userbyid(s.stxowner) AS statistics_owner,
-                   stat.expr,
-                   sd.stxdinherit AS inherited,
-                   (stat.a).stanullfrac AS null_frac,
-                   (stat.a).stawidth AS avg_width,
-                   (stat.a).stadistinct AS n_distinct,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 1 THEN (stat.a).stavalues1
-                       WHEN (stat.a).stakind2 = 1 THEN (stat.a).stavalues2
-                       WHEN (stat.a).stakind3 = 1 THEN (stat.a).stavalues3
-                       WHEN (stat.a).stakind4 = 1 THEN (stat.a).stavalues4
-                       WHEN (stat.a).stakind5 = 1 THEN (stat.a).stavalues5
-                   END) AS most_common_vals,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 1 THEN (stat.a).stanumbers1
-                       WHEN (stat.a).stakind2 = 1 THEN (stat.a).stanumbers2
-                       WHEN (stat.a).stakind3 = 1 THEN (stat.a).stanumbers3
-                       WHEN (stat.a).stakind4 = 1 THEN (stat.a).stanumbers4
-                       WHEN (stat.a).stakind5 = 1 THEN (stat.a).stanumbers5
-                   END) AS most_common_freqs,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 2 THEN (stat.a).stavalues1
-                       WHEN (stat.a).stakind2 = 2 THEN (stat.a).stavalues2
-                       WHEN (stat.a).stakind3 = 2 THEN (stat.a).stavalues3
-                       WHEN (stat.a).stakind4 = 2 THEN (stat.a).stavalues4
-                       WHEN (stat.a).stakind5 = 2 THEN (stat.a).stavalues5
-                   END) AS histogram_bounds,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 3 THEN (stat.a).stanumbers1[1]
-                       WHEN (stat.a).stakind2 = 3 THEN (stat.a).stanumbers2[1]
-                       WHEN (stat.a).stakind3 = 3 THEN (stat.a).stanumbers3[1]
-                       WHEN (stat.a).stakind4 = 3 THEN (stat.a).stanumbers4[1]
-                       WHEN (stat.a).stakind5 = 3 THEN (stat.a).stanumbers5[1]
-                   END) correlation,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 4 THEN (stat.a).stavalues1
-                       WHEN (stat.a).stakind2 = 4 THEN (stat.a).stavalues2
-                       WHEN (stat.a).stakind3 = 4 THEN (stat.a).stavalues3
-                       WHEN (stat.a).stakind4 = 4 THEN (stat.a).stavalues4
-                       WHEN (stat.a).stakind5 = 4 THEN (stat.a).stavalues5
-                   END) AS most_common_elems,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 4 THEN (stat.a).stanumbers1
-                       WHEN (stat.a).stakind2 = 4 THEN (stat.a).stanumbers2
-                       WHEN (stat.a).stakind3 = 4 THEN (stat.a).stanumbers3
-                       WHEN (stat.a).stakind4 = 4 THEN (stat.a).stanumbers4
-                       WHEN (stat.a).stakind5 = 4 THEN (stat.a).stanumbers5
-                   END) AS most_common_elem_freqs,
-                   (CASE
-                       WHEN (stat.a).stakind1 = 5 THEN (stat.a).stanumbers1
-                       WHEN (stat.a).stakind2 = 5 THEN (stat.a).stanumbers2
-                       WHEN (stat.a).stakind3 = 5 THEN (stat.a).stanumbers3
-                       WHEN (stat.a).stakind4 = 5 THEN (stat.a).stanumbers4
-                       WHEN (stat.a).stakind5 = 5 THEN (stat.a).stanumbers5
-                   END) AS elem_count_histogram
-            FROM pg_statistic_ext s JOIN pg_class c ON (c.oid = s.stxrelid)
-                 LEFT JOIN pg_statistic_ext_data sd ON (s.oid = sd.stxoid)
-                 LEFT JOIN pg_namespace cn ON (cn.oid = c.relnamespace)
-                 LEFT JOIN pg_namespace sn ON (sn.oid = s.stxnamespace)
-                 JOIN LATERAL (
-                     SELECT unnest(pg_get_statisticsobjdef_expressions(s.oid)) AS expr,
-                            unnest(sd.stxdexpr)::pg_statistic AS a
-                 ) stat ON (stat.expr IS NOT NULL)
-            WHERE pg_has_role(c.relowner, 'USAGE')
-            AND (c.relrowsecurity = false OR NOT row_security_active(c.oid));
-    END IF;
-END $$;

compute_tools/src/migrations/tests/0012-fix-CVE-2024-4317.sql

@@ -1,5 +0,0 @@
--- Testing that this migration actually works would require spinning up a
--- Postgres instance running on a vulnerable version. Let's trust that the
--- Postgres community created a SQL fix that actually works.
-
-SELECT 1;

compute_tools/src/spec.rs

@@ -6,12 +6,12 @@ use compute_api::responses::{
     ComputeConfig, ControlPlaneComputeStatus, ControlPlaneConfigResponse,
 };
 use reqwest::StatusCode;
-use tokio_postgres::{Client, Config};
+use tokio_postgres::Client;
 use tracing::{error, info, instrument};
 
 use crate::config;
 use crate::metrics::{CPLANE_REQUESTS_TOTAL, CPlaneRequestRPC, UNKNOWN_HTTP_STATUS};
-use crate::migration::{Migration, MigrationRunner};
+use crate::migration::MigrationRunner;
 use crate::params::PG_HBA_ALL_MD5;
 
 // Do control plane request and return response if any. In case of error it
@@ -169,7 +169,7 @@ pub async fn handle_neon_extension_upgrade(client: &mut Client) -> Result<()> {
 }
 
 #[instrument(skip_all)]
-pub async fn handle_migrations(config: Config) -> Result<()> {
+pub async fn handle_migrations(client: &mut Client) -> Result<()> {
     info!("handle migrations");
 
     // !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
@@ -178,50 +178,30 @@ pub async fn handle_migrations(config: Config) -> Result<()> {
 
     // Add new migrations in numerical order.
     let migrations = [
-        Migration::Cluster(include_str!(
-            "./migrations/0001-neon_superuser_bypass_rls.sql"
-        )),
-        Migration::Cluster(include_str!("./migrations/0002-alter_roles.sql")),
-        Migration::Cluster(include_str!(
-            "./migrations/0003-grant_pg_create_subscription_to_neon_superuser.sql"
-        )),
-        Migration::Cluster(include_str!(
-            "./migrations/0004-grant_pg_monitor_to_neon_superuser.sql"
-        )),
-        Migration::Cluster(include_str!(
-            "./migrations/0005-grant_all_on_tables_to_neon_superuser.sql"
-        )),
-        Migration::Cluster(include_str!(
-            "./migrations/0006-grant_all_on_sequences_to_neon_superuser.sql"
-        )),
-        Migration::Cluster(include_str!(
+        include_str!("./migrations/0001-neon_superuser_bypass_rls.sql"),
+        include_str!("./migrations/0002-alter_roles.sql"),
+        include_str!("./migrations/0003-grant_pg_create_subscription_to_neon_superuser.sql"),
+        include_str!("./migrations/0004-grant_pg_monitor_to_neon_superuser.sql"),
+        include_str!("./migrations/0005-grant_all_on_tables_to_neon_superuser.sql"),
+        include_str!("./migrations/0006-grant_all_on_sequences_to_neon_superuser.sql"),
+        include_str!(
             "./migrations/0007-grant_all_on_tables_to_neon_superuser_with_grant_option.sql"
-        )),
-        Migration::Cluster(include_str!(
+        ),
+        include_str!(
             "./migrations/0008-grant_all_on_sequences_to_neon_superuser_with_grant_option.sql"
-        )),
-        Migration::Cluster(include_str!(
-            "./migrations/0009-revoke_replication_for_previously_allowed_roles.sql"
-        )),
-        Migration::Cluster(include_str!(
+        ),
+        include_str!("./migrations/0009-revoke_replication_for_previously_allowed_roles.sql"),
+        include_str!(
             "./migrations/0010-grant_snapshot_synchronization_funcs_to_neon_superuser.sql"
-        )),
-        Migration::Cluster(include_str!(
+        ),
+        include_str!(
             "./migrations/0011-grant_pg_show_replication_origin_status_to_neon_superuser.sql"
-        )),
-        Migration::PerDatabase(include_str!("./migrations/0012-fix-CVE-2024-4317.sql")),
+        ),
     ];
 
-    let runner = match MigrationRunner::new(config, &migrations) {
-        Ok(runner) => runner,
-        Err(e) => {
-            error!("Failed to construct a migration runner: {}", e);
-            return Err(e);
-        }
-    };
+    MigrationRunner::new(client, &migrations)
+        .run_migrations()
+        .await?;
 
-    runner.run_migrations().await.map_err(|e| {
-        error!("Failed to run the migrations: {}", e);
-        e
-    })
+    Ok(())
 }

control_plane/src/bin/neon_local.rs

@@ -18,12 +18,11 @@ use anyhow::{Context, Result, anyhow, bail};
 use clap::Parser;
 use compute_api::spec::ComputeMode;
 use control_plane::endpoint::ComputeControlPlane;
+use control_plane::endpoint_storage::{ENDPOINT_STORAGE_DEFAULT_PORT, EndpointStorage};
 use control_plane::local_env::{
-    InitForceMode, LocalEnv, NeonBroker, NeonLocalInitConf, NeonLocalInitPageserverConf,
-    ObjectStorageConf, SafekeeperConf,
+    EndpointStorageConf, InitForceMode, LocalEnv, NeonBroker, NeonLocalInitConf,
+    NeonLocalInitPageserverConf, SafekeeperConf,
 };
-use control_plane::object_storage::OBJECT_STORAGE_DEFAULT_PORT;
-use control_plane::object_storage::ObjectStorage;
 use control_plane::pageserver::PageServerNode;
 use control_plane::safekeeper::SafekeeperNode;
 use control_plane::storage_controller::{
@@ -93,7 +92,7 @@ enum NeonLocalCmd {
     #[command(subcommand)]
     Safekeeper(SafekeeperCmd),
     #[command(subcommand)]
-    ObjectStorage(ObjectStorageCmd),
+    EndpointStorage(EndpointStorageCmd),
     #[command(subcommand)]
     Endpoint(EndpointCmd),
     #[command(subcommand)]
@@ -460,14 +459,14 @@ enum SafekeeperCmd {
 
 #[derive(clap::Subcommand)]
 #[clap(about = "Manage object storage")]
-enum ObjectStorageCmd {
-    Start(ObjectStorageStartCmd),
-    Stop(ObjectStorageStopCmd),
+enum EndpointStorageCmd {
+    Start(EndpointStorageStartCmd),
+    Stop(EndpointStorageStopCmd),
 }
 
 #[derive(clap::Args)]
 #[clap(about = "Start object storage")]
-struct ObjectStorageStartCmd {
+struct EndpointStorageStartCmd {
     #[clap(short = 't', long, help = "timeout until we fail the command")]
     #[arg(default_value = "10s")]
     start_timeout: humantime::Duration,
@@ -475,7 +474,7 @@ struct ObjectStorageStartCmd {
 
 #[derive(clap::Args)]
 #[clap(about = "Stop object storage")]
-struct ObjectStorageStopCmd {
+struct EndpointStorageStopCmd {
     #[arg(value_enum, default_value = "fast")]
     #[clap(
         short = 'm',
@@ -797,7 +796,9 @@ fn main() -> Result<()> {
             }
             NeonLocalCmd::StorageBroker(subcmd) => rt.block_on(handle_storage_broker(&subcmd, env)),
             NeonLocalCmd::Safekeeper(subcmd) => rt.block_on(handle_safekeeper(&subcmd, env)),
-            NeonLocalCmd::ObjectStorage(subcmd) => rt.block_on(handle_object_storage(&subcmd, env)),
+            NeonLocalCmd::EndpointStorage(subcmd) => {
+                rt.block_on(handle_endpoint_storage(&subcmd, env))
+            }
             NeonLocalCmd::Endpoint(subcmd) => rt.block_on(handle_endpoint(&subcmd, env)),
             NeonLocalCmd::Mappings(subcmd) => handle_mappings(&subcmd, env),
         };
@@ -1014,8 +1015,8 @@ fn handle_init(args: &InitCmdArgs) -> anyhow::Result<LocalEnv> {
                     }
                 })
                 .collect(),
-            object_storage: ObjectStorageConf {
-                port: OBJECT_STORAGE_DEFAULT_PORT,
+            endpoint_storage: EndpointStorageConf {
+                port: ENDPOINT_STORAGE_DEFAULT_PORT,
             },
             pg_distrib_dir: None,
             neon_distrib_dir: None,
@@ -1735,12 +1736,15 @@ async fn handle_safekeeper(subcmd: &SafekeeperCmd, env: &local_env::LocalEnv) ->
     Ok(())
 }
 
-async fn handle_object_storage(subcmd: &ObjectStorageCmd, env: &local_env::LocalEnv) -> Result<()> {
-    use ObjectStorageCmd::*;
-    let storage = ObjectStorage::from_env(env);
+async fn handle_endpoint_storage(
+    subcmd: &EndpointStorageCmd,
+    env: &local_env::LocalEnv,
+) -> Result<()> {
+    use EndpointStorageCmd::*;
+    let storage = EndpointStorage::from_env(env);
 
     // In tests like test_forward_compatibility or test_graceful_cluster_restart
-    // old neon binaries (without object_storage) are present
+    // old neon binaries (without endpoint_storage) are present
     if !storage.bin.exists() {
         eprintln!(
             "{} binary not found. Ignore if this is a compatibility test",
@@ -1750,13 +1754,13 @@ async fn handle_object_storage(subcmd: &ObjectStorageCmd, env: &local_env::Local
     }
 
     match subcmd {
-        Start(ObjectStorageStartCmd { start_timeout }) => {
+        Start(EndpointStorageStartCmd { start_timeout }) => {
             if let Err(e) = storage.start(start_timeout).await {
-                eprintln!("object_storage start failed: {e}");
+                eprintln!("endpoint_storage start failed: {e}");
                 exit(1);
             }
         }
-        Stop(ObjectStorageStopCmd { stop_mode }) => {
+        Stop(EndpointStorageStopCmd { stop_mode }) => {
             let immediate = match stop_mode {
                 StopMode::Fast => false,
                 StopMode::Immediate => true,
@@ -1866,10 +1870,10 @@ async fn handle_start_all_impl(
         }
 
         js.spawn(async move {
-            ObjectStorage::from_env(env)
+            EndpointStorage::from_env(env)
                 .start(&retry_timeout)
                 .await
-                .map_err(|e| e.context("start object_storage"))
+                .map_err(|e| e.context("start endpoint_storage"))
         });
     })();
 
@@ -1968,9 +1972,9 @@ async fn try_stop_all(env: &local_env::LocalEnv, immediate: bool) {
         }
     }
 
-    let storage = ObjectStorage::from_env(env);
+    let storage = EndpointStorage::from_env(env);
     if let Err(e) = storage.stop(immediate) {
-        eprintln!("object_storage stop failed: {:#}", e);
+        eprintln!("endpoint_storage stop failed: {:#}", e);
     }
 
     for ps_conf in &env.pageservers {

control_plane/src/endpoint_storage.rs

@@ -1,34 +1,33 @@
 use crate::background_process::{self, start_process, stop_process};
 use crate::local_env::LocalEnv;
-use anyhow::anyhow;
 use anyhow::{Context, Result};
 use camino::Utf8PathBuf;
 use std::io::Write;
 use std::time::Duration;
 
 /// Directory within .neon which will be used by default for LocalFs remote storage.
-pub const OBJECT_STORAGE_REMOTE_STORAGE_DIR: &str = "local_fs_remote_storage/object_storage";
-pub const OBJECT_STORAGE_DEFAULT_PORT: u16 = 9993;
+pub const ENDPOINT_STORAGE_REMOTE_STORAGE_DIR: &str = "local_fs_remote_storage/endpoint_storage";
+pub const ENDPOINT_STORAGE_DEFAULT_PORT: u16 = 9993;
 
-pub struct ObjectStorage {
+pub struct EndpointStorage {
     pub bin: Utf8PathBuf,
     pub data_dir: Utf8PathBuf,
     pub pemfile: Utf8PathBuf,
     pub port: u16,
 }
 
-impl ObjectStorage {
-    pub fn from_env(env: &LocalEnv) -> ObjectStorage {
-        ObjectStorage {
-            bin: Utf8PathBuf::from_path_buf(env.object_storage_bin()).unwrap(),
-            data_dir: Utf8PathBuf::from_path_buf(env.object_storage_data_dir()).unwrap(),
+impl EndpointStorage {
+    pub fn from_env(env: &LocalEnv) -> EndpointStorage {
+        EndpointStorage {
+            bin: Utf8PathBuf::from_path_buf(env.endpoint_storage_bin()).unwrap(),
+            data_dir: Utf8PathBuf::from_path_buf(env.endpoint_storage_data_dir()).unwrap(),
             pemfile: Utf8PathBuf::from_path_buf(env.public_key_path.clone()).unwrap(),
-            port: env.object_storage.port,
+            port: env.endpoint_storage.port,
         }
     }
 
     fn config_path(&self) -> Utf8PathBuf {
-        self.data_dir.join("object_storage.json")
+        self.data_dir.join("endpoint_storage.json")
     }
 
     fn listen_addr(&self) -> Utf8PathBuf {
@@ -49,7 +48,7 @@ impl ObjectStorage {
         let cfg = Cfg {
             listen: self.listen_addr(),
             pemfile: parent.join(self.pemfile.clone()),
-            local_path: parent.join(OBJECT_STORAGE_REMOTE_STORAGE_DIR),
+            local_path: parent.join(ENDPOINT_STORAGE_REMOTE_STORAGE_DIR),
             r#type: "LocalFs".to_string(),
         };
         std::fs::create_dir_all(self.config_path().parent().unwrap())?;
@@ -59,24 +58,19 @@ impl ObjectStorage {
     }
 
     pub async fn start(&self, retry_timeout: &Duration) -> Result<()> {
-        println!("Starting s3 proxy at {}", self.listen_addr());
+        println!("Starting endpoint_storage at {}", self.listen_addr());
         std::io::stdout().flush().context("flush stdout")?;
 
         let process_status_check = || async {
-            tokio::time::sleep(Duration::from_millis(500)).await;
-            let res = reqwest::Client::new()
-                .get(format!("http://{}/metrics", self.listen_addr()))
-                .send()
-                .await;
-            match res {
-                Ok(response) if response.status().is_success() => Ok(true),
-                Ok(_) => Err(anyhow!("Failed to query /metrics")),
-                Err(e) => Err(anyhow!("Failed to check node status: {e}")),
+            let res = reqwest::Client::new().get(format!("http://{}/metrics", self.listen_addr()));
+            match res.send().await {
+                Ok(res) => Ok(res.status().is_success()),
+                Err(_) => Ok(false),
             }
         };
 
         let res = start_process(
-            "object_storage",
+            "endpoint_storage",
             &self.data_dir.clone().into_std_path_buf(),
             &self.bin.clone().into_std_path_buf(),
             vec![self.config_path().to_string()],
@@ -94,14 +88,14 @@ impl ObjectStorage {
     }
 
     pub fn stop(&self, immediate: bool) -> anyhow::Result<()> {
-        stop_process(immediate, "object_storage", &self.pid_file())
+        stop_process(immediate, "endpoint_storage", &self.pid_file())
     }
 
     fn log_file(&self) -> Utf8PathBuf {
-        self.data_dir.join("object_storage.log")
+        self.data_dir.join("endpoint_storage.log")
     }
 
     fn pid_file(&self) -> Utf8PathBuf {
-        self.data_dir.join("object_storage.pid")
+        self.data_dir.join("endpoint_storage.pid")
     }
 }

control_plane/src/lib.rs

@@ -9,8 +9,8 @@
 mod background_process;
 pub mod broker;
 pub mod endpoint;
+pub mod endpoint_storage;
 pub mod local_env;
-pub mod object_storage;
 pub mod pageserver;
 pub mod postgresql_conf;
 pub mod safekeeper;

control_plane/src/local_env.rs

@@ -19,7 +19,7 @@ use serde::{Deserialize, Serialize};
 use utils::auth::encode_from_key_file;
 use utils::id::{NodeId, TenantId, TenantTimelineId, TimelineId};
 
-use crate::object_storage::{OBJECT_STORAGE_REMOTE_STORAGE_DIR, ObjectStorage};
+use crate::endpoint_storage::{ENDPOINT_STORAGE_REMOTE_STORAGE_DIR, EndpointStorage};
 use crate::pageserver::{PAGESERVER_REMOTE_STORAGE_DIR, PageServerNode};
 use crate::safekeeper::SafekeeperNode;
 
@@ -72,7 +72,7 @@ pub struct LocalEnv {
 
     pub safekeepers: Vec<SafekeeperConf>,
 
-    pub object_storage: ObjectStorageConf,
+    pub endpoint_storage: EndpointStorageConf,
 
     // Control plane upcall API for pageserver: if None, we will not run storage_controller  If set, this will
     // be propagated into each pageserver's configuration.
@@ -110,7 +110,7 @@ pub struct OnDiskConfig {
     )]
     pub pageservers: Vec<PageServerConf>,
     pub safekeepers: Vec<SafekeeperConf>,
-    pub object_storage: ObjectStorageConf,
+    pub endpoint_storage: EndpointStorageConf,
     pub control_plane_api: Option<Url>,
     pub control_plane_hooks_api: Option<Url>,
     pub control_plane_compute_hook_api: Option<Url>,
@@ -144,7 +144,7 @@ pub struct NeonLocalInitConf {
     pub storage_controller: Option<NeonStorageControllerConf>,
     pub pageservers: Vec<NeonLocalInitPageserverConf>,
     pub safekeepers: Vec<SafekeeperConf>,
-    pub object_storage: ObjectStorageConf,
+    pub endpoint_storage: EndpointStorageConf,
     pub control_plane_api: Option<Url>,
     pub control_plane_hooks_api: Option<Url>,
     pub generate_local_ssl_certs: bool,
@@ -152,7 +152,7 @@ pub struct NeonLocalInitConf {
 
 #[derive(Serialize, Default, Deserialize, PartialEq, Eq, Clone, Debug)]
 #[serde(default)]
-pub struct ObjectStorageConf {
+pub struct EndpointStorageConf {
     pub port: u16,
 }
 
@@ -413,8 +413,8 @@ impl LocalEnv {
         self.pg_dir(pg_version, "lib")
     }
 
-    pub fn object_storage_bin(&self) -> PathBuf {
-        self.neon_distrib_dir.join("object_storage")
+    pub fn endpoint_storage_bin(&self) -> PathBuf {
+        self.neon_distrib_dir.join("endpoint_storage")
     }
 
     pub fn pageserver_bin(&self) -> PathBuf {
@@ -450,8 +450,8 @@ impl LocalEnv {
         self.base_data_dir.join("safekeepers").join(data_dir_name)
     }
 
-    pub fn object_storage_data_dir(&self) -> PathBuf {
-        self.base_data_dir.join("object_storage")
+    pub fn endpoint_storage_data_dir(&self) -> PathBuf {
+        self.base_data_dir.join("endpoint_storage")
     }
 
     pub fn get_pageserver_conf(&self, id: NodeId) -> anyhow::Result<&PageServerConf> {
@@ -615,7 +615,7 @@ impl LocalEnv {
                 control_plane_compute_hook_api: _,
                 branch_name_mappings,
                 generate_local_ssl_certs,
-                object_storage,
+                endpoint_storage,
             } = on_disk_config;
             LocalEnv {
                 base_data_dir: repopath.to_owned(),
@@ -632,7 +632,7 @@ impl LocalEnv {
                 control_plane_hooks_api,
                 branch_name_mappings,
                 generate_local_ssl_certs,
-                object_storage,
+                endpoint_storage,
             }
         };
 
@@ -742,7 +742,7 @@ impl LocalEnv {
                 control_plane_compute_hook_api: None,
                 branch_name_mappings: self.branch_name_mappings.clone(),
                 generate_local_ssl_certs: self.generate_local_ssl_certs,
-                object_storage: self.object_storage.clone(),
+                endpoint_storage: self.endpoint_storage.clone(),
             },
         )
     }
@@ -849,7 +849,7 @@ impl LocalEnv {
             control_plane_api,
             generate_local_ssl_certs,
             control_plane_hooks_api,
-            object_storage,
+            endpoint_storage,
         } = conf;
 
         // Find postgres binaries.
@@ -901,7 +901,7 @@ impl LocalEnv {
             control_plane_hooks_api,
             branch_name_mappings: Default::default(),
             generate_local_ssl_certs,
-            object_storage,
+            endpoint_storage,
         };
 
         if generate_local_ssl_certs {
@@ -929,13 +929,13 @@ impl LocalEnv {
                 .context("pageserver init failed")?;
         }
 
-        ObjectStorage::from_env(&env)
+        EndpointStorage::from_env(&env)
             .init()
             .context("object storage init failed")?;
 
         // setup remote remote location for default LocalFs remote storage
         std::fs::create_dir_all(env.base_data_dir.join(PAGESERVER_REMOTE_STORAGE_DIR))?;
-        std::fs::create_dir_all(env.base_data_dir.join(OBJECT_STORAGE_REMOTE_STORAGE_DIR))?;
+        std::fs::create_dir_all(env.base_data_dir.join(ENDPOINT_STORAGE_REMOTE_STORAGE_DIR))?;
 
         env.persist_config()
     }

endpoint_storage/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "object_storage"
+name = "endpoint_storage"
 version = "0.0.1"
 edition.workspace = true
 license.workspace = true

endpoint_storage/src/app.rs

@@ -2,7 +2,7 @@ use anyhow::anyhow;
 use axum::body::{Body, Bytes};
 use axum::response::{IntoResponse, Response};
 use axum::{Router, http::StatusCode};
-use object_storage::{PrefixS3Path, S3Path, Storage, bad_request, internal_error, not_found, ok};
+use endpoint_storage::{PrefixS3Path, S3Path, Storage, bad_request, internal_error, not_found, ok};
 use remote_storage::TimeoutOrCancel;
 use remote_storage::{DownloadError, DownloadOpts, GenericRemoteStorage, RemotePath};
 use std::{sync::Arc, time::SystemTime, time::UNIX_EPOCH};
@@ -46,12 +46,12 @@ async fn metrics() -> Result {
 
 async fn get(S3Path { path }: S3Path, state: State) -> Result {
     info!(%path, "downloading");
-    let download_err = |e| {
-        if let DownloadError::NotFound = e {
-            info!(%path, %e, "downloading"); // 404 is not an issue of _this_ service
+    let download_err = |err| {
+        if let DownloadError::NotFound = err {
+            info!(%path, %err, "downloading"); // 404 is not an issue of _this_ service
             return not_found(&path);
         }
-        internal_error(e, &path, "downloading")
+        internal_error(err, &path, "downloading")
     };
     let cancel = state.cancel.clone();
     let opts = &DownloadOpts::default();
@@ -249,7 +249,7 @@ mod tests {
         };
 
         let proxy = Storage {
-            auth: object_storage::JwtAuth::new(TEST_PUB_KEY_ED25519).unwrap(),
+            auth: endpoint_storage::JwtAuth::new(TEST_PUB_KEY_ED25519).unwrap(),
             storage,
             cancel: cancel.clone(),
             max_upload_file_limit: usize::MAX,
@@ -343,14 +343,14 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
         TimelineId::from_array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3, 4, 5, 7]);
     const ENDPOINT_ID: &str = "ep-winter-frost-a662z3vg";
     fn token() -> String {
-        let claims = object_storage::Claims {
+        let claims = endpoint_storage::Claims {
             tenant_id: TENANT_ID,
             timeline_id: TIMELINE_ID,
             endpoint_id: ENDPOINT_ID.into(),
             exp: u64::MAX,
         };
         let key = jsonwebtoken::EncodingKey::from_ed_pem(TEST_PRIV_KEY_ED25519).unwrap();
-        let header = jsonwebtoken::Header::new(object_storage::VALIDATION_ALGO);
+        let header = jsonwebtoken::Header::new(endpoint_storage::VALIDATION_ALGO);
         jsonwebtoken::encode(&header, &claims, &key).unwrap()
     }
 
@@ -364,7 +364,10 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
             vec![TIMELINE_ID.to_string(), TimelineId::generate().to_string()],
             vec![ENDPOINT_ID, "ep-ololo"]
         )
-        .skip(1);
+        // first one is fully valid path, second path is valid for GET as
+        // read paths may have different endpoint if tenant and timeline matches
+        // (needed for prewarming RO->RW replica)
+        .skip(2);
 
         for ((uri, method), (tenant, timeline, endpoint)) in iproduct!(routes(), args) {
             info!(%uri, %method, %tenant, %timeline, %endpoint);
@@ -475,6 +478,16 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
         requests_chain(chain.into_iter(), |_| token()).await;
     }
 
+    #[testlog(tokio::test)]
+    async fn read_other_endpoint_data() {
+        let uri = format!("/{TENANT_ID}/{TIMELINE_ID}/other_endpoint/key");
+        let chain = vec![
+            (uri.clone(), "GET", "", StatusCode::NOT_FOUND, false),
+            (uri.clone(), "PUT", "", StatusCode::UNAUTHORIZED, false),
+        ];
+        requests_chain(chain.into_iter(), |_| token()).await;
+    }
+
     fn delete_prefix_token(uri: &str) -> String {
         use serde::Serialize;
         let parts = uri.split("/").collect::<Vec<&str>>();
@@ -482,7 +495,7 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
         struct PrefixClaims {
             tenant_id: TenantId,
             timeline_id: Option<TimelineId>,
-            endpoint_id: Option<object_storage::EndpointId>,
+            endpoint_id: Option<endpoint_storage::EndpointId>,
             exp: u64,
         }
         let claims = PrefixClaims {
@@ -492,7 +505,7 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
             exp: u64::MAX,
         };
         let key = jsonwebtoken::EncodingKey::from_ed_pem(TEST_PRIV_KEY_ED25519).unwrap();
-        let header = jsonwebtoken::Header::new(object_storage::VALIDATION_ALGO);
+        let header = jsonwebtoken::Header::new(endpoint_storage::VALIDATION_ALGO);
         jsonwebtoken::encode(&header, &claims, &key).unwrap()
     }
 

endpoint_storage/src/lib.rs

@@ -169,10 +169,19 @@ impl FromRequestParts<Arc<Storage>> for S3Path {
             .auth
             .decode(bearer.token())
             .map_err(|e| bad_request(e, "decoding token"))?;
+
+        // Read paths may have different endpoint ids. For readonly -> readwrite replica
+        // prewarming, endpoint must read other endpoint's data.
+        let endpoint_id = if parts.method == axum::http::Method::GET {
+            claims.endpoint_id.clone()
+        } else {
+            path.endpoint_id.clone()
+        };
+
         let route = Claims {
             tenant_id: path.tenant_id,
             timeline_id: path.timeline_id,
-            endpoint_id: path.endpoint_id.clone(),
+            endpoint_id,
             exp: claims.exp,
         };
         if route != claims {

endpoint_storage/src/main.rs

@@ -1,4 +1,4 @@
-//! `object_storage` is a service which provides API for uploading and downloading
+//! `endpoint_storage` is a service which provides API for uploading and downloading
 //! files. It is used by compute and control plane for accessing LFC prewarm data.
 //! This service is deployed either as a separate component or as part of compute image
 //! for large computes.
@@ -33,7 +33,7 @@ async fn main() -> anyhow::Result<()> {
 
     let config: String = std::env::args().skip(1).take(1).collect();
     if config.is_empty() {
-        anyhow::bail!("Usage: object_storage config.json")
+        anyhow::bail!("Usage: endpoint_storage config.json")
     }
     info!("Reading config from {config}");
     let config = std::fs::read_to_string(config.clone())?;
@@ -41,7 +41,7 @@ async fn main() -> anyhow::Result<()> {
     info!("Reading pemfile from {}", config.pemfile.clone());
     let pemfile = std::fs::read(config.pemfile.clone())?;
     info!("Loading public key from {}", config.pemfile.clone());
-    let auth = object_storage::JwtAuth::new(&pemfile)?;
+    let auth = endpoint_storage::JwtAuth::new(&pemfile)?;
 
     let listener = tokio::net::TcpListener::bind(config.listen).await.unwrap();
     info!("listening on {}", listener.local_addr().unwrap());
@@ -50,7 +50,7 @@ async fn main() -> anyhow::Result<()> {
     let cancel = tokio_util::sync::CancellationToken::new();
     app::check_storage_permissions(&storage, cancel.clone()).await?;
 
-    let proxy = std::sync::Arc::new(object_storage::Storage {
+    let proxy = std::sync::Arc::new(endpoint_storage::Storage {
         auth,
         storage,
         cancel: cancel.clone(),

libs/compute_api/src/spec.rs

@@ -242,13 +242,22 @@ impl RemoteExtSpec {
 
         match self.extension_data.get(real_ext_name) {
             Some(_ext_data) => {
+                // We have decided to use the Go naming convention due to Kubernetes.
+
+                let arch = match std::env::consts::ARCH {
+                    "x86_64" => "amd64",
+                    "aarch64" => "arm64",
+                    arch => arch,
+                };
+
                 // Construct the path to the extension archive
                 // BUILD_TAG/PG_MAJOR_VERSION/extensions/EXTENSION_NAME.tar.zst
                 //
                 // Keep it in sync with path generation in
                 // https://github.com/neondatabase/build-custom-extensions/tree/main
-                let archive_path_str =
-                    format!("{build_tag}/{pg_major_version}/extensions/{real_ext_name}.tar.zst");
+                let archive_path_str = format!(
+                    "{build_tag}/{arch}/{pg_major_version}/extensions/{real_ext_name}.tar.zst"
+                );
                 Ok((
                     real_ext_name.to_string(),
                     RemotePath::from_string(&archive_path_str)?,

pageserver/src/consumption_metrics/metrics.rs

@@ -263,7 +263,9 @@ where
     while let Some((tenant_id, tenant)) = tenants.next().await {
         let mut tenant_resident_size = 0;
 
-        for timeline in tenant.list_timelines() {
+        let timelines = tenant.list_timelines();
+        let timelines_len = timelines.len();
+        for timeline in timelines {
             let timeline_id = timeline.timeline_id;
 
             match TimelineSnapshot::collect(&timeline, ctx) {
@@ -289,6 +291,11 @@ where
             tenant_resident_size += timeline.resident_physical_size();
         }
 
+        if timelines_len == 0 {
+            // Force set it to 1 byte to avoid not being reported -- all timelines are offloaded.
+            tenant_resident_size = 1;
+        }
+
         let snap = TenantSnapshot::collect(&tenant, tenant_resident_size);
         snap.to_metrics(tenant_id, Utc::now(), cache, &mut current_metrics);
     }

pageserver/src/tenant/timeline.rs

@@ -1285,6 +1285,10 @@ impl Timeline {
         reconstruct_state: &mut ValuesReconstructState,
         ctx: &RequestContext,
     ) -> Result<BTreeMap<Key, Result<Bytes, PageReconstructError>>, GetVectoredError> {
+        if query.is_empty() {
+            return Ok(BTreeMap::default());
+        }
+
         let read_path = if self.conf.enable_read_path_debugging || ctx.read_path_debug() {
             Some(ReadPath::new(
                 query.total_keyspace(),

pgxn/neon/pagestore_smgr.c

@@ -803,7 +803,13 @@ neon_create(SMgrRelation reln, ForkNumber forkNum, bool isRedo)
 
 		case RELPERSISTENCE_TEMP:
 		case RELPERSISTENCE_UNLOGGED:
+#ifdef DEBUG_COMPARE_LOCAL
+			mdcreate(reln, forkNum, forkNum == INIT_FORKNUM || isRedo);
+			if (forkNum == MAIN_FORKNUM)
+				mdcreate(reln, INIT_FORKNUM, true);
+#else
 			mdcreate(reln, forkNum, isRedo);
+#endif
 			return;
 
 		default:
@@ -1973,6 +1979,10 @@ neon_start_unlogged_build(SMgrRelation reln)
 		case RELPERSISTENCE_UNLOGGED:
 			unlogged_build_rel = reln;
 			unlogged_build_phase = UNLOGGED_BUILD_NOT_PERMANENT;
+#ifdef DEBUG_COMPARE_LOCAL
+			if (!IsParallelWorker())
+				mdcreate(reln, INIT_FORKNUM, true);
+#endif
 			return;
 
 		default:
@@ -1995,12 +2005,14 @@ neon_start_unlogged_build(SMgrRelation reln)
 	 * FIXME: should we pass isRedo true to create the tablespace dir if it
 	 * doesn't exist? Is it needed?
 	 */
-#ifndef DEBUG_COMPARE_LOCAL
  	if (!IsParallelWorker())
+	{
+#ifndef DEBUG_COMPARE_LOCAL
 		mdcreate(reln, MAIN_FORKNUM, false);
 #else
-	mdcreate(reln, INIT_FORKNUM, false);
+		mdcreate(reln, INIT_FORKNUM, true);
 #endif
+	}
 }
 
 /*
@@ -2099,12 +2111,12 @@ neon_end_unlogged_build(SMgrRelation reln)
 #ifndef DEBUG_COMPARE_LOCAL
 			/* use isRedo == true, so that we drop it immediately */
 			mdunlink(rinfob, forknum, true);
-#else
-			mdunlink(rinfob, INIT_FORKNUM, true);
 #endif
 		}
+#ifdef DEBUG_COMPARE_LOCAL
+		mdunlink(rinfob, INIT_FORKNUM, true);
+#endif
 	}
-
 	unlogged_build_rel = NULL;
 	unlogged_build_phase = UNLOGGED_BUILD_NOT_IN_PROGRESS;
 }

proxy/src/lib.rs

@@ -91,6 +91,7 @@ mod jemalloc;
 mod logging;
 mod metrics;
 mod parse;
+mod pglb;
 mod protocol2;
 mod proxy;
 mod rate_limiter;

proxy/src/pglb/inprocess.rs

@@ -0,0 +1,193 @@
+#![allow(dead_code, reason = "TODO: work in progress")]
+
+use std::pin::{Pin, pin};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::task::{Context, Poll};
+use std::{fmt, io};
+
+use tokio::io::{AsyncRead, AsyncWrite, DuplexStream, ReadBuf};
+use tokio::sync::mpsc;
+
+const STREAM_CHANNEL_SIZE: usize = 16;
+const MAX_STREAM_BUFFER_SIZE: usize = 4096;
+
+#[derive(Debug)]
+pub struct Connection {
+    stream_sender: mpsc::Sender<Stream>,
+    stream_receiver: mpsc::Receiver<Stream>,
+    stream_id_counter: Arc<AtomicUsize>,
+}
+
+impl Connection {
+    pub fn new() -> (Connection, Connection) {
+        let (sender_a, receiver_a) = mpsc::channel(STREAM_CHANNEL_SIZE);
+        let (sender_b, receiver_b) = mpsc::channel(STREAM_CHANNEL_SIZE);
+
+        let stream_id_counter = Arc::new(AtomicUsize::new(1));
+
+        let conn_a = Connection {
+            stream_sender: sender_a,
+            stream_receiver: receiver_b,
+            stream_id_counter: Arc::clone(&stream_id_counter),
+        };
+        let conn_b = Connection {
+            stream_sender: sender_b,
+            stream_receiver: receiver_a,
+            stream_id_counter,
+        };
+
+        (conn_a, conn_b)
+    }
+
+    #[inline]
+    fn next_stream_id(&self) -> StreamId {
+        StreamId(self.stream_id_counter.fetch_add(1, Ordering::Relaxed))
+    }
+
+    #[tracing::instrument(skip_all, fields(stream_id = tracing::field::Empty, err))]
+    pub async fn open_stream(&self) -> io::Result<Stream> {
+        let (local, remote) = tokio::io::duplex(MAX_STREAM_BUFFER_SIZE);
+        let stream_id = self.next_stream_id();
+        tracing::Span::current().record("stream_id", stream_id.0);
+
+        let local = Stream {
+            inner: local,
+            id: stream_id,
+        };
+        let remote = Stream {
+            inner: remote,
+            id: stream_id,
+        };
+
+        self.stream_sender
+            .send(remote)
+            .await
+            .map_err(io::Error::other)?;
+
+        Ok(local)
+    }
+
+    #[tracing::instrument(skip_all, fields(stream_id = tracing::field::Empty, err))]
+    pub async fn accept_stream(&mut self) -> io::Result<Option<Stream>> {
+        Ok(self.stream_receiver.recv().await.inspect(|stream| {
+            tracing::Span::current().record("stream_id", stream.id.0);
+        }))
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct StreamId(usize);
+
+impl fmt::Display for StreamId {
+    #[inline]
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+// TODO: Proper closing. Currently Streams can outlive their Connections.
+// Carry WeakSender and check strong_count?
+#[derive(Debug)]
+pub struct Stream {
+    inner: DuplexStream,
+    id: StreamId,
+}
+
+impl Stream {
+    #[inline]
+    pub fn id(&self) -> StreamId {
+        self.id
+    }
+}
+
+impl AsyncRead for Stream {
+    #[tracing::instrument(level = "debug", skip_all, fields(stream_id = %self.id))]
+    #[inline]
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        pin!(&mut self.inner).poll_read(cx, buf)
+    }
+}
+
+impl AsyncWrite for Stream {
+    #[tracing::instrument(level = "debug", skip_all, fields(stream_id = %self.id))]
+    #[inline]
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, io::Error>> {
+        pin!(&mut self.inner).poll_write(cx, buf)
+    }
+
+    #[tracing::instrument(level = "debug", skip_all, fields(stream_id = %self.id))]
+    #[inline]
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        pin!(&mut self.inner).poll_flush(cx)
+    }
+
+    #[tracing::instrument(level = "debug", skip_all, fields(stream_id = %self.id))]
+    #[inline]
+    fn poll_shutdown(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), io::Error>> {
+        pin!(&mut self.inner).poll_shutdown(cx)
+    }
+
+    #[tracing::instrument(level = "debug", skip_all, fields(stream_id = %self.id))]
+    #[inline]
+    fn poll_write_vectored(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[io::IoSlice<'_>],
+    ) -> Poll<Result<usize, io::Error>> {
+        pin!(&mut self.inner).poll_write_vectored(cx, bufs)
+    }
+
+    #[inline]
+    fn is_write_vectored(&self) -> bool {
+        self.inner.is_write_vectored()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_simple_roundtrip() {
+        let (client, mut server) = Connection::new();
+
+        let server_task = tokio::spawn(async move {
+            while let Some(mut stream) = server.accept_stream().await.unwrap() {
+                tokio::spawn(async move {
+                    let mut buf = [0; 64];
+                    loop {
+                        match stream.read(&mut buf).await.unwrap() {
+                            0 => break,
+                            n => stream.write(&buf[..n]).await.unwrap(),
+                        };
+                    }
+                });
+            }
+        });
+
+        let mut stream = client.open_stream().await.unwrap();
+        stream.write_all(b"hello!").await.unwrap();
+        let mut buf = [0; 64];
+        let n = stream.read(&mut buf).await.unwrap();
+        assert_eq!(n, 6);
+        assert_eq!(&buf[..n], b"hello!");
+
+        drop(stream);
+        drop(client);
+        server_task.await.unwrap();
+    }
+}

proxy/src/pglb/mod.rs

@@ -0,0 +1 @@
+pub mod inprocess;

safekeeper/src/bin/safekeeper.rs

@@ -14,6 +14,7 @@ use clap::{ArgAction, Parser};
 use futures::future::BoxFuture;
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
+use http_utils::tls_certs::ReloadingCertificateResolver;
 use metrics::set_build_info_metric;
 use remote_storage::RemoteStorageConfig;
 use safekeeper::defaults::{
@@ -23,8 +24,8 @@ use safekeeper::defaults::{
     DEFAULT_SSL_CERT_RELOAD_PERIOD, DEFAULT_SSL_KEY_FILE,
 };
 use safekeeper::{
-    BROKER_RUNTIME, GlobalTimelines, HTTP_RUNTIME, SafeKeeperConf, WAL_SERVICE_RUNTIME, broker,
-    control_file, http, wal_backup, wal_service,
+    BACKGROUND_RUNTIME, BROKER_RUNTIME, GlobalTimelines, HTTP_RUNTIME, SafeKeeperConf,
+    WAL_SERVICE_RUNTIME, broker, control_file, http, wal_backup, wal_service,
 };
 use sd_notify::NotifyState;
 use storage_broker::{DEFAULT_ENDPOINT, Uri};
@@ -215,16 +216,21 @@ struct Args {
     ssl_cert_file: Utf8PathBuf,
     /// Period to reload certificate and private key from files.
     #[arg(long, value_parser = humantime::parse_duration, default_value = DEFAULT_SSL_CERT_RELOAD_PERIOD)]
-    pub ssl_cert_reload_period: Duration,
+    ssl_cert_reload_period: Duration,
     /// Trusted root CA certificates to use in https APIs.
     #[arg(long)]
-    pub ssl_ca_file: Option<Utf8PathBuf>,
+    ssl_ca_file: Option<Utf8PathBuf>,
     /// Flag to use https for requests to peer's safekeeper API.
     #[arg(long)]
-    pub use_https_safekeeper_api: bool,
+    use_https_safekeeper_api: bool,
     /// Path to the JWT auth token used to authenticate with other safekeepers.
     #[arg(long)]
     auth_token_path: Option<Utf8PathBuf>,
+    /// Enable TLS in WAL service API.
+    /// Does not force TLS: the client negotiates TLS usage during the handshake.
+    /// Uses key and certificate from ssl_key_file/ssl_cert_file.
+    #[arg(long)]
+    enable_tls_wal_service_api: bool,
 }
 
 // Like PathBufValueParser, but allows empty string.
@@ -418,6 +424,7 @@ async fn main() -> anyhow::Result<()> {
         ssl_cert_reload_period: args.ssl_cert_reload_period,
         ssl_ca_certs,
         use_https_safekeeper_api: args.use_https_safekeeper_api,
+        enable_tls_wal_service_api: args.enable_tls_wal_service_api,
     });
 
     // initialize sentry if SENTRY_DSN is provided
@@ -517,6 +524,36 @@ async fn start_safekeeper(conf: Arc<SafeKeeperConf>) -> Result<()> {
         info!("running in current thread runtime");
     }
 
+    let tls_server_config = if conf.listen_https_addr.is_some() || conf.enable_tls_wal_service_api {
+        let ssl_key_file = conf.ssl_key_file.clone();
+        let ssl_cert_file = conf.ssl_cert_file.clone();
+        let ssl_cert_reload_period = conf.ssl_cert_reload_period;
+
+        // Create resolver in BACKGROUND_RUNTIME, so the background certificate reloading
+        // task is run in this runtime.
+        let cert_resolver = current_thread_rt
+            .as_ref()
+            .unwrap_or_else(|| BACKGROUND_RUNTIME.handle())
+            .spawn(async move {
+                ReloadingCertificateResolver::new(
+                    "main",
+                    &ssl_key_file,
+                    &ssl_cert_file,
+                    ssl_cert_reload_period,
+                )
+                .await
+            })
+            .await??;
+
+        let config = rustls::ServerConfig::builder()
+            .with_no_client_auth()
+            .with_cert_resolver(cert_resolver);
+
+        Some(Arc::new(config))
+    } else {
+        None
+    };
+
     let wal_service_handle = current_thread_rt
         .as_ref()
         .unwrap_or_else(|| WAL_SERVICE_RUNTIME.handle())
@@ -524,6 +561,9 @@ async fn start_safekeeper(conf: Arc<SafeKeeperConf>) -> Result<()> {
             conf.clone(),
             pg_listener,
             Scope::SafekeeperData,
+            conf.enable_tls_wal_service_api
+                .then(|| tls_server_config.clone())
+                .flatten(),
             global_timelines.clone(),
         ))
         // wrap with task name for error reporting
@@ -552,6 +592,9 @@ async fn start_safekeeper(conf: Arc<SafeKeeperConf>) -> Result<()> {
                 conf.clone(),
                 pg_listener_tenant_only,
                 Scope::Tenant,
+                conf.enable_tls_wal_service_api
+                    .then(|| tls_server_config.clone())
+                    .flatten(),
                 global_timelines.clone(),
             ))
             // wrap with task name for error reporting
@@ -577,6 +620,7 @@ async fn start_safekeeper(conf: Arc<SafeKeeperConf>) -> Result<()> {
             .spawn(http::task_main_https(
                 conf.clone(),
                 https_listener,
+                tls_server_config.expect("tls_server_config is set earlier if https is enabled"),
                 global_timelines.clone(),
             ))
             .map(|res| ("HTTPS service main".to_owned(), res));

safekeeper/src/http/mod.rs

@@ -1,7 +1,6 @@
 pub mod routes;
 use std::sync::Arc;
 
-use http_utils::tls_certs::ReloadingCertificateResolver;
 pub use routes::make_router;
 pub use safekeeper_api::models;
 use tokio_util::sync::CancellationToken;
@@ -28,21 +27,10 @@ pub async fn task_main_http(
 pub async fn task_main_https(
     conf: Arc<SafeKeeperConf>,
     https_listener: std::net::TcpListener,
+    tls_config: Arc<rustls::ServerConfig>,
     global_timelines: Arc<GlobalTimelines>,
 ) -> anyhow::Result<()> {
-    let cert_resolver = ReloadingCertificateResolver::new(
-        "main",
-        &conf.ssl_key_file,
-        &conf.ssl_cert_file,
-        conf.ssl_cert_reload_period,
-    )
-    .await?;
-
-    let server_config = rustls::ServerConfig::builder()
-        .with_no_client_auth()
-        .with_cert_resolver(cert_resolver);
-
-    let tls_acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(server_config));
+    let tls_acceptor = tokio_rustls::TlsAcceptor::from(tls_config);
 
     let router = make_router(conf, global_timelines)
         .build()

safekeeper/src/lib.rs

@@ -122,6 +122,7 @@ pub struct SafeKeeperConf {
     pub ssl_cert_reload_period: Duration,
     pub ssl_ca_certs: Vec<Pem>,
     pub use_https_safekeeper_api: bool,
+    pub enable_tls_wal_service_api: bool,
 }
 
 impl SafeKeeperConf {
@@ -172,6 +173,7 @@ impl SafeKeeperConf {
             ssl_cert_reload_period: Duration::from_secs(60),
             ssl_ca_certs: Vec::new(),
             use_https_safekeeper_api: false,
+            enable_tls_wal_service_api: false,
         }
     }
 }
@@ -209,3 +211,12 @@ pub static WAL_BACKUP_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
         .build()
         .expect("Failed to create WAL backup runtime")
 });
+
+pub static BACKGROUND_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
+    tokio::runtime::Builder::new_multi_thread()
+        .thread_name("background worker")
+        .worker_threads(1) // there is only one task now (ssl certificate reloading), having more threads doesn't make sense
+        .enable_all()
+        .build()
+        .expect("Failed to create background runtime")
+});

safekeeper/src/wal_service.rs

@@ -29,6 +29,7 @@ pub async fn task_main(
     conf: Arc<SafeKeeperConf>,
     pg_listener: std::net::TcpListener,
     allowed_auth_scope: Scope,
+    tls_config: Option<Arc<rustls::ServerConfig>>,
     global_timelines: Arc<GlobalTimelines>,
 ) -> anyhow::Result<()> {
     // Tokio's from_std won't do this for us, per its comment.
@@ -43,9 +44,10 @@ pub async fn task_main(
         let conf = conf.clone();
         let conn_id = issue_connection_id(&mut connection_count);
         let global_timelines = global_timelines.clone();
+        let tls_config = tls_config.clone();
         tokio::spawn(
             async move {
-                if let Err(err) = handle_socket(socket, conf, conn_id, allowed_auth_scope, global_timelines).await {
+                if let Err(err) = handle_socket(socket, conf, conn_id, allowed_auth_scope, tls_config, global_timelines).await {
                     error!("connection handler exited: {}", err);
                 }
             }
@@ -61,6 +63,7 @@ async fn handle_socket(
     conf: Arc<SafeKeeperConf>,
     conn_id: ConnectionId,
     allowed_auth_scope: Scope,
+    tls_config: Option<Arc<rustls::ServerConfig>>,
     global_timelines: Arc<GlobalTimelines>,
 ) -> Result<(), QueryError> {
     socket.set_nodelay(true)?;
@@ -110,7 +113,8 @@ async fn handle_socket(
         auth_pair,
         global_timelines,
     );
-    let pgbackend = PostgresBackend::new_from_io(socket_fd, socket, peer_addr, auth_type, None)?;
+    let pgbackend =
+        PostgresBackend::new_from_io(socket_fd, socket, peer_addr, auth_type, tls_config)?;
     // libpq protocol between safekeeper and walproposer / pageserver
     // We don't use shutdown.
     pgbackend

safekeeper/tests/walproposer_sim/safekeeper.rs

@@ -185,6 +185,7 @@ pub fn run_server(os: NodeOs, disk: Arc<SafekeeperDisk>) -> Result<()> {
         ssl_cert_reload_period: Duration::ZERO,
         ssl_ca_certs: Vec::new(),
         use_https_safekeeper_api: false,
+        enable_tls_wal_service_api: false,
     };
 
     let mut global = GlobalMap::new(disk, conf.clone())?;

storage_controller/src/service/safekeeper_service.rs

@@ -151,11 +151,39 @@ impl Service {
             "Got {} non-successful responses from initial creation request of total {total_result_count} responses",
             remaining.len()
         );
-        if remaining.len() >= 2 {
+        let target_sk_count = timeline_persistence.sk_set.len();
+        let quorum_size = match target_sk_count {
+            0 => {
+                return Err(ApiError::InternalServerError(anyhow::anyhow!(
+                    "timeline configured without any safekeepers",
+                )));
+            }
+            1 | 2 => {
+                #[cfg(feature = "testing")]
+                {
+                    // In test settings, it is allowed to have one or two safekeepers
+                    target_sk_count
+                }
+                #[cfg(not(feature = "testing"))]
+                {
+                    // The region is misconfigured: we need at least three safekeepers to be configured
+                    // in order to schedule work to them
+                    tracing::warn!(
+                        "couldn't find at least 3 safekeepers for timeline, found: {:?}",
+                        timeline_persistence.sk_set
+                    );
+                    return Err(ApiError::InternalServerError(anyhow::anyhow!(
+                        "couldn't find at least 3 safekeepers to put timeline to"
+                    )));
+                }
+            }
+            _ => target_sk_count / 2 + 1,
+        };
+        let success_count = target_sk_count - remaining.len();
+        if success_count < quorum_size {
             // Failure
             return Err(ApiError::InternalServerError(anyhow::anyhow!(
-                "not enough successful reconciliations to reach quorum, please retry: {} errored",
-                remaining.len()
+                "not enough successful reconciliations to reach quorum size: {success_count} of {quorum_size} of total {target_sk_count}"
             )));
         }
 
@@ -492,8 +520,6 @@ impl Service {
     pub(crate) async fn safekeepers_for_new_timeline(
         &self,
     ) -> Result<Vec<SafekeeperInfo>, ApiError> {
-        // Number of safekeepers in different AZs we are looking for
-        let wanted_count = 3;
         let mut all_safekeepers = {
             let locked = self.inner.read().unwrap();
             locked
@@ -532,6 +558,19 @@ impl Service {
                 sk.1.id.0,
             )
         });
+        // Number of safekeepers in different AZs we are looking for
+        let wanted_count = match all_safekeepers.len() {
+            0 => {
+                return Err(ApiError::InternalServerError(anyhow::anyhow!(
+                    "couldn't find any active safekeeper for new timeline",
+                )));
+            }
+            // Have laxer requirements on testig mode as we don't want to
+            // spin up three safekeepers for every single test
+            #[cfg(feature = "testing")]
+            1 | 2 => all_safekeepers.len(),
+            _ => 3,
+        };
         let mut sks = Vec::new();
         let mut azs = HashSet::new();
         for (_sk_util, sk_info, az_id) in all_safekeepers.iter() {

test_runner/fixtures/neon_cli.py

@@ -417,14 +417,14 @@ class NeonLocalCli(AbstractNeonCli):
             cmd.append(f"--instance-id={instance_id}")
         return self.raw_cli(cmd)
 
-    def object_storage_start(self, timeout_in_seconds: int | None = None):
-        cmd = ["object-storage", "start"]
+    def endpoint_storage_start(self, timeout_in_seconds: int | None = None):
+        cmd = ["endpoint-storage", "start"]
         if timeout_in_seconds is not None:
             cmd.append(f"--start-timeout={timeout_in_seconds}s")
         return self.raw_cli(cmd)
 
-    def object_storage_stop(self, immediate: bool):
-        cmd = ["object-storage", "stop"]
+    def endpoint_storage_stop(self, immediate: bool):
+        cmd = ["endpoint-storage", "stop"]
         if immediate:
             cmd.extend(["-m", "immediate"])
         return self.raw_cli(cmd)

test_runner/fixtures/neon_fixtures.py

@@ -1029,7 +1029,7 @@ class NeonEnvBuilder:
 
             self.env.broker.assert_no_errors()
 
-            self.env.object_storage.assert_no_errors()
+            self.env.endpoint_storage.assert_no_errors()
 
         try:
             self.overlay_cleanup_teardown()
@@ -1126,7 +1126,7 @@ class NeonEnv:
             pagectl_env_vars["RUST_LOG"] = self.rust_log_override
         self.pagectl = Pagectl(extra_env=pagectl_env_vars, binpath=self.neon_binpath)
 
-        self.object_storage = ObjectStorage(self)
+        self.endpoint_storage = EndpointStorage(self)
 
         # The URL for the pageserver to use as its control_plane_api config
         if config.storage_controller_port_override is not None:
@@ -1183,7 +1183,7 @@ class NeonEnv:
             },
             "safekeepers": [],
             "pageservers": [],
-            "object_storage": {"port": self.port_distributor.get_port()},
+            "endpoint_storage": {"port": self.port_distributor.get_port()},
             "generate_local_ssl_certs": self.generate_local_ssl_certs,
         }
 
@@ -1420,7 +1420,7 @@ class NeonEnv:
                 self.storage_controller.on_safekeeper_deploy(sk_id, body)
                 self.storage_controller.safekeeper_scheduling_policy(sk_id, "Active")
 
-        self.object_storage.start(timeout_in_seconds=timeout_in_seconds)
+        self.endpoint_storage.start(timeout_in_seconds=timeout_in_seconds)
 
     def stop(self, immediate=False, ps_assert_metric_no_errors=False, fail_on_endpoint_errors=True):
         """
@@ -1439,7 +1439,7 @@ class NeonEnv:
         except Exception as e:
             raise_later = e
 
-        self.object_storage.stop(immediate=immediate)
+        self.endpoint_storage.stop(immediate=immediate)
 
         # Stop storage controller before pageservers: we don't want it to spuriously
         # detect a pageserver "failure" during test teardown
@@ -2660,24 +2660,24 @@ class NeonStorageController(MetricsGetter, LogUtils):
         self.stop(immediate=True)
 
 
-class ObjectStorage(LogUtils):
+class EndpointStorage(LogUtils):
     def __init__(self, env: NeonEnv):
-        service_dir = env.repo_dir / "object_storage"
-        super().__init__(logfile=service_dir / "object_storage.log")
-        self.conf_path = service_dir / "object_storage.json"
+        service_dir = env.repo_dir / "endpoint_storage"
+        super().__init__(logfile=service_dir / "endpoint_storage.log")
+        self.conf_path = service_dir / "endpoint_storage.json"
         self.env = env
 
     def base_url(self):
         return json.loads(self.conf_path.read_text())["listen"]
 
     def start(self, timeout_in_seconds: int | None = None):
-        self.env.neon_cli.object_storage_start(timeout_in_seconds)
+        self.env.neon_cli.endpoint_storage_start(timeout_in_seconds)
 
     def stop(self, immediate: bool = False):
-        self.env.neon_cli.object_storage_stop(immediate)
+        self.env.neon_cli.endpoint_storage_stop(immediate)
 
     def assert_no_errors(self):
-        assert_no_errors(self.logfile, "object_storage", [])
+        assert_no_errors(self.logfile, "endpoint_storage", [])
 
 
 class NeonProxiedStorageController(NeonStorageController):

test_runner/performance/test_physical_replication.py

@@ -65,7 +65,7 @@ def test_ro_replica_lag(
     project = neon_api.create_project(pg_version)
     project_id = project["project"]["id"]
     log.info("Project ID: %s", project_id)
-    log.info("Primary endpoint ID: %s", project["project"]["endpoints"][0]["id"])
+    log.info("Primary endpoint ID: %s", project["endpoints"][0]["id"])
     neon_api.wait_for_operation_to_finish(project_id)
     error_occurred = False
     try:
@@ -198,7 +198,7 @@ def test_replication_start_stop(
     project = neon_api.create_project(pg_version)
     project_id = project["project"]["id"]
     log.info("Project ID: %s", project_id)
-    log.info("Primary endpoint ID: %s", project["project"]["endpoints"][0]["id"])
+    log.info("Primary endpoint ID: %s", project["endpoints"][0]["id"])
     neon_api.wait_for_operation_to_finish(project_id)
     try:
         branch_id = project["branch"]["id"]

test_runner/regress/test_compute_migrations.py

@@ -75,7 +75,6 @@ def test_compute_migrations_retry(neon_simple_env: NeonEnv, compute_migrations_d
             pattern = rf"Running migration id={i}"
 
         endpoint.log_contains(pattern)
-        endpoint.log_contains(rf"Finished migration id={i}")
 
 
 @pytest.mark.parametrize(

test_runner/regress/test_download_extensions.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import os
+import platform
 import shutil
 import tarfile
 from typing import TYPE_CHECKING
@@ -58,7 +59,18 @@ def test_remote_extensions(
     extensions_endpoint = f"http://{host}:{port}/pg-ext-s3-gateway"
 
     build_tag = os.environ.get("BUILD_TAG", "latest")
-    archive_route = f"{build_tag}/v{pg_version}/extensions/test_extension.tar.zst"
+
+    # We have decided to use the Go naming convention due to Kubernetes.
+    arch = platform.machine()
+    match arch:
+        case "aarch64":
+            arch = "arm64"
+        case "x86_64":
+            arch = "amd64"
+        case _:
+            pass
+
+    archive_route = f"{build_tag}/{arch}/v{pg_version}/extensions/test_extension.tar.zst"
     tarball = test_output_dir / "test_extension.tar"
     extension_dir = (
         base_dir / "test_runner" / "regress" / "data" / "test_remote_extensions" / "test_extension"

test_runner/regress/test_endpoint_storage.py

@@ -8,7 +8,7 @@ from jwcrypto import jwk, jwt
 
 
 @pytest.mark.asyncio
-async def test_object_storage_insert_retrieve_delete(neon_simple_env: NeonEnv):
+async def test_endpoint_storage_insert_retrieve_delete(neon_simple_env: NeonEnv):
     """
     Inserts, retrieves, and deletes test file using a JWT token
     """
@@ -31,7 +31,7 @@ async def test_object_storage_insert_retrieve_delete(neon_simple_env: NeonEnv):
     token.make_signed_token(key)
     token = token.serialize()
 
-    base_url = env.object_storage.base_url()
+    base_url = env.endpoint_storage.base_url()
     key = f"http://{base_url}/{tenant_id}/{timeline_id}/{endpoint_id}/key"
     headers = {"Authorization": f"Bearer {token}"}
     log.info(f"cache key url {key}")

test_runner/regress/test_neon_cli.py

@@ -138,7 +138,7 @@ def test_cli_start_stop(neon_env_builder: NeonEnvBuilder):
     env.neon_cli.pageserver_stop(env.pageserver.id)
     env.neon_cli.safekeeper_stop()
     env.neon_cli.storage_controller_stop(False)
-    env.neon_cli.object_storage_stop(False)
+    env.neon_cli.endpoint_storage_stop(False)
     env.neon_cli.storage_broker_stop()
 
     # Keep NeonEnv state up to date, it usually owns starting/stopping services
@@ -185,7 +185,7 @@ def test_cli_start_stop_multi(neon_env_builder: NeonEnvBuilder):
     env.neon_cli.safekeeper_stop(neon_env_builder.safekeepers_id_start + 1)
     env.neon_cli.safekeeper_stop(neon_env_builder.safekeepers_id_start + 2)
 
-    env.neon_cli.object_storage_stop(False)
+    env.neon_cli.endpoint_storage_stop(False)
 
     # Stop this to get out of the way of the following `start`
     env.neon_cli.storage_controller_stop(False)

test_runner/regress/test_storage_controller.py

@@ -95,7 +95,7 @@ def test_storage_controller_smoke(
     env.pageservers[1].start()
     for sk in env.safekeepers:
         sk.start()
-    env.object_storage.start()
+    env.endpoint_storage.start()
 
     # The pageservers we started should have registered with the sharding service on startup
     nodes = env.storage_controller.node_list()
@@ -347,7 +347,7 @@ def prepare_onboarding_env(
     env = neon_env_builder.init_configs()
     env.broker.start()
     env.storage_controller.start()
-    env.object_storage.start()
+    env.endpoint_storage.start()
 
     # This is the pageserver where we'll initially create the tenant.  Run it in emergency
     # mode so that it doesn't talk to storage controller, and do not register it.
@@ -1612,16 +1612,18 @@ def test_storage_controller_heartbeats(
     env = neon_env_builder.init_configs()
     env.start()
 
-    # Default log allow list permits connection errors, but this test will use error responses on
-    # the utilization endpoint.
-    env.storage_controller.allowed_errors.append(
-        ".*Call to node.*management API.*failed.*failpoint.*"
-    )
-    # The server starts listening to the socket before sending re-attach request,
-    # but it starts serving HTTP only when re-attach is completed.
-    # If re-attach is slow (last scenario), storcon's heartbeat requests will time out.
-    env.storage_controller.allowed_errors.append(
-        ".*Call to node.*management API.*failed.* Timeout.*"
+    env.storage_controller.allowed_errors.extend(
+        [
+            # Default log allow list permits connection errors, but this test will use error responses on
+            # the utilization endpoint.
+            ".*Call to node.*management API.*failed.*failpoint.*",
+            # The server starts listening to the socket before sending re-attach request,
+            # but it starts serving HTTP only when re-attach is completed.
+            # If re-attach is slow (last scenario), storcon's heartbeat requests will time out.
+            ".*Call to node.*management API.*failed.* Timeout.*",
+            # We will intentionally cause reconcile errors
+            ".*Reconcile error.*",
+        ]
     )
 
     # Initially we have two online pageservers
@@ -4240,6 +4242,63 @@ def test_storcon_create_delete_sk_down(
     wait_until(timeline_deleted_on_sk)
 
 
+@run_only_on_default_postgres("PG version is not interesting here")
+@pytest.mark.parametrize("num_safekeepers", [1, 2, 3])
+@pytest.mark.parametrize("deletetion_subject", [DeletionSubject.TENANT, DeletionSubject.TIMELINE])
+def test_storcon_few_sk(
+    neon_env_builder: NeonEnvBuilder,
+    num_safekeepers: int,
+    deletetion_subject: DeletionSubject,
+):
+    """
+    Test that the storcon can create and delete tenants and timelines with a limited/special number of safekeepers
+      - num_safekeepers: number of safekeepers.
+      - deletion_subject: test that both single timeline and whole tenant deletion work.
+    """
+
+    neon_env_builder.num_safekeepers = num_safekeepers
+    safekeeper_list = list(range(1, num_safekeepers + 1))
+    neon_env_builder.storage_controller_config = {
+        "timelines_onto_safekeepers": True,
+    }
+    env = neon_env_builder.init_start()
+
+    tenant_id = TenantId.generate()
+    timeline_id = TimelineId.generate()
+    env.create_tenant(tenant_id, timeline_id)
+    child_timeline_id = env.create_branch("child_of_main", tenant_id)
+
+    env.safekeepers[0].assert_log_contains(f"creating new timeline {tenant_id}/{timeline_id}")
+
+    config_lines = [
+        "neon.safekeeper_proto_version = 3",
+    ]
+    with env.endpoints.create("main", tenant_id=tenant_id, config_lines=config_lines) as ep:
+        # endpoint should start.
+        ep.start(safekeeper_generation=1, safekeepers=safekeeper_list)
+        ep.safe_psql("CREATE TABLE IF NOT EXISTS t(key int, value text)")
+
+    with env.endpoints.create(
+        "child_of_main", tenant_id=tenant_id, config_lines=config_lines
+    ) as ep:
+        # endpoint should start.
+        ep.start(safekeeper_generation=1, safekeepers=safekeeper_list)
+        ep.safe_psql("CREATE TABLE IF NOT EXISTS t(key int, value text)")
+
+    if deletetion_subject is DeletionSubject.TENANT:
+        env.storage_controller.pageserver_api().tenant_delete(tenant_id)
+    else:
+        env.storage_controller.pageserver_api().timeline_delete(tenant_id, child_timeline_id)
+
+    # ensure that there is log msgs for the third safekeeper too
+    def timeline_deleted_on_sk():
+        env.safekeepers[0].assert_log_contains(
+            f"deleting timeline {tenant_id}/{child_timeline_id} from disk"
+        )
+
+    wait_until(timeline_deleted_on_sk)
+
+
 @pytest.mark.parametrize("wrong_az", [True, False])
 def test_storage_controller_graceful_migration(neon_env_builder: NeonEnvBuilder, wrong_az: bool):
     """
@@ -4461,3 +4520,56 @@ def test_storage_controller_migrate_with_pageserver_restart(
         "shards": [{"node_id": int(secondary.id), "shard_number": 0}],
         "preferred_az": DEFAULT_AZ_ID,
     }
+
+
+def test_storage_controller_shard_scheduling_policy_essential(neon_env_builder: NeonEnvBuilder):
+    """
+    Check if essential scheduling policy works as expected.
+    """
+    neon_env_builder.num_pageservers = 2
+    env = neon_env_builder.init_configs()
+    env.start()
+
+    env.storage_controller.tenant_create(env.initial_tenant)
+    env.storage_controller.tenant_policy_update(
+        env.initial_tenant,
+        {
+            "placement": {"Attached": 1},
+            "scheduling": "Essential",
+        },
+    )
+    env.storage_controller.reconcile_until_idle()
+
+    # Ensure that the tenant is attached to both: one is primary, the other is secondary
+    pageserver_1_attachments = (
+        env.pageservers[0].http_client().tenant_list_locations()["tenant_shards"]
+    )
+    pageserver_2_attachments = (
+        env.pageservers[1].http_client().tenant_list_locations()["tenant_shards"]
+    )
+    assert len(pageserver_1_attachments) == 1
+    assert len(pageserver_2_attachments) == 1
+    primary_pageserver = None
+    if pageserver_1_attachments[0][1]["mode"] == "AttachedSingle":
+        primary_pageserver = 0
+        assert pageserver_2_attachments[0][1]["mode"] == "Secondary"
+    elif pageserver_1_attachments[0][1]["mode"] == "Secondary":
+        primary_pageserver = 1
+        assert pageserver_2_attachments[0][1]["mode"] == "AttachedSingle"
+    else:
+        assert False, "unreachable"
+    secondary_pageserver = 1 - primary_pageserver
+
+    # # Ensure the tenant gets attached to the secondary pageserver
+    # env.pageservers[primary_pageserver].stop()
+    # env.storage_controller.node_configure(
+    #     env.pageservers[primary_pageserver].id, {"availability": "Offline"}
+    # )
+    env.storage_controller.node_drain(env.pageservers[primary_pageserver].id)
+    env.storage_controller.reconcile_until_idle()
+    assert (
+        env.pageservers[secondary_pageserver]
+        .http_client()
+        .tenant_list_locations()["tenant_shards"][0][1]["mode"]
+        == "AttachedSingle"
+    )