abstract byte handling

This makes it easier to add a different client implementation alongside
the current one. I started working on a new gRPC-based protocol to
replace the libpq protocol, which will introduce a new function like
`client_libpq`, but for the new protocol.

It's a little more readable with less indentation anyway.

.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/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/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/binary/pg_sni_router.rs

@@ -297,7 +297,7 @@ async fn handle_client(
     // Starting from here we only proxy the client's traffic.
     info!("performing the proxy pass...");
 
-    match copy_bidirectional_client_compute(&mut tls_stream, &mut client).await {
+    match copy_bidirectional_client_compute(&mut tls_stream, &mut client, |_, _| {}).await {
         Ok(_) => Ok(()),
         Err(ErrorSource::Client(err)) => Err(err).context("client"),
         Err(ErrorSource::Compute(err)) => Err(err).context("compute"),

proxy/src/binary/proxy.rs

@@ -24,6 +24,7 @@ use crate::config::{
 use crate::context::parquet::ParquetUploadArgs;
 use crate::http::health_server::AppMetrics;
 use crate::metrics::Metrics;
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::rate_limiter::{
     EndpointRateLimiter, LeakyBucketConfig, RateBucketInfo, WakeComputeRateLimiter,
 };
@@ -418,6 +419,8 @@ pub async fn run() -> anyhow::Result<()> {
         64,
     ));
 
+    let conntracking = Arc::new(ConnectionTracking::default());
+
     // client facing tasks. these will exit on error or on cancellation
     // cancellation returns Ok(())
     let mut client_tasks = JoinSet::new();
@@ -431,6 +434,7 @@ pub async fn run() -> anyhow::Result<()> {
                     cancellation_token.clone(),
                     cancellation_handler.clone(),
                     endpoint_rate_limiter.clone(),
+                    conntracking.clone(),
                 ));
             }
 
@@ -453,6 +457,7 @@ pub async fn run() -> anyhow::Result<()> {
                     proxy_listener,
                     cancellation_token.clone(),
                     cancellation_handler.clone(),
+                    conntracking.clone(),
                 ));
             }
         }

proxy/src/console_redirect_proxy.rs

@@ -13,6 +13,7 @@ use crate::error::ReportableError;
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
 use crate::protocol2::{ConnectHeader, ConnectionInfo, read_proxy_protocol};
 use crate::proxy::connect_compute::{TcpMechanism, connect_to_compute};
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::proxy::handshake::{HandshakeData, handshake};
 use crate::proxy::passthrough::ProxyPassthrough;
 use crate::proxy::{
@@ -25,6 +26,7 @@ pub async fn task_main(
     listener: tokio::net::TcpListener,
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandler>,
+    conntracking: Arc<ConnectionTracking>,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("proxy has shut down");
@@ -50,6 +52,7 @@ pub async fn task_main(
         let session_id = uuid::Uuid::new_v4();
         let cancellation_handler = Arc::clone(&cancellation_handler);
         let cancellations = cancellations.clone();
+        let conntracking = Arc::clone(&conntracking);
 
         debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
 
@@ -111,6 +114,7 @@ pub async fn task_main(
                 socket,
                 conn_gauge,
                 cancellations,
+                conntracking,
             )
             .instrument(ctx.span())
             .boxed()
@@ -167,6 +171,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     stream: S,
     conn_gauge: NumClientConnectionsGuard<'static>,
     cancellations: tokio_util::task::task_tracker::TaskTracker,
+    conntracking: Arc<ConnectionTracking>,
 ) -> Result<Option<ProxyPassthrough<S>>, ClientRequestError> {
     debug!(
         protocol = %ctx.protocol(),
@@ -264,6 +269,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
         compute: node,
         session_id: ctx.session_id(),
         cancel: session,
+        conntracking,
         _req: request_gauge,
         _conn: conn_gauge,
     }))

proxy/src/metrics.rs

@@ -200,8 +200,10 @@ pub enum HttpDirection {
 #[derive(FixedCardinalityLabel, Copy, Clone)]
 #[label(singleton = "direction")]
 pub enum Direction {
-    Tx,
-    Rx,
+    #[label(rename = "tx")]
+    ComputeToClient,
+    #[label(rename = "rx")]
+    ClientToCompute,
 }
 
 #[derive(FixedCardinalityLabel, Clone, Copy, Debug)]

proxy/src/proxy/conntrack.rs

@@ -0,0 +1,564 @@
+use std::fmt;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::time::SystemTime;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct ConnId(usize);
+
+#[derive(Default)]
+pub struct ConnectionTracking {
+    conns: clashmap::ClashMap<ConnId, (ConnectionState, SystemTime)>,
+}
+
+impl ConnectionTracking {
+    pub fn new_tracker(self: &Arc<Self>) -> ConnectionTracker<Conn> {
+        let conn_id = self.new_conn_id();
+        ConnectionTracker::new(Conn {
+            conn_id,
+            tracking: Arc::clone(self),
+        })
+    }
+
+    fn new_conn_id(&self) -> ConnId {
+        static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
+        let id = ConnId(NEXT_ID.fetch_add(1, Ordering::Relaxed));
+        self.conns
+            .insert(id, (ConnectionState::Idle, SystemTime::now()));
+        id
+    }
+
+    fn update(&self, conn_id: ConnId, new_state: ConnectionState) {
+        let new_timestamp = SystemTime::now();
+        let old_state = self.conns.insert(conn_id, (new_state, new_timestamp));
+
+        if let Some((old_state, _old_timestamp)) = old_state {
+            tracing::debug!(?conn_id, %old_state, %new_state, "conntrack: update");
+        } else {
+            tracing::debug!(?conn_id, %new_state, "conntrack: update");
+        }
+    }
+
+    fn remove(&self, conn_id: ConnId) {
+        if let Some((_, (old_state, _old_timestamp))) = self.conns.remove(&conn_id) {
+            tracing::debug!(?conn_id, %old_state, "conntrack: remove");
+        }
+    }
+}
+
+pub struct Conn {
+    conn_id: ConnId,
+    tracking: Arc<ConnectionTracking>,
+}
+
+impl StateChangeObserver for Conn {
+    fn change(&mut self, _old_state: ConnectionState, new_state: ConnectionState) {
+        match new_state {
+            ConnectionState::Init
+            | ConnectionState::Idle
+            | ConnectionState::Transaction
+            | ConnectionState::Busy
+            | ConnectionState::Unknown => self.tracking.update(self.conn_id, new_state),
+            ConnectionState::Closed => self.tracking.remove(self.conn_id),
+        }
+    }
+}
+
+/// Called by `ConnectionTracker` whenever the `ConnectionState` changed.
+pub trait StateChangeObserver {
+    /// Called iff the connection's state changed.
+    fn change(&mut self, old_state: ConnectionState, new_state: ConnectionState);
+}
+
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
+#[repr(u8)]
+pub enum ConnectionState {
+    #[default]
+    Init = 0,
+    Idle = 1,
+    Transaction = 2,
+    Busy = 3,
+    Closed = 4,
+    Unknown = 5,
+}
+
+impl fmt::Display for ConnectionState {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            ConnectionState::Init => f.write_str("init"),
+            ConnectionState::Idle => f.write_str("idle"),
+            ConnectionState::Transaction => f.write_str("transaction"),
+            ConnectionState::Busy => f.write_str("busy"),
+            ConnectionState::Closed => f.write_str("closed"),
+            ConnectionState::Unknown => f.write_str("unknown"),
+        }
+    }
+}
+
+/// Tracks the `ConnectionState` of a connection by inspecting the frontend and
+/// backend stream and reacting to specific messages. Used in combination with
+/// two `TrackedStream`s.
+pub struct ConnectionTracker<SCO: StateChangeObserver> {
+    state: ConnectionState,
+    observer: SCO,
+}
+
+impl<SCO: StateChangeObserver> Drop for ConnectionTracker<SCO> {
+    fn drop(&mut self) {
+        self.observer.change(self.state, ConnectionState::Closed);
+    }
+}
+
+impl<SCO: StateChangeObserver> ConnectionTracker<SCO> {
+    pub fn new(observer: SCO) -> Self {
+        ConnectionTracker {
+            state: ConnectionState::default(),
+            observer,
+        }
+    }
+
+    pub fn frontend_message_tag(&mut self, tag: Tag) {
+        self.update_state(|old_state| Self::state_from_frontend_tag(old_state, tag));
+    }
+
+    pub fn backend_message_tag(&mut self, tag: Tag) {
+        self.update_state(|old_state| Self::state_from_backend_tag(old_state, tag));
+    }
+
+    fn update_state(&mut self, new_state_fn: impl FnOnce(ConnectionState) -> ConnectionState) {
+        let old_state = self.state;
+        let new_state = new_state_fn(old_state);
+        if old_state != new_state {
+            self.observer.change(old_state, new_state);
+            self.state = new_state;
+        }
+    }
+
+    fn state_from_frontend_tag(_old_state: ConnectionState, fe_tag: Tag) -> ConnectionState {
+        // Most activity from the client puts connection into busy state.
+        // Only the server can put a connection back into idle state.
+        match fe_tag {
+            Tag::Start | Tag::ReadyForQuery(_) | Tag::Message(_) => ConnectionState::Busy,
+            Tag::End => ConnectionState::Closed,
+            Tag::Lost => ConnectionState::Unknown,
+        }
+    }
+
+    fn state_from_backend_tag(old_state: ConnectionState, be_tag: Tag) -> ConnectionState {
+        match be_tag {
+            // Check for RFQ and put connection into idle or idle in transaction state.
+            Tag::ReadyForQuery(b'I') => ConnectionState::Idle,
+            Tag::ReadyForQuery(b'T') => ConnectionState::Transaction,
+            Tag::ReadyForQuery(b'E') => ConnectionState::Transaction,
+            // We can't put a connection into idle state for unknown RFQ status.
+            Tag::ReadyForQuery(_) => ConnectionState::Unknown,
+            // Ignore out-fo message from the server.
+            Tag::NOTICE | Tag::NOTIFICATION_RESPONSE | Tag::PARAMETER_STATUS => old_state,
+            // All other activity from server puts connection into busy state.
+            Tag::Start | Tag::Message(_) => ConnectionState::Busy,
+
+            Tag::End => ConnectionState::Closed,
+            Tag::Lost => ConnectionState::Unknown,
+        }
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum Tag {
+    Message(u8),
+    ReadyForQuery(u8),
+    Start,
+    End,
+    Lost,
+}
+
+impl Tag {
+    const READY_FOR_QUERY: Tag = Tag::Message(b'Z');
+    const NOTICE: Tag = Tag::Message(b'N');
+    const NOTIFICATION_RESPONSE: Tag = Tag::Message(b'A');
+    const PARAMETER_STATUS: Tag = Tag::Message(b'S');
+}
+
+impl fmt::Display for Tag {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match *self {
+            Tag::Start => f.write_str("start"),
+            Tag::End => f.write_str("end"),
+            Tag::Lost => f.write_str("lost"),
+            Tag::Message(tag) => write!(f, "'{}'", tag as char),
+            Tag::ReadyForQuery(status) => write!(f, "ReadyForQuery:'{}'", status as char),
+        }
+    }
+}
+
+pub trait TagObserver {
+    fn observe(&mut self, tag: Tag);
+}
+
+impl<F: FnMut(Tag)> TagObserver for F {
+    fn observe(&mut self, tag: Tag) {
+        (self)(tag);
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub(super) enum StreamScannerState {
+    #[allow(dead_code)]
+    /// Initial state when no message has been read and we are looling for a
+    /// message without a tag.
+    Start,
+    /// Read a message tag.
+    Tag,
+    /// Read the length bytes and calculate the total length.
+    Length {
+        tag: Tag,
+        /// Number of bytes missing to know the full length of the message: 0..=4
+        length_bytes_missing: usize,
+        /// Total length of the message (without tag) that is calculated as we
+        /// read the bytes for the length.
+        calculated_length: usize,
+    },
+    /// Read (= skip) the payload.
+    Payload {
+        tag: Tag,
+        /// If this is the first time payload bytes are read. Important when
+        /// inspecting specific messages, like ReadyForQuery.
+        first: bool,
+        /// Number of payload bytes left to read before looking for a new tag.
+        bytes_to_skip: usize,
+    },
+    /// Stream was terminated.
+    End,
+    /// Stream ended up in a lost state. We only stop tracking the stream, not
+    /// interrupt it.
+    Lost,
+}
+
+impl StreamScannerState {
+    pub(super) fn scan_bytes<TO: TagObserver>(&mut self, mut buf: &[u8], observer: &mut TO) {
+        use StreamScannerState as S;
+
+        if matches!(*self, S::End | S::Lost) {
+            return;
+        }
+        if buf.is_empty() {
+            match *self {
+                S::Start | S::Tag => {
+                    observer.observe(Tag::End);
+                    *self = S::End;
+                    return;
+                }
+                S::Length { .. } | S::Payload { .. } => {
+                    observer.observe(Tag::Lost);
+                    *self = S::Lost;
+                    return;
+                }
+                S::End | S::Lost => unreachable!(),
+            }
+        }
+
+        while !buf.is_empty() {
+            match *self {
+                S::Start => {
+                    *self = S::Length {
+                        tag: Tag::Start,
+                        length_bytes_missing: 4,
+                        calculated_length: 0,
+                    };
+                }
+
+                S::Tag => {
+                    let tag = buf.first().copied().expect("buf not empty");
+                    buf = &buf[1..];
+
+                    *self = S::Length {
+                        tag: Tag::Message(tag),
+                        length_bytes_missing: 4,
+                        calculated_length: 0,
+                    };
+                }
+
+                S::Length {
+                    tag,
+                    mut length_bytes_missing,
+                    mut calculated_length,
+                } => {
+                    let consume = length_bytes_missing.min(buf.len());
+
+                    let (length_bytes, remainder) = buf.split_at(consume);
+                    for b in length_bytes {
+                        calculated_length <<= 8;
+                        calculated_length |= *b as usize;
+                    }
+                    buf = remainder;
+
+                    length_bytes_missing -= consume;
+                    if length_bytes_missing == 0 {
+                        let Some(bytes_to_skip) = calculated_length.checked_sub(4) else {
+                            observer.observe(Tag::Lost);
+                            *self = S::Lost;
+                            return;
+                        };
+
+                        if bytes_to_skip == 0 {
+                            observer.observe(tag);
+                            *self = S::Tag;
+                        } else {
+                            *self = S::Payload {
+                                tag,
+                                first: true,
+                                bytes_to_skip,
+                            };
+                        }
+                    } else {
+                        *self = S::Length {
+                            tag,
+                            length_bytes_missing,
+                            calculated_length,
+                        };
+                    }
+                }
+
+                S::Payload {
+                    tag,
+                    first,
+                    mut bytes_to_skip,
+                } => {
+                    let consume = bytes_to_skip.min(buf.len());
+                    bytes_to_skip -= consume;
+                    if bytes_to_skip == 0 {
+                        if tag == Tag::READY_FOR_QUERY && first && consume == 1 {
+                            let status = buf.first().copied().expect("buf not empty");
+                            observer.observe(Tag::ReadyForQuery(status));
+                        } else {
+                            observer.observe(tag);
+                        }
+                        *self = S::Tag;
+                    } else {
+                        *self = S::Payload {
+                            tag,
+                            first: false,
+                            bytes_to_skip,
+                        };
+                    }
+                    buf = &buf[consume..];
+                }
+
+                S::End | S::Lost => unreachable!(),
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::cell::RefCell;
+    use std::io;
+    use std::pin::Pin;
+    use std::rc::Rc;
+    use std::sync::{Arc, Mutex};
+    use std::task::{Context, Poll};
+
+    use tokio::io::{AsyncRead, ReadBuf};
+    use tokio::io::{AsyncReadExt, BufReader};
+
+    use super::*;
+
+    #[test]
+    fn test_stream_scanner() {
+        let tags = Rc::new(RefCell::new(Vec::new()));
+        let observer_tags = tags.clone();
+        let mut observer = move |tag| {
+            observer_tags.borrow_mut().push(tag);
+        };
+        let mut state = StreamScannerState::Start;
+
+        state.scan_bytes(&[0, 0], &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[]);
+        assert_eq!(
+            state,
+            StreamScannerState::Length {
+                tag: Tag::Start,
+                length_bytes_missing: 2,
+                calculated_length: 0,
+            }
+        );
+
+        state.scan_bytes(&[0x01, 0x01, 0x00], &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[]);
+        assert_eq!(
+            state,
+            StreamScannerState::Payload {
+                tag: Tag::Start,
+                first: false,
+                bytes_to_skip: 0x00000101 - 4 - 1,
+            }
+        );
+
+        state.scan_bytes(vec![0; 0x00000101 - 4 - 1 - 1].as_slice(), &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[]);
+        assert_eq!(
+            state,
+            StreamScannerState::Payload {
+                tag: Tag::Start,
+                first: false,
+                bytes_to_skip: 1,
+            }
+        );
+
+        state.scan_bytes(&[0x00, b'A', 0x00, 0x00, 0x00, 0x08], &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[Tag::Start]);
+        assert_eq!(
+            state,
+            StreamScannerState::Payload {
+                tag: Tag::Message(b'A'),
+                first: true,
+                bytes_to_skip: 4,
+            }
+        );
+
+        state.scan_bytes(&[0, 0, 0, 0], &mut observer);
+        assert_eq!(tags.borrow().as_slice(), &[Tag::Start, Tag::Message(b'A')]);
+        assert_eq!(state, StreamScannerState::Tag);
+
+        state.scan_bytes(&[b'Z', 0x00, 0x00, 0x00, 0x05, b'T'], &mut observer);
+        assert_eq!(
+            tags.borrow().as_slice(),
+            &[Tag::Start, Tag::Message(b'A'), Tag::ReadyForQuery(b'T')]
+        );
+        assert_eq!(state, StreamScannerState::Tag);
+
+        state.scan_bytes(&[], &mut observer);
+        assert_eq!(
+            tags.borrow().as_slice(),
+            &[
+                Tag::Start,
+                Tag::Message(b'A'),
+                Tag::ReadyForQuery(b'T'),
+                Tag::End
+            ]
+        );
+        assert_eq!(state, StreamScannerState::End);
+    }
+
+    #[tokio::test]
+    async fn test_connection_tracker() {
+        let transitions: Arc<Mutex<Vec<(ConnectionState, ConnectionState)>>> = Arc::default();
+        struct Observer(Arc<Mutex<Vec<(ConnectionState, ConnectionState)>>>);
+        impl StateChangeObserver for Observer {
+            fn change(&mut self, old_state: ConnectionState, new_state: ConnectionState) {
+                self.0.lock().unwrap().push((old_state, new_state));
+            }
+        }
+        let mut tracker = ConnectionTracker::new(Observer(transitions.clone()));
+
+        let stream = BufReader::new(
+            &[
+                0, 0, 0, 4, // Init
+                b'Z', 0, 0, 0, 5, b'I', // Init -> Idle
+                b'x', 0, 0, 0, 4, // Idle -> Busy
+                b'Z', 0, 0, 0, 5, b'I', // Busy -> Idle
+            ][..],
+        );
+        // AsyncRead
+        let mut stream = TrackedStream::new(stream, |tag| tracker.backend_message_tag(tag));
+
+        let mut readbuf = [0; 2];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 2);
+        assert_eq!(&readbuf, &[0, 0,]);
+        assert!(transitions.lock().unwrap().is_empty());
+
+        let mut readbuf = [0; 2];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 2);
+        assert_eq!(&readbuf, &[0, 4]);
+        assert_eq!(
+            transitions.lock().unwrap().as_slice(),
+            &[(ConnectionState::Init, ConnectionState::Busy)]
+        );
+
+        let mut readbuf = [0; 6];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 6);
+        assert_eq!(&readbuf, &[b'Z', 0, 0, 0, 5, b'I']);
+        assert_eq!(
+            transitions.lock().unwrap().as_slice(),
+            &[
+                (ConnectionState::Init, ConnectionState::Busy),
+                (ConnectionState::Busy, ConnectionState::Idle),
+            ]
+        );
+
+        let mut readbuf = [0; 5];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 5);
+        assert_eq!(&readbuf, &[b'x', 0, 0, 0, 4]);
+        assert_eq!(
+            transitions.lock().unwrap().as_slice(),
+            &[
+                (ConnectionState::Init, ConnectionState::Busy),
+                (ConnectionState::Busy, ConnectionState::Idle),
+                (ConnectionState::Idle, ConnectionState::Busy),
+            ]
+        );
+
+        let mut readbuf = [0; 6];
+        let n = stream.read_exact(&mut readbuf).await.unwrap();
+        assert_eq!(n, 6);
+        assert_eq!(&readbuf, &[b'Z', 0, 0, 0, 5, b'I']);
+        assert_eq!(
+            transitions.lock().unwrap().as_slice(),
+            &[
+                (ConnectionState::Init, ConnectionState::Busy),
+                (ConnectionState::Busy, ConnectionState::Idle),
+                (ConnectionState::Idle, ConnectionState::Busy),
+                (ConnectionState::Busy, ConnectionState::Idle),
+            ]
+        );
+    }
+
+    pub struct TrackedStream<S, TO> {
+        stream: S,
+        observer: TO,
+        state: StreamScannerState,
+    }
+
+    impl<S: Unpin, TO> Unpin for TrackedStream<S, TO> {}
+
+    impl<S: AsyncRead + Unpin, TO: TagObserver> TrackedStream<S, TO> {
+        pub const fn new(stream: S, observer: TO) -> Self {
+            TrackedStream {
+                stream,
+                observer,
+                state: StreamScannerState::Start,
+            }
+        }
+    }
+
+    impl<S: AsyncRead + Unpin, TO: TagObserver> AsyncRead for TrackedStream<S, TO> {
+        #[inline]
+        fn poll_read(
+            self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+            buf: &mut ReadBuf<'_>,
+        ) -> Poll<io::Result<()>> {
+            let Self {
+                stream,
+                observer,
+                state,
+            } = Pin::into_inner(self);
+
+            let old_len = buf.filled().len();
+            match Pin::new(stream).poll_read(cx, buf) {
+                Poll::Ready(Ok(())) => {
+                    let new_len = buf.filled().len();
+                    state.scan_bytes(&buf.filled()[old_len..new_len], observer);
+                    Poll::Ready(Ok(()))
+                }
+                Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+                Poll::Pending => Poll::Pending,
+            }
+        }
+    }
+}

proxy/src/proxy/copy_bidirectional.rs

@@ -6,6 +6,8 @@ use std::task::{Context, Poll, ready};
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tracing::info;
 
+use crate::metrics::Direction;
+
 #[derive(Debug)]
 enum TransferState {
     Running(CopyBuffer),
@@ -43,6 +45,8 @@ pub enum ErrorSource {
 fn transfer_one_direction<A, B>(
     cx: &mut Context<'_>,
     state: &mut TransferState,
+    direction: Direction,
+    conn_tracker: &mut impl for<'a> FnMut(Direction, &'a [u8]),
     r: &mut A,
     w: &mut B,
 ) -> Poll<Result<u64, ErrorDirection>>
@@ -55,7 +59,8 @@ where
     loop {
         match state {
             TransferState::Running(buf) => {
-                let count = ready!(buf.poll_copy(cx, r.as_mut(), w.as_mut()))?;
+                let count =
+                    ready!(buf.poll_copy(cx, direction, conn_tracker, r.as_mut(), w.as_mut()))?;
                 *state = TransferState::ShuttingDown(count);
             }
             TransferState::ShuttingDown(count) => {
@@ -71,6 +76,7 @@ where
 pub async fn copy_bidirectional_client_compute<Client, Compute>(
     client: &mut Client,
     compute: &mut Compute,
+    mut conn_tracker: impl for<'a> FnMut(Direction, &'a [u8]),
 ) -> Result<(u64, u64), ErrorSource>
 where
     Client: AsyncRead + AsyncWrite + Unpin + ?Sized,
@@ -80,12 +86,24 @@ where
     let mut compute_to_client = TransferState::Running(CopyBuffer::new());
 
     poll_fn(|cx| {
-        let mut client_to_compute_result =
-            transfer_one_direction(cx, &mut client_to_compute, client, compute)
-                .map_err(ErrorSource::from_client)?;
-        let mut compute_to_client_result =
-            transfer_one_direction(cx, &mut compute_to_client, compute, client)
-                .map_err(ErrorSource::from_compute)?;
+        let mut client_to_compute_result = transfer_one_direction(
+            cx,
+            &mut client_to_compute,
+            Direction::ClientToCompute,
+            &mut conn_tracker,
+            client,
+            compute,
+        )
+        .map_err(ErrorSource::from_client)?;
+        let mut compute_to_client_result = transfer_one_direction(
+            cx,
+            &mut compute_to_client,
+            Direction::ComputeToClient,
+            &mut conn_tracker,
+            compute,
+            client,
+        )
+        .map_err(ErrorSource::from_compute)?;
 
         // TODO: 1 info log, with a enum label for close direction.
 
@@ -95,9 +113,15 @@ where
                 info!("Compute is done, terminate client");
                 // Initiate shutdown
                 client_to_compute = TransferState::ShuttingDown(buf.amt);
-                client_to_compute_result =
-                    transfer_one_direction(cx, &mut client_to_compute, client, compute)
-                        .map_err(ErrorSource::from_client)?;
+                client_to_compute_result = transfer_one_direction(
+                    cx,
+                    &mut client_to_compute,
+                    Direction::ClientToCompute,
+                    &mut conn_tracker,
+                    client,
+                    compute,
+                )
+                .map_err(ErrorSource::from_client)?;
             }
         }
 
@@ -107,9 +131,15 @@ where
                 info!("Client is done, terminate compute");
                 // Initiate shutdown
                 compute_to_client = TransferState::ShuttingDown(buf.amt);
-                compute_to_client_result =
-                    transfer_one_direction(cx, &mut compute_to_client, compute, client)
-                        .map_err(ErrorSource::from_compute)?;
+                compute_to_client_result = transfer_one_direction(
+                    cx,
+                    &mut compute_to_client,
+                    Direction::ComputeToClient,
+                    &mut conn_tracker,
+                    compute,
+                    client,
+                )
+                .map_err(ErrorSource::from_compute)?;
             }
         }
 
@@ -148,6 +178,8 @@ impl CopyBuffer {
     fn poll_fill_buf<R>(
         &mut self,
         cx: &mut Context<'_>,
+        direction: Direction,
+        conn_tracker: &mut impl for<'a> FnMut(Direction, &'a [u8]),
         reader: Pin<&mut R>,
     ) -> Poll<io::Result<()>>
     where
@@ -158,6 +190,8 @@ impl CopyBuffer {
         buf.set_filled(me.cap);
 
         let res = reader.poll_read(cx, &mut buf);
+        conn_tracker(direction, &buf.filled()[me.cap..]);
+
         if let Poll::Ready(Ok(())) = res {
             let filled_len = buf.filled().len();
             me.read_done = me.cap == filled_len;
@@ -169,6 +203,8 @@ impl CopyBuffer {
     fn poll_write_buf<R, W>(
         &mut self,
         cx: &mut Context<'_>,
+        direction: Direction,
+        conn_tracker: &mut impl for<'a> FnMut(Direction, &'a [u8]),
         mut reader: Pin<&mut R>,
         mut writer: Pin<&mut W>,
     ) -> Poll<Result<usize, ErrorDirection>>
@@ -182,7 +218,8 @@ impl CopyBuffer {
                 // Top up the buffer towards full if we can read a bit more
                 // data - this should improve the chances of a large write
                 if !me.read_done && me.cap < me.buf.len() {
-                    ready!(me.poll_fill_buf(cx, reader.as_mut())).map_err(ErrorDirection::Read)?;
+                    ready!(me.poll_fill_buf(cx, direction, conn_tracker, reader.as_mut()))
+                        .map_err(ErrorDirection::Read)?;
                 }
                 Poll::Pending
             }
@@ -193,6 +230,8 @@ impl CopyBuffer {
     pub(super) fn poll_copy<R, W>(
         &mut self,
         cx: &mut Context<'_>,
+        direction: Direction,
+        conn_tracker: &mut impl for<'a> FnMut(Direction, &'a [u8]),
         mut reader: Pin<&mut R>,
         mut writer: Pin<&mut W>,
     ) -> Poll<Result<u64, ErrorDirection>>
@@ -204,7 +243,7 @@ impl CopyBuffer {
             // If there is some space left in our buffer, then we try to read some
             // data to continue, thus maximizing the chances of a large write.
             if self.cap < self.buf.len() && !self.read_done {
-                match self.poll_fill_buf(cx, reader.as_mut()) {
+                match self.poll_fill_buf(cx, direction, conn_tracker, reader.as_mut()) {
                     Poll::Ready(Ok(())) => (),
                     Poll::Ready(Err(err)) => return Poll::Ready(Err(ErrorDirection::Read(err))),
                     Poll::Pending => {
@@ -227,7 +266,13 @@ impl CopyBuffer {
 
             // If our buffer has some data, let's write it out!
             while self.pos < self.cap {
-                let i = ready!(self.poll_write_buf(cx, reader.as_mut(), writer.as_mut()))?;
+                let i = ready!(self.poll_write_buf(
+                    cx,
+                    direction,
+                    conn_tracker,
+                    reader.as_mut(),
+                    writer.as_mut()
+                ))?;
                 if i == 0 {
                     return Poll::Ready(Err(ErrorDirection::Write(io::Error::new(
                         io::ErrorKind::WriteZero,
@@ -278,9 +323,10 @@ mod tests {
         compute_client.write_all(b"Neon").await.unwrap();
         compute_client.shutdown().await.unwrap();
 
-        let result = copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy)
-            .await
-            .unwrap();
+        let result =
+            copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy, |_, _| {})
+                .await
+                .unwrap();
 
         // Assert correct transferred amounts
         let (client_to_compute_count, compute_to_client_count) = result;
@@ -301,9 +347,10 @@ mod tests {
             .await
             .unwrap();
 
-        let result = copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy)
-            .await
-            .unwrap();
+        let result =
+            copy_bidirectional_client_compute(&mut client_proxy, &mut compute_proxy, |_, _| {})
+                .await
+                .unwrap();
 
         // Assert correct transferred amounts
         let (client_to_compute_count, compute_to_client_count) = result;

proxy/src/proxy/mod.rs

@@ -2,6 +2,7 @@
 mod tests;
 
 pub(crate) mod connect_compute;
+pub mod conntrack;
 mod copy_bidirectional;
 pub(crate) mod handshake;
 pub(crate) mod passthrough;
@@ -30,6 +31,7 @@ use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
 use crate::protocol2::{ConnectHeader, ConnectionInfo, ConnectionInfoExtra, read_proxy_protocol};
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::proxy::handshake::{HandshakeData, handshake};
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::stream::{PqStream, Stream};
@@ -60,6 +62,7 @@ pub async fn task_main(
     cancellation_token: CancellationToken,
     cancellation_handler: Arc<CancellationHandler>,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
+    conntracking: Arc<ConnectionTracking>,
 ) -> anyhow::Result<()> {
     scopeguard::defer! {
         info!("proxy has shut down");
@@ -85,6 +88,7 @@ pub async fn task_main(
         let session_id = uuid::Uuid::new_v4();
         let cancellation_handler = Arc::clone(&cancellation_handler);
         let cancellations = cancellations.clone();
+        let conntracking = Arc::clone(&conntracking);
 
         debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
         let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
@@ -149,6 +153,7 @@ pub async fn task_main(
                 endpoint_rate_limiter2,
                 conn_gauge,
                 cancellations,
+                conntracking,
             )
             .instrument(ctx.span())
             .boxed()
@@ -268,6 +273,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     conn_gauge: NumClientConnectionsGuard<'static>,
     cancellations: tokio_util::task::task_tracker::TaskTracker,
+    conntracking: Arc<ConnectionTracking>,
 ) -> Result<Option<ProxyPassthrough<S>>, ClientRequestError> {
     debug!(
         protocol = %ctx.protocol(),
@@ -409,6 +415,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
         compute: node,
         session_id: ctx.session_id(),
         cancel: session,
+        conntracking,
         _req: request_gauge,
         _conn: conn_gauge,
     }))

proxy/src/proxy/passthrough.rs

@@ -1,7 +1,8 @@
+use std::sync::Arc;
+
 use smol_str::SmolStr;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::debug;
-use utils::measured_stream::MeasuredStream;
 
 use super::copy_bidirectional::ErrorSource;
 use crate::cancellation;
@@ -9,16 +10,19 @@ use crate::compute::PostgresConnection;
 use crate::config::ComputeConfig;
 use crate::control_plane::messages::MetricsAuxInfo;
 use crate::metrics::{Direction, Metrics, NumClientConnectionsGuard, NumConnectionRequestsGuard};
+use crate::proxy::conntrack::{ConnectionTracking, StreamScannerState};
+use crate::proxy::copy_bidirectional::copy_bidirectional_client_compute;
 use crate::stream::Stream;
 use crate::usage_metrics::{Ids, MetricCounterRecorder, USAGE_METRICS};
 
 /// Forward bytes in both directions (client <-> compute).
 #[tracing::instrument(skip_all)]
 pub(crate) async fn proxy_pass(
-    client: impl AsyncRead + AsyncWrite + Unpin,
-    compute: impl AsyncRead + AsyncWrite + Unpin,
+    mut client: impl AsyncRead + AsyncWrite + Unpin,
+    mut compute: impl AsyncRead + AsyncWrite + Unpin,
     aux: MetricsAuxInfo,
     private_link_id: Option<SmolStr>,
+    conntracking: &Arc<ConnectionTracking>,
 ) -> Result<(), ErrorSource> {
     // we will report ingress at a later date
     let usage_tx = USAGE_METRICS.register(Ids {
@@ -27,35 +31,35 @@ pub(crate) async fn proxy_pass(
         private_link_id,
     });
 
+    let mut conn_tracker = conntracking.new_tracker();
+
     let metrics = &Metrics::get().proxy.io_bytes;
-    let m_sent = metrics.with_labels(Direction::Tx);
-    let mut client = MeasuredStream::new(
-        client,
-        |_| {},
-        |cnt| {
-            // Number of bytes we sent to the client (outbound).
-            metrics.get_metric(m_sent).inc_by(cnt as u64);
-            usage_tx.record_egress(cnt as u64);
-        },
-    );
+    let m_sent = metrics.with_labels(Direction::ComputeToClient);
+    let m_recv = metrics.with_labels(Direction::ClientToCompute);
 
-    let m_recv = metrics.with_labels(Direction::Rx);
-    let mut compute = MeasuredStream::new(
-        compute,
-        |_| {},
-        |cnt| {
-            // Number of bytes the client sent to the compute node (inbound).
-            metrics.get_metric(m_recv).inc_by(cnt as u64);
-            usage_tx.record_ingress(cnt as u64);
-        },
-    );
+    let mut client_to_compute = StreamScannerState::Tag;
+    let mut compute_to_client = StreamScannerState::Tag;
 
-    // Starting from here we only proxy the client's traffic.
     debug!("performing the proxy pass...");
-    let _ = crate::proxy::copy_bidirectional::copy_bidirectional_client_compute(
-        &mut client,
-        &mut compute,
-    )
+
+    let _ = copy_bidirectional_client_compute(&mut client, &mut compute, |direction, bytes| {
+        match direction {
+            Direction::ClientToCompute => {
+                client_to_compute
+                    .scan_bytes(bytes, &mut |tag| conn_tracker.frontend_message_tag(tag));
+
+                metrics.get_metric(m_recv).inc_by(bytes.len() as u64);
+                usage_tx.record_ingress(bytes.len() as u64);
+            }
+            Direction::ComputeToClient => {
+                compute_to_client
+                    .scan_bytes(bytes, &mut |tag| conn_tracker.backend_message_tag(tag));
+
+                metrics.get_metric(m_sent).inc_by(bytes.len() as u64);
+                usage_tx.record_egress(bytes.len() as u64);
+            }
+        }
+    })
     .await?;
 
     Ok(())
@@ -68,6 +72,7 @@ pub(crate) struct ProxyPassthrough<S> {
     pub(crate) session_id: uuid::Uuid,
     pub(crate) private_link_id: Option<SmolStr>,
     pub(crate) cancel: cancellation::Session,
+    pub(crate) conntracking: Arc<ConnectionTracking>,
 
     pub(crate) _req: NumConnectionRequestsGuard<'static>,
     pub(crate) _conn: NumClientConnectionsGuard<'static>,
@@ -83,6 +88,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> ProxyPassthrough<S> {
             self.compute.stream,
             self.aux,
             self.private_link_id,
+            &self.conntracking,
         )
         .await;
         if let Err(err) = self

proxy/src/serverless/mod.rs

@@ -50,6 +50,7 @@ use crate::context::RequestContext;
 use crate::ext::TaskExt;
 use crate::metrics::Metrics;
 use crate::protocol2::{ChainRW, ConnectHeader, ConnectionInfo, read_proxy_protocol};
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::proxy::run_until_cancelled;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::serverless::backend::PoolingBackend;
@@ -124,6 +125,9 @@ pub async fn task_main(
     connections.close(); // allows `connections.wait to complete`
 
     let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
+
+    let conntracking = Arc::new(ConnectionTracking::default());
+
     while let Some(res) = run_until_cancelled(ws_listener.accept(), &cancellation_token).await {
         let (conn, peer_addr) = res.context("could not accept TCP stream")?;
         if let Err(e) = conn.set_nodelay(true) {
@@ -153,6 +157,8 @@ pub async fn task_main(
         let cancellation_handler = cancellation_handler.clone();
         let endpoint_rate_limiter = endpoint_rate_limiter.clone();
         let cancellations = cancellations.clone();
+        let conntracking = Arc::clone(&conntracking);
+
         connections.spawn(
             async move {
                 let conn_token2 = conn_token.clone();
@@ -185,6 +191,7 @@ pub async fn task_main(
                     cancellation_handler,
                     endpoint_rate_limiter,
                     conn_token,
+                    conntracking,
                     conn,
                     conn_info,
                     session_id,
@@ -309,6 +316,7 @@ async fn connection_handler(
     cancellation_handler: Arc<CancellationHandler>,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     cancellation_token: CancellationToken,
+    conntracking: Arc<ConnectionTracking>,
     conn: AsyncRW,
     conn_info: ConnectionInfo,
     session_id: uuid::Uuid,
@@ -347,6 +355,7 @@ async fn connection_handler(
             // `request_handler` is not cancel safe. It expects to be cancelled only at specific times.
             // By spawning the future, we ensure it never gets cancelled until it decides to.
             let cancellations = cancellations.clone();
+            let conntracking = Arc::clone(&conntracking);
             let handler = connections.spawn(
                 request_handler(
                     req,
@@ -359,6 +368,7 @@ async fn connection_handler(
                     http_request_token,
                     endpoint_rate_limiter.clone(),
                     cancellations,
+                    conntracking,
                 )
                 .in_current_span()
                 .map_ok_or_else(api_error_into_response, |r| r),
@@ -407,6 +417,7 @@ async fn request_handler(
     http_cancellation_token: CancellationToken,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     cancellations: TaskTracker,
+    conntracking: Arc<ConnectionTracking>,
 ) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ApiError> {
     let host = request
         .headers()
@@ -452,6 +463,7 @@ async fn request_handler(
                     endpoint_rate_limiter,
                     host,
                     cancellations,
+                    conntracking,
                 )
                 .await
                 {

proxy/src/serverless/websocket.rs

@@ -17,6 +17,7 @@ use crate::config::ProxyConfig;
 use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::Metrics;
+use crate::proxy::conntrack::ConnectionTracking;
 use crate::proxy::{ClientMode, ErrorSource, handle_client};
 use crate::rate_limiter::EndpointRateLimiter;
 
@@ -133,6 +134,7 @@ pub(crate) async fn serve_websocket(
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     hostname: Option<String>,
     cancellations: tokio_util::task::task_tracker::TaskTracker,
+    conntracking: Arc<ConnectionTracking>,
 ) -> anyhow::Result<()> {
     let websocket = websocket.await?;
     let websocket = WebSocketServer::after_handshake(TokioIo::new(websocket));
@@ -152,6 +154,7 @@ pub(crate) async fn serve_websocket(
         endpoint_rate_limiter,
         conn_gauge,
         cancellations,
+        conntracking,
     ))
     .await;
 

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)