Merge pull request #9921 from neondatabase/rc/release-proxy/2024-11-28

Proxy release 2024-11-28

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

@@ -36,8 +36,8 @@ inputs:
     description: 'Region name for real s3 tests'
     required: false
     default: ''
-  rerun_failed:
-    description: 'Whether to rerun failed tests'
+  rerun_flaky:
+    description: 'Whether to rerun flaky tests'
     required: false
     default: 'false'
   pg_version:
@@ -108,7 +108,7 @@ runs:
         COMPATIBILITY_SNAPSHOT_DIR: /tmp/compatibility_snapshot_pg${{ inputs.pg_version }}
         ALLOW_BACKWARD_COMPATIBILITY_BREAKAGE: contains(github.event.pull_request.labels.*.name, 'backward compatibility breakage')
         ALLOW_FORWARD_COMPATIBILITY_BREAKAGE: contains(github.event.pull_request.labels.*.name, 'forward compatibility breakage')
-        RERUN_FAILED: ${{ inputs.rerun_failed }}
+        RERUN_FLAKY: ${{ inputs.rerun_flaky }}
         PG_VERSION: ${{ inputs.pg_version }}
       shell: bash -euxo pipefail {0}
       run: |
@@ -154,8 +154,15 @@ runs:
           EXTRA_PARAMS="--out-dir $PERF_REPORT_DIR $EXTRA_PARAMS"
         fi
 
-        if [ "${RERUN_FAILED}" == "true" ]; then
-          EXTRA_PARAMS="--reruns 2 $EXTRA_PARAMS"
+        if [ "${RERUN_FLAKY}" == "true" ]; then
+          mkdir -p $TEST_OUTPUT
+          poetry run ./scripts/flaky_tests.py "${TEST_RESULT_CONNSTR}" \
+                                              --days 7 \
+                                              --output "$TEST_OUTPUT/flaky.json" \
+                                              --pg-version "${DEFAULT_PG_VERSION}" \
+                                              --build-type "${BUILD_TYPE}"
+
+          EXTRA_PARAMS="--flaky-tests-json $TEST_OUTPUT/flaky.json $EXTRA_PARAMS"
         fi
 
         # We use pytest-split plugin to run benchmarks in parallel on different CI runners

.github/workflows/_build-and-test-locally.yml

@@ -293,7 +293,7 @@ jobs:
           run_with_real_s3: true
           real_s3_bucket: neon-github-ci-tests
           real_s3_region: eu-central-1
-          rerun_failed: true
+          rerun_flaky: true
           pg_version: ${{ matrix.pg_version }}
         env:
           TEST_RESULT_CONNSTR: ${{ secrets.REGRESS_TEST_RESULT_CONNSTR_NEW }}

Makefile

@@ -147,8 +147,6 @@ postgres-%: postgres-configure-% \
 	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/$*/contrib/pg_prewarm install
 	+@echo "Compiling pg_buffercache $*"
 	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/$*/contrib/pg_buffercache install
-	+@echo "Compiling pg_visibility $*"
-	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/$*/contrib/pg_visibility install
 	+@echo "Compiling pageinspect $*"
 	$(MAKE) -C $(POSTGRES_INSTALL_DIR)/build/$*/contrib/pageinspect install
 	+@echo "Compiling amcheck $*"

build-tools.Dockerfile

@@ -57,9 +57,9 @@ RUN mkdir -p /pgcopydb/bin && \
     mkdir -p /pgcopydb/lib && \
     chmod -R 755 /pgcopydb && \
     chown -R nonroot:nonroot /pgcopydb
-
-COPY --from=pgcopydb_builder /usr/lib/postgresql/16/bin/pgcopydb /pgcopydb/bin/pgcopydb
-COPY --from=pgcopydb_builder /pgcopydb/lib/libpq.so.5 /pgcopydb/lib/libpq.so.5
+        
+COPY --from=pgcopydb_builder /usr/lib/postgresql/16/bin/pgcopydb /pgcopydb/bin/pgcopydb 
+COPY --from=pgcopydb_builder /pgcopydb/lib/libpq.so.5 /pgcopydb/lib/libpq.so.5 
 
 # System deps
 #
@@ -258,14 +258,14 @@ WORKDIR /home/nonroot
 
 # Rust
 # Please keep the version of llvm (installed above) in sync with rust llvm (`rustc --version --verbose | grep LLVM`)
-ENV RUSTC_VERSION=1.83.0
+ENV RUSTC_VERSION=1.82.0
 ENV RUSTUP_HOME="/home/nonroot/.rustup"
 ENV PATH="/home/nonroot/.cargo/bin:${PATH}"
 ARG RUSTFILT_VERSION=0.2.1
-ARG CARGO_HAKARI_VERSION=0.9.33
-ARG CARGO_DENY_VERSION=0.16.2
-ARG CARGO_HACK_VERSION=0.6.33
-ARG CARGO_NEXTEST_VERSION=0.9.85
+ARG CARGO_HAKARI_VERSION=0.9.30
+ARG CARGO_DENY_VERSION=0.16.1
+ARG CARGO_HACK_VERSION=0.6.31
+ARG CARGO_NEXTEST_VERSION=0.9.72
 RUN curl -sSO https://static.rust-lang.org/rustup/dist/$(uname -m)-unknown-linux-gnu/rustup-init && whoami && \
 	chmod +x rustup-init && \
 	./rustup-init -y --default-toolchain ${RUSTC_VERSION} && \
@@ -289,7 +289,7 @@ RUN whoami \
     && cargo --version --verbose \
     && rustup --version --verbose \
     && rustc --version --verbose \
-    && clang --version
+    && clang --version 
 
 RUN if [ "${DEBIAN_VERSION}" = "bookworm" ]; then \
     LD_LIBRARY_PATH=/pgcopydb/lib /pgcopydb/bin/pgcopydb --version; \

compute_tools/src/catalog.rs

@@ -1,3 +1,4 @@
+use compute_api::responses::CatalogObjects;
 use futures::Stream;
 use postgres::NoTls;
 use std::{path::Path, process::Stdio, result::Result, sync::Arc};
@@ -12,8 +13,7 @@ use tokio_util::codec::{BytesCodec, FramedRead};
 use tracing::warn;
 
 use crate::compute::ComputeNode;
-use crate::pg_helpers::{get_existing_dbs_async, get_existing_roles_async, postgres_conf_for_db};
-use compute_api::responses::CatalogObjects;
+use crate::pg_helpers::{get_existing_dbs_async, get_existing_roles_async};
 
 pub async fn get_dbs_and_roles(compute: &Arc<ComputeNode>) -> anyhow::Result<CatalogObjects> {
     let connstr = compute.connstr.clone();
@@ -43,8 +43,6 @@ pub enum SchemaDumpError {
     DatabaseDoesNotExist,
     #[error("Failed to execute pg_dump.")]
     IO(#[from] std::io::Error),
-    #[error("Unexpected error.")]
-    Unexpected,
 }
 
 // It uses the pg_dump utility to dump the schema of the specified database.
@@ -62,38 +60,11 @@ pub async fn get_database_schema(
     let pgbin = &compute.pgbin;
     let basepath = Path::new(pgbin).parent().unwrap();
     let pgdump = basepath.join("pg_dump");
-
-    // Replace the DB in the connection string and disable it to parts.
-    // This is the only option to handle DBs with special characters.
-    let conf =
-        postgres_conf_for_db(&compute.connstr, dbname).map_err(|_| SchemaDumpError::Unexpected)?;
-    let host = conf
-        .get_hosts()
-        .first()
-        .ok_or(SchemaDumpError::Unexpected)?;
-    let host = match host {
-        tokio_postgres::config::Host::Tcp(ip) => ip.to_string(),
-        #[cfg(unix)]
-        tokio_postgres::config::Host::Unix(path) => path.to_string_lossy().to_string(),
-    };
-    let port = conf
-        .get_ports()
-        .first()
-        .ok_or(SchemaDumpError::Unexpected)?;
-    let user = conf.get_user().ok_or(SchemaDumpError::Unexpected)?;
-    let dbname = conf.get_dbname().ok_or(SchemaDumpError::Unexpected)?;
-
+    let mut connstr = compute.connstr.clone();
+    connstr.set_path(dbname);
     let mut cmd = Command::new(pgdump)
-        // XXX: this seems to be the only option to deal with DBs with `=` in the name
-        // See <https://www.postgresql.org/message-id/flat/20151023003445.931.91267%40wrigleys.postgresql.org>
-        .env("PGDATABASE", dbname)
-        .arg("--host")
-        .arg(host)
-        .arg("--port")
-        .arg(port.to_string())
-        .arg("--username")
-        .arg(user)
         .arg("--schema-only")
+        .arg(connstr.as_str())
         .stdout(Stdio::piped())
         .stderr(Stdio::piped())
         .kill_on_drop(true)

compute_tools/src/compute.rs

@@ -34,8 +34,9 @@ use utils::measured_stream::MeasuredReader;
 use nix::sys::signal::{kill, Signal};
 use remote_storage::{DownloadError, RemotePath};
 use tokio::spawn;
+use url::Url;
 
-use crate::installed_extensions::get_installed_extensions;
+use crate::installed_extensions::get_installed_extensions_sync;
 use crate::local_proxy;
 use crate::pg_helpers::*;
 use crate::spec::*;
@@ -815,32 +816,30 @@ impl ComputeNode {
         Ok(())
     }
 
-    async fn get_maintenance_client(
-        conf: &tokio_postgres::Config,
-    ) -> Result<tokio_postgres::Client> {
-        let mut conf = conf.clone();
+    async fn get_maintenance_client(url: &Url) -> Result<tokio_postgres::Client> {
+        let mut connstr = url.clone();
 
-        conf.application_name("apply_config");
+        connstr
+            .query_pairs_mut()
+            .append_pair("application_name", "apply_config");
 
-        let (client, conn) = match conf.connect(NoTls).await {
-            // If connection fails, it may be the old node with `zenith_admin` superuser.
-            //
-            // In this case we need to connect with old `zenith_admin` name
-            // and create new user. We cannot simply rename connected user,
-            // but we can create a new one and grant it all privileges.
+        let (client, conn) = match tokio_postgres::connect(connstr.as_str(), NoTls).await {
             Err(e) => match e.code() {
                 Some(&SqlState::INVALID_PASSWORD)
                 | Some(&SqlState::INVALID_AUTHORIZATION_SPECIFICATION) => {
-                    // Connect with zenith_admin if cloud_admin could not authenticate
+                    // connect with zenith_admin if cloud_admin could not authenticate
                     info!(
                         "cannot connect to postgres: {}, retrying with `zenith_admin` username",
                         e
                     );
-                    let mut zenith_admin_conf = postgres::config::Config::from(conf.clone());
-                    zenith_admin_conf.user("zenith_admin");
+                    let mut zenith_admin_connstr = connstr.clone();
+
+                    zenith_admin_connstr
+                        .set_username("zenith_admin")
+                        .map_err(|_| anyhow::anyhow!("invalid connstr"))?;
 
                     let mut client =
-                        zenith_admin_conf.connect(NoTls)
+                        Client::connect(zenith_admin_connstr.as_str(), NoTls)
                             .context("broken cloud_admin credential: tried connecting with cloud_admin but could not authenticate, and zenith_admin does not work either")?;
 
                     // Disable forwarding so that users don't get a cloud_admin role
@@ -854,8 +853,8 @@ impl ComputeNode {
 
                     drop(client);
 
-                    // Reconnect with connstring with expected name
-                    conf.connect(NoTls).await?
+                    // reconnect with connstring with expected name
+                    tokio_postgres::connect(connstr.as_str(), NoTls).await?
                 }
                 _ => return Err(e.into()),
             },
@@ -886,7 +885,7 @@ impl ComputeNode {
     pub fn apply_spec_sql(
         &self,
         spec: Arc<ComputeSpec>,
-        conf: Arc<tokio_postgres::Config>,
+        url: Arc<Url>,
         concurrency: usize,
     ) -> Result<()> {
         let rt = tokio::runtime::Builder::new_multi_thread()
@@ -898,7 +897,7 @@ impl ComputeNode {
 
         rt.block_on(async {
             // Proceed with post-startup configuration. Note, that order of operations is important.
-            let client = Self::get_maintenance_client(&conf).await?;
+            let client = Self::get_maintenance_client(&url).await?;
             let spec = spec.clone();
 
             let databases = get_existing_dbs_async(&client).await?;
@@ -932,7 +931,7 @@ impl ComputeNode {
                 RenameAndDeleteDatabases,
                 CreateAndAlterDatabases,
             ] {
-                info!("Applying phase {:?}", &phase);
+                debug!("Applying phase {:?}", &phase);
                 apply_operations(
                     spec.clone(),
                     ctx.clone(),
@@ -943,7 +942,6 @@ impl ComputeNode {
                 .await?;
             }
 
-            info!("Applying RunInEachDatabase phase");
             let concurrency_token = Arc::new(tokio::sync::Semaphore::new(concurrency));
 
             let db_processes = spec
@@ -957,7 +955,7 @@ impl ComputeNode {
                     let spec = spec.clone();
                     let ctx = ctx.clone();
                     let jwks_roles = jwks_roles.clone();
-                    let mut conf = conf.as_ref().clone();
+                    let mut url = url.as_ref().clone();
                     let concurrency_token = concurrency_token.clone();
                     let db = db.clone();
 
@@ -966,14 +964,14 @@ impl ComputeNode {
                     match &db {
                         DB::SystemDB => {}
                         DB::UserDB(db) => {
-                            conf.dbname(db.name.as_str());
+                            url.set_path(db.name.as_str());
                         }
                     }
 
-                    let conf = Arc::new(conf);
+                    let url = Arc::new(url);
                     let fut = Self::apply_spec_sql_db(
                         spec.clone(),
-                        conf,
+                        url,
                         ctx.clone(),
                         jwks_roles.clone(),
                         concurrency_token.clone(),
@@ -1019,7 +1017,7 @@ impl ComputeNode {
     /// semaphore.  The caller has to make sure the semaphore isn't exhausted.
     async fn apply_spec_sql_db(
         spec: Arc<ComputeSpec>,
-        conf: Arc<tokio_postgres::Config>,
+        url: Arc<Url>,
         ctx: Arc<tokio::sync::RwLock<MutableApplyContext>>,
         jwks_roles: Arc<HashSet<String>>,
         concurrency_token: Arc<tokio::sync::Semaphore>,
@@ -1048,7 +1046,7 @@ impl ComputeNode {
                 // that database.
                 || async {
                     if client_conn.is_none() {
-                        let db_client = Self::get_maintenance_client(&conf).await?;
+                        let db_client = Self::get_maintenance_client(&url).await?;
                         client_conn.replace(db_client);
                     }
                     let client = client_conn.as_ref().unwrap();
@@ -1063,16 +1061,34 @@ impl ComputeNode {
         Ok::<(), anyhow::Error>(())
     }
 
-    /// Choose how many concurrent connections to use for applying the spec changes.
-    pub fn max_service_connections(
-        &self,
-        compute_state: &ComputeState,
-        spec: &ComputeSpec,
-    ) -> usize {
-        // If the cluster is in Init state we don't have to deal with user connections,
+    /// Do initial configuration of the already started Postgres.
+    #[instrument(skip_all)]
+    pub fn apply_config(&self, compute_state: &ComputeState) -> Result<()> {
+        // If connection fails,
+        // it may be the old node with `zenith_admin` superuser.
+        //
+        // In this case we need to connect with old `zenith_admin` name
+        // and create new user. We cannot simply rename connected user,
+        // but we can create a new one and grant it all privileges.
+        let mut url = self.connstr.clone();
+        url.query_pairs_mut()
+            .append_pair("application_name", "apply_config");
+
+        let url = Arc::new(url);
+        let spec = Arc::new(
+            compute_state
+                .pspec
+                .as_ref()
+                .expect("spec must be set")
+                .spec
+                .clone(),
+        );
+
+        // Choose how many concurrent connections to use for applying the spec changes.
+        // If the cluster is not currently Running we don't have to deal with user connections,
         // and can thus use all `max_connections` connection slots. However, that's generally not
         // very efficient, so we generally still limit it to a smaller number.
-        if compute_state.status == ComputeStatus::Init {
+        let max_concurrent_connections = if compute_state.status != ComputeStatus::Running {
             // If the settings contain 'max_connections', use that as template
             if let Some(config) = spec.cluster.settings.find("max_connections") {
                 config.parse::<usize>().ok()
@@ -1128,29 +1144,10 @@ impl ComputeNode {
                 .map(|val| if val > 1 { val - 1 } else { 1 })
                 .last()
                 .unwrap_or(3)
-        }
-    }
-
-    /// Do initial configuration of the already started Postgres.
-    #[instrument(skip_all)]
-    pub fn apply_config(&self, compute_state: &ComputeState) -> Result<()> {
-        let mut conf = tokio_postgres::Config::from_str(self.connstr.as_str()).unwrap();
-        conf.application_name("apply_config");
-
-        let conf = Arc::new(conf);
-        let spec = Arc::new(
-            compute_state
-                .pspec
-                .as_ref()
-                .expect("spec must be set")
-                .spec
-                .clone(),
-        );
-
-        let max_concurrent_connections = self.max_service_connections(compute_state, &spec);
+        };
 
         // Merge-apply spec & changes to PostgreSQL state.
-        self.apply_spec_sql(spec.clone(), conf.clone(), max_concurrent_connections)?;
+        self.apply_spec_sql(spec.clone(), url.clone(), max_concurrent_connections)?;
 
         if let Some(ref local_proxy) = &spec.clone().local_proxy_config {
             info!("configuring local_proxy");
@@ -1159,11 +1156,12 @@ impl ComputeNode {
 
         // Run migrations separately to not hold up cold starts
         thread::spawn(move || {
-            let conf = conf.as_ref().clone();
-            let mut conf = postgres::config::Config::from(conf);
-            conf.application_name("migrations");
+            let mut connstr = url.as_ref().clone();
+            connstr
+                .query_pairs_mut()
+                .append_pair("application_name", "migrations");
 
-            let mut client = conf.connect(NoTls)?;
+            let mut client = Client::connect(connstr.as_str(), NoTls)?;
             handle_migrations(&mut client).context("apply_config handle_migrations")
         });
 
@@ -1224,28 +1222,21 @@ impl ComputeNode {
         let pgdata_path = Path::new(&self.pgdata);
         let postgresql_conf_path = pgdata_path.join("postgresql.conf");
         config::write_postgres_conf(&postgresql_conf_path, &spec, None)?;
-
-        // TODO(ololobus): We need a concurrency during reconfiguration as well,
-        // but DB is already running and used by user. We can easily get out of
-        // `max_connections` limit, and the current code won't handle that.
-        // let compute_state = self.state.lock().unwrap().clone();
-        // let max_concurrent_connections = self.max_service_connections(&compute_state, &spec);
-        let max_concurrent_connections = 1;
-
-        // Temporarily reset max_cluster_size in config
+        // temporarily reset max_cluster_size in config
         // to avoid the possibility of hitting the limit, while we are reconfiguring:
-        // creating new extensions, roles, etc.
+        // creating new extensions, roles, etc...
         config::with_compute_ctl_tmp_override(pgdata_path, "neon.max_cluster_size=-1", || {
             self.pg_reload_conf()?;
 
             if spec.mode == ComputeMode::Primary {
-                let mut conf = tokio_postgres::Config::from_str(self.connstr.as_str()).unwrap();
-                conf.application_name("apply_config");
-                let conf = Arc::new(conf);
+                let mut url = self.connstr.clone();
+                url.query_pairs_mut()
+                    .append_pair("application_name", "apply_config");
+                let url = Arc::new(url);
 
                 let spec = Arc::new(spec.clone());
 
-                self.apply_spec_sql(spec, conf, max_concurrent_connections)?;
+                self.apply_spec_sql(spec, url, 1)?;
             }
 
             Ok(())
@@ -1371,17 +1362,7 @@ impl ComputeNode {
 
             let connstr = self.connstr.clone();
             thread::spawn(move || {
-                let res = get_installed_extensions(&connstr);
-                match res {
-                    Ok(extensions) => {
-                        info!(
-                            "[NEON_EXT_STAT] {}",
-                            serde_json::to_string(&extensions)
-                                .expect("failed to serialize extensions list")
-                        );
-                    }
-                    Err(err) => error!("could not get installed extensions: {err:?}"),
-                }
+                get_installed_extensions_sync(connstr).context("get_installed_extensions")
             });
         }
 

compute_tools/src/http/api.rs

@@ -296,12 +296,7 @@ async fn routes(req: Request<Body>, compute: &Arc<ComputeNode>) -> Response<Body
             }
 
             let connstr = compute.connstr.clone();
-            let res = task::spawn_blocking(move || {
-                installed_extensions::get_installed_extensions(&connstr)
-            })
-            .await
-            .unwrap();
-
+            let res = crate::installed_extensions::get_installed_extensions(connstr).await;
             match res {
                 Ok(res) => render_json(Body::from(serde_json::to_string(&res).unwrap())),
                 Err(e) => render_json_error(

compute_tools/src/installed_extensions.rs

@@ -2,16 +2,17 @@ use compute_api::responses::{InstalledExtension, InstalledExtensions};
 use metrics::proto::MetricFamily;
 use std::collections::HashMap;
 use std::collections::HashSet;
+use tracing::info;
+use url::Url;
 
 use anyhow::Result;
 use postgres::{Client, NoTls};
+use tokio::task;
 
 use metrics::core::Collector;
 use metrics::{register_uint_gauge_vec, UIntGaugeVec};
 use once_cell::sync::Lazy;
 
-use crate::pg_helpers::postgres_conf_for_db;
-
 /// We don't reuse get_existing_dbs() just for code clarity
 /// and to make database listing query here more explicit.
 ///
@@ -41,51 +42,75 @@ fn list_dbs(client: &mut Client) -> Result<Vec<String>> {
 ///
 /// Same extension can be installed in multiple databases with different versions,
 /// we only keep the highest and lowest version across all databases.
-pub fn get_installed_extensions(connstr: &url::Url) -> Result<InstalledExtensions> {
-    let mut client = Client::connect(connstr.as_str(), NoTls)?;
-    let databases: Vec<String> = list_dbs(&mut client)?;
+pub async fn get_installed_extensions(connstr: Url) -> Result<InstalledExtensions> {
+    let mut connstr = connstr.clone();
 
-    let mut extensions_map: HashMap<String, InstalledExtension> = HashMap::new();
-    for db in databases.iter() {
-        let config = postgres_conf_for_db(connstr, db)?;
-        let mut db_client = config.connect(NoTls)?;
-        let extensions: Vec<(String, String)> = db_client
-            .query(
-                "SELECT extname, extversion FROM pg_catalog.pg_extension;",
-                &[],
-            )?
-            .iter()
-            .map(|row| (row.get("extname"), row.get("extversion")))
-            .collect();
+    task::spawn_blocking(move || {
+        let mut client = Client::connect(connstr.as_str(), NoTls)?;
+        let databases: Vec<String> = list_dbs(&mut client)?;
 
-        for (extname, v) in extensions.iter() {
-            let version = v.to_string();
+        let mut extensions_map: HashMap<String, InstalledExtension> = HashMap::new();
+        for db in databases.iter() {
+            connstr.set_path(db);
+            let mut db_client = Client::connect(connstr.as_str(), NoTls)?;
+            let extensions: Vec<(String, String)> = db_client
+                .query(
+                    "SELECT extname, extversion FROM pg_catalog.pg_extension;",
+                    &[],
+                )?
+                .iter()
+                .map(|row| (row.get("extname"), row.get("extversion")))
+                .collect();
 
-            // increment the number of databases where the version of extension is installed
-            INSTALLED_EXTENSIONS
-                .with_label_values(&[extname, &version])
-                .inc();
+            for (extname, v) in extensions.iter() {
+                let version = v.to_string();
 
-            extensions_map
-                .entry(extname.to_string())
-                .and_modify(|e| {
-                    e.versions.insert(version.clone());
-                    // count the number of databases where the extension is installed
-                    e.n_databases += 1;
-                })
-                .or_insert(InstalledExtension {
-                    extname: extname.to_string(),
-                    versions: HashSet::from([version.clone()]),
-                    n_databases: 1,
-                });
+                // increment the number of databases where the version of extension is installed
+                INSTALLED_EXTENSIONS
+                    .with_label_values(&[extname, &version])
+                    .inc();
+
+                extensions_map
+                    .entry(extname.to_string())
+                    .and_modify(|e| {
+                        e.versions.insert(version.clone());
+                        // count the number of databases where the extension is installed
+                        e.n_databases += 1;
+                    })
+                    .or_insert(InstalledExtension {
+                        extname: extname.to_string(),
+                        versions: HashSet::from([version.clone()]),
+                        n_databases: 1,
+                    });
+            }
         }
-    }
 
-    let res = InstalledExtensions {
-        extensions: extensions_map.values().cloned().collect(),
-    };
+        let res = InstalledExtensions {
+            extensions: extensions_map.values().cloned().collect(),
+        };
 
-    Ok(res)
+        Ok(res)
+    })
+    .await?
+}
+
+// Gather info about installed extensions
+pub fn get_installed_extensions_sync(connstr: Url) -> Result<()> {
+    let rt = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .expect("failed to create runtime");
+    let result = rt
+        .block_on(crate::installed_extensions::get_installed_extensions(
+            connstr,
+        ))
+        .expect("failed to get installed extensions");
+
+    info!(
+        "[NEON_EXT_STAT] {}",
+        serde_json::to_string(&result).expect("failed to serialize extensions list")
+    );
+    Ok(())
 }
 
 static INSTALLED_EXTENSIONS: Lazy<UIntGaugeVec> = Lazy::new(|| {

compute_tools/src/pg_helpers.rs

@@ -6,7 +6,6 @@ use std::io::{BufRead, BufReader};
 use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
 use std::process::Child;
-use std::str::FromStr;
 use std::thread::JoinHandle;
 use std::time::{Duration, Instant};
 
@@ -14,10 +13,8 @@ use anyhow::{bail, Result};
 use futures::StreamExt;
 use ini::Ini;
 use notify::{RecursiveMode, Watcher};
-use postgres::config::Config;
 use tokio::io::AsyncBufReadExt;
 use tokio::time::timeout;
-use tokio_postgres;
 use tokio_postgres::NoTls;
 use tracing::{debug, error, info, instrument};
 
@@ -545,11 +542,3 @@ async fn handle_postgres_logs_async(stderr: tokio::process::ChildStderr) -> Resu
 
     Ok(())
 }
-
-/// `Postgres::config::Config` handles database names with whitespaces
-/// and special characters properly.
-pub fn postgres_conf_for_db(connstr: &url::Url, dbname: &str) -> Result<Config> {
-    let mut conf = Config::from_str(connstr.as_str())?;
-    conf.dbname(dbname);
-    Ok(conf)
-}

libs/remote_storage/src/azure_blob.rs

@@ -220,11 +220,6 @@ impl AzureBlobStorage {
         let started_at = ScopeGuard::into_inner(started_at);
         let outcome = match &download {
             Ok(_) => AttemptOutcome::Ok,
-            // At this level in the stack 404 and 304 responses do not indicate an error.
-            // There's expected cases when a blob may not exist or hasn't been modified since
-            // the last get (e.g. probing for timeline indices and heatmap downloads).
-            // Callers should handle errors if they are unexpected.
-            Err(DownloadError::NotFound | DownloadError::Unmodified) => AttemptOutcome::Ok,
             Err(_) => AttemptOutcome::Err,
         };
         crate::metrics::BUCKET_METRICS

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

@@ -454,10 +454,6 @@ pub(super) async fn handle_walreceiver_connection(
                     timeline.get_last_record_lsn()
                 );
 
-                if let Some(lsn) = next_record_lsn {
-                    last_rec_lsn = lsn;
-                }
-
                 Some(streaming_lsn)
             }
 

poetry.lock

@@ -2563,18 +2563,18 @@ pytest = "*"
 
 [[package]]
 name = "pytest-rerunfailures"
-version = "15.0"
+version = "13.0"
 description = "pytest plugin to re-run tests to eliminate flaky failures"
 optional = false
-python-versions = ">=3.9"
+python-versions = ">=3.7"
 files = [
-    {file = "pytest-rerunfailures-15.0.tar.gz", hash = "sha256:2d9ac7baf59f4c13ac730b47f6fa80e755d1ba0581da45ce30b72fb3542b4474"},
-    {file = "pytest_rerunfailures-15.0-py3-none-any.whl", hash = "sha256:dd150c4795c229ef44320adc9a0c0532c51b78bb7a6843a8c53556b9a611df1a"},
+    {file = "pytest-rerunfailures-13.0.tar.gz", hash = "sha256:e132dbe420bc476f544b96e7036edd0a69707574209b6677263c950d19b09199"},
+    {file = "pytest_rerunfailures-13.0-py3-none-any.whl", hash = "sha256:34919cb3fcb1f8e5d4b940aa75ccdea9661bade925091873b7c6fa5548333069"},
 ]
 
 [package.dependencies]
 packaging = ">=17.1"
-pytest = ">=7.4,<8.2.2 || >8.2.2"
+pytest = ">=7"
 
 [[package]]
 name = "pytest-split"
@@ -3524,4 +3524,4 @@ cffi = ["cffi (>=1.11)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "426c385df93f578ba3537c40a269535e27fbcca1978b3cf266096ecbc298c6a9"
+content-hash = "21debe1116843e5d14bdf37d6e265c68c63a98a64ba04ec8b8a02af2e8d9f486"

proxy/src/cancellation.rs

@@ -99,17 +99,16 @@ impl<P: CancellationPublisher> CancellationHandler<P> {
     /// Try to cancel a running query for the corresponding connection.
     /// If the cancellation key is not found, it will be published to Redis.
     /// check_allowed - if true, check if the IP is allowed to cancel the query
-    /// return Result primarily for tests
     pub(crate) async fn cancel_session(
         &self,
         key: CancelKeyData,
         session_id: Uuid,
-        peer_addr: IpAddr,
+        peer_addr: &IpAddr,
         check_allowed: bool,
     ) -> Result<(), CancelError> {
         // TODO: check for unspecified address is only for backward compatibility, should be removed
         if !peer_addr.is_unspecified() {
-            let subnet_key = match peer_addr {
+            let subnet_key = match *peer_addr {
                 IpAddr::V4(ip) => IpNet::V4(Ipv4Net::new_assert(ip, 24).trunc()), // use defaut mask here
                 IpAddr::V6(ip) => IpNet::V6(Ipv6Net::new_assert(ip, 64).trunc()),
             };
@@ -142,11 +141,9 @@ impl<P: CancellationPublisher> CancellationHandler<P> {
                 return Ok(());
             }
 
-            match self.client.try_publish(key, session_id, peer_addr).await {
+            match self.client.try_publish(key, session_id, *peer_addr).await {
                 Ok(()) => {} // do nothing
                 Err(e) => {
-                    // log it here since cancel_session could be spawned in a task
-                    tracing::error!("failed to publish cancellation key: {key}, error: {e}");
                     return Err(CancelError::IO(std::io::Error::new(
                         std::io::ErrorKind::Other,
                         e.to_string(),
@@ -157,10 +154,8 @@ impl<P: CancellationPublisher> CancellationHandler<P> {
         };
 
         if check_allowed
-            && !check_peer_addr_is_in_list(&peer_addr, cancel_closure.ip_allowlist.as_slice())
+            && !check_peer_addr_is_in_list(peer_addr, cancel_closure.ip_allowlist.as_slice())
         {
-            // log it here since cancel_session could be spawned in a task
-            tracing::warn!("IP is not allowed to cancel the query: {key}");
             return Err(CancelError::IpNotAllowed);
         }
 
@@ -311,7 +306,7 @@ mod tests {
                     cancel_key: 0,
                 },
                 Uuid::new_v4(),
-                "127.0.0.1".parse().unwrap(),
+                &("127.0.0.1".parse().unwrap()),
                 true,
             )
             .await

proxy/src/console_redirect_proxy.rs

@@ -35,7 +35,6 @@ pub async fn task_main(
     socket2::SockRef::from(&listener).set_keepalive(true)?;
 
     let connections = tokio_util::task::task_tracker::TaskTracker::new();
-    let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
 
     while let Some(accept_result) =
         run_until_cancelled(listener.accept(), &cancellation_token).await
@@ -49,7 +48,6 @@ pub async fn task_main(
 
         let session_id = uuid::Uuid::new_v4();
         let cancellation_handler = Arc::clone(&cancellation_handler);
-        let cancellations = cancellations.clone();
 
         debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
 
@@ -98,7 +96,6 @@ pub async fn task_main(
                 cancellation_handler,
                 socket,
                 conn_gauge,
-                cancellations,
             )
             .instrument(ctx.span())
             .boxed()
@@ -130,12 +127,10 @@ pub async fn task_main(
     }
 
     connections.close();
-    cancellations.close();
     drop(listener);
 
     // Drain connections
     connections.wait().await;
-    cancellations.wait().await;
 
     Ok(())
 }
@@ -147,7 +142,6 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     cancellation_handler: Arc<CancellationHandlerMain>,
     stream: S,
     conn_gauge: NumClientConnectionsGuard<'static>,
-    cancellations: tokio_util::task::task_tracker::TaskTracker,
 ) -> Result<Option<ProxyPassthrough<CancellationHandlerMainInternal, S>>, ClientRequestError> {
     debug!(
         protocol = %ctx.protocol(),
@@ -167,26 +161,15 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
         match tokio::time::timeout(config.handshake_timeout, do_handshake).await?? {
             HandshakeData::Startup(stream, params) => (stream, params),
             HandshakeData::Cancel(cancel_key_data) => {
-                // spawn a task to cancel the session, but don't wait for it
-                cancellations.spawn({
-                    let cancellation_handler_clone = Arc::clone(&cancellation_handler);
-                    let session_id = ctx.session_id();
-                    let peer_ip = ctx.peer_addr();
-                    async move {
-                        drop(
-                            cancellation_handler_clone
-                                .cancel_session(
-                                    cancel_key_data,
-                                    session_id,
-                                    peer_ip,
-                                    config.authentication_config.ip_allowlist_check_enabled,
-                                )
-                                .await,
-                        );
-                    }
-                });
-
-                return Ok(None);
+                return Ok(cancellation_handler
+                    .cancel_session(
+                        cancel_key_data,
+                        ctx.session_id(),
+                        &ctx.peer_addr(),
+                        config.authentication_config.ip_allowlist_check_enabled,
+                    )
+                    .await
+                    .map(|()| None)?)
             }
         };
     drop(pause);

proxy/src/proxy/mod.rs

@@ -69,7 +69,6 @@ pub async fn task_main(
     socket2::SockRef::from(&listener).set_keepalive(true)?;
 
     let connections = tokio_util::task::task_tracker::TaskTracker::new();
-    let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
 
     while let Some(accept_result) =
         run_until_cancelled(listener.accept(), &cancellation_token).await
@@ -83,7 +82,6 @@ pub async fn task_main(
 
         let session_id = uuid::Uuid::new_v4();
         let cancellation_handler = Arc::clone(&cancellation_handler);
-        let cancellations = cancellations.clone();
 
         debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
         let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
@@ -135,7 +133,6 @@ pub async fn task_main(
                 ClientMode::Tcp,
                 endpoint_rate_limiter2,
                 conn_gauge,
-                cancellations,
             )
             .instrument(ctx.span())
             .boxed()
@@ -167,12 +164,10 @@ pub async fn task_main(
     }
 
     connections.close();
-    cancellations.close();
     drop(listener);
 
     // Drain connections
     connections.wait().await;
-    cancellations.wait().await;
 
     Ok(())
 }
@@ -255,7 +250,6 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     mode: ClientMode,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     conn_gauge: NumClientConnectionsGuard<'static>,
-    cancellations: tokio_util::task::task_tracker::TaskTracker,
 ) -> Result<Option<ProxyPassthrough<CancellationHandlerMainInternal, S>>, ClientRequestError> {
     debug!(
         protocol = %ctx.protocol(),
@@ -276,26 +270,15 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
         match tokio::time::timeout(config.handshake_timeout, do_handshake).await?? {
             HandshakeData::Startup(stream, params) => (stream, params),
             HandshakeData::Cancel(cancel_key_data) => {
-                // spawn a task to cancel the session, but don't wait for it
-                cancellations.spawn({
-                    let cancellation_handler_clone = Arc::clone(&cancellation_handler);
-                    let session_id = ctx.session_id();
-                    let peer_ip = ctx.peer_addr();
-                    async move {
-                        drop(
-                            cancellation_handler_clone
-                                .cancel_session(
-                                    cancel_key_data,
-                                    session_id,
-                                    peer_ip,
-                                    config.authentication_config.ip_allowlist_check_enabled,
-                                )
-                                .await,
-                        );
-                    }
-                });
-
-                return Ok(None);
+                return Ok(cancellation_handler
+                    .cancel_session(
+                        cancel_key_data,
+                        ctx.session_id(),
+                        &ctx.peer_addr(),
+                        config.authentication_config.ip_allowlist_check_enabled,
+                    )
+                    .await
+                    .map(|()| None)?)
             }
         };
     drop(pause);

proxy/src/redis/notifications.rs

@@ -149,7 +149,7 @@ impl<C: ProjectInfoCache + Send + Sync + 'static> MessageHandler<C> {
                     .cancel_session(
                         cancel_session.cancel_key_data,
                         uuid::Uuid::nil(),
-                        peer_addr,
+                        &peer_addr,
                         cancel_session.peer_addr.is_some(),
                     )
                     .await

proxy/src/serverless/mod.rs

@@ -132,7 +132,6 @@ pub async fn task_main(
     let connections = tokio_util::task::task_tracker::TaskTracker::new();
     connections.close(); // allows `connections.wait to complete`
 
-    let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
     while let Some(res) = run_until_cancelled(ws_listener.accept(), &cancellation_token).await {
         let (conn, peer_addr) = res.context("could not accept TCP stream")?;
         if let Err(e) = conn.set_nodelay(true) {
@@ -161,7 +160,6 @@ pub async fn task_main(
         let connections2 = connections.clone();
         let cancellation_handler = cancellation_handler.clone();
         let endpoint_rate_limiter = endpoint_rate_limiter.clone();
-        let cancellations = cancellations.clone();
         connections.spawn(
             async move {
                 let conn_token2 = conn_token.clone();
@@ -190,7 +188,6 @@ pub async fn task_main(
                     config,
                     backend,
                     connections2,
-                    cancellations,
                     cancellation_handler,
                     endpoint_rate_limiter,
                     conn_token,
@@ -316,7 +313,6 @@ async fn connection_handler(
     config: &'static ProxyConfig,
     backend: Arc<PoolingBackend>,
     connections: TaskTracker,
-    cancellations: TaskTracker,
     cancellation_handler: Arc<CancellationHandlerMain>,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     cancellation_token: CancellationToken,
@@ -357,7 +353,6 @@ async fn connection_handler(
 
             // `request_handler` is not cancel safe. It expects to be cancelled only at specific times.
             // By spawning the future, we ensure it never gets cancelled until it decides to.
-            let cancellations = cancellations.clone();
             let handler = connections.spawn(
                 request_handler(
                     req,
@@ -369,7 +364,6 @@ async fn connection_handler(
                     conn_info2.clone(),
                     http_request_token,
                     endpoint_rate_limiter.clone(),
-                    cancellations,
                 )
                 .in_current_span()
                 .map_ok_or_else(api_error_into_response, |r| r),
@@ -417,7 +411,6 @@ async fn request_handler(
     // used to cancel in-flight HTTP requests. not used to cancel websockets
     http_cancellation_token: CancellationToken,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
-    cancellations: TaskTracker,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ApiError> {
     let host = request
         .headers()
@@ -443,7 +436,6 @@ async fn request_handler(
         let (response, websocket) = framed_websockets::upgrade::upgrade(&mut request)
             .map_err(|e| ApiError::BadRequest(e.into()))?;
 
-        let cancellations = cancellations.clone();
         ws_connections.spawn(
             async move {
                 if let Err(e) = websocket::serve_websocket(
@@ -454,7 +446,6 @@ async fn request_handler(
                     cancellation_handler,
                     endpoint_rate_limiter,
                     host,
-                    cancellations,
                 )
                 .await
                 {

proxy/src/serverless/websocket.rs

@@ -123,7 +123,6 @@ impl<S: AsyncRead + AsyncWrite + Unpin> AsyncBufRead for WebSocketRw<S> {
     }
 }
 
-#[allow(clippy::too_many_arguments)]
 pub(crate) async fn serve_websocket(
     config: &'static ProxyConfig,
     auth_backend: &'static crate::auth::Backend<'static, ()>,
@@ -132,7 +131,6 @@ pub(crate) async fn serve_websocket(
     cancellation_handler: Arc<CancellationHandlerMain>,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     hostname: Option<String>,
-    cancellations: tokio_util::task::task_tracker::TaskTracker,
 ) -> anyhow::Result<()> {
     let websocket = websocket.await?;
     let websocket = WebSocketServer::after_handshake(TokioIo::new(websocket));
@@ -151,7 +149,6 @@ pub(crate) async fn serve_websocket(
         ClientMode::Websockets { hostname },
         endpoint_rate_limiter,
         conn_gauge,
-        cancellations,
     ))
     .await;
 

pyproject.toml

@@ -33,7 +33,7 @@ types-psutil = "^5.9.5.12"
 types-toml = "^0.10.8.6"
 pytest-httpserver = "^1.0.8"
 aiohttp = "3.10.11"
-pytest-rerunfailures = "^15.0"
+pytest-rerunfailures = "^13.0"
 types-pytest-lazy-fixture = "^0.6.3.3"
 pytest-split = "^0.8.1"
 zstandard = "^0.21.0"

rust-toolchain.toml

@@ -1,5 +1,5 @@
 [toolchain]
-channel = "1.83.0"
+channel = "1.82.0"
 profile = "default"
 # The default profile includes rustc, rust-std, cargo, rust-docs, rustfmt and clippy.
 # https://rust-lang.github.io/rustup/concepts/profiles.html

safekeeper/src/handler.rs

@@ -212,9 +212,8 @@ impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
                 );
 
             if let Some(shard) = self.shard.as_ref() {
-                if let Some(slug) = shard.shard_slug().strip_prefix("-") {
-                    tracing::Span::current().record("shard", tracing::field::display(slug));
-                }
+                tracing::Span::current()
+                    .record("shard", tracing::field::display(shard.shard_slug()));
             }
 
             Ok(())

safekeeper/src/wal_service.rs

@@ -44,7 +44,7 @@ pub async fn task_main(
                     error!("connection handler exited: {}", err);
                 }
             }
-            .instrument(info_span!("", cid = %conn_id, ttid = field::Empty, application_name = field::Empty, shard = field::Empty)),
+            .instrument(info_span!("", cid = %conn_id, ttid = field::Empty, application_name = field::Empty)),
         );
     }
 }

scripts/flaky_tests.py

@@ -0,0 +1,147 @@
+#! /usr/bin/env python3
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+import psycopg2
+import psycopg2.extras
+import toml
+
+if TYPE_CHECKING:
+    from typing import Any
+
+FLAKY_TESTS_QUERY = """
+    SELECT
+        DISTINCT parent_suite, suite, name
+    FROM results
+    WHERE
+        started_at > CURRENT_DATE - INTERVAL '%s' day
+        AND (
+            (status IN ('failed', 'broken') AND reference = 'refs/heads/main')
+            OR flaky
+        )
+    ;
+"""
+
+
+def main(args: argparse.Namespace):
+    connstr = args.connstr
+    interval_days = args.days
+    output = args.output
+
+    build_type = args.build_type
+    pg_version = args.pg_version
+
+    res: defaultdict[str, defaultdict[str, dict[str, bool]]]
+    res = defaultdict(lambda: defaultdict(dict))
+
+    try:
+        logging.info("connecting to the database...")
+        with psycopg2.connect(connstr, connect_timeout=30) as conn:
+            with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
+                logging.info("fetching flaky tests...")
+                cur.execute(FLAKY_TESTS_QUERY, (interval_days,))
+                rows = cur.fetchall()
+    except psycopg2.OperationalError as exc:
+        logging.error("cannot fetch flaky tests from the DB due to an error", exc)
+        rows = []
+
+    # If a test run has non-default PAGESERVER_VIRTUAL_FILE_IO_ENGINE (i.e. not empty, not tokio-epoll-uring),
+    # use it to parametrize test name along with build_type and pg_version
+    #
+    # See test_runner/fixtures/parametrize.py for details
+    if (io_engine := os.getenv("PAGESERVER_VIRTUAL_FILE_IO_ENGINE", "")) not in (
+        "",
+        "tokio-epoll-uring",
+    ):
+        pageserver_virtual_file_io_engine_parameter = f"-{io_engine}"
+    else:
+        pageserver_virtual_file_io_engine_parameter = ""
+
+    # re-use existing records of flaky tests from before parametrization by compaction_algorithm
+    def get_pageserver_default_tenant_config_compaction_algorithm() -> dict[str, Any] | None:
+        """Duplicated from parametrize.py"""
+        toml_table = os.getenv("PAGESERVER_DEFAULT_TENANT_CONFIG_COMPACTION_ALGORITHM")
+        if toml_table is None:
+            return None
+        v = toml.loads(toml_table)
+        assert isinstance(v, dict)
+        return v
+
+    pageserver_default_tenant_config_compaction_algorithm_parameter = ""
+    if (
+        explicit_default := get_pageserver_default_tenant_config_compaction_algorithm()
+    ) is not None:
+        pageserver_default_tenant_config_compaction_algorithm_parameter = (
+            f"-{explicit_default['kind']}"
+        )
+
+    for row in rows:
+        # We don't want to automatically rerun tests in a performance suite
+        if row["parent_suite"] != "test_runner.regress":
+            continue
+
+        if row["name"].endswith("]"):
+            parametrized_test = row["name"].replace(
+                "[",
+                f"[{build_type}-pg{pg_version}{pageserver_virtual_file_io_engine_parameter}{pageserver_default_tenant_config_compaction_algorithm_parameter}-",
+            )
+        else:
+            parametrized_test = f"{row['name']}[{build_type}-pg{pg_version}{pageserver_virtual_file_io_engine_parameter}{pageserver_default_tenant_config_compaction_algorithm_parameter}]"
+
+        res[row["parent_suite"]][row["suite"]][parametrized_test] = True
+
+        logging.info(
+            f"\t{row['parent_suite'].replace('.', '/')}/{row['suite']}.py::{parametrized_test}"
+        )
+
+    logging.info(f"saving results to {output.name}")
+    json.dump(res, output, indent=2)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Detect flaky tests in the last N days")
+    parser.add_argument(
+        "--output",
+        type=argparse.FileType("w"),
+        default="flaky.json",
+        help="path to output json file (default: flaky.json)",
+    )
+    parser.add_argument(
+        "--days",
+        required=False,
+        default=10,
+        type=int,
+        help="how many days to look back for flaky tests (default: 10)",
+    )
+    parser.add_argument(
+        "--build-type",
+        required=True,
+        type=str,
+        help="for which build type to create list of flaky tests (debug or release)",
+    )
+    parser.add_argument(
+        "--pg-version",
+        required=True,
+        type=int,
+        help="for which Postgres version to create list of flaky tests (14, 15, etc.)",
+    )
+    parser.add_argument(
+        "connstr",
+        help="connection string to the test results database",
+    )
+    args = parser.parse_args()
+
+    level = logging.INFO
+    logging.basicConfig(
+        format="%(message)s",
+        level=level,
+    )
+
+    main(args)

test_runner/conftest.py

@@ -13,5 +13,5 @@ pytest_plugins = (
     "fixtures.pg_stats",
     "fixtures.compare_fixtures",
     "fixtures.slow",
-    "fixtures.reruns",
+    "fixtures.flaky",
 )

test_runner/fixtures/endpoint/http.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import urllib.parse
-
 import requests
 from requests.adapters import HTTPAdapter
 
@@ -22,9 +20,7 @@ class EndpointHttpClient(requests.Session):
         return res.json()
 
     def database_schema(self, database: str):
-        res = self.get(
-            f"http://localhost:{self.port}/database_schema?database={urllib.parse.quote(database, safe='')}"
-        )
+        res = self.get(f"http://localhost:{self.port}/database_schema?database={database}")
         res.raise_for_status()
         return res.text
 

test_runner/fixtures/flaky.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import json
+from collections.abc import MutableMapping
+from pathlib import Path
+from typing import TYPE_CHECKING, cast
+
+import pytest
+from _pytest.config import Config
+from _pytest.config.argparsing import Parser
+from allure_commons.types import LabelType
+from allure_pytest.utils import allure_name, allure_suite_labels
+
+from fixtures.log_helper import log
+
+if TYPE_CHECKING:
+    from collections.abc import MutableMapping
+    from typing import Any
+
+
+"""
+The plugin reruns flaky tests.
+It uses `pytest.mark.flaky` provided by `pytest-rerunfailures` plugin and flaky tests detected by `scripts/flaky_tests.py`
+
+Note: the logic of getting flaky tests is extracted to a separate script to avoid running it for each of N xdist workers
+"""
+
+
+def pytest_addoption(parser: Parser):
+    parser.addoption(
+        "--flaky-tests-json",
+        action="store",
+        type=Path,
+        help="Path to json file with flaky tests generated by scripts/flaky_tests.py",
+    )
+
+
+def pytest_collection_modifyitems(config: Config, items: list[pytest.Item]):
+    if not config.getoption("--flaky-tests-json"):
+        return
+
+    # Any error with getting flaky tests aren't critical, so just do not rerun any tests
+    flaky_json = config.getoption("--flaky-tests-json")
+    if not flaky_json.exists():
+        return
+
+    content = flaky_json.read_text()
+    try:
+        flaky_tests = json.loads(content)
+    except ValueError:
+        log.error(f"Can't parse {content} as json")
+        return
+
+    for item in items:
+        # Use the same logic for constructing test name as Allure does (we store allure-provided data in DB)
+        # Ref https://github.com/allure-framework/allure-python/blob/2.13.1/allure-pytest/src/listener.py#L98-L100
+        allure_labels = dict(allure_suite_labels(item))
+        parent_suite = str(allure_labels.get(LabelType.PARENT_SUITE))
+        suite = str(allure_labels.get(LabelType.SUITE))
+        params = item.callspec.params if hasattr(item, "callspec") else {}
+        name = allure_name(item, params)
+
+        if flaky_tests.get(parent_suite, {}).get(suite, {}).get(name, False):
+            # Rerun 3 times = 1 original run + 2 reruns
+            log.info(f"Marking {item.nodeid} as flaky. It will be rerun up to 3 times")
+            item.add_marker(pytest.mark.flaky(reruns=2))
+
+            # pytest-rerunfailures is not compatible with pytest-timeout (timeout is not set for reruns),
+            #   we can workaround it by setting `timeout_func_only` to True[1].
+            # Unfortunately, setting `timeout_func_only = True` globally in pytest.ini is broken[2],
+            #   but we still can do it using pytest marker.
+            #
+            # - [1] https://github.com/pytest-dev/pytest-rerunfailures/issues/99
+            # - [2] https://github.com/pytest-dev/pytest-timeout/issues/142
+            timeout_marker = item.get_closest_marker("timeout")
+            if timeout_marker is not None:
+                kwargs = cast("MutableMapping[str, Any]", timeout_marker.kwargs)
+                kwargs["func_only"] = True

test_runner/fixtures/neon_fixtures.py

@@ -3934,35 +3934,6 @@ class Endpoint(PgProtocol, LogUtils):
             log.info(json.dumps(dict(data_dict, **kwargs)))
             json.dump(dict(data_dict, **kwargs), file, indent=4)
 
-    def respec_deep(self, **kwargs: Any) -> None:
-        """
-        Update the endpoint.json file taking into account nested keys.
-        It does one level deep update. Should enough for most cases.
-        Distinct method from respec() to do not break existing functionality.
-        NOTE: This method also updates the spec.json file, not endpoint.json.
-        We need it because neon_local also writes to spec.json, so intended
-        use-case is i) start endpoint with some config, ii) respec_deep(),
-        iii) call reconfigure() to apply the changes.
-        """
-        config_path = os.path.join(self.endpoint_path(), "spec.json")
-        with open(config_path) as f:
-            data_dict: dict[str, Any] = json.load(f)
-
-        log.info("Current compute spec: %s", json.dumps(data_dict, indent=4))
-
-        for key, value in kwargs.items():
-            if isinstance(value, dict):
-                if key not in data_dict:
-                    data_dict[key] = value
-                else:
-                    data_dict[key] = {**data_dict[key], **value}
-            else:
-                data_dict[key] = value
-
-        with open(config_path, "w") as file:
-            log.info("Updating compute spec to: %s", json.dumps(data_dict, indent=4))
-            json.dump(data_dict, file, indent=4)
-
     # Please note: Migrations only run if pg_skip_catalog_updates is false
     def wait_for_migrations(self, num_migrations: int = 11):
         with self.cursor() as cur:

test_runner/fixtures/paths.py

@@ -30,7 +30,7 @@ def get_test_dir(request: FixtureRequest, top_output_dir: Path, prefix: str | No
     test_name = request.node.name
     test_dir = top_output_dir / f"{prefix or ''}{test_name.replace('/', '-')}"
 
-    # We rerun failed tests multiple times, use a separate directory for each run.
+    # We rerun flaky tests multiple times, use a separate directory for each run.
     if (suffix := getattr(request.node, "execution_count", None)) is not None:
         test_dir = test_dir.parent / f"{test_dir.name}-{suffix}"
 

test_runner/fixtures/reruns.py

@@ -1,31 +0,0 @@
-from __future__ import annotations
-
-from collections.abc import MutableMapping
-from typing import TYPE_CHECKING, cast
-
-import pytest
-
-if TYPE_CHECKING:
-    from collections.abc import MutableMapping
-    from typing import Any
-
-    from _pytest.config import Config
-
-
-def pytest_collection_modifyitems(config: Config, items: list[pytest.Item]):
-    # pytest-rerunfailures is not compatible with pytest-timeout (timeout is not set for reruns),
-    #   we can workaround it by setting `timeout_func_only` to True[1].
-    # Unfortunately, setting `timeout_func_only = True` globally in pytest.ini is broken[2],
-    #   but we still can do it using pytest marker.
-    #
-    # - [1] https://github.com/pytest-dev/pytest-rerunfailures/issues/99
-    # - [2] https://github.com/pytest-dev/pytest-timeout/issues/142
-
-    if not config.getoption("--reruns"):
-        return
-
-    for item in items:
-        timeout_marker = item.get_closest_marker("timeout")
-        if timeout_marker is not None:
-            kwargs = cast("MutableMapping[str, Any]", timeout_marker.kwargs)
-            kwargs["func_only"] = True

test_runner/regress/test_compute_catalog.py

@@ -3,60 +3,13 @@ from __future__ import annotations
 import requests
 from fixtures.neon_fixtures import NeonEnv
 
-TEST_DB_NAMES = [
-    {
-        "name": "neondb",
-        "owner": "cloud_admin",
-    },
-    {
-        "name": "db with spaces",
-        "owner": "cloud_admin",
-    },
-    {
-        "name": "db with%20spaces ",
-        "owner": "cloud_admin",
-    },
-    {
-        "name": "db with whitespaces	",
-        "owner": "cloud_admin",
-    },
-    {
-        "name": "injective db with spaces'; SELECT pg_sleep(10);",
-        "owner": "cloud_admin",
-    },
-    {
-        "name": "db with #pound-sign and &ampersands=true",
-        "owner": "cloud_admin",
-    },
-    {
-        "name": "db with emoji 🌍",
-        "owner": "cloud_admin",
-    },
-]
-
 
 def test_compute_catalog(neon_simple_env: NeonEnv):
-    """
-    Create a bunch of databases with tricky names and test that we can list them
-    and dump via API.
-    """
     env = neon_simple_env
 
-    endpoint = env.endpoints.create_start("main")
-
-    # Update the spec.json file to include new databases
-    # and reconfigure the endpoint to create some test databases.
-    endpoint.respec_deep(
-        **{
-            "skip_pg_catalog_updates": False,
-            "cluster": {
-                "databases": TEST_DB_NAMES,
-            },
-        }
-    )
-    endpoint.reconfigure()
-
+    endpoint = env.endpoints.create_start("main", config_lines=["log_min_messages=debug1"])
     client = endpoint.http_client()
+
     objects = client.dbs_and_roles()
 
     # Assert that 'cloud_admin' role exists in the 'roles' list
@@ -69,24 +22,9 @@ def test_compute_catalog(neon_simple_env: NeonEnv):
         db["name"] == "postgres" for db in objects["databases"]
     ), "The 'postgres' database is missing"
 
-    # Check other databases
-    for test_db in TEST_DB_NAMES:
-        db = next((db for db in objects["databases"] if db["name"] == test_db["name"]), None)
-        assert db is not None, f"The '{test_db['name']}' database is missing"
-        assert (
-            db["owner"] == test_db["owner"]
-        ), f"The '{test_db['name']}' database has incorrect owner"
+    ddl = client.database_schema(database="postgres")
 
-        ddl = client.database_schema(database=test_db["name"])
-
-        # Check that it looks like a valid PostgreSQL dump
-        assert "-- PostgreSQL database dump" in ddl
-
-        # Check that it doesn't contain health_check and migration traces.
-        # They are only created in system `postgres` database, so by checking
-        # that we ensure that we dump right databases.
-        assert "health_check" not in ddl, f"The '{test_db['name']}' database contains health_check"
-        assert "migration" not in ddl, f"The '{test_db['name']}' database contains migrations data"
+    assert "-- PostgreSQL database dump" in ddl
 
     try:
         client.database_schema(database="nonexistentdb")
@@ -95,44 +33,3 @@ def test_compute_catalog(neon_simple_env: NeonEnv):
         assert (
             e.response.status_code == 404
         ), f"Expected 404 status code, but got {e.response.status_code}"
-
-
-def test_compute_create_databases(neon_simple_env: NeonEnv):
-    """
-    Test that compute_ctl can create and work with databases with special
-    characters (whitespaces, %, tabs, etc.) in the name.
-    """
-    env = neon_simple_env
-
-    # Create and start endpoint so that neon_local put all the generated
-    # stuff into the spec.json file.
-    endpoint = env.endpoints.create_start("main")
-
-    # Update the spec.json file to include new databases
-    # and reconfigure the endpoint to apply the changes.
-    endpoint.respec_deep(
-        **{
-            "skip_pg_catalog_updates": False,
-            "cluster": {
-                "databases": TEST_DB_NAMES,
-            },
-        }
-    )
-    endpoint.reconfigure()
-
-    for db in TEST_DB_NAMES:
-        # Check that database has a correct name in the system catalog
-        with endpoint.cursor() as cursor:
-            cursor.execute("SELECT datname FROM pg_database WHERE datname = %s", (db["name"],))
-            catalog_db = cursor.fetchone()
-            assert catalog_db is not None
-            assert len(catalog_db) == 1
-            assert catalog_db[0] == db["name"]
-
-        # Check that we can connect to this database without any issues
-        with endpoint.cursor(dbname=db["name"]) as cursor:
-            cursor.execute("SELECT * FROM current_database()")
-            curr_db = cursor.fetchone()
-            assert curr_db is not None
-            assert len(curr_db) == 1
-            assert curr_db[0] == db["name"]