fix missing registration

## Problem

We run benchmarks in batches (five parallel jobs on different runners).
If any test in a batch fails, we won’t upload any results for that
batch, even for the tests that passed.

## Summary of changes
- Move the results upload to a separate step in the run-python-test-set
action, and execute this step even if tests fail.

.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

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

@@ -349,7 +349,7 @@ jobs:
       contents: read
       statuses: write
     needs: [ build-neon ]
-    runs-on: ${{ fromJSON(format('["self-hosted", "{0}"]', inputs.arch == 'arm64' && 'large-arm64' || 'large')) }}
+    runs-on: ${{ fromJSON(format('["self-hosted", "{0}"]', inputs.arch == 'arm64' && 'large-arm64' || 'large-metal')) }}
     container:
       image: ${{ inputs.build-tools-image }}
       credentials:

.github/workflows/_meta.yml

@@ -165,5 +165,5 @@ jobs:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           CURRENT_SHA: ${{ github.sha }}
         run: |
-          RELEASE_PR_RUN_ID=$(gh api "/repos/${GITHUB_REPOSITORY}/actions/runs?head_sha=$CURRENT_SHA" | jq '[.workflow_runs[] | select(.name == "Build and Test") | select(.head_branch | test("^rc/release(-(proxy|compute))?/[0-9]{4}-[0-9]{2}-[0-9]{2}$"; "s"))] | first | .id // ("Failed to find Build and Test run from  RC PR!" | halt_error(1))')
+          RELEASE_PR_RUN_ID=$(gh api "/repos/${GITHUB_REPOSITORY}/actions/runs?head_sha=$CURRENT_SHA" | jq '[.workflow_runs[] | select(.name == "Build and Test") | select(.head_branch | test("^rc/release.*$"; "s"))] | first | .id // ("Failed to find Build and Test run from  RC PR!" | halt_error(1))')
           echo "release-pr-run-id=$RELEASE_PR_RUN_ID" | tee -a $GITHUB_OUTPUT

.github/workflows/fast-forward.yml

@@ -27,15 +27,17 @@ jobs:
       - name: Fast forwarding
         uses: sequoia-pgp/fast-forward@ea7628bedcb0b0b96e94383ada458d812fca4979
         # See https://docs.github.com/en/graphql/reference/enums#mergestatestatus
-        if: ${{ github.event.pull_request.mergeable_state  == 'clean' }}
+        if: ${{ contains(fromJSON('["clean", "unstable"]'), github.event.pull_request.mergeable_state) }}
         with:
           merge: true
           comment: on-error
           github_token: ${{ secrets.CI_ACCESS_TOKEN }}
 
       - name: Comment if mergeable_state is not clean
-        if: ${{ github.event.pull_request.mergeable_state  != 'clean' }}
+        if: ${{ !contains(fromJSON('["clean", "unstable"]'), github.event.pull_request.mergeable_state) }}
+        env:
+          GH_TOKEN: ${{ secrets.CI_ACCESS_TOKEN }}
         run: |
           gh pr comment ${{ github.event.pull_request.number }} \
             --repo "${GITHUB_REPOSITORY}" \
-            --body "Not trying to forward pull-request, because \`mergeable_state\` is \`${{ github.event.pull_request.mergeable_state }}\`, not \`clean\`."
+            --body "Not trying to forward pull-request, because \`mergeable_state\` is \`${{ github.event.pull_request.mergeable_state }}\`, not \`clean\` or \`unstable\`."

.github/workflows/random-ops-test.yml

@@ -0,0 +1,93 @@
+name: Random Operations Test
+
+on:
+  schedule:
+    # * is a special character in YAML so you have to quote this string
+    #          ┌───────────── minute (0 - 59)
+    #          │  ┌───────────── hour (0 - 23)
+    #          │  │  ┌───────────── day of the month (1 - 31)
+    #          │  │  │ ┌───────────── month (1 - 12 or JAN-DEC)
+    #          │  │  │ │ ┌───────────── day of the week (0 - 6 or SUN-SAT)
+    - cron:  '23 */2 * * *' # runs every 2 hours
+  workflow_dispatch:
+    inputs:
+      random_seed:
+        type: number
+        description: 'The random seed'
+        required: false
+        default: 0
+      num_operations:
+        type: number
+        description: "The number of operations to test"
+        default: 250
+
+defaults:
+  run:
+    shell: bash -euxo pipefail {0}
+
+permissions: {}
+
+env:
+  DEFAULT_PG_VERSION: 16
+  PLATFORM: neon-captest-new
+  AWS_DEFAULT_REGION: eu-central-1
+
+jobs:
+  run-random-rests:
+    env:
+      POSTGRES_DISTRIB_DIR: /tmp/neon/pg_install
+    runs-on: small
+    permissions:
+      id-token: write
+      statuses: write
+
+    strategy:
+      fail-fast: false
+      matrix:
+        pg-version: [16, 17]
+
+    container:
+      image: ghcr.io/neondatabase/build-tools:pinned-bookworm
+      credentials:
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+      options: --init
+    steps:
+      - name: Harden the runner (Audit all outbound calls)
+        uses: step-security/harden-runner@4d991eb9b905ef189e4c376166672c3f2f230481 # v2.11.0
+        with:
+          egress-policy: audit
+
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+
+      - name: Download Neon artifact
+        uses: ./.github/actions/download
+        with:
+          name: neon-${{ runner.os }}-${{ runner.arch }}-release-artifact
+          path: /tmp/neon/
+          prefix: latest
+          aws-oicd-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+
+      - name: Run tests
+        uses: ./.github/actions/run-python-test-set
+        with:
+          build_type: remote
+          test_selection: random_ops
+          run_in_parallel: false
+          extra_params: -m remote_cluster
+          pg_version: ${{ matrix.pg-version }}
+          aws-oicd-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+        env:
+          NEON_API_KEY: ${{ secrets.NEON_STAGING_API_KEY }}
+          RANDOM_SEED: ${{ inputs.random_seed }}
+          NUM_OPERATIONS: ${{ inputs.num_operations }}
+
+      - name: Create Allure report
+        if: ${{ !cancelled() }}
+        id: create-allure-report
+        uses: ./.github/actions/allure-report-generate
+        with:
+          store-test-results-into-db: true
+          aws-oicd-role-arn: ${{ vars.DEV_AWS_OIDC_ROLE_ARN }}
+        env:
+          REGRESS_TEST_RESULT_CONNSTR_NEW: ${{ secrets.REGRESS_TEST_RESULT_CONNSTR_NEW }}

Cargo.lock

@@ -4285,6 +4285,7 @@ dependencies = [
  "pageserver_api",
  "pageserver_client",
  "pageserver_compaction",
+ "pem",
  "pin-project-lite",
  "postgres-protocol",
  "postgres-types",
@@ -6001,6 +6002,7 @@ dependencies = [
  "once_cell",
  "pageserver_api",
  "parking_lot 0.12.1",
+ "pem",
  "postgres-protocol",
  "postgres_backend",
  "postgres_ffi",

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_anon.patch

@@ -1,265 +0,0 @@
-commit 00aa659afc9c7336ab81036edec3017168aabf40
-Author: Heikki Linnakangas <heikki@neon.tech>
-Date:   Tue Nov 12 16:59:19 2024 +0200
-
-    Temporarily disable test that depends on timezone
-
-diff --git a/tests/expected/generalization.out b/tests/expected/generalization.out
-index 23ef5fa..9e60deb 100644
---- a/ext-src/pg_anon-src/tests/expected/generalization.out
-+++ b/ext-src/pg_anon-src/tests/expected/generalization.out
-@@ -284,12 +284,9 @@ SELECT anon.generalize_tstzrange('19041107','century');
-  ["Tue Jan 01 00:00:00 1901 PST","Mon Jan 01 00:00:00 2001 PST")
- (1 row)
- 
--SELECT anon.generalize_tstzrange('19041107','millennium');
--                      generalize_tstzrange                       
-------------------------------------------------------------------
-- ["Thu Jan 01 00:00:00 1001 PST","Mon Jan 01 00:00:00 2001 PST")
--(1 row)
--
-+-- temporarily disabled, see:
-+-- https://gitlab.com/dalibo/postgresql_anonymizer/-/commit/199f0a392b37c59d92ae441fb8f037e094a11a52#note_2148017485
-+--SELECT anon.generalize_tstzrange('19041107','millennium');
- -- generalize_daterange
- SELECT anon.generalize_daterange('19041107');
-   generalize_daterange   
-diff --git a/tests/sql/generalization.sql b/tests/sql/generalization.sql
-index b868344..b4fc977 100644
---- a/ext-src/pg_anon-src/tests/sql/generalization.sql
-+++ b/ext-src/pg_anon-src/tests/sql/generalization.sql
-@@ -61,7 +61,9 @@ SELECT anon.generalize_tstzrange('19041107','month');
- SELECT anon.generalize_tstzrange('19041107','year');
- SELECT anon.generalize_tstzrange('19041107','decade');
- SELECT anon.generalize_tstzrange('19041107','century');
--SELECT anon.generalize_tstzrange('19041107','millennium');
-+-- temporarily disabled, see:
-+-- https://gitlab.com/dalibo/postgresql_anonymizer/-/commit/199f0a392b37c59d92ae441fb8f037e094a11a52#note_2148017485
-+--SELECT anon.generalize_tstzrange('19041107','millennium');
- 
- -- generalize_daterange
- SELECT anon.generalize_daterange('19041107');
-
-commit 7dd414ee75f2875cffb1d6ba474df1f135a6fc6f
-Author: Alexey Masterov <alexeymasterov@neon.tech>
-Date:   Fri May 31 06:34:26 2024 +0000
-
-    These alternative expected files were added to consider the neon features
-
-diff --git a/ext-src/pg_anon-src/tests/expected/permissions_masked_role_1.out b/ext-src/pg_anon-src/tests/expected/permissions_masked_role_1.out
-new file mode 100644
-index 0000000..2539cfd
---- /dev/null
-+++ b/ext-src/pg_anon-src/tests/expected/permissions_masked_role_1.out
-@@ -0,0 +1,101 @@
-+BEGIN;
-+CREATE EXTENSION anon CASCADE;
-+NOTICE:  installing required extension "pgcrypto"
-+SELECT anon.init();
-+ init 
-+------
-+ t
-+(1 row)
-+
-+CREATE ROLE mallory_the_masked_user;
-+SECURITY LABEL FOR anon ON ROLE mallory_the_masked_user IS 'MASKED';
-+CREATE TABLE t1(i INT);
-+ALTER TABLE t1 ADD COLUMN t TEXT;
-+SECURITY LABEL FOR anon ON COLUMN t1.t
-+IS 'MASKED WITH VALUE NULL';
-+INSERT INTO t1 VALUES (1,'test');
-+--
-+-- We're checking the owner's permissions
-+--
-+-- see
-+-- https://postgresql-anonymizer.readthedocs.io/en/latest/SECURITY/#permissions
-+--
-+SET ROLE mallory_the_masked_user;
-+SELECT anon.pseudo_first_name(0) IS NOT NULL;
-+ ?column? 
-+----------
-+ t
-+(1 row)
-+
-+-- SHOULD FAIL
-+DO $$
-+BEGIN
-+  PERFORM anon.init();
-+  EXCEPTION WHEN insufficient_privilege
-+  THEN RAISE NOTICE 'insufficient_privilege';
-+END$$;
-+NOTICE:  insufficient_privilege
-+-- SHOULD FAIL
-+DO $$
-+BEGIN
-+  PERFORM anon.anonymize_table('t1');
-+  EXCEPTION WHEN insufficient_privilege
-+  THEN RAISE NOTICE 'insufficient_privilege';
-+END$$;
-+NOTICE:  insufficient_privilege
-+-- SHOULD FAIL
-+SAVEPOINT fail_start_engine;
-+SELECT anon.start_dynamic_masking();
-+ERROR:  Only supersusers can start the dynamic masking engine.
-+CONTEXT:  PL/pgSQL function anon.start_dynamic_masking(boolean) line 18 at RAISE
-+ROLLBACK TO fail_start_engine;
-+RESET ROLE;
-+SELECT anon.start_dynamic_masking();
-+ start_dynamic_masking 
-+-----------------------
-+ t
-+(1 row)
-+
-+SET ROLE mallory_the_masked_user;
-+SELECT * FROM mask.t1;
-+ i | t 
-+---+---
-+ 1 | 
-+(1 row)
-+
-+-- SHOULD FAIL
-+DO $$
-+BEGIN
-+  SELECT * FROM public.t1;
-+  EXCEPTION WHEN insufficient_privilege
-+  THEN RAISE NOTICE 'insufficient_privilege';
-+END$$;
-+NOTICE:  insufficient_privilege
-+-- SHOULD FAIL
-+SAVEPOINT fail_stop_engine;
-+SELECT anon.stop_dynamic_masking();
-+ERROR:  Only supersusers can stop the dynamic masking engine.
-+CONTEXT:  PL/pgSQL function anon.stop_dynamic_masking() line 18 at RAISE
-+ROLLBACK TO fail_stop_engine;
-+RESET ROLE;
-+SELECT anon.stop_dynamic_masking();
-+NOTICE:  The previous priviledges of 'mallory_the_masked_user' are not restored. You need to grant them manually.
-+ stop_dynamic_masking 
-+----------------------
-+ t
-+(1 row)
-+
-+SET ROLE mallory_the_masked_user;
-+SELECT COUNT(*)=1 FROM anon.pg_masking_rules;
-+ ?column? 
-+----------
-+ t
-+(1 row)
-+
-+-- SHOULD FAIL
-+SAVEPOINT fail_seclabel_on_role;
-+SECURITY LABEL FOR anon ON ROLE mallory_the_masked_user IS NULL;
-+ERROR:  permission denied
-+DETAIL:  The current user must have the CREATEROLE attribute.
-+ROLLBACK TO fail_seclabel_on_role;
-+ROLLBACK;
-diff --git a/ext-src/pg_anon-src/tests/expected/permissions_owner_1.out b/ext-src/pg_anon-src/tests/expected/permissions_owner_1.out
-new file mode 100644
-index 0000000..8b090fe
---- /dev/null
-+++ b/ext-src/pg_anon-src/tests/expected/permissions_owner_1.out
-@@ -0,0 +1,104 @@
-+BEGIN;
-+CREATE EXTENSION anon CASCADE;
-+NOTICE:  installing required extension "pgcrypto"
-+SELECT anon.init();
-+ init 
-+------
-+ t
-+(1 row)
-+
-+CREATE ROLE oscar_the_owner;
-+ALTER DATABASE :DBNAME OWNER TO oscar_the_owner;
-+CREATE ROLE mallory_the_masked_user;
-+SECURITY LABEL FOR anon ON ROLE mallory_the_masked_user IS 'MASKED';
-+--
-+-- We're checking the owner's permissions
-+--
-+-- see
-+-- https://postgresql-anonymizer.readthedocs.io/en/latest/SECURITY/#permissions
-+--
-+SET ROLE oscar_the_owner;
-+SELECT anon.pseudo_first_name(0) IS NOT NULL;
-+ ?column? 
-+----------
-+ t
-+(1 row)
-+
-+-- SHOULD FAIL
-+DO $$
-+BEGIN
-+  PERFORM anon.init();
-+  EXCEPTION WHEN insufficient_privilege
-+  THEN RAISE NOTICE 'insufficient_privilege';
-+END$$;
-+NOTICE:  insufficient_privilege
-+CREATE TABLE t1(i INT);
-+ALTER TABLE t1 ADD COLUMN t TEXT;
-+SECURITY LABEL FOR anon ON COLUMN t1.t
-+IS 'MASKED WITH VALUE NULL';
-+INSERT INTO t1 VALUES (1,'test');
-+SELECT anon.anonymize_table('t1');
-+ anonymize_table 
-+-----------------
-+ t
-+(1 row)
-+
-+SELECT * FROM t1;
-+ i | t 
-+---+---
-+ 1 | 
-+(1 row)
-+
-+UPDATE t1 SET t='test' WHERE i=1;
-+-- SHOULD FAIL
-+SAVEPOINT fail_start_engine;
-+SELECT anon.start_dynamic_masking();
-+ start_dynamic_masking 
-+-----------------------
-+ t
-+(1 row)
-+
-+ROLLBACK TO fail_start_engine;
-+RESET ROLE;
-+SELECT anon.start_dynamic_masking();
-+ start_dynamic_masking 
-+-----------------------
-+ t
-+(1 row)
-+
-+SET ROLE oscar_the_owner;
-+SELECT * FROM t1;
-+ i |  t   
-+---+------
-+ 1 | test
-+(1 row)
-+
-+--SELECT * FROM mask.t1;
-+-- SHOULD FAIL
-+SAVEPOINT fail_stop_engine;
-+SELECT anon.stop_dynamic_masking();
-+ERROR:  permission denied for schema mask
-+CONTEXT:  SQL statement "DROP VIEW mask.t1;"
-+PL/pgSQL function anon.mask_drop_view(oid) line 3 at EXECUTE
-+SQL statement "SELECT anon.mask_drop_view(oid)
-+  FROM pg_catalog.pg_class
-+  WHERE relnamespace=quote_ident(pg_catalog.current_setting('anon.sourceschema'))::REGNAMESPACE
-+  AND relkind IN ('r','p','f')"
-+PL/pgSQL function anon.stop_dynamic_masking() line 22 at PERFORM
-+ROLLBACK TO fail_stop_engine;
-+RESET ROLE;
-+SELECT anon.stop_dynamic_masking();
-+NOTICE:  The previous priviledges of 'mallory_the_masked_user' are not restored. You need to grant them manually.
-+ stop_dynamic_masking 
-+----------------------
-+ t
-+(1 row)
-+
-+SET ROLE oscar_the_owner;
-+-- SHOULD FAIL
-+SAVEPOINT fail_seclabel_on_role;
-+SECURITY LABEL FOR anon ON ROLE mallory_the_masked_user IS NULL;
-+ERROR:  permission denied
-+DETAIL:  The current user must have the CREATEROLE attribute.
-+ROLLBACK TO fail_seclabel_on_role;
-+ROLLBACK;

compute_tools/src/compute.rs

@@ -641,7 +641,26 @@ impl ComputeNode {
 
                 let log_directory_path = Path::new(&self.params.pgdata).join("log");
                 let log_directory_path = log_directory_path.to_string_lossy().to_string();
-                configure_audit_rsyslog(log_directory_path.clone(), "hipaa", &remote_endpoint)?;
+
+                // Add project_id,endpoint_id tag to identify the logs.
+                //
+                // These ids are passed from cplane,
+                // for backwards compatibility (old computes that don't have them),
+                // we set them to None.
+                // TODO: Clean up this code when all computes have them.
+                let tag: Option<String> = match (
+                    pspec.spec.project_id.as_deref(),
+                    pspec.spec.endpoint_id.as_deref(),
+                ) {
+                    (Some(project_id), Some(endpoint_id)) => {
+                        Some(format!("{project_id}/{endpoint_id}"))
+                    }
+                    (Some(project_id), None) => Some(format!("{project_id}/None")),
+                    (None, Some(endpoint_id)) => Some(format!("None,{endpoint_id}")),
+                    (None, None) => None,
+                };
+
+                configure_audit_rsyslog(log_directory_path.clone(), tag, &remote_endpoint)?;
 
                 // Launch a background task to clean up the audit logs
                 launch_pgaudit_gc(log_directory_path);

compute_tools/src/rsyslog.rs

@@ -50,13 +50,13 @@ fn restart_rsyslog() -> Result<()> {
 
 pub fn configure_audit_rsyslog(
     log_directory: String,
-    tag: &str,
+    tag: Option<String>,
     remote_endpoint: &str,
 ) -> Result<()> {
     let config_content: String = format!(
         include_str!("config_template/compute_audit_rsyslog_template.conf"),
         log_directory = log_directory,
-        tag = tag,
+        tag = tag.unwrap_or("".to_string()),
         remote_endpoint = remote_endpoint
     );
 

libs/pageserver_api/src/config.rs

@@ -682,10 +682,10 @@ pub mod tenant_conf_defaults {
     pub const DEFAULT_COMPACTION_SHARD_ANCESTOR: bool = true;
 
     // This value needs to be tuned to avoid OOM. We have 3/4*CPUs threads for L0 compaction, that's
-    // 3/4*16=9 on most of our pageservers. Compacting 20 layers requires about 1 GB memory (could
-    // be reduced later by optimizing L0 hole calculation to avoid loading all keys into memory). So
-    // with this config, we can get a maximum peak compaction usage of 9 GB.
-    pub const DEFAULT_COMPACTION_UPPER_LIMIT: usize = 20;
+    // 3/4*8=6 on most of our pageservers. Compacting 10 layers requires a maximum of
+    // DEFAULT_CHECKPOINT_DISTANCE*10 memory, that's 2560MB. So with this config, we can get a maximum peak
+    // compaction usage of 15360MB.
+    pub const DEFAULT_COMPACTION_UPPER_LIMIT: usize = 10;
     // Enable L0 compaction pass and semaphore by default. L0 compaction must be responsive to avoid
     // read amp.
     pub const DEFAULT_COMPACTION_L0_FIRST: bool = true;
@@ -702,8 +702,11 @@ pub mod tenant_conf_defaults {
     // Relevant: https://github.com/neondatabase/neon/issues/3394
     pub const DEFAULT_GC_PERIOD: &str = "1 hr";
     pub const DEFAULT_IMAGE_CREATION_THRESHOLD: usize = 3;
-    // If there are more than threshold * compaction_threshold (that is 3 * 10 in the default config) L0 layers, image
-    // layer creation will end immediately. Set to 0 to disable.
+    // Currently, any value other than 0 will trigger image layer creation preemption immediately with L0 backpressure
+    // without looking at the exact number of L0 layers.
+    // It was expected to have the following behavior:
+    // > If there are more than threshold * compaction_threshold (that is 3 * 10 in the default config) L0 layers, image
+    // > layer creation will end immediately. Set to 0 to disable.
     pub const DEFAULT_IMAGE_CREATION_PREEMPT_THRESHOLD: usize = 3;
     pub const DEFAULT_PITR_INTERVAL: &str = "7 days";
     pub const DEFAULT_WALRECEIVER_CONNECT_TIMEOUT: &str = "10 seconds";

pageserver/Cargo.toml

@@ -78,6 +78,7 @@ metrics.workspace = true
 pageserver_api.workspace = true
 pageserver_client.workspace = true # for ResponseErrorMessageExt TOOD refactor that
 pageserver_compaction.workspace = true
+pem.workspace = true
 postgres_connection.workspace = true
 postgres_ffi.workspace = true
 pq_proto.workspace = true

pageserver/pagebench/src/cmd/getpage_latest_lsn.rs

@@ -68,6 +68,13 @@ pub(crate) struct Args {
     targets: Option<Vec<TenantTimelineId>>,
 }
 
+/// State shared by all clients
+#[derive(Debug)]
+struct SharedState {
+    start_work_barrier: tokio::sync::Barrier,
+    live_stats: LiveStats,
+}
+
 #[derive(Debug, Default)]
 struct LiveStats {
     completed_requests: AtomicU64,
@@ -240,24 +247,26 @@ async fn main_impl(
         all_ranges
     };
 
-    let live_stats = Arc::new(LiveStats::default());
-
     let num_live_stats_dump = 1;
     let num_work_sender_tasks = args.num_clients.get() * timelines.len();
     let num_main_impl = 1;
 
-    let start_work_barrier = Arc::new(tokio::sync::Barrier::new(
-        num_live_stats_dump + num_work_sender_tasks + num_main_impl,
-    ));
+    let shared_state = Arc::new(SharedState {
+        start_work_barrier: tokio::sync::Barrier::new(
+            num_live_stats_dump + num_work_sender_tasks + num_main_impl,
+        ),
+        live_stats: LiveStats::default(),
+    });
+    let cancel = CancellationToken::new();
 
+    let ss = shared_state.clone();
     tokio::spawn({
-        let stats = Arc::clone(&live_stats);
-        let start_work_barrier = Arc::clone(&start_work_barrier);
         async move {
-            start_work_barrier.wait().await;
+            ss.start_work_barrier.wait().await;
             loop {
                 let start = std::time::Instant::now();
                 tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+                let stats = &ss.live_stats;
                 let completed_requests = stats.completed_requests.swap(0, Ordering::Relaxed);
                 let missed = stats.missed.swap(0, Ordering::Relaxed);
                 let elapsed = start.elapsed();
@@ -270,14 +279,12 @@ async fn main_impl(
         }
     });
 
-    let cancel = CancellationToken::new();
-
     let rps_period = args
         .per_client_rate
         .map(|rps_limit| Duration::from_secs_f64(1.0 / (rps_limit as f64)));
     let make_worker: &dyn Fn(WorkerId) -> Pin<Box<dyn Send + Future<Output = ()>>> = &|worker_id| {
-        let live_stats = live_stats.clone();
-        let start_work_barrier = start_work_barrier.clone();
+        let ss = shared_state.clone();
+        let cancel = cancel.clone();
         let ranges: Vec<KeyRange> = all_ranges
             .iter()
             .filter(|r| r.timeline == worker_id.timeline)
@@ -287,85 +294,8 @@ async fn main_impl(
             rand::distributions::weighted::WeightedIndex::new(ranges.iter().map(|v| v.len()))
                 .unwrap();
 
-        let cancel = cancel.clone();
         Box::pin(async move {
-            let client =
-                pageserver_client::page_service::Client::new(args.page_service_connstring.clone())
-                    .await
-                    .unwrap();
-            let mut client = client
-                .pagestream(worker_id.timeline.tenant_id, worker_id.timeline.timeline_id)
-                .await
-                .unwrap();
-
-            start_work_barrier.wait().await;
-            let client_start = Instant::now();
-            let mut ticks_processed = 0;
-            let mut inflight = VecDeque::new();
-            while !cancel.is_cancelled() {
-                // Detect if a request took longer than the RPS rate
-                if let Some(period) = &rps_period {
-                    let periods_passed_until_now =
-                        usize::try_from(client_start.elapsed().as_micros() / period.as_micros())
-                            .unwrap();
-
-                    if periods_passed_until_now > ticks_processed {
-                        live_stats.missed((periods_passed_until_now - ticks_processed) as u64);
-                    }
-                    ticks_processed = periods_passed_until_now;
-                }
-
-                while inflight.len() < args.queue_depth.get() {
-                    let start = Instant::now();
-                    let req = {
-                        let mut rng = rand::thread_rng();
-                        let r = &ranges[weights.sample(&mut rng)];
-                        let key: i128 = rng.gen_range(r.start..r.end);
-                        let key = Key::from_i128(key);
-                        assert!(key.is_rel_block_key());
-                        let (rel_tag, block_no) = key
-                            .to_rel_block()
-                            .expect("we filter non-rel-block keys out above");
-                        PagestreamGetPageRequest {
-                            hdr: PagestreamRequest {
-                                reqid: 0,
-                                request_lsn: if rng.gen_bool(args.req_latest_probability) {
-                                    Lsn::MAX
-                                } else {
-                                    r.timeline_lsn
-                                },
-                                not_modified_since: r.timeline_lsn,
-                            },
-                            rel: rel_tag,
-                            blkno: block_no,
-                        }
-                    };
-                    client.getpage_send(req).await.unwrap();
-                    inflight.push_back(start);
-                }
-
-                let start = inflight.pop_front().unwrap();
-                client.getpage_recv().await.unwrap();
-                let end = Instant::now();
-                live_stats.request_done();
-                ticks_processed += 1;
-                STATS.with(|stats| {
-                    stats
-                        .borrow()
-                        .lock()
-                        .unwrap()
-                        .observe(end.duration_since(start))
-                        .unwrap();
-                });
-
-                if let Some(period) = &rps_period {
-                    let next_at = client_start
-                        + Duration::from_micros(
-                            (ticks_processed) as u64 * u64::try_from(period.as_micros()).unwrap(),
-                        );
-                    tokio::time::sleep_until(next_at.into()).await;
-                }
-            }
+            client_libpq(args, worker_id, ss, cancel, rps_period, ranges, weights).await
         })
     };
 
@@ -387,7 +317,7 @@ async fn main_impl(
     };
 
     info!("waiting for everything to become ready");
-    start_work_barrier.wait().await;
+    shared_state.start_work_barrier.wait().await;
     info!("work started");
     if let Some(runtime) = args.runtime {
         tokio::time::sleep(runtime.into()).await;
@@ -416,3 +346,91 @@ async fn main_impl(
 
     anyhow::Ok(())
 }
+
+async fn client_libpq(
+    args: &Args,
+    worker_id: WorkerId,
+    shared_state: Arc<SharedState>,
+    cancel: CancellationToken,
+    rps_period: Option<Duration>,
+    ranges: Vec<KeyRange>,
+    weights: rand::distributions::weighted::WeightedIndex<i128>,
+) {
+    let client = pageserver_client::page_service::Client::new(args.page_service_connstring.clone())
+        .await
+        .unwrap();
+    let mut client = client
+        .pagestream(worker_id.timeline.tenant_id, worker_id.timeline.timeline_id)
+        .await
+        .unwrap();
+
+    shared_state.start_work_barrier.wait().await;
+    let client_start = Instant::now();
+    let mut ticks_processed = 0;
+    let mut inflight = VecDeque::new();
+    while !cancel.is_cancelled() {
+        // Detect if a request took longer than the RPS rate
+        if let Some(period) = &rps_period {
+            let periods_passed_until_now =
+                usize::try_from(client_start.elapsed().as_micros() / period.as_micros()).unwrap();
+
+            if periods_passed_until_now > ticks_processed {
+                shared_state
+                    .live_stats
+                    .missed((periods_passed_until_now - ticks_processed) as u64);
+            }
+            ticks_processed = periods_passed_until_now;
+        }
+
+        while inflight.len() < args.queue_depth.get() {
+            let start = Instant::now();
+            let req = {
+                let mut rng = rand::thread_rng();
+                let r = &ranges[weights.sample(&mut rng)];
+                let key: i128 = rng.gen_range(r.start..r.end);
+                let key = Key::from_i128(key);
+                assert!(key.is_rel_block_key());
+                let (rel_tag, block_no) = key
+                    .to_rel_block()
+                    .expect("we filter non-rel-block keys out above");
+                PagestreamGetPageRequest {
+                    hdr: PagestreamRequest {
+                        reqid: 0,
+                        request_lsn: if rng.gen_bool(args.req_latest_probability) {
+                            Lsn::MAX
+                        } else {
+                            r.timeline_lsn
+                        },
+                        not_modified_since: r.timeline_lsn,
+                    },
+                    rel: rel_tag,
+                    blkno: block_no,
+                }
+            };
+            client.getpage_send(req).await.unwrap();
+            inflight.push_back(start);
+        }
+
+        let start = inflight.pop_front().unwrap();
+        client.getpage_recv().await.unwrap();
+        let end = Instant::now();
+        shared_state.live_stats.request_done();
+        ticks_processed += 1;
+        STATS.with(|stats| {
+            stats
+                .borrow()
+                .lock()
+                .unwrap()
+                .observe(end.duration_since(start))
+                .unwrap();
+        });
+
+        if let Some(period) = &rps_period {
+            let next_at = client_start
+                + Duration::from_micros(
+                    (ticks_processed) as u64 * u64::try_from(period.as_micros()).unwrap(),
+                );
+            tokio::time::sleep_until(next_at.into()).await;
+        }
+    }
+}

pageserver/src/bin/pageserver.rs

@@ -416,8 +416,18 @@ fn start_pageserver(
     // The storage_broker::connect call needs to happen inside a tokio runtime thread.
     let broker_client = WALRECEIVER_RUNTIME
         .block_on(async {
+            let tls_config = storage_broker::ClientTlsConfig::new().ca_certificates(
+                conf.ssl_ca_certs
+                    .iter()
+                    .map(pem::encode)
+                    .map(storage_broker::Certificate::from_pem),
+            );
             // Note: we do not attempt connecting here (but validate endpoints sanity).
-            storage_broker::connect(conf.broker_endpoint.clone(), conf.broker_keepalive_interval)
+            storage_broker::connect(
+                conf.broker_endpoint.clone(),
+                conf.broker_keepalive_interval,
+                tls_config,
+            )
         })
         .with_context(|| {
             format!(

pageserver/src/config.rs

@@ -17,9 +17,10 @@ use once_cell::sync::OnceCell;
 use pageserver_api::config::{DiskUsageEvictionTaskConfig, MaxVectoredReadBytes};
 use pageserver_api::models::ImageCompressionAlgorithm;
 use pageserver_api::shard::TenantShardId;
+use pem::Pem;
 use postgres_backend::AuthType;
 use remote_storage::{RemotePath, RemoteStorageConfig};
-use reqwest::{Certificate, Url};
+use reqwest::Url;
 use storage_broker::Uri;
 use utils::id::{NodeId, TimelineId};
 use utils::logging::{LogFormat, SecretString};
@@ -67,8 +68,8 @@ pub struct PageServerConf {
     /// Period to reload certificate and private key from files.
     /// Default: 60s.
     pub ssl_cert_reload_period: Duration,
-    /// Trusted root CA certificates to use in https APIs.
-    pub ssl_ca_certs: Vec<Certificate>,
+    /// Trusted root CA certificates to use in https APIs in PEM format.
+    pub ssl_ca_certs: Vec<Pem>,
 
     /// Current availability zone. Used for traffic metrics.
     pub availability_zone: Option<String>,
@@ -497,7 +498,10 @@ impl PageServerConf {
             ssl_ca_certs: match ssl_ca_file {
                 Some(ssl_ca_file) => {
                     let buf = std::fs::read(ssl_ca_file)?;
-                    Certificate::from_pem_bundle(&buf)?
+                    pem::parse_many(&buf)?
+                        .into_iter()
+                        .filter(|pem| pem.tag() == "CERTIFICATE")
+                        .collect()
                 }
                 None => Vec::new(),
             },

pageserver/src/controller_upcall_client.rs

@@ -8,6 +8,7 @@ use pageserver_api::upcall_api::{
     ReAttachRequest, ReAttachResponse, ReAttachResponseTenant, ValidateRequest,
     ValidateRequestTenant, ValidateResponse,
 };
+use reqwest::Certificate;
 use serde::Serialize;
 use serde::de::DeserializeOwned;
 use tokio_util::sync::CancellationToken;
@@ -76,8 +77,8 @@ impl StorageControllerUpcallClient {
             client = client.default_headers(headers);
         }
 
-        for ssl_ca_cert in &conf.ssl_ca_certs {
-            client = client.add_root_certificate(ssl_ca_cert.clone());
+        for cert in &conf.ssl_ca_certs {
+            client = client.add_root_certificate(Certificate::from_der(cert.contents())?);
         }
 
         Ok(Some(Self {

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(),

pageserver/src/tenant/timeline/compaction.rs

@@ -7,7 +7,7 @@
 use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque};
 use std::ops::{Deref, Range};
 use std::sync::Arc;
-use std::time::{Duration, Instant, SystemTime};
+use std::time::{Duration, Instant};
 
 use super::layer_manager::LayerManager;
 use super::{
@@ -119,25 +119,32 @@ pub struct GcCompactionMetaStatistics {
     /// The layer size after compaction.
     pub after_compaction_layer_size: u64,
     /// The start time of the meta job.
-    pub start_time: Option<SystemTime>,
+    pub start_time: Option<chrono::DateTime<chrono::Utc>>,
     /// The end time of the meta job.
-    pub end_time: Option<SystemTime>,
+    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
     /// The duration of the meta job.
     pub duration_secs: f64,
     /// The id of the meta job.
     pub meta_job_id: GcCompactionJobId,
     /// The LSN below which the layers are compacted, used to compute the statistics.
     pub below_lsn: Lsn,
+    /// The retention ratio of the meta job (after_compaction_layer_size / before_compaction_layer_size)
+    pub retention_ratio: f64,
 }
 
 impl GcCompactionMetaStatistics {
     fn finalize(&mut self) {
-        let end_time = SystemTime::now();
+        let end_time = chrono::Utc::now();
         if let Some(start_time) = self.start_time {
-            if let Ok(duration) = end_time.duration_since(start_time) {
-                self.duration_secs = duration.as_secs_f64();
+            if end_time > start_time {
+                let delta = end_time - start_time;
+                if let Ok(std_dur) = delta.to_std() {
+                    self.duration_secs = std_dur.as_secs_f64();
+                }
             }
         }
+        self.retention_ratio = self.after_compaction_layer_size as f64
+            / (self.before_compaction_layer_size as f64 + 1.0);
         self.end_time = Some(end_time);
     }
 }
@@ -520,7 +527,7 @@ impl GcCompactionQueue {
                 }
                 guard.meta_statistics = Some(GcCompactionMetaStatistics {
                     meta_job_id: id,
-                    start_time: Some(SystemTime::now()),
+                    start_time: Some(chrono::Utc::now()),
                     before_compaction_layer_size: layer_size,
                     below_lsn: expected_l2_lsn,
                     total_sub_compaction_jobs: jobs_len,

pageserver/src/tenant/timeline/import_pgdata/upcall_api.rs

@@ -1,6 +1,6 @@
 //! FIXME: most of this is copy-paste from mgmt_api.rs ; dedupe into a `reqwest_utils::Client` crate.
 use pageserver_client::mgmt_api::{Error, ResponseErrorMessageExt};
-use reqwest::Method;
+use reqwest::{Certificate, Method};
 use serde::{Deserialize, Serialize};
 use tokio_util::sync::CancellationToken;
 use tracing::error;
@@ -34,7 +34,7 @@ impl Client {
         };
         let mut http_client = reqwest::Client::builder();
         for cert in &conf.ssl_ca_certs {
-            http_client = http_client.add_root_certificate(cert.clone());
+            http_client = http_client.add_root_certificate(Certificate::from_der(cert.contents())?);
         }
         let http_client = http_client.build()?;
 

proxy/src/serverless/backend.rs

@@ -16,9 +16,9 @@ use tracing::field::display;
 use tracing::{debug, info};
 
 use super::AsyncRW;
-use super::conn_pool::poll_client;
+use super::conn_pool::poll_client_generic;
 use super::conn_pool_lib::{Client, ConnInfo, EndpointConnPool, GlobalConnPool};
-use super::http_conn_pool::{self, HttpConnPool, Send, poll_http2_client};
+use super::http_conn_pool::{self, HttpConnPool};
 use super::local_conn_pool::{self, EXT_NAME, EXT_SCHEMA, EXT_VERSION, LocalConnPool};
 use crate::auth::backend::local::StaticAuthRules;
 use crate::auth::backend::{ComputeCredentials, ComputeUserInfo};
@@ -42,10 +42,9 @@ use crate::rate_limiter::EndpointRateLimiter;
 use crate::types::{EndpointId, Host, LOCAL_PROXY_SUFFIX};
 
 pub(crate) struct PoolingBackend {
-    pub(crate) http_conn_pool: Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
+    pub(crate) http_conn_pool: Arc<GlobalConnPool<HttpConnPool>>,
     pub(crate) local_pool: Arc<LocalConnPool<postgres_client::Client>>,
-    pub(crate) pool:
-        Arc<GlobalConnPool<postgres_client::Client, EndpointConnPool<postgres_client::Client>>>,
+    pub(crate) pool: Arc<GlobalConnPool<EndpointConnPool<postgres_client::Client>>>,
 
     pub(crate) config: &'static ProxyConfig,
     pub(crate) auth_backend: &'static crate::auth::Backend<'static, ()>,
@@ -212,7 +211,7 @@ impl PoolingBackend {
             None
         } else {
             debug!("pool: looking for an existing connection");
-            self.pool.get(ctx, &conn_info)?
+            self.pool.get(ctx, &conn_info)
         };
 
         if let Some(client) = maybe_client {
@@ -246,9 +245,9 @@ impl PoolingBackend {
         &self,
         ctx: &RequestContext,
         conn_info: ConnInfo,
-    ) -> Result<http_conn_pool::Client<Send>, HttpConnError> {
+    ) -> Result<http_conn_pool::Client, HttpConnError> {
         debug!("pool: looking for an existing connection");
-        if let Ok(Some(client)) = self.http_conn_pool.get(ctx, &conn_info) {
+        if let Some(client) = self.http_conn_pool.get(ctx, &conn_info) {
             return Ok(client);
         }
 
@@ -532,7 +531,7 @@ impl ShouldRetryWakeCompute for LocalProxyConnError {
 }
 
 struct TokioMechanism {
-    pool: Arc<GlobalConnPool<postgres_client::Client, EndpointConnPool<postgres_client::Client>>>,
+    pool: Arc<GlobalConnPool<EndpointConnPool<postgres_client::Client>>>,
     conn_info: ConnInfo,
     conn_id: uuid::Uuid,
 
@@ -578,7 +577,7 @@ impl ConnectMechanism for TokioMechanism {
             info!("latency={}, query_id={}", ctx.get_proxy_latency(), query_id);
         }
 
-        Ok(poll_client(
+        Ok(poll_client_generic(
             self.pool.clone(),
             ctx,
             self.conn_info.clone(),
@@ -593,7 +592,7 @@ impl ConnectMechanism for TokioMechanism {
 }
 
 struct HyperMechanism {
-    pool: Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
+    pool: Arc<GlobalConnPool<HttpConnPool>>,
     conn_info: ConnInfo,
     conn_id: uuid::Uuid,
 
@@ -603,7 +602,7 @@ struct HyperMechanism {
 
 #[async_trait]
 impl ConnectMechanism for HyperMechanism {
-    type Connection = http_conn_pool::Client<Send>;
+    type Connection = http_conn_pool::Client;
     type ConnectError = HttpConnError;
     type Error = HttpConnError;
 
@@ -639,15 +638,26 @@ impl ConnectMechanism for HyperMechanism {
             info!("latency={}, query_id={}", ctx.get_proxy_latency(), query_id);
         }
 
-        Ok(poll_http2_client(
+        let client = poll_client_generic(
             self.pool.clone(),
             ctx,
-            &self.conn_info,
+            self.conn_info.clone(),
             client,
             connection,
             self.conn_id,
             node_info.aux.clone(),
-        ))
+        );
+
+        // auth-broker -> local-proxy clients don't return to the pool, since
+        // they are multiplexing and cloneable. So instead we insert it once here.
+        if let Some(endpoint) = self.conn_info.endpoint_cache_key() {
+            self.pool
+                .get_or_create_endpoint_pool(&endpoint)
+                .write()
+                .register(&client);
+        }
+
+        Ok(client)
     }
 
     fn update_connect_config(&self, _config: &mut compute::ConnCfg) {}

proxy/src/serverless/conn_pool.rs

@@ -11,7 +11,7 @@ use smallvec::SmallVec;
 use tokio::net::TcpStream;
 use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
-use tracing::{Instrument, error, info, info_span, warn};
+use tracing::{error, info, info_span, warn};
 #[cfg(test)]
 use {
     super::conn_pool_lib::GlobalConnPoolOptions,
@@ -20,8 +20,7 @@ use {
 };
 
 use super::conn_pool_lib::{
-    Client, ClientDataEnum, ClientInnerCommon, ClientInnerExt, ConnInfo, EndpointConnPool,
-    GlobalConnPool,
+    ClientDataEnum, ClientInnerCommon, ConnInfo, EndpointConnPoolExt, GlobalConnPool,
 };
 use crate::context::RequestContext;
 use crate::control_plane::messages::MetricsAuxInfo;
@@ -29,6 +28,7 @@ use crate::metrics::Metrics;
 use crate::tls::postgres_rustls::MakeRustlsConnect;
 
 type TlsStream = <MakeRustlsConnect as MakeTlsConnect<TcpStream>>::Stream;
+pub(super) type Conn = postgres_client::Connection<TcpStream, TlsStream>;
 
 #[derive(Debug, Clone)]
 pub(crate) struct ConnInfoWithAuth {
@@ -56,20 +56,20 @@ impl fmt::Display for ConnInfo {
     }
 }
 
-pub(crate) fn poll_client<C: ClientInnerExt>(
-    global_pool: Arc<GlobalConnPool<C, EndpointConnPool<C>>>,
+pub(crate) fn poll_client_generic<P: EndpointConnPoolExt>(
+    global_pool: Arc<GlobalConnPool<P>>,
     ctx: &RequestContext,
     conn_info: ConnInfo,
-    client: C,
-    mut connection: postgres_client::Connection<TcpStream, TlsStream>,
+    client: P::ClientInner,
+    connection: P::Connection,
     conn_id: uuid::Uuid,
     aux: MetricsAuxInfo,
-) -> Client<C> {
+) -> P::Client {
     let conn_gauge = Metrics::get().proxy.db_connections.guard(ctx.protocol());
-    let mut session_id = ctx.session_id();
+    let session_id = ctx.session_id();
     let (tx, mut rx) = tokio::sync::watch::channel(session_id);
 
-    let span = info_span!(parent: None, "connection", %conn_id);
+    let span = info_span!(parent: None, "connection", %conn_id, %session_id);
     let cold_start_info = ctx.cold_start_info();
     span.in_scope(|| {
         info!(cold_start_info = cold_start_info.as_str(), %conn_info, %session_id, "new connection");
@@ -85,27 +85,30 @@ pub(crate) fn poll_client<C: ClientInnerExt>(
     let cancel = CancellationToken::new();
     let cancelled = cancel.clone().cancelled_owned();
 
-    tokio::spawn(
-    async move {
+    tokio::spawn(async move {
         let _conn_gauge = conn_gauge;
         let mut idle_timeout = pin!(tokio::time::sleep(idle));
         let mut cancelled = pin!(cancelled);
+        let mut connection = pin!(P::spawn_conn(connection));
 
         poll_fn(move |cx| {
+            let _enter = span.enter();
+
             if cancelled.as_mut().poll(cx).is_ready() {
                 info!("connection dropped");
-                return Poll::Ready(())
+                return Poll::Ready(());
             }
 
             match rx.has_changed() {
                 Ok(true) => {
-                    session_id = *rx.borrow_and_update();
-                    info!(%session_id, "changed session");
+                    let session_id = *rx.borrow_and_update();
+                    span.record("session_id", tracing::field::display(session_id));
+                    info!("changed session");
                     idle_timeout.as_mut().reset(Instant::now() + idle);
                 }
                 Err(_) => {
                     info!("connection dropped");
-                    return Poll::Ready(())
+                    return Poll::Ready(());
                 }
                 _ => {}
             }
@@ -117,48 +120,25 @@ pub(crate) fn poll_client<C: ClientInnerExt>(
                 if let Some(pool) = pool.clone().upgrade() {
                     // remove client from pool - should close the connection if it's idle.
                     // does nothing if the client is currently checked-out and in-use
-                    if pool.write().remove_client(db_user.clone(), conn_id) {
+                    if pool.write().remove_conn(db_user.clone(), conn_id) {
                         info!("idle connection removed");
                     }
                 }
             }
 
-            loop {
-                let message = ready!(connection.poll_message(cx));
-
-                match message {
-                    Some(Ok(AsyncMessage::Notice(notice))) => {
-                        info!(%session_id, "notice: {}", notice);
-                    }
-                    Some(Ok(AsyncMessage::Notification(notif))) => {
-                        warn!(%session_id, pid = notif.process_id(), channel = notif.channel(), "notification received");
-                    }
-                    Some(Ok(_)) => {
-                        warn!(%session_id, "unknown message");
-                    }
-                    Some(Err(e)) => {
-                        error!(%session_id, "connection error: {}", e);
-                        break
-                    }
-                    None => {
-                        info!("connection closed");
-                        break
-                    }
-                }
-            }
+            ready!(connection.as_mut().poll(cx));
 
             // remove from connection pool
             if let Some(pool) = pool.clone().upgrade() {
-                if pool.write().remove_client(db_user.clone(), conn_id) {
+                if pool.write().remove_conn(db_user.clone(), conn_id) {
                     info!("closed connection removed");
                 }
             }
 
             Poll::Ready(())
-        }).await;
-
-    }
-    .instrument(span));
+        })
+        .await;
+    });
     let inner = ClientInnerCommon {
         inner: client,
         aux,
@@ -169,7 +149,42 @@ pub(crate) fn poll_client<C: ClientInnerExt>(
         }),
     };
 
-    Client::new(inner, conn_info, pool_clone)
+    P::wrap_client(inner, conn_info, pool_clone)
+}
+
+pub async fn poll_tokio_postgres_conn_really(mut connection: Conn) {
+    poll_fn(move |cx| {
+        loop {
+            let message = ready!(connection.poll_message(cx));
+
+            match message {
+                Some(Ok(AsyncMessage::Notice(notice))) => {
+                    info!("notice: {}", notice);
+                }
+                Some(Ok(AsyncMessage::Notification(notif))) => {
+                    warn!(
+                        pid = notif.process_id(),
+                        channel = notif.channel(),
+                        "notification received"
+                    );
+                }
+                Some(Ok(_)) => {
+                    warn!("unknown message");
+                }
+                Some(Err(e)) => {
+                    error!("connection error: {}", e);
+                    break;
+                }
+                None => {
+                    info!("connection closed");
+                    break;
+                }
+            }
+        }
+
+        Poll::Ready(())
+    })
+    .await;
 }
 
 #[derive(Clone)]
@@ -179,11 +194,11 @@ pub(crate) struct ClientDataRemote {
 }
 
 impl ClientDataRemote {
-    pub fn session(&mut self) -> &mut tokio::sync::watch::Sender<uuid::Uuid> {
-        &mut self.session
+    pub fn session(&self) -> &tokio::sync::watch::Sender<uuid::Uuid> {
+        &self.session
     }
 
-    pub fn cancel(&mut self) {
+    pub fn cancel(&self) {
         self.cancel.cancel();
     }
 }
@@ -195,6 +210,7 @@ mod tests {
     use super::*;
     use crate::proxy::NeonOptions;
     use crate::serverless::cancel_set::CancelSet;
+    use crate::serverless::conn_pool_lib::{Client, ClientInnerExt};
     use crate::types::{BranchId, EndpointId, ProjectId};
 
     struct MockClient(Arc<AtomicBool>);

proxy/src/serverless/conn_pool_lib.rs

@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::marker::PhantomData;
 use std::ops::Deref;
 use std::sync::atomic::{self, AtomicUsize};
 use std::sync::{Arc, Weak};
@@ -12,11 +11,10 @@ use rand::Rng;
 use smol_str::ToSmolStr;
 use tracing::{Span, debug, info};
 
-use super::backend::HttpConnError;
-use super::conn_pool::ClientDataRemote;
-use super::http_conn_pool::ClientDataHttp;
+use super::conn_pool::{ClientDataRemote, poll_tokio_postgres_conn_really};
 use super::local_conn_pool::ClientDataLocal;
 use crate::auth::backend::ComputeUserInfo;
+use crate::config::HttpConfig;
 use crate::context::RequestContext;
 use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo};
 use crate::metrics::{HttpEndpointPoolsGuard, Metrics};
@@ -51,7 +49,6 @@ impl ConnInfo {
 pub(crate) enum ClientDataEnum {
     Remote(ClientDataRemote),
     Local(ClientDataLocal),
-    Http(ClientDataHttp),
 }
 
 #[derive(Clone)]
@@ -64,14 +61,9 @@ pub(crate) struct ClientInnerCommon<C: ClientInnerExt> {
 
 impl<C: ClientInnerExt> Drop for ClientInnerCommon<C> {
     fn drop(&mut self) {
-        match &mut self.data {
-            ClientDataEnum::Remote(remote_data) => {
-                remote_data.cancel();
-            }
-            ClientDataEnum::Local(local_data) => {
-                local_data.cancel();
-            }
-            ClientDataEnum::Http(_http_data) => (),
+        match &self.data {
+            ClientDataEnum::Remote(remote_data) => remote_data.cancel(),
+            ClientDataEnum::Local(local_data) => local_data.cancel(),
         }
     }
 }
@@ -81,8 +73,8 @@ impl<C: ClientInnerExt> ClientInnerCommon<C> {
         self.conn_id
     }
 
-    pub(crate) fn get_data(&mut self) -> &mut ClientDataEnum {
-        &mut self.data
+    pub(crate) fn get_data(&self) -> &ClientDataEnum {
+        &self.data
     }
 }
 
@@ -326,12 +318,70 @@ impl<C: ClientInnerExt> DbUserConn<C> for DbUserConnPool<C> {
     }
 }
 
-pub(crate) trait EndpointConnPoolExt<C: ClientInnerExt> {
+pub(crate) trait EndpointConnPoolExt: Send + Sync + 'static {
+    type Client;
+    type ClientInner: ClientInnerExt;
+    type Connection: Send + 'static;
+
+    fn create(config: &HttpConfig, global_connections_count: Arc<AtomicUsize>) -> Self;
+    fn wrap_client(
+        inner: ClientInnerCommon<Self::ClientInner>,
+        conn_info: ConnInfo,
+        pool: Weak<RwLock<Self>>,
+    ) -> Self::Client;
+
+    fn get_conn_entry(
+        &mut self,
+        db_user: (DbName, RoleName),
+    ) -> Option<ClientInnerCommon<Self::ClientInner>>;
+    fn remove_conn(&mut self, db_user: (DbName, RoleName), conn_id: uuid::Uuid) -> bool;
+
+    fn spawn_conn(conn: Self::Connection) -> impl Future<Output = ()> + Send + 'static;
+
     fn clear_closed(&mut self) -> usize;
     fn total_conns(&self) -> usize;
 }
 
-impl<C: ClientInnerExt> EndpointConnPoolExt<C> for EndpointConnPool<C> {
+impl<C: ClientInnerExt> EndpointConnPoolExt for EndpointConnPool<C> {
+    type Client = Client<C>;
+    type ClientInner = C;
+    type Connection = super::conn_pool::Conn;
+
+    fn create(config: &HttpConfig, global_connections_count: Arc<AtomicUsize>) -> Self {
+        EndpointConnPool {
+            pools: HashMap::new(),
+            total_conns: 0,
+            max_conns: config.pool_options.max_conns_per_endpoint,
+            _guard: Metrics::get().proxy.http_endpoint_pools.guard(),
+            global_connections_count,
+            global_pool_size_max_conns: config.pool_options.max_total_conns,
+            pool_name: String::from("remote"),
+        }
+    }
+
+    fn wrap_client(
+        client: ClientInnerCommon<Self::ClientInner>,
+        conn_info: ConnInfo,
+        pool: Weak<RwLock<Self>>,
+    ) -> Self::Client {
+        Client::new(client, conn_info.clone(), pool)
+    }
+
+    fn get_conn_entry(
+        &mut self,
+        db_user: (DbName, RoleName),
+    ) -> Option<ClientInnerCommon<Self::ClientInner>> {
+        Some(self.get_conn_entry(db_user)?.conn)
+    }
+
+    fn remove_conn(&mut self, db_user: (DbName, RoleName), conn_id: uuid::Uuid) -> bool {
+        self.remove_client(db_user, conn_id)
+    }
+
+    async fn spawn_conn(conn: Self::Connection) {
+        poll_tokio_postgres_conn_really(conn).await;
+    }
+
     fn clear_closed(&mut self) -> usize {
         let mut clients_removed: usize = 0;
         for db_pool in self.pools.values_mut() {
@@ -345,10 +395,9 @@ impl<C: ClientInnerExt> EndpointConnPoolExt<C> for EndpointConnPool<C> {
     }
 }
 
-pub(crate) struct GlobalConnPool<C, P>
+pub(crate) struct GlobalConnPool<P>
 where
-    C: ClientInnerExt,
-    P: EndpointConnPoolExt<C>,
+    P: EndpointConnPoolExt,
 {
     // endpoint -> per-endpoint connection pool
     //
@@ -367,8 +416,6 @@ where
     pub(crate) global_connections_count: Arc<AtomicUsize>,
 
     pub(crate) config: &'static crate::config::HttpConfig,
-
-    _marker: PhantomData<C>,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -391,10 +438,9 @@ pub struct GlobalConnPoolOptions {
     pub max_total_conns: usize,
 }
 
-impl<C, P> GlobalConnPool<C, P>
+impl<P> GlobalConnPool<P>
 where
-    C: ClientInnerExt,
-    P: EndpointConnPoolExt<C>,
+    P: EndpointConnPoolExt,
 {
     pub(crate) fn new(config: &'static crate::config::HttpConfig) -> Arc<Self> {
         let shards = config.pool_options.pool_shards;
@@ -403,7 +449,6 @@ where
             global_pool_size: AtomicUsize::new(0),
             config,
             global_connections_count: Arc::new(AtomicUsize::new(0)),
-            _marker: PhantomData,
         })
     }
 
@@ -492,80 +537,72 @@ where
     }
 }
 
-impl<C: ClientInnerExt> GlobalConnPool<C, EndpointConnPool<C>> {
+impl<P: EndpointConnPoolExt> GlobalConnPool<P> {
     pub(crate) fn get(
         self: &Arc<Self>,
         ctx: &RequestContext,
         conn_info: &ConnInfo,
-    ) -> Result<Option<Client<C>>, HttpConnError> {
-        let mut client: Option<ClientInnerCommon<C>> = None;
-        let Some(endpoint) = conn_info.endpoint_cache_key() else {
-            return Ok(None);
-        };
+    ) -> Option<P::Client> {
+        let endpoint = conn_info.endpoint_cache_key()?;
 
-        let endpoint_pool = self.get_or_create_endpoint_pool(&endpoint);
-        if let Some(entry) = endpoint_pool
+        let endpoint_pool = self.get_endpoint_pool(&endpoint)?;
+        let client = endpoint_pool
             .write()
-            .get_conn_entry(conn_info.db_and_user())
-        {
-            client = Some(entry.conn);
-        }
+            .get_conn_entry(conn_info.db_and_user())?;
+
         let endpoint_pool = Arc::downgrade(&endpoint_pool);
 
-        // ok return cached connection if found and establish a new one otherwise
-        if let Some(mut client) = client {
-            if client.inner.is_closed() {
-                info!("pool: cached connection '{conn_info}' is closed, opening a new one");
-                return Ok(None);
-            }
-            tracing::Span::current()
-                .record("conn_id", tracing::field::display(client.get_conn_id()));
-            tracing::Span::current().record(
-                "pid",
-                tracing::field::display(client.inner.get_process_id()),
-            );
-            debug!(
-                cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
-                "pool: reusing connection '{conn_info}'"
-            );
-
-            match client.get_data() {
-                ClientDataEnum::Local(data) => {
-                    data.session().send(ctx.session_id())?;
-                }
-
-                ClientDataEnum::Remote(data) => {
-                    data.session().send(ctx.session_id())?;
-                }
-                ClientDataEnum::Http(_) => (),
-            }
-
-            ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit);
-            ctx.success();
-            return Ok(Some(Client::new(client, conn_info.clone(), endpoint_pool)));
+        if client.inner.is_closed() {
+            info!("pool: cached connection '{conn_info}' is closed, opening a new one");
+            return None;
         }
-        Ok(None)
+
+        tracing::Span::current().record("conn_id", tracing::field::display(client.get_conn_id()));
+        tracing::Span::current().record(
+            "pid",
+            tracing::field::display(client.inner.get_process_id()),
+        );
+        debug!(
+            cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
+            "pool: reusing connection '{conn_info}'"
+        );
+
+        match client.get_data() {
+            ClientDataEnum::Local(data) => {
+                data.session().send(ctx.session_id()).ok()?;
+            }
+            ClientDataEnum::Remote(data) => {
+                data.session().send(ctx.session_id()).ok()?;
+            }
+        }
+
+        ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit);
+        Some(P::wrap_client(client, conn_info.clone(), endpoint_pool))
+    }
+}
+
+impl<P: EndpointConnPoolExt> GlobalConnPool<P> {
+    pub(crate) fn get_endpoint_pool(
+        self: &Arc<Self>,
+        endpoint: &EndpointCacheKey,
+    ) -> Option<Arc<RwLock<P>>> {
+        Some(self.global_pool.get(endpoint)?.clone())
     }
 
     pub(crate) fn get_or_create_endpoint_pool(
         self: &Arc<Self>,
         endpoint: &EndpointCacheKey,
-    ) -> Arc<RwLock<EndpointConnPool<C>>> {
+    ) -> Arc<RwLock<P>> {
         // fast path
         if let Some(pool) = self.global_pool.get(endpoint) {
             return pool.clone();
         }
 
         // slow path
-        let new_pool = Arc::new(RwLock::new(EndpointConnPool {
-            pools: HashMap::new(),
-            total_conns: 0,
-            max_conns: self.config.pool_options.max_conns_per_endpoint,
-            _guard: Metrics::get().proxy.http_endpoint_pools.guard(),
-            global_connections_count: self.global_connections_count.clone(),
-            global_pool_size_max_conns: self.config.pool_options.max_total_conns,
-            pool_name: String::from("remote"),
-        }));
+        let new_pool = Arc::new(RwLock::new(P::create(
+            self.config,
+            self.global_connections_count.clone(),
+        )));
 
         // find or create a pool for this endpoint
         let mut created = false;
@@ -592,6 +629,7 @@ impl<C: ClientInnerExt> GlobalConnPool<C, EndpointConnPool<C>> {
         pool
     }
 }
+
 pub(crate) struct Client<C: ClientInnerExt> {
     span: Span,
     inner: Option<ClientInnerCommon<C>>,

proxy/src/serverless/http_conn_pool.rs

@@ -4,32 +4,25 @@ use std::sync::{Arc, Weak};
 
 use hyper::client::conn::http2;
 use hyper_util::rt::{TokioExecutor, TokioIo};
-use parking_lot::RwLock;
 use smol_str::ToSmolStr;
-use tracing::{Instrument, debug, error, info, info_span};
+use tracing::{error, info};
 
 use super::AsyncRW;
-use super::backend::HttpConnError;
 use super::conn_pool_lib::{
-    ClientDataEnum, ClientInnerCommon, ClientInnerExt, ConnInfo, ConnPoolEntry,
-    EndpointConnPoolExt, GlobalConnPool,
+    ClientInnerCommon, ClientInnerExt, ConnInfo, ConnPoolEntry, EndpointConnPoolExt,
 };
+use crate::config::HttpConfig;
 use crate::context::RequestContext;
-use crate::control_plane::messages::{ColdStartInfo, MetricsAuxInfo};
 use crate::metrics::{HttpEndpointPoolsGuard, Metrics};
 use crate::protocol2::ConnectionInfoExtra;
-use crate::types::EndpointCacheKey;
 use crate::usage_metrics::{Ids, MetricCounter, USAGE_METRICS};
 
 pub(crate) type Send = http2::SendRequest<hyper::body::Incoming>;
 pub(crate) type Connect = http2::Connection<TokioIo<AsyncRW>, hyper::body::Incoming, TokioExecutor>;
 
-#[derive(Clone)]
-pub(crate) struct ClientDataHttp();
-
 // Per-endpoint connection pool
 // Number of open connections is limited by the `max_conns_per_endpoint`.
-pub(crate) struct HttpConnPool<C: ClientInnerExt + Clone> {
+pub(crate) struct HttpConnPool {
     // TODO(conrad):
     // either we should open more connections depending on stream count
     // (not exposed by hyper, need our own counter)
@@ -39,13 +32,13 @@ pub(crate) struct HttpConnPool<C: ClientInnerExt + Clone> {
     // seems somewhat redundant though.
     //
     // Probably we should run a semaphore and just the single conn. TBD.
-    conns: VecDeque<ConnPoolEntry<C>>,
+    conns: VecDeque<ConnPoolEntry<Send>>,
     _guard: HttpEndpointPoolsGuard<'static>,
     global_connections_count: Arc<AtomicUsize>,
 }
 
-impl<C: ClientInnerExt + Clone> HttpConnPool<C> {
-    fn get_conn_entry(&mut self) -> Option<ConnPoolEntry<C>> {
+impl HttpConnPool {
+    fn get_conn_entry(&mut self) -> Option<ConnPoolEntry<Send>> {
         let Self { conns, .. } = self;
 
         loop {
@@ -83,9 +76,59 @@ impl<C: ClientInnerExt + Clone> HttpConnPool<C> {
         }
         removed > 0
     }
+
+    pub fn register(&mut self, client: &Client) {
+        self.conns.push_back(ConnPoolEntry {
+            conn: client.inner.clone(),
+            _last_access: std::time::Instant::now(),
+        });
+    }
 }
 
-impl<C: ClientInnerExt + Clone> EndpointConnPoolExt<C> for HttpConnPool<C> {
+impl EndpointConnPoolExt for HttpConnPool {
+    type Client = Client;
+    type ClientInner = Send;
+    type Connection = Connect;
+
+    fn create(_config: &HttpConfig, global_connections_count: Arc<AtomicUsize>) -> Self {
+        HttpConnPool {
+            conns: VecDeque::new(),
+            _guard: Metrics::get().proxy.http_endpoint_pools.guard(),
+            global_connections_count,
+        }
+    }
+
+    fn wrap_client(
+        inner: ClientInnerCommon<Self::ClientInner>,
+        _conn_info: ConnInfo,
+        _pool: Weak<parking_lot::RwLock<Self>>,
+    ) -> Self::Client {
+        Client::new(inner)
+    }
+
+    fn get_conn_entry(
+        &mut self,
+        _db_user: (crate::types::DbName, crate::types::RoleName),
+    ) -> Option<ClientInnerCommon<Self::ClientInner>> {
+        Some(self.get_conn_entry()?.conn)
+    }
+
+    fn remove_conn(
+        &mut self,
+        _db_user: (crate::types::DbName, crate::types::RoleName),
+        conn_id: uuid::Uuid,
+    ) -> bool {
+        self.remove_conn(conn_id)
+    }
+
+    async fn spawn_conn(conn: Self::Connection) {
+        let res = conn.await;
+        match res {
+            Ok(()) => info!("connection closed"),
+            Err(e) => error!("connection error: {e:?}"),
+        }
+    }
+
     fn clear_closed(&mut self) -> usize {
         let Self { conns, .. } = self;
         let old_len = conns.len();
@@ -100,7 +143,7 @@ impl<C: ClientInnerExt + Clone> EndpointConnPoolExt<C> for HttpConnPool<C> {
     }
 }
 
-impl<C: ClientInnerExt + Clone> Drop for HttpConnPool<C> {
+impl Drop for HttpConnPool {
     fn drop(&mut self) {
         if !self.conns.is_empty() {
             self.global_connections_count
@@ -114,154 +157,12 @@ impl<C: ClientInnerExt + Clone> Drop for HttpConnPool<C> {
     }
 }
 
-impl<C: ClientInnerExt + Clone> GlobalConnPool<C, HttpConnPool<C>> {
-    #[expect(unused_results)]
-    pub(crate) fn get(
-        self: &Arc<Self>,
-        ctx: &RequestContext,
-        conn_info: &ConnInfo,
-    ) -> Result<Option<Client<C>>, HttpConnError> {
-        let result: Result<Option<Client<C>>, HttpConnError>;
-        let Some(endpoint) = conn_info.endpoint_cache_key() else {
-            result = Ok(None);
-            return result;
-        };
-        let endpoint_pool = self.get_or_create_endpoint_pool(&endpoint);
-        let Some(client) = endpoint_pool.write().get_conn_entry() else {
-            result = Ok(None);
-            return result;
-        };
-
-        tracing::Span::current().record("conn_id", tracing::field::display(client.conn.conn_id));
-        debug!(
-            cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
-            "pool: reusing connection '{conn_info}'"
-        );
-        ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit);
-        ctx.success();
-
-        Ok(Some(Client::new(client.conn.clone())))
-    }
-
-    fn get_or_create_endpoint_pool(
-        self: &Arc<Self>,
-        endpoint: &EndpointCacheKey,
-    ) -> Arc<RwLock<HttpConnPool<C>>> {
-        // fast path
-        if let Some(pool) = self.global_pool.get(endpoint) {
-            return pool.clone();
-        }
-
-        // slow path
-        let new_pool = Arc::new(RwLock::new(HttpConnPool {
-            conns: VecDeque::new(),
-            _guard: Metrics::get().proxy.http_endpoint_pools.guard(),
-            global_connections_count: self.global_connections_count.clone(),
-        }));
-
-        // find or create a pool for this endpoint
-        let mut created = false;
-        let pool = self
-            .global_pool
-            .entry(endpoint.clone())
-            .or_insert_with(|| {
-                created = true;
-                new_pool
-            })
-            .clone();
-
-        // log new global pool size
-        if created {
-            let global_pool_size = self
-                .global_pool_size
-                .fetch_add(1, atomic::Ordering::Relaxed)
-                + 1;
-            info!(
-                "pool: created new pool for '{endpoint}', global pool size now {global_pool_size}"
-            );
-        }
-
-        pool
-    }
+pub(crate) struct Client {
+    pub(crate) inner: ClientInnerCommon<Send>,
 }
 
-pub(crate) fn poll_http2_client(
-    global_pool: Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
-    ctx: &RequestContext,
-    conn_info: &ConnInfo,
-    client: Send,
-    connection: Connect,
-    conn_id: uuid::Uuid,
-    aux: MetricsAuxInfo,
-) -> Client<Send> {
-    let conn_gauge = Metrics::get().proxy.db_connections.guard(ctx.protocol());
-    let session_id = ctx.session_id();
-
-    let span = info_span!(parent: None, "connection", %conn_id);
-    let cold_start_info = ctx.cold_start_info();
-    span.in_scope(|| {
-        info!(cold_start_info = cold_start_info.as_str(), %conn_info, %session_id, "new connection");
-    });
-
-    let pool = match conn_info.endpoint_cache_key() {
-        Some(endpoint) => {
-            let pool = global_pool.get_or_create_endpoint_pool(&endpoint);
-            let client = ClientInnerCommon {
-                inner: client.clone(),
-                aux: aux.clone(),
-                conn_id,
-                data: ClientDataEnum::Http(ClientDataHttp()),
-            };
-            pool.write().conns.push_back(ConnPoolEntry {
-                conn: client,
-                _last_access: std::time::Instant::now(),
-            });
-            Metrics::get()
-                .proxy
-                .http_pool_opened_connections
-                .get_metric()
-                .inc();
-
-            Arc::downgrade(&pool)
-        }
-        None => Weak::new(),
-    };
-
-    tokio::spawn(
-        async move {
-            let _conn_gauge = conn_gauge;
-            let res = connection.await;
-            match res {
-                Ok(()) => info!("connection closed"),
-                Err(e) => error!(%session_id, "connection error: {e:?}"),
-            }
-
-            // remove from connection pool
-            if let Some(pool) = pool.clone().upgrade() {
-                if pool.write().remove_conn(conn_id) {
-                    info!("closed connection removed");
-                }
-            }
-        }
-        .instrument(span),
-    );
-
-    let client = ClientInnerCommon {
-        inner: client,
-        aux,
-        conn_id,
-        data: ClientDataEnum::Http(ClientDataHttp()),
-    };
-
-    Client::new(client)
-}
-
-pub(crate) struct Client<C: ClientInnerExt + Clone> {
-    pub(crate) inner: ClientInnerCommon<C>,
-}
-
-impl<C: ClientInnerExt + Clone> Client<C> {
-    pub(self) fn new(inner: ClientInnerCommon<C>) -> Self {
+impl Client {
+    pub(self) fn new(inner: ClientInnerCommon<Send>) -> Self {
         Self { inner }
     }
 

proxy/src/serverless/local_conn_pool.rs

@@ -53,11 +53,11 @@ pub(crate) struct ClientDataLocal {
 }
 
 impl ClientDataLocal {
-    pub fn session(&mut self) -> &mut tokio::sync::watch::Sender<uuid::Uuid> {
-        &mut self.session
+    pub fn session(&self) -> &tokio::sync::watch::Sender<uuid::Uuid> {
+        &self.session
     }
 
-    pub fn cancel(&mut self) {
+    pub fn cancel(&self) {
         self.cancel.cancel();
     }
 }
@@ -99,7 +99,7 @@ impl<C: ClientInnerExt> LocalConnPool<C> {
             .map(|entry| entry.conn);
 
         // ok return cached connection if found and establish a new one otherwise
-        if let Some(mut client) = client {
+        if let Some(client) = client {
             if client.inner.is_closed() {
                 info!("local_pool: cached connection '{conn_info}' is closed, opening a new one");
                 return Ok(None);
@@ -120,11 +120,9 @@ impl<C: ClientInnerExt> LocalConnPool<C> {
                 ClientDataEnum::Local(data) => {
                     data.session().send(ctx.session_id())?;
                 }
-
                 ClientDataEnum::Remote(data) => {
                     data.session().send(ctx.session_id())?;
                 }
-                ClientDataEnum::Http(_) => (),
             }
 
             ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit);

safekeeper/Cargo.toml

@@ -55,6 +55,7 @@ tokio-util = { workspace = true }
 tracing.workspace = true
 url.workspace = true
 metrics.workspace = true
+pem.workspace = true
 postgres_backend.workspace = true
 postgres_ffi.workspace = true
 pq_proto.workspace = true

safekeeper/src/bin/safekeeper.rs

@@ -16,7 +16,6 @@ use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
 use metrics::set_build_info_metric;
 use remote_storage::RemoteStorageConfig;
-use reqwest::Certificate;
 use safekeeper::defaults::{
     DEFAULT_CONTROL_FILE_SAVE_INTERVAL, DEFAULT_EVICTION_MIN_RESIDENT, DEFAULT_HEARTBEAT_TIMEOUT,
     DEFAULT_HTTP_LISTEN_ADDR, DEFAULT_MAX_OFFLOADER_LAG_BYTES, DEFAULT_PARTIAL_BACKUP_CONCURRENCY,
@@ -373,7 +372,10 @@ async fn main() -> anyhow::Result<()> {
         Some(ssl_ca_file) => {
             tracing::info!("Using ssl root CA file: {ssl_ca_file:?}");
             let buf = tokio::fs::read(ssl_ca_file).await?;
-            Certificate::from_pem_bundle(&buf)?
+            pem::parse_many(&buf)?
+                .into_iter()
+                .filter(|pem| pem.tag() == "CERTIFICATE")
+                .collect()
         }
         None => Vec::new(),
     };

safekeeper/src/broker.rs

@@ -24,6 +24,15 @@ use crate::{GlobalTimelines, SafeKeeperConf};
 const RETRY_INTERVAL_MSEC: u64 = 1000;
 const PUSH_INTERVAL_MSEC: u64 = 1000;
 
+fn make_tls_config(conf: &SafeKeeperConf) -> storage_broker::ClientTlsConfig {
+    storage_broker::ClientTlsConfig::new().ca_certificates(
+        conf.ssl_ca_certs
+            .iter()
+            .map(pem::encode)
+            .map(storage_broker::Certificate::from_pem),
+    )
+}
+
 /// Push once in a while data about all active timelines to the broker.
 async fn push_loop(
     conf: Arc<SafeKeeperConf>,
@@ -37,8 +46,11 @@ async fn push_loop(
 
     let active_timelines_set = global_timelines.get_global_broker_active_set();
 
-    let mut client =
-        storage_broker::connect(conf.broker_endpoint.clone(), conf.broker_keepalive_interval)?;
+    let mut client = storage_broker::connect(
+        conf.broker_endpoint.clone(),
+        conf.broker_keepalive_interval,
+        make_tls_config(&conf),
+    )?;
     let push_interval = Duration::from_millis(PUSH_INTERVAL_MSEC);
 
     let outbound = async_stream::stream! {
@@ -81,8 +93,11 @@ async fn pull_loop(
     global_timelines: Arc<GlobalTimelines>,
     stats: Arc<BrokerStats>,
 ) -> Result<()> {
-    let mut client =
-        storage_broker::connect(conf.broker_endpoint.clone(), conf.broker_keepalive_interval)?;
+    let mut client = storage_broker::connect(
+        conf.broker_endpoint.clone(),
+        conf.broker_keepalive_interval,
+        make_tls_config(&conf),
+    )?;
 
     // TODO: subscribe only to local timelines instead of all
     let request = SubscribeSafekeeperInfoRequest {
@@ -134,8 +149,11 @@ async fn discover_loop(
     global_timelines: Arc<GlobalTimelines>,
     stats: Arc<BrokerStats>,
 ) -> Result<()> {
-    let mut client =
-        storage_broker::connect(conf.broker_endpoint.clone(), conf.broker_keepalive_interval)?;
+    let mut client = storage_broker::connect(
+        conf.broker_endpoint.clone(),
+        conf.broker_keepalive_interval,
+        make_tls_config(&conf),
+    )?;
 
     let request = SubscribeByFilterRequest {
         types: vec![TypeSubscription {

safekeeper/src/http/routes.rs

@@ -14,6 +14,7 @@ use http_utils::json::{json_request, json_response};
 use http_utils::request::{ensure_no_body, parse_query_param, parse_request_param};
 use http_utils::{RequestExt, RouterBuilder};
 use hyper::{Body, Request, Response, StatusCode};
+use pem::Pem;
 use postgres_ffi::WAL_SEGMENT_SIZE;
 use safekeeper_api::models::{
     AcceptorStateStatus, PullTimelineRequest, SafekeeperStatus, SkTimelineInfo, TenantDeleteResult,
@@ -230,14 +231,20 @@ async fn timeline_pull_handler(mut request: Request<Body>) -> Result<Response<Bo
     let conf = get_conf(&request);
     let global_timelines = get_global_timelines(&request);
 
-    let resp = pull_timeline::handle_request(
-        data,
-        conf.sk_auth_token.clone(),
-        conf.ssl_ca_certs.clone(),
-        global_timelines,
-    )
-    .await
-    .map_err(ApiError::InternalServerError)?;
+    let ca_certs = conf
+        .ssl_ca_certs
+        .iter()
+        .map(Pem::contents)
+        .map(reqwest::Certificate::from_der)
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(|e| {
+            ApiError::InternalServerError(anyhow::anyhow!("failed to parse CA certs: {e}"))
+        })?;
+
+    let resp =
+        pull_timeline::handle_request(data, conf.sk_auth_token.clone(), ca_certs, global_timelines)
+            .await
+            .map_err(ApiError::InternalServerError)?;
     json_response(StatusCode::OK, resp)
 }
 

safekeeper/src/lib.rs

@@ -6,8 +6,8 @@ use std::time::Duration;
 
 use camino::Utf8PathBuf;
 use once_cell::sync::Lazy;
+use pem::Pem;
 use remote_storage::RemoteStorageConfig;
-use reqwest::Certificate;
 use storage_broker::Uri;
 use tokio::runtime::Runtime;
 use utils::auth::SwappableJwtAuth;
@@ -120,7 +120,7 @@ pub struct SafeKeeperConf {
     pub ssl_key_file: Utf8PathBuf,
     pub ssl_cert_file: Utf8PathBuf,
     pub ssl_cert_reload_period: Duration,
-    pub ssl_ca_certs: Vec<Certificate>,
+    pub ssl_ca_certs: Vec<Pem>,
     pub use_https_safekeeper_api: bool,
 }
 

safekeeper/src/recovery.rs

@@ -8,6 +8,7 @@ use std::time::SystemTime;
 use anyhow::{Context, bail};
 use futures::StreamExt;
 use postgres_protocol::message::backend::ReplicationMessage;
+use reqwest::Certificate;
 use safekeeper_api::Term;
 use safekeeper_api::membership::INVALID_GENERATION;
 use safekeeper_api::models::{PeerInfo, TimelineStatus};
@@ -241,7 +242,7 @@ async fn recover(
 
     let mut client = reqwest::Client::builder();
     for cert in &conf.ssl_ca_certs {
-        client = client.add_root_certificate(cert.clone());
+        client = client.add_root_certificate(Certificate::from_der(cert.contents())?);
     }
     let client = client
         .build()

storage_broker/benches/rps.rs

@@ -87,7 +87,12 @@ fn tli_from_u64(i: u64) -> Vec<u8> {
 async fn subscribe(client: Option<BrokerClientChannel>, counter: Arc<AtomicU64>, i: u64) {
     let mut client = match client {
         Some(c) => c,
-        None => storage_broker::connect(DEFAULT_ENDPOINT, Duration::from_secs(5)).unwrap(),
+        None => storage_broker::connect(
+            DEFAULT_ENDPOINT,
+            Duration::from_secs(5),
+            storage_broker::ClientTlsConfig::new(),
+        )
+        .unwrap(),
     };
 
     let ttid = ProtoTenantTimelineId {
@@ -119,7 +124,12 @@ async fn subscribe(client: Option<BrokerClientChannel>, counter: Arc<AtomicU64>,
 async fn publish(client: Option<BrokerClientChannel>, n_keys: u64) {
     let mut client = match client {
         Some(c) => c,
-        None => storage_broker::connect(DEFAULT_ENDPOINT, Duration::from_secs(5)).unwrap(),
+        None => storage_broker::connect(
+            DEFAULT_ENDPOINT,
+            Duration::from_secs(5),
+            storage_broker::ClientTlsConfig::new(),
+        )
+        .unwrap(),
     };
     let mut counter: u64 = 0;
 
@@ -164,7 +174,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     }
     let h = tokio::spawn(progress_reporter(counters.clone()));
 
-    let c = storage_broker::connect(DEFAULT_ENDPOINT, Duration::from_secs(5)).unwrap();
+    let c = storage_broker::connect(
+        DEFAULT_ENDPOINT,
+        Duration::from_secs(5),
+        storage_broker::ClientTlsConfig::new(),
+    )
+    .unwrap();
 
     for i in 0..args.num_subs {
         let c = Some(c.clone());

storage_broker/src/lib.rs

@@ -4,7 +4,7 @@ use proto::TenantTimelineId as ProtoTenantTimelineId;
 use proto::broker_service_client::BrokerServiceClient;
 use tonic::Status;
 use tonic::codegen::StdError;
-use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
+use tonic::transport::{Channel, Endpoint};
 use utils::id::{TenantId, TenantTimelineId, TimelineId};
 
 // Code generated by protobuf.
@@ -20,6 +20,7 @@ pub mod metrics;
 
 // Re-exports to avoid direct tonic dependency in user crates.
 pub use hyper::Uri;
+pub use tonic::transport::{Certificate, ClientTlsConfig};
 pub use tonic::{Code, Request, Streaming};
 
 pub const DEFAULT_LISTEN_ADDR: &str = "127.0.0.1:50051";
@@ -38,7 +39,11 @@ pub type BrokerClientChannel = BrokerServiceClient<Channel>;
 //
 // NB: this function is not async, but still must be run on a tokio runtime thread
 // because that's a requirement of tonic_endpoint.connect_lazy()'s Channel::new call.
-pub fn connect<U>(endpoint: U, keepalive_interval: Duration) -> anyhow::Result<BrokerClientChannel>
+pub fn connect<U>(
+    endpoint: U,
+    keepalive_interval: Duration,
+    tls_config: ClientTlsConfig,
+) -> anyhow::Result<BrokerClientChannel>
 where
     U: std::convert::TryInto<Uri>,
     U::Error: std::error::Error + Send + Sync + 'static,
@@ -54,8 +59,7 @@ where
         rustls::crypto::ring::default_provider()
             .install_default()
             .ok();
-        let tls = ClientTlsConfig::new();
-        tonic_endpoint = tonic_endpoint.tls_config(tls)?;
+        tonic_endpoint = tonic_endpoint.tls_config(tls_config)?;
     }
     tonic_endpoint = tonic_endpoint
         .http2_keep_alive_interval(keepalive_interval)

storage_controller/src/persistence.rs

@@ -126,6 +126,7 @@ pub(crate) enum DatabaseOperation {
     InsertTimelineReconcile,
     RemoveTimelineReconcile,
     ListTimelineReconcile,
+    ListTimelineReconcileStartup,
 }
 
 #[must_use]
@@ -1521,23 +1522,41 @@ impl Persistence {
         .await
     }
 
-    /// Load pending operations from db.
-    pub(crate) async fn list_pending_ops(
+    /// Load pending operations from db, joined together with timeline data.
+    pub(crate) async fn list_pending_ops_with_timelines(
         &self,
-    ) -> DatabaseResult<Vec<TimelinePendingOpPersistence>> {
+    ) -> DatabaseResult<Vec<(TimelinePendingOpPersistence, Option<TimelinePersistence>)>> {
         use crate::schema::safekeeper_timeline_pending_ops::dsl;
+        use crate::schema::timelines;
 
         let timeline_from_db = self
-            .with_measured_conn(DatabaseOperation::ListTimelineReconcile, move |conn| {
-                Box::pin(async move {
-                    let from_db: Vec<TimelinePendingOpPersistence> =
-                        dsl::safekeeper_timeline_pending_ops.load(conn).await?;
-                    Ok(from_db)
-                })
-            })
+            .with_measured_conn(
+                DatabaseOperation::ListTimelineReconcileStartup,
+                move |conn| {
+                    Box::pin(async move {
+                        let from_db: Vec<(TimelinePendingOpPersistence, Option<TimelineFromDb>)> =
+                            dsl::safekeeper_timeline_pending_ops
+                                .left_join(
+                                    timelines::table.on(timelines::tenant_id
+                                        .eq(dsl::tenant_id)
+                                        .and(timelines::timeline_id.eq(dsl::timeline_id))),
+                                )
+                                .select((
+                                    TimelinePendingOpPersistence::as_select(),
+                                    Option::<TimelineFromDb>::as_select(),
+                                ))
+                                .load(conn)
+                                .await?;
+                        Ok(from_db)
+                    })
+                },
+            )
             .await?;
 
-        Ok(timeline_from_db)
+        Ok(timeline_from_db
+            .into_iter()
+            .map(|(op, tl_opt)| (op, tl_opt.map(|tl_opt| tl_opt.into_persistence())))
+            .collect())
     }
     /// List pending operations for a given timeline (including tenant-global ones)
     pub(crate) async fn list_pending_ops_for_timeline(
@@ -1580,7 +1599,7 @@ impl Persistence {
 
         let tenant_id = &tenant_id;
         let timeline_id = &timeline_id;
-        self.with_measured_conn(DatabaseOperation::ListTimelineReconcile, move |conn| {
+        self.with_measured_conn(DatabaseOperation::RemoveTimelineReconcile, move |conn| {
             let timeline_id_str = timeline_id.map(|tid| tid.to_string()).unwrap_or_default();
             Box::pin(async move {
                 diesel::delete(dsl::safekeeper_timeline_pending_ops)

storage_controller/src/service.rs

@@ -824,9 +824,13 @@ impl Service {
             let mut locked = self.inner.write().unwrap();
             locked.become_leader();
 
+            for (sk_id, _sk) in locked.safekeepers.clone().iter() {
+                locked.safekeeper_reconcilers.start_reconciler(*sk_id, self);
+            }
+
             locked
                 .safekeeper_reconcilers
-                .schedule_request_vec(self, sk_schedule_requests);
+                .schedule_request_vec(sk_schedule_requests);
         }
 
         // TODO: if any tenant's intent now differs from its loaded generation_pageserver, we should clear that

storage_controller/src/service/safekeeper_reconciler.rs

@@ -30,31 +30,35 @@ impl SafekeeperReconcilers {
             reconcilers: HashMap::new(),
         }
     }
-    pub(crate) fn schedule_request_vec(
-        &mut self,
-        service: &Arc<Service>,
-        reqs: Vec<ScheduleRequest>,
-    ) {
+    /// Adds a safekeeper-specific reconciler.
+    /// Can be called multiple times, but it needs to be called at least once
+    /// for every new safekeeper added.
+    pub(crate) fn start_reconciler(&mut self, node_id: NodeId, service: &Arc<Service>) {
+        self.reconcilers.entry(node_id).or_insert_with(|| {
+            SafekeeperReconciler::spawn(self.cancel.child_token(), service.clone())
+        });
+    }
+    /// Stop a safekeeper-specific reconciler.
+    /// Stops the reconciler, cancelling all ongoing tasks.
+    pub(crate) fn stop_reconciler(&mut self, node_id: NodeId) {
+        if let Some(handle) = self.reconcilers.remove(&node_id) {
+            handle.cancel.cancel();
+        }
+    }
+    pub(crate) fn schedule_request_vec(&self, reqs: Vec<ScheduleRequest>) {
         tracing::info!(
             "Scheduling {} pending safekeeper ops loaded from db",
             reqs.len()
         );
         for req in reqs {
-            self.schedule_request(service, req);
+            self.schedule_request(req);
         }
     }
-    pub(crate) fn schedule_request(&mut self, service: &Arc<Service>, req: ScheduleRequest) {
+    pub(crate) fn schedule_request(&self, req: ScheduleRequest) {
         let node_id = req.safekeeper.get_id();
-        let reconciler_handle = self.reconcilers.entry(node_id).or_insert_with(|| {
-            SafekeeperReconciler::spawn(self.cancel.child_token(), service.clone())
-        });
+        let reconciler_handle = self.reconcilers.get(&node_id).unwrap();
         reconciler_handle.schedule_reconcile(req);
     }
-    pub(crate) fn cancel_safekeeper(&mut self, node_id: NodeId) {
-        if let Some(handle) = self.reconcilers.remove(&node_id) {
-            handle.cancel.cancel();
-        }
-    }
     /// Cancel ongoing reconciles for the given timeline
     ///
     /// Specifying `None` here only removes reconciles for the tenant-global reconciliation,
@@ -78,9 +82,12 @@ pub(crate) async fn load_schedule_requests(
     service: &Arc<Service>,
     safekeepers: &HashMap<NodeId, Safekeeper>,
 ) -> anyhow::Result<Vec<ScheduleRequest>> {
-    let pending_ops = service.persistence.list_pending_ops().await?;
-    let mut res = Vec::with_capacity(pending_ops.len());
-    for op_persist in pending_ops {
+    let pending_ops_timelines = service
+        .persistence
+        .list_pending_ops_with_timelines()
+        .await?;
+    let mut res = Vec::with_capacity(pending_ops_timelines.len());
+    for (op_persist, timeline_persist) in pending_ops_timelines {
         let node_id = NodeId(op_persist.sk_id as u64);
         let Some(sk) = safekeepers.get(&node_id) else {
             // This shouldn't happen, at least the safekeeper should exist as decomissioned.
@@ -102,16 +109,12 @@ pub(crate) async fn load_schedule_requests(
             SafekeeperTimelineOpKind::Delete => Vec::new(),
             SafekeeperTimelineOpKind::Exclude => Vec::new(),
             SafekeeperTimelineOpKind::Pull => {
-                // TODO this code is super hacky, it doesn't take migrations into account
-                let Some(timeline_id) = timeline_id else {
+                if timeline_id.is_none() {
+                    // We only do this extra check (outside of timeline_persist check) to give better error msgs
                     anyhow::bail!(
                         "timeline_id is empty for `pull` schedule request for {tenant_id}"
                     );
                 };
-                let timeline_persist = service
-                    .persistence
-                    .get_timeline(tenant_id, timeline_id)
-                    .await?;
                 let Some(timeline_persist) = timeline_persist else {
                     // This shouldn't happen, the timeline should still exist
                     tracing::warn!(
@@ -163,6 +166,7 @@ pub(crate) struct ScheduleRequest {
     pub(crate) kind: SafekeeperTimelineOpKind,
 }
 
+/// Handle to per safekeeper reconciler.
 struct ReconcilerHandle {
     tx: UnboundedSender<(ScheduleRequest, CancellationToken)>,
     ongoing_tokens: Arc<ClashMap<(TenantId, Option<TimelineId>), CancellationToken>>,
@@ -170,7 +174,10 @@ struct ReconcilerHandle {
 }
 
 impl ReconcilerHandle {
-    /// Obtain a new token slot, cancelling any existing reconciliations for that timeline
+    /// Obtain a new token slot, cancelling any existing reconciliations for
+    /// that timeline. It is not useful to have >1 operation per <tenant_id,
+    /// timeline_id, safekeeper>, hence scheduling op cancels current one if it
+    /// exists.
     fn new_token_slot(
         &self,
         tenant_id: TenantId,
@@ -305,15 +312,16 @@ impl SafekeeperReconciler {
             SafekeeperTimelineOpKind::Delete => {
                 let tenant_id = req.tenant_id;
                 if let Some(timeline_id) = req.timeline_id {
-                    let deleted = self.reconcile_inner(
-                        req,
-                        async |client| client.delete_timeline(tenant_id, timeline_id).await,
-                        |_resp| {
-                            tracing::info!(%tenant_id, %timeline_id, "deleted timeline from {req_host}");
-                        },
-                        req_cancel,
-                    )
-                    .await;
+                    let deleted = self
+                        .reconcile_inner(
+                            req,
+                            async |client| client.delete_timeline(tenant_id, timeline_id).await,
+                            |_resp| {
+                                tracing::info!("deleted timeline from {req_host}");
+                            },
+                            req_cancel,
+                        )
+                        .await;
                     if deleted {
                         self.delete_timeline_from_db(tenant_id, timeline_id).await;
                     }
@@ -344,12 +352,13 @@ impl SafekeeperReconciler {
         {
             Ok(list) => {
                 if !list.is_empty() {
-                    tracing::info!(%tenant_id, %timeline_id, "not deleting timeline from db as there is {} open reconciles", list.len());
+                    // duplicate the timeline_id here because it might be None in the reconcile context
+                    tracing::info!(%timeline_id, "not deleting timeline from db as there is {} open reconciles", list.len());
                     return;
                 }
             }
             Err(e) => {
-                tracing::warn!(%tenant_id, %timeline_id, "couldn't query pending ops: {e}");
+                tracing::warn!(%timeline_id, "couldn't query pending ops: {e}");
                 return;
             }
         }

storage_controller/src/service/safekeeper_service.rs

@@ -46,6 +46,7 @@ impl Service {
             .map(SecretString::from);
         let mut joinset = JoinSet::new();
 
+        // Prepare membership::Configuration from choosen safekeepers.
         let safekeepers = {
             let locked = self.inner.read().unwrap();
             locked.safekeepers.clone()
@@ -205,7 +206,7 @@ impl Service {
             tenant_id: tenant_id.to_string(),
             timeline_id: timeline_id.to_string(),
             start_lsn: start_lsn.into(),
-            generation: 0,
+            generation: 1,
             sk_set: sks_persistence.clone(),
             new_sk_set: None,
             cplane_notified_generation: 0,
@@ -254,7 +255,7 @@ impl Service {
             self.persistence.insert_pending_op(pending_op).await?;
         }
         if !remaining.is_empty() {
-            let mut locked = self.inner.write().unwrap();
+            let locked = self.inner.read().unwrap();
             for remaining_id in remaining {
                 let Some(sk) = locked.safekeepers.get(&remaining_id) else {
                     return Err(ApiError::InternalServerError(anyhow::anyhow!(
@@ -290,7 +291,7 @@ impl Service {
                     generation: timeline_persist.generation as u32,
                     kind: crate::persistence::SafekeeperTimelineOpKind::Pull,
                 };
-                locked.safekeeper_reconcilers.schedule_request(self, req);
+                locked.safekeeper_reconcilers.schedule_request(req);
             }
         }
 
@@ -357,7 +358,7 @@ impl Service {
             let pending_op = TimelinePendingOpPersistence {
                 tenant_id: tenant_id.to_string(),
                 timeline_id: timeline_id.to_string(),
-                generation: tl.generation,
+                generation: i32::MAX,
                 op_kind: SafekeeperTimelineOpKind::Delete,
                 sk_id: *sk_id,
             };
@@ -365,7 +366,7 @@ impl Service {
             self.persistence.insert_pending_op(pending_op).await?;
         }
         {
-            let mut locked = self.inner.write().unwrap();
+            let locked = self.inner.read().unwrap();
             for sk_id in all_sks {
                 let sk_id = NodeId(*sk_id as u64);
                 let Some(sk) = locked.safekeepers.get(&sk_id) else {
@@ -383,7 +384,7 @@ impl Service {
                     generation: tl.generation as u32,
                     kind: SafekeeperTimelineOpKind::Delete,
                 };
-                locked.safekeeper_reconcilers.schedule_request(self, req);
+                locked.safekeeper_reconcilers.schedule_request(req);
             }
         }
         Ok(())
@@ -482,7 +483,7 @@ impl Service {
                 tenant_id,
                 timeline_id: None,
             };
-            locked.safekeeper_reconcilers.schedule_request(self, req);
+            locked.safekeeper_reconcilers.schedule_request(req);
         }
         Ok(())
     }
@@ -579,7 +580,7 @@ impl Service {
     }
 
     pub(crate) async fn upsert_safekeeper(
-        &self,
+        self: &Arc<Service>,
         record: crate::persistence::SafekeeperUpsert,
     ) -> Result<(), ApiError> {
         let node_id = NodeId(record.id as u64);
@@ -618,6 +619,9 @@ impl Service {
                     );
                 }
             }
+            locked
+                .safekeeper_reconcilers
+                .start_reconciler(node_id, self);
             locked.safekeepers = Arc::new(safekeepers);
             metrics::METRICS_REGISTRY
                 .metrics_group
@@ -638,7 +642,7 @@ impl Service {
     }
 
     pub(crate) async fn set_safekeeper_scheduling_policy(
-        &self,
+        self: &Arc<Service>,
         id: i64,
         scheduling_policy: SkSchedulingPolicy,
     ) -> Result<(), DatabaseError> {
@@ -656,9 +660,13 @@ impl Service {
             sk.set_scheduling_policy(scheduling_policy);
 
             match scheduling_policy {
-                SkSchedulingPolicy::Active => (),
+                SkSchedulingPolicy::Active => {
+                    locked
+                        .safekeeper_reconcilers
+                        .start_reconciler(node_id, self);
+                }
                 SkSchedulingPolicy::Decomissioned | SkSchedulingPolicy::Pause => {
-                    locked.safekeeper_reconcilers.cancel_safekeeper(node_id);
+                    locked.safekeeper_reconcilers.stop_reconciler(node_id);
                 }
             }
 

test_runner/fixtures/neon_api.py

@@ -22,19 +22,62 @@ def connection_parameters_to_env(params: dict[str, str]) -> dict[str, str]:
     }
 
 
+# Some API calls not yet implemented.
+# You may want to copy not-yet-implemented methods from the PR https://github.com/neondatabase/neon/pull/11305
 class NeonAPI:
     def __init__(self, neon_api_key: str, neon_api_base_url: str):
         self.__neon_api_key = neon_api_key
         self.__neon_api_base_url = neon_api_base_url.strip("/")
+        self.retry_if_possible = False
+        self.attempts = 10
+        self.sleep_before_retry = 1
+        self.retries524 = 0
+        self.retries4xx = 0
 
     def __request(self, method: str | bytes, endpoint: str, **kwargs: Any) -> requests.Response:
-        if "headers" not in kwargs:
-            kwargs["headers"] = {}
+        kwargs["headers"] = kwargs.get("headers", {})
         kwargs["headers"]["Authorization"] = f"Bearer {self.__neon_api_key}"
 
-        resp = requests.request(method, f"{self.__neon_api_base_url}{endpoint}", **kwargs)
-        log.debug("%s %s returned a %d: %s", method, endpoint, resp.status_code, resp.text)
-        resp.raise_for_status()
+        for attempt in range(self.attempts):
+            retry = False
+            resp = requests.request(method, f"{self.__neon_api_base_url}{endpoint}", **kwargs)
+            if resp.status_code >= 400:
+                log.error(
+                    "%s %s returned a %d: %s",
+                    method,
+                    endpoint,
+                    resp.status_code,
+                    resp.text if resp.status_code != 524 else "CloudFlare error page",
+                )
+            else:
+                log.debug("%s %s returned a %d: %s", method, endpoint, resp.status_code, resp.text)
+            if not self.retry_if_possible:
+                resp.raise_for_status()
+                break
+            elif resp.status_code >= 400:
+                if resp.status_code == 422:
+                    if resp.json()["message"] == "branch not ready yet":
+                        retry = True
+                        self.retries4xx += 1
+                elif resp.status_code == 423 and resp.json()["message"] in {
+                    "endpoint is in some transitive state, could not suspend",
+                    "project already has running conflicting operations, scheduling of new ones is prohibited",
+                }:
+                    retry = True
+                    self.retries4xx += 1
+                elif resp.status_code == 524:
+                    log.info("The request was timed out, trying to get operations")
+                    retry = True
+                    self.retries524 += 1
+            if retry:
+                log.info("Retrying, attempt %s/%s", attempt + 1, self.attempts)
+                time.sleep(self.sleep_before_retry)
+                continue
+            else:
+                resp.raise_for_status()
+            break
+        else:
+            raise RuntimeError("Max retry count is reached")
 
         return resp
 
@@ -101,6 +144,96 @@ class NeonAPI:
 
         return cast("dict[str, Any]", resp.json())
 
+    def create_branch(
+        self,
+        project_id: str,
+        branch_name: str | None = None,
+        parent_id: str | None = None,
+        parent_lsn: str | None = None,
+        parent_timestamp: str | None = None,
+        protected: bool | None = None,
+        archived: bool | None = None,
+        init_source: str | None = None,
+        add_endpoint=True,
+    ) -> dict[str, Any]:
+        data: dict[str, Any] = {}
+        if add_endpoint:
+            data["endpoints"] = [{"type": "read_write"}]
+        data["branch"] = {}
+        if parent_id:
+            data["branch"]["parent_id"] = parent_id
+        if branch_name:
+            data["branch"]["name"] = branch_name
+        if parent_lsn is not None:
+            data["branch"]["parent_lsn"] = parent_lsn
+        if parent_timestamp is not None:
+            data["branch"]["parent_timestamp"] = parent_timestamp
+        if protected is not None:
+            data["branch"]["protected"] = protected
+        if init_source is not None:
+            data["branch"]["init_source"] = init_source
+        if archived is not None:
+            data["branch"]["archived"] = archived
+        if not data["branch"]:
+            data.pop("branch")
+        resp = self.__request(
+            "POST",
+            f"/projects/{project_id}/branches",
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/json",
+            },
+            json=data,
+        )
+        return cast("dict[str, Any]", resp.json())
+
+    def get_branch_details(self, project_id: str, branch_id: str) -> dict[str, Any]:
+        resp = self.__request(
+            "GET",
+            f"/projects/{project_id}/branches/{branch_id}",
+            headers={
+                "Accept": "application/json",
+            },
+        )
+        return cast("dict[str, Any]", resp.json())
+
+    def delete_branch(self, project_id: str, branch_id: str) -> dict[str, Any]:
+        resp = self.__request(
+            "DELETE",
+            f"/projects/{project_id}/branches/{branch_id}",
+            headers={
+                "Accept": "application/json",
+            },
+        )
+        return cast("dict[str, Any]", resp.json())
+
+    def restore_branch(
+        self,
+        project_id: str,
+        branch_id: str,
+        source_branch_id: str,
+        source_lsn: str | None,
+        source_timestamp: str | None,
+        preserve_under_name: str | None,
+    ):
+        data = {"source_branch_id": source_branch_id}
+        if source_lsn:
+            data["source_lsn"] = source_lsn
+        if source_timestamp:
+            data["source_timestamp"] = source_timestamp
+        if preserve_under_name:
+            data["preserve_under_name"] = preserve_under_name
+        log.info("Data: %s", data)
+        resp = self.__request(
+            "POST",
+            f"/projects/{project_id}/branches/{branch_id}/restore",
+            headers={
+                "Accept": "application/json",
+            },
+            json=data,
+        )
+        return cast("dict[str, Any]", resp.json())
+
     def start_endpoint(
         self,
         project_id: str,
@@ -176,6 +309,10 @@ class NeonAPI:
 
         return cast("dict[str, Any]", resp.json())
 
+    def delete_endpoint(self, project_id: str, endpoint_id: str) -> dict[str, Any]:
+        resp = self.__request("DELETE", f"/projects/{project_id}/endpoints/{endpoint_id}")
+        return cast("dict[str,Any]", resp.json())
+
     def get_connection_uri(
         self,
         project_id: str,

test_runner/fixtures/neon_fixtures.py

@@ -3185,6 +3185,7 @@ class PgBin:
         command: list[str],
         env: Env | None = None,
         cwd: str | Path | None = None,
+        stderr_pipe: Any | None = None,
     ) -> subprocess.Popen[Any]:
         """
         Run one of the postgres binaries, not waiting for it to finish
@@ -3202,7 +3203,9 @@ class PgBin:
         log.info(f"Running command '{' '.join(command)}'")
         env = self._build_env(env)
         self._log_env(env)
-        return subprocess.Popen(command, env=env, cwd=cwd, stdout=subprocess.PIPE, text=True)
+        return subprocess.Popen(
+            command, env=env, cwd=cwd, stdout=subprocess.PIPE, stderr=stderr_pipe, text=True
+        )
 
     def run(
         self,

test_runner/random_ops/README.md

@@ -0,0 +1,93 @@
+# Random Operations Test for Neon Stability
+
+## Problem Statement
+
+Neon needs robust testing of Neon's stability to ensure reliability for users. The random operations test addresses this by continuously exercising the API with unpredictable sequences of operations, helping to identify edge cases and potential issues that might not be caught by deterministic tests.
+
+### Key Components
+
+#### 1. Class Structure
+
+The test implements three main classes to model the Neon architecture:
+
+- **NeonProject**: Represents a Neon project and manages the lifecycle of branches and endpoints
+- **NeonBranch**: Represents a branch within a project, with methods for creating child branches, endpoints, and performing point-in-time restores
+- **NeonEndpoint**: Represents an endpoint (connection point) for a branch, with methods for managing benchmarks
+
+#### 2. Operations Tested
+
+The test randomly performs the following operations with weighted probabilities:
+
+- **Creating branches** 
+- **Deleting branches**
+- **Adding read-only endpoints**
+- **Deleting read-only endpoints**
+- **Restoring branches to random points in time**
+
+#### 3. Load Generation
+
+Each branch and endpoint is loaded with `pgbench` to simulate real database workloads during testing. This ensures that the operations are performed against branches with actual data and ongoing transactions.
+
+#### 4. Error Handling
+
+The test includes robust error handling for various scenarios:
+- Branch limit exceeded
+- Connection timeouts
+- Control plane timeouts (HTTP 524 errors)
+- Benchmark failures
+
+#### 5. CI Integration
+
+The test is integrated into the CI pipeline via a GitHub workflow that runs daily, ensuring continuous validation of API stability.
+
+## How It Works
+
+1. The test creates a Neon project using the Public API
+2. It initializes the main branch with pgbench data
+3. It performs random operations according to the weighted probabilities
+4. During each operation, it checks that all running benchmarks are still operational
+5. The test cleans up by deleting the project at the end
+
+## Configuration
+
+The test can be configured with:
+- `RANDOM_SEED`: Set a specific random seed for reproducible test runs
+- `NEON_API_KEY`: API key for authentication
+- `NEON_API_BASE_URL`: Base URL for the API (defaults to staging environment)
+- `NUM_OPERATIONS`: The number of operations to be performed
+
+## Running the Test
+
+The test is designed to run in the CI environment but can also be executed locally:
+
+```bash
+NEON_API_KEY=your_api_key ./scripts/pytest test_runner/random_ops/test_random_ops.py -m remote_cluster
+```
+
+To run with a specific random seed for reproducibility:
+
+```bash
+RANDOM_SEED=12345 NEON_API_KEY=your_api_key ./scripts/pytest test_runner/random_ops/test_random_ops.py -m remote_cluster
+```
+
+To run with the custom number of operations:
+
+```bash
+NUM_OPERATIONS=500 NEON_API_KEY=your_api_key ./scripts/pytest test_runner/random_ops/test_random_ops.py -m remote_cluster
+```
+
+## Benefits
+
+This test provides several key benefits:
+1. **Comprehensive API testing**: Exercises multiple API endpoints in combination
+2. **Edge case discovery**: Random sequences may uncover issues not found in deterministic tests
+3. **Stability validation**: Continuous execution helps ensure long-term API reliability
+4. **Regression prevention**: Detects if new changes break existing API functionality
+
+## Future Improvements
+
+Potential enhancements to the test could include:
+1. Adding more API operations, e.g. `reset_to_parent`, `snapshot`, etc 
+2. Implementing more sophisticated load patterns
+3. Adding metrics collection to measure API performance
+4. Extending test duration for longer-term stability validation

test_runner/random_ops/test_random_ops.py

@@ -0,0 +1,463 @@
+"""
+Run the random API tests on the cloud instance of Neon
+"""
+
+from __future__ import annotations
+
+import os
+import random
+import subprocess
+import time
+from datetime import UTC, datetime, timedelta
+from typing import TYPE_CHECKING, Any
+
+import pytest
+from fixtures.log_helper import log
+from requests import HTTPError
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from fixtures.neon_api import NeonAPI
+    from fixtures.neon_fixtures import PgBin
+    from fixtures.pg_version import PgVersion
+
+
+class NeonEndpoint:
+    """
+    Neon Endpoint
+    Gets the output of the API call of an endpoint creation
+    """
+
+    def __init__(self, project: NeonProject, endpoint: dict[str, Any]):
+        self.project: NeonProject = project
+        self.id: str = endpoint["id"]
+        # The branch endpoint belongs to
+        self.branch: NeonBranch = project.branches[endpoint["branch_id"]]
+        self.type: str = endpoint["type"]
+        # add itself to the list of endpoints of the branch
+        self.branch.endpoints[self.id] = self
+        self.project.endpoints[self.id] = self
+        self.host: str = endpoint["host"]
+        self.benchmark: subprocess.Popen[Any] | None = None
+        # The connection environment is used when running benchmark
+        self.connect_env: dict[str, str] | None = None
+        if self.branch.connect_env:
+            self.connect_env = self.branch.connect_env.copy()
+            self.connect_env["PGHOST"] = self.host
+
+    def delete(self):
+        self.project.delete_endpoint(self.id)
+
+    def start_benchmark(self, clients=10):
+        return self.project.start_benchmark(self.id, clients=clients)
+
+    def check_benchmark(self):
+        self.project.check_benchmark(self.id)
+
+    def terminate_benchmark(self):
+        self.project.terminate_benchmark(self.id)
+
+
+class NeonBranch:
+    """
+    Neon Branch
+    Gets the output of the API call of the Neon Public API call of a branch creation as a first parameter
+    is_reset defines if the branch is a reset one i.e. created as a result of the reset API Call
+    """
+
+    def __init__(self, project, branch: dict[str, Any], is_reset=False):
+        self.id: str = branch["branch"]["id"]
+        self.desc = branch
+        self.project: NeonProject = project
+        self.neon_api: NeonAPI = project.neon_api
+        self.project_id: str = branch["branch"]["project_id"]
+        self.parent: NeonBranch | None = (
+            self.project.branches[branch["branch"]["parent_id"]]
+            if "parent_id" in branch["branch"]
+            else None
+        )
+        if is_reset:
+            self.project.reset_branches.add(self.id)
+        elif self.parent:
+            self.project.leaf_branches[self.id] = self
+        if self.parent is not None and self.parent.id in self.project.leaf_branches:
+            self.project.leaf_branches.pop(self.parent.id)
+        self.project.branches[self.id] = self
+        self.children: dict[str, NeonBranch] = {}
+        if self.parent is not None:
+            self.parent.children[self.id] = self
+        self.endpoints: dict[str, NeonEndpoint] = {}
+        self.connection_parameters: dict[str, str] | None = (
+            branch["connection_uris"][0]["connection_parameters"]
+            if "connection_uris" in branch
+            else None
+        )
+        self.benchmark: subprocess.Popen[Any] | None = None
+        self.updated_at: datetime = datetime.fromisoformat(branch["branch"]["updated_at"])
+        self.connect_env: dict[str, str] | None = None
+        if self.connection_parameters:
+            self.connect_env = {
+                "PGHOST": self.connection_parameters["host"],
+                "PGUSER": self.connection_parameters["role"],
+                "PGDATABASE": self.connection_parameters["database"],
+                "PGPASSWORD": self.connection_parameters["password"],
+                "PGSSLMODE": "require",
+            }
+
+    def __str__(self):
+        """
+        Prints the branch's name with all the predecessors
+        (r) means the branch is a reset one
+        """
+        return f"{self.id}{'(r)' if self.id in self.project.reset_branches else ''}, parent: {self.parent}"
+
+    def create_child_branch(self) -> NeonBranch | None:
+        return self.project.create_branch(self.id)
+
+    def create_ro_endpoint(self) -> NeonEndpoint:
+        return NeonEndpoint(
+            self.project,
+            self.neon_api.create_endpoint(self.project_id, self.id, "read_only", {})["endpoint"],
+        )
+
+    def delete(self) -> None:
+        self.project.delete_branch(self.id)
+
+    def start_benchmark(self, clients=10) -> subprocess.Popen[Any]:
+        return self.project.start_benchmark(self.id, clients=clients)
+
+    def check_benchmark(self) -> None:
+        self.project.check_benchmark(self.id)
+
+    def terminate_benchmark(self) -> None:
+        self.project.terminate_benchmark(self.id)
+
+    def restore_random_time(self) -> None:
+        """
+        Does PITR, i.e. calls the reset API call on the same branch to the random time in the past
+        """
+        min_time = self.updated_at + timedelta(seconds=1)
+        max_time = datetime.now(UTC) - timedelta(seconds=1)
+        target_time = (min_time + (max_time - min_time) * random.random()).replace(microsecond=0)
+        res = self.restore(
+            self.id,
+            source_timestamp=target_time.isoformat().replace("+00:00", "Z"),
+            preserve_under_name=self.project.gen_restore_name(),
+        )
+        if res is None:
+            return
+        self.updated_at = datetime.fromisoformat(res["branch"]["updated_at"])
+        parent_id: str = res["branch"]["parent_id"]
+        # Creates an object for the parent branch
+        # After the reset operation a new parent branch is created
+        parent = NeonBranch(
+            self.project, self.neon_api.get_branch_details(self.project_id, parent_id), True
+        )
+        self.project.branches[parent_id] = parent
+        self.parent = parent
+        parent.children[self.id] = self
+        self.project.wait()
+
+    def restore(
+        self,
+        source_branch_id: str,
+        source_lsn: str | None = None,
+        source_timestamp: str | None = None,
+        preserve_under_name: str | None = None,
+    ) -> dict[str, Any] | None:
+        endpoints = [ep for ep in self.endpoints.values() if ep.type == "read_only"]
+        # Terminate all the benchmarks running to prevent errors. Errors in benchmark during pgbench are expected
+        for ep in endpoints:
+            ep.terminate_benchmark()
+        self.terminate_benchmark()
+        try:
+            res: dict[str, Any] = self.neon_api.restore_branch(
+                self.project_id,
+                self.id,
+                source_branch_id,
+                source_lsn,
+                source_timestamp,
+                preserve_under_name,
+            )
+        except HTTPError as he:
+            if (
+                he.response.status_code == 422
+                and he.response.json()["code"] == "BRANCHES_LIMIT_EXCEEDED"
+            ):
+                log.info("Branch limit exceeded, skipping")
+                return None
+            else:
+                raise HTTPError(he) from he
+        self.project.wait()
+        self.start_benchmark()
+        for ep in endpoints:
+            ep.start_benchmark()
+        return res
+
+
+class NeonProject:
+    """
+    The project object
+    Calls the Public API to create a Neon Project
+    """
+
+    def __init__(self, neon_api: NeonAPI, pg_bin: PgBin, pg_version: PgVersion):
+        self.neon_api = neon_api
+        self.pg_bin = pg_bin
+        proj = self.neon_api.create_project(
+            pg_version, f"Automatic random API test {os.getenv('GITHUB_RUN_ID')}"
+        )
+        self.id: str = proj["project"]["id"]
+        self.name: str = proj["project"]["name"]
+        self.connection_uri: str = proj["connection_uris"][0]["connection_uri"]
+        self.connection_parameters: dict[str, str] = proj["connection_uris"][0][
+            "connection_parameters"
+        ]
+        self.pg_version: PgVersion = pg_version
+        # Leaf branches are the branches, which do not have children
+        self.leaf_branches: dict[str, NeonBranch] = {}
+        self.branches: dict[str, NeonBranch] = {}
+        self.reset_branches: set[str] = set()
+        self.main_branch: NeonBranch = NeonBranch(self, proj)
+        self.main_branch.connection_parameters = self.connection_parameters
+        self.endpoints: dict[str, NeonEndpoint] = {}
+        for endpoint in proj["endpoints"]:
+            NeonEndpoint(self, endpoint)
+        self.neon_api.wait_for_operation_to_finish(self.id)
+        self.benchmarks: dict[str, subprocess.Popen[Any]] = {}
+        self.restore_num: int = 0
+        self.restart_pgbench_on_console_errors: bool = False
+
+    def delete(self):
+        self.neon_api.delete_project(self.id)
+
+    def create_branch(self, parent_id: str | None = None) -> NeonBranch | None:
+        self.wait()
+        try:
+            branch_def = self.neon_api.create_branch(self.id, parent_id=parent_id)
+        except HTTPError as he:
+            if (
+                he.response.status_code == 422
+                and he.response.json()["code"] == "BRANCHES_LIMIT_EXCEEDED"
+            ):
+                log.info("Branch limit exceeded, skipping")
+                return None
+            else:
+                raise HTTPError(he) from he
+        new_branch = NeonBranch(self, branch_def)
+        self.wait()
+        return new_branch
+
+    def delete_branch(self, branch_id: str) -> None:
+        parent = self.branches[branch_id].parent
+        if not parent or branch_id == self.main_branch.id:
+            raise RuntimeError("Cannot delete the main branch")
+        if branch_id not in self.leaf_branches and branch_id not in self.reset_branches:
+            raise RuntimeError(f"The branch {branch_id}, probably, has ancestors")
+        if branch_id not in self.branches:
+            raise RuntimeError(f"The branch with id {branch_id} is not found")
+        endpoints_to_delete = [
+            ep for ep in self.branches[branch_id].endpoints.values() if ep.type == "read_only"
+        ]
+        for ep in endpoints_to_delete:
+            ep.delete()
+        if branch_id not in self.reset_branches:
+            self.terminate_benchmark(branch_id)
+        self.neon_api.delete_branch(self.id, branch_id)
+        if len(parent.children) == 1 and parent.id != self.main_branch.id:
+            self.leaf_branches[parent.id] = parent
+        parent.children.pop(branch_id)
+        if branch_id in self.leaf_branches:
+            self.leaf_branches.pop(branch_id)
+        else:
+            self.reset_branches.remove(branch_id)
+        self.branches.pop(branch_id)
+        self.wait()
+        if parent.id in self.reset_branches:
+            parent.delete()
+
+    def delete_endpoint(self, endpoint_id: str) -> None:
+        self.terminate_benchmark(endpoint_id)
+        self.neon_api.delete_endpoint(self.id, endpoint_id)
+        self.endpoints[endpoint_id].branch.endpoints.pop(endpoint_id)
+        self.endpoints.pop(endpoint_id)
+        self.wait()
+
+    def start_benchmark(self, target: str, clients: int = 10) -> subprocess.Popen[Any]:
+        if target in self.benchmarks:
+            raise RuntimeError(f"Benchmark was already started for {target}")
+        is_endpoint = target.startswith("ep")
+        read_only = is_endpoint and self.endpoints[target].type == "read_only"
+        cmd = ["pgbench", f"-c{clients}", "-T10800", "-Mprepared"]
+        if read_only:
+            cmd.extend(["-S", "-n"])
+        target_object = self.endpoints[target] if is_endpoint else self.branches[target]
+        if target_object.connect_env is None:
+            raise RuntimeError(f"The connection environment is not defined for {target}")
+        log.info(
+            "running pgbench on %s, cmd: %s, host: %s",
+            target,
+            cmd,
+            target_object.connect_env["PGHOST"],
+        )
+        pgbench = self.pg_bin.run_nonblocking(
+            cmd, env=target_object.connect_env, stderr_pipe=subprocess.PIPE
+        )
+        self.benchmarks[target] = pgbench
+        target_object.benchmark = pgbench
+        time.sleep(2)
+        return pgbench
+
+    def check_all_benchmarks(self) -> None:
+        for target in tuple(self.benchmarks.keys()):
+            self.check_benchmark(target)
+
+    def check_benchmark(self, target) -> None:
+        rc = self.benchmarks[target].poll()
+        if rc is not None:
+            _, err = self.benchmarks[target].communicate()
+            log.error("STDERR: %s", err)
+            # if the benchmark failed due to irresponsible Control plane,
+            # just restart it
+            if self.restart_pgbench_on_console_errors and (
+                "ERROR:  Couldn't connect to compute node" in err
+                or "ERROR:  Console request failed" in err
+            ):
+                log.info("Restarting benchmark for %s", target)
+                self.benchmarks.pop(target)
+                self.start_benchmark(target)
+                return
+            raise RuntimeError(f"The benchmark for {target} ended with code {rc}")
+
+    def terminate_benchmark(self, target):
+        log.info("Terminating the benchmark %s", target)
+        target_endpoint = target.startswith("ep")
+        self.check_benchmark(target)
+        self.benchmarks[target].terminate()
+        self.benchmarks.pop(target)
+        if target_endpoint:
+            self.endpoints[target].benchmark = None
+        else:
+            self.branches[target].benchmark = None
+
+    def wait(self):
+        """
+        Wait for all the operations to be finished
+        """
+        return self.neon_api.wait_for_operation_to_finish(self.id)
+
+    def gen_restore_name(self):
+        self.restore_num += 1
+        return f"restore{self.restore_num}"
+
+
+@pytest.fixture()
+def setup_class(
+    pg_version: PgVersion,
+    pg_bin: PgBin,
+    neon_api: NeonAPI,
+):
+    neon_api.retry_if_possible = True
+    project = NeonProject(neon_api, pg_bin, pg_version)
+    log.info("Created a project with id %s, name %s", project.id, project.name)
+    yield pg_bin, project
+    log.info("Retried 524 errors: %s", neon_api.retries524)
+    log.info("Retried 4xx errors: %s", neon_api.retries4xx)
+    if neon_api.retries524 > 0:
+        print(f"::warning::Retried on 524 error {neon_api.retries524} times")
+    if neon_api.retries4xx > 0:
+        print(f"::warning::Retried on 4xx error {neon_api.retries4xx} times")
+    log.info("Removing the project")
+    project.delete()
+
+
+def do_action(project: NeonProject, action: str) -> None:
+    """
+    Runs the action
+    """
+    log.info("Action: %s", action)
+    if action == "new_branch":
+        log.info("Trying to create a new branch")
+        parent = project.branches[
+            random.choice(list(set(project.branches.keys()) - project.reset_branches))
+        ]
+        log.info("Parent: %s", parent)
+        child = parent.create_child_branch()
+        if child is None:
+            return
+        log.info("Created branch %s", child)
+        child.start_benchmark()
+    elif action == "delete_branch":
+        if project.leaf_branches:
+            target = random.choice(list(project.leaf_branches.values()))
+            log.info("Trying to delete branch %s", target)
+            target.delete()
+        else:
+            log.info("Leaf branches not found, skipping")
+    elif action == "new_ro_endpoint":
+        ep = random.choice(
+            [br for br in project.branches.values() if br.id not in project.reset_branches]
+        ).create_ro_endpoint()
+        log.info("Created the RO endpoint with id %s branch: %s", ep.id, ep.branch.id)
+        ep.start_benchmark()
+    elif action == "delete_ro_endpoint":
+        ro_endpoints: list[NeonEndpoint] = [
+            endpoint for endpoint in project.endpoints.values() if endpoint.type == "read_only"
+        ]
+        if ro_endpoints:
+            target_ep: NeonEndpoint = random.choice(ro_endpoints)
+            target_ep.delete()
+            log.info("endpoint %s deleted", target_ep.id)
+        else:
+            log.info("no read_only endpoints present, skipping")
+    elif action == "restore_random_time":
+        if project.leaf_branches:
+            br: NeonBranch = random.choice(list(project.leaf_branches.values()))
+            log.info("Restore %s", br)
+            br.restore_random_time()
+        else:
+            log.info("No leaf branches found")
+    else:
+        raise ValueError(f"The action {action} is unknown")
+
+
+@pytest.mark.timeout(7200)
+@pytest.mark.remote_cluster
+def test_api_random(
+    setup_class,
+    pg_distrib_dir: Path,
+    test_output_dir: Path,
+):
+    """
+    Run the random API tests
+    """
+    if seed_env := os.getenv("RANDOM_SEED"):
+        seed = int(seed_env)
+    else:
+        seed = 0
+    if seed == 0:
+        seed = int(time.time())
+    log.info("Using random seed: %s", seed)
+    random.seed(seed)
+    pg_bin, project = setup_class
+    # Here we can assign weights
+    ACTIONS = (
+        ("new_branch", 1.5),
+        ("new_ro_endpoint", 1.4),
+        ("delete_ro_endpoint", 0.8),
+        ("delete_branch", 1.0),
+        ("restore_random_time", 1.2),
+    )
+    if num_ops_env := os.getenv("NUM_OPERATIONS"):
+        num_operations = int(num_ops_env)
+    else:
+        num_operations = 250
+    pg_bin.run(["pgbench", "-i", "-I", "dtGvp", "-s100"], env=project.main_branch.connect_env)
+    for _ in range(num_operations):
+        log.info("Starting action #%s", _ + 1)
+        do_action(
+            project, random.choices([a[0] for a in ACTIONS], weights=[w[1] for w in ACTIONS])[0]
+        )
+        project.check_all_benchmarks()
+    assert True

test_runner/regress/test_storage_controller.py

@@ -4114,13 +4114,29 @@ def test_storage_controller_location_conf_equivalence(neon_env_builder: NeonEnvB
     assert reconciles_after_restart == 0
 
 
+class RestartStorcon(Enum):
+    RESTART = "restart"
+    ONLINE = "online"
+
+
+class DeletionSubject(Enum):
+    TIMELINE = "timeline"
+    TENANT = "tenant"
+
+
 @run_only_on_default_postgres("PG version is not interesting here")
-@pytest.mark.parametrize("restart_storcon", [True, False])
-def test_storcon_create_delete_sk_down(neon_env_builder: NeonEnvBuilder, restart_storcon: bool):
+@pytest.mark.parametrize("restart_storcon", [RestartStorcon.RESTART, RestartStorcon.ONLINE])
+@pytest.mark.parametrize("deletetion_subject", [DeletionSubject.TENANT, DeletionSubject.TIMELINE])
+def test_storcon_create_delete_sk_down(
+    neon_env_builder: NeonEnvBuilder,
+    restart_storcon: RestartStorcon,
+    deletetion_subject: DeletionSubject,
+):
     """
     Test that the storcon can create and delete tenants and timelines with a safekeeper being down.
-      - restart_storcon: tests whether the pending ops are persisted.
+      - restart_storcon: tests that the pending ops are persisted.
         if we don't restart, we test that we don't require it to come from the db.
+      - deletion_subject: test that both single timeline and whole tenant deletion work.
     """
 
     neon_env_builder.num_safekeepers = 3
@@ -4143,6 +4159,7 @@ def test_storcon_create_delete_sk_down(neon_env_builder: NeonEnvBuilder, restart
     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[1].assert_log_contains(f"creating new timeline {tenant_id}/{timeline_id}")
     env.safekeepers[2].assert_log_contains(f"creating new timeline {tenant_id}/{timeline_id}")
@@ -4155,7 +4172,7 @@ def test_storcon_create_delete_sk_down(neon_env_builder: NeonEnvBuilder, restart
         ]
     )
 
-    if restart_storcon:
+    if restart_storcon == RestartStorcon.RESTART:
         # Restart the storcon to check that we persist operations
         env.storage_controller.stop()
         env.storage_controller.start()
@@ -4168,6 +4185,13 @@ def test_storcon_create_delete_sk_down(neon_env_builder: NeonEnvBuilder, restart
         ep.start(safekeeper_generation=1, safekeepers=[1, 2, 3])
         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=[1, 2, 3])
+        ep.safe_psql("CREATE TABLE IF NOT EXISTS t(key int, value text)")
+
     env.storage_controller.assert_log_contains("writing pending op for sk id 1")
     env.safekeepers[0].start()
 
@@ -4176,25 +4200,31 @@ def test_storcon_create_delete_sk_down(neon_env_builder: NeonEnvBuilder, restart
         env.safekeepers[0].assert_log_contains(
             f"pulling timeline {tenant_id}/{timeline_id} from safekeeper"
         )
+        env.safekeepers[0].assert_log_contains(
+            f"pulling timeline {tenant_id}/{child_timeline_id} from safekeeper"
+        )
 
     wait_until(logged_contains_on_sk)
 
     env.safekeepers[1].stop()
 
-    env.storage_controller.pageserver_api().tenant_delete(tenant_id)
+    if deletetion_subject == 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 the safekeeper deleted the timeline
     def timeline_deleted_on_active_sks():
         env.safekeepers[0].assert_log_contains(
-            f"deleting timeline {tenant_id}/{timeline_id} from disk"
+            f"deleting timeline {tenant_id}/{child_timeline_id} from disk"
         )
         env.safekeepers[2].assert_log_contains(
-            f"deleting timeline {tenant_id}/{timeline_id} from disk"
+            f"deleting timeline {tenant_id}/{child_timeline_id} from disk"
         )
 
     wait_until(timeline_deleted_on_active_sks)
 
-    if restart_storcon:
+    if restart_storcon == RestartStorcon.RESTART:
         # Restart the storcon to check that we persist operations
         env.storage_controller.stop()
         env.storage_controller.start()
@@ -4204,7 +4234,7 @@ def test_storcon_create_delete_sk_down(neon_env_builder: NeonEnvBuilder, restart
     # ensure that there is log msgs for the third safekeeper too
     def timeline_deleted_on_sk():
         env.safekeepers[1].assert_log_contains(
-            f"deleting timeline {tenant_id}/{timeline_id} from disk"
+            f"deleting timeline {tenant_id}/{child_timeline_id} from disk"
         )
 
     wait_until(timeline_deleted_on_sk)