Change l0 definition (currently panics)

Refactors Compute::prepare_and_run. It's split into subroutines
differently, to make it easier to attach tracing spans to the
different stages. The high-level logic for waiting for Postgres to
exit is moved to the caller.

Replace 'env_logger' with 'tracing', and add `#instrument` directives
to different stages fo the startup process. This is a fairly
mechanical change, except for the changes in 'spec.rs'. 'spec.rs'
contained some complicated formatting, where parts of log messages
were printed directly to stdout with `print`s. That was a bit messed
up because the log normally goes to stderr, but those lines were
printed to stdout. In our docker images, stderr and stdout both go to
the same place so you wouldn't notice, but I don't think it was
intentional.

This changes the log format to the default
'tracing_subscriber::format' format. It's different from the Postgres
log format, however, and because both compute_tools and Postgres print
to the same log, it's now a mix of two different formats.  I'm not
sure how the Grafana log parsing pipeline can handle that. If it's a
problem, we can build custom formatter to change the compute_tools log
format to be the same as Postgres's, like it was before this commit,
or we can change the Postgres log format to match tracing_formatter's,
or we can start printing compute_tool's log output to a different
destination than Postgres

.config/hakari.toml

@@ -4,7 +4,7 @@
 hakari-package = "workspace_hack"
 
 # Format for `workspace-hack = ...` lines in other Cargo.tomls. Requires cargo-hakari 0.9.8 or above.
-dep-format-version = "2"
+dep-format-version = "3"
 
 # Setting workspace.resolver = "2" in the root Cargo.toml is HIGHLY recommended.
 # Hakari works much better with the new feature resolver.

.github/ansible/prod.ap-southeast-1.hosts.yaml

@@ -6,6 +6,8 @@ storage:
     broker_endpoint: http://storage-broker-lb.epsilon.ap-southeast-1.internal.aws.neon.tech:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"

.github/ansible/prod.eu-central-1.hosts.yaml

@@ -6,6 +6,8 @@ storage:
     broker_endpoint: http://storage-broker-lb.gamma.eu-central-1.internal.aws.neon.tech:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"

.github/ansible/prod.us-east-2.hosts.yaml

@@ -6,6 +6,8 @@ storage:
     broker_endpoint: http://storage-broker-lb.delta.us-east-2.internal.aws.neon.tech:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"
@@ -34,4 +36,4 @@ storage:
           ansible_host:  i-06d113fb73bfddeb0
         safekeeper-2.us-east-2.aws.neon.tech:
           ansible_host:  i-09f66c8e04afff2e8
-          
+

.github/ansible/prod.us-west-2.hosts.yaml

@@ -6,6 +6,8 @@ storage:
     broker_endpoint: http://storage-broker-lb.eta.us-west-2.internal.aws.neon.tech:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"

.github/ansible/production.hosts.yaml

@@ -7,6 +7,8 @@ storage:
     broker_endpoint: http://storage-broker.prod.local:50051
     pageserver_config_stub:
       pg_distrib_dir: /usr/local
+      metric_collection_endpoint: http://console-release.local/billing/api/v1/usage_events
+      metric_collection_interval: 10min
       remote_storage:
         bucket_name: "{{ bucket_name }}"
         bucket_region: "{{ bucket_region }}"

.github/ansible/staging.eu-west-1.hosts.yaml

@@ -18,7 +18,7 @@ storage:
     ansible_aws_ssm_region: eu-west-1
     ansible_aws_ssm_bucket_name: neon-dev-storage-eu-west-1
     console_region_id: aws-eu-west-1
-    sentry_environment: development
+    sentry_environment: staging
 
   children:
     pageservers:

.github/ansible/staging.us-east-2.hosts.yaml

@@ -18,7 +18,7 @@ storage:
     ansible_aws_ssm_region: us-east-2
     ansible_aws_ssm_bucket_name: neon-staging-storage-us-east-2
     console_region_id: aws-us-east-2
-    sentry_environment: development
+    sentry_environment: staging
 
   children:
     pageservers:

.github/helm-values/dev-eu-west-1-zeta.neon-proxy-scram.yaml

@@ -8,8 +8,10 @@ settings:
   authBackend: "console"
   authEndpoint: "http://console-staging.local/management/api/v2"
   domain: "*.eu-west-1.aws.neon.build"
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"
   wssPort: 8443
+  metricCollectionEndpoint: "http://console-staging.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "1min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:

.github/helm-values/dev-eu-west-1-zeta.neon-storage-broker.yaml

@@ -49,4 +49,4 @@ extraManifests:
           - "{{ .Release.Namespace }}"
 
 settings:
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"

.github/helm-values/dev-us-east-2-beta.neon-proxy-link.yaml

@@ -8,7 +8,9 @@ settings:
   authBackend: "link"
   authEndpoint: "https://console.stage.neon.tech/authenticate_proxy_request/"
   uri: "https://console.stage.neon.tech/psql_session/"
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"
+  metricCollectionEndpoint: "http://console-staging.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "1min"
 
 # -- Additional labels for neon-proxy-link pods
 podLabels:

.github/helm-values/dev-us-east-2-beta.neon-proxy-scram-legacy.yaml

@@ -8,8 +8,10 @@ settings:
   authBackend: "console"
   authEndpoint: "http://console-staging.local/management/api/v2"
   domain: "*.cloud.stage.neon.tech"
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"
   wssPort: 8443
+  metricCollectionEndpoint: "http://console-staging.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "1min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:

.github/helm-values/dev-us-east-2-beta.neon-proxy-scram.yaml

@@ -8,8 +8,10 @@ settings:
   authBackend: "console"
   authEndpoint: "http://console-staging.local/management/api/v2"
   domain: "*.us-east-2.aws.neon.build"
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"
   wssPort: 8443
+  metricCollectionEndpoint: "http://console-staging.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "1min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:

.github/helm-values/dev-us-east-2-beta.neon-storage-broker.yaml

@@ -49,4 +49,4 @@ extraManifests:
           - "{{ .Release.Namespace }}"
 
 settings:
-  sentryEnvironment: "development"
+  sentryEnvironment: "staging"

.github/helm-values/prod-ap-southeast-1-epsilon.neon-proxy-scram.yaml

@@ -10,6 +10,8 @@ settings:
   domain: "*.ap-southeast-1.aws.neon.tech"
   sentryEnvironment: "production"
   wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:

.github/helm-values/prod-eu-central-1-gamma.neon-proxy-scram.yaml

@@ -10,6 +10,8 @@ settings:
   domain: "*.eu-central-1.aws.neon.tech"
   sentryEnvironment: "production"
   wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:

.github/helm-values/prod-us-east-2-delta.neon-proxy-scram.yaml

@@ -10,6 +10,8 @@ settings:
   domain: "*.us-east-2.aws.neon.tech"
   sentryEnvironment: "production"
   wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:

.github/helm-values/prod-us-west-2-eta.neon-proxy-scram.yaml

@@ -10,6 +10,8 @@ settings:
   domain: "*.us-west-2.aws.neon.tech"
   sentryEnvironment: "production"
   wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 # -- Additional labels for neon-proxy pods
 podLabels:

.github/helm-values/production.proxy-scram.yaml

@@ -4,6 +4,8 @@ settings:
   domain: "*.cloud.neon.tech"
   sentryEnvironment: "production"
   wssPort: 8443
+  metricCollectionEndpoint: "http://console-release.local/billing/api/v1/usage_events"
+  metricCollectionInterval: "10min"
 
 podLabels:
   zenith_service: proxy-scram

.github/workflows/benchmarking.yml

@@ -489,3 +489,108 @@ jobs:
         slack-message: "Periodic TPC-H perf testing ${{ matrix.platform }}: ${{ job.status }}\n${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
       env:
         SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
+
+  user-examples-compare:
+    if: success() || failure()
+    needs: [ tpch-compare ]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        # neon-captest-prefetch: We have pre-created projects with prefetch enabled
+        # rds-aurora: Aurora Postgres Serverless v2 with autoscaling from 0.5 to 2 ACUs
+        # rds-postgres: RDS Postgres db.m5.large instance (2 vCPU, 8 GiB) with gp3 EBS storage
+        platform: [ neon-captest-prefetch, rds-postgres, rds-aurora ]
+
+    env:
+      POSTGRES_DISTRIB_DIR: /tmp/neon/pg_install
+      DEFAULT_PG_VERSION: 14
+      TEST_OUTPUT: /tmp/test_output
+      BUILD_TYPE: remote
+      SAVE_PERF_REPORT: ${{ github.event.inputs.save_perf_report || ( github.ref == 'refs/heads/main' ) }}
+      PLATFORM: ${{ matrix.platform }}
+
+    runs-on: [ self-hosted, us-east-2, x64 ]
+    container:
+      image: 369495373322.dkr.ecr.eu-central-1.amazonaws.com/rust:pinned
+      options: --init
+
+    timeout-minutes: 360 # 6h
+
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Download Neon artifact
+      uses: ./.github/actions/download
+      with:
+        name: neon-${{ runner.os }}-release-artifact
+        path: /tmp/neon/
+        prefix: latest
+
+    - name: Add Postgres binaries to PATH
+      run: |
+        ${POSTGRES_DISTRIB_DIR}/v${DEFAULT_PG_VERSION}/bin/pgbench --version
+        echo "${POSTGRES_DISTRIB_DIR}/v${DEFAULT_PG_VERSION}/bin" >> $GITHUB_PATH
+
+    - name: Set up Connection String
+      id: set-up-connstr
+      run: |
+        case "${PLATFORM}" in
+          neon-captest-prefetch)
+            CONNSTR=${{ secrets.BENCHMARK_USER_EXAMPLE_CAPTEST_CONNSTR }}
+            ;;
+          rds-aurora)
+            CONNSTR=${{ secrets.BENCHMARK_USER_EXAMPLE_RDS_AURORA_CONNSTR }}
+            ;;
+          rds-postgres)
+            CONNSTR=${{ secrets.BENCHMARK_USER_EXAMPLE_RDS_POSTGRES_CONNSTR }}
+            ;;
+          *)
+            echo 2>&1 "Unknown PLATFORM=${PLATFORM}. Allowed only 'neon-captest-prefetch', 'rds-aurora', or 'rds-postgres'"
+            exit 1
+            ;;
+        esac
+
+        echo "connstr=${CONNSTR}" >> $GITHUB_OUTPUT
+
+        psql ${CONNSTR} -c "SELECT version();"
+
+    - name: Set database options
+      if: matrix.platform == 'neon-captest-prefetch'
+      run: |
+        DB_NAME=$(psql ${BENCHMARK_CONNSTR} --no-align --quiet -t -c "SELECT current_database()")
+
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET enable_seqscan_prefetch=on"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET effective_io_concurrency=32"
+        psql ${BENCHMARK_CONNSTR} -c "ALTER DATABASE ${DB_NAME} SET maintenance_io_concurrency=32"
+      env:
+        BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
+
+    - name: Run user examples
+      uses: ./.github/actions/run-python-test-set
+      with:
+        build_type: ${{ env.BUILD_TYPE }}
+        test_selection: performance/test_perf_olap.py
+        run_in_parallel: false
+        save_perf_report: ${{ env.SAVE_PERF_REPORT }}
+        extra_params: -m remote_cluster --timeout 21600 -k test_user_examples
+      env:
+        VIP_VAP_ACCESS_TOKEN: "${{ secrets.VIP_VAP_ACCESS_TOKEN }}"
+        PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
+        BENCHMARK_CONNSTR: ${{ steps.set-up-connstr.outputs.connstr }}
+
+    - name: Create Allure report
+      if: success() || failure()
+      uses: ./.github/actions/allure-report
+      with:
+        action: generate
+        build_type: ${{ env.BUILD_TYPE }}
+
+    - name: Post to a Slack channel
+      if: ${{ github.event.schedule && failure() }}
+      uses: slackapi/slack-github-action@v1
+      with:
+        channel-id: "C033QLM5P7D" # dev-staging-stream
+        slack-message: "Periodic TPC-H perf testing ${{ matrix.platform }}: ${{ job.status }}\n${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
+      env:
+        SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}

.github/workflows/build_and_test.yml

@@ -595,6 +595,8 @@ jobs:
     defaults:
       run:
         shell: sh -eu {0}
+    env:
+      VM_INFORMANT_VERSION: 0.1.1
 
     steps:
       - name: Downloading latest vm-builder
@@ -606,9 +608,22 @@ jobs:
         run: |
           docker pull 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}}
 
+      - name: Downloading VM informant version ${{ env.VM_INFORMANT_VERSION }}
+        run: |
+          curl -fL https://github.com/neondatabase/autoscaling/releases/download/${{ env.VM_INFORMANT_VERSION }}/vm-informant -o vm-informant
+          chmod +x vm-informant
+
+      - name: Adding VM informant to compute-node image
+        run: |
+          ID=$(docker create 369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}})
+          docker cp vm-informant $ID:/bin/vm-informant
+          docker commit $ID temp-vm-compute-node
+          docker rm -f $ID
+
       - name: Build vm image
         run: |
-          ./vm-builder -src=369495373322.dkr.ecr.eu-central-1.amazonaws.com/compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}} -dst=369495373322.dkr.ecr.eu-central-1.amazonaws.com/vm-compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}}
+          # note: as of 2023-01-12, vm-builder requires a trailing ":latest" for local images
+          ./vm-builder -src=temp-vm-compute-node:latest -dst=369495373322.dkr.ecr.eu-central-1.amazonaws.com/vm-compute-node-${{ matrix.version }}:${{needs.tag.outputs.build-tag}}
 
       - name: Pushing vm-compute-node image
         run: |

.gitignore

@@ -1,7 +1,5 @@
 /pg_install
 /target
-/tmp_check
-/tmp_check_cli
 __pycache__/
 test_output/
 .vscode

Cargo.toml

@@ -10,6 +10,143 @@ members = [
     "libs/*",
 ]
 
+[workspace.package]
+edition = "2021"
+license = "Apache-2.0"
+
+## All dependency versions, used in the project
+[workspace.dependencies]
+anyhow = { version = "1.0", features = ["backtrace"] }
+async-stream = "0.3"
+async-trait = "0.1"
+atty = "0.2.14"
+aws-config = { version = "0.51.0", default-features = false, features=["rustls"] }
+aws-sdk-s3 = "0.21.0"
+aws-smithy-http = "0.51.0"
+aws-types = "0.51.0"
+base64 = "0.13.0"
+bincode = "1.3"
+bindgen = "0.61"
+bstr = "1.0"
+byteorder = "1.4"
+bytes = "1.0"
+chrono = { version = "0.4", default-features = false, features = ["clock"] }
+clap = "4.0"
+close_fds = "0.3.2"
+comfy-table = "6.1"
+const_format = "0.2"
+crc32c = "0.6"
+crossbeam-utils = "0.8.5"
+fail = "0.5.0"
+fs2 = "0.4.3"
+futures = "0.3"
+futures-core = "0.3"
+futures-util = "0.3"
+git-version = "0.3"
+hashbrown = "0.13"
+hex = "0.4"
+hex-literal = "0.3"
+hmac = "0.12.1"
+hostname = "0.3.1"
+humantime = "2.1"
+humantime-serde = "1.1.1"
+hyper = "0.14"
+hyper-tungstenite = "0.9"
+itertools = "0.10"
+jsonwebtoken = "8"
+libc = "0.2"
+md5 = "0.7.0"
+memoffset = "0.8"
+nix = "0.26"
+notify = "5.0.0"
+num-traits = "0.2.15"
+once_cell = "1.13"
+parking_lot = "0.12"
+pin-project-lite = "0.2"
+prometheus = {version = "0.13", default_features=false, features = ["process"]} # removes protobuf dependency
+prost = "0.11"
+rand = "0.8"
+regex = "1.4"
+reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"] }
+routerify = "3"
+rstar = "0.9.3"
+rustls = "0.20"
+rustls-pemfile = "1"
+rustls-split = "0.3"
+scopeguard = "1.1"
+sentry = { version = "0.29", default-features = false, features = ["backtrace", "contexts", "panic", "rustls", "reqwest" ] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1"
+serde_with = "2.0"
+sha2 = "0.10.2"
+signal-hook = "0.3"
+socket2 = "0.4.4"
+strum = "0.24"
+strum_macros = "0.24"
+svg_fmt = "0.4.1"
+tar = "0.4"
+thiserror = "1.0"
+tls-listener = { version = "0.6", features = ["rustls", "hyper-h1"] }
+tokio = { version = "1.17", features = ["macros"] }
+tokio-postgres-rustls = "0.9.0"
+tokio-rustls = "0.23"
+tokio-stream = "0.1"
+tokio-util = { version = "0.7", features = ["io"] }
+toml = "0.5"
+toml_edit = { version = "0.17", features = ["easy"] }
+tonic = {version = "0.8", features = ["tls", "tls-roots"]}
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+url = "2.2"
+uuid = { version = "1.2", features = ["v4", "serde"] }
+walkdir = "2.3.2"
+webpki-roots = "0.22.5"
+x509-parser = "0.14"
+
+## TODO replace this with tracing
+env_logger = "0.10"
+log = "0.4"
+
+## TODO switch when the new release is made
+amplify_num = { git = "https://github.com/rust-amplify/rust-amplify.git", tag = "v4.0.0-beta.1" }
+
+## Libraries from neondatabase/ git forks, ideally with changes to be upstreamed
+postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+tokio-tar = { git = "https://github.com/neondatabase/tokio-tar.git", rev="404df61437de0feef49ba2ccdbdd94eb8ad6e142" }
+
+## Local libraries
+consumption_metrics = { version = "0.1", path = "./libs/consumption_metrics/" }
+metrics = { version = "0.1", path = "./libs/metrics/" }
+pageserver_api = { version = "0.1", path = "./libs/pageserver_api/" }
+postgres_connection = { version = "0.1", path = "./libs/postgres_connection/" }
+postgres_ffi = { version = "0.1", path = "./libs/postgres_ffi/" }
+pq_proto = { version = "0.1", path = "./libs/pq_proto/" }
+remote_storage = { version = "0.1", path = "./libs/remote_storage/" }
+safekeeper_api = { version = "0.1", path = "./libs/safekeeper_api" }
+storage_broker = { version = "0.1", path = "./storage_broker/" } # Note: main broker code is inside the binary crate, so linking with the library shouldn't be heavy.
+tenant_size_model = { version = "0.1", path = "./libs/tenant_size_model/" }
+utils = { version = "0.1", path = "./libs/utils/" }
+
+## Common library dependency
+workspace_hack = { version = "0.1", path = "./workspace_hack/" }
+
+## Build dependencies
+criterion = "0.4"
+rcgen = "0.10"
+rstest = "0.16"
+tempfile = "3.2"
+tonic-build = "0.8"
+
+# This is only needed for proxy's tests.
+# TODO: we should probably fork `tokio-postgres-rustls` instead.
+[patch.crates-io]
+tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+
+################# Binary contents sections
+
 [profile.release]
 # This is useful for profiling and, to some extent, debug.
 # Besides, debug info should not affect the performance.
@@ -70,9 +207,3 @@ inherits = "release"
 debug = false # true = 2 = all symbols, 1 = line only
 opt-level = "z"
 lto = true
-
-
-# This is only needed for proxy's tests.
-# TODO: we should probably fork `tokio-postgres-rustls` instead.
-[patch.crates-io]
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }

Dockerfile.compute-node-v14

@@ -34,7 +34,8 @@ RUN cd postgres && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/bloom.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrowlocks.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/intagg.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/earthdistance.control
 
 #########################################################################################
 #
@@ -62,8 +63,7 @@ RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.1.tar.gz && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_tiger_geocoder.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_sfcgal.control
 
 #########################################################################################
 #

Dockerfile.compute-node-v15

@@ -34,7 +34,8 @@ RUN cd postgres && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/bloom.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgrowlocks.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/intagg.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/pgstattuple.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/earthdistance.control
 
 #########################################################################################
 #
@@ -62,8 +63,7 @@ RUN wget https://download.osgeo.org/postgis/source/postgis-3.3.1.tar.gz && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_tiger_geocoder.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_topology.control && \
     echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer.control && \
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control && \
+    echo 'trusted = true' >> /usr/local/pgsql/share/extension/address_standardizer_data_us.control
-    echo 'trusted = true' >> /usr/local/pgsql/share/extension/postgis_sfcgal.control
 
 #########################################################################################
 #

compute_tools/Cargo.toml

@@ -1,24 +1,25 @@
 [package]
 name = "compute_tools"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
+anyhow.workspace = true
-chrono = { version = "0.4", default-features = false, features = ["clock"] }
+chrono.workspace = true
-clap = "4.0"
+clap.workspace = true
-env_logger = "0.9"
+futures.workspace = true
-futures = "0.3.13"
+hyper = { workspace = true, features = ["full"] }
-hyper = { version = "0.14", features = ["full"] }
+notify.workspace = true
-log = { version = "0.4", features = ["std", "serde"] }
+postgres.workspace = true
-notify = "5.0.0"
+regex.workspace = true
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+serde.workspace = true
-regex = "1"
+serde_json.workspace = true
-serde = { version = "1.0", features = ["derive"] }
+tar.workspace = true
-serde_json = "1"
+tokio = { workspace = true, features = ["rt", "rt-multi-thread"] }
-tar = "0.4"
+tokio-postgres.workspace = true
-tokio = { version = "1.17", features = ["macros", "rt", "rt-multi-thread"] }
+tracing.workspace = true
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+tracing-subscriber.workspace = true
-url = "2.2.2"
+url.workspace = true
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
+
+workspace_hack.workspace = true

compute_tools/README.md

@@ -19,6 +19,10 @@ Also `compute_ctl` spawns two separate service threads:
 - `http-endpoint` runs a Hyper HTTP API server, which serves readiness and the
   last activity requests.
 
+If the `vm-informant` binary is present at `/bin/vm-informant`, it will also be started. For VM
+compute nodes, `vm-informant` communicates with the VM autoscaling system. It coordinates
+downscaling and (eventually) will request immediate upscaling under resource pressure.
+
 Usage example:
 ```sh
 compute_ctl -D /var/db/postgres/compute \

compute_tools/src/bin/compute_ctl.rs

@@ -18,6 +18,10 @@
 //! - `http-endpoint` runs a Hyper HTTP API server, which serves readiness and the
 //!   last activity requests.
 //!
+//! If the `vm-informant` binary is present at `/bin/vm-informant`, it will also be started. For VM
+//! compute nodes, `vm-informant` communicates with the VM autoscaling system. It coordinates
+//! downscaling and (eventually) will request immediate upscaling under resource pressure.
+//!
 //! Usage example:
 //! ```sh
 //! compute_ctl -D /var/db/postgres/compute \
@@ -36,10 +40,11 @@ use std::{thread, time::Duration};
 use anyhow::{Context, Result};
 use chrono::Utc;
 use clap::Arg;
-use log::{error, info};
+use tracing::{error, info};
 
 use compute_tools::compute::{ComputeMetrics, ComputeNode, ComputeState, ComputeStatus};
 use compute_tools::http::api::launch_http_server;
+use compute_tools::informant::spawn_vm_informant_if_present;
 use compute_tools::logger::*;
 use compute_tools::monitor::launch_monitor;
 use compute_tools::params::*;
@@ -48,7 +53,6 @@ use compute_tools::spec::*;
 use url::Url;
 
 fn main() -> Result<()> {
-    // TODO: re-use `utils::logging` later
     init_logger(DEFAULT_LOG_LEVEL)?;
 
     let matches = cli().get_matches();
@@ -114,30 +118,48 @@ fn main() -> Result<()> {
     // requests, while configuration is still in progress.
     let _http_handle = launch_http_server(&compute).expect("cannot launch http endpoint thread");
     let _monitor_handle = launch_monitor(&compute).expect("cannot launch compute monitor thread");
+    // Also spawn the thread responsible for handling the VM informant -- if it's present
+    let _vm_informant_handle = spawn_vm_informant_if_present().expect("cannot launch VM informant");
 
-    // Run compute (Postgres) and hang waiting on it.
+    // Start Postgres
-    match compute.prepare_and_run() {
+    let mut delay_exit = false;
-        Ok(ec) => {
+    let mut exit_code = None;
-            let code = ec.code().unwrap_or(1);
+    let pg = match compute.start_compute() {
-            info!("Postgres exited with code {}, shutting down", code);
+        Ok(pg) => Some(pg),
-            exit(code)
+        Err(err) => {
-        }
+            error!("could not start the compute node: {:?}", err);
-        Err(error) => {
-            error!("could not start the compute node: {:?}", error);
-
             let mut state = compute.state.write().unwrap();
-            state.error = Some(format!("{:?}", error));
+            state.error = Some(format!("{:?}", err));
             state.status = ComputeStatus::Failed;
             drop(state);
-
+            delay_exit = true;
-            // Keep serving HTTP requests, so the cloud control plane was able to
+            None
-            // get the actual error.
-            info!("giving control plane 30s to collect the error before shutdown");
-            thread::sleep(Duration::from_secs(30));
-            info!("shutting down");
-            Err(error)
         }
+    };
+
+    // Wait for the child Postgres process forever. In this state Ctrl+C will
+    // propagate to Postgres and it will be shut down as well.
+    if let Some(mut pg) = pg {
+        let ecode = pg
+            .wait()
+            .expect("failed to start waiting on Postgres process");
+        info!("Postgres exited with code {}, shutting down", ecode);
+        exit_code = ecode.code()
     }
+
+    if let Err(err) = compute.check_for_core_dumps() {
+        error!("error while checking for core dumps: {err:?}");
+    }
+
+    // If launch failed, keep serving HTTP requests for a while, so the cloud
+    // control plane can get the actual error.
+    if delay_exit {
+        info!("giving control plane 30s to collect the error before shutdown");
+        thread::sleep(Duration::from_secs(30));
+        info!("shutting down");
+    }
+
+    exit(exit_code.unwrap_or(1))
 }
 
 fn cli() -> clap::Command {

compute_tools/src/checker.rs

@@ -1,10 +1,11 @@
 use anyhow::{anyhow, Result};
-use log::error;
 use postgres::Client;
 use tokio_postgres::NoTls;
+use tracing::{error, instrument};
 
 use crate::compute::ComputeNode;
 
+#[instrument(skip_all)]
 pub fn create_writability_check_data(client: &mut Client) -> Result<()> {
     let query = "
     CREATE TABLE IF NOT EXISTS health_check (
@@ -21,6 +22,7 @@ pub fn create_writability_check_data(client: &mut Client) -> Result<()> {
     Ok(())
 }
 
+#[instrument(skip_all)]
 pub async fn check_writability(compute: &ComputeNode) -> Result<()> {
     let (client, connection) = tokio_postgres::connect(compute.connstr.as_str(), NoTls).await?;
     if client.is_closed() {

compute_tools/src/compute.rs

@@ -17,15 +17,15 @@
 use std::fs;
 use std::os::unix::fs::PermissionsExt;
 use std::path::Path;
-use std::process::{Command, ExitStatus, Stdio};
+use std::process::{Command, Stdio};
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::RwLock;
 
 use anyhow::{Context, Result};
 use chrono::{DateTime, Utc};
-use log::{info, warn};
 use postgres::{Client, NoTls};
 use serde::{Serialize, Serializer};
+use tracing::{info, instrument, warn};
 
 use crate::checker::create_writability_check_data;
 use crate::config;
@@ -121,6 +121,7 @@ impl ComputeNode {
 
     // Get basebackup from the libpq connection to pageserver using `connstr` and
     // unarchive it to `pgdata` directory overriding all its previous content.
+    #[instrument(skip(self))]
     fn get_basebackup(&self, lsn: &str) -> Result<()> {
         let start_time = Utc::now();
 
@@ -154,6 +155,7 @@ impl ComputeNode {
 
     // Run `postgres` in a special mode with `--sync-safekeepers` argument
     // and return the reported LSN back to the caller.
+    #[instrument(skip(self))]
     fn sync_safekeepers(&self) -> Result<String> {
         let start_time = Utc::now();
 
@@ -196,6 +198,7 @@ impl ComputeNode {
 
     /// Do all the preparations like PGDATA directory creation, configuration,
     /// safekeepers sync, basebackup, etc.
+    #[instrument(skip(self))]
     pub fn prepare_pgdata(&self) -> Result<()> {
         let spec = &self.spec;
         let pgdata_path = Path::new(&self.pgdata);
@@ -229,9 +232,8 @@ impl ComputeNode {
 
     /// Start Postgres as a child process and manage DBs/roles.
     /// After that this will hang waiting on the postmaster process to exit.
-    pub fn run(&self) -> Result<ExitStatus> {
+    #[instrument(skip(self))]
-        let start_time = Utc::now();
+    pub fn start_postgres(&self) -> Result<std::process::Child> {
-
         let pgdata_path = Path::new(&self.pgdata);
 
         // Run postgres as a child process.
@@ -242,6 +244,11 @@ impl ComputeNode {
 
         wait_for_postgres(&mut pg, pgdata_path)?;
 
+        Ok(pg)
+    }
+
+    #[instrument(skip(self))]
+    pub fn apply_config(&self) -> Result<()> {
         // If connection fails,
         // it may be the old node with `zenith_admin` superuser.
         //
@@ -279,8 +286,34 @@ impl ComputeNode {
 
         // 'Close' connection
         drop(client);
-        let startup_end_time = Utc::now();
 
+        info!(
+            "finished configuration of compute for project {}",
+            self.spec.cluster.cluster_id
+        );
+
+        Ok(())
+    }
+
+    #[instrument(skip(self))]
+    pub fn start_compute(&self) -> Result<std::process::Child> {
+        info!(
+            "starting compute for project {}, operation {}, tenant {}, timeline {}",
+            self.spec.cluster.cluster_id,
+            self.spec.operation_uuid.as_ref().unwrap(),
+            self.tenant,
+            self.timeline,
+        );
+
+        self.prepare_pgdata()?;
+
+        let start_time = Utc::now();
+
+        let pg = self.start_postgres()?;
+
+        self.apply_config()?;
+
+        let startup_end_time = Utc::now();
         self.metrics.config_ms.store(
             startup_end_time
                 .signed_duration_since(start_time)
@@ -300,34 +333,7 @@ impl ComputeNode {
 
         self.set_status(ComputeStatus::Running);
 
-        info!(
+        Ok(pg)
-            "finished configuration of compute for project {}",
-            self.spec.cluster.cluster_id
-        );
-
-        // Wait for child Postgres process basically forever. In this state Ctrl+C
-        // will propagate to Postgres and it will be shut down as well.
-        let ecode = pg
-            .wait()
-            .expect("failed to start waiting on Postgres process");
-
-        self.check_for_core_dumps()
-            .expect("failed to check for core dumps");
-
-        Ok(ecode)
-    }
-
-    pub fn prepare_and_run(&self) -> Result<ExitStatus> {
-        info!(
-            "starting compute for project {}, operation {}, tenant {}, timeline {}",
-            self.spec.cluster.cluster_id,
-            self.spec.operation_uuid.as_ref().unwrap(),
-            self.tenant,
-            self.timeline,
-        );
-
-        self.prepare_pgdata()?;
-        self.run()
     }
 
     // Look for core dumps and collect backtraces.
@@ -340,7 +346,7 @@ impl ComputeNode {
     //
     // Use that as a default location and pattern, except macos where core dumps are written
     // to /cores/ directory by default.
-    fn check_for_core_dumps(&self) -> Result<()> {
+    pub fn check_for_core_dumps(&self) -> Result<()> {
         let core_dump_dir = match std::env::consts::OS {
             "macos" => Path::new("/cores/"),
             _ => Path::new(&self.pgdata),

compute_tools/src/http/api.rs

@@ -6,8 +6,8 @@ use std::thread;
 use anyhow::Result;
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Method, Request, Response, Server, StatusCode};
-use log::{error, info};
 use serde_json;
+use tracing::{error, info};
 
 use crate::compute::ComputeNode;
 

compute_tools/src/informant.rs

@@ -0,0 +1,50 @@
+use std::path::Path;
+use std::process;
+use std::thread;
+use std::time::Duration;
+use tracing::{info, warn};
+
+use anyhow::{Context, Result};
+
+const VM_INFORMANT_PATH: &str = "/bin/vm-informant";
+const RESTART_INFORMANT_AFTER_MILLIS: u64 = 5000;
+
+/// Launch a thread to start the VM informant if it's present (and restart, on failure)
+pub fn spawn_vm_informant_if_present() -> Result<Option<thread::JoinHandle<()>>> {
+    let exists = Path::new(VM_INFORMANT_PATH)
+        .try_exists()
+        .context("could not check if path exists")?;
+
+    if !exists {
+        return Ok(None);
+    }
+
+    Ok(Some(
+        thread::Builder::new()
+            .name("run-vm-informant".into())
+            .spawn(move || run_informant())?,
+    ))
+}
+
+fn run_informant() -> ! {
+    let restart_wait = Duration::from_millis(RESTART_INFORMANT_AFTER_MILLIS);
+
+    info!("starting VM informant");
+
+    loop {
+        let mut cmd = process::Command::new(VM_INFORMANT_PATH);
+        // Block on subprocess:
+        let result = cmd.status();
+
+        match result {
+            Err(e) => warn!("failed to run VM informant at {VM_INFORMANT_PATH:?}: {e}"),
+            Ok(status) if !status.success() => {
+                warn!("{VM_INFORMANT_PATH} exited with code {status:?}, retrying")
+            }
+            Ok(_) => info!("{VM_INFORMANT_PATH} ended gracefully (unexpectedly). Retrying"),
+        }
+
+        // Wait before retrying
+        thread::sleep(restart_wait);
+    }
+}

compute_tools/src/lib.rs

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

compute_tools/src/logger.rs

@@ -1,42 +1,20 @@
-use std::io::Write;
-
 use anyhow::Result;
-use chrono::Utc;
+use tracing_subscriber::layer::SubscriberExt;
-use env_logger::{Builder, Env};
+use tracing_subscriber::prelude::*;
-
-macro_rules! info_println {
-    ($($tts:tt)*) => {
-        if log_enabled!(Level::Info) {
-            println!($($tts)*);
-        }
-    }
-}
-
-macro_rules! info_print {
-    ($($tts:tt)*) => {
-        if log_enabled!(Level::Info) {
-            print!($($tts)*);
-        }
-    }
-}
 
 /// Initialize `env_logger` using either `default_level` or
 /// `RUST_LOG` environment variable as default log level.
 pub fn init_logger(default_level: &str) -> Result<()> {
-    let env = Env::default().filter_or("RUST_LOG", default_level);
+    let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
+        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(default_level));
 
-    Builder::from_env(env)
+    let fmt_layer = tracing_subscriber::fmt::layer()
-        .format(|buf, record| {
+        .with_target(false)
-            let thread_handle = std::thread::current();
+        .with_writer(std::io::stderr);
-            writeln!(
+
-                buf,
+    tracing_subscriber::registry()
-                "{} [{}] {}: {}",
+        .with(env_filter)
-                Utc::now().format("%Y-%m-%d %H:%M:%S%.3f %Z"),
+        .with(fmt_layer)
-                thread_handle.name().unwrap_or("main"),
-                record.level(),
-                record.args()
-            )
-        })
         .init();
 
     Ok(())

compute_tools/src/monitor.rs

@@ -3,8 +3,8 @@ use std::{thread, time};
 
 use anyhow::Result;
 use chrono::{DateTime, Utc};
-use log::{debug, info};
 use postgres::{Client, NoTls};
+use tracing::{debug, info};
 
 use crate::compute::ComputeNode;
 

compute_tools/src/pg_helpers.rs

@@ -11,6 +11,7 @@ use anyhow::{bail, Result};
 use notify::{RecursiveMode, Watcher};
 use postgres::{Client, Transaction};
 use serde::Deserialize;
+use tracing::{debug, instrument};
 
 const POSTGRES_WAIT_TIMEOUT: Duration = Duration::from_millis(60 * 1000); // milliseconds
 
@@ -229,6 +230,7 @@ pub fn get_existing_dbs(client: &mut Client) -> Result<Vec<Database>> {
 /// Wait for Postgres to become ready to accept connections. It's ready to
 /// accept connections when the state-field in `pgdata/postmaster.pid` says
 /// 'ready'.
+#[instrument(skip(pg))]
 pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
     let pid_path = pgdata.join("postmaster.pid");
 
@@ -287,18 +289,18 @@ pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
         }
 
         let res = rx.recv_timeout(Duration::from_millis(100));
-        log::debug!("woken up by notify: {res:?}");
+        debug!("woken up by notify: {res:?}");
         // If there are multiple events in the channel already, we only need to be
         // check once. Swallow the extra events before we go ahead to check the
         // pid file.
         while let Ok(res) = rx.try_recv() {
-            log::debug!("swallowing extra event: {res:?}");
+            debug!("swallowing extra event: {res:?}");
         }
 
         // Check that we can open pid file first.
         if let Ok(file) = File::open(&pid_path) {
             if !postmaster_pid_seen {
-                log::debug!("postmaster.pid appeared");
+                debug!("postmaster.pid appeared");
                 watcher
                     .unwatch(pgdata)
                     .expect("Failed to remove pgdata dir watch");
@@ -314,7 +316,7 @@ pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
             // Pid file could be there and we could read it, but it could be empty, for example.
             if let Some(Ok(line)) = last_line {
                 let status = line.trim();
-                log::debug!("last line of postmaster.pid: {status:?}");
+                debug!("last line of postmaster.pid: {status:?}");
 
                 // Now Postgres is ready to accept connections
                 if status == "ready" {
@@ -330,7 +332,7 @@ pub fn wait_for_postgres(pg: &mut Child, pgdata: &Path) -> Result<()> {
         }
     }
 
-    log::info!("PostgreSQL is now running, continuing to configure it");
+    tracing::info!("PostgreSQL is now running, continuing to configure it");
 
     Ok(())
 }

compute_tools/src/spec.rs

@@ -1,12 +1,11 @@
 use std::path::Path;
 use std::str::FromStr;
-use std::time::Instant;
 
 use anyhow::Result;
-use log::{info, log_enabled, warn, Level};
 use postgres::config::Config;
 use postgres::{Client, NoTls};
 use serde::Deserialize;
+use tracing::{info, info_span, instrument, span_enabled, warn, Level};
 
 use crate::compute::ComputeNode;
 use crate::config;
@@ -80,23 +79,25 @@ pub fn update_pg_hba(pgdata_path: &Path) -> Result<()> {
 
 /// Given a cluster spec json and open transaction it handles roles creation,
 /// deletion and update.
+#[instrument(skip_all)]
 pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     let mut xact = client.transaction()?;
     let existing_roles: Vec<Role> = get_existing_roles(&mut xact)?;
 
     // Print a list of existing Postgres roles (only in debug mode)
-    info!("postgres roles:");
+    if span_enabled!(Level::INFO) {
-    for r in &existing_roles {
+        info!("postgres roles:");
-        info_println!(
+        for r in &existing_roles {
-            "{} - {}:{}",
+            info!(
-            " ".repeat(27 + 5),
+                "    - {}:{}",
-            r.name,
+                r.name,
-            if r.encrypted_password.is_some() {
+                if r.encrypted_password.is_some() {
-                "[FILTERED]"
+                    "[FILTERED]"
-            } else {
+                } else {
-                "(null)"
+                    "(null)"
-            }
+                }
-        );
+            );
+        }
     }
 
     // Process delta operations first
@@ -137,58 +138,68 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     info!("cluster spec roles:");
     for role in &spec.cluster.roles {
         let name = &role.name;
-
-        info_print!(
-            "{} - {}:{}",
-            " ".repeat(27 + 5),
-            name,
-            if role.encrypted_password.is_some() {
-                "[FILTERED]"
-            } else {
-                "(null)"
-            }
-        );
-
         // XXX: with a limited number of roles it is fine, but consider making it a HashMap
         let pg_role = existing_roles.iter().find(|r| r.name == *name);
 
-        if let Some(r) = pg_role {
+        enum RoleAction {
-            let mut update_role = false;
+            None,
-
+            Update,
+            Create,
+        }
+        let action = if let Some(r) = pg_role {
             if (r.encrypted_password.is_none() && role.encrypted_password.is_some())
                 || (r.encrypted_password.is_some() && role.encrypted_password.is_none())
             {
-                update_role = true;
+                RoleAction::Update
             } else if let Some(pg_pwd) = &r.encrypted_password {
                 // Check whether password changed or not (trim 'md5:' prefix first)
-                update_role = pg_pwd[3..] != *role.encrypted_password.as_ref().unwrap();
+                if pg_pwd[3..] != *role.encrypted_password.as_ref().unwrap() {
+                    RoleAction::Update
+                } else {
+                    RoleAction::None
+                }
+            } else {
+                RoleAction::None
             }
+        } else {
+            RoleAction::Create
+        };
 
-            if update_role {
+        match action {
+            RoleAction::None => {}
+            RoleAction::Update => {
                 let mut query: String = format!("ALTER ROLE {} ", name.pg_quote());
-                info_print!(" -> update");
-
                 query.push_str(&role.to_pg_options());
                 xact.execute(query.as_str(), &[])?;
             }
-        } else {
+            RoleAction::Create => {
-            info!("role name: '{}'", &name);
+                let mut query: String = format!("CREATE ROLE {} ", name.pg_quote());
-            let mut query: String = format!("CREATE ROLE {} ", name.pg_quote());
+                info!("role create query: '{}'", &query);
-            info!("role create query: '{}'", &query);
+                query.push_str(&role.to_pg_options());
-            info_print!(" -> create");
+                xact.execute(query.as_str(), &[])?;
 
-            query.push_str(&role.to_pg_options());
+                let grant_query = format!(
-            xact.execute(query.as_str(), &[])?;
+                    "GRANT pg_read_all_data, pg_write_all_data TO {}",
-
+                    name.pg_quote()
-            let grant_query = format!(
+                );
-                "GRANT pg_read_all_data, pg_write_all_data TO {}",
+                xact.execute(grant_query.as_str(), &[])?;
-                name.pg_quote()
+                info!("role grant query: '{}'", &grant_query);
-            );
+            }
-            xact.execute(grant_query.as_str(), &[])?;
-            info!("role grant query: '{}'", &grant_query);
         }
 
-        info_print!("\n");
+        if span_enabled!(Level::INFO) {
+            let pwd = if role.encrypted_password.is_some() {
+                "[FILTERED]"
+            } else {
+                "(null)"
+            };
+            let action_str = match action {
+                RoleAction::None => "",
+                RoleAction::Create => " -> create",
+                RoleAction::Update => " -> update",
+            };
+            info!("   - {}:{}{}", name, pwd, action_str);
+        }
     }
 
     xact.commit()?;
@@ -197,6 +208,7 @@ pub fn handle_roles(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
 }
 
 /// Reassign all dependent objects and delete requested roles.
+#[instrument(skip_all)]
 pub fn handle_role_deletions(node: &ComputeNode, client: &mut Client) -> Result<()> {
     if let Some(ops) = &node.spec.delta_operations {
         // First, reassign all dependent objects to db owners.
@@ -261,13 +273,16 @@ fn reassign_owned_objects(node: &ComputeNode, role_name: &PgIdent) -> Result<()>
 /// like `CREATE DATABASE` and `DROP DATABASE` do not support it. Statement-level
 /// atomicity should be enough here due to the order of operations and various checks,
 /// which together provide us idempotency.
+#[instrument(skip_all)]
 pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     let existing_dbs: Vec<Database> = get_existing_dbs(client)?;
 
     // Print a list of existing Postgres databases (only in debug mode)
-    info!("postgres databases:");
+    if span_enabled!(Level::INFO) {
-    for r in &existing_dbs {
+        info!("postgres databases:");
-        info_println!("{} - {}:{}", " ".repeat(27 + 5), r.name, r.owner);
+        for r in &existing_dbs {
+            info!("    {}:{}", r.name, r.owner);
+        }
     }
 
     // Process delta operations first
@@ -310,13 +325,15 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
     for db in &spec.cluster.databases {
         let name = &db.name;
 
-        info_print!("{} - {}:{}", " ".repeat(27 + 5), db.name, db.owner);
-
         // XXX: with a limited number of databases it is fine, but consider making it a HashMap
         let pg_db = existing_dbs.iter().find(|r| r.name == *name);
 
-        let start_time = Instant::now();
+        enum DatabaseAction {
-        if let Some(r) = pg_db {
+            None,
+            Update,
+            Create,
+        }
+        let action = if let Some(r) = pg_db {
             // XXX: db owner name is returned as quoted string from Postgres,
             // when quoting is needed.
             let new_owner = if r.owner.starts_with('"') {
@@ -326,29 +343,42 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
             };
 
             if new_owner != r.owner {
+                // Update the owner
+                DatabaseAction::Update
+            } else {
+                DatabaseAction::None
+            }
+        } else {
+            DatabaseAction::Create
+        };
+
+        match action {
+            DatabaseAction::None => {}
+            DatabaseAction::Update => {
                 let query: String = format!(
                     "ALTER DATABASE {} OWNER TO {}",
                     name.pg_quote(),
                     db.owner.pg_quote()
                 );
-                info_print!(" -> update");
+                let _ = info_span!("executing", query).entered();
-
                 client.execute(query.as_str(), &[])?;
-                let elapsed = start_time.elapsed().as_millis();
-                info_print!(" ({} ms)", elapsed);
             }
-        } else {
+            DatabaseAction::Create => {
-            let mut query: String = format!("CREATE DATABASE {} ", name.pg_quote());
+                let mut query: String = format!("CREATE DATABASE {} ", name.pg_quote());
-            info_print!(" -> create");
+                query.push_str(&db.to_pg_options());
+                let _ = info_span!("executing", query).entered();
+                client.execute(query.as_str(), &[])?;
+            }
+        };
 
-            query.push_str(&db.to_pg_options());
+        if span_enabled!(Level::INFO) {
-            client.execute(query.as_str(), &[])?;
+            let action_str = match action {
-
+                DatabaseAction::None => "",
-            let elapsed = start_time.elapsed().as_millis();
+                DatabaseAction::Create => " -> create",
-            info_print!(" ({} ms)", elapsed);
+                DatabaseAction::Update => " -> update",
+            };
+            info!("   - {}:{}{}", db.name, db.owner, action_str);
         }
-
-        info_print!("\n");
     }
 
     Ok(())
@@ -356,6 +386,7 @@ pub fn handle_databases(spec: &ComputeSpec, client: &mut Client) -> Result<()> {
 
 /// Grant CREATE ON DATABASE to the database owner and do some other alters and grants
 /// to allow users creating trusted extensions and re-creating `public` schema, for example.
+#[instrument(skip_all)]
 pub fn handle_grants(node: &ComputeNode, client: &mut Client) -> Result<()> {
     let spec = &node.spec;
 

control_plane/.gitignore

@@ -1 +0,0 @@
-tmp_check/

control_plane/Cargo.toml

@@ -1,32 +1,31 @@
 [package]
 name = "control_plane"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
+anyhow.workspace = true
-clap = "4.0"
+clap.workspace = true
-comfy-table = "6.1"
+comfy-table.workspace = true
-git-version = "0.3.5"
+git-version.workspace = true
-nix = "0.25"
+nix.workspace = true
-once_cell = "1.13.0"
+once_cell.workspace = true
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev = "43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres.workspace = true
-regex = "1"
+regex.workspace = true
-reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
+reqwest = { workspace = true, features = ["blocking", "json"] }
-serde = { version = "1.0", features = ["derive"] }
+serde.workspace = true
-serde_with = "2.0"
+serde_with.workspace = true
-tar = "0.4.38"
+tar.workspace = true
-thiserror = "1"
+thiserror.workspace = true
-toml = "0.5"
+toml.workspace = true
-url = "2.2.2"
+url.workspace = true
-
 # Note: Do not directly depend on pageserver or safekeeper; use pageserver_api or safekeeper_api
 # instead, so that recompile times are better.
-pageserver_api = { path = "../libs/pageserver_api" }
+pageserver_api.workspace = true
-postgres_connection = { path = "../libs/postgres_connection" }
+safekeeper_api.workspace = true
-safekeeper_api = { path = "../libs/safekeeper_api" }
+postgres_connection.workspace = true
-# Note: main broker code is inside the binary crate, so linking with the library shouldn't be heavy.
+storage_broker.workspace = true
-storage_broker = { version = "0.1", path = "../storage_broker" }
+utils.workspace = true
-utils = { path = "../libs/utils" }
+
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
+workspace_hack.workspace = true

deny.toml

@@ -52,7 +52,7 @@ name = "ring"
 version = "*"
 expression = "MIT AND ISC AND OpenSSL"
 license-files = [
-    { path = "LICENSE", hash = 0xbd0eed23 },
+    { path = "LICENSE", hash = 0xbd0eed23 }
 ]
 
 [licenses.private]

docs/consumption_metrics.md

@@ -0,0 +1,115 @@
+### Overview
+Pageserver and proxy periodically collect consumption metrics and push them to a HTTP endpoint.
+
+This doc describes current implementation details.
+For design details see [the RFC](./rfcs/021-metering.md) and [the discussion on Github](https://github.com/neondatabase/neon/pull/2884).
+
+- The metrics are collected in a separate thread, and the collection interval and endpoint are configurable.
+
+- Metrics are cached, so that we don't send unchanged metrics on every iteration.
+
+- Metrics are sent in batches of 1000 (see CHUNK_SIZE const) metrics max with no particular grouping guarantees.
+
+batch format is
+```json
+
+{ "events" : [metric1, metric2, ...]]}
+
+```
+See metric format examples below.
+
+- All metrics values are in bytes, unless otherwise specified.
+
+- Currently no retries are implemented.
+
+### Pageserver metrics
+
+#### Configuration
+The endpoint and the collection interval are specified in the pageserver config file (or can be passed as command line arguments):
+`metric_collection_endpoint` defaults to None, which means that metric collection is disabled by default.
+`metric_collection_interval` defaults to 10min
+
+#### Metrics
+
+Currently, the following metrics are collected:
+
+- `written_size`
+
+Amount of WAL produced , by a timeline, i.e. last_record_lsn
+This is an absolute, per-timeline metric.
+
+- `resident_size`
+
+Size of all the layer files in the tenant's directory on disk on the pageserver.
+This is an absolute, per-tenant metric.
+
+- `remote_storage_size`
+
+Size of the remote storage (S3) directory.
+This is an absolute, per-tenant metric.
+
+- `timeline_logical_size`
+Logical size of the data in the timeline
+This is an absolute, per-timeline metric.
+
+- `synthetic_storage_size`
+Size of all tenant's branches including WAL
+This is the same metric that `tenant/{tenant_id}/size` endpoint returns.
+This is an absolute, per-tenant metric.
+
+Synthetic storage size is calculated in a separate thread, so it might be slightly outdated.
+
+#### Format example
+
+```json
+{
+"metric": "remote_storage_size",
+"type": "absolute",
+"time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"tenant_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+"timeline_id": "a03ebb4f5922a1c56ff7485cc8854143",
+}
+```
+
+`idempotency_key` is a unique key for each metric, so that we can deduplicate metrics.
+It is a combination of the time, node_id and a random number.
+
+### Proxy consumption metrics
+
+#### Configuration
+The endpoint and the collection interval can be passed as command line arguments for proxy:
+`metric_collection_endpoint` no default, which means that metric collection is disabled by default.
+`metric_collection_interval` no default
+
+#### Metrics
+
+Currently, only one proxy metric is collected:
+
+- `proxy_io_bytes_per_client`
+Outbound traffic per client.
+This is an incremental, per-endpoint metric.
+
+#### Format example
+
+```json
+{
+"metric": "proxy_io_bytes_per_client",
+"type": "incremental",
+"start_time": "2022-12-28T11:07:19.317310284Z",
+"stop_time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"endpoint_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+}
+```
+
+The metric is incremental, so the value is the difference between the current and the previous value.
+If there is no previous value, the value, the value is the current value and the `start_time` equals `stop_time`.
+
+### TODO
+
+- [ ] Handle errors better: currently if one tenant fails to gather metrics, the whole iteration fails and metrics are not sent for any tenant.
+- [ ] Add retries
+- [ ] Tune the interval

docs/rfcs/021-metering.md

@@ -0,0 +1,186 @@
+# Consumption tracking
+
+
+# Goals
+
+This proposal is made with two mostly but not entirely overlapping goals:
+
+* Collect info that is needed for consumption-based billing
+* Cross-check AWS bills
+
+
+# Metrics
+
+There are six metrics to collect:
+
+* CPU time. Wall clock seconds * the current number of cores. We have a fixed ratio of memory to cores, so the current memory size is the function of the number of cores. Measured per each `endpoint`.
+
+* Traffic. In/out traffic on the proxy. Measured per each `endpoint`.
+
+* Written size. Amount of data we write. That is different from both traffic and storage size, as only during the writing we
+
+  a) occupy some disk bandwidth on safekeepers
+
+  b) necessarily cross AZ boundaries delivering WAL to all safekeepers
+
+  Each timeline/branch has at most one writer, so the data is collected per branch.
+
+* Synthetic storage size. That is what is exposed now with pageserver's `/v1/tenant/{}/size`. Looks like now it is per-tenant. (Side note: can we make it per branch to show as branch physical size in UI?)
+
+* Real storage size. That is the size of the tenant directory on the pageservers disk. Per-tenant.
+
+* S3 storage size. That is the size of the tenant data on S3. Per-tenant.
+
+That info should be enough to build an internal model that predicts AWS price (hence tracking `written data` and `real storage size`). As for the billing model we probably can get away with mentioning only `CPU time`, `synthetic storage size`, and `traffic` consumption.
+
+# Services participating in metrics collection
+
+## Proxy
+
+For actual implementation details check `/docs/consumption_metrics.md`
+
+Proxy is the only place that knows about traffic flow, so it tracks it and reports it with quite a small interval, let's say 1 minute. A small interval is needed here since the proxy is stateless, and any restart will reset accumulated consumption. Also proxy should report deltas since the last report, not an absolute value of the counter. Such kind of events is easier to integrate over a period of time to get the amount of traffic during some time interval.
+
+Example event:
+
+```json
+{
+"metric": "proxy_io_bytes_per_client",
+"type": "incremental",
+"start_time": "2022-12-28T11:07:19.317310284Z",
+"stop_time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"endpoint_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+}
+```
+
+Since we report deltas over some period of time, it makes sense to include `event_start_time`/`event_stop_time` where `event_start_time` is the time of the previous report. That will allow us to identify metering gaps better (e.g., failed send/delivery).
+
+When there is no active connection proxy can avoid reporting anything. Also, deltas are additive, so several console instances serving the same user and endpoint can report traffic without coordination.
+
+## Console
+
+The console knows about start/stop events, so it knows the amount of CPU time allocated to each endpoint. It also knows about operation successes and failures and can avoid billing clients after unsuccessful 'suspend' events. The console doesn't know the current compute size within the allowed limits on the endpoint. So with CPU time, we do the following:
+
+* While we don't yet have the autoscaling console can report `cpu time` as the number of seconds since the last `start_compute` event.
+
+* When we have autoscaling, `autoscaler-agent` can report `cpu time`*`compute_units_count` in the same increments as the proxy reports traffic.
+
+Example event:
+
+```json
+{
+    "metric": "effective_compute_seconds",
+    "type": "increment",
+    "endpoint_id": "blazing-warrior-34",
+    "event_start_time": ...,
+    "event_stop_time": ...,
+    "value": 12345454,
+}
+```
+
+I'd also suggest reporting one value, `cpu time`*`compute_units_count`, instead of two separate fields as it makes event schema simpler (it is possible to treat it the same way as traffic) and preserves additivity.
+
+## Pageserver
+
+For actual implementation details check `/docs/consumption_metrics.md`
+
+Pageserver knows / has access to / can calculate the rest of the metrics:
+
+* Written size -- that is basically `last_received_lsn`,
+* Synthetic storage size -- there is a way to calculate it, albeit a costly one,
+* Real storage size -- there is a way to calculate it using a layer map or filesystem,
+* S3 storage size -- can calculate it by S3 API calls
+
+Some of those metrics are expensive to calculate, so the reporting period here is driven mainly by implementation details. We can set it to, for example, once per hour. Not a big deal since the pageserver is stateful, and all metrics can be reported as an absolute value, not increments. At the same time, a smaller reporting period improves UX, so it would be good to have something more real-time.
+
+`written size` is primarily a safekeeper-related metric, but since it is available on both pageserver and safekeeper, we can avoid reporting anything from the safekeeper.
+
+Example event:
+
+```json
+{
+"metric": "remote_storage_size",
+"type": "absolute",
+"time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"tenant_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+"timeline_id": "a03ebb4f5922a1c56ff7485cc8854143",
+}
+```
+
+# Data collection
+
+## Push vs. pull
+
+We already have pull-based Prometheus metrics, so it is tempting to use them here too. However, in our setup, it is hard to tell when some metric changes. For example, garbage collection will constantly free some disk space over a week, even if the project is down for that week. We could also iterate through all existing tenants/branches/endpoints, but that means some amount of code to do that properly and most likely we will end up with some per-metric hacks in the collector to cut out some of the tenants that are surely not changing that metric.
+
+With the push model, it is easier to publish data only about actively changing metrics -- pageserver knows when it performs s3 offloads, garbage collection and starts/stops consuming data from the safekeeper; proxy knows about connected clients; console / autoscaler-agent knows about active cpu time.
+
+Hence, let's go with a push-based model.
+
+## Common bus vs. proxying through the console
+
+We can implement such push systems in a few ways:
+
+a. Each component pushes its metrics to the "common bus", namely segment, Kafka, or something similar. That approach scales well, but it would be harder to test it locally, will introduce new dependencies, we will have to distribute secrets for that connection to all of the components, etc. We would also have to loop back some of the events and their aggregates to the console, as we want to show some that metrics to the user in real-time.
+
+b. Each component can call HTTP `POST` with its events to the console, and the console can forward it to the segment for later integration with metronome / orb / onebill / etc. With that approach, only the console has to speak with segment. Also since that data passes through the console, the console can save the latest metrics values, so there is no need for constant feedback of that events back from the segment.
+
+# Implementation
+
+Each (proxy|pageserver|autoscaler-agent) sends consumption events to the single endpoint in the console:
+
+```json
+POST /usage_events HTTP/1.1
+Content-Type: application/json
+
+[
+{
+"metric": "remote_storage_size",
+"type": "absolute",
+"time": "2022-12-28T11:07:19.317310284Z",
+"idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
+"value": 12345454,
+"tenant_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
+"timeline_id": "a03ebb4f5922a1c56ff7485cc8854143",
+},
+...
+]
+```
+
+![data flow](./images/metering.jpg)
+
+Events could be either:
+* `incremental` -- change in consumption since the previous event or service restart. That is `effective_cpu_seconds`, `traffic_in_bytes`, and `traffic_out_bytes`.
+* `absolute` -- that is the current value of a metric. All of the size-related metrics are absolute.
+
+Each service can post events at its own pace and bundle together data from different tenants/endpoints.
+
+The console algorithm upon receive of events could be the following:
+
+1. Create and send a segment event with the same content (possibly enriching it with tenant/timeline data for endpoint-based events).
+2. Update the latest state of per-tenant and per-endpoint metrics in the database.
+3. Check whether any of that metrics is above the allowed threshold and stop the project if necessary.
+
+Since all the data comes in batches, we can do the batch update to reduce the number of queries in the database. Proxy traffic is probably the most frequent metric, so with batching, we will have extra `number_of_proxies` requests to the database each minute. This is most likely fine for now but will generate many dead tuples in the console database. If that is the case, we can change step 2 to the following:
+
+2.1. Check if there $tenant_$metric / $endpoint_$metric key in Redis
+
+2.2. If no stored value is found and the metric is incremental, then fetch the current value from DWH (which keeps aggregated value for all the events) and publish it.
+
+2.3. Publish a new value (absolute metric) or add an increment to the stored value (incremental metric)
+
+## Consumption watchdog
+
+Since all the data goes through the console, we don't have to run any background thread/coroutines to check whether consumption is within the allowed limits. We only change consumption with `POST /usage_events`, so limit checks could be applied in the same handler.
+
+## Extensibility
+
+If we need to add a new metric (e.g. s3 traffic or something else), the console code should, by default, process it and publish segment event, even if the metric name is unknown to the console.
+
+## Naming & schema
+
+Each metric name should end up with units -- now `_seconds` and `_bytes`, and segment event should always have `tenant_id` and `timeline_id`/`endpoint_id` where applicable.

docs/sourcetree.md

@@ -18,10 +18,6 @@ Intended to be used in integration tests and in CLI tools for local installation
 Documentation of the Neon features and concepts.
 Now it is mostly dev documentation.
 
-`/monitoring`:
-
-TODO
-
 `/pageserver`:
 
 Neon storage service.
@@ -98,6 +94,13 @@ cargo hakari manage-deps
 
 If you don't have hakari installed (`error: no such subcommand: hakari`), install it by running `cargo install cargo-hakari`.
 
+### Checking Rust 3rd-parties
+[Cargo deny](https://embarkstudios.github.io/cargo-deny/index.html) is a cargo plugin that lets us lint project's dependency graph to ensure all dependencies conform to requirements. It detects security issues, matches licenses, and ensures crates only come from trusted sources.
+
+```bash
+cargo deny check
+```
+
 ## Using Python
 Note that Debian/Ubuntu Python packages are stale, as it commonly happens,
 so manual installation of dependencies is not recommended.

libs/consumption_metrics/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "consumption_metrics"
+version = "0.1.0"
+edition = "2021"
+license = "Apache-2.0"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+anyhow = "1.0.68"
+chrono = { version = "0.4", default-features = false, features = ["clock", "serde"] }
+rand = "0.8.3"
+serde = "1.0.152"
+serde_with = "2.1.0"
+utils = { version = "0.1.0", path = "../utils" }
+workspace_hack = { version = "0.1.0", path = "../../workspace_hack" }

libs/consumption_metrics/src/lib.rs

@@ -0,0 +1,50 @@
+//!
+//! Shared code for consumption metics collection
+//!
+use chrono::{DateTime, Utc};
+use rand::Rng;
+use serde::Serialize;
+
+#[derive(Serialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
+#[serde(tag = "type")]
+pub enum EventType {
+    #[serde(rename = "absolute")]
+    Absolute { time: DateTime<Utc> },
+    #[serde(rename = "incremental")]
+    Incremental {
+        start_time: DateTime<Utc>,
+        stop_time: DateTime<Utc>,
+    },
+}
+
+#[derive(Serialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
+pub struct Event<Extra> {
+    #[serde(flatten)]
+    #[serde(rename = "type")]
+    pub kind: EventType,
+
+    pub metric: &'static str,
+    pub idempotency_key: String,
+    pub value: u64,
+
+    #[serde(flatten)]
+    pub extra: Extra,
+}
+
+pub fn idempotency_key(node_id: String) -> String {
+    format!(
+        "{}-{}-{:04}",
+        Utc::now(),
+        node_id,
+        rand::thread_rng().gen_range(0..=9999)
+    )
+}
+
+pub const CHUNK_SIZE: usize = 1000;
+
+// Just a wrapper around a slice of events
+// to serialize it as `{"events" : [ ] }
+#[derive(serde::Serialize)]
+pub struct EventChunk<'a, T> {
+    pub events: &'a [T],
+}

libs/metrics/Cargo.toml

@@ -1,11 +1,12 @@
 [package]
 name = "metrics"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-prometheus = {version = "0.13", default_features=false, features = ["process"]} # removes protobuf dependency
+prometheus.workspace = true
-libc = "0.2"
+libc.workspace = true
-once_cell = "1.13.0"
+once_cell.workspace = true
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+
+workspace_hack.workspace = true

libs/pageserver_api/Cargo.toml

@@ -1,17 +1,17 @@
 [package]
 name = "pageserver_api"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-serde = { version = "1.0", features = ["derive"] }
+serde.workspace = true
-serde_with = "2.0"
+serde_with.workspace = true
-const_format = "0.2.21"
+const_format.workspace = true
-anyhow = { version = "1.0", features = ["backtrace"] }
+anyhow.workspace = true
-bytes = "1.0.1"
+bytes.workspace = true
-byteorder = "1.4.3"
+byteorder.workspace = true
+utils.workspace = true
+postgres_ffi.workspace = true
 
-utils = { path = "../utils" }
+workspace_hack.workspace = true
-postgres_ffi = { path = "../postgres_ffi" }
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }

libs/pageserver_api/src/models.rs

@@ -1,4 +1,4 @@
-use std::num::NonZeroU64;
+use std::num::{NonZeroU64, NonZeroUsize};
 
 use byteorder::{BigEndian, ReadBytesExt};
 use serde::{Deserialize, Serialize};
@@ -210,6 +210,11 @@ pub struct TimelineInfo {
     pub state: TimelineState,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct DownloadRemoteLayersTaskSpawnRequest {
+    pub max_concurrent_downloads: NonZeroUsize,
+}
+
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct DownloadRemoteLayersTaskInfo {
     pub task_id: String,

libs/postgres_connection/Cargo.toml

@@ -1,18 +1,17 @@
 [package]
 name = "postgres_connection"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
-
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-anyhow = "1.0"
+anyhow.workspace = true
-itertools = "0.10.3"
+itertools.workspace = true
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev = "43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres.workspace = true
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+tokio-postgres.workspace = true
-url = "2.2.2"
+url.workspace = true
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+
+workspace_hack.workspace = true
 
 [dev-dependencies]
-once_cell = "1.13.0"
+once_cell.workspace = true

libs/postgres_ffi/Cargo.toml

@@ -1,30 +1,31 @@
 [package]
 name = "postgres_ffi"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-rand = "0.8.3"
+rand.workspace = true
-regex = "1.4.5"
+regex.workspace = true
-bytes = "1.0.1"
+bytes.workspace = true
-byteorder = "1.4.3"
+byteorder.workspace = true
-anyhow = "1.0"
+anyhow.workspace = true
-crc32c = "0.6.0"
+crc32c.workspace = true
-hex = "0.4.3"
+hex.workspace = true
-once_cell = "1.13.0"
+once_cell.workspace = true
-log = "0.4.14"
+log.workspace = true
-memoffset = "0.7"
+memoffset.workspace = true
-thiserror = "1.0"
+thiserror.workspace = true
-serde = { version = "1.0", features = ["derive"] }
+serde.workspace = true
-utils = { path = "../utils" }
+utils.workspace = true
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+
+workspace_hack.workspace = true
 
 [dev-dependencies]
-env_logger = "0.9"
+env_logger.workspace = true
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres.workspace = true
 wal_craft = { path = "wal_craft" }
 
 [build-dependencies]
-anyhow = "1.0"
+anyhow.workspace = true
-bindgen = "0.61"
+bindgen.workspace = true

libs/postgres_ffi/wal_craft/Cargo.toml

@@ -1,17 +1,17 @@
 [package]
 name = "wal_craft"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-anyhow = "1.0"
+anyhow.workspace = true
-clap = "4.0"
+clap.workspace = true
-env_logger = "0.9"
+env_logger.workspace = true
-log = "0.4"
+log.workspace = true
-once_cell = "1.13.0"
+once_cell.workspace = true
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres.workspace = true
-postgres_ffi = { path = "../" }
+postgres_ffi.workspace = true
-tempfile = "3.2"
+tempfile.workspace = true
-workspace_hack = { version = "0.1", path = "../../../workspace_hack" }
+
+workspace_hack.workspace = true

libs/pq_proto/Cargo.toml

@@ -1,18 +1,18 @@
 [package]
 name = "pq_proto"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
+anyhow.workspace = true
-bytes = "1.0.1"
+bytes.workspace = true
-pin-project-lite = "0.2.7"
+pin-project-lite.workspace = true
-postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres-protocol.workspace = true
-rand = "0.8.3"
+rand.workspace = true
-serde = { version = "1.0", features = ["derive"] }
+serde.workspace = true
-tokio = { version = "1.17", features = ["macros"] }
+tokio.workspace = true
-tracing = "0.1"
+tracing.workspace = true
-thiserror = "1.0"
+thiserror.workspace = true
 
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+workspace_hack.workspace = true

libs/remote_storage/Cargo.toml

@@ -1,28 +1,28 @@
 [package]
 name = "remote_storage"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-anyhow = { version = "1.0", features = ["backtrace"] }
+anyhow.workspace = true
-async-trait = "0.1"
+async-trait.workspace = true
-metrics = { version = "0.1", path = "../metrics" }
+once_cell.workspace = true
-utils = { version = "0.1", path = "../utils" }
+aws-smithy-http.workspace = true
-once_cell = "1.13.0"
+aws-types.workspace = true
-aws-smithy-http = "0.51.0"
+aws-config.workspace = true
-aws-types = "0.51.0"
+aws-sdk-s3.workspace = true
-aws-config = { version = "0.51.0", default-features = false, features=["rustls"] }
+hyper = { workspace = true, features = ["stream"] }
-aws-sdk-s3 = "0.21.0"
+serde.workspace = true
-hyper = { version = "0.14", features = ["stream"] }
+serde_json.workspace = true
-serde = { version = "1.0", features = ["derive"] }
+tokio = { workspace = true, features = ["sync", "fs", "io-util"] }
-serde_json = "1"
+tokio-util.workspace = true
-tokio = { version = "1.17", features = ["sync", "macros", "fs", "io-util"] }
+toml_edit.workspace = true
-tokio-util = { version = "0.7", features = ["io"] }
+tracing.workspace = true
-toml_edit = { version = "0.14", features = ["easy"] }
+metrics.workspace = true
-tracing = "0.1.27"
+utils.workspace = true
 
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+workspace_hack.workspace = true
 
 [dev-dependencies]
-tempfile = "3.2"
+tempfile.workspace = true

libs/safekeeper_api/Cargo.toml

@@ -1,13 +1,13 @@
 [package]
 name = "safekeeper_api"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-serde = { version = "1.0", features = ["derive"] }
+serde.workspace = true
-serde_with = "2.0"
+serde_with.workspace = true
-const_format = "0.2.21"
+const_format.workspace = true
+utils.workspace = true
 
-utils = { path = "../utils" }
+workspace_hack.workspace = true
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }

libs/tenant_size_model/Cargo.toml

@@ -1,9 +1,11 @@
 [package]
 name = "tenant_size_model"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
 publish = false
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+anyhow.workspace = true
+
+workspace_hack.workspace = true

libs/tenant_size_model/src/lib.rs

@@ -1,6 +1,8 @@
 use std::borrow::Cow;
 use std::collections::HashMap;
 
+use anyhow::Context;
+
 /// Pricing model or history size builder.
 ///
 /// Maintains knowledge of the branches and their modifications. Generic over the branch name key
@@ -134,7 +136,7 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
         size: Option<u64>,
     ) where
         K: std::borrow::Borrow<Q>,
-        Q: std::hash::Hash + Eq,
+        Q: std::hash::Hash + Eq + std::fmt::Debug,
     {
         let lastseg_id = *self.branches.get(branch).unwrap();
         let newseg_id = self.segments.len();
@@ -214,20 +216,24 @@ impl<K: std::hash::Hash + Eq + 'static> Storage<K> {
     }
 
     /// Panics if the parent branch cannot be found.
-    pub fn branch<Q: ?Sized>(&mut self, parent: &Q, name: K)
+    pub fn branch<Q: ?Sized>(&mut self, parent: &Q, name: K) -> anyhow::Result<()>
     where
-        K: std::borrow::Borrow<Q>,
+        K: std::borrow::Borrow<Q> + std::fmt::Debug,
-        Q: std::hash::Hash + Eq,
+        Q: std::hash::Hash + Eq + std::fmt::Debug,
     {
         // Find the right segment
-        let branchseg_id = *self
+        let branchseg_id = *self.branches.get(parent).with_context(|| {
-            .branches
+            format!(
-            .get(parent)
+                "should had found the parent {:?} by key. in branches {:?}",
-            .expect("should had found the parent by key");
+                parent, self.branches
+            )
+        })?;
+
         let _branchseg = &mut self.segments[branchseg_id];
 
         // Create branch name for it
         self.branches.insert(name, branchseg_id);
+        Ok(())
     }
 
     pub fn calculate(&mut self, retention_period: u64) -> SegmentSize {

libs/tenant_size_model/src/main.rs

@@ -38,7 +38,7 @@ fn scenario_2() -> (Vec<Segment>, SegmentSize) {
     }
 
     // Branch
-    storage.branch("main", "child");
+    storage.branch("main", "child").unwrap();
     storage.update("child", 1_000);
 
     // More updates on parent
@@ -63,7 +63,7 @@ fn scenario_3() -> (Vec<Segment>, SegmentSize) {
     }
 
     // Branch
-    storage.branch("main", "child");
+    storage.branch("main", "child").unwrap();
     storage.update("child", 1_000);
 
     // More updates on parent
@@ -90,7 +90,7 @@ fn scenario_4() -> (Vec<Segment>, SegmentSize) {
     }
 
     // Branch
-    storage.branch("main", "child");
+    storage.branch("main", "child").unwrap();
     storage.update("child", 1_000);
 
     // More updates on parent
@@ -106,10 +106,10 @@ fn scenario_4() -> (Vec<Segment>, SegmentSize) {
 fn scenario_5() -> (Vec<Segment>, SegmentSize) {
     let mut storage = Storage::new("a");
     storage.insert("a", 5000);
-    storage.branch("a", "b");
+    storage.branch("a", "b").unwrap();
     storage.update("b", 4000);
     storage.update("a", 2000);
-    storage.branch("a", "c");
+    storage.branch("a", "c").unwrap();
     storage.insert("c", 4000);
     storage.insert("a", 2000);
 
@@ -133,12 +133,12 @@ fn scenario_6() -> (Vec<Segment>, SegmentSize) {
 
     let mut storage = Storage::new(None);
 
-    storage.branch(&None, branches[0]); // at 0
+    storage.branch(&None, branches[0]).unwrap(); // at 0
     storage.modify_branch(&branches[0], NO_OP, 108951064, 43696128); // at 108951064
-    storage.branch(&branches[0], branches[1]); // at 108951064
+    storage.branch(&branches[0], branches[1]).unwrap(); // at 108951064
     storage.modify_branch(&branches[1], NO_OP, 15560408, -1851392); // at 124511472
     storage.modify_branch(&branches[0], NO_OP, 174464360, -1531904); // at 283415424
-    storage.branch(&branches[0], branches[2]); // at 283415424
+    storage.branch(&branches[0], branches[2]).unwrap(); // at 283415424
     storage.modify_branch(&branches[2], NO_OP, 15906192, 8192); // at 299321616
     storage.modify_branch(&branches[0], NO_OP, 18909976, 32768); // at 302325400
 

libs/utils/Cargo.toml

@@ -1,48 +1,49 @@
 [package]
 name = "utils"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-sentry = { version = "0.29.0", default-features = false, features = ["backtrace", "contexts", "panic", "rustls", "reqwest" ] }
+sentry.workspace = true
-async-trait = "0.1"
+async-trait.workspace = true
-anyhow = "1.0"
+anyhow.workspace = true
-bincode = "1.3"
+bincode.workspace = true
-bytes = "1.0.1"
+bytes.workspace = true
-hyper = { version = "0.14.7", features = ["full"] }
+hyper = { workspace = true, features = ["full"] }
-routerify = "3"
+routerify.workspace = true
-serde = { version = "1.0", features = ["derive"] }
+serde.workspace = true
-serde_json = "1"
+serde_json.workspace = true
-thiserror = "1.0"
+thiserror.workspace = true
-tokio = { version = "1.17", features = ["macros"]}
+tokio.workspace = true
-tokio-rustls = "0.23"
+tokio-rustls.workspace = true
-tracing = "0.1"
+tracing.workspace = true
-tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
+tracing-subscriber = { workspace = true, features = ["json"] }
-nix = "0.25"
+nix.workspace = true
-signal-hook = "0.3.10"
+signal-hook.workspace = true
-rand = "0.8.3"
+rand.workspace = true
-jsonwebtoken = "8"
+jsonwebtoken.workspace = true
-hex = { version = "0.4.3", features = ["serde"] }
+hex = { workspace = true, features = ["serde"] }
-rustls = "0.20.2"
+rustls.workspace = true
-rustls-split = "0.3.0"
+rustls-split.workspace = true
-git-version = "0.3.5"
+git-version.workspace = true
-serde_with = "2.0"
+serde_with.workspace = true
-once_cell = "1.13.0"
+once_cell.workspace = true
-strum = "0.24"
+strum.workspace = true
-strum_macros = "0.24"
+strum_macros.workspace = true
 
-metrics = { path = "../metrics" }
+metrics.workspace = true
-pq_proto = { path = "../pq_proto" }
+pq_proto.workspace = true
-workspace_hack = { version = "0.1", path = "../../workspace_hack" }
+
+workspace_hack.workspace = true
 
 [dev-dependencies]
-byteorder = "1.4.3"
+byteorder.workspace = true
-bytes = "1.0.1"
+bytes.workspace = true
-hex-literal = "0.3"
+hex-literal.workspace = true
-tempfile = "3.2"
+tempfile.workspace = true
-criterion = "0.4"
+criterion.workspace = true
-rustls-pemfile = "1"
+rustls-pemfile.workspace = true
 
 [[bench]]
 name = "benchmarks"

libs/utils/src/logging.rs

@@ -8,6 +8,7 @@ use strum_macros::{EnumString, EnumVariantNames};
 pub enum LogFormat {
     Plain,
     Json,
+    Test,
 }
 
 impl LogFormat {
@@ -39,6 +40,7 @@ pub fn init(log_format: LogFormat) -> anyhow::Result<()> {
     match log_format {
         LogFormat::Json => base_logger.json().init(),
         LogFormat::Plain => base_logger.init(),
+        LogFormat::Test => base_logger.with_test_writer().init(),
     }
 
     Ok(())

libs/utils/src/postgres_backend_async.rs

@@ -7,12 +7,12 @@ use crate::postgres_backend::AuthType;
 use anyhow::Context;
 use bytes::{Buf, Bytes, BytesMut};
 use pq_proto::{BeMessage, ConnectionError, FeMessage, FeStartupPacket, SQLSTATE_INTERNAL_ERROR};
-use std::future::Future;
 use std::io;
 use std::net::SocketAddr;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::Poll;
+use std::{future::Future, task::ready};
 use tracing::{debug, error, info, trace};
 
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
@@ -253,12 +253,9 @@ impl PostgresBackend {
         cx: &mut std::task::Context<'_>,
     ) -> Poll<Result<(), std::io::Error>> {
         while self.buf_out.has_remaining() {
-            match Pin::new(&mut self.stream).poll_write(cx, self.buf_out.chunk()) {
+            match ready!(Pin::new(&mut self.stream).poll_write(cx, self.buf_out.chunk())) {
-                Poll::Ready(Ok(bytes_written)) => {
+                Ok(bytes_written) => self.buf_out.advance(bytes_written),
-                    self.buf_out.advance(bytes_written);
+                Err(err) => return Poll::Ready(Err(err)),
-                }
-                Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
-                Poll::Pending => return Poll::Pending,
             }
         }
         Poll::Ready(Ok(()))
@@ -573,10 +570,9 @@ impl<'a> AsyncWrite for CopyDataWriter<'a> {
         // It's not strictly required to flush between each message, but makes it easier
         // to view in wireshark, and usually the messages that the callers write are
         // decently-sized anyway.
-        match this.pgb.poll_write_buf(cx) {
+        match ready!(this.pgb.poll_write_buf(cx)) {
-            Poll::Ready(Ok(())) => {}
+            Ok(()) => {}
-            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+            Err(err) => return Poll::Ready(Err(err)),
-            Poll::Pending => return Poll::Pending,
         }
 
         // CopyData
@@ -593,10 +589,9 @@ impl<'a> AsyncWrite for CopyDataWriter<'a> {
         cx: &mut std::task::Context<'_>,
     ) -> Poll<Result<(), std::io::Error>> {
         let this = self.get_mut();
-        match this.pgb.poll_write_buf(cx) {
+        match ready!(this.pgb.poll_write_buf(cx)) {
-            Poll::Ready(Ok(())) => {}
+            Ok(()) => {}
-            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+            Err(err) => return Poll::Ready(Err(err)),
-            Poll::Pending => return Poll::Pending,
         }
         this.pgb.poll_flush(cx)
     }
@@ -605,10 +600,9 @@ impl<'a> AsyncWrite for CopyDataWriter<'a> {
         cx: &mut std::task::Context<'_>,
     ) -> Poll<Result<(), std::io::Error>> {
         let this = self.get_mut();
-        match this.pgb.poll_write_buf(cx) {
+        match ready!(this.pgb.poll_write_buf(cx)) {
-            Poll::Ready(Ok(())) => {}
+            Ok(()) => {}
-            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+            Err(err) => return Poll::Ready(Err(err)),
-            Poll::Pending => return Poll::Pending,
         }
         this.pgb.poll_flush(cx)
     }

pageserver/Cargo.toml

@@ -1,8 +1,8 @@
 [package]
 name = "pageserver"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [features]
 default = []
@@ -11,68 +11,68 @@ default = []
 testing = ["fail/failpoints"]
 
 [dependencies]
-amplify_num = { git = "https://github.com/hlinnaka/rust-amplify.git", branch = "unsigned-int-perf" }
+amplify_num.workspace = true
-anyhow = { version = "1.0", features = ["backtrace"] }
+anyhow.workspace = true
-async-stream = "0.3"
+async-stream.workspace = true
-async-trait = "0.1"
+async-trait.workspace = true
-byteorder = "1.4.3"
+byteorder.workspace = true
-bytes = "1.0.1"
+bytes.workspace = true
-chrono = { version = "0.4.23", default-features = false, features = ["clock", "serde"] }
+chrono = { workspace = true, features = ["serde"] }
-clap = { version = "4.0", features = ["string"] }
+clap = { workspace = true, features = ["string"] }
-close_fds = "0.3.2"
+close_fds.workspace = true
-const_format = "0.2.21"
+const_format.workspace = true
-crc32c = "0.6.0"
+consumption_metrics.workspace = true
-crossbeam-utils = "0.8.5"
+crc32c.workspace = true
-fail = "0.5.0"
+crossbeam-utils.workspace = true
-futures = "0.3.13"
+fail.workspace = true
-git-version = "0.3.5"
+futures.workspace = true
-hex = "0.4.3"
+git-version.workspace = true
-humantime = "2.1.0"
+hex.workspace = true
-humantime-serde = "1.1.1"
+humantime.workspace = true
-hyper = "0.14"
+humantime-serde.workspace = true
-itertools = "0.10.3"
+hyper.workspace = true
-nix = "0.25"
+itertools.workspace = true
-num-traits = "0.2.15"
+nix.workspace = true
-once_cell = "1.13.0"
+num-traits.workspace = true
-pin-project-lite = "0.2.7"
+once_cell.workspace = true
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+pin-project-lite.workspace = true
-postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres.workspace = true
-postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres-protocol.workspace = true
-rand = "0.8.3"
+postgres-types.workspace = true
-regex = "1.4.5"
+rand.workspace = true
-rstar = "0.9.3"
+regex.workspace = true
-scopeguard = "1.1.0"
+rstar.workspace = true
-serde = { version = "1.0", features = ["derive"] }
+scopeguard.workspace = true
-serde_json = { version = "1.0", features = ["raw_value"] }
+serde.workspace = true
-serde_with = "2.0"
+serde_json = { workspace = true, features = ["raw_value"] }
-signal-hook = "0.3.10"
+serde_with.workspace = true
-svg_fmt = "0.4.1"
+signal-hook.workspace = true
-tokio-tar = { git = "https://github.com/neondatabase/tokio-tar.git", rev="404df61437de0feef49ba2ccdbdd94eb8ad6e142" }
+svg_fmt.workspace = true
-thiserror = "1.0"
+tokio-tar.workspace = true
-tokio = { version = "1.17", features = ["process", "sync", "macros", "fs", "rt", "io-util", "time"] }
+thiserror.workspace = true
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+tokio = { workspace = true, features = ["process", "sync", "fs", "rt", "io-util", "time"] }
-tokio-util = { version = "0.7.3", features = ["io", "io-util"] }
+tokio-postgres.workspace = true
-toml_edit = { version = "0.14", features = ["easy"] }
+tokio-util.workspace = true
-tracing = "0.1.36"
+toml_edit.workspace = true
-url = "2"
+tracing.workspace = true
-walkdir = "2.3.2"
+url.workspace = true
-
+walkdir.workspace = true
-metrics = { path = "../libs/metrics" }
+metrics.workspace = true
-pageserver_api = { path = "../libs/pageserver_api" }
+pageserver_api.workspace = true
-postgres_connection = { path = "../libs/postgres_connection" }
+postgres_connection.workspace = true
-postgres_ffi = { path = "../libs/postgres_ffi" }
+postgres_ffi.workspace = true
-pq_proto = { path = "../libs/pq_proto" }
+pq_proto.workspace = true
-remote_storage = { path = "../libs/remote_storage" }
+remote_storage.workspace = true
-storage_broker = { version = "0.1", path = "../storage_broker" }
+storage_broker.workspace = true
-tenant_size_model = { path = "../libs/tenant_size_model" }
+tenant_size_model.workspace = true
-utils = { path = "../libs/utils" }
+utils.workspace = true
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
+workspace_hack.workspace = true
-reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"] }
+reqwest.workspace = true
 
 [dev-dependencies]
-criterion = "0.4"
+criterion.workspace = true
-hex-literal = "0.3"
+hex-literal.workspace = true
-tempfile = "3.2"
+tempfile.workspace = true
 
 [[bench]]
 name = "bench_layer_map"

pageserver/benches/bench_walredo.rs

@@ -30,33 +30,44 @@ fn redo_scenarios(c: &mut Criterion) {
     let conf = PageServerConf::dummy_conf(repo_dir.path().to_path_buf());
     let conf = Box::leak(Box::new(conf));
     let tenant_id = TenantId::generate();
-    // std::fs::create_dir_all(conf.tenant_path(&tenant_id)).unwrap();
+
-    let mut manager = PostgresRedoManager::new(conf, tenant_id);
+    let manager = PostgresRedoManager::new(conf, tenant_id);
-    manager.launch_process(14).unwrap();
 
     let manager = Arc::new(manager);
 
+    tracing::info!("executing first");
+    short().execute(&manager).unwrap();
+    tracing::info!("first executed");
+
     let thread_counts = [1, 2, 4, 8, 16];
 
-    for thread_count in thread_counts {
+    let mut group = c.benchmark_group("short");
-        c.bench_with_input(
+    group.sampling_mode(criterion::SamplingMode::Flat);
-            BenchmarkId::new("short-50record", thread_count),
-            &thread_count,
-            |b, thread_count| {
-                add_multithreaded_walredo_requesters(b, *thread_count, &manager, short, 50);
-            },
-        );
-    }
 
     for thread_count in thread_counts {
-        c.bench_with_input(
+        group.bench_with_input(
-            BenchmarkId::new("medium-10record", thread_count),
+            BenchmarkId::new("short", thread_count),
             &thread_count,
             |b, thread_count| {
-                add_multithreaded_walredo_requesters(b, *thread_count, &manager, medium, 10);
+                add_multithreaded_walredo_requesters(b, *thread_count, &manager, short);
             },
         );
     }
+    drop(group);
+
+    let mut group = c.benchmark_group("medium");
+    group.sampling_mode(criterion::SamplingMode::Flat);
+
+    for thread_count in thread_counts {
+        group.bench_with_input(
+            BenchmarkId::new("medium", thread_count),
+            &thread_count,
+            |b, thread_count| {
+                add_multithreaded_walredo_requesters(b, *thread_count, &manager, medium);
+            },
+        );
+    }
+    drop(group);
 }
 
 /// Sets up `threads` number of requesters to `request_redo`, with the given input.
@@ -65,46 +76,66 @@ fn add_multithreaded_walredo_requesters(
     threads: u32,
     manager: &Arc<PostgresRedoManager>,
     input_factory: fn() -> Request,
-    request_repeats: usize,
 ) {
-    b.iter_batched_ref(
+    assert_ne!(threads, 0);
-        || {
-            // barrier for all of the threads, and the benchmarked thread
-            let barrier = Arc::new(Barrier::new(threads as usize + 1));
 
-            let jhs = (0..threads)
+    if threads == 1 {
-                .map(|_| {
+        b.iter_batched_ref(
-                    std::thread::spawn({
+            || Some(input_factory()),
-                        let manager = manager.clone();
+            |input| execute_all(input.take(), manager),
-                        let barrier = barrier.clone();
+            criterion::BatchSize::PerIteration,
-                        move || {
+        );
-                            let input = std::iter::repeat(input_factory())
+    } else {
-                                .take(request_repeats)
+        let (work_tx, work_rx) = std::sync::mpsc::sync_channel(threads as usize);
-                                .collect::<Vec<_>>();
 
-                            barrier.wait();
+        let work_rx = std::sync::Arc::new(std::sync::Mutex::new(work_rx));
 
-                            execute_all(input, &manager).unwrap();
+        let barrier = Arc::new(Barrier::new(threads as usize + 1));
 
-                            barrier.wait();
+        let jhs = (0..threads)
+            .map(|_| {
+                std::thread::spawn({
+                    let manager = manager.clone();
+                    let barrier = barrier.clone();
+                    let work_rx = work_rx.clone();
+                    move || loop {
+                        // queue up and wait if we want to go another round
+                        if work_rx.lock().unwrap().recv().is_err() {
+                            break;
                         }
-                    })
+
+                        let input = Some(input_factory());
+
+                        barrier.wait();
+
+                        execute_all(input, &manager).unwrap();
+
+                        barrier.wait();
+                    }
                 })
-                .collect::<Vec<_>>();
+            })
+            .collect::<Vec<_>>();
 
-            (barrier, JoinOnDrop(jhs))
+        let _jhs = JoinOnDrop(jhs);
-        },
-        |input| {
-            let barrier = &input.0;
 
-            // start the work
+        b.iter_batched(
-            barrier.wait();
+            || {
+                for _ in 0..threads {
+                    work_tx.send(()).unwrap()
+                }
+            },
+            |()| {
+                // start the work
+                barrier.wait();
 
-            // wait for work to complete
+                // wait for work to complete
-            barrier.wait();
+                barrier.wait();
-        },
+            },
-        criterion::BatchSize::PerIteration,
+            criterion::BatchSize::PerIteration,
-    );
+        );
+
+        drop(work_tx);
+    }
 }
 
 struct JoinOnDrop(Vec<std::thread::JoinHandle<()>>);
@@ -121,7 +152,10 @@ impl Drop for JoinOnDrop {
     }
 }
 
-fn execute_all(input: Vec<Request>, manager: &PostgresRedoManager) -> Result<(), WalRedoError> {
+fn execute_all<I>(input: I, manager: &PostgresRedoManager) -> Result<(), WalRedoError>
+where
+    I: IntoIterator<Item = Request>,
+{
     // just fire all requests as fast as possible
     input.into_iter().try_for_each(|req| {
         let page = req.execute(manager)?;
@@ -143,6 +177,7 @@ macro_rules! lsn {
     }};
 }
 
+/// Short payload, 1132 bytes.
 // pg_records are copypasted from log, where they are put with Debug impl of Bytes, which uses \0
 // for null bytes.
 #[allow(clippy::octal_escapes)]
@@ -172,6 +207,7 @@ fn short() -> Request {
     }
 }
 
+/// Medium sized payload, serializes as 26393 bytes.
 // see [`short`]
 #[allow(clippy::octal_escapes)]
 fn medium() -> Request {

pageserver/src/bin/pageserver.rs

@@ -336,6 +336,7 @@ fn start_pageserver(conf: &'static PageServerConf) -> anyhow::Result<()> {
                     pageserver::consumption_metrics::collect_metrics(
                         metric_collection_endpoint,
                         conf.metric_collection_interval,
+                        conf.synthetic_size_calculation_interval,
                         conf.id,
                     )
                     .instrument(info_span!("metrics_collection"))

pageserver/src/config.rs

@@ -59,6 +59,8 @@ pub mod defaults {
 
     pub const DEFAULT_METRIC_COLLECTION_INTERVAL: &str = "10 min";
     pub const DEFAULT_METRIC_COLLECTION_ENDPOINT: Option<reqwest::Url> = None;
+    pub const DEFAULT_SYNTHETIC_SIZE_CALCULATION_INTERVAL: &str = "10 min";
+
     ///
     /// Default built-in configuration file.
     ///
@@ -83,6 +85,7 @@ pub mod defaults {
 #concurrent_tenant_size_logical_size_queries = '{DEFAULT_CONCURRENT_TENANT_SIZE_LOGICAL_SIZE_QUERIES}'
 
 #metric_collection_interval = '{DEFAULT_METRIC_COLLECTION_INTERVAL}'
+#synthetic_size_calculation_interval = '{DEFAULT_SYNTHETIC_SIZE_CALCULATION_INTERVAL}'
 
 # [tenant_config]
 #checkpoint_distance = {DEFAULT_CHECKPOINT_DISTANCE} # in bytes
@@ -152,6 +155,7 @@ pub struct PageServerConf {
     // How often to collect metrics and send them to the metrics endpoint.
     pub metric_collection_interval: Duration,
     pub metric_collection_endpoint: Option<Url>,
+    pub synthetic_size_calculation_interval: Duration,
 
     pub test_remote_failures: u64,
 }
@@ -215,6 +219,7 @@ struct PageServerConfigBuilder {
 
     metric_collection_interval: BuilderValue<Duration>,
     metric_collection_endpoint: BuilderValue<Option<Url>>,
+    synthetic_size_calculation_interval: BuilderValue<Duration>,
 
     test_remote_failures: BuilderValue<u64>,
 }
@@ -255,6 +260,10 @@ impl Default for PageServerConfigBuilder {
                 DEFAULT_METRIC_COLLECTION_INTERVAL,
             )
             .expect("cannot parse default metric collection interval")),
+            synthetic_size_calculation_interval: Set(humantime::parse_duration(
+                DEFAULT_SYNTHETIC_SIZE_CALCULATION_INTERVAL,
+            )
+            .expect("cannot parse default synthetic size calculation interval")),
             metric_collection_endpoint: Set(DEFAULT_METRIC_COLLECTION_ENDPOINT),
 
             test_remote_failures: Set(0),
@@ -342,6 +351,14 @@ impl PageServerConfigBuilder {
         self.metric_collection_endpoint = BuilderValue::Set(metric_collection_endpoint)
     }
 
+    pub fn synthetic_size_calculation_interval(
+        &mut self,
+        synthetic_size_calculation_interval: Duration,
+    ) {
+        self.synthetic_size_calculation_interval =
+            BuilderValue::Set(synthetic_size_calculation_interval)
+    }
+
     pub fn test_remote_failures(&mut self, fail_first: u64) {
         self.test_remote_failures = BuilderValue::Set(fail_first);
     }
@@ -399,6 +416,9 @@ impl PageServerConfigBuilder {
             metric_collection_endpoint: self
                 .metric_collection_endpoint
                 .ok_or(anyhow!("missing metric_collection_endpoint"))?,
+            synthetic_size_calculation_interval: self
+                .synthetic_size_calculation_interval
+                .ok_or(anyhow!("missing synthetic_size_calculation_interval"))?,
             test_remote_failures: self
                 .test_remote_failures
                 .ok_or(anyhow!("missing test_remote_failuers"))?,
@@ -577,7 +597,8 @@ impl PageServerConf {
                     let endpoint = parse_toml_string(key, item)?.parse().context("failed to parse metric_collection_endpoint")?;
                     builder.metric_collection_endpoint(Some(endpoint));
                 },
-
+                "synthetic_size_calculation_interval" =>
+                    builder.synthetic_size_calculation_interval(parse_toml_duration(key, item)?),
                 "test_remote_failures" => builder.test_remote_failures(parse_toml_u64(key, item)?),
                 _ => bail!("unrecognized pageserver option '{key}'"),
             }
@@ -672,11 +693,6 @@ impl PageServerConf {
         Ok(t_conf)
     }
 
-    #[cfg(test)]
-    pub fn test_repo_dir(test_name: &str) -> PathBuf {
-        PathBuf::from(format!("../tmp_check/test_{test_name}"))
-    }
-
     pub fn dummy_conf(repo_dir: PathBuf) -> Self {
         let pg_distrib_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../pg_install");
 
@@ -701,6 +717,7 @@ impl PageServerConf {
             concurrent_tenant_size_logical_size_queries: ConfigurableSemaphore::default(),
             metric_collection_interval: Duration::from_secs(60),
             metric_collection_endpoint: defaults::DEFAULT_METRIC_COLLECTION_ENDPOINT,
+            synthetic_size_calculation_interval: Duration::from_secs(60),
             test_remote_failures: 0,
         }
     }
@@ -834,6 +851,7 @@ id = 10
 
 metric_collection_interval = '222 s'
 metric_collection_endpoint = 'http://localhost:80/metrics'
+synthetic_size_calculation_interval = '333 s'
 log_format = 'json'
 
 "#;
@@ -880,6 +898,9 @@ log_format = 'json'
                     defaults::DEFAULT_METRIC_COLLECTION_INTERVAL
                 )?,
                 metric_collection_endpoint: defaults::DEFAULT_METRIC_COLLECTION_ENDPOINT,
+                synthetic_size_calculation_interval: humantime::parse_duration(
+                    defaults::DEFAULT_SYNTHETIC_SIZE_CALCULATION_INTERVAL
+                )?,
                 test_remote_failures: 0,
             },
             "Correct defaults should be used when no config values are provided"
@@ -926,6 +947,7 @@ log_format = 'json'
                 concurrent_tenant_size_logical_size_queries: ConfigurableSemaphore::default(),
                 metric_collection_interval: Duration::from_secs(222),
                 metric_collection_endpoint: Some(Url::parse("http://localhost:80/metrics")?),
+                synthetic_size_calculation_interval: Duration::from_secs(333),
                 test_remote_failures: 0,
             },
             "Should be able to parse all basic config values correctly"

pageserver/src/consumption_metrics.rs

@@ -3,154 +3,74 @@
 //! and push them to a HTTP endpoint.
 //! Cache metrics to send only the updated ones.
 //!
-
+use crate::task_mgr::{self, TaskKind, BACKGROUND_RUNTIME};
-use anyhow;
-use tracing::*;
-use utils::id::NodeId;
-use utils::id::TimelineId;
-
-use crate::task_mgr;
 use crate::tenant::mgr;
+use anyhow;
+use chrono::Utc;
+use consumption_metrics::{idempotency_key, Event, EventChunk, EventType, CHUNK_SIZE};
 use pageserver_api::models::TenantState;
-use utils::id::TenantId;
+use reqwest::Url;
-
+use serde::Serialize;
-use serde::{Deserialize, Serialize};
 use serde_with::{serde_as, DisplayFromStr};
 use std::collections::HashMap;
-use std::fmt;
-use std::str::FromStr;
 use std::time::Duration;
+use tracing::*;
+use utils::id::{NodeId, TenantId, TimelineId};
 
-use chrono::{DateTime, Utc};
+const WRITTEN_SIZE: &str = "written_size";
-use rand::Rng;
+const SYNTHETIC_STORAGE_SIZE: &str = "synthetic_storage_size";
-use reqwest::Url;
+const RESIDENT_SIZE: &str = "resident_size";
+const REMOTE_STORAGE_SIZE: &str = "remote_storage_size";
+const TIMELINE_LOGICAL_SIZE: &str = "timeline_logical_size";
 
-/// ConsumptionMetric struct that defines the format for one metric entry
-/// i.e.
-///
-/// ```json
-/// {
-/// "metric": "remote_storage_size",
-/// "type": "absolute",
-/// "tenant_id": "5d07d9ce9237c4cd845ea7918c0afa7d",
-/// "timeline_id": "a03ebb4f5922a1c56ff7485cc8854143",
-/// "time": "2022-12-28T11:07:19.317310284Z",
-/// "idempotency_key": "2022-12-28 11:07:19.317310324 UTC-1-4019",
-/// "value": 12345454,
-/// }
-/// ```
 #[serde_as]
-#[derive(Serialize, Deserialize, Debug, Clone, Eq, PartialEq, Ord, PartialOrd)]
+#[derive(Serialize)]
-pub struct ConsumptionMetric {
+struct Ids {
-    pub metric: ConsumptionMetricKind,
-    #[serde(rename = "type")]
-    pub metric_type: &'static str,
     #[serde_as(as = "DisplayFromStr")]
-    pub tenant_id: TenantId,
+    tenant_id: TenantId,
     #[serde_as(as = "Option<DisplayFromStr>")]
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub timeline_id: Option<TimelineId>,
-    pub time: DateTime<Utc>,
-    pub idempotency_key: String,
-    pub value: u64,
-}
-
-impl ConsumptionMetric {
-    pub fn new_absolute<R: Rng + ?Sized>(
-        metric: ConsumptionMetricKind,
-        tenant_id: TenantId,
-        timeline_id: Option<TimelineId>,
-        value: u64,
-        node_id: NodeId,
-        rng: &mut R,
-    ) -> Self {
-        Self {
-            metric,
-            metric_type: "absolute",
-            tenant_id,
-            timeline_id,
-            time: Utc::now(),
-            // key that allows metric collector to distinguish unique events
-            idempotency_key: format!("{}-{}-{:04}", Utc::now(), node_id, rng.gen_range(0..=9999)),
-            value,
-        }
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, Serialize, Deserialize)]
-#[serde(rename_all = "snake_case")]
-pub enum ConsumptionMetricKind {
-    /// Amount of WAL produced , by a timeline, i.e. last_record_lsn
-    /// This is an absolute, per-timeline metric.
-    WrittenSize,
-    /// Size of all tenant branches including WAL
-    /// This is an absolute, per-tenant metric.
-    /// This is the same metric that tenant/tenant_id/size endpoint returns.
-    SyntheticStorageSize,
-    /// Size of all the layer files in the tenant's directory on disk on the pageserver.
-    /// This is an absolute, per-tenant metric.
-    /// See also prometheus metric RESIDENT_PHYSICAL_SIZE.
-    ResidentSize,
-    /// Size of the remote storage (S3) directory.
-    /// This is an absolute, per-tenant metric.
-    RemoteStorageSize,
-    /// Logical size of the data in the timeline
-    /// This is an absolute, per-timeline metric
-    TimelineLogicalSize,
-}
-
-impl FromStr for ConsumptionMetricKind {
-    type Err = anyhow::Error;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "written_size" => Ok(Self::WrittenSize),
-            "synthetic_storage_size" => Ok(Self::SyntheticStorageSize),
-            "resident_size" => Ok(Self::ResidentSize),
-            "remote_storage_size" => Ok(Self::RemoteStorageSize),
-            "timeline_logical_size" => Ok(Self::TimelineLogicalSize),
-            _ => anyhow::bail!("invalid value \"{s}\" for metric type"),
-        }
-    }
-}
-
-impl fmt::Display for ConsumptionMetricKind {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.write_str(match self {
-            ConsumptionMetricKind::WrittenSize => "written_size",
-            ConsumptionMetricKind::SyntheticStorageSize => "synthetic_storage_size",
-            ConsumptionMetricKind::ResidentSize => "resident_size",
-            ConsumptionMetricKind::RemoteStorageSize => "remote_storage_size",
-            ConsumptionMetricKind::TimelineLogicalSize => "timeline_logical_size",
-        })
-    }
-}
-
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-pub struct ConsumptionMetricsKey {
-    tenant_id: TenantId,
     timeline_id: Option<TimelineId>,
-    metric: ConsumptionMetricKind,
 }
 
-#[derive(serde::Serialize)]
+/// Key that uniquely identifies the object, this metric describes.
-struct EventChunk<'a> {
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-    events: &'a [ConsumptionMetric],
+pub struct PageserverConsumptionMetricsKey {
+    pub tenant_id: TenantId,
+    pub timeline_id: Option<TimelineId>,
+    pub metric: &'static str,
 }
 
 /// Main thread that serves metrics collection
 pub async fn collect_metrics(
     metric_collection_endpoint: &Url,
     metric_collection_interval: Duration,
+    synthetic_size_calculation_interval: Duration,
     node_id: NodeId,
 ) -> anyhow::Result<()> {
     let mut ticker = tokio::time::interval(metric_collection_interval);
 
     info!("starting collect_metrics");
 
+    // spin up background worker that caclulates tenant sizes
+    task_mgr::spawn(
+        BACKGROUND_RUNTIME.handle(),
+        TaskKind::CalculateSyntheticSize,
+        None,
+        None,
+        "synthetic size calculation",
+        true,
+        async move {
+            calculate_synthetic_size_worker(synthetic_size_calculation_interval)
+                .instrument(info_span!("synthetic_size_worker"))
+                .await?;
+            Ok(())
+        },
+    );
+
     // define client here to reuse it for all requests
     let client = reqwest::Client::new();
-    let mut cached_metrics: HashMap<ConsumptionMetricsKey, u64> = HashMap::new();
+    let mut cached_metrics: HashMap<PageserverConsumptionMetricsKey, u64> = HashMap::new();
 
     loop {
         tokio::select! {
@@ -159,7 +79,10 @@ pub async fn collect_metrics(
                 return Ok(());
             },
             _ = ticker.tick() => {
-                collect_metrics_task(&client, &mut cached_metrics, metric_collection_endpoint, node_id).await?;
+                if let Err(err) = collect_metrics_iteration(&client, &mut cached_metrics, metric_collection_endpoint, node_id).await
+                {
+                    error!("metrics collection failed: {err:?}");
+                }
             }
         }
     }
@@ -169,15 +92,20 @@ pub async fn collect_metrics(
 ///
 /// Gather per-tenant and per-timeline metrics and send them to the `metric_collection_endpoint`.
 /// Cache metrics to avoid sending the same metrics multiple times.
-pub async fn collect_metrics_task(
+///
+/// TODO
+/// - refactor this function (chunking+sending part) to reuse it in proxy module;
+/// - improve error handling. Now if one tenant fails to collect metrics,
+/// the whole iteration fails and metrics for other tenants are not collected.
+pub async fn collect_metrics_iteration(
     client: &reqwest::Client,
-    cached_metrics: &mut HashMap<ConsumptionMetricsKey, u64>,
+    cached_metrics: &mut HashMap<PageserverConsumptionMetricsKey, u64>,
     metric_collection_endpoint: &reqwest::Url,
     node_id: NodeId,
 ) -> anyhow::Result<()> {
-    let mut current_metrics: Vec<(ConsumptionMetricsKey, u64)> = Vec::new();
+    let mut current_metrics: Vec<(PageserverConsumptionMetricsKey, u64)> = Vec::new();
     trace!(
-        "starting collect_metrics_task. metric_collection_endpoint: {}",
+        "starting collect_metrics_iteration. metric_collection_endpoint: {}",
         metric_collection_endpoint
     );
 
@@ -201,10 +129,10 @@ pub async fn collect_metrics_task(
                 let timeline_written_size = u64::from(timeline.get_last_record_lsn());
 
                 current_metrics.push((
-                    ConsumptionMetricsKey {
+                    PageserverConsumptionMetricsKey {
                         tenant_id,
                         timeline_id: Some(timeline.timeline_id),
-                        metric: ConsumptionMetricKind::WrittenSize,
+                        metric: WRITTEN_SIZE,
                     },
                     timeline_written_size,
                 ));
@@ -213,10 +141,10 @@ pub async fn collect_metrics_task(
                 // Only send timeline logical size when it is fully calculated.
                 if is_exact {
                     current_metrics.push((
-                        ConsumptionMetricsKey {
+                        PageserverConsumptionMetricsKey {
                             tenant_id,
                             timeline_id: Some(timeline.timeline_id),
-                            metric: ConsumptionMetricKind::TimelineLogicalSize,
+                            metric: TIMELINE_LOGICAL_SIZE,
                         },
                         timeline_logical_size,
                     ));
@@ -234,24 +162,34 @@ pub async fn collect_metrics_task(
         );
 
         current_metrics.push((
-            ConsumptionMetricsKey {
+            PageserverConsumptionMetricsKey {
                 tenant_id,
                 timeline_id: None,
-                metric: ConsumptionMetricKind::ResidentSize,
+                metric: RESIDENT_SIZE,
             },
             tenant_resident_size,
         ));
 
         current_metrics.push((
-            ConsumptionMetricsKey {
+            PageserverConsumptionMetricsKey {
                 tenant_id,
                 timeline_id: None,
-                metric: ConsumptionMetricKind::RemoteStorageSize,
+                metric: REMOTE_STORAGE_SIZE,
             },
             tenant_remote_size,
         ));
 
-        // TODO add SyntheticStorageSize metric
+        // Note that this metric is calculated in a separate bgworker
+        // Here we only use cached value, which may lag behind the real latest one
+        let tenant_synthetic_size = tenant.get_cached_synthetic_size();
+        current_metrics.push((
+            PageserverConsumptionMetricsKey {
+                tenant_id,
+                timeline_id: None,
+                metric: SYNTHETIC_STORAGE_SIZE,
+            },
+            tenant_synthetic_size,
+        ));
     }
 
     // Filter metrics
@@ -267,35 +205,29 @@ pub async fn collect_metrics_task(
 
     // Send metrics.
     // Split into chunks of 1000 metrics to avoid exceeding the max request size
-    const CHUNK_SIZE: usize = 1000;
     let chunks = current_metrics.chunks(CHUNK_SIZE);
 
-    let mut chunk_to_send: Vec<ConsumptionMetric> = Vec::with_capacity(1000);
+    let mut chunk_to_send: Vec<Event<Ids>> = Vec::with_capacity(CHUNK_SIZE);
 
     for chunk in chunks {
         chunk_to_send.clear();
 
-        // this code block is needed to convince compiler
+        // enrich metrics with type,timestamp and idempotency key before sending
-        // that rng is not reused aroung await point
+        chunk_to_send.extend(chunk.iter().map(|(curr_key, curr_val)| Event {
-        {
+            kind: EventType::Absolute { time: Utc::now() },
-            // enrich metrics with timestamp and metric_kind before sending
+            metric: curr_key.metric,
-            let mut rng = rand::thread_rng();
+            idempotency_key: idempotency_key(node_id.to_string()),
-            chunk_to_send.extend(chunk.iter().map(|(curr_key, curr_val)| {
+            value: *curr_val,
-                ConsumptionMetric::new_absolute(
+            extra: Ids {
-                    curr_key.metric,
+                tenant_id: curr_key.tenant_id,
-                    curr_key.tenant_id,
+                timeline_id: curr_key.timeline_id,
-                    curr_key.timeline_id,
+            },
-                    *curr_val,
+        }));
-                    node_id,
-                    &mut rng,
-                )
-            }));
-        }
 
         let chunk_json = serde_json::value::to_raw_value(&EventChunk {
             events: &chunk_to_send,
         })
-        .expect("ConsumptionMetric should not fail serialization");
+        .expect("PageserverConsumptionMetric should not fail serialization");
 
         let res = client
             .post(metric_collection_endpoint.clone())
@@ -322,3 +254,39 @@ pub async fn collect_metrics_task(
 
     Ok(())
 }
+
+/// Caclculate synthetic size for each active tenant
+pub async fn calculate_synthetic_size_worker(
+    synthetic_size_calculation_interval: Duration,
+) -> anyhow::Result<()> {
+    info!("starting calculate_synthetic_size_worker");
+
+    let mut ticker = tokio::time::interval(synthetic_size_calculation_interval);
+
+    loop {
+        tokio::select! {
+            _ = task_mgr::shutdown_watcher() => {
+                return Ok(());
+            },
+        _ = ticker.tick() => {
+
+                let tenants = mgr::list_tenants().await;
+                // iterate through list of Active tenants and collect metrics
+                for (tenant_id, tenant_state) in tenants {
+
+                    if tenant_state != TenantState::Active {
+                        continue;
+                    }
+
+                    if let Ok(tenant) = mgr::get_tenant(tenant_id, true).await
+                    {
+                        if let Err(e) = tenant.calculate_synthetic_size().await {
+                            error!("failed to calculate synthetic size for tenant {}: {}", tenant_id, e);
+                        }
+                    }
+
+                }
+            }
+        }
+    }
+}

pageserver/src/http/routes.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use anyhow::{anyhow, Context, Result};
 use hyper::StatusCode;
 use hyper::{Body, Request, Response, Uri};
+use pageserver_api::models::DownloadRemoteLayersTaskSpawnRequest;
 use remote_storage::GenericRemoteStorage;
 use tokio_util::sync::CancellationToken;
 use tracing::*;
@@ -788,10 +789,11 @@ async fn timeline_checkpoint_handler(request: Request<Body>) -> Result<Response<
 }
 
 async fn timeline_download_remote_layers_handler_post(
-    request: Request<Body>,
+    mut request: Request<Body>,
 ) -> Result<Response<Body>, ApiError> {
     let tenant_id: TenantId = parse_request_param(&request, "tenant_id")?;
     let timeline_id: TimelineId = parse_request_param(&request, "timeline_id")?;
+    let body: DownloadRemoteLayersTaskSpawnRequest = json_request(&mut request).await?;
     check_permission(&request, Some(tenant_id))?;
 
     let tenant = mgr::get_tenant(tenant_id, true)
@@ -800,7 +802,7 @@ async fn timeline_download_remote_layers_handler_post(
     let timeline = tenant
         .get_timeline(timeline_id, true)
         .map_err(ApiError::NotFound)?;
-    match timeline.spawn_download_all_remote_layers().await {
+    match timeline.spawn_download_all_remote_layers(body).await {
         Ok(st) => json_response(StatusCode::ACCEPTED, st),
         Err(st) => json_response(StatusCode::CONFLICT, st),
     }

pageserver/src/task_mgr.rs

@@ -220,6 +220,8 @@ pub enum TaskKind {
 
     // task that drives downloading layers
     DownloadAllRemoteLayers,
+    // Task that calculates synthetis size for all active tenants
+    CalculateSyntheticSize,
 }
 
 #[derive(Default)]

pageserver/src/tenant.rs

@@ -38,6 +38,8 @@ use std::path::Path;
 use std::path::PathBuf;
 use std::process::Command;
 use std::process::Stdio;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 use std::sync::MutexGuard;
 use std::sync::{Mutex, RwLock};
@@ -139,6 +141,7 @@ pub struct Tenant {
 
     /// Cached logical sizes updated updated on each [`Tenant::gather_size_inputs`].
     cached_logical_sizes: tokio::sync::Mutex<HashMap<(TimelineId, Lsn), u64>>,
+    cached_synthetic_tenant_size: Arc<AtomicU64>,
 }
 
 /// A timeline with some of its files on disk, being initialized.
@@ -438,8 +441,16 @@ struct RemoteStartupData {
 
 impl Tenant {
     /// Yet another helper for timeline initialization.
-    /// Contains common part for `load_local_timeline` and `load_remote_timeline`
+    /// Contains the common part of `load_local_timeline` and `load_remote_timeline`.
-    async fn setup_timeline(
+    ///
+    /// - Initializes the Timeline struct and inserts it into the tenant's hash map
+    /// - Scans the local timeline directory for layer files and builds the layer map
+    /// - Downloads remote index file and adds remote files to the layer map
+    /// - Schedules remote upload tasks for any files that are present locally but missing from remote storage.
+    ///
+    /// If the operation fails, the timeline is left in the tenant's hash map in Broken state. On success,
+    /// it is marked as Active.
+    async fn timeline_init_and_sync(
         &self,
         timeline_id: TimelineId,
         remote_client: Option<RemoteTimelineClient>,
@@ -482,10 +493,7 @@ impl Tenant {
             // But we shouldnt start walreceiver before we have all the data locally, because working walreceiver
             // will ingest data which may require looking at the layers which are not yet available locally
             match timeline.initialize_with_lock(&mut timelines_accessor, true, false) {
-                Ok(initialized_timeline) => {
+                Ok(new_timeline) => new_timeline,
-                    timelines_accessor.insert(timeline_id, initialized_timeline.clone());
-                    Ok(initialized_timeline)
-                }
                 Err(e) => {
                     error!("Failed to initialize timeline {tenant_id}/{timeline_id}: {e:?}");
                     // FIXME using None is a hack, it wont hurt, just ugly.
@@ -501,16 +509,14 @@ impl Tenant {
                             None,
                         )
                         .with_context(|| {
-                            format!(
+                            format!("creating broken timeline data for {tenant_id}/{timeline_id}")
-                            "Failed to crate broken timeline data for {tenant_id}/{timeline_id}"
-                        )
                         })?;
                     broken_timeline.set_state(TimelineState::Broken);
                     timelines_accessor.insert(timeline_id, broken_timeline);
-                    Err(e)
+                    return Err(e);
                 }
             }
-        }?;
+        };
 
         if self.remote_storage.is_some() {
             // Reconcile local state with remote storage, downloading anything that's
@@ -783,7 +789,7 @@ impl Tenant {
         // cannot be older than the local one
         let local_metadata = None;
 
-        self.setup_timeline(
+        self.timeline_init_and_sync(
             timeline_id,
             Some(remote_client),
             Some(RemoteStartupData {
@@ -1048,7 +1054,7 @@ impl Tenant {
             None => None,
         };
 
-        self.setup_timeline(
+        self.timeline_init_and_sync(
             timeline_id,
             remote_client,
             remote_startup_data,
@@ -1722,6 +1728,7 @@ impl Tenant {
             remote_storage,
             state,
             cached_logical_sizes: tokio::sync::Mutex::new(HashMap::new()),
+            cached_synthetic_tenant_size: Arc::new(AtomicU64::new(0)),
         }
     }
 
@@ -2359,6 +2366,24 @@ impl Tenant {
 
         size::gather_inputs(self, logical_sizes_at_once, &mut shared_cache).await
     }
+
+    /// Calculate synthetic tenant size
+    /// This is periodically called by background worker.
+    /// result is cached in tenant struct
+    #[instrument(skip_all, fields(tenant_id=%self.tenant_id))]
+    pub async fn calculate_synthetic_size(&self) -> anyhow::Result<u64> {
+        let inputs = self.gather_size_inputs().await?;
+
+        let size = inputs.calculate()?;
+
+        self.cached_synthetic_tenant_size
+            .store(size, Ordering::Relaxed);
+
+        Ok(size)
+    }
+    pub fn get_cached_synthetic_size(&self) -> u64 {
+        self.cached_synthetic_tenant_size.load(Ordering::Relaxed)
+    }
 }
 
 fn remove_timeline_and_uninit_mark(timeline_dir: &Path, uninit_mark: &Path) -> anyhow::Result<()> {
@@ -2601,9 +2626,11 @@ where
 #[cfg(test)]
 pub mod harness {
     use bytes::{Bytes, BytesMut};
-    use once_cell::sync::Lazy;
+    use once_cell::sync::OnceCell;
-    use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
+    use std::sync::Arc;
     use std::{fs, path::PathBuf};
+    use tempfile::TempDir;
+    use utils::logging;
     use utils::lsn::Lsn;
 
     use crate::{
@@ -2634,8 +2661,6 @@ pub mod harness {
         buf.freeze()
     }
 
-    static LOCK: Lazy<RwLock<()>> = Lazy::new(|| RwLock::new(()));
-
     impl From<TenantConf> for TenantConfOpt {
         fn from(tenant_conf: TenantConf) -> Self {
             Self {
@@ -2656,36 +2681,31 @@ pub mod harness {
         }
     }
 
-    pub struct TenantHarness<'a> {
+    /// The harness saves some boilerplate and provides a way to create functional tenant
+    /// without running pageserver binary. It uses temporary directory to store data in it.
+    /// Tempdir gets removed on harness drop.
+    pub struct TenantHarness {
+        // keep the struct to not to remove tmp dir during the test
+        _temp_repo_dir: TempDir,
         pub conf: &'static PageServerConf,
         pub tenant_conf: TenantConf,
         pub tenant_id: TenantId,
-
-        pub lock_guard: (
-            Option<RwLockReadGuard<'a, ()>>,
-            Option<RwLockWriteGuard<'a, ()>>,
-        ),
     }
 
-    impl<'a> TenantHarness<'a> {
+    static LOG_HANDLE: OnceCell<()> = OnceCell::new();
-        pub fn create(test_name: &'static str) -> anyhow::Result<Self> {
-            Self::create_internal(test_name, false)
-        }
-        pub fn create_exclusive(test_name: &'static str) -> anyhow::Result<Self> {
-            Self::create_internal(test_name, true)
-        }
-        fn create_internal(test_name: &'static str, exclusive: bool) -> anyhow::Result<Self> {
-            let lock_guard = if exclusive {
-                (None, Some(LOCK.write().unwrap()))
-            } else {
-                (Some(LOCK.read().unwrap()), None)
-            };
 
-            let repo_dir = PageServerConf::test_repo_dir(test_name);
+    impl TenantHarness {
-            let _ = fs::remove_dir_all(&repo_dir);
+        pub fn new() -> anyhow::Result<Self> {
-            fs::create_dir_all(&repo_dir)?;
+            LOG_HANDLE.get_or_init(|| {
+                logging::init(logging::LogFormat::Test).expect("Failed to init test logging")
+            });
 
-            let conf = PageServerConf::dummy_conf(repo_dir);
+            let temp_repo_dir = tempfile::tempdir()?;
+            // `TempDir` uses a randomly generated subdirectory of a system tmp dir,
+            // so far it's enough to take care of concurrently running tests.
+            let repo_dir = temp_repo_dir.path();
+
+            let conf = PageServerConf::dummy_conf(repo_dir.to_path_buf());
             // Make a static copy of the config. This can never be free'd, but that's
             // OK in a test.
             let conf: &'static PageServerConf = Box::leak(Box::new(conf));
@@ -2703,10 +2723,10 @@ pub mod harness {
             fs::create_dir_all(conf.timelines_path(&tenant_id))?;
 
             Ok(Self {
+                _temp_repo_dir: temp_repo_dir,
                 conf,
                 tenant_conf,
                 tenant_id,
-                lock_guard,
             })
         }
 
@@ -2800,7 +2820,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_basic() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_basic")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -2833,9 +2854,8 @@ mod tests {
 
     #[tokio::test]
     async fn no_duplicate_timelines() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("no_duplicate_timelines")?
+        let harness = TenantHarness::new()?;
-            .load()
+        let tenant = harness.load().await;
-            .await;
         let _ = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -2866,7 +2886,8 @@ mod tests {
     ///
     #[tokio::test]
     async fn test_branch() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_branch")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -2963,10 +2984,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_prohibit_branch_creation_on_garbage_collected_data() -> anyhow::Result<()> {
-        let tenant =
+        let harness = TenantHarness::new()?;
-            TenantHarness::create("test_prohibit_branch_creation_on_garbage_collected_data")?
+        let tenant = harness.load().await;
-                .load()
-                .await;
         let tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -3001,9 +3020,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_prohibit_branch_creation_on_pre_initdb_lsn() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_prohibit_branch_creation_on_pre_initdb_lsn")?
+        let harness = TenantHarness::new()?;
-            .load()
+        let tenant = harness.load().await;
-            .await;
 
         tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0x50), DEFAULT_PG_VERSION)?
@@ -3052,9 +3070,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_retain_data_in_parent_which_is_needed_for_child() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_retain_data_in_parent_which_is_needed_for_child")?
+        let harness = TenantHarness::new()?;
-            .load()
+        let tenant = harness.load().await;
-            .await;
         let tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -3076,9 +3093,8 @@ mod tests {
     }
     #[tokio::test]
     async fn test_parent_keeps_data_forever_after_branching() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_parent_keeps_data_forever_after_branching")?
+        let harness = TenantHarness::new()?;
-            .load()
+        let tenant = harness.load().await;
-            .await;
         let tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -3109,8 +3125,7 @@ mod tests {
 
     #[tokio::test]
     async fn timeline_load() -> anyhow::Result<()> {
-        const TEST_NAME: &str = "timeline_load";
+        let harness = TenantHarness::new()?;
-        let harness = TenantHarness::create(TEST_NAME)?;
         {
             let tenant = harness.load().await;
             let tline = tenant
@@ -3129,8 +3144,7 @@ mod tests {
 
     #[tokio::test]
     async fn timeline_load_with_ancestor() -> anyhow::Result<()> {
-        const TEST_NAME: &str = "timeline_load_with_ancestor";
+        let harness = TenantHarness::new()?;
-        let harness = TenantHarness::create(TEST_NAME)?;
         // create two timelines
         {
             let tenant = harness.load().await;
@@ -3168,8 +3182,7 @@ mod tests {
 
     #[tokio::test]
     async fn corrupt_metadata() -> anyhow::Result<()> {
-        const TEST_NAME: &str = "corrupt_metadata";
+        let harness = TenantHarness::new()?;
-        let harness = TenantHarness::create(TEST_NAME)?;
         let tenant = harness.load().await;
 
         tenant
@@ -3210,7 +3223,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_images() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_images")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -3277,7 +3291,8 @@ mod tests {
     //
     #[tokio::test]
     async fn test_bulk_insert() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_bulk_insert")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -3321,7 +3336,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_random_updates() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_random_updates")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -3394,9 +3410,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_traverse_branches() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_traverse_branches")?
+        let harness = TenantHarness::new()?;
-            .load()
+        let tenant = harness.load().await;
-            .await;
         let mut tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;
@@ -3480,9 +3495,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_traverse_ancestors() -> anyhow::Result<()> {
-        let tenant = TenantHarness::create("test_traverse_ancestors")?
+        let harness = TenantHarness::new()?;
-            .load()
+        let tenant = harness.load().await;
-            .await;
         let mut tline = tenant
             .create_empty_timeline(TIMELINE_ID, Lsn(0), DEFAULT_PG_VERSION)?
             .initialize()?;

pageserver/src/tenant/ephemeral_file.rs

@@ -76,7 +76,7 @@ impl EphemeralFile {
         })
     }
 
-    fn fill_buffer(&self, buf: &mut [u8], blkno: u32) -> Result<(), io::Error> {
+    fn fill_buffer(&self, buf: &mut [u8], blkno: u32) -> io::Result<()> {
         let mut off = 0;
         while off < PAGE_SZ {
             let n = self
@@ -277,7 +277,7 @@ impl Drop for EphemeralFile {
     }
 }
 
-pub fn writeback(file_id: u64, blkno: u32, buf: &[u8]) -> Result<(), io::Error> {
+pub fn writeback(file_id: u64, blkno: u32, buf: &[u8]) -> io::Result<()> {
     if let Some(file) = EPHEMERAL_FILES.read().unwrap().files.get(&file_id) {
         match file.write_all_at(buf, blkno as u64 * PAGE_SZ as u64) {
             Ok(_) => Ok(()),
@@ -332,25 +332,17 @@ mod tests {
     use super::*;
     use crate::tenant::blob_io::{BlobCursor, BlobWriter};
     use crate::tenant::block_io::BlockCursor;
+    use crate::tenant::harness::TenantHarness;
     use rand::{seq::SliceRandom, thread_rng, RngCore};
     use std::fs;
     use std::str::FromStr;
 
-    fn harness(
+    fn harness() -> Result<(TenantHarness, TimelineId), io::Error> {
-        test_name: &str,
+        let harness = TenantHarness::new().expect("Failed to create tenant harness");
-    ) -> Result<(&'static PageServerConf, TenantId, TimelineId), io::Error> {
-        let repo_dir = PageServerConf::test_repo_dir(test_name);
-        let _ = fs::remove_dir_all(&repo_dir);
-        let conf = PageServerConf::dummy_conf(repo_dir);
-        // Make a static copy of the config. This can never be free'd, but that's
-        // OK in a test.
-        let conf: &'static PageServerConf = Box::leak(Box::new(conf));
-
-        let tenant_id = TenantId::from_str("11000000000000000000000000000000").unwrap();
         let timeline_id = TimelineId::from_str("22000000000000000000000000000000").unwrap();
-        fs::create_dir_all(conf.timeline_path(&timeline_id, &tenant_id))?;
+        fs::create_dir_all(harness.timeline_path(&timeline_id))?;
 
-        Ok((conf, tenant_id, timeline_id))
+        Ok((harness, timeline_id))
     }
 
     // Helper function to slurp contents of a file, starting at the current position,
@@ -367,10 +359,10 @@ mod tests {
     }
 
     #[test]
-    fn test_ephemeral_files() -> Result<(), io::Error> {
+    fn test_ephemeral_files() -> io::Result<()> {
-        let (conf, tenant_id, timeline_id) = harness("ephemeral_files")?;
+        let (harness, timeline_id) = harness()?;
 
-        let file_a = EphemeralFile::create(conf, tenant_id, timeline_id)?;
+        let file_a = EphemeralFile::create(harness.conf, harness.tenant_id, timeline_id)?;
 
         file_a.write_all_at(b"foo", 0)?;
         assert_eq!("foo", read_string(&file_a, 0, 20)?);
@@ -381,7 +373,7 @@ mod tests {
         // Open a lot of files, enough to cause some page evictions.
         let mut efiles = Vec::new();
         for fileno in 0..100 {
-            let efile = EphemeralFile::create(conf, tenant_id, timeline_id)?;
+            let efile = EphemeralFile::create(harness.conf, harness.tenant_id, timeline_id)?;
             efile.write_all_at(format!("file {}", fileno).as_bytes(), 0)?;
             assert_eq!(format!("file {}", fileno), read_string(&efile, 0, 10)?);
             efiles.push((fileno, efile));
@@ -398,10 +390,10 @@ mod tests {
     }
 
     #[test]
-    fn test_ephemeral_blobs() -> Result<(), io::Error> {
+    fn test_ephemeral_blobs() -> io::Result<()> {
-        let (conf, tenant_id, timeline_id) = harness("ephemeral_blobs")?;
+        let (harness, timeline_id) = harness()?;
 
-        let mut file = EphemeralFile::create(conf, tenant_id, timeline_id)?;
+        let mut file = EphemeralFile::create(harness.conf, harness.tenant_id, timeline_id)?;
 
         let pos_foo = file.write_blob(b"foo")?;
         assert_eq!(b"foo", file.block_cursor().read_blob(pos_foo)?.as_slice());

pageserver/src/tenant/layer_map.rs

@@ -365,11 +365,27 @@ where
         }
     }
 
+    /// Whether a layer is L0.
+    ///
+    /// A layer is L0 if it's sparse enough that compacting it
+    /// is always more efficient than reimaging over it.
+    fn is_l0(layer: &L) -> bool {
+        // 1 TB key range, very conservative
+        let max_l0_range: i128 = 1024 * 1024 * 1024 * 1024 / (postgres_ffi::BLCKSZ as i128);
+        let kr = layer.get_key_range();
+        // TODO Use `collect_keyspace` result to only count used space covered.
+        //      Otherwise there's a chance that with the current `max_l0_range`
+        //      all layers will be classified as l0.
+        // Divide by 10 to avoid overflow
+        let size = kr.end.to_i128() / 10 - kr.start.to_i128() / 10;
+        size > max_l0_range / 10
+    }
+
     ///
     /// Insert an on-disk layer
     ///
     pub fn insert_historic(&mut self, layer: Arc<L>) {
-        if layer.get_key_range() == (Key::MIN..Key::MAX) {
+        if Self::is_l0(&layer) {
             self.l0_delta_layers.push(layer.clone());
         }
         self.historic_layers.insert(LayerRTreeObject::new(layer));
@@ -382,7 +398,7 @@ where
     /// This should be called when the corresponding file on disk has been deleted.
     ///
     pub fn remove_historic(&mut self, layer: Arc<L>) {
-        if layer.get_key_range() == (Key::MIN..Key::MAX) {
+        if Self::is_l0(&layer) {
             let len_before = self.l0_delta_layers.len();
 
             // FIXME: ptr_eq might fail to return true for 'dyn'
@@ -571,9 +587,7 @@ where
 
             // We ignore level0 delta layers. Unless the whole keyspace fits
             // into one partition
-            if !range_eq(key_range, &(Key::MIN..Key::MAX))
+            if !range_eq(key_range, &(Key::MIN..Key::MAX)) && Self::is_l0(&l) {
-                && range_eq(&l.get_key_range(), &(Key::MIN..Key::MAX))
-            {
                 continue;
             }
 

pageserver/src/tenant/remote_timeline_client.rs

@@ -135,7 +135,7 @@
 //! - Initiate upload queue with that [`IndexPart`].
 //! - Reschedule all lost operations by comparing the local filesystem state
 //!   and remote state as per [`IndexPart`]. This is done in
-//!   [`Timeline::setup_timeline`] and [`Timeline::reconcile_with_remote`].
+//!   [`Timeline::timeline_init_and_sync`] and [`Timeline::reconcile_with_remote`].
 //!
 //! Note that if we crash during file deletion between the index update
 //! that removes the file from the list of files, and deleting the remote file,
@@ -1064,7 +1064,7 @@ mod tests {
     // Test scheduling
     #[test]
     fn upload_scheduling() -> anyhow::Result<()> {
-        let harness = TenantHarness::create("upload_scheduling")?;
+        let harness = TenantHarness::new()?;
         let timeline_path = harness.timeline_path(&TIMELINE_ID);
         std::fs::create_dir_all(&timeline_path)?;
 

pageserver/src/tenant/size.rs

@@ -44,6 +44,116 @@ struct TimelineInputs {
     next_gc_cutoff: Lsn,
 }
 
+// Adjust BranchFrom sorting so that we always process ancestor
+// before descendants. This is needed to correctly calculate size of
+// descendant timelines.
+//
+// Note that we may have multiple BranchFroms at the same LSN, so we
+// need to sort them in the tree order.
+//
+// see updates_sort_with_branches_at_same_lsn test below
+fn sort_updates_in_tree_order(updates: Vec<Update>) -> anyhow::Result<Vec<Update>> {
+    let mut sorted_updates = Vec::with_capacity(updates.len());
+    let mut known_timelineids = HashSet::new();
+    let mut i = 0;
+    while i < updates.len() {
+        let curr_upd = &updates[i];
+
+        if let Command::BranchFrom(parent_id) = curr_upd.command {
+            let parent_id = match parent_id {
+                Some(parent_id) if known_timelineids.contains(&parent_id) => {
+                    // we have already processed ancestor
+                    // process this BranchFrom Update normally
+                    known_timelineids.insert(curr_upd.timeline_id);
+                    sorted_updates.push(*curr_upd);
+                    i += 1;
+                    continue;
+                }
+                None => {
+                    known_timelineids.insert(curr_upd.timeline_id);
+                    sorted_updates.push(*curr_upd);
+                    i += 1;
+                    continue;
+                }
+                Some(parent_id) => parent_id,
+            };
+
+            let mut j = i;
+
+            // we have not processed ancestor yet.
+            // there is a chance that it is at the same Lsn
+            if !known_timelineids.contains(&parent_id) {
+                let mut curr_lsn_branchfroms: HashMap<TimelineId, Vec<(TimelineId, usize)>> =
+                    HashMap::new();
+
+                // inspect all branchpoints at the same lsn
+                while j < updates.len() && updates[j].lsn == curr_upd.lsn {
+                    let lookahead_upd = &updates[j];
+                    j += 1;
+
+                    if let Command::BranchFrom(lookahead_parent_id) = lookahead_upd.command {
+                        match lookahead_parent_id {
+                            Some(lookahead_parent_id)
+                                if !known_timelineids.contains(&lookahead_parent_id) =>
+                            {
+                                // we have not processed ancestor yet
+                                // store it for later
+                                let es =
+                                    curr_lsn_branchfroms.entry(lookahead_parent_id).or_default();
+                                es.push((lookahead_upd.timeline_id, j));
+                            }
+                            _ => {
+                                // we have already processed ancestor
+                                // process this BranchFrom Update normally
+                                known_timelineids.insert(lookahead_upd.timeline_id);
+                                sorted_updates.push(*lookahead_upd);
+                            }
+                        }
+                    }
+                }
+
+                // process BranchFroms in the tree order
+                // check that we don't have a cycle if somet entry is orphan
+                // (this should not happen, but better to be safe)
+                let mut processed_some_entry = true;
+                while processed_some_entry {
+                    processed_some_entry = false;
+
+                    curr_lsn_branchfroms.retain(|parent_id, branchfroms| {
+                        if known_timelineids.contains(parent_id) {
+                            for (timeline_id, j) in branchfroms {
+                                known_timelineids.insert(*timeline_id);
+                                sorted_updates.push(updates[*j - 1]);
+                            }
+                            processed_some_entry = true;
+                            false
+                        } else {
+                            true
+                        }
+                    });
+                }
+
+                if !curr_lsn_branchfroms.is_empty() {
+                    // orphans are expected to be rare and transient between tenant reloads
+                    // for example, an broken ancestor without the child branch being broken.
+                    anyhow::bail!(
+                        "orphan branch(es) detected in BranchFroms: {curr_lsn_branchfroms:?}"
+                    );
+                }
+            }
+
+            assert!(j > i);
+            i = j;
+        } else {
+            // not a BranchFrom, keep the same order
+            sorted_updates.push(*curr_upd);
+            i += 1;
+        }
+    }
+
+    Ok(sorted_updates)
+}
+
 /// Gathers the inputs for the tenant sizing model.
 ///
 /// Tenant size does not consider the latest state, but only the state until next_gc_cutoff, which
@@ -267,7 +377,11 @@ pub(super) async fn gather_inputs(
     // for branch points, which come as multiple updates at the same LSN, the Command::Update
     // is needed before a branch is made out of that branch Command::BranchFrom. this is
     // handled by the variant order in `Command`.
+    //
     updates.sort_unstable();
+    // And another sort to handle Command::BranchFrom ordering
+    // in case when there are multiple branches at the same LSN.
+    let sorted_updates = sort_updates_in_tree_order(updates)?;
 
     let retention_period = match max_cutoff_distance {
         Some(max) => max.0,
@@ -277,7 +391,7 @@ pub(super) async fn gather_inputs(
     };
 
     Ok(ModelInputs {
-        updates,
+        updates: sorted_updates,
         retention_period,
         timeline_inputs,
     })
@@ -295,6 +409,7 @@ impl ModelInputs {
                 command: op,
                 timeline_id,
             } = update;
+
             let Lsn(now) = *lsn;
             match op {
                 Command::Update(sz) => {
@@ -304,7 +419,8 @@ impl ModelInputs {
                     storage.insert_point(&Some(*timeline_id), "".into(), now, None);
                 }
                 Command::BranchFrom(parent) => {
-                    storage.branch(parent, Some(*timeline_id));
+                    // This branch command may fail if it cannot find a parent to branch from.
+                    storage.branch(parent, Some(*timeline_id))?;
                 }
             }
         }
@@ -372,6 +488,7 @@ async fn calculate_logical_size(
 
     let size_res = timeline
         .spawn_ondemand_logical_size_calculation(lsn)
+        .instrument(info_span!("spawn_ondemand_logical_size_calculation"))
         .await?;
     Ok(TimelineAtLsnSizeResult(timeline, lsn, size_res))
 }
@@ -463,3 +580,137 @@ fn verify_size_for_multiple_branches() {
 
     assert_eq!(inputs.calculate().unwrap(), 36_409_872);
 }
+
+#[test]
+fn updates_sort_with_branches_at_same_lsn() {
+    use std::str::FromStr;
+    use Command::{BranchFrom, EndOfBranch};
+
+    macro_rules! lsn {
+        ($e:expr) => {
+            Lsn::from_str($e).unwrap()
+        };
+    }
+
+    let ids = [
+        TimelineId::from_str("00000000000000000000000000000000").unwrap(),
+        TimelineId::from_str("11111111111111111111111111111111").unwrap(),
+        TimelineId::from_str("22222222222222222222222222222222").unwrap(),
+        TimelineId::from_str("33333333333333333333333333333333").unwrap(),
+        TimelineId::from_str("44444444444444444444444444444444").unwrap(),
+    ];
+
+    // issue https://github.com/neondatabase/neon/issues/3179
+    let commands = vec![
+        Update {
+            lsn: lsn!("0/0"),
+            command: BranchFrom(None),
+            timeline_id: ids[0],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: Command::Update(25387008),
+            timeline_id: ids[0],
+        },
+        // next three are wrongly sorted, because
+        // ids[1] is branched from before ids[1] exists
+        // and ids[2] is branched from before ids[2] exists
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[1])),
+            timeline_id: ids[3],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[0])),
+            timeline_id: ids[2],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[2])),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CA85B8"),
+            command: Command::Update(28925952),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CD85B8"),
+            command: Command::Update(29024256),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CD85B8"),
+            command: BranchFrom(Some(ids[1])),
+            timeline_id: ids[4],
+        },
+        Update {
+            lsn: lsn!("0/22DCE70"),
+            command: Command::Update(32546816),
+            timeline_id: ids[3],
+        },
+        Update {
+            lsn: lsn!("0/230CE70"),
+            command: EndOfBranch,
+            timeline_id: ids[3],
+        },
+    ];
+
+    let expected = vec![
+        Update {
+            lsn: lsn!("0/0"),
+            command: BranchFrom(None),
+            timeline_id: ids[0],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: Command::Update(25387008),
+            timeline_id: ids[0],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[0])),
+            timeline_id: ids[2],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[2])),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/169AD58"),
+            command: BranchFrom(Some(ids[1])),
+            timeline_id: ids[3],
+        },
+        Update {
+            lsn: lsn!("0/1CA85B8"),
+            command: Command::Update(28925952),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CD85B8"),
+            command: Command::Update(29024256),
+            timeline_id: ids[1],
+        },
+        Update {
+            lsn: lsn!("0/1CD85B8"),
+            command: BranchFrom(Some(ids[1])),
+            timeline_id: ids[4],
+        },
+        Update {
+            lsn: lsn!("0/22DCE70"),
+            command: Command::Update(32546816),
+            timeline_id: ids[3],
+        },
+        Update {
+            lsn: lsn!("0/230CE70"),
+            command: EndOfBranch,
+            timeline_id: ids[3],
+        },
+    ];
+
+    let sorted_commands = sort_updates_in_tree_order(commands).unwrap();
+
+    assert_eq!(sorted_commands, expected);
+}

pageserver/src/tenant/timeline.rs

@@ -3,12 +3,12 @@
 use anyhow::{anyhow, bail, ensure, Context};
 use bytes::Bytes;
 use fail::fail_point;
-use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
 use pageserver_api::models::{
-    DownloadRemoteLayersTaskInfo, DownloadRemoteLayersTaskState, TimelineState,
+    DownloadRemoteLayersTaskInfo, DownloadRemoteLayersTaskSpawnRequest,
+    DownloadRemoteLayersTaskState, TimelineState,
 };
 use tokio::sync::{oneshot, watch, Semaphore, TryAcquireError};
 use tokio_util::sync::CancellationToken;
@@ -3116,6 +3116,7 @@ impl Timeline {
 
     pub async fn spawn_download_all_remote_layers(
         self: Arc<Self>,
+        request: DownloadRemoteLayersTaskSpawnRequest,
     ) -> Result<DownloadRemoteLayersTaskInfo, DownloadRemoteLayersTaskInfo> {
         let mut status_guard = self.download_all_remote_layers_task_info.write().unwrap();
         if let Some(st) = &*status_guard {
@@ -3139,7 +3140,7 @@ impl Timeline {
             "download all remote layers task",
             false,
             async move {
-                self_clone.download_all_remote_layers().await;
+                self_clone.download_all_remote_layers(request).await;
                 let mut status_guard = self_clone.download_all_remote_layers_task_info.write().unwrap();
                  match &mut *status_guard {
                     None => {
@@ -3171,15 +3172,23 @@ impl Timeline {
         Ok(initial_info)
     }
 
-    async fn download_all_remote_layers(self: &Arc<Self>) {
+    async fn download_all_remote_layers(
-        let mut downloads: FuturesUnordered<_> = {
+        self: &Arc<Self>,
+        request: DownloadRemoteLayersTaskSpawnRequest,
+    ) {
+        let mut downloads = Vec::new();
+        {
             let layers = self.layers.read().unwrap();
             layers
                 .iter_historic_layers()
                 .filter_map(|l| l.downcast_remote_layer())
                 .map(|l| self.download_remote_layer(l))
-                .collect()
+                .for_each(|dl| downloads.push(dl))
-        };
+        }
+        let total_layer_count = downloads.len();
+        // limit download concurrency as specified in request
+        let downloads = futures::stream::iter(downloads);
+        let mut downloads = downloads.buffer_unordered(request.max_concurrent_downloads.get());
 
         macro_rules! lock_status {
             ($st:ident) => {
@@ -3200,7 +3209,7 @@ impl Timeline {
 
         {
             lock_status!(st);
-            st.total_layer_count = downloads.len().try_into().unwrap();
+            st.total_layer_count = total_layer_count as u64;
         }
         loop {
             tokio::select! {

pageserver/src/virtual_file.rs

@@ -525,12 +525,13 @@ mod tests {
         })
     }
 
-    fn test_files<OF, FD>(testname: &str, openfunc: OF) -> Result<(), Error>
+    fn test_files<OF, FD>(test_name: &str, openfunc: OF) -> Result<(), Error>
     where
         FD: Read + Write + Seek + FileExt,
         OF: Fn(&Path, &OpenOptions) -> Result<FD, std::io::Error>,
     {
-        let testdir = crate::config::PageServerConf::test_repo_dir(testname);
+        let temp_repo_dir = tempfile::tempdir()?;
+        let testdir = temp_repo_dir.path().join(test_name);
         std::fs::create_dir_all(&testdir)?;
 
         let path_a = testdir.join("file_a");
@@ -632,7 +633,8 @@ mod tests {
         const THREADS: usize = 100;
         const SAMPLE: [u8; SIZE] = [0xADu8; SIZE];
 
-        let testdir = crate::config::PageServerConf::test_repo_dir("vfile_concurrency");
+        let temp_repo_dir = tempfile::tempdir()?;
+        let testdir = temp_repo_dir.path().join("vfile_concurrency");
         std::fs::create_dir_all(&testdir)?;
 
         // Create a test file.

pageserver/src/walingest.rs

@@ -1146,7 +1146,8 @@ mod tests {
 
     #[tokio::test]
     async fn test_relsize() -> Result<()> {
-        let tenant = TenantHarness::create("test_relsize")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = create_test_timeline(&tenant, TIMELINE_ID, DEFAULT_PG_VERSION)?;
         let mut walingest = init_walingest_test(&tline).await?;
 
@@ -1323,7 +1324,8 @@ mod tests {
     // and then created it again within the same layer.
     #[tokio::test]
     async fn test_drop_extend() -> Result<()> {
-        let tenant = TenantHarness::create("test_drop_extend")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = create_test_timeline(&tenant, TIMELINE_ID, DEFAULT_PG_VERSION)?;
         let mut walingest = init_walingest_test(&tline).await?;
 
@@ -1376,7 +1378,8 @@ mod tests {
     // and then extended it again within the same layer.
     #[tokio::test]
     async fn test_truncate_extend() -> Result<()> {
-        let tenant = TenantHarness::create("test_truncate_extend")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = create_test_timeline(&tenant, TIMELINE_ID, DEFAULT_PG_VERSION)?;
         let mut walingest = init_walingest_test(&tline).await?;
 
@@ -1497,7 +1500,8 @@ mod tests {
     /// split into multiple 1 GB segments in Postgres.
     #[tokio::test]
     async fn test_large_rel() -> Result<()> {
-        let tenant = TenantHarness::create("test_large_rel")?.load().await;
+        let harness = TenantHarness::new()?;
+        let tenant = harness.load().await;
         let tline = create_test_timeline(&tenant, TIMELINE_ID, DEFAULT_PG_VERSION)?;
         let mut walingest = init_walingest_test(&tline).await?;
 

pageserver/src/walreceiver.rs

@@ -116,6 +116,10 @@ impl<E: Clone> TaskHandle<E> {
         let join_handle = WALRECEIVER_RUNTIME.spawn(async move {
             events_sender.send(TaskStateUpdate::Started).ok();
             task(events_sender, cancellation_clone).await
+            // events_sender is dropped at some point during the .await above.
+            // But the task is still running on WALRECEIVER_RUNTIME.
+            // That is the window when `!jh.is_finished()`
+            // is true inside `fn next_task_event()` below.
         });
 
         TaskHandle {
@@ -132,7 +136,23 @@ impl<E: Clone> TaskHandle<E> {
                 TaskEvent::End(match self.join_handle.as_mut() {
                     Some(jh) => {
                         if !jh.is_finished() {
-                            warn!("sender is dropped while join handle is still alive");
+                            // Barring any implementation errors in this module, we can
+                            // only arrive here while the task that executes the future
+                            // passed to `Self::spawn()` is still execution. Cf the comment
+                            // in Self::spawn().
+                            //
+                            // This was logging at warning level in earlier versions, presumably
+                            // to leave some breadcrumbs in case we had an implementation
+                            // error that would would make us get stuck in `jh.await`.
+                            //
+                            // There hasn't been such a bug so far.
+                            // But in a busy system, e.g., during pageserver restart,
+                            // we arrive here often enough that the warning-level logs
+                            // became a distraction.
+                            // So, tone them down to info-level.
+                            //
+                            // XXX: rewrite this module to eliminate the race condition.
+                            info!("sender is dropped while join handle is still alive");
                         }
 
                         let res = jh

pageserver/src/walreceiver/connection_manager.rs

@@ -846,7 +846,7 @@ mod tests {
 
     #[tokio::test]
     async fn no_connection_no_candidate() -> anyhow::Result<()> {
-        let harness = TenantHarness::create("no_connection_no_candidate")?;
+        let harness = TenantHarness::new()?;
         let mut state = dummy_state(&harness).await;
         let now = Utc::now().naive_utc();
 
@@ -879,7 +879,7 @@ mod tests {
 
     #[tokio::test]
     async fn connection_no_candidate() -> anyhow::Result<()> {
-        let harness = TenantHarness::create("connection_no_candidate")?;
+        let harness = TenantHarness::new()?;
         let mut state = dummy_state(&harness).await;
         let now = Utc::now().naive_utc();
 
@@ -942,7 +942,7 @@ mod tests {
 
     #[tokio::test]
     async fn no_connection_candidate() -> anyhow::Result<()> {
-        let harness = TenantHarness::create("no_connection_candidate")?;
+        let harness = TenantHarness::new()?;
         let mut state = dummy_state(&harness).await;
         let now = Utc::now().naive_utc();
 
@@ -1001,7 +1001,7 @@ mod tests {
 
     #[tokio::test]
     async fn candidate_with_many_connection_failures() -> anyhow::Result<()> {
-        let harness = TenantHarness::create("candidate_with_many_connection_failures")?;
+        let harness = TenantHarness::new()?;
         let mut state = dummy_state(&harness).await;
         let now = Utc::now().naive_utc();
 
@@ -1041,7 +1041,7 @@ mod tests {
 
     #[tokio::test]
     async fn lsn_wal_over_threshhold_current_candidate() -> anyhow::Result<()> {
-        let harness = TenantHarness::create("lsn_wal_over_threshcurrent_candidate")?;
+        let harness = TenantHarness::new()?;
         let mut state = dummy_state(&harness).await;
         let current_lsn = Lsn(100_000).align();
         let now = Utc::now().naive_utc();
@@ -1105,7 +1105,7 @@ mod tests {
 
     #[tokio::test]
     async fn timeout_connection_threshhold_current_candidate() -> anyhow::Result<()> {
-        let harness = TenantHarness::create("timeout_connection_threshhold_current_candidate")?;
+        let harness = TenantHarness::new()?;
         let mut state = dummy_state(&harness).await;
         let current_lsn = Lsn(100_000).align();
         let now = Utc::now().naive_utc();
@@ -1166,7 +1166,7 @@ mod tests {
 
     #[tokio::test]
     async fn timeout_wal_over_threshhold_current_candidate() -> anyhow::Result<()> {
-        let harness = TenantHarness::create("timeout_wal_over_threshhold_current_candidate")?;
+        let harness = TenantHarness::new()?;
         let mut state = dummy_state(&harness).await;
         let current_lsn = Lsn(100_000).align();
         let new_lsn = Lsn(100_100).align();
@@ -1232,7 +1232,7 @@ mod tests {
 
     const DUMMY_SAFEKEEPER_HOST: &str = "safekeeper_connstr";
 
-    async fn dummy_state(harness: &TenantHarness<'_>) -> WalreceiverState {
+    async fn dummy_state(harness: &TenantHarness) -> WalreceiverState {
         WalreceiverState {
             id: TenantTimelineId {
                 tenant_id: harness.tenant_id,

pageserver/src/walredo.rs

@@ -626,24 +626,20 @@ impl PostgresRedoProcess {
 
         // Create empty data directory for wal-redo postgres, deleting old one first.
         if datadir.exists() {
-            info!(
+            info!("old temporary datadir {datadir:?} exists, removing");
-                "old temporary datadir {} exists, removing",
+            fs::remove_dir_all(&datadir).map_err(|e| {
-                datadir.display()
+                Error::new(
-            );
+                    e.kind(),
-            fs::remove_dir_all(&datadir)?;
+                    format!("Old temporary dir {datadir:?} removal failure: {e}"),
+                )
+            })?;
         }
-        let pg_bin_dir_path = conf.pg_bin_dir(pg_version).map_err(|e| {
+        let pg_bin_dir_path = conf
-            Error::new(
+            .pg_bin_dir(pg_version)
-                ErrorKind::Other,
+            .map_err(|e| Error::new(ErrorKind::Other, format!("incorrect pg_bin_dir path: {e}")))?;
-                format!("incorrect pg_bin_dir path: {}", e),
+        let pg_lib_dir_path = conf
-            )
+            .pg_lib_dir(pg_version)
-        })?;
+            .map_err(|e| Error::new(ErrorKind::Other, format!("incorrect pg_lib_dir path: {e}")))?;
-        let pg_lib_dir_path = conf.pg_lib_dir(pg_version).map_err(|e| {
-            Error::new(
-                ErrorKind::Other,
-                format!("incorrect pg_lib_dir path: {}", e),
-            )
-        })?;
 
         info!("running initdb in {}", datadir.display());
         let initdb = Command::new(pg_bin_dir_path.join("initdb"))
@@ -1010,3 +1006,110 @@ fn build_get_page_msg(tag: BufferTag, buf: &mut Vec<u8>) {
     tag.ser_into(buf)
         .expect("serialize BufferTag should always succeed");
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{PostgresRedoManager, WalRedoManager};
+    use crate::repository::Key;
+    use crate::{config::PageServerConf, walrecord::NeonWalRecord};
+    use bytes::Bytes;
+    use std::str::FromStr;
+    use utils::{id::TenantId, lsn::Lsn};
+
+    #[test]
+    fn short_v14_redo() {
+        let expected = std::fs::read("fixtures/short_v14_redo.page").unwrap();
+
+        let h = RedoHarness::new().unwrap();
+
+        let page = h
+            .manager
+            .request_redo(
+                Key {
+                    field1: 0,
+                    field2: 1663,
+                    field3: 13010,
+                    field4: 1259,
+                    field5: 0,
+                    field6: 0,
+                },
+                Lsn::from_str("0/16E2408").unwrap(),
+                None,
+                short_records(),
+                14,
+            )
+            .unwrap();
+
+        assert_eq!(&expected, &*page);
+    }
+
+    #[test]
+    fn short_v14_fails_for_wrong_key_but_returns_zero_page() {
+        let h = RedoHarness::new().unwrap();
+
+        let page = h
+            .manager
+            .request_redo(
+                Key {
+                    field1: 0,
+                    field2: 1663,
+                    // key should be 13010
+                    field3: 13130,
+                    field4: 1259,
+                    field5: 0,
+                    field6: 0,
+                },
+                Lsn::from_str("0/16E2408").unwrap(),
+                None,
+                short_records(),
+                14,
+            )
+            .unwrap();
+
+        // TODO: there will be some stderr printout, which is forwarded to tracing that could
+        // perhaps be captured as long as it's in the same thread.
+        assert_eq!(page, crate::ZERO_PAGE);
+    }
+
+    #[allow(clippy::octal_escapes)]
+    fn short_records() -> Vec<(Lsn, NeonWalRecord)> {
+        vec![
+            (
+                Lsn::from_str("0/16A9388").unwrap(),
+                NeonWalRecord::Postgres {
+                    will_init: true,
+                    rec: Bytes::from_static(b"j\x03\0\0\0\x04\0\0\xe8\x7fj\x01\0\0\0\0\0\n\0\0\xd0\x16\x13Y\0\x10\0\04\x03\xd4\0\x05\x7f\x06\0\0\xd22\0\0\xeb\x04\0\0\0\0\0\0\xff\x03\0\0\0\0\x80\xeca\x01\0\0\x01\0\xd4\0\xa0\x1d\0 \x04 \0\0\0\0/\0\x01\0\xa0\x9dX\x01\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0.\0\x01\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\00\x9f\x9a\x01P\x9e\xb2\x01\0\x04\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x02\0!\0\x01\x08 \xff\xff\xff?\0\0\0\0\0\0@\0\0another_table\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x98\x08\0\0\x02@\0\0\0\0\0\0\n\0\0\0\x02\0\0\0\0@\0\0\0\0\0\0\0\0\0\0\0\0\x80\xbf\0\0\0\0\0\0\0\0\0\0pr\x01\0\0\0\0\0\0\0\0\x01d\0\0\0\0\0\0\x04\0\0\x01\0\0\0\0\0\0\0\x0c\x02\0\0\0\0\0\0\0\0\0\0\0\0\0\0/\0!\x80\x03+ \xff\xff\xff\x7f\0\0\0\0\0\xdf\x04\0\0pg_type\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x0b\0\0\0G\0\0\0\0\0\0\0\n\0\0\0\x02\0\0\0\0\0\0\0\0\0\0\0\x0e\0\0\0\0@\x16D\x0e\0\0\0K\x10\0\0\x01\0pr \0\0\0\0\0\0\0\0\x01n\0\0\0\0\0\xd6\x02\0\0\x01\0\0\0[\x01\0\0\0\0\0\0\0\t\x04\0\0\x02\0\0\0\x01\0\0\0\n\0\0\0\n\0\0\0\x7f\0\0\0\0\0\0\0\n\0\0\0\x02\0\0\0\0\0\0C\x01\0\0\x15\x01\0\0\0\0\0\0\0\0\0\0\0\0\0\0.\0!\x80\x03+ \xff\xff\xff\x7f\0\0\0\0\0;\n\0\0pg_statistic\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x0b\0\0\0\xfd.\0\0\0\0\0\0\n\0\0\0\x02\0\0\0;\n\0\0\0\0\0\0\x13\0\0\0\0\0\xcbC\x13\0\0\0\x18\x0b\0\0\x01\0pr\x1f\0\0\0\0\0\0\0\0\x01n\0\0\0\0\0\xd6\x02\0\0\x01\0\0\0C\x01\0\0\0\0\0\0\0\t\x04\0\0\x01\0\0\0\x01\0\0\0\n\0\0\0\n\0\0\0\x7f\0\0\0\0\0\0\x02\0\x01")
+                }
+            ),
+            (
+                Lsn::from_str("0/16D4080").unwrap(),
+                NeonWalRecord::Postgres {
+                    will_init: false,
+                    rec: Bytes::from_static(b"\xbc\0\0\0\0\0\0\0h?m\x01\0\0\0\0p\n\0\09\x08\xa3\xea\0 \x8c\0\x7f\x06\0\0\xd22\0\0\xeb\x04\0\0\0\0\0\0\xff\x02\0@\0\0another_table\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\x98\x08\0\0\x02@\0\0\0\0\0\0\n\0\0\0\x02\0\0\0\0@\0\0\0\0\0\0\x05\0\0\0\0@zD\x05\0\0\0\0\0\0\0\0\0pr\x01\0\0\0\0\0\0\0\0\x01d\0\0\0\0\0\0\x04\0\0\x01\0\0\0\x02\0")
+                }
+            )
+        ]
+    }
+
+    struct RedoHarness {
+        // underscored because unused, except for removal at drop
+        _repo_dir: tempfile::TempDir,
+        manager: PostgresRedoManager,
+    }
+
+    impl RedoHarness {
+        fn new() -> anyhow::Result<Self> {
+            let repo_dir = tempfile::tempdir()?;
+            let conf = PageServerConf::dummy_conf(repo_dir.path().to_path_buf());
+            let conf = Box::leak(Box::new(conf));
+            let tenant_id = TenantId::generate();
+
+            let manager = PostgresRedoManager::new(conf, tenant_id);
+
+            Ok(RedoHarness {
+                _repo_dir: repo_dir,
+                manager,
+            })
+        }
+    }
+}

pgxn/neon/pagestore_client.h

@@ -52,7 +52,7 @@ typedef struct
 #define NEON_TAG "[NEON_SMGR] "
 #define neon_log(tag, fmt, ...) ereport(tag,                                  \
 										(errmsg(NEON_TAG fmt, ##__VA_ARGS__), \
-										 errhidestmt(true), errhidecontext(true)))
+										 errhidestmt(true), errhidecontext(true), internalerrposition(0)))
 
 /*
  * supertype of all the Neon*Request structs below

pgxn/neon/pagestore_smgr.c

@@ -287,12 +287,13 @@ compact_prefetch_buffers(void)
 
 	/*
 	 * Here we have established:
-	 * slots < search_ring_index may be unused (not scanned)
+	 *   slots < search_ring_index have an unknown state (not scanned)
-	 * slots >= search_ring_index and <= empty_ring_index are unused
+	 *   slots >= search_ring_index and <= empty_ring_index are unused
-	 * slots > empty_ring_index are in use, or outside our buffer's range.
+	 *   slots > empty_ring_index are in use, or outside our buffer's range.
+	 * ... unless search_ring_index <= ring_last
 	 * 
 	 * Therefore, there is a gap of at least one unused items between
-	 * search_ring_index and empty_ring_index, which grows as we hit
+	 * search_ring_index and empty_ring_index (both inclusive), which grows as we hit
 	 * more unused items while moving backwards through the array.
 	 */
 
@@ -302,6 +303,7 @@ compact_prefetch_buffers(void)
 		PrefetchRequest *target_slot;
 		bool		found;
 
+		/* update search index to an unprocessed entry */
 		search_ring_index--;
 
 		source_slot = GetPrfSlot(search_ring_index);
@@ -309,6 +311,7 @@ compact_prefetch_buffers(void)
 		if (source_slot->status == PRFS_UNUSED)
 			continue;
 
+		/* slot is used -- start moving slot */
 		target_slot = GetPrfSlot(empty_ring_index);
 
 		Assert(source_slot->status == PRFS_RECEIVED);
@@ -328,16 +331,22 @@ compact_prefetch_buffers(void)
 		/* Adjust the location of our known-empty slot */
 		empty_ring_index--;
 
+		/* empty the moved slot */
 		source_slot->status = PRFS_UNUSED;
 		source_slot->buftag = (BufferTag) {0};
 		source_slot->response = NULL;
 		source_slot->my_ring_index = 0;
 		source_slot->effective_request_lsn = 0;
 
+		/* update bookkeeping */
 		n_moved++;
 	}
 
-	if (MyPState->ring_last != empty_ring_index)
+	/*
+	 * Only when we've moved slots we can expect trailing unused slots,
+	 * so only then we clean up trailing unused slots.
+	 */
+	if (n_moved > 0)
 	{
 		prefetch_cleanup_trailing_unused();
 		return true;

proxy/Cargo.toml

@@ -1,58 +1,63 @@
 [package]
 name = "proxy"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-anyhow = "1.0"
+anyhow.workspace = true
-atty = "0.2.14"
+atty.workspace = true
-base64 = "0.13.0"
+base64.workspace = true
-bstr = "1.0"
+bstr.workspace = true
-bytes = { version = "1.0.1", features = ['serde'] }
+bytes = {workspace = true, features = ['serde'] }
-clap = "4.0"
+clap.workspace = true
-futures = "0.3.13"
+chrono.workspace = true
-git-version = "0.3.5"
+consumption_metrics.workspace = true
-hashbrown = "0.12"
+futures.workspace = true
-hex = "0.4.3"
+git-version.workspace = true
-hmac = "0.12.1"
+hashbrown.workspace = true
-hyper = "0.14"
+hex.workspace = true
-hyper-tungstenite = "0.8.1"
+hmac.workspace = true
-itertools = "0.10.3"
+hyper.workspace = true
-md5 = "0.7.0"
+hyper-tungstenite.workspace = true
-once_cell = "1.13.0"
+itertools.workspace = true
-parking_lot = "0.12"
+md5.workspace = true
-pin-project-lite = "0.2.7"
+once_cell.workspace = true
-rand = "0.8.3"
+parking_lot.workspace = true
-regex = "1.4.5"
+pin-project-lite.workspace = true
-reqwest = { version = "0.11", default-features = false, features = [ "json", "rustls-tls" ] }
+rand.workspace = true
-routerify = "3"
+regex.workspace = true
-rustls = "0.20.0"
+reqwest = { workspace = true, features = [ "json" ] }
-rustls-pemfile = "1"
+routerify.workspace = true
-scopeguard = "1.1.0"
+rustls.workspace = true
-serde = "1"
+rustls-pemfile.workspace = true
-serde_json = "1"
+scopeguard.workspace = true
-sha2 = "0.10.2"
+serde.workspace = true
-socket2 = "0.4.4"
+serde_json.workspace = true
-thiserror = "1.0.30"
+sha2.workspace = true
-tokio = { version = "1.17", features = ["macros"] }
+socket2.workspace = true
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+thiserror.workspace = true
-tokio-rustls = "0.23.0"
+tokio.workspace = true
-tls-listener = { version = "0.5.1", features = ["rustls", "hyper-h1"] }
+tokio-postgres.workspace = true
-tracing = "0.1.36"
+tokio-rustls.workspace = true
-tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+tls-listener.workspace = true
-url = "2.2.2"
+tracing.workspace = true
-uuid = { version = "1.2", features = ["v4", "serde"] }
+tracing-subscriber.workspace = true
-webpki-roots = "0.22.5"
+url.workspace = true
-x509-parser = "0.14"
+uuid.workspace = true
+webpki-roots.workspace = true
+x509-parser.workspace = true
+metrics.workspace = true
+pq_proto.workspace = true
+utils.workspace = true
+prometheus.workspace = true
+humantime.workspace = true
+hostname.workspace = true
 
-metrics = { path = "../libs/metrics" }
+workspace_hack.workspace = true
-pq_proto = { path = "../libs/pq_proto" }
-utils = { path = "../libs/utils" }
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
 
 [dev-dependencies]
-async-trait = "0.1"
+async-trait.workspace = true
-rcgen = "0.10"
+rcgen.workspace = true
-rstest = "0.15"
+rstest.workspace = true
-tokio-postgres-rustls = "0.9.0"
+tokio-postgres-rustls.workspace = true

proxy/src/auth.rs

@@ -42,9 +42,9 @@ pub enum AuthErrorImpl {
     MalformedPassword(&'static str),
 
     #[error(
-        "Project ID is not specified. \
+        "Endpoint ID is not specified. \
         Either please upgrade the postgres client library (libpq) for SNI support \
-        or pass the project ID (first part of the domain name) as a parameter: '?options=project%3D<project-id>'. \
+        or pass the endpoint ID (first part of the domain name) as a parameter: '?options=project%3D<endpoint-id>'. \
         See more at https://neon.tech/sni"
     )]
     MissingProjectName,

proxy/src/config.rs

@@ -5,6 +5,12 @@ use std::sync::Arc;
 pub struct ProxyConfig {
     pub tls_config: Option<TlsConfig>,
     pub auth_backend: auth::BackendType<'static, ()>,
+    pub metric_collection_config: Option<MetricCollectionConfig>,
+}
+
+pub struct MetricCollectionConfig {
+    pub endpoint: reqwest::Url,
+    pub interval: std::time::Duration,
 }
 
 pub struct TlsConfig {

proxy/src/http/websocket.rs

@@ -13,7 +13,7 @@ use std::convert::Infallible;
 use std::future::ready;
 use std::pin::Pin;
 use std::sync::Arc;
-use std::task::{Context, Poll};
+use std::task::{ready, Context, Poll};
 use tls_listener::TlsListener;
 
 use tokio::io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
@@ -104,10 +104,9 @@ impl AsyncRead for WebSocketRW {
             return Poll::Ready(Ok(()));
         }
 
-        let inner_buf = match self.as_mut().poll_fill_buf(cx) {
+        let inner_buf = match ready!(self.as_mut().poll_fill_buf(cx)) {
-            Poll::Ready(Ok(buf)) => buf,
+            Ok(buf) => buf,
-            Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+            Err(err) => return Poll::Ready(Err(err)),
-            Poll::Pending => return Poll::Pending,
         };
         let len = std::cmp::min(inner_buf.len(), buf.remaining());
         buf.put_slice(&inner_buf[..len]);
@@ -124,8 +123,8 @@ impl AsyncBufRead for WebSocketRW {
                 let buf = self.project().chunk.as_ref().unwrap().chunk();
                 return Poll::Ready(Ok(buf));
             } else {
-                match self.as_mut().project().stream.poll_next(cx) {
+                match ready!(self.as_mut().project().stream.poll_next(cx)) {
-                    Poll::Ready(Some(Ok(message))) => match message {
+                    Some(Ok(message)) => match message {
                         Message::Text(_) => {}
                         Message::Binary(chunk) => {
                             *self.as_mut().project().chunk = Some(Bytes::from(chunk));
@@ -142,9 +141,8 @@ impl AsyncBufRead for WebSocketRW {
                             unreachable!();
                         }
                     },
-                    Poll::Ready(Some(Err(err))) => return Poll::Ready(Err(ws_err_into(err))),
+                    Some(Err(err)) => return Poll::Ready(Err(ws_err_into(err))),
-                    Poll::Ready(None) => return Poll::Ready(Ok(&[])),
+                    None => return Poll::Ready(Ok(&[])),
-                    Poll::Pending => return Poll::Pending,
                 }
             }
         }

proxy/src/main.rs

@@ -11,6 +11,7 @@ mod config;
 mod console;
 mod error;
 mod http;
+mod metrics;
 mod mgmt;
 mod parse;
 mod proxy;
@@ -20,14 +21,14 @@ mod stream;
 mod url;
 mod waiters;
 
+use ::metrics::set_build_info_metric;
 use anyhow::{bail, Context};
 use clap::{self, Arg};
 use config::ProxyConfig;
 use futures::FutureExt;
-use metrics::set_build_info_metric;
 use std::{borrow::Cow, future::Future, net::SocketAddr};
 use tokio::{net::TcpListener, task::JoinError};
-use tracing::info;
+use tracing::{info, info_span, Instrument};
 use utils::project_git_version;
 use utils::sentry_init::{init_sentry, release_name};
 
@@ -65,6 +66,22 @@ async fn main() -> anyhow::Result<()> {
     let mgmt_address: SocketAddr = arg_matches.get_one::<String>("mgmt").unwrap().parse()?;
     let http_address: SocketAddr = arg_matches.get_one::<String>("http").unwrap().parse()?;
 
+    let metric_collection_config = match
+    (
+        arg_matches.get_one::<String>("metric-collection-endpoint"),
+        arg_matches.get_one::<String>("metric-collection-interval"),
+    ) {
+
+        (Some(endpoint), Some(interval)) => {
+            Some(config::MetricCollectionConfig {
+                endpoint: endpoint.parse()?,
+                interval: humantime::parse_duration(interval)?,
+            })
+        }
+        (None, None) => None,
+        _ => bail!("either both or neither metric-collection-endpoint and metric-collection-interval must be specified"),
+    };
+
     let auth_backend = match arg_matches
         .get_one::<String>("auth-backend")
         .unwrap()
@@ -95,6 +112,7 @@ async fn main() -> anyhow::Result<()> {
     let config: &ProxyConfig = Box::leak(Box::new(ProxyConfig {
         tls_config,
         auth_backend,
+        metric_collection_config,
     }));
 
     info!("Version: {GIT_VERSION}");
@@ -126,6 +144,21 @@ async fn main() -> anyhow::Result<()> {
         )));
     }
 
+    if let Some(metric_collection_config) = &config.metric_collection_config {
+        let hostname = hostname::get()?
+            .into_string()
+            .map_err(|e| anyhow::anyhow!("failed to get hostname {e:?}"))?;
+
+        tasks.push(tokio::spawn(
+            metrics::collect_metrics(
+                &metric_collection_config.endpoint,
+                metric_collection_config.interval,
+                hostname,
+            )
+            .instrument(info_span!("collect_metrics")),
+        ));
+    }
+
     let tasks = tasks.into_iter().map(flatten_err);
 
     set_build_info_metric(GIT_VERSION);
@@ -199,6 +232,16 @@ fn cli() -> clap::Command {
                 .alias("ssl-cert") // backwards compatibility
                 .help("path to TLS cert for client postgres connections"),
         )
+        .arg(
+            Arg::new("metric-collection-endpoint")
+                .long("metric-collection-endpoint")
+                .help("metric collection HTTP endpoint"),
+        )
+        .arg(
+            Arg::new("metric-collection-interval")
+                .long("metric-collection-interval")
+                .help("metric collection interval"),
+        )
 }
 
 #[test]

proxy/src/metrics.rs

@@ -0,0 +1,196 @@
+//!
+//! Periodically collect proxy consumption metrics
+//! and push them to a HTTP endpoint.
+//!
+use chrono::{DateTime, Utc};
+use consumption_metrics::{idempotency_key, Event, EventChunk, EventType, CHUNK_SIZE};
+use serde::Serialize;
+use std::{collections::HashMap, time::Duration};
+use tracing::{debug, error, log::info, trace};
+
+const PROXY_IO_BYTES_PER_CLIENT: &str = "proxy_io_bytes_per_client";
+
+///
+/// Key that uniquely identifies the object, this metric describes.
+/// Currently, endpoint_id is enough, but this may change later,
+/// so keep it in a named struct.
+///
+/// Both the proxy and the ingestion endpoint will live in the same region (or cell)
+/// so while the project-id is unique across regions the whole pipeline will work correctly
+/// because we enrich the event with project_id in the control-plane endpoint.
+///
+#[derive(Eq, Hash, PartialEq, Serialize)]
+pub struct Ids {
+    pub endpoint_id: String,
+}
+
+pub async fn collect_metrics(
+    metric_collection_endpoint: &reqwest::Url,
+    metric_collection_interval: Duration,
+    hostname: String,
+) -> anyhow::Result<()> {
+    scopeguard::defer! {
+        info!("collect_metrics has shut down");
+    }
+
+    let mut ticker = tokio::time::interval(metric_collection_interval);
+
+    info!(
+        "starting collect_metrics. metric_collection_endpoint: {}",
+        metric_collection_endpoint
+    );
+
+    // define client here to reuse it for all requests
+    let client = reqwest::Client::new();
+    let mut cached_metrics: HashMap<Ids, (u64, DateTime<Utc>)> = HashMap::new();
+
+    loop {
+        tokio::select! {
+            _ = ticker.tick() => {
+
+                match collect_metrics_iteration(&client, &mut cached_metrics, metric_collection_endpoint, hostname.clone()).await
+                {
+                    Err(e) => {
+                        error!("Failed to send consumption metrics: {} ", e);
+                    },
+                    Ok(_) => { trace!("collect_metrics_iteration completed successfully") },
+                }
+            }
+        }
+    }
+}
+
+pub fn gather_proxy_io_bytes_per_client() -> Vec<(Ids, (u64, DateTime<Utc>))> {
+    let mut current_metrics: Vec<(Ids, (u64, DateTime<Utc>))> = Vec::new();
+    let metrics = prometheus::default_registry().gather();
+
+    for m in metrics {
+        if m.get_name() == "proxy_io_bytes_per_client" {
+            for ms in m.get_metric() {
+                let direction = ms
+                    .get_label()
+                    .iter()
+                    .find(|l| l.get_name() == "direction")
+                    .unwrap()
+                    .get_value();
+
+                // Only collect metric for outbound traffic
+                if direction == "tx" {
+                    let endpoint_id = ms
+                        .get_label()
+                        .iter()
+                        .find(|l| l.get_name() == "endpoint_id")
+                        .unwrap()
+                        .get_value();
+                    let value = ms.get_counter().get_value() as u64;
+
+                    debug!("endpoint_id:val - {}: {}", endpoint_id, value);
+                    current_metrics.push((
+                        Ids {
+                            endpoint_id: endpoint_id.to_string(),
+                        },
+                        (value, Utc::now()),
+                    ));
+                }
+            }
+        }
+    }
+
+    current_metrics
+}
+
+pub async fn collect_metrics_iteration(
+    client: &reqwest::Client,
+    cached_metrics: &mut HashMap<Ids, (u64, DateTime<Utc>)>,
+    metric_collection_endpoint: &reqwest::Url,
+    hostname: String,
+) -> anyhow::Result<()> {
+    info!(
+        "starting collect_metrics_iteration. metric_collection_endpoint: {}",
+        metric_collection_endpoint
+    );
+
+    let current_metrics = gather_proxy_io_bytes_per_client();
+
+    let metrics_to_send: Vec<Event<Ids>> = current_metrics
+        .iter()
+        .filter_map(|(curr_key, (curr_val, curr_time))| {
+            let mut start_time = *curr_time;
+            let mut value = *curr_val;
+
+            if let Some((prev_val, prev_time)) = cached_metrics.get(curr_key) {
+                // Only send metrics updates if the metric has changed
+                if curr_val - prev_val > 0 {
+                    value = curr_val - prev_val;
+                    start_time = *prev_time;
+                } else {
+                    return None;
+                }
+            };
+
+            Some(Event {
+                kind: EventType::Incremental {
+                    start_time,
+                    stop_time: *curr_time,
+                },
+                metric: PROXY_IO_BYTES_PER_CLIENT,
+                idempotency_key: idempotency_key(hostname.clone()),
+                value,
+                extra: Ids {
+                    endpoint_id: curr_key.endpoint_id.clone(),
+                },
+            })
+        })
+        .collect();
+
+    if metrics_to_send.is_empty() {
+        trace!("no new metrics to send");
+        return Ok(());
+    }
+
+    // Send metrics.
+    // Split into chunks of 1000 metrics to avoid exceeding the max request size
+    for chunk in metrics_to_send.chunks(CHUNK_SIZE) {
+        let chunk_json = serde_json::value::to_raw_value(&EventChunk { events: chunk })
+            .expect("ProxyConsumptionMetric should not fail serialization");
+
+        let res = client
+            .post(metric_collection_endpoint.clone())
+            .json(&chunk_json)
+            .send()
+            .await;
+
+        let res = match res {
+            Ok(x) => x,
+            Err(err) => {
+                error!("failed to send metrics: {:?}", err);
+                continue;
+            }
+        };
+
+        if res.status().is_success() {
+            // update cached metrics after they were sent successfully
+            for send_metric in chunk {
+                let stop_time = match send_metric.kind {
+                    EventType::Incremental { stop_time, .. } => stop_time,
+                    _ => unreachable!(),
+                };
+
+                cached_metrics
+                    .entry(Ids {
+                        endpoint_id: send_metric.extra.endpoint_id.clone(),
+                    })
+                    // update cached value (add delta) and time
+                    .and_modify(|e| {
+                        e.0 += send_metric.value;
+                        e.1 = stop_time
+                    })
+                    // cache new metric
+                    .or_insert((send_metric.value, stop_time));
+            }
+        } else {
+            error!("metrics endpoint refused the sent metrics: {:?}", res);
+        }
+    }
+    Ok(())
+}

safekeeper/Cargo.toml

@@ -1,48 +1,48 @@
 [package]
 name = "safekeeper"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [dependencies]
-async-stream = "0.3"
+async-stream.workspace = true
-anyhow = "1.0"
+anyhow.workspace = true
-async-trait = "0.1"
+async-trait.workspace = true
-byteorder = "1.4.3"
+byteorder.workspace = true
-bytes = "1.0.1"
+bytes.workspace = true
-clap = { version = "4.0", features = ["derive"] }
+clap = { workspace = true, features = ["derive"] }
-const_format = "0.2.21"
+const_format.workspace = true
-crc32c = "0.6.0"
+crc32c.workspace = true
-fs2 = "0.4.3"
+fs2.workspace = true
-git-version = "0.3.5"
+git-version.workspace = true
-hex = "0.4.3"
+hex.workspace = true
-humantime = "2.1.0"
+humantime.workspace = true
-hyper = "0.14"
+hyper.workspace = true
-nix = "0.25"
+nix.workspace = true
-once_cell = "1.13.0"
+once_cell.workspace = true
-parking_lot = "0.12.1"
+parking_lot.workspace = true
-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres.workspace = true
-postgres-protocol = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+postgres-protocol.workspace = true
-regex = "1.4.5"
+regex.workspace = true
-serde = { version = "1.0", features = ["derive"] }
+serde.workspace = true
-serde_json = "1"
+serde_json.workspace = true
-serde_with = "2.0"
+serde_with.workspace = true
-signal-hook = "0.3.10"
+signal-hook.workspace = true
-thiserror = "1"
+thiserror.workspace = true
-tokio = { version = "1.17", features = ["macros", "fs"] }
+tokio = { workspace = true, features = ["fs"] }
-tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", rev="43e6db254a97fdecbce33d8bc0890accfd74495e" }
+tokio-postgres.workspace = true
-toml_edit = { version = "0.14", features = ["easy"] }
+toml_edit.workspace = true
-tracing = "0.1.27"
+tracing.workspace = true
-url = "2.2.2"
+url.workspace = true
+metrics.workspace = true
+postgres_ffi.workspace = true
+pq_proto.workspace = true
+remote_storage.workspace = true
+safekeeper_api.workspace = true
+storage_broker.workspace = true
+utils.workspace = true
 
-metrics = { path = "../libs/metrics" }
+workspace_hack.workspace = true
-postgres_ffi = { path = "../libs/postgres_ffi" }
-pq_proto = { path = "../libs/pq_proto" }
-remote_storage = { path = "../libs/remote_storage" }
-safekeeper_api = { path = "../libs/safekeeper_api" }
-storage_broker = { version = "0.1", path = "../storage_broker" }
-utils = { path = "../libs/utils" }
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
 
 [dev-dependencies]
-tempfile = "3.2"
+tempfile.workspace = true

safekeeper/src/timeline.rs

@@ -466,7 +466,7 @@ impl Timeline {
             Ok(_) => Ok(()),
             Err(e) => {
                 // Bootstrap failed, cancel timeline and remove timeline directory.
-                self.cancel();
+                self.cancel(shared_state);
 
                 if let Err(fs_err) = std::fs::remove_dir_all(&self.timeline_dir) {
                     warn!(
@@ -487,20 +487,23 @@ impl Timeline {
         shared_state: &mut MutexGuard<SharedState>,
     ) -> Result<(bool, bool)> {
         let was_active = shared_state.active;
-        self.cancel();
+        self.cancel(shared_state);
         let dir_existed = delete_dir(&self.timeline_dir)?;
         Ok((dir_existed, was_active))
     }
 
     /// Cancel timeline to prevent further usage. Background tasks will stop
     /// eventually after receiving cancellation signal.
-    fn cancel(&self) {
+    fn cancel(&self, shared_state: &mut MutexGuard<SharedState>) {
-        info!("Timeline {} is cancelled", self.ttid);
+        info!("timeline {} is cancelled", self.ttid);
         let _ = self.cancellation_tx.send(true);
         let res = self.wal_backup_launcher_tx.blocking_send(self.ttid);
         if let Err(e) = res {
             error!("Failed to send stop signal to wal_backup_launcher: {}", e);
         }
+        // Close associated FDs. Nobody will be able to touch timeline data once
+        // it is cancelled, so WAL storage won't be opened again.
+        shared_state.sk.wal_store.close();
     }
 
     /// Returns if timeline is cancelled.
@@ -537,10 +540,6 @@ impl Timeline {
     /// De-register compute connection, shutting down timeline activity if
     /// pageserver doesn't need catchup.
     pub fn on_compute_disconnect(&self) -> Result<()> {
-        if self.is_cancelled() {
-            bail!(TimelineError::Cancelled(self.ttid));
-        }
-
         let is_wal_backup_action_pending: bool;
         {
             let mut shared_state = self.write_shared_state();

safekeeper/src/wal_storage.rs

@@ -55,6 +55,12 @@ pub trait Storage {
     /// that without timeline lock.
     fn remove_up_to(&self) -> Box<dyn Fn(XLogSegNo) -> Result<()>>;
 
+    /// Release resources associated with the storage -- technically, close FDs.
+    /// Currently we don't remove timelines until restart (#3146), so need to
+    /// spare descriptors. This would be useful for temporary tli detach as
+    /// well.
+    fn close(&mut self) {}
+
     /// Get metrics for this timeline.
     fn get_metrics(&self) -> WalStorageMetrics;
 }
@@ -401,6 +407,11 @@ impl Storage for PhysicalStorage {
         })
     }
 
+    fn close(&mut self) {
+        // close happens in destructor
+        let _open_file = self.file.take();
+    }
+
     fn get_metrics(&self) -> WalStorageMetrics {
         self.metrics.clone()
     }

storage_broker/Cargo.toml

@@ -1,38 +1,38 @@
 [package]
 name = "storage_broker"
 version = "0.1.0"
-edition = "2021"
+edition.workspace = true
-license = "Apache-2.0"
+license.workspace = true
 
 [features]
 bench = []
 
 [dependencies]
-anyhow = "1.0"
+anyhow.workspace = true
-async-stream = "0.3"
+async-stream.workspace = true
-bytes = "1.0"
+bytes.workspace = true
-clap = { version = "4.0", features = ["derive"] }
+clap = { workspace = true, features = ["derive"] }
-const_format = "0.2.21"
+const_format.workspace = true
-futures = "0.3"
+futures.workspace = true
-futures-core = "0.3"
+futures-core.workspace = true
-futures-util = "0.3"
+futures-util.workspace = true
-git-version = "0.3.5"
+git-version.workspace = true
-humantime = "2.1.0"
+humantime.workspace = true
-hyper = {version = "0.14.14", features = ["full"]}
+hyper = { workspace = true, features = ["full"] }
-once_cell = "1.13.0"
+once_cell.workspace = true
-parking_lot = "0.12"
+parking_lot.workspace = true
-prost = "0.11"
+prost.workspace = true
-tonic = {version = "0.8", features = ["tls", "tls-roots"]}
+tonic.workspace = true
-tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
+tokio = { workspace = true, features = ["rt-multi-thread"] }
-tokio-stream = "0.1"
+tokio-stream.workspace = true
-tracing = "0.1.27"
+tracing.workspace = true
+metrics.workspace = true
+utils.workspace = true
 
-metrics = { path = "../libs/metrics" }
+workspace_hack.workspace = true
-utils = { path = "../libs/utils" }
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
 
 [build-dependencies]
-tonic-build = "0.8"
+tonic-build.workspace = true
 
 [[bench]]
 name = "rps"

test_runner/fixtures/neon_fixtures.py

@@ -22,6 +22,7 @@ from itertools import chain, product
 from pathlib import Path
 from types import TracebackType
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union, cast
+from urllib.parse import urlparse
 
 import asyncpg
 import backoff  # type: ignore
@@ -1350,11 +1351,18 @@ class PageserverHttpClient(requests.Session):
         assert res_json is None
 
     def timeline_spawn_download_remote_layers(
-        self, tenant_id: TenantId, timeline_id: TimelineId
+        self,
+        tenant_id: TenantId,
+        timeline_id: TimelineId,
+        max_concurrent_downloads: int,
     ) -> dict[str, Any]:
 
+        body = {
+            "max_concurrent_downloads": max_concurrent_downloads,
+        }
         res = self.post(
             f"http://localhost:{self.port}/v1/tenant/{tenant_id}/timeline/{timeline_id}/download_remote_layers",
+            json=body,
         )
         self.verbose_error(res)
         res_json = res.json()
@@ -1388,10 +1396,13 @@ class PageserverHttpClient(requests.Session):
         self,
         tenant_id: TenantId,
         timeline_id: TimelineId,
+        max_concurrent_downloads: int,
         errors_ok=False,
         at_least_one_download=True,
     ):
-        res = self.timeline_spawn_download_remote_layers(tenant_id, timeline_id)
+        res = self.timeline_spawn_download_remote_layers(
+            tenant_id, timeline_id, max_concurrent_downloads
+        )
         while True:
             completed = self.timeline_poll_download_remote_layers_status(
                 tenant_id, timeline_id, res, poll_state="Completed"
@@ -2323,6 +2334,8 @@ class NeonProxy(PgProtocol):
         http_port: int,
         mgmt_port: int,
         auth_backend: NeonProxy.AuthBackend,
+        metric_collection_endpoint: Optional[str] = None,
+        metric_collection_interval: Optional[str] = None,
     ):
         host = "127.0.0.1"
         super().__init__(dsn=auth_backend.default_conn_url, host=host, port=proxy_port)
@@ -2333,6 +2346,8 @@ class NeonProxy(PgProtocol):
         self.proxy_port = proxy_port
         self.mgmt_port = mgmt_port
         self.auth_backend = auth_backend
+        self.metric_collection_endpoint = metric_collection_endpoint
+        self.metric_collection_interval = metric_collection_interval
         self._popen: Optional[subprocess.Popen[bytes]] = None
 
     def start(self) -> NeonProxy:
@@ -2344,6 +2359,16 @@ class NeonProxy(PgProtocol):
             *["--mgmt", f"{self.host}:{self.mgmt_port}"],
             *self.auth_backend.extra_args(),
         ]
+
+        if (
+            self.metric_collection_endpoint is not None
+            and self.metric_collection_interval is not None
+        ):
+            args += [
+                *["--metric-collection-endpoint", self.metric_collection_endpoint],
+                *["--metric-collection-interval", self.metric_collection_interval],
+            ]
+
         self._popen = subprocess.Popen(args)
         self._wait_until_ready()
         return self
@@ -2357,6 +2382,25 @@ class NeonProxy(PgProtocol):
         request_result.raise_for_status()
         return request_result.text
 
+    @staticmethod
+    def get_session_id(uri_prefix, uri_line):
+        assert uri_prefix in uri_line
+
+        url_parts = urlparse(uri_line)
+        psql_session_id = url_parts.path[1:]
+        assert psql_session_id.isalnum(), "session_id should only contain alphanumeric chars"
+
+        return psql_session_id
+
+    @staticmethod
+    async def find_auth_link(link_auth_uri, proc):
+        for _ in range(100):
+            line = (await proc.stderr.readline()).decode("utf-8").strip()
+            log.info(f"psql line: {line}")
+            if link_auth_uri in line:
+                log.info(f"SUCCESS, found auth url: {line}")
+                return line
+
     def __enter__(self) -> NeonProxy:
         return self
 
@@ -2371,6 +2415,46 @@ class NeonProxy(PgProtocol):
             # it's a child process. This is mostly to clean up in between different tests.
             self._popen.kill()
 
+    @staticmethod
+    async def activate_link_auth(
+        local_vanilla_pg, proxy_with_metric_collector, psql_session_id, create_user=True
+    ):
+
+        pg_user = "proxy"
+
+        if create_user:
+            log.info("creating a new user for link auth test")
+            local_vanilla_pg.start()
+            local_vanilla_pg.safe_psql(f"create user {pg_user} with login superuser")
+
+        db_info = json.dumps(
+            {
+                "session_id": psql_session_id,
+                "result": {
+                    "Success": {
+                        "host": local_vanilla_pg.default_options["host"],
+                        "port": local_vanilla_pg.default_options["port"],
+                        "dbname": local_vanilla_pg.default_options["dbname"],
+                        "user": pg_user,
+                        "aux": {
+                            "project_id": "test_project_id",
+                            "endpoint_id": "test_endpoint_id",
+                            "branch_id": "test_branch_id",
+                        },
+                    }
+                },
+            }
+        )
+
+        log.info("sending session activation message")
+        psql = await PSQL(
+            host=proxy_with_metric_collector.host,
+            port=proxy_with_metric_collector.mgmt_port,
+        ).run(db_info)
+        assert psql.stdout is not None
+        out = (await psql.stdout.read()).decode("utf-8").strip()
+        assert out == "ok"
+
 
 @pytest.fixture(scope="function")
 def link_proxy(port_distributor: PortDistributor, neon_binpath: Path) -> Iterator[NeonProxy]:

test_runner/performance/test_perf_olap.py

@@ -144,3 +144,42 @@ def test_tpch(query: LabelledQuery, remote_compare: RemoteCompare):
     """
 
     run_psql(remote_compare, query, times=1)
+
+
+@pytest.mark.remote_cluster
+def test_user_examples(remote_compare: RemoteCompare):
+    query = LabelledQuery(
+        "Q1",
+        r"""
+        SELECT
+            v20.c2263 AS v1,
+            v19.c2484 AS v2,
+            DATE_TRUNC('month', v18.c37)::DATE AS v3,
+            (ARRAY_AGG(c1840 order by v18.c37))[1] AS v4,
+            (ARRAY_AGG(c1841 order by v18.c37 DESC))[1] AS v5,
+            SUM(v17.c1843) AS v6,
+            SUM(v17.c1844) AS v7,
+            SUM(v17.c1848) AS v8,
+            SUM(v17.c1845) AS v9,
+            SUM(v17.c1846) AS v10,
+            SUM(v17.c1861) AS v11,
+            SUM(v17.c1860) AS v12,
+            SUM(v17.c1869) AS v13,
+            SUM(v17.c1856) AS v14,
+            SUM(v17.c1855) AS v15,
+            SUM(v17.c1854) AS v16
+        FROM
+            s3.t266 v17
+            INNER JOIN s1.t41 v18 ON v18.c34 = v17.c1836
+            INNER JOIN s3.t571 v19 ON v19.c2482 = v17.c1834
+            INNER JOIN s3.t331 v20 ON v20.c2261 = v17.c1835
+        WHERE
+            (v17.c1835 = 4) AND
+            (v18.c37 >= '2019-03-01') AND
+            (v17.c1833 = 2)
+        GROUP BY v1, v2, v3
+        ORDER BY v1, v2, v3
+        LIMIT 199;
+        """,
+    )
+    run_psql(remote_compare, query, times=3)

test_runner/pg_clients/csharp/npgsql/Dockerfile

@@ -1,4 +1,4 @@
-FROM mcr.microsoft.com/dotnet/sdk:6.0 AS build
+FROM mcr.microsoft.com/dotnet/sdk:7.0 AS build
 WORKDIR /source
 
 COPY *.csproj .
@@ -7,7 +7,7 @@ RUN dotnet restore
 COPY . .
 RUN dotnet publish -c release -o /app --no-restore
 
-FROM mcr.microsoft.com/dotnet/runtime:6.0
+FROM mcr.microsoft.com/dotnet/runtime:7.0
 WORKDIR /app
 COPY --from=build /app .
 

test_runner/pg_clients/csharp/npgsql/csharp-npgsql.csproj

@@ -2,13 +2,13 @@
 
   <PropertyGroup>
     <OutputType>Exe</OutputType>
-    <TargetFramework>net6.0</TargetFramework>
+    <TargetFramework>net7.0</TargetFramework>
     <ImplicitUsings>enable</ImplicitUsings>
     <Nullable>enable</Nullable>
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Npgsql" Version="6.0.5" />
+    <PackageReference Include="Npgsql" Version="7.0.0" />
   </ItemGroup>
 
 </Project>

test_runner/pg_clients/java/jdbc/Dockerfile

@@ -4,7 +4,7 @@ WORKDIR /source
 COPY . .
 
 WORKDIR /app
-RUN curl --output postgresql.jar https://jdbc.postgresql.org/download/postgresql-42.4.0.jar && \
+RUN curl --output postgresql.jar https://jdbc.postgresql.org/download/postgresql-42.5.1.jar && \
     javac -d /app /source/Example.java
 
 CMD ["java", "-cp", "/app/postgresql.jar:.", "Example"]

test_runner/pg_clients/python/asyncpg/Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.10
+FROM python:3.11
 WORKDIR /source
 
 COPY . .

test_runner/pg_clients/python/asyncpg/requirements.txt

@@ -1 +1 @@
-asyncpg==0.25.0
+asyncpg==0.27.0