standardise logging

.dockerignore

@@ -4,7 +4,6 @@
 !Cargo.lock
 !Cargo.toml
 !Makefile
-!postgres.mk
 !rust-toolchain.toml
 !scripts/ninstall.sh
 !docker-compose/run-tests.sh

.github/workflows/build-macos.yml

@@ -94,6 +94,11 @@ jobs:
         run: |
           make "neon-pg-ext-${{ matrix.postgres-version }}" -j$(sysctl -n hw.ncpu)
 
+      - name: Get postgres headers ${{ matrix.postgres-version }}
+        if: steps.cache_pg.outputs.cache-hit != 'true'
+        run: |
+          make postgres-headers-${{ matrix.postgres-version }} -j$(sysctl -n hw.ncpu)
+
       - name: Upload "pg_install/${{ matrix.postgres-version }}" artifact
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
         with:
@@ -135,12 +140,6 @@ jobs:
           name: pg_install--v17
           path: pg_install/v17
 
-      # `actions/download-artifact` doesn't preserve permissions:
-      # https://github.com/actions/download-artifact?tab=readme-ov-file#permission-loss
-      - name: Make pg_install/v*/bin/* executable
-        run: |
-          chmod +x pg_install/v*/bin/*
-
       - name: Cache walproposer-lib
         id: cache_walproposer_lib
         uses: actions/cache@5a3ec84eff668545956fd18022155c47e93e2684 # v4.2.3
@@ -168,7 +167,7 @@ jobs:
       - name: Build walproposer-lib (only for v17)
         if: steps.cache_walproposer_lib.outputs.cache-hit != 'true'
         run:
-          make walproposer-lib -j$(sysctl -n hw.ncpu) PG_INSTALL_CACHED=1
+          make walproposer-lib -j$(sysctl -n hw.ncpu)
 
       - name: Upload "build/walproposer-lib" artifact
         uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2

.github/workflows/build_and_test.yml

@@ -69,7 +69,7 @@ jobs:
           submodules: true
 
       - name: Check for file changes
-        uses: dorny/paths-filter@de90cc6fb38fc0963ad72b210f1f284cd68cea36  # v3.0.2
+        uses: step-security/paths-filter@v3
         id: files-changed
         with:
           token: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/large_oltp_benchmark.yml

@@ -153,7 +153,7 @@ jobs:
         PERF_TEST_RESULT_CONNSTR: "${{ secrets.PERF_TEST_RESULT_CONNSTR }}"
 
     - name: Benchmark database maintenance
-      if: ${{ matrix.test_maintenance }}
+      if: ${{ matrix.test_maintenance == 'true' }}
       uses: ./.github/actions/run-python-test-set
       with:
         build_type: ${{ env.BUILD_TYPE }}

.github/workflows/neon_extra_builds.yml

@@ -53,7 +53,7 @@ jobs:
           submodules: true
 
       - name: Check for Postgres changes
-        uses: dorny/paths-filter@1441771bbfdd59dcd748680ee64ebd8faab1a242  #v3
+        uses: step-security/paths-filter@v3
         id: files_changed
         with:
           token: ${{ github.token }}

.github/workflows/pre-merge-checks.yml

@@ -34,7 +34,7 @@ jobs:
 
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
 
-      - uses: tj-actions/changed-files@ed68ef82c095e0d48ec87eccea555d944a631a4c # v46.0.5
+      - uses: step-security/changed-files@3dbe17c78367e7d60f00d78ae6781a35be47b4a1 # v45.0.1
         id: python-src
         with:
           files: |
@@ -45,7 +45,7 @@ jobs:
             poetry.lock
             pyproject.toml
 
-      - uses: tj-actions/changed-files@ed68ef82c095e0d48ec87eccea555d944a631a4c # v46.0.5
+      - uses: step-security/changed-files@3dbe17c78367e7d60f00d78ae6781a35be47b4a1 # v45.0.1
         id: rust-src
         with:
           files: |

.github/workflows/proxy-benchmark.yml

@@ -60,23 +60,22 @@ jobs:
         } >> "$GITHUB_ENV"
 
     - name: Run proxy-bench
-      run: ${PROXY_BENCH_PATH}/run.sh
+      run: ./${PROXY_BENCH_PATH}/run.sh
 
     - name: Ingest Bench Results # neon repo script
-      if: always()
+      if: success()
       run: |
         mkdir -p $TEST_OUTPUT
         python $NEON_DIR/scripts/proxy_bench_results_ingest.py --out $TEST_OUTPUT
 
     - name: Push Metrics to Proxy perf database
-      if: always()
+      if: success()
       env:
         PERF_TEST_RESULT_CONNSTR: "${{ secrets.PROXY_TEST_RESULT_CONNSTR }}"
         REPORT_FROM: $TEST_OUTPUT
       run: $NEON_DIR/scripts/generate_and_push_perf_report.sh
 
     - name: Docker cleanup
-      if: always()
       run: docker compose down
 
     - name: Notify Failure

Cargo.lock

@@ -1279,7 +1279,6 @@ dependencies = [
  "remote_storage",
  "serde",
  "serde_json",
- "url",
  "utils",
 ]
 
@@ -1317,7 +1316,6 @@ dependencies = [
  "opentelemetry",
  "opentelemetry_sdk",
  "p256 0.13.2",
- "pageserver_page_api",
  "postgres",
  "postgres_initdb",
  "postgres_versioninfo",
@@ -1337,7 +1335,6 @@ dependencies = [
  "tokio-postgres",
  "tokio-stream",
  "tokio-util",
- "tonic 0.13.1",
  "tower 0.5.2",
  "tower-http",
  "tower-otel",
@@ -4411,7 +4408,6 @@ dependencies = [
  "postgres_backend",
  "postgres_ffi_types",
  "postgres_versioninfo",
- "posthog_client_lite",
  "rand 0.8.5",
  "remote_storage",
  "reqwest",
@@ -4422,7 +4418,6 @@ dependencies = [
  "strum",
  "strum_macros",
  "thiserror 1.0.69",
- "tracing",
  "tracing-utils",
  "utils",
 ]
@@ -4479,14 +4474,12 @@ dependencies = [
  "bytes",
  "futures",
  "pageserver_api",
- "postgres_ffi_types",
+ "postgres_ffi",
  "prost 0.13.5",
- "prost-types 0.13.5",
  "strum",
  "strum_macros",
  "thiserror 1.0.69",
  "tokio",
- "tokio-util",
  "tonic 0.13.1",
  "tonic-build",
  "utils",
@@ -5159,7 +5152,7 @@ dependencies = [
  "petgraph",
  "prettyplease",
  "prost 0.13.5",
- "prost-types 0.13.5",
+ "prost-types 0.13.3",
  "regex",
  "syn 2.0.100",
  "tempfile",
@@ -5202,9 +5195,9 @@ dependencies = [
 
 [[package]]
 name = "prost-types"
-version = "0.13.5"
+version = "0.13.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16"
+checksum = "4759aa0d3a6232fb8dbdb97b61de2c20047c68aca932c7ed76da9d788508d670"
 dependencies = [
  "prost 0.13.5",
 ]
@@ -6811,7 +6804,6 @@ dependencies = [
  "chrono",
  "clap",
  "clashmap",
- "compute_api",
  "control_plane",
  "cron",
  "diesel",
@@ -6823,7 +6815,6 @@ dependencies = [
  "hex",
  "http-utils",
  "humantime",
- "humantime-serde",
  "hyper 0.14.30",
  "itertools 0.10.5",
  "json-structural-diff",
@@ -6834,7 +6825,6 @@ dependencies = [
  "pageserver_api",
  "pageserver_client",
  "postgres_connection",
- "posthog_client_lite",
  "rand 0.8.5",
  "regex",
  "reqwest",
@@ -7645,7 +7635,7 @@ dependencies = [
  "prettyplease",
  "proc-macro2",
  "prost-build 0.13.3",
- "prost-types 0.13.5",
+ "prost-types 0.13.3",
  "quote",
  "syn 2.0.100",
 ]
@@ -7657,7 +7647,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f9687bd5bfeafebdded2356950f278bba8226f0b32109537c4253406e09aafe1"
 dependencies = [
  "prost 0.13.5",
- "prost-types 0.13.5",
+ "prost-types 0.13.3",
  "tokio",
  "tokio-stream",
  "tonic 0.13.1",
@@ -8686,6 +8676,7 @@ dependencies = [
  "num-iter",
  "num-rational",
  "num-traits",
+ "once_cell",
  "p256 0.13.2",
  "parquet",
  "prettyplease",

Cargo.toml

@@ -152,7 +152,6 @@ pprof = { version = "0.14", features = ["criterion", "flamegraph", "frame-pointe
 procfs = "0.16"
 prometheus = {version = "0.13", default-features=false, features = ["process"]} # removes protobuf dependency
 prost = "0.13.5"
-prost-types = "0.13.5"
 rand = "0.8"
 redis = { version = "0.29.2", features = ["tokio-rustls-comp", "keep-alive"] }
 regex = "1.10.2"
@@ -200,7 +199,7 @@ tokio-postgres-rustls = "0.12.0"
 tokio-rustls = { version = "0.26.0", default-features = false, features = ["tls12", "ring"]}
 tokio-stream = "0.1"
 tokio-tar = "0.3"
-tokio-util = { version = "0.7.10", features = ["io", "io-util", "rt"] }
+tokio-util = { version = "0.7.10", features = ["io", "rt"] }
 toml = "0.8"
 toml_edit = "0.22"
 tonic = { version = "0.13.1", default-features = false, features = ["channel", "codegen", "gzip", "prost", "router", "server", "tls-ring", "tls-native-roots", "zstd"] }

Dockerfile

@@ -40,7 +40,6 @@ COPY --chown=nonroot vendor/postgres-v16 vendor/postgres-v16
 COPY --chown=nonroot vendor/postgres-v17 vendor/postgres-v17
 COPY --chown=nonroot pgxn pgxn
 COPY --chown=nonroot Makefile Makefile
-COPY --chown=nonroot postgres.mk postgres.mk
 COPY --chown=nonroot scripts/ninstall.sh scripts/ninstall.sh
 
 ENV BUILD_TYPE=release

Makefile

@@ -4,14 +4,11 @@ ROOT_PROJECT_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
 # managers.
 POSTGRES_INSTALL_DIR ?= $(ROOT_PROJECT_DIR)/pg_install/
 
-# Supported PostgreSQL versions
-POSTGRES_VERSIONS = v17 v16 v15 v14
-
 # CARGO_BUILD_FLAGS: Extra flags to pass to `cargo build`. `--locked`
 # and `--features testing` are popular examples.
 #
-# CARGO_PROFILE: Set to override the cargo profile to use. By default,
-# it is derived from BUILD_TYPE.
+# CARGO_PROFILE: You can also set to override the cargo profile to
+# use. By default, it is derived from BUILD_TYPE.
 
 # All intermediate build artifacts are stored here.
 BUILD_DIR := build
@@ -98,24 +95,91 @@ CACHEDIR_TAG_CONTENTS := "Signature: 8a477f597d28d172789f06886806bc55"
 # Top level Makefile to build Neon and PostgreSQL
 #
 .PHONY: all
-all: neon postgres-install neon-pg-ext
+all: neon postgres neon-pg-ext
 
 ### Neon Rust bits
 #
 # The 'postgres_ffi' depends on the Postgres headers.
 .PHONY: neon
-neon: postgres-headers-install walproposer-lib cargo-target-dir
+neon: postgres-headers walproposer-lib cargo-target-dir
 	+@echo "Compiling Neon"
 	$(CARGO_CMD_PREFIX) cargo build $(CARGO_BUILD_FLAGS) $(CARGO_PROFILE)
-
 .PHONY: cargo-target-dir
 cargo-target-dir:
 	# https://github.com/rust-lang/cargo/issues/14281
 	mkdir -p target
 	test -e target/CACHEDIR.TAG || echo "$(CACHEDIR_TAG_CONTENTS)" > target/CACHEDIR.TAG
 
+### PostgreSQL parts
+# Some rules are duplicated for Postgres v14 and 15. We may want to refactor
+# to avoid the duplication in the future, but it's tolerable for now.
+#
+$(BUILD_DIR)/%/config.status:
+	mkdir -p $(BUILD_DIR)
+	test -e $(BUILD_DIR)/CACHEDIR.TAG || echo "$(CACHEDIR_TAG_CONTENTS)" > $(BUILD_DIR)/CACHEDIR.TAG
+
+	+@echo "Configuring Postgres $* build"
+	@test -s $(ROOT_PROJECT_DIR)/vendor/postgres-$*/configure || { \
+		echo "\nPostgres submodule not found in $(ROOT_PROJECT_DIR)/vendor/postgres-$*/, execute "; \
+		echo "'git submodule update --init --recursive --depth 2 --progress .' in project root.\n"; \
+		exit 1; }
+	mkdir -p $(BUILD_DIR)/$*
+
+	VERSION=$*; \
+	EXTRA_VERSION=$$(cd $(ROOT_PROJECT_DIR)/vendor/postgres-$$VERSION && git rev-parse HEAD); \
+	(cd $(BUILD_DIR)/$$VERSION && \
+	env PATH="$(EXTRA_PATH_OVERRIDES):$$PATH" $(ROOT_PROJECT_DIR)/vendor/postgres-$$VERSION/configure \
+		CFLAGS='$(PG_CFLAGS)' LDFLAGS='$(PG_LDFLAGS)' \
+		$(PG_CONFIGURE_OPTS) --with-extra-version=" ($$EXTRA_VERSION)" \
+		--prefix=$(abspath $(POSTGRES_INSTALL_DIR))/$$VERSION > configure.log)
+
+# nicer alias to run 'configure'
+# Note: I've been unable to use templates for this part of our configuration.
+# I'm not sure why it wouldn't work, but this is the only place (apart from
+# the "build-all-versions" entry points) where direct mention of PostgreSQL
+# versions is used.
+.PHONY: postgres-configure-v17
+postgres-configure-v17: $(BUILD_DIR)/v17/config.status
+.PHONY: postgres-configure-v16
+postgres-configure-v16: $(BUILD_DIR)/v16/config.status
+.PHONY: postgres-configure-v15
+postgres-configure-v15: $(BUILD_DIR)/v15/config.status
+.PHONY: postgres-configure-v14
+postgres-configure-v14: $(BUILD_DIR)/v14/config.status
+
+# Install the PostgreSQL header files into $(POSTGRES_INSTALL_DIR)/<version>/include
+.PHONY: postgres-headers-%
+postgres-headers-%: postgres-configure-%
+	+@echo "Installing PostgreSQL $* headers"
+	$(MAKE) -C $(BUILD_DIR)/$*/src/include MAKELEVEL=0 install
+
+# Compile and install PostgreSQL
+.PHONY: postgres-%
+postgres-%: postgres-configure-% \
+		  postgres-headers-% # to prevent `make install` conflicts with neon's `postgres-headers`
+	+@echo "Compiling PostgreSQL $*"
+	$(MAKE) -C $(BUILD_DIR)/$* MAKELEVEL=0 install
+	+@echo "Compiling pg_prewarm $*"
+	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pg_prewarm install
+	+@echo "Compiling pg_buffercache $*"
+	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pg_buffercache install
+	+@echo "Compiling pg_visibility $*"
+	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pg_visibility install
+	+@echo "Compiling pageinspect $*"
+	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pageinspect install
+	+@echo "Compiling pg_trgm $*"
+	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pg_trgm install
+	+@echo "Compiling amcheck $*"
+	$(MAKE) -C $(BUILD_DIR)/$*/contrib/amcheck install
+	+@echo "Compiling test_decoding $*"
+	$(MAKE) -C $(BUILD_DIR)/$*/contrib/test_decoding install
+
+.PHONY: postgres-check-%
+postgres-check-%: postgres-%
+	$(MAKE) -C $(BUILD_DIR)/$* MAKELEVEL=0 check
+
 .PHONY: neon-pg-ext-%
-neon-pg-ext-%: postgres-install-%
+neon-pg-ext-%: postgres-%
 	+@echo "Compiling neon-specific Postgres extensions for $*"
 	mkdir -p $(BUILD_DIR)/pgxn-$*
 	$(MAKE) PG_CONFIG=$(POSTGRES_INSTALL_DIR)/$*/bin/pg_config COPT='$(COPT)' \
@@ -154,14 +218,39 @@ ifeq ($(UNAME_S),Linux)
 		pg_crc32c.o
 endif
 
-# Shorthand to call neon-pg-ext-% target for all Postgres versions
 .PHONY: neon-pg-ext
-neon-pg-ext: $(foreach pg_version,$(POSTGRES_VERSIONS),neon-pg-ext-$(pg_version))
+neon-pg-ext: \
+	neon-pg-ext-v14 \
+	neon-pg-ext-v15 \
+	neon-pg-ext-v16 \
+	neon-pg-ext-v17
+
+# shorthand to build all Postgres versions
+.PHONY: postgres
+postgres: \
+	postgres-v14 \
+	postgres-v15 \
+	postgres-v16 \
+	postgres-v17
+
+.PHONY: postgres-headers
+postgres-headers: \
+	postgres-headers-v14 \
+	postgres-headers-v15 \
+	postgres-headers-v16 \
+	postgres-headers-v17
+
+.PHONY: postgres-check
+postgres-check: \
+	postgres-check-v14 \
+	postgres-check-v15 \
+	postgres-check-v16 \
+	postgres-check-v17
 
 # This removes everything
 .PHONY: distclean
 distclean:
-	$(RM) -r $(POSTGRES_INSTALL_DIR) $(BUILD_DIR)
+	$(RM) -r $(POSTGRES_INSTALL_DIR)
 	$(CARGO_CMD_PREFIX) cargo clean
 
 .PHONY: fmt
@@ -209,19 +298,3 @@ neon-pgindent: postgres-v17-pg-bsd-indent neon-pg-ext-v17
 .PHONY: setup-pre-commit-hook
 setup-pre-commit-hook:
 	ln -s -f $(ROOT_PROJECT_DIR)/pre-commit.py .git/hooks/pre-commit
-
-# Targets for building PostgreSQL are defined in postgres.mk.
-#
-# But if the caller has indicated that PostgreSQL is already
-# installed, by setting the PG_INSTALL_CACHED variable, skip it.
-ifdef PG_INSTALL_CACHED
-postgres-install: skip-install
-$(foreach pg_version,$(POSTGRES_VERSIONS),postgres-install-$(pg_version)): skip-install
-postgres-headers-install:
-	+@echo "Skipping installation of PostgreSQL headers because PG_INSTALL_CACHED is set"
-skip-install:
-	+@echo "Skipping PostgreSQL installation because PG_INSTALL_CACHED is set"
-
-else
-include postgres.mk
-endif

build-tools.Dockerfile

@@ -165,7 +165,6 @@ RUN curl -fsSL \
     && rm sql_exporter.tar.gz
 
 # protobuf-compiler (protoc)
-# Keep the version the same as in compute/compute-node.Dockerfile
 ENV PROTOC_VERSION=25.1
 RUN curl -fsSL "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-$(uname -m | sed 's/aarch64/aarch_64/g').zip" -o "protoc.zip" \
     && unzip -q protoc.zip -d protoc \
@@ -180,7 +179,7 @@ RUN curl -sL "https://github.com/peak/s5cmd/releases/download/v${S5CMD_VERSION}/
     && mv s5cmd /usr/local/bin/s5cmd
 
 # LLVM
-ENV LLVM_VERSION=20
+ENV LLVM_VERSION=19
 RUN curl -fsSL 'https://apt.llvm.org/llvm-snapshot.gpg.key' | apt-key add - \
     && echo "deb http://apt.llvm.org/${DEBIAN_VERSION}/ llvm-toolchain-${DEBIAN_VERSION}-${LLVM_VERSION} main" > /etc/apt/sources.list.d/llvm.stable.list \
     && apt update \
@@ -293,7 +292,7 @@ WORKDIR /home/nonroot
 
 # Rust
 # Please keep the version of llvm (installed above) in sync with rust llvm (`rustc --version --verbose | grep LLVM`)
-ENV RUSTC_VERSION=1.88.0
+ENV RUSTC_VERSION=1.87.0
 ENV RUSTUP_HOME="/home/nonroot/.rustup"
 ENV PATH="/home/nonroot/.cargo/bin:${PATH}"
 ARG RUSTFILT_VERSION=0.2.1

compute/compute-node.Dockerfile

@@ -115,9 +115,6 @@ ARG EXTENSIONS=all
 FROM $BASE_IMAGE_SHA AS build-deps
 ARG DEBIAN_VERSION
 
-# Keep in sync with build-tools.Dockerfile
-ENV PROTOC_VERSION=25.1
-
 # Use strict mode for bash to catch errors early
 SHELL ["/bin/bash", "-euo", "pipefail", "-c"]
 
@@ -152,14 +149,8 @@ RUN case $DEBIAN_VERSION in \
     libclang-dev \
     jsonnet \
     $VERSION_INSTALLS \
-    && apt clean && rm -rf /var/lib/apt/lists/* \
-    && useradd -ms /bin/bash nonroot -b /home \
-    # Install protoc from binary release, since Debian's versions are too old.
-    && curl -fsSL "https://github.com/protocolbuffers/protobuf/releases/download/v${PROTOC_VERSION}/protoc-${PROTOC_VERSION}-linux-$(uname -m | sed 's/aarch64/aarch_64/g').zip" -o "protoc.zip" \
-    && unzip -q protoc.zip -d protoc \
-    && mv protoc/bin/protoc /usr/local/bin/protoc \
-    && mv protoc/include/google /usr/local/include/google \
-    && rm -rf protoc.zip protoc
+    && apt clean && rm -rf /var/lib/apt/lists/* && \
+    useradd -ms /bin/bash nonroot -b /home
 
 #########################################################################################
 #
@@ -1179,7 +1170,7 @@ COPY --from=pgrag-src /ext-src/ /ext-src/
 # Install it using virtual environment, because Python 3.11 (the default version on Debian 12 (Bookworm)) complains otherwise
 WORKDIR /ext-src/onnxruntime-src
 RUN apt update && apt install --no-install-recommends --no-install-suggests -y \
-    python3 python3-pip python3-venv && \
+    python3 python3-pip python3-venv protobuf-compiler && \
     apt clean && rm -rf /var/lib/apt/lists/* && \
     python3 -m venv venv && \
     . venv/bin/activate && \

compute_tools/Cargo.toml

@@ -38,7 +38,6 @@ once_cell.workspace = true
 opentelemetry.workspace = true
 opentelemetry_sdk.workspace = true
 p256 = { version = "0.13", features = ["pem"] }
-pageserver_page_api.workspace = true
 postgres.workspace = true
 regex.workspace = true
 reqwest = { workspace = true, features = ["json"] }
@@ -54,7 +53,6 @@ tokio = { workspace = true, features = ["rt", "rt-multi-thread"] }
 tokio-postgres.workspace = true
 tokio-util.workspace = true
 tokio-stream.workspace = true
-tonic.workspace = true
 tower-otel.workspace = true
 tracing.workspace = true
 tracing-opentelemetry.workspace = true

compute_tools/src/bin/compute_ctl.rs

@@ -36,8 +36,6 @@
 use std::ffi::OsString;
 use std::fs::File;
 use std::process::exit;
-use std::sync::Arc;
-use std::sync::atomic::AtomicU64;
 use std::sync::mpsc;
 use std::thread;
 use std::time::Duration;
@@ -192,9 +190,7 @@ fn main() -> Result<()> {
             cgroup: cli.cgroup,
             #[cfg(target_os = "linux")]
             vm_monitor_addr: cli.vm_monitor_addr,
-            installed_extensions_collection_interval: Arc::new(AtomicU64::new(
-                cli.installed_extensions_collection_interval,
-            )),
+            installed_extensions_collection_interval: cli.installed_extensions_collection_interval,
         },
         config,
     )?;

compute_tools/src/compute.rs

@@ -6,7 +6,7 @@ use compute_api::responses::{
     LfcPrewarmState, TlsConfig,
 };
 use compute_api::spec::{
-    ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, PageserverProtocol, PgIdent,
+    ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, PgIdent,
 };
 use futures::StreamExt;
 use futures::future::join_all;
@@ -15,17 +15,17 @@ use itertools::Itertools;
 use nix::sys::signal::{Signal, kill};
 use nix::unistd::Pid;
 use once_cell::sync::Lazy;
-use pageserver_page_api::{self as page_api, BaseBackupCompression};
 use postgres;
 use postgres::NoTls;
 use postgres::error::SqlState;
 use remote_storage::{DownloadError, RemotePath};
 use std::collections::{HashMap, HashSet};
+use std::net::SocketAddr;
 use std::os::unix::fs::{PermissionsExt, symlink};
 use std::path::Path;
 use std::process::{Command, Stdio};
 use std::str::FromStr;
-use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::{Arc, Condvar, Mutex, RwLock};
 use std::time::{Duration, Instant};
 use std::{env, fs};
@@ -36,7 +36,6 @@ use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 use utils::measured_stream::MeasuredReader;
 use utils::pid_file;
-use utils::shard::{ShardCount, ShardIndex, ShardNumber};
 
 use crate::configurator::launch_configurator;
 use crate::disk_quota::set_disk_quota;
@@ -70,7 +69,6 @@ pub static BUILD_TAG: Lazy<String> = Lazy::new(|| {
         .unwrap_or(BUILD_TAG_DEFAULT)
         .to_string()
 });
-const DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL: u64 = 3600;
 
 /// Static configuration params that don't change after startup. These mostly
 /// come from the CLI args, or are derived from them.
@@ -104,7 +102,7 @@ pub struct ComputeNodeParams {
     pub remote_ext_base_url: Option<Url>,
 
     /// Interval for installed extensions collection
-    pub installed_extensions_collection_interval: Arc<AtomicU64>,
+    pub installed_extensions_collection_interval: u64,
 }
 
 /// Compute node info shared across several `compute_ctl` threads.
@@ -127,9 +125,6 @@ pub struct ComputeNode {
     // key: ext_archive_name, value: started download time, download_completed?
     pub ext_download_progress: RwLock<HashMap<String, (DateTime<Utc>, bool)>>,
     pub compute_ctl_config: ComputeCtlConfig,
-
-    /// Handle to the extension stats collection task
-    extension_stats_task: Mutex<Option<tokio::task::JoinHandle<()>>>,
 }
 
 // store some metrics about download size that might impact startup time
@@ -223,8 +218,7 @@ pub struct ParsedSpec {
     pub pageserver_connstr: String,
     pub safekeeper_connstrings: Vec<String>,
     pub storage_auth_token: Option<String>,
-    /// k8s dns name and port
-    pub endpoint_storage_addr: Option<String>,
+    pub endpoint_storage_addr: Option<SocketAddr>,
     pub endpoint_storage_token: Option<String>,
 }
 
@@ -319,10 +313,13 @@ impl TryFrom<ComputeSpec> for ParsedSpec {
                 .or(Err("invalid timeline id"))?
         };
 
-        let endpoint_storage_addr: Option<String> = spec
+        let endpoint_storage_addr: Option<SocketAddr> = spec
             .endpoint_storage_addr
             .clone()
-            .or_else(|| spec.cluster.settings.find("neon.endpoint_storage_addr"));
+            .or_else(|| spec.cluster.settings.find("neon.endpoint_storage_addr"))
+            .unwrap_or_default()
+            .parse()
+            .ok();
         let endpoint_storage_token = spec
             .endpoint_storage_token
             .clone()
@@ -432,7 +429,6 @@ impl ComputeNode {
             state_changed: Condvar::new(),
             ext_download_progress: RwLock::new(HashMap::new()),
             compute_ctl_config: config.compute_ctl_config,
-            extension_stats_task: Mutex::new(None),
         })
     }
 
@@ -520,9 +516,6 @@ impl ComputeNode {
             None
         };
 
-        // Terminate the extension stats collection task
-        this.terminate_extension_stats_task();
-
         // Terminate the vm_monitor so it releases the file watcher on
         // /sys/fs/cgroup/neon-postgres.
         // Note: the vm-monitor only runs on linux because it requires cgroups.
@@ -1005,80 +998,13 @@ impl ComputeNode {
         Ok(())
     }
 
-    /// Fetches a basebackup from the Pageserver using the compute state's Pageserver connstring and
-    /// unarchives it to `pgdata` directory, replacing any existing contents.
+    // Get basebackup from the libpq connection to pageserver using `connstr` and
+    // unarchive it to `pgdata` directory overriding all its previous content.
     #[instrument(skip_all, fields(%lsn))]
     fn try_get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
         let spec = compute_state.pspec.as_ref().expect("spec must be set");
+        let start_time = Instant::now();
 
-        let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
-        let started = Instant::now();
-
-        let (connected, size) = match PageserverProtocol::from_connstring(shard0_connstr)? {
-            PageserverProtocol::Libpq => self.try_get_basebackup_libpq(spec, lsn)?,
-            PageserverProtocol::Grpc => self.try_get_basebackup_grpc(spec, lsn)?,
-        };
-
-        let mut state = self.state.lock().unwrap();
-        state.metrics.pageserver_connect_micros =
-            connected.duration_since(started).as_micros() as u64;
-        state.metrics.basebackup_bytes = size as u64;
-        state.metrics.basebackup_ms = started.elapsed().as_millis() as u64;
-
-        Ok(())
-    }
-
-    /// Fetches a basebackup via gRPC. The connstring must use grpc://. Returns the timestamp when
-    /// the connection was established, and the (compressed) size of the basebackup.
-    fn try_get_basebackup_grpc(&self, spec: &ParsedSpec, lsn: Lsn) -> Result<(Instant, usize)> {
-        let shard0_connstr = spec
-            .pageserver_connstr
-            .split(',')
-            .next()
-            .unwrap()
-            .to_string();
-        let shard_index = match spec.pageserver_connstr.split(',').count() as u8 {
-            0 | 1 => ShardIndex::unsharded(),
-            count => ShardIndex::new(ShardNumber(0), ShardCount(count)),
-        };
-
-        let (reader, connected) = tokio::runtime::Handle::current().block_on(async move {
-            let mut client = page_api::Client::new(
-                shard0_connstr,
-                spec.tenant_id,
-                spec.timeline_id,
-                shard_index,
-                spec.storage_auth_token.clone(),
-                None, // NB: base backups use payload compression
-            )
-            .await?;
-            let connected = Instant::now();
-            let reader = client
-                .get_base_backup(page_api::GetBaseBackupRequest {
-                    lsn: (lsn != Lsn(0)).then_some(lsn),
-                    compression: BaseBackupCompression::Gzip,
-                    replica: spec.spec.mode != ComputeMode::Primary,
-                    full: false,
-                })
-                .await?;
-            anyhow::Ok((reader, connected))
-        })?;
-
-        let mut reader = MeasuredReader::new(tokio_util::io::SyncIoBridge::new(reader));
-
-        // Set `ignore_zeros` so that unpack() reads the entire stream and doesn't just stop at the
-        // end-of-archive marker. If the server errors, the tar::Builder drop handler will write an
-        // end-of-archive marker before the error is emitted, and we would not see the error.
-        let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&mut reader));
-        ar.set_ignore_zeros(true);
-        ar.unpack(&self.params.pgdata)?;
-
-        Ok((connected, reader.get_byte_count()))
-    }
-
-    /// Fetches a basebackup via libpq. The connstring must use postgresql://. Returns the timestamp
-    /// when the connection was established, and the (compressed) size of the basebackup.
-    fn try_get_basebackup_libpq(&self, spec: &ParsedSpec, lsn: Lsn) -> Result<(Instant, usize)> {
         let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
         let mut config = postgres::Config::from_str(shard0_connstr)?;
 
@@ -1092,14 +1018,16 @@ impl ComputeNode {
         }
 
         config.application_name("compute_ctl");
-        config.options(&format!(
-            "-c neon.compute_mode={}",
-            spec.spec.mode.to_type_str()
-        ));
+        if let Some(spec) = &compute_state.pspec {
+            config.options(&format!(
+                "-c neon.compute_mode={}",
+                spec.spec.mode.to_type_str()
+            ));
+        }
 
         // Connect to pageserver
         let mut client = config.connect(NoTls)?;
-        let connected = Instant::now();
+        let pageserver_connect_micros = start_time.elapsed().as_micros() as u64;
 
         let basebackup_cmd = match lsn {
             Lsn(0) => {
@@ -1136,13 +1064,16 @@ impl ComputeNode {
         // Set `ignore_zeros` so that unpack() reads all the Copy data and
         // doesn't stop at the end-of-archive marker. Otherwise, if the server
         // sends an Error after finishing the tarball, we will not notice it.
-        // The tar::Builder drop handler will write an end-of-archive marker
-        // before emitting the error, and we would not see it otherwise.
         let mut ar = tar::Archive::new(flate2::read::GzDecoder::new(&mut bufreader));
         ar.set_ignore_zeros(true);
         ar.unpack(&self.params.pgdata)?;
 
-        Ok((connected, measured_reader.get_byte_count()))
+        // Report metrics
+        let mut state = self.state.lock().unwrap();
+        state.metrics.pageserver_connect_micros = pageserver_connect_micros;
+        state.metrics.basebackup_bytes = measured_reader.get_byte_count() as u64;
+        state.metrics.basebackup_ms = start_time.elapsed().as_millis() as u64;
+        Ok(())
     }
 
     // Gets the basebackup in a retry loop
@@ -1679,8 +1610,6 @@ impl ComputeNode {
             tls_config = self.compute_ctl_config.tls.clone();
         }
 
-        self.update_installed_extensions_collection_interval(&spec);
-
         let max_concurrent_connections = self.max_service_connections(compute_state, &spec);
 
         // Merge-apply spec & changes to PostgreSQL state.
@@ -1745,8 +1674,6 @@ impl ComputeNode {
 
         let tls_config = self.tls_config(&spec);
 
-        self.update_installed_extensions_collection_interval(&spec);
-
         if let Some(ref pgbouncer_settings) = spec.pgbouncer_settings {
             info!("tuning pgbouncer");
 
@@ -2351,20 +2278,10 @@ LIMIT 100",
     }
 
     pub fn spawn_extension_stats_task(&self) {
-        // Cancel any existing task
-        if let Some(handle) = self.extension_stats_task.lock().unwrap().take() {
-            handle.abort();
-        }
-
         let conf = self.tokio_conn_conf.clone();
-        let atomic_interval = self.params.installed_extensions_collection_interval.clone();
-        let mut installed_extensions_collection_interval =
-            2 * atomic_interval.load(std::sync::atomic::Ordering::SeqCst);
-        info!(
-            "[NEON_EXT_SPAWN] Spawning background installed extensions worker with Timeout: {}",
-            installed_extensions_collection_interval
-        );
-        let handle = tokio::spawn(async move {
+        let installed_extensions_collection_interval =
+            self.params.installed_extensions_collection_interval;
+        tokio::spawn(async move {
             // An initial sleep is added to ensure that two collections don't happen at the same time.
             // The first collection happens during compute startup.
             tokio::time::sleep(tokio::time::Duration::from_secs(
@@ -2377,48 +2294,8 @@ LIMIT 100",
             loop {
                 interval.tick().await;
                 let _ = installed_extensions(conf.clone()).await;
-                // Acquire a read lock on the compute spec and then update the interval if necessary
-                interval = tokio::time::interval(tokio::time::Duration::from_secs(std::cmp::max(
-                    installed_extensions_collection_interval,
-                    2 * atomic_interval.load(std::sync::atomic::Ordering::SeqCst),
-                )));
-                installed_extensions_collection_interval = interval.period().as_secs();
             }
         });
-
-        // Store the new task handle
-        *self.extension_stats_task.lock().unwrap() = Some(handle);
-    }
-
-    fn terminate_extension_stats_task(&self) {
-        if let Some(handle) = self.extension_stats_task.lock().unwrap().take() {
-            handle.abort();
-        }
-    }
-
-    fn update_installed_extensions_collection_interval(&self, spec: &ComputeSpec) {
-        // Update the interval for collecting installed extensions statistics
-        // If the value is -1, we never suspend so set the value to default collection.
-        // If the value is 0, it means default, we will just continue to use the default.
-        if spec.suspend_timeout_seconds == -1 || spec.suspend_timeout_seconds == 0 {
-            info!(
-                "[NEON_EXT_INT_UPD] Spec Timeout: {}, New Timeout: {}",
-                spec.suspend_timeout_seconds, DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL
-            );
-            self.params.installed_extensions_collection_interval.store(
-                DEFAULT_INSTALLED_EXTENSIONS_COLLECTION_INTERVAL,
-                std::sync::atomic::Ordering::SeqCst,
-            );
-        } else {
-            info!(
-                "[NEON_EXT_INT_UPD] Spec Timeout: {}",
-                spec.suspend_timeout_seconds
-            );
-            self.params.installed_extensions_collection_interval.store(
-                spec.suspend_timeout_seconds as u64,
-                std::sync::atomic::Ordering::SeqCst,
-            );
-        }
     }
 }
 

compute_tools/src/lsn_lease.rs

@@ -4,9 +4,7 @@ use std::thread;
 use std::time::{Duration, SystemTime};
 
 use anyhow::{Result, bail};
-use compute_api::spec::{ComputeMode, PageserverProtocol};
-use itertools::Itertools as _;
-use pageserver_page_api as page_api;
+use compute_api::spec::ComputeMode;
 use postgres::{NoTls, SimpleQueryMessage};
 use tracing::{info, warn};
 use utils::id::{TenantId, TimelineId};
@@ -78,17 +76,25 @@ fn acquire_lsn_lease_with_retry(
 
     loop {
         // Note: List of pageservers is dynamic, need to re-read configs before each attempt.
-        let (connstrings, auth) = {
+        let configs = {
             let state = compute.state.lock().unwrap();
+
             let spec = state.pspec.as_ref().expect("spec must be set");
-            (
-                spec.pageserver_connstr.clone(),
-                spec.storage_auth_token.clone(),
-            )
+
+            let conn_strings = spec.pageserver_connstr.split(',');
+
+            conn_strings
+                .map(|connstr| {
+                    let mut config = postgres::Config::from_str(connstr).expect("Invalid connstr");
+                    if let Some(storage_auth_token) = &spec.storage_auth_token {
+                        config.password(storage_auth_token.clone());
+                    }
+                    config
+                })
+                .collect::<Vec<_>>()
         };
 
-        let result =
-            try_acquire_lsn_lease(&connstrings, auth.as_deref(), tenant_id, timeline_id, lsn);
+        let result = try_acquire_lsn_lease(tenant_id, timeline_id, lsn, &configs);
         match result {
             Ok(Some(res)) => {
                 return Ok(res);
@@ -110,104 +116,68 @@ fn acquire_lsn_lease_with_retry(
     }
 }
 
-/// Tries to acquire LSN leases on all Pageserver shards.
+/// Tries to acquire an LSN lease through PS page_service API.
 fn try_acquire_lsn_lease(
-    connstrings: &str,
-    auth: Option<&str>,
     tenant_id: TenantId,
     timeline_id: TimelineId,
     lsn: Lsn,
+    configs: &[postgres::Config],
 ) -> Result<Option<SystemTime>> {
-    let connstrings = connstrings.split(',').collect_vec();
-    let shard_count = connstrings.len();
-    let mut leases = Vec::new();
-
-    for (shard_number, &connstring) in connstrings.iter().enumerate() {
-        let tenant_shard_id = match shard_count {
-            0 | 1 => TenantShardId::unsharded(tenant_id),
-            shard_count => TenantShardId {
-                tenant_id,
-                shard_number: ShardNumber(shard_number as u8),
-                shard_count: ShardCount::new(shard_count as u8),
-            },
+    fn get_valid_until(
+        config: &postgres::Config,
+        tenant_shard_id: TenantShardId,
+        timeline_id: TimelineId,
+        lsn: Lsn,
+    ) -> Result<Option<SystemTime>> {
+        let mut client = config.connect(NoTls)?;
+        let cmd = format!("lease lsn {tenant_shard_id} {timeline_id} {lsn} ");
+        let res = client.simple_query(&cmd)?;
+        let msg = match res.first() {
+            Some(msg) => msg,
+            None => bail!("empty response"),
+        };
+        let row = match msg {
+            SimpleQueryMessage::Row(row) => row,
+            _ => bail!("error parsing lsn lease response"),
         };
 
-        let lease = match PageserverProtocol::from_connstring(connstring)? {
-            PageserverProtocol::Libpq => {
-                acquire_lsn_lease_libpq(connstring, auth, tenant_shard_id, timeline_id, lsn)?
-            }
-            PageserverProtocol::Grpc => {
-                acquire_lsn_lease_grpc(connstring, auth, tenant_shard_id, timeline_id, lsn)?
-            }
-        };
-        leases.push(lease);
+        // Note: this will be None if a lease is explicitly not granted.
+        let valid_until_str = row.get("valid_until");
+
+        let valid_until = valid_until_str.map(|s| {
+            SystemTime::UNIX_EPOCH
+                .checked_add(Duration::from_millis(u128::from_str(s).unwrap() as u64))
+                .expect("Time larger than max SystemTime could handle")
+        });
+        Ok(valid_until)
     }
 
-    Ok(leases.into_iter().min().flatten())
-}
+    let shard_count = configs.len();
 
-/// Acquires an LSN lease on a single shard, using the libpq API. The connstring must use a
-/// postgresql:// scheme.
-fn acquire_lsn_lease_libpq(
-    connstring: &str,
-    auth: Option<&str>,
-    tenant_shard_id: TenantShardId,
-    timeline_id: TimelineId,
-    lsn: Lsn,
-) -> Result<Option<SystemTime>> {
-    let mut config = postgres::Config::from_str(connstring)?;
-    if let Some(auth) = auth {
-        config.password(auth);
-    }
-    let mut client = config.connect(NoTls)?;
-    let cmd = format!("lease lsn {tenant_shard_id} {timeline_id} {lsn} ");
-    let res = client.simple_query(&cmd)?;
-    let msg = match res.first() {
-        Some(msg) => msg,
-        None => bail!("empty response"),
-    };
-    let row = match msg {
-        SimpleQueryMessage::Row(row) => row,
-        _ => bail!("error parsing lsn lease response"),
+    let valid_until = if shard_count > 1 {
+        configs
+            .iter()
+            .enumerate()
+            .map(|(shard_number, config)| {
+                let tenant_shard_id = TenantShardId {
+                    tenant_id,
+                    shard_count: ShardCount::new(shard_count as u8),
+                    shard_number: ShardNumber(shard_number as u8),
+                };
+                get_valid_until(config, tenant_shard_id, timeline_id, lsn)
+            })
+            .collect::<Result<Vec<Option<SystemTime>>>>()?
+            .into_iter()
+            .min()
+            .unwrap()
+    } else {
+        get_valid_until(
+            &configs[0],
+            TenantShardId::unsharded(tenant_id),
+            timeline_id,
+            lsn,
+        )?
     };
 
-    // Note: this will be None if a lease is explicitly not granted.
-    let valid_until_str = row.get("valid_until");
-
-    let valid_until = valid_until_str.map(|s| {
-        SystemTime::UNIX_EPOCH
-            .checked_add(Duration::from_millis(u128::from_str(s).unwrap() as u64))
-            .expect("Time larger than max SystemTime could handle")
-    });
     Ok(valid_until)
 }
-
-/// Acquires an LSN lease on a single shard, using the gRPC API. The connstring must use a
-/// grpc:// scheme.
-fn acquire_lsn_lease_grpc(
-    connstring: &str,
-    auth: Option<&str>,
-    tenant_shard_id: TenantShardId,
-    timeline_id: TimelineId,
-    lsn: Lsn,
-) -> Result<Option<SystemTime>> {
-    tokio::runtime::Handle::current().block_on(async move {
-        let mut client = page_api::Client::new(
-            connstring.to_string(),
-            tenant_shard_id.tenant_id,
-            timeline_id,
-            tenant_shard_id.to_index(),
-            auth.map(String::from),
-            None,
-        )
-        .await?;
-
-        let req = page_api::LeaseLsnRequest { lsn };
-        match client.lease_lsn(req).await {
-            Ok(expires) => Ok(Some(expires)),
-            // Lease couldn't be acquired because the LSN has been garbage collected.
-            Err(err) if err.code() == tonic::Code::FailedPrecondition => Ok(None),
-            Err(err) => Err(err.into()),
-        }
-    })
-}

compute_tools/tests/cluster_spec.json

@@ -3,8 +3,7 @@
 
   "timestamp": "2021-05-23T18:25:43.511Z",
   "operation_uuid": "0f657b36-4b0f-4a2d-9c2e-1dcd615e7d8b",
-  "suspend_timeout_seconds": 3600,
-  
+
   "cluster": {
     "cluster_id": "test-cluster-42",
     "name": "Zenith Test",

control_plane/src/bin/neon_local.rs

@@ -16,9 +16,9 @@ use std::time::Duration;
 use anyhow::{Context, Result, anyhow, bail};
 use clap::Parser;
 use compute_api::requests::ComputeClaimsScope;
-use compute_api::spec::{ComputeMode, PageserverProtocol};
+use compute_api::spec::ComputeMode;
 use control_plane::broker::StorageBroker;
-use control_plane::endpoint::{ComputeControlPlane, EndpointTerminateMode};
+use control_plane::endpoint::{ComputeControlPlane, EndpointTerminateMode, PageserverProtocol};
 use control_plane::endpoint_storage::{ENDPOINT_STORAGE_DEFAULT_ADDR, EndpointStorage};
 use control_plane::local_env;
 use control_plane::local_env::{
@@ -1649,9 +1649,7 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
             // If --safekeepers argument is given, use only the listed
             // safekeeper nodes; otherwise all from the env.
             let safekeepers = parse_safekeepers(&args.safekeepers)?;
-            endpoint
-                .reconfigure(Some(pageservers), None, safekeepers, None)
-                .await?;
+            endpoint.reconfigure(pageservers, None, safekeepers).await?;
         }
         EndpointCmd::Stop(args) => {
             let endpoint_id = &args.endpoint_id;

control_plane/src/endpoint.rs

@@ -56,8 +56,8 @@ use compute_api::responses::{
     TlsConfig,
 };
 use compute_api::spec::{
-    Cluster, ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, Database, PageserverProtocol,
-    PgIdent, RemoteExtSpec, Role,
+    Cluster, ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, Database, PgIdent,
+    RemoteExtSpec, Role,
 };
 use jsonwebtoken::jwk::{
     AlgorithmParameters, CommonParameters, EllipticCurve, Jwk, JwkSet, KeyAlgorithm, KeyOperations,
@@ -373,6 +373,29 @@ impl std::fmt::Display for EndpointTerminateMode {
     }
 }
 
+/// Protocol used to connect to a Pageserver.
+#[derive(Clone, Copy, Debug)]
+pub enum PageserverProtocol {
+    Libpq,
+    Grpc,
+}
+
+impl PageserverProtocol {
+    /// Returns the URL scheme for the protocol, used in connstrings.
+    pub fn scheme(&self) -> &'static str {
+        match self {
+            Self::Libpq => "postgresql",
+            Self::Grpc => "grpc",
+        }
+    }
+}
+
+impl Display for PageserverProtocol {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.scheme())
+    }
+}
+
 impl Endpoint {
     fn from_dir_entry(entry: std::fs::DirEntry, env: &LocalEnv) -> Result<Endpoint> {
         if !entry.file_type()?.is_dir() {
@@ -780,7 +803,6 @@ impl Endpoint {
                 endpoint_storage_addr: Some(endpoint_storage_addr),
                 endpoint_storage_token: Some(endpoint_storage_token),
                 autoprewarm: false,
-                suspend_timeout_seconds: -1, // Only used in neon_local.
             };
 
             // this strange code is needed to support respec() in tests
@@ -975,11 +997,12 @@ impl Endpoint {
 
     pub async fn reconfigure(
         &self,
-        pageservers: Option<Vec<(PageserverProtocol, Host, u16)>>,
+        pageservers: Vec<(PageserverProtocol, Host, u16)>,
         stripe_size: Option<ShardStripeSize>,
         safekeepers: Option<Vec<NodeId>>,
-        safekeeper_generation: Option<SafekeeperGeneration>,
     ) -> Result<()> {
+        anyhow::ensure!(!pageservers.is_empty(), "no pageservers provided");
+
         let (mut spec, compute_ctl_config) = {
             let config_path = self.endpoint_path().join("config.json");
             let file = std::fs::File::open(config_path)?;
@@ -991,24 +1014,16 @@ impl Endpoint {
         let postgresql_conf = self.read_postgresql_conf()?;
         spec.cluster.postgresql_conf = Some(postgresql_conf);
 
-        // If pageservers are not specified, don't change them.
-        if let Some(pageservers) = pageservers {
-            anyhow::ensure!(!pageservers.is_empty(), "no pageservers provided");
-
-            let pageserver_connstr = Self::build_pageserver_connstr(&pageservers);
-            spec.pageserver_connstring = Some(pageserver_connstr);
-            if stripe_size.is_some() {
-                spec.shard_stripe_size = stripe_size.map(|s| s.0 as usize);
-            }
+        let pageserver_connstr = Self::build_pageserver_connstr(&pageservers);
+        spec.pageserver_connstring = Some(pageserver_connstr);
+        if stripe_size.is_some() {
+            spec.shard_stripe_size = stripe_size.map(|s| s.0 as usize);
         }
 
         // If safekeepers are not specified, don't change them.
         if let Some(safekeepers) = safekeepers {
             let safekeeper_connstrings = self.build_safekeepers_connstrs(safekeepers)?;
             spec.safekeeper_connstrings = safekeeper_connstrings;
-            if let Some(g) = safekeeper_generation {
-                spec.safekeepers_generation = Some(g.into_inner());
-            }
         }
 
         let client = reqwest::Client::builder()
@@ -1046,24 +1061,6 @@ impl Endpoint {
         }
     }
 
-    pub async fn reconfigure_pageservers(
-        &self,
-        pageservers: Vec<(PageserverProtocol, Host, u16)>,
-        stripe_size: Option<ShardStripeSize>,
-    ) -> Result<()> {
-        self.reconfigure(Some(pageservers), stripe_size, None, None)
-            .await
-    }
-
-    pub async fn reconfigure_safekeepers(
-        &self,
-        safekeepers: Vec<NodeId>,
-        generation: SafekeeperGeneration,
-    ) -> Result<()> {
-        self.reconfigure(None, None, Some(safekeepers), Some(generation))
-            .await
-    }
-
     pub async fn stop(
         &self,
         mode: EndpointTerminateMode,

control_plane/src/local_env.rs

@@ -12,7 +12,6 @@ use std::{env, fs};
 
 use anyhow::{Context, bail};
 use clap::ValueEnum;
-use pageserver_api::config::PostHogConfig;
 use pem::Pem;
 use postgres_backend::AuthType;
 use reqwest::{Certificate, Url};
@@ -212,9 +211,7 @@ pub struct NeonStorageControllerConf {
 
     pub use_local_compute_notifications: bool,
 
-    pub timeline_safekeeper_count: Option<usize>,
-
-    pub posthog_config: Option<PostHogConfig>,
+    pub timeline_safekeeper_count: Option<i64>,
 
     pub kick_secondary_downloads: Option<bool>,
 }
@@ -248,7 +245,6 @@ impl Default for NeonStorageControllerConf {
             use_https_safekeeper_api: false,
             use_local_compute_notifications: true,
             timeline_safekeeper_count: None,
-            posthog_config: None,
             kick_secondary_downloads: None,
         }
     }

control_plane/src/storage_controller.rs

@@ -638,28 +638,10 @@ impl StorageController {
             args.push("--timelines-onto-safekeepers".to_string());
         }
 
-        // neon_local is used in test environments where we often have less than 3 safekeepers.
-        if self.config.timeline_safekeeper_count.is_some() || self.env.safekeepers.len() < 3 {
-            let sk_cnt = self
-                .config
-                .timeline_safekeeper_count
-                .unwrap_or(self.env.safekeepers.len());
-
+        if let Some(sk_cnt) = self.config.timeline_safekeeper_count {
             args.push(format!("--timeline-safekeeper-count={sk_cnt}"));
         }
 
-        let mut envs = vec![
-            ("LD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
-            ("DYLD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
-        ];
-
-        if let Some(posthog_config) = &self.config.posthog_config {
-            envs.push((
-                "POSTHOG_CONFIG".to_string(),
-                serde_json::to_string(posthog_config)?,
-            ));
-        }
-
         println!("Starting storage controller");
 
         background_process::start_process(
@@ -667,7 +649,10 @@ impl StorageController {
             &instance_dir,
             &self.env.storage_controller_bin(),
             args,
-            envs,
+            vec![
+                ("LD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
+                ("DYLD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
+            ],
             background_process::InitialPidFile::Create(self.pid_file(start_args.instance_id)),
             &start_args.start_timeout,
             || async {

docker-compose/compute_wrapper/var/db/postgres/configs/config.json

@@ -4,7 +4,6 @@
 
         "timestamp": "2022-10-12T18:00:00.000Z",
         "operation_uuid": "0f657b36-4b0f-4a2d-9c2e-1dcd615e7d8c",
-        "suspend_timeout_seconds": -1,
 
         "cluster": {
             "cluster_id": "docker_compose",

docs/rfcs/040-Endpoint-Persistent-Unlogged-Files-Storage.md

@@ -1,396 +0,0 @@
-# Memo: Endpoint Persistent Unlogged Files Storage
-Created on 2024-11-05
-Implemented on N/A
-
-## Summary
-A design for a storage system that allows storage of files required to make
-Neon's Endpoints have a better experience at or after a reboot.
-
-## Motivation
-Several systems inside PostgreSQL (and Neon) need some persistent storage for
-optimal workings across reboots and restarts, but still work without.
-Examples are the query-level statistics files of `pg_stat_statements` in
-`pg_stat/pg_stat_statements.stat`, and `pg_prewarm`'s `autoprewarm.blocks`.
-We need a storage system that can store and manage these files for each
-Endpoint, without necessarily granting users access to an unlimited storage
-device.
-
-## Goals
-- Store known files for Endpoints with reasonable persistence.  
-  _Data loss in this service, while annoying and bad for UX, won't lose any
-  customer's data._
-
-## Non Goals (if relevant)
-- This storage system does not need branching, file versioning, or other such
-  features. The files are as ephemeral to the timeline of the data as the
-  Endpoints that host the data.
-- This storage system does not need to store _all_ user files, only 'known'
-  user files.
-- This storage system does not need to be hosted fully inside Computes.  
-  _Instead, this will be a separate component similar to Pageserver,
-  SafeKeeper, the S3 proxy used for dynamically loaded extensions, etc._
-
-## Impacted components
-- Compute needs new code to load and store these files in its lifetime.
-- Control Plane needs to consider this new storage system when signalling
-  the deletion of an Endpoint, Timeline, or Tenant.
-- Control Plane needs to consider this new storage system when it resets
-  or re-assigns an endpoint's timeline/branch state.
-
-A new service is created: the Endpoint Persistent Unlogged Files Storage
-service.  This could be integrated in e.g. Pageserver or Control Plane, or a
-separately hosted service.
-
-## Proposed implementation
-Endpoint-related data files are managed by a newly designed service (which
-optionally is integrated in an existing service like Pageserver or Control
-Plane), which stores data directly into S3 or any blob storage of choice.
-
-Upon deletion of the Endpoint, or reassignment of the endpoint to a different
-branch, this ephemeral data is dropped: the data stored may not match the
-state of the branch's data after reassignment, and on endpoint deletion the
-data won't have any use to the user.
-
-Compute gets credentials (JWT token with Tenant, Timeline & Endpoint claims)
-which it can use to authenticate to this new service and retrieve and store
-data associated with this endpoint.  This limited scope reduces leaks of data
-across endpoints and timeline resets, and limits the ability of endpoints to
-mess with other endpoints' data.
-
-The path of this endpoint data in S3 is initially as follows:
-
-    s3://<regional-epufs-bucket>/
-      tenants/
-        <hex-tenant-id>/
-          tenants/
-            <hex-timeline-id>/
-              endpoints/
-                <endpoint-id>/
-                  pgdata/
-                    <file_path_in_pgdatadir>
-
-For other blob storages an equivalent or similar path can be constructed.
-
-### Reliability, failure modes and corner cases (if relevant)
-Reliability is important, but not critical to the workings of Neon.  The data
-stored in this service will, when lost, reduce performance, but won't be a
-cause of permanent data loss - only operational metadata is stored.
-
-Most, if not all, blob storage services have sufficiently high persistence
-guarantees to cater our need for persistence and uptime. The only concern with
-blob storages is that the access latency is generally higher than local disk,
-but for the object types stored (cache state, ...) I don't think this will be
-much of an issue.
-
-### Interaction/Sequence diagram (if relevant)
-
-In these diagrams you can replace S3 with any persistent storage device of
-choice, but S3 is chosen as representative name: The well-known and short name
-of AWS' blob storage. Azure Blob Storage should work too, but it has a much
-longer name making it less practical for the diagrams.
-
-Write data:
-
-```http
-POST /tenants/<tenant-id>/timelines/<tl-id>/endpoints/<endpoint-id>/pgdata/<the-pgdata-path>
-Host: epufs.svc.neon.local
-
-<<<
-
-200 OK
-{
-  "version": "<opaque>", # opaque file version token, changes when the file contents change
-  "size": <bytes>,
-}
-```
-
-```mermaid
-sequenceDiagram
-    autonumber
-    participant co as Compute
-    participant ep as EPUFS
-    participant s3 as Blob Storage
-
-    co-->ep: Connect with credentials
-    co->>+ep: Store Unlogged Persistent File
-    opt is authenticated
-        ep->>s3: Write UPF to S3
-    end
-    ep->>-co: OK / Failure / Auth Failure
-    co-->ep: Cancel connection
-```
-
-Read data: (optional with cache-relevant request parameters, e.g. If-Modified-Since)
-```http
-GET /tenants/<tenant-id>/timelines/<tl-id>/endpoints/<endpoint-id>/pgdata/<the-pgdata-path>
-Host: epufs.svc.neon.local
-
-<<<
-
-200 OK
-
-<file data>
-```
-
-```mermaid
-sequenceDiagram
-    autonumber
-    participant co as Compute
-    participant ep as EPUFS
-    participant s3 as Blob Storage
-
-    co->>+ep: Read Unlogged Persistent File
-    opt is authenticated
-        ep->>+s3: Request UPF from storage
-        s3->>-ep: Receive UPF from storage
-    end
-    ep->>-co: OK(response) / Failure(storage, auth, ...)
-```
-
-Compute Startup:
-```mermaid
-sequenceDiagram
-    autonumber
-    participant co as Compute
-    participant ps as Pageserver
-    participant ep as EPUFS
-    participant es as Extension server
-
-    note over co: Bind endpoint ep-xxx
-    par Get basebackup
-        co->>+ps: Request basebackup @ LSN
-        ps-)ps: Construct basebackup
-        ps->>-co: Receive basebackup TAR @ LSN
-    and Get startup-critical Unlogged Persistent Files
-        co->>+ep: Get all UPFs of endpoint ep-xxx
-        ep-)ep: Retrieve and gather all UPFs
-        ep->>-co: TAR of UPFs
-    and Get startup-critical extensions
-        loop For every startup-critical extension
-            co->>es: Get critical extension
-            es->>co: Receive critical extension
-        end
-    end
-    note over co: Start compute
-```
-
-CPlane ops:
-```http
-DELETE /tenants/<tenant-id>/timelines/<timeline-id>/endpoints/<endpoint-id>
-Host: epufs.svc.neon.local
-
-<<<
-
-200 OK
-{
-  "tenant": "<tenant-id>",
-  "timeline": "<timeline-id>",
-  "endpoint": "<endpoint-id>",
-  "deleted": {
-    "files": <count>,
-    "bytes": <count>,
-  },
-}
-```
-
-```http
-DELETE /tenants/<tenant-id>/timelines/<timeline-id>
-Host: epufs.svc.neon.local
-
-<<<
-
-200 OK
-{
-  "tenant": "<tenant-id>",
-  "timeline": "<timeline-id>",
-  "deleted": {
-    "files": <count>,
-    "bytes": <count>,
-  },
-}
-```
-
-```http
-DELETE /tenants/<tenant-id>
-Host: epufs.svc.neon.local
-
-<<<
-
-200 OK
-{
-  "tenant": "<tenant-id>",
-  "deleted": {
-    "files": <count>,
-    "bytes": <count>,
-  },
-}
-```
-
-```mermaid
-sequenceDiagram
-    autonumber
-    participant cp as Control Plane
-    participant ep as EPUFS
-    participant s3 as Blob Storage
-
-    alt Tenant deleted
-        cp-)ep: Tenant deleted
-        loop For every object associated with removed tenant
-            ep->>s3: Remove data of deleted tenant from Storage
-        end
-        opt
-            ep-)cp: Tenant cleanup complete
-        end
-    alt Timeline deleted
-        cp-)ep: Timeline deleted
-        loop For every object associated with removed timeline
-            ep->>s3: Remove data of deleted timeline from Storage
-        end
-        opt
-            ep-)cp: Timeline cleanup complete
-        end
-    else Endpoint reassigned or removed
-        cp->>+ep: Endpoint reassigned
-        loop For every object associated with reassigned/removed endpoint
-            ep->>s3: Remove data from Storage
-        end
-        ep->>-cp: Cleanup complete
-    end
-```
-
-### Scalability (if relevant)
-
-Provisionally:  As this service is going to be part of compute startup, this
-service should be able to quickly respond to all requests.  Therefore this
-service is deployed to every AZ we host Computes in, and Computes communicate
-(generally) only to the EPUFS endpoint of the AZ they're hosted in.
-
-Local caching of frequently restarted endpoints' data or metadata may be
-needed for best performance.  However, due to the regional nature of stored
-data but zonal nature of the service deployment, we should be careful when we
-implement any local caching, as it is possible that computes in AZ 1 will
-update data originally written and thus cached by AZ 2.  Cache version tests
-and invalidation is therefore required if we want to roll out caching to this
-service, which is too broad a scope for an MVC.  This is why caching is left
-out of scope for this RFC, and should be considered separately after this RFC
-is implemented.
-
-### Security implications (if relevant)
-This service must be able to authenticate users at least by Tenant ID,
-Timeline ID and Endpoint ID. This will use the existing JWT infrastructure of
-Compute, which will be upgraded to the extent needed to support Timeline- and
-Endpoint-based claims.
-
-The service requires unlimited access to (a prefix of) a blob storage bucket,
-and thus must be hosted outside the Compute VM sandbox.
-
-A service that generates pre-signed request URLs for Compute to download the
-data from that URL is likely problematic, too:  Compute would be able to write
-unlimited data to the bucket, or exfiltrate this signed URL to get read/write
-access to specific objects in this bucket, which would still effectively give
-users access to the S3 bucket (but with improved access logging).
-
-There may be a use case for transferring data associated with one endpoint to
-another endpoint (e.g. to make one endpoint warm its caches with the state of
-another endpoint), but that's not currently in scope, and specific needs may
-be solved through out-of-line communication of data or pre-signed URLs.
-
-### Unresolved questions (if relevant)
-Caching of files is not in the implementation scope of the document, but
-should at some future point be considered to maximize performance.
-
-## Alternative implementation (if relevant)
-Several ideas have come up to solve this issue:
-
-### Use AUXfile
-One prevalent idea was to WAL-log the files using our AUXfile mechanism.
-
-Benefits:
-
-+ We already have this storage mechanism
-
-Demerits:
-
-- It isn't available on read replicas
-- Additional WAL will be consumed during shutdown and after the shutdown
-  checkpoint, which needs PG modifications to work without panics.
-- It increases the data we need to manage in our versioned storage, thus
-  causing higher storage costs with higher retention due to duplication at
-  the storage layer.
-
-### Sign URLs for read/write operations, instead of proxying them
-
-Benefits:
-
-+ The service can be implemented with a much reduced IO budget
-
-Demerits:
-
-- Users could get access to these signed credentials
-- Not all blob storage services may implement URL signing
-
-### Give endpoints each their own directly accessed block volume
-
-Benefits:
-
-+ Easier to integrate for PostgreSQL
-
-Demerits:
-
-- Little control on data size and contents
-- Potentially problematic as we'd need to store data all across the pgdata
-  directory.
-- EBS is not a good candidate
-   - Attaches in 10s of seconds, if not more; i.e. too cold to start
-   - Shared EBS volumes are a no-go, as you'd have to schedule the endpoint
-     with users of the same EBS volumes, which can't work with VM migration
-   - EBS storage costs are very high (>80$/kilotenant when using a
-     volume/tenant)
-   - EBS volumes can't be mounted across AZ boundaries
-- Bucket per endpoint is unfeasible
-   - S3 buckets are priced at $20/month per 1k, which we could better spend
-     on developers.
-   - Allocating service accounts takes time (100s of ms), and service accounts
-     are a limited resource, too; so they're not a good candidate to allocate
-     on a per-endpoint basis.
-   - Giving credentials limited to prefix has similar issues as the pre-signed
-     URL approach.
-   - Bucket DNS lookup will fill DNS caches and put pressure on DNS lookup
-     much more than our current systems would.
-- Volumes bound by hypervisor are unlikely
-   - This requires significant investment and increased software on the
-     hypervisor.
-   - It is unclear if we can attach volumes after boot, i.e. for pooled
-     instances.
-
-### Put the files into a table
-
-Benefits:
-
- + Mostly already available in PostgreSQL
-
-Demerits:
-
- - Uses WAL
-   - Can't be used after shutdown checkpoint
-   - Needs a RW endpoint, and table & catalog access to write to this data
- - Gets hit with DB size limitations
- - Depending on user acces:
-   - Inaccessible:  
-     The user doesn't have control over database size caused by
-     these systems.
-   - Accessible:  
-     The user can corrupt these files and cause the system to crash while
-     user-corrupted files are present, thus increasing on-call overhead.
-
-## Definition of Done (if relevant)
-
-This project is done if we have:
-
-- One S3 bucket equivalent per region, which stores this per-endpoint data.
-- A new service endpoint in at least every AZ, which indirectly grants
-  endpoints access to the data stored for these endpoints in these buckets.
-- Compute writes & reads temp-data at shutdown and startup, respectively, for
-  at least the pg_prewarm or lfc_prewarm state files.
-- Cleanup of endpoint data is triggered when the endpoint is deleted or is
-  detached from its current timeline.

libs/compute_api/Cargo.toml

@@ -12,7 +12,6 @@ jsonwebtoken.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 regex.workspace = true
-url.workspace = true
 
 utils = { path = "../utils" }
 remote_storage = { version = "0.1", path = "../remote_storage/" }

libs/compute_api/src/spec.rs

@@ -4,14 +4,11 @@
 //! provide it by calling the compute_ctl's `/compute_ctl` endpoint, or
 //! compute_ctl can fetch it by calling the control plane's API.
 use std::collections::HashMap;
-use std::fmt::Display;
 
-use anyhow::anyhow;
 use indexmap::IndexMap;
 use regex::Regex;
 use remote_storage::RemotePath;
 use serde::{Deserialize, Serialize};
-use url::Url;
 use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 
@@ -184,11 +181,6 @@ pub struct ComputeSpec {
     /// Download LFC state from endpoint_storage and pass it to Postgres on startup
     #[serde(default)]
     pub autoprewarm: bool,
-
-    /// Suspend timeout in seconds.
-    ///
-    /// We use this value to derive other values, such as the installed extensions metric.
-    pub suspend_timeout_seconds: i64,
 }
 
 /// Feature flag to signal `compute_ctl` to enable certain experimental functionality.
@@ -437,47 +429,6 @@ pub struct JwksSettings {
     pub jwt_audience: Option<String>,
 }
 
-/// Protocol used to connect to a Pageserver. Parsed from the connstring scheme.
-#[derive(Clone, Copy, Debug, Default)]
-pub enum PageserverProtocol {
-    /// The original protocol based on libpq and COPY. Uses postgresql:// or postgres:// scheme.
-    #[default]
-    Libpq,
-    /// A newer, gRPC-based protocol. Uses grpc:// scheme.
-    Grpc,
-}
-
-impl PageserverProtocol {
-    /// Parses the protocol from a connstring scheme. Defaults to Libpq if no scheme is given.
-    /// Errors if the connstring is an invalid URL.
-    pub fn from_connstring(connstring: &str) -> anyhow::Result<Self> {
-        let scheme = match Url::parse(connstring) {
-            Ok(url) => url.scheme().to_lowercase(),
-            Err(url::ParseError::RelativeUrlWithoutBase) => return Ok(Self::default()),
-            Err(err) => return Err(anyhow!("invalid connstring URL: {err}")),
-        };
-        match scheme.as_str() {
-            "postgresql" | "postgres" => Ok(Self::Libpq),
-            "grpc" => Ok(Self::Grpc),
-            scheme => Err(anyhow!("invalid protocol scheme: {scheme}")),
-        }
-    }
-
-    /// Returns the URL scheme for the protocol, for use in connstrings.
-    pub fn scheme(&self) -> &'static str {
-        match self {
-            Self::Libpq => "postgresql",
-            Self::Grpc => "grpc",
-        }
-    }
-}
-
-impl Display for PageserverProtocol {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_str(self.scheme())
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::fs::File;

libs/compute_api/tests/cluster_spec.json

@@ -3,7 +3,6 @@
 
     "timestamp": "2021-05-23T18:25:43.511Z",
     "operation_uuid": "0f657b36-4b0f-4a2d-9c2e-1dcd615e7d8b",
-    "suspend_timeout_seconds": 3600,
 
     "cluster": {
         "cluster_id": "test-cluster-42",

libs/pageserver_api/Cargo.toml

@@ -19,7 +19,6 @@ byteorder.workspace = true
 utils.workspace = true
 postgres_ffi_types.workspace = true
 postgres_versioninfo.workspace = true
-posthog_client_lite.workspace = true
 enum-map.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
@@ -30,13 +29,12 @@ humantime-serde.workspace = true
 chrono = { workspace = true, features = ["serde"] }
 itertools.workspace = true
 storage_broker.workspace = true
-camino = { workspace = true, features = ["serde1"] }
+camino = {workspace = true, features = ["serde1"]}
 remote_storage.workspace = true
 postgres_backend.workspace = true
-nix = { workspace = true, optional = true }
+nix = {workspace = true, optional = true}
 reqwest.workspace = true
 rand.workspace = true
-tracing.workspace = true
 tracing-utils.workspace = true
 once_cell.workspace = true
 

libs/pageserver_api/src/config.rs

@@ -4,7 +4,6 @@ use camino::Utf8PathBuf;
 mod tests;
 
 use const_format::formatcp;
-use posthog_client_lite::PostHogClientConfig;
 pub const DEFAULT_PG_LISTEN_PORT: u16 = 64000;
 pub const DEFAULT_PG_LISTEN_ADDR: &str = formatcp!("127.0.0.1:{DEFAULT_PG_LISTEN_PORT}");
 pub const DEFAULT_HTTP_LISTEN_PORT: u16 = 9898;
@@ -64,66 +63,25 @@ impl Display for NodeMetadata {
     }
 }
 
-/// PostHog integration config. This is used in pageserver, storcon, and neon_local.
-/// Ensure backward compatibility when adding new fields.
+/// PostHog integration config.
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub struct PostHogConfig {
     /// PostHog project ID
-    #[serde(default)]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub project_id: Option<String>,
+    pub project_id: String,
     /// Server-side (private) API key
-    #[serde(default)]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub server_api_key: Option<String>,
+    pub server_api_key: String,
     /// Client-side (public) API key
-    #[serde(default)]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub client_api_key: Option<String>,
+    pub client_api_key: String,
     /// Private API URL
-    #[serde(default)]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub private_api_url: Option<String>,
+    pub private_api_url: String,
     /// Public API URL
-    #[serde(default)]
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub public_api_url: Option<String>,
-    /// Refresh interval for the feature flag spec.
-    /// The storcon will push the feature flag spec to the pageserver. If the pageserver does not receive
-    /// the spec for `refresh_interval`, it will fetch the spec from the PostHog API.
-    #[serde(default)]
+    pub public_api_url: String,
+    /// Refresh interval for the feature flag spec
     #[serde(skip_serializing_if = "Option::is_none")]
     #[serde(with = "humantime_serde")]
     pub refresh_interval: Option<Duration>,
 }
 
-impl PostHogConfig {
-    pub fn try_into_posthog_config(self) -> Result<PostHogClientConfig, &'static str> {
-        let Some(project_id) = self.project_id else {
-            return Err("project_id is required");
-        };
-        let Some(server_api_key) = self.server_api_key else {
-            return Err("server_api_key is required");
-        };
-        let Some(client_api_key) = self.client_api_key else {
-            return Err("client_api_key is required");
-        };
-        let Some(private_api_url) = self.private_api_url else {
-            return Err("private_api_url is required");
-        };
-        let Some(public_api_url) = self.public_api_url else {
-            return Err("public_api_url is required");
-        };
-        Ok(PostHogClientConfig {
-            project_id,
-            server_api_key,
-            client_api_key,
-            private_api_url,
-            public_api_url,
-        })
-    }
-}
-
 /// `pageserver.toml`
 ///
 /// We use serde derive with `#[serde(default)]` to generate a deserializer
@@ -409,9 +367,6 @@ pub struct BasebackupCacheConfig {
     // TODO(diko): support max_entry_size_bytes.
     // pub max_entry_size_bytes: u64,
     pub max_size_entries: usize,
-    /// Size of the channel used to send prepare requests to the basebackup cache worker.
-    /// If exceeded, new prepare requests will be dropped.
-    pub prepare_channel_size: usize,
 }
 
 impl Default for BasebackupCacheConfig {
@@ -420,8 +375,7 @@ impl Default for BasebackupCacheConfig {
             cleanup_period: Duration::from_secs(60),
             max_total_size_bytes: 1024 * 1024 * 1024, // 1 GiB
             // max_entry_size_bytes: 16 * 1024 * 1024,   // 16 MiB
-            max_size_entries: 10000,
-            prepare_channel_size: 100,
+            max_size_entries: 1000,
         }
     }
 }

libs/pageserver_api/src/controller_api.rs

@@ -546,11 +546,6 @@ pub struct TimelineImportRequest {
     pub sk_set: Vec<NodeId>,
 }
 
-#[derive(serde::Serialize, serde::Deserialize, Clone)]
-pub struct TimelineSafekeeperMigrateRequest {
-    pub new_sk_set: Vec<NodeId>,
-}
-
 #[cfg(test)]
 mod test {
     use serde_json;

libs/pageserver_api/src/models.rs

@@ -21,9 +21,7 @@ use utils::{completion, serde_system_time};
 
 use crate::config::Ratio;
 use crate::key::{CompactKey, Key};
-use crate::shard::{
-    DEFAULT_STRIPE_SIZE, ShardCount, ShardIdentity, ShardStripeSize, TenantShardId,
-};
+use crate::shard::{DEFAULT_STRIPE_SIZE, ShardCount, ShardStripeSize, TenantShardId};
 
 /// The state of a tenant in this pageserver.
 ///
@@ -477,7 +475,7 @@ pub struct TenantShardSplitResponse {
 }
 
 /// Parameters that apply to all shards in a tenant.  Used during tenant creation.
-#[derive(Clone, Copy, Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug)]
 #[serde(deny_unknown_fields)]
 pub struct ShardParameters {
     pub count: ShardCount,
@@ -499,15 +497,6 @@ impl Default for ShardParameters {
     }
 }
 
-impl From<ShardIdentity> for ShardParameters {
-    fn from(identity: ShardIdentity) -> Self {
-        Self {
-            count: identity.count,
-            stripe_size: identity.stripe_size,
-        }
-    }
-}
-
 #[derive(Debug, Default, Clone, Eq, PartialEq)]
 pub enum FieldPatch<T> {
     Upsert(T),

libs/pageserver_api/src/shard.rs

@@ -37,7 +37,6 @@ use std::hash::{Hash, Hasher};
 pub use ::utils::shard::*;
 use postgres_ffi_types::forknum::INIT_FORKNUM;
 use serde::{Deserialize, Serialize};
-use utils::critical;
 
 use crate::key::Key;
 use crate::models::ShardParameters;
@@ -180,7 +179,7 @@ impl ShardIdentity {
 
     /// For use when creating ShardIdentity instances for new shards, where a creation request
     /// specifies the ShardParameters that apply to all shards.
-    pub fn from_params(number: ShardNumber, params: ShardParameters) -> Self {
+    pub fn from_params(number: ShardNumber, params: &ShardParameters) -> Self {
         Self {
             number,
             count: params.count,
@@ -189,17 +188,6 @@ impl ShardIdentity {
         }
     }
 
-    /// Asserts that the given shard identities are equal. Changes to shard parameters will likely
-    /// result in data corruption.
-    pub fn assert_equal(&self, other: ShardIdentity) {
-        if self != &other {
-            // TODO: for now, we're conservative and just log errors in production. Turn this into a
-            // real assertion when we're confident it doesn't misfire, and also reject requests that
-            // attempt to change it with an error response.
-            critical!("shard identity mismatch: {self:?} != {other:?}");
-        }
-    }
-
     fn is_broken(&self) -> bool {
         self.layout == LAYOUT_BROKEN
     }

libs/posthog_client_lite/src/background_loop.rs

@@ -1,22 +1,17 @@
 //! A background loop that fetches feature flags from PostHog and updates the feature store.
 
-use std::{
-    sync::Arc,
-    time::{Duration, SystemTime},
-};
+use std::{sync::Arc, time::Duration};
 
 use arc_swap::ArcSwap;
 use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, info_span};
 
-use crate::{
-    CaptureEvent, FeatureStore, LocalEvaluationResponse, PostHogClient, PostHogClientConfig,
-};
+use crate::{CaptureEvent, FeatureStore, PostHogClient, PostHogClientConfig};
 
 /// A background loop that fetches feature flags from PostHog and updates the feature store.
 pub struct FeatureResolverBackgroundLoop {
     posthog_client: PostHogClient,
-    feature_store: ArcSwap<(SystemTime, Arc<FeatureStore>)>,
+    feature_store: ArcSwap<FeatureStore>,
     cancel: CancellationToken,
 }
 
@@ -24,35 +19,11 @@ impl FeatureResolverBackgroundLoop {
     pub fn new(config: PostHogClientConfig, shutdown_pageserver: CancellationToken) -> Self {
         Self {
             posthog_client: PostHogClient::new(config),
-            feature_store: ArcSwap::new(Arc::new((
-                SystemTime::UNIX_EPOCH,
-                Arc::new(FeatureStore::new()),
-            ))),
+            feature_store: ArcSwap::new(Arc::new(FeatureStore::new())),
             cancel: shutdown_pageserver,
         }
     }
 
-    /// Update the feature store with a new feature flag spec bypassing the normal refresh loop.
-    pub fn update(&self, spec: String) -> anyhow::Result<()> {
-        let resp: LocalEvaluationResponse = serde_json::from_str(&spec)?;
-        self.update_feature_store_nofail(resp, "http_propagate");
-        Ok(())
-    }
-
-    fn update_feature_store_nofail(&self, resp: LocalEvaluationResponse, source: &'static str) {
-        let project_id = self.posthog_client.config.project_id.parse::<u64>().ok();
-        match FeatureStore::new_with_flags(resp.flags, project_id) {
-            Ok(feature_store) => {
-                self.feature_store
-                    .store(Arc::new((SystemTime::now(), Arc::new(feature_store))));
-                tracing::info!("Feature flag updated from {}", source);
-            }
-            Err(e) => {
-                tracing::warn!("Cannot process feature flag spec from {}: {}", source, e);
-            }
-        }
-    }
-
     pub fn spawn(
         self: Arc<Self>,
         handle: &tokio::runtime::Handle,
@@ -76,17 +47,6 @@ impl FeatureResolverBackgroundLoop {
                         _ = ticker.tick() => {}
                         _ = cancel.cancelled() => break
                     }
-                    {
-                        let last_update = this.feature_store.load().0;
-                        if let Ok(elapsed) = last_update.elapsed() {
-                            if elapsed < refresh_period {
-                                tracing::debug!(
-                                    "Skipping feature flag refresh because it's too soon"
-                                );
-                                continue;
-                            }
-                        }
-                    }
                     let resp = match this
                         .posthog_client
                         .get_feature_flags_local_evaluation()
@@ -98,7 +58,16 @@ impl FeatureResolverBackgroundLoop {
                             continue;
                         }
                     };
-                    this.update_feature_store_nofail(resp, "refresh_loop");
+                    let project_id = this.posthog_client.config.project_id.parse::<u64>().ok();
+                    match FeatureStore::new_with_flags(resp.flags, project_id) {
+                        Ok(feature_store) => {
+                            this.feature_store.store(Arc::new(feature_store));
+                            tracing::info!("Feature flag updated");
+                        }
+                        Err(e) => {
+                            tracing::warn!("Cannot process feature flag spec: {}", e);
+                        }
+                    }
                 }
                 tracing::info!("PostHog feature resolver stopped");
             }
@@ -123,6 +92,6 @@ impl FeatureResolverBackgroundLoop {
     }
 
     pub fn feature_store(&self) -> Arc<FeatureStore> {
-        self.feature_store.load().1.clone()
+        self.feature_store.load_full()
     }
 }

libs/posthog_client_lite/src/lib.rs

@@ -544,8 +544,17 @@ impl PostHogClient {
         self.config.server_api_key.starts_with("phs_")
     }
 
-    /// Get the raw JSON spec, same as `get_feature_flags_local_evaluation` but without parsing.
-    pub async fn get_feature_flags_local_evaluation_raw(&self) -> anyhow::Result<String> {
+    /// Fetch the feature flag specs from the server.
+    ///
+    /// This is unfortunately an undocumented API at:
+    /// - <https://posthog.com/docs/api/feature-flags#get-api-projects-project_id-feature_flags-local_evaluation>
+    /// - <https://posthog.com/docs/feature-flags/local-evaluation>
+    ///
+    /// The handling logic in [`FeatureStore`] mostly follows the Python API implementation.
+    /// See `_compute_flag_locally` in <https://github.com/PostHog/posthog-python/blob/master/posthog/client.py>
+    pub async fn get_feature_flags_local_evaluation(
+        &self,
+    ) -> anyhow::Result<LocalEvaluationResponse> {
         // BASE_URL/api/projects/:project_id/feature_flags/local_evaluation
         // with bearer token of self.server_api_key
         // OR
@@ -579,22 +588,7 @@ impl PostHogClient {
                 body
             ));
         }
-        Ok(body)
-    }
-
-    /// Fetch the feature flag specs from the server.
-    ///
-    /// This is unfortunately an undocumented API at:
-    /// - <https://posthog.com/docs/api/feature-flags#get-api-projects-project_id-feature_flags-local_evaluation>
-    /// - <https://posthog.com/docs/feature-flags/local-evaluation>
-    ///
-    /// The handling logic in [`FeatureStore`] mostly follows the Python API implementation.
-    /// See `_compute_flag_locally` in <https://github.com/PostHog/posthog-python/blob/master/posthog/client.py>
-    pub async fn get_feature_flags_local_evaluation(
-        &self,
-    ) -> Result<LocalEvaluationResponse, anyhow::Error> {
-        let raw = self.get_feature_flags_local_evaluation_raw().await?;
-        Ok(serde_json::from_str(&raw)?)
+        Ok(serde_json::from_str(&body)?)
     }
 
     /// Capture an event. This will only be used to report the feature flag usage back to PostHog, though

libs/proxy/tokio-postgres2/src/connect.rs

@@ -1,12 +1,13 @@
 use std::net::IpAddr;
 
 use postgres_protocol2::message::backend::Message;
+use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::TcpStream;
 use tokio::sync::mpsc;
 
 use crate::client::SocketConfig;
 use crate::codec::BackendMessage;
-use crate::config::Host;
+use crate::config::{Host, SslMode};
 use crate::connect_raw::connect_raw;
 use crate::connect_socket::connect_socket;
 use crate::connect_tls::connect_tls;
@@ -46,13 +47,7 @@ where
 {
     let socket = connect_socket(host_addr, host, port, config.connect_timeout).await?;
     let stream = connect_tls(socket, config.ssl_mode, tls).await?;
-    let RawConnection {
-        stream,
-        parameters,
-        delayed_notice,
-        process_id,
-        secret_key,
-    } = connect_raw(stream, config).await?;
+    let raw = connect_raw(stream, config).await?;
 
     let socket_config = SocketConfig {
         host_addr,
@@ -61,24 +56,46 @@ where
         connect_timeout: config.connect_timeout,
     };
 
-    let (client_tx, conn_rx) = mpsc::unbounded_channel();
-    let (conn_tx, client_rx) = mpsc::channel(4);
-    let client = Client::new(
-        client_tx,
-        client_rx,
-        socket_config,
-        config.ssl_mode,
-        process_id,
-        secret_key,
-    );
-
-    // delayed notices are always sent as "Async" messages.
-    let delayed = delayed_notice
-        .into_iter()
-        .map(|m| BackendMessage::Async(Message::NoticeResponse(m)))
-        .collect();
-
-    let connection = Connection::new(stream, delayed, parameters, conn_tx, conn_rx);
-
-    Ok((client, connection))
+    Ok(raw.into_managed_conn(socket_config, config.ssl_mode))
+}
+
+impl<S, T> RawConnection<S, T>
+where
+    S: AsyncRead + AsyncWrite + Unpin,
+    T: AsyncRead + AsyncWrite + Unpin,
+{
+    pub fn into_managed_conn(
+        self,
+        socket_config: SocketConfig,
+        ssl_mode: SslMode,
+    ) -> (Client, Connection<S, T>) {
+        let RawConnection {
+            stream,
+            parameters,
+            delayed_notice,
+            process_id,
+            secret_key,
+        } = self;
+
+        let (client_tx, conn_rx) = mpsc::unbounded_channel();
+        let (conn_tx, client_rx) = mpsc::channel(4);
+        let client = Client::new(
+            client_tx,
+            client_rx,
+            socket_config,
+            ssl_mode,
+            process_id,
+            secret_key,
+        );
+
+        // delayed notices are always sent as "Async" messages.
+        let delayed = delayed_notice
+            .into_iter()
+            .map(|m| BackendMessage::Async(Message::NoticeResponse(m)))
+            .collect();
+
+        let connection = Connection::new(stream, delayed, parameters, conn_tx, conn_rx);
+
+        (client, connection)
+    }
 }

libs/safekeeper_api/src/models.rs

@@ -210,7 +210,7 @@ pub struct TimelineStatus {
 }
 
 /// Request to switch membership configuration.
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Serialize, Deserialize)]
 #[serde(transparent)]
 pub struct TimelineMembershipSwitchRequest {
     pub mconf: Configuration,
@@ -221,8 +221,6 @@ pub struct TimelineMembershipSwitchRequest {
 pub struct TimelineMembershipSwitchResponse {
     pub previous_conf: Configuration,
     pub current_conf: Configuration,
-    pub term: Term,
-    pub flush_lsn: Lsn,
 }
 
 #[derive(Clone, Copy, Serialize, Deserialize)]

libs/utils/src/sync/gate.rs

@@ -86,14 +86,6 @@ pub enum GateError {
     GateClosed,
 }
 
-impl GateError {
-    pub fn is_cancel(&self) -> bool {
-        match self {
-            GateError::GateClosed => true,
-        }
-    }
-}
-
 impl Default for Gate {
     fn default() -> Self {
         Self {

pageserver/client/src/mgmt_api.rs

@@ -844,13 +844,4 @@ impl Client {
             .await
             .map_err(Error::ReceiveBody)
     }
-
-    pub async fn update_feature_flag_spec(&self, spec: String) -> Result<()> {
-        let uri = format!("{}/v1/feature_flag_spec", self.mgmt_api_endpoint);
-        self.request(Method::POST, uri, spec)
-            .await?
-            .json()
-            .await
-            .map_err(Error::ReceiveBody)
-    }
 }

pageserver/page_api/Cargo.toml

@@ -9,14 +9,12 @@ anyhow.workspace = true
 bytes.workspace = true
 futures.workspace = true
 pageserver_api.workspace = true
-postgres_ffi_types.workspace = true
+postgres_ffi.workspace = true
 prost.workspace = true
-prost-types.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
 thiserror.workspace = true
 tokio.workspace = true
-tokio-util.workspace = true
 tonic.workspace = true
 utils.workspace = true
 workspace_hack.workspace = true

pageserver/page_api/proto/page_service.proto

@@ -35,8 +35,6 @@
 syntax = "proto3";
 package page_api;
 
-import "google/protobuf/timestamp.proto";
-
 service PageService {
   // Returns whether a relation exists.
   rpc CheckRelExists(CheckRelExistsRequest) returns (CheckRelExistsResponse);
@@ -66,10 +64,6 @@ service PageService {
 
   // Fetches an SLRU segment.
   rpc GetSlruSegment (GetSlruSegmentRequest) returns (GetSlruSegmentResponse);
-
-  // Acquires or extends a lease on the given LSN. This guarantees that the Pageserver won't garbage
-  // collect the LSN until the lease expires. Must be acquired on all relevant shards.
-  rpc LeaseLsn (LeaseLsnRequest) returns (LeaseLsnResponse);
 }
 
 // The LSN a request should read at.
@@ -116,19 +110,6 @@ message GetBaseBackupRequest {
   bool replica = 2;
   // If true, include relation files in the base backup. Mainly for debugging and tests.
   bool full = 3;
-  // Compression algorithm to use. Base backups send a compressed payload instead of using gRPC
-  // compression, so that we can cache compressed backups on the server.
-  BaseBackupCompression compression = 4;
-}
-
-// Base backup compression algorithms.
-enum BaseBackupCompression {
-  // Unknown algorithm. Used when clients send an unsupported algorithm.
-  BASE_BACKUP_COMPRESSION_UNKNOWN = 0;
-  // No compression.
-  BASE_BACKUP_COMPRESSION_NONE = 1;
-  // GZIP compression.
-  BASE_BACKUP_COMPRESSION_GZIP = 2;
 }
 
 // Base backup response chunk, returned as an ordered stream.
@@ -258,17 +239,3 @@ message GetSlruSegmentRequest {
 message GetSlruSegmentResponse {
   bytes segment = 1;
 }
-
-// Acquires or extends a lease on the given LSN. This guarantees that the Pageserver won't garbage
-// collect the LSN until the lease expires. Must be acquired on all relevant shards.
-message LeaseLsnRequest {
-  // The LSN to lease. Can't be 0 or below the current GC cutoff.
-  uint64 lsn = 1;
-}
-
-// Lease acquisition response. If the lease could not be granted because the LSN has already been
-// garbage collected, a FailedPrecondition status will be returned instead.
-message LeaseLsnResponse {
-  // The lease expiration time.
-  google.protobuf.Timestamp expires = 1;
-}

pageserver/page_api/src/client.rs

@@ -1,7 +1,8 @@
-use anyhow::Result;
-use futures::{Stream, StreamExt as _, TryStreamExt as _};
-use tokio::io::AsyncRead;
-use tokio_util::io::StreamReader;
+use std::convert::TryInto;
+
+use bytes::Bytes;
+use futures::TryStreamExt;
+use futures::{Stream, StreamExt};
 use tonic::metadata::AsciiMetadataValue;
 use tonic::metadata::errors::InvalidMetadataValue;
 use tonic::transport::Channel;
@@ -11,6 +12,8 @@ use utils::id::TenantId;
 use utils::id::TimelineId;
 use utils::shard::ShardIndex;
 
+use anyhow::Result;
+
 use crate::model;
 use crate::proto;
 
@@ -66,7 +69,6 @@ impl tonic::service::Interceptor for AuthInterceptor {
         Ok(req)
     }
 }
-
 #[derive(Clone)]
 pub struct Client {
     client: proto::PageServiceClient<
@@ -93,6 +95,7 @@ impl Client {
 
         if let Some(compression) = compression {
             // TODO: benchmark this (including network latency).
+            // TODO: consider enabling compression by default.
             client = client
                 .accept_compressed(compression)
                 .send_compressed(compression);
@@ -118,15 +121,22 @@ impl Client {
     pub async fn get_base_backup(
         &mut self,
         req: model::GetBaseBackupRequest,
-    ) -> Result<impl AsyncRead + use<>, tonic::Status> {
-        let req = proto::GetBaseBackupRequest::from(req);
-        let chunks = self.client.get_base_backup(req).await?.into_inner();
-        let reader = StreamReader::new(
-            chunks
-                .map_ok(|resp| resp.chunk)
-                .map_err(std::io::Error::other),
-        );
-        Ok(reader)
+    ) -> Result<impl Stream<Item = Result<Bytes, tonic::Status>> + 'static, tonic::Status> {
+        let proto_req = proto::GetBaseBackupRequest::from(req);
+
+        let response_stream: Streaming<proto::GetBaseBackupResponseChunk> =
+            self.client.get_base_backup(proto_req).await?.into_inner();
+
+        // TODO: Consider dechunking internally
+        let domain_stream = response_stream.map(|chunk_res| {
+            chunk_res.and_then(|proto_chunk| {
+                proto_chunk.try_into().map_err(|e| {
+                    tonic::Status::internal(format!("Failed to convert response chunk: {e}"))
+                })
+            })
+        });
+
+        Ok(domain_stream)
     }
 
     /// Returns the total size of a database, as # of bytes.
@@ -187,17 +197,4 @@ impl Client {
         let response = self.client.get_slru_segment(proto_req).await?;
         Ok(response.into_inner().try_into()?)
     }
-
-    /// Acquires or extends a lease on the given LSN. This guarantees that the Pageserver won't
-    /// garbage collect the LSN until the lease expires. Must be acquired on all relevant shards.
-    ///
-    /// Returns the lease expiration time, or a FailedPrecondition status if the lease could not be
-    /// acquired because the LSN has already been garbage collected.
-    pub async fn lease_lsn(
-        &mut self,
-        req: model::LeaseLsnRequest,
-    ) -> Result<model::LeaseLsnResponse, tonic::Status> {
-        let req = proto::LeaseLsnRequest::from(req);
-        Ok(self.client.lease_lsn(req).await?.into_inner().try_into()?)
-    }
 }

pageserver/page_api/src/model.rs

@@ -16,11 +16,10 @@
 //! stream combinators without dealing with errors, and avoids validating the same message twice.
 
 use std::fmt::Display;
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use bytes::Bytes;
-use postgres_ffi_types::Oid;
-// TODO: split out Lsn, RelTag, SlruKind and other basic types to a separate crate, to avoid
+use postgres_ffi::Oid;
+// TODO: split out Lsn, RelTag, SlruKind, Oid and other basic types to a separate crate, to avoid
 // pulling in all of their other crate dependencies when building the client.
 use utils::lsn::Lsn;
 
@@ -192,21 +191,15 @@ pub struct GetBaseBackupRequest {
     pub replica: bool,
     /// If true, include relation files in the base backup. Mainly for debugging and tests.
     pub full: bool,
-    /// Compression algorithm to use. Base backups send a compressed payload instead of using gRPC
-    /// compression, so that we can cache compressed backups on the server.
-    pub compression: BaseBackupCompression,
 }
 
-impl TryFrom<proto::GetBaseBackupRequest> for GetBaseBackupRequest {
-    type Error = ProtocolError;
-
-    fn try_from(pb: proto::GetBaseBackupRequest) -> Result<Self, Self::Error> {
-        Ok(Self {
+impl From<proto::GetBaseBackupRequest> for GetBaseBackupRequest {
+    fn from(pb: proto::GetBaseBackupRequest) -> Self {
+        Self {
             lsn: (pb.lsn != 0).then_some(Lsn(pb.lsn)),
             replica: pb.replica,
             full: pb.full,
-            compression: pb.compression.try_into()?,
-        })
+        }
     }
 }
 
@@ -216,55 +209,10 @@ impl From<GetBaseBackupRequest> for proto::GetBaseBackupRequest {
             lsn: request.lsn.unwrap_or_default().0,
             replica: request.replica,
             full: request.full,
-            compression: request.compression.into(),
         }
     }
 }
 
-/// Base backup compression algorithm.
-#[derive(Clone, Copy, Debug)]
-pub enum BaseBackupCompression {
-    None,
-    Gzip,
-}
-
-impl TryFrom<proto::BaseBackupCompression> for BaseBackupCompression {
-    type Error = ProtocolError;
-
-    fn try_from(pb: proto::BaseBackupCompression) -> Result<Self, Self::Error> {
-        match pb {
-            proto::BaseBackupCompression::Unknown => Err(ProtocolError::invalid("compression", pb)),
-            proto::BaseBackupCompression::None => Ok(Self::None),
-            proto::BaseBackupCompression::Gzip => Ok(Self::Gzip),
-        }
-    }
-}
-
-impl TryFrom<i32> for BaseBackupCompression {
-    type Error = ProtocolError;
-
-    fn try_from(compression: i32) -> Result<Self, Self::Error> {
-        proto::BaseBackupCompression::try_from(compression)
-            .map_err(|_| ProtocolError::invalid("compression", compression))
-            .and_then(Self::try_from)
-    }
-}
-
-impl From<BaseBackupCompression> for proto::BaseBackupCompression {
-    fn from(compression: BaseBackupCompression) -> Self {
-        match compression {
-            BaseBackupCompression::None => Self::None,
-            BaseBackupCompression::Gzip => Self::Gzip,
-        }
-    }
-}
-
-impl From<BaseBackupCompression> for i32 {
-    fn from(compression: BaseBackupCompression) -> Self {
-        proto::BaseBackupCompression::from(compression).into()
-    }
-}
-
 pub type GetBaseBackupResponseChunk = Bytes;
 
 impl TryFrom<proto::GetBaseBackupResponseChunk> for GetBaseBackupResponseChunk {
@@ -704,54 +652,3 @@ impl From<GetSlruSegmentResponse> for proto::GetSlruSegmentResponse {
 
 // SlruKind is defined in pageserver_api::reltag.
 pub type SlruKind = pageserver_api::reltag::SlruKind;
-
-/// Acquires or extends a lease on the given LSN. This guarantees that the Pageserver won't garbage
-/// collect the LSN until the lease expires.
-pub struct LeaseLsnRequest {
-    /// The LSN to lease.
-    pub lsn: Lsn,
-}
-
-impl TryFrom<proto::LeaseLsnRequest> for LeaseLsnRequest {
-    type Error = ProtocolError;
-
-    fn try_from(pb: proto::LeaseLsnRequest) -> Result<Self, Self::Error> {
-        if pb.lsn == 0 {
-            return Err(ProtocolError::Missing("lsn"));
-        }
-        Ok(Self { lsn: Lsn(pb.lsn) })
-    }
-}
-
-impl From<LeaseLsnRequest> for proto::LeaseLsnRequest {
-    fn from(request: LeaseLsnRequest) -> Self {
-        Self { lsn: request.lsn.0 }
-    }
-}
-
-/// Lease expiration time. If the lease could not be granted because the LSN has already been
-/// garbage collected, a FailedPrecondition status will be returned instead.
-pub type LeaseLsnResponse = SystemTime;
-
-impl TryFrom<proto::LeaseLsnResponse> for LeaseLsnResponse {
-    type Error = ProtocolError;
-
-    fn try_from(pb: proto::LeaseLsnResponse) -> Result<Self, Self::Error> {
-        let expires = pb.expires.ok_or(ProtocolError::Missing("expires"))?;
-        UNIX_EPOCH
-            .checked_add(Duration::new(expires.seconds as u64, expires.nanos as u32))
-            .ok_or_else(|| ProtocolError::invalid("expires", expires))
-    }
-}
-
-impl From<LeaseLsnResponse> for proto::LeaseLsnResponse {
-    fn from(response: LeaseLsnResponse) -> Self {
-        let expires = response.duration_since(UNIX_EPOCH).unwrap_or_default();
-        Self {
-            expires: Some(prost_types::Timestamp {
-                seconds: expires.as_secs() as i64,
-                nanos: expires.subsec_nanos() as i32,
-            }),
-        }
-    }
-}

pageserver/pagebench/src/cmd/basebackup.rs

@@ -317,7 +317,6 @@ impl Client for LibpqClient {
 /// A gRPC Pageserver client.
 struct GrpcClient {
     inner: page_api::Client,
-    compression: page_api::BaseBackupCompression,
 }
 
 impl GrpcClient {
@@ -332,14 +331,10 @@ impl GrpcClient {
             ttid.timeline_id,
             ShardIndex::unsharded(),
             None,
-            None, // NB: uses payload compression
+            compression.then_some(tonic::codec::CompressionEncoding::Zstd),
         )
         .await?;
-        let compression = match compression {
-            true => page_api::BaseBackupCompression::Gzip,
-            false => page_api::BaseBackupCompression::None,
-        };
-        Ok(Self { inner, compression })
+        Ok(Self { inner })
     }
 }
 
@@ -353,8 +348,10 @@ impl Client for GrpcClient {
             lsn,
             replica: false,
             full: false,
-            compression: self.compression,
         };
-        Ok(Box::pin(self.inner.get_base_backup(req).await?))
+        let stream = self.inner.get_base_backup(req).await?;
+        Ok(Box::pin(StreamReader::new(
+            stream.map_err(std::io::Error::other),
+        )))
     }
 }

pageserver/src/basebackup.rs

@@ -14,7 +14,6 @@ use std::fmt::Write as FmtWrite;
 use std::time::{Instant, SystemTime};
 
 use anyhow::{Context, anyhow};
-use async_compression::tokio::write::GzipEncoder;
 use bytes::{BufMut, Bytes, BytesMut};
 use fail::fail_point;
 use pageserver_api::key::{Key, rel_block_to_key};
@@ -26,7 +25,8 @@ use postgres_ffi::{
 };
 use postgres_ffi_types::constants::{DEFAULTTABLESPACE_OID, GLOBALTABLESPACE_OID};
 use postgres_ffi_types::forknum::{INIT_FORKNUM, MAIN_FORKNUM};
-use tokio::io::{self, AsyncWrite, AsyncWriteExt as _};
+use tokio::io;
+use tokio::io::AsyncWrite;
 use tokio_tar::{Builder, EntryType, Header};
 use tracing::*;
 use utils::lsn::Lsn;
@@ -97,7 +97,6 @@ impl From<BasebackupError> for tonic::Status {
 ///  * When working without safekeepers. In this situation it is important to match the lsn
 ///    we are taking basebackup on with the lsn that is used in pageserver's walreceiver
 ///    to start the replication.
-#[allow(clippy::too_many_arguments)]
 pub async fn send_basebackup_tarball<'a, W>(
     write: &'a mut W,
     timeline: &'a Timeline,
@@ -105,7 +104,6 @@ pub async fn send_basebackup_tarball<'a, W>(
     prev_lsn: Option<Lsn>,
     full_backup: bool,
     replica: bool,
-    gzip_level: Option<async_compression::Level>,
     ctx: &'a RequestContext,
 ) -> Result<(), BasebackupError>
 where
@@ -124,7 +122,7 @@ where
     // prev_lsn value; that happens if the timeline was just branched from
     // an old LSN and it doesn't have any WAL of its own yet. We will set
     // prev_lsn to Lsn(0) if we cannot provide the correct value.
-    let (backup_prev, lsn) = if let Some(req_lsn) = req_lsn {
+    let (backup_prev, backup_lsn) = if let Some(req_lsn) = req_lsn {
         // Backup was requested at a particular LSN. The caller should've
         // already checked that it's a valid LSN.
 
@@ -145,7 +143,7 @@ where
     };
 
     // Consolidate the derived and the provided prev_lsn values
-    let prev_record_lsn = if let Some(provided_prev_lsn) = prev_lsn {
+    let prev_lsn = if let Some(provided_prev_lsn) = prev_lsn {
         if backup_prev != Lsn(0) && backup_prev != provided_prev_lsn {
             return Err(BasebackupError::Server(anyhow!(
                 "backup_prev {backup_prev} != provided_prev_lsn {provided_prev_lsn}"
@@ -157,55 +155,30 @@ where
     };
 
     info!(
-        "taking basebackup lsn={lsn}, prev_lsn={prev_record_lsn} \
-        (full_backup={full_backup}, replica={replica}, gzip={gzip_level:?})",
-    );
-    let span = info_span!("send_tarball", backup_lsn=%lsn);
-
-    let io_concurrency = IoConcurrency::spawn_from_conf(
-        timeline.conf.get_vectored_concurrent_io,
-        timeline
-            .gate
-            .enter()
-            .map_err(|_| BasebackupError::Shutdown)?,
+        "taking basebackup lsn={}, prev_lsn={} (full_backup={}, replica={})",
+        backup_lsn, prev_lsn, full_backup, replica
     );
 
-    if let Some(gzip_level) = gzip_level {
-        let mut encoder = GzipEncoder::with_quality(write, gzip_level);
-        Basebackup {
-            ar: Builder::new_non_terminated(&mut encoder),
-            timeline,
-            lsn,
-            prev_record_lsn,
-            full_backup,
-            replica,
-            ctx,
-            io_concurrency,
-        }
+    let basebackup = Basebackup {
+        ar: Builder::new_non_terminated(write),
+        timeline,
+        lsn: backup_lsn,
+        prev_record_lsn: prev_lsn,
+        full_backup,
+        replica,
+        ctx,
+        io_concurrency: IoConcurrency::spawn_from_conf(
+            timeline.conf.get_vectored_concurrent_io,
+            timeline
+                .gate
+                .enter()
+                .map_err(|_| BasebackupError::Shutdown)?,
+        ),
+    };
+    basebackup
         .send_tarball()
-        .instrument(span)
-        .await?;
-        encoder
-            .shutdown()
-            .await
-            .map_err(|err| BasebackupError::Client(err, "gzip"))?;
-    } else {
-        Basebackup {
-            ar: Builder::new_non_terminated(write),
-            timeline,
-            lsn,
-            prev_record_lsn,
-            full_backup,
-            replica,
-            ctx,
-            io_concurrency,
-        }
-        .send_tarball()
-        .instrument(span)
-        .await?;
-    }
-
-    Ok(())
+        .instrument(info_span!("send_tarball", backup_lsn=%backup_lsn))
+        .await
 }
 
 /// This is short-living object only for the time of tarball creation,

pageserver/src/basebackup_cache.rs

@@ -1,12 +1,13 @@
 use std::{collections::HashMap, sync::Arc};
 
 use anyhow::Context;
+use async_compression::tokio::write::GzipEncoder;
 use camino::{Utf8Path, Utf8PathBuf};
 use metrics::core::{AtomicU64, GenericCounter};
 use pageserver_api::{config::BasebackupCacheConfig, models::TenantState};
 use tokio::{
     io::{AsyncWriteExt, BufWriter},
-    sync::mpsc::{Receiver, Sender, error::TrySendError},
+    sync::mpsc::{UnboundedReceiver, UnboundedSender},
 };
 use tokio_util::sync::CancellationToken;
 use utils::{
@@ -19,8 +20,8 @@ use crate::{
     basebackup::send_basebackup_tarball,
     context::{DownloadBehavior, RequestContext},
     metrics::{
-        BASEBACKUP_CACHE_ENTRIES, BASEBACKUP_CACHE_PREPARE, BASEBACKUP_CACHE_PREPARE_QUEUE_SIZE,
-        BASEBACKUP_CACHE_READ, BASEBACKUP_CACHE_SIZE,
+        BASEBACKUP_CACHE_ENTRIES, BASEBACKUP_CACHE_PREPARE, BASEBACKUP_CACHE_READ,
+        BASEBACKUP_CACHE_SIZE,
     },
     task_mgr::TaskKind,
     tenant::{
@@ -35,8 +36,8 @@ pub struct BasebackupPrepareRequest {
     pub lsn: Lsn,
 }
 
-pub type BasebackupPrepareSender = Sender<BasebackupPrepareRequest>;
-pub type BasebackupPrepareReceiver = Receiver<BasebackupPrepareRequest>;
+pub type BasebackupPrepareSender = UnboundedSender<BasebackupPrepareRequest>;
+pub type BasebackupPrepareReceiver = UnboundedReceiver<BasebackupPrepareRequest>;
 
 #[derive(Clone)]
 struct CacheEntry {
@@ -60,65 +61,40 @@ struct CacheEntry {
 /// and ~1 RPS for get requests.
 pub struct BasebackupCache {
     data_dir: Utf8PathBuf,
-    config: Option<BasebackupCacheConfig>,
 
     entries: std::sync::Mutex<HashMap<TenantTimelineId, CacheEntry>>,
 
-    prepare_sender: BasebackupPrepareSender,
-
     read_hit_count: GenericCounter<AtomicU64>,
     read_miss_count: GenericCounter<AtomicU64>,
     read_err_count: GenericCounter<AtomicU64>,
-
-    prepare_skip_count: GenericCounter<AtomicU64>,
 }
 
 impl BasebackupCache {
-    /// Create a new BasebackupCache instance.
-    /// Also returns a BasebackupPrepareReceiver which is needed to start
-    /// the background task.
-    /// The cache is initialized from the data_dir in the background task.
-    /// The cache will return `None` for any get requests until the initialization is complete.
-    /// The background task is spawned separately using [`Self::spawn_background_task`]
-    /// to avoid a circular dependency between the cache and the tenant manager.
-    pub fn new(
+    /// Creates a BasebackupCache and spawns the background task.
+    /// The initialization of the cache is performed in the background and does not
+    /// block the caller. The cache will return `None` for any get requests until
+    /// initialization is complete.
+    pub fn spawn(
+        runtime_handle: &tokio::runtime::Handle,
         data_dir: Utf8PathBuf,
         config: Option<BasebackupCacheConfig>,
-    ) -> (Arc<Self>, BasebackupPrepareReceiver) {
-        let chan_size = config.as_ref().map(|c| c.max_size_entries).unwrap_or(1);
-
-        let (prepare_sender, prepare_receiver) = tokio::sync::mpsc::channel(chan_size);
-
+        prepare_receiver: BasebackupPrepareReceiver,
+        tenant_manager: Arc<TenantManager>,
+        cancel: CancellationToken,
+    ) -> Arc<Self> {
         let cache = Arc::new(BasebackupCache {
             data_dir,
-            config,
+
             entries: std::sync::Mutex::new(HashMap::new()),
-            prepare_sender,
 
             read_hit_count: BASEBACKUP_CACHE_READ.with_label_values(&["hit"]),
             read_miss_count: BASEBACKUP_CACHE_READ.with_label_values(&["miss"]),
             read_err_count: BASEBACKUP_CACHE_READ.with_label_values(&["error"]),
-
-            prepare_skip_count: BASEBACKUP_CACHE_PREPARE.with_label_values(&["skip"]),
         });
 
-        (cache, prepare_receiver)
-    }
-
-    /// Spawns the background task.
-    /// The background task initializes the cache from the disk,
-    /// processes prepare requests, and cleans up outdated cache entries.
-    /// Noop if the cache is disabled (config is None).
-    pub fn spawn_background_task(
-        self: Arc<Self>,
-        runtime_handle: &tokio::runtime::Handle,
-        prepare_receiver: BasebackupPrepareReceiver,
-        tenant_manager: Arc<TenantManager>,
-        cancel: CancellationToken,
-    ) {
-        if let Some(config) = self.config.clone() {
+        if let Some(config) = config {
             let background = BackgroundTask {
-                c: self,
+                c: cache.clone(),
 
                 config,
                 tenant_manager,
@@ -133,45 +109,8 @@ impl BasebackupCache {
             };
             runtime_handle.spawn(background.run(prepare_receiver));
         }
-    }
 
-    /// Send a basebackup prepare request to the background task.
-    /// The basebackup will be prepared asynchronously, it does not block the caller.
-    /// The request will be skipped if any cache limits are exceeded.
-    pub fn send_prepare(&self, tenant_shard_id: TenantShardId, timeline_id: TimelineId, lsn: Lsn) {
-        let req = BasebackupPrepareRequest {
-            tenant_shard_id,
-            timeline_id,
-            lsn,
-        };
-
-        BASEBACKUP_CACHE_PREPARE_QUEUE_SIZE.inc();
-        let res = self.prepare_sender.try_send(req);
-
-        if let Err(e) = res {
-            BASEBACKUP_CACHE_PREPARE_QUEUE_SIZE.dec();
-            self.prepare_skip_count.inc();
-            match e {
-                TrySendError::Full(_) => {
-                    // Basebackup prepares are pretty rare, normally we should not hit this.
-                    tracing::info!(
-                        tenant_id = %tenant_shard_id.tenant_id,
-                        %timeline_id,
-                        %lsn,
-                        "Basebackup prepare channel is full, skipping the request"
-                    );
-                }
-                TrySendError::Closed(_) => {
-                    // Normal during shutdown, not critical.
-                    tracing::info!(
-                        tenant_id = %tenant_shard_id.tenant_id,
-                        %timeline_id,
-                        %lsn,
-                        "Basebackup prepare channel is closed, skipping the request"
-                    );
-                }
-            }
-        }
+        cache
     }
 
     /// Gets a basebackup entry from the cache.
@@ -184,10 +123,6 @@ impl BasebackupCache {
         timeline_id: TimelineId,
         lsn: Lsn,
     ) -> Option<tokio::fs::File> {
-        if !self.is_enabled() {
-            return None;
-        }
-
         // Fast path. Check if the entry exists using the in-memory state.
         let tti = TenantTimelineId::new(tenant_id, timeline_id);
         if self.entries.lock().unwrap().get(&tti).map(|e| e.lsn) != Some(lsn) {
@@ -215,10 +150,6 @@ impl BasebackupCache {
         }
     }
 
-    pub fn is_enabled(&self) -> bool {
-        self.config.is_some()
-    }
-
     // Private methods.
 
     fn entry_filename(tenant_id: TenantId, timeline_id: TimelineId, lsn: Lsn) -> String {
@@ -436,7 +367,6 @@ impl BackgroundTask {
         loop {
             tokio::select! {
                 Some(req) = prepare_receiver.recv() => {
-                    BASEBACKUP_CACHE_PREPARE_QUEUE_SIZE.dec();
                     if let Err(err) = self.prepare_basebackup(
                         req.tenant_shard_id,
                         req.timeline_id,
@@ -664,6 +594,13 @@ impl BackgroundTask {
         let file = tokio::fs::File::create(entry_tmp_path).await?;
         let mut writer = BufWriter::new(file);
 
+        let mut encoder = GzipEncoder::with_quality(
+            &mut writer,
+            // Level::Best because compression is not on the hot path of basebackup requests.
+            // The decompression is almost not affected by the compression level.
+            async_compression::Level::Best,
+        );
+
         // We may receive a request before the WAL record is applied to the timeline.
         // Wait for the requested LSN to be applied.
         timeline
@@ -676,19 +613,17 @@ impl BackgroundTask {
             .await?;
 
         send_basebackup_tarball(
-            &mut writer,
+            &mut encoder,
             timeline,
             Some(req_lsn),
             None,
             false,
             false,
-            // Level::Best because compression is not on the hot path of basebackup requests.
-            // The decompression is almost not affected by the compression level.
-            Some(async_compression::Level::Best),
             &ctx,
         )
         .await?;
 
+        encoder.shutdown().await?;
         writer.flush().await?;
         writer.into_inner().sync_all().await?;
 

pageserver/src/bin/pageserver.rs

@@ -569,10 +569,8 @@ fn start_pageserver(
         pageserver::l0_flush::L0FlushGlobalState::new(conf.l0_flush.clone());
 
     // Scan the local 'tenants/' directory and start loading the tenants
-    let (basebackup_cache, basebackup_prepare_receiver) = BasebackupCache::new(
-        conf.basebackup_cache_dir(),
-        conf.basebackup_cache_config.clone(),
-    );
+    let (basebackup_prepare_sender, basebackup_prepare_receiver) =
+        tokio::sync::mpsc::unbounded_channel();
     let deletion_queue_client = deletion_queue.new_client();
     let background_purges = mgr::BackgroundPurges::default();
 
@@ -584,7 +582,7 @@ fn start_pageserver(
             remote_storage: remote_storage.clone(),
             deletion_queue_client,
             l0_flush_global_state,
-            basebackup_cache: Arc::clone(&basebackup_cache),
+            basebackup_prepare_sender,
             feature_resolver: feature_resolver.clone(),
         },
         shutdown_pageserver.clone(),
@@ -592,8 +590,10 @@ fn start_pageserver(
     let tenant_manager = Arc::new(tenant_manager);
     BACKGROUND_RUNTIME.block_on(mgr::init_tenant_mgr(tenant_manager.clone(), order))?;
 
-    basebackup_cache.spawn_background_task(
+    let basebackup_cache = BasebackupCache::spawn(
         BACKGROUND_RUNTIME.handle(),
+        conf.basebackup_cache_dir(),
+        conf.basebackup_cache_config.clone(),
         basebackup_prepare_receiver,
         Arc::clone(&tenant_manager),
         shutdown_pageserver.child_token(),
@@ -806,6 +806,7 @@ fn start_pageserver(
         } else {
             None
         },
+        basebackup_cache,
     );
 
     // Spawn a Pageserver gRPC server task. It will spawn separate tasks for

pageserver/src/bin/test_helper_slow_client_reads.rs

@@ -37,7 +37,7 @@ async fn main() -> anyhow::Result<()> {
                 not_modified_since: Lsn(23),
             },
             batch_key: 42,
-            message: format!("message {msg}"),
+            message: format!("message {}", msg),
         }));
         let Ok(res) = tokio::time::timeout(Duration::from_secs(10), fut).await else {
             eprintln!("pipe seems full");

pageserver/src/config.rs

@@ -762,40 +762,4 @@ mod tests {
         let result = PageServerConf::parse_and_validate(NodeId(0), config_toml, &workdir);
         assert_eq!(result.is_ok(), is_valid);
     }
-
-    #[test]
-    fn test_config_posthog_config_is_valid() {
-        let input = r#"
-            control_plane_api = "http://localhost:6666"
-
-            [posthog_config]
-            server_api_key = "phs_AAA"
-            client_api_key = "phc_BBB"
-            project_id = "000"
-            private_api_url = "https://us.posthog.com"
-            public_api_url = "https://us.i.posthog.com"
-        "#;
-        let config_toml = toml_edit::de::from_str::<pageserver_api::config::ConfigToml>(input)
-            .expect("posthogconfig is valid");
-        let workdir = Utf8PathBuf::from("/nonexistent");
-        PageServerConf::parse_and_validate(NodeId(0), config_toml, &workdir)
-            .expect("parse_and_validate");
-    }
-
-    #[test]
-    fn test_config_posthog_incomplete_config_is_valid() {
-        let input = r#"
-            control_plane_api = "http://localhost:6666"
-
-            [posthog_config]
-            server_api_key = "phs_AAA"
-            private_api_url = "https://us.posthog.com"
-            public_api_url = "https://us.i.posthog.com"
-        "#;
-        let config_toml = toml_edit::de::from_str::<pageserver_api::config::ConfigToml>(input)
-            .expect("posthogconfig is valid");
-        let workdir = Utf8PathBuf::from("/nonexistent");
-        PageServerConf::parse_and_validate(NodeId(0), config_toml, &workdir)
-            .expect("parse_and_validate");
-    }
 }

pageserver/src/feature_resolver.rs

@@ -3,7 +3,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
 use arc_swap::ArcSwap;
 use pageserver_api::config::NodeMetadata;
 use posthog_client_lite::{
-    CaptureEvent, FeatureResolverBackgroundLoop, PostHogEvaluationError,
+    CaptureEvent, FeatureResolverBackgroundLoop, PostHogClientConfig, PostHogEvaluationError,
     PostHogFlagFilterPropertyValue,
 };
 use remote_storage::RemoteStorageKind;
@@ -31,13 +31,6 @@ impl FeatureResolver {
         }
     }
 
-    pub fn update(&self, spec: String) -> anyhow::Result<()> {
-        if let Some(inner) = &self.inner {
-            inner.update(spec)?;
-        }
-        Ok(())
-    }
-
     pub fn spawn(
         conf: &PageServerConf,
         shutdown_pageserver: CancellationToken,
@@ -45,24 +38,16 @@ impl FeatureResolver {
     ) -> anyhow::Result<Self> {
         // DO NOT block in this function: make it return as fast as possible to avoid startup delays.
         if let Some(posthog_config) = &conf.posthog_config {
-            let posthog_client_config = match posthog_config.clone().try_into_posthog_config() {
-                Ok(config) => config,
-                Err(e) => {
-                    tracing::warn!(
-                        "invalid posthog config, skipping posthog integration: {}",
-                        e
-                    );
-                    return Ok(FeatureResolver {
-                        inner: None,
-                        internal_properties: None,
-                        force_overrides_for_testing: Arc::new(ArcSwap::new(Arc::new(
-                            HashMap::new(),
-                        ))),
-                    });
-                }
-            };
-            let inner =
-                FeatureResolverBackgroundLoop::new(posthog_client_config, shutdown_pageserver);
+            let inner = FeatureResolverBackgroundLoop::new(
+                PostHogClientConfig {
+                    server_api_key: posthog_config.server_api_key.clone(),
+                    client_api_key: posthog_config.client_api_key.clone(),
+                    project_id: posthog_config.project_id.clone(),
+                    private_api_url: posthog_config.private_api_url.clone(),
+                    public_api_url: posthog_config.public_api_url.clone(),
+                },
+                shutdown_pageserver,
+            );
             let inner = Arc::new(inner);
 
             // The properties shared by all tenants on this pageserver.

pageserver/src/http/routes.rs

@@ -1893,13 +1893,9 @@ async fn update_tenant_config_handler(
     let location_conf = LocationConf::attached_single(
         new_tenant_conf.clone(),
         tenant.get_generation(),
-        ShardParameters::from(tenant.get_shard_identity()),
+        &ShardParameters::default(),
     );
 
-    tenant
-        .get_shard_identity()
-        .assert_equal(location_conf.shard); // not strictly necessary since we construct it above
-
     crate::tenant::TenantShard::persist_tenant_config(state.conf, &tenant_shard_id, &location_conf)
         .await
         .map_err(|e| ApiError::InternalServerError(anyhow::anyhow!(e)))?;
@@ -1941,13 +1937,9 @@ async fn patch_tenant_config_handler(
     let location_conf = LocationConf::attached_single(
         updated,
         tenant.get_generation(),
-        ShardParameters::from(tenant.get_shard_identity()),
+        &ShardParameters::default(),
     );
 
-    tenant
-        .get_shard_identity()
-        .assert_equal(location_conf.shard); // not strictly necessary since we construct it above
-
     crate::tenant::TenantShard::persist_tenant_config(state.conf, &tenant_shard_id, &location_conf)
         .await
         .map_err(|e| ApiError::InternalServerError(anyhow::anyhow!(e)))?;
@@ -3751,20 +3743,6 @@ async fn force_override_feature_flag_for_testing_delete(
     json_response(StatusCode::OK, ())
 }
 
-async fn update_feature_flag_spec(
-    mut request: Request<Body>,
-    _cancel: CancellationToken,
-) -> Result<Response<Body>, ApiError> {
-    check_permission(&request, None)?;
-    let body = json_request(&mut request).await?;
-    let state = get_state(&request);
-    state
-        .feature_resolver
-        .update(body)
-        .map_err(ApiError::InternalServerError)?;
-    json_response(StatusCode::OK, ())
-}
-
 /// Common functionality of all the HTTP API handlers.
 ///
 /// - Adds a tracing span to each request (by `request_span`)
@@ -4150,8 +4128,5 @@ pub fn make_router(
         .delete("/v1/feature_flag/:flag_key", |r| {
             testing_api_handler("force override feature flag - delete", r, force_override_feature_flag_for_testing_delete)
         })
-        .post("/v1/feature_flag_spec", |r| {
-            api_handler(r, update_feature_flag_spec)
-        })
         .any(handler_404))
 }

pageserver/src/metrics.rs

@@ -4439,14 +4439,6 @@ pub(crate) static BASEBACKUP_CACHE_SIZE: Lazy<UIntGauge> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
-pub(crate) static BASEBACKUP_CACHE_PREPARE_QUEUE_SIZE: Lazy<UIntGauge> = Lazy::new(|| {
-    register_uint_gauge!(
-        "pageserver_basebackup_cache_prepare_queue_size",
-        "Number of requests in the basebackup prepare channel"
-    )
-    .expect("failed to define a metric")
-});
-
 static PAGESERVER_CONFIG_IGNORED_ITEMS: Lazy<UIntGaugeVec> = Lazy::new(|| {
     register_uint_gauge_vec!(
         "pageserver_config_ignored_items",

pageserver/src/page_service.rs

@@ -12,9 +12,9 @@ use std::task::{Context, Poll};
 use std::time::{Duration, Instant, SystemTime};
 use std::{io, str};
 
-use anyhow::{Context as _, bail};
+use anyhow::{Context as _, anyhow, bail};
+use async_compression::tokio::write::GzipEncoder;
 use bytes::{Buf as _, BufMut as _, BytesMut};
-use chrono::Utc;
 use futures::future::BoxFuture;
 use futures::{FutureExt, Stream};
 use itertools::Itertools;
@@ -63,6 +63,7 @@ use utils::{failpoint_support, span_record};
 
 use crate::auth::check_permission;
 use crate::basebackup::{self, BasebackupError};
+use crate::basebackup_cache::BasebackupCache;
 use crate::config::PageServerConf;
 use crate::context::{
     DownloadBehavior, PerfInstrumentFutureExt, RequestContext, RequestContextBuilder,
@@ -137,6 +138,7 @@ pub fn spawn(
     perf_trace_dispatch: Option<Dispatch>,
     tcp_listener: tokio::net::TcpListener,
     tls_config: Option<Arc<rustls::ServerConfig>>,
+    basebackup_cache: Arc<BasebackupCache>,
 ) -> Listener {
     let cancel = CancellationToken::new();
     let libpq_ctx = RequestContext::todo_child(
@@ -158,6 +160,7 @@ pub fn spawn(
             conf.pg_auth_type,
             tls_config,
             conf.page_service_pipelining.clone(),
+            basebackup_cache,
             libpq_ctx,
             cancel.clone(),
         )
@@ -216,6 +219,7 @@ pub async fn libpq_listener_main(
     auth_type: AuthType,
     tls_config: Option<Arc<rustls::ServerConfig>>,
     pipelining_config: PageServicePipeliningConfig,
+    basebackup_cache: Arc<BasebackupCache>,
     listener_ctx: RequestContext,
     listener_cancel: CancellationToken,
 ) -> Connections {
@@ -259,6 +263,7 @@ pub async fn libpq_listener_main(
                     auth_type,
                     tls_config.clone(),
                     pipelining_config.clone(),
+                    Arc::clone(&basebackup_cache),
                     connection_ctx,
                     connections_cancel.child_token(),
                     gate_guard,
@@ -301,6 +306,7 @@ async fn page_service_conn_main(
     auth_type: AuthType,
     tls_config: Option<Arc<rustls::ServerConfig>>,
     pipelining_config: PageServicePipeliningConfig,
+    basebackup_cache: Arc<BasebackupCache>,
     connection_ctx: RequestContext,
     cancel: CancellationToken,
     gate_guard: GateGuard,
@@ -366,6 +372,7 @@ async fn page_service_conn_main(
         pipelining_config,
         conf.get_vectored_concurrent_io,
         perf_span_fields,
+        basebackup_cache,
         connection_ctx,
         cancel.clone(),
         gate_guard,
@@ -419,6 +426,8 @@ struct PageServerHandler {
     pipelining_config: PageServicePipeliningConfig,
     get_vectored_concurrent_io: GetVectoredConcurrentIo,
 
+    basebackup_cache: Arc<BasebackupCache>,
+
     gate_guard: GateGuard,
 }
 
@@ -904,6 +913,7 @@ impl PageServerHandler {
         pipelining_config: PageServicePipeliningConfig,
         get_vectored_concurrent_io: GetVectoredConcurrentIo,
         perf_span_fields: ConnectionPerfSpanFields,
+        basebackup_cache: Arc<BasebackupCache>,
         connection_ctx: RequestContext,
         cancel: CancellationToken,
         gate_guard: GateGuard,
@@ -917,6 +927,7 @@ impl PageServerHandler {
             cancel,
             pipelining_config,
             get_vectored_concurrent_io,
+            basebackup_cache,
             gate_guard,
         }
     }
@@ -2602,16 +2613,26 @@ impl PageServerHandler {
                 prev_lsn,
                 full_backup,
                 replica,
-                None,
                 &ctx,
             )
             .await?;
         } else {
             let mut writer = BufWriter::new(pgb.copyout_writer());
 
-            let cached = timeline
-                .get_cached_basebackup_if_enabled(lsn, prev_lsn, full_backup, replica, gzip)
-                .await;
+            let cached = {
+                // Basebackup is cached only for this combination of parameters.
+                if timeline.is_basebackup_cache_enabled()
+                    && gzip
+                    && lsn.is_some()
+                    && prev_lsn.is_none()
+                {
+                    self.basebackup_cache
+                        .get(tenant_id, timeline_id, lsn.unwrap())
+                        .await
+                } else {
+                    None
+                }
+            };
 
             if let Some(mut cached) = cached {
                 from_cache = true;
@@ -2620,6 +2641,31 @@ impl PageServerHandler {
                     .map_err(|err| {
                         BasebackupError::Client(err, "handle_basebackup_request,cached,copy")
                     })?;
+            } else if gzip {
+                let mut encoder = GzipEncoder::with_quality(
+                    &mut writer,
+                    // NOTE using fast compression because it's on the critical path
+                    //      for compute startup. For an empty database, we get
+                    //      <100KB with this method. The Level::Best compression method
+                    //      gives us <20KB, but maybe we should add basebackup caching
+                    //      on compute shutdown first.
+                    async_compression::Level::Fastest,
+                );
+                basebackup::send_basebackup_tarball(
+                    &mut encoder,
+                    &timeline,
+                    lsn,
+                    prev_lsn,
+                    full_backup,
+                    replica,
+                    &ctx,
+                )
+                .await?;
+                // shutdown the encoder to ensure the gzip footer is written
+                encoder
+                    .shutdown()
+                    .await
+                    .map_err(|e| QueryError::Disconnected(ConnectionError::Io(e)))?;
             } else {
                 basebackup::send_basebackup_tarball(
                     &mut writer,
@@ -2628,11 +2674,6 @@ impl PageServerHandler {
                     prev_lsn,
                     full_backup,
                     replica,
-                    // NB: using fast compression because it's on the critical path for compute
-                    // startup. For an empty database, we get <100KB with this method. The
-                    // Level::Best compression method gives us <20KB, but maybe we should add
-                    // basebackup caching on compute shutdown first.
-                    gzip.then_some(async_compression::Level::Fastest),
                     &ctx,
                 )
                 .await?;
@@ -3512,7 +3553,7 @@ impl proto::PageService for GrpcPageServiceHandler {
         if timeline.is_archived() == Some(true) {
             return Err(tonic::Status::failed_precondition("timeline is archived"));
         }
-        let req: page_api::GetBaseBackupRequest = req.into_inner().try_into()?;
+        let req: page_api::GetBaseBackupRequest = req.into_inner().into();
 
         span_record!(lsn=?req.lsn);
 
@@ -3538,50 +3579,20 @@ impl proto::PageService for GrpcPageServiceHandler {
         let span = Span::current();
         let (mut simplex_read, mut simplex_write) = tokio::io::simplex(CHUNK_SIZE);
         let jh = tokio::spawn(async move {
-            let gzip_level = match req.compression {
-                page_api::BaseBackupCompression::None => None,
-                // NB: using fast compression because it's on the critical path for compute
-                // startup. For an empty database, we get <100KB with this method. The
-                // Level::Best compression method gives us <20KB, but maybe we should add
-                // basebackup caching on compute shutdown first.
-                page_api::BaseBackupCompression::Gzip => Some(async_compression::Level::Fastest),
-            };
-
-            // Check for a cached basebackup.
-            let cached = timeline
-                .get_cached_basebackup_if_enabled(
-                    req.lsn,
-                    None,
-                    req.full,
-                    req.replica,
-                    gzip_level.is_some(),
-                )
-                .await;
-
-            let result = if let Some(mut cached) = cached {
-                // If we have a cached basebackup, send it.
-                tokio::io::copy(&mut cached, &mut simplex_write)
-                    .await
-                    .map(|_| ())
-                    .map_err(|err| BasebackupError::Client(err, "cached,copy"))
-            } else {
-                basebackup::send_basebackup_tarball(
-                    &mut simplex_write,
-                    &timeline,
-                    req.lsn,
-                    None,
-                    req.full,
-                    req.replica,
-                    gzip_level,
-                    &ctx,
-                )
-                .instrument(span) // propagate request span
-                .await
-            };
-            simplex_write
-                .shutdown()
-                .await
-                .map_err(|err| BasebackupError::Client(err, "simplex_write"))?;
+            let result = basebackup::send_basebackup_tarball(
+                &mut simplex_write,
+                &timeline,
+                req.lsn,
+                None,
+                req.full,
+                req.replica,
+                &ctx,
+            )
+            .instrument(span) // propagate request span
+            .await;
+            simplex_write.shutdown().await.map_err(|err| {
+                BasebackupError::Server(anyhow!("simplex shutdown failed: {err}"))
+            })?;
             result
         });
 
@@ -3761,36 +3772,6 @@ impl proto::PageService for GrpcPageServiceHandler {
         let resp: page_api::GetSlruSegmentResponse = resp.segment;
         Ok(tonic::Response::new(resp.into()))
     }
-
-    #[instrument(skip_all, fields(lsn))]
-    async fn lease_lsn(
-        &self,
-        req: tonic::Request<proto::LeaseLsnRequest>,
-    ) -> Result<tonic::Response<proto::LeaseLsnResponse>, tonic::Status> {
-        let timeline = self.get_request_timeline(&req).await?;
-        let ctx = self.ctx.with_scope_timeline(&timeline);
-
-        // Validate and convert the request, and decorate the span.
-        let req: page_api::LeaseLsnRequest = req.into_inner().try_into()?;
-
-        span_record!(lsn=%req.lsn);
-
-        // Attempt to acquire a lease. Return FailedPrecondition if the lease could not be granted.
-        let lease_length = timeline.get_lsn_lease_length();
-        let expires = match timeline.renew_lsn_lease(req.lsn, lease_length, &ctx) {
-            Ok(lease) => lease.valid_until,
-            Err(err) => return Err(tonic::Status::failed_precondition(format!("{err}"))),
-        };
-
-        // TODO: is this spammy? Move it compute-side?
-        info!(
-            "acquired lease for {} until {}",
-            req.lsn,
-            chrono::DateTime::<Utc>::from(expires).to_rfc3339()
-        );
-
-        Ok(tonic::Response::new(expires.into()))
-    }
 }
 
 /// gRPC middleware layer that handles observability concerns:

pageserver/src/pgdatadir_mapping.rs

@@ -3015,7 +3015,7 @@ mod tests {
         // This shard will get the even blocks
         let shard = ShardIdentity::from_params(
             ShardNumber(0),
-            ShardParameters {
+            &ShardParameters {
                 count: ShardCount(2),
                 stripe_size: ShardStripeSize(1),
             },

pageserver/src/tenant.rs

@@ -80,7 +80,7 @@ use self::timeline::uninit::{TimelineCreateGuard, TimelineExclusionError, Uninit
 use self::timeline::{
     EvictionTaskTenantState, GcCutoffs, TimelineDeleteProgress, TimelineResources, WaitLsnError,
 };
-use crate::basebackup_cache::BasebackupCache;
+use crate::basebackup_cache::BasebackupPrepareSender;
 use crate::config::PageServerConf;
 use crate::context;
 use crate::context::RequestContextBuilder;
@@ -162,7 +162,7 @@ pub struct TenantSharedResources {
     pub remote_storage: GenericRemoteStorage,
     pub deletion_queue_client: DeletionQueueClient,
     pub l0_flush_global_state: L0FlushGlobalState,
-    pub basebackup_cache: Arc<BasebackupCache>,
+    pub basebackup_prepare_sender: BasebackupPrepareSender,
     pub feature_resolver: FeatureResolver,
 }
 
@@ -331,7 +331,7 @@ pub struct TenantShard {
     deletion_queue_client: DeletionQueueClient,
 
     /// A channel to send async requests to prepare a basebackup for the basebackup cache.
-    basebackup_cache: Arc<BasebackupCache>,
+    basebackup_prepare_sender: BasebackupPrepareSender,
 
     /// Cached logical sizes updated updated on each [`TenantShard::gather_size_inputs`].
     cached_logical_sizes: tokio::sync::Mutex<HashMap<(TimelineId, Lsn), u64>>,
@@ -1363,7 +1363,7 @@ impl TenantShard {
             remote_storage,
             deletion_queue_client,
             l0_flush_global_state,
-            basebackup_cache,
+            basebackup_prepare_sender,
             feature_resolver,
         } = resources;
 
@@ -1380,7 +1380,7 @@ impl TenantShard {
             remote_storage.clone(),
             deletion_queue_client,
             l0_flush_global_state,
-            basebackup_cache,
+            basebackup_prepare_sender,
             feature_resolver,
         ));
 
@@ -3872,10 +3872,6 @@ impl TenantShard {
         &self.tenant_shard_id
     }
 
-    pub(crate) fn get_shard_identity(&self) -> ShardIdentity {
-        self.shard_identity
-    }
-
     pub(crate) fn get_shard_stripe_size(&self) -> ShardStripeSize {
         self.shard_identity.stripe_size
     }
@@ -4384,7 +4380,7 @@ impl TenantShard {
         remote_storage: GenericRemoteStorage,
         deletion_queue_client: DeletionQueueClient,
         l0_flush_global_state: L0FlushGlobalState,
-        basebackup_cache: Arc<BasebackupCache>,
+        basebackup_prepare_sender: BasebackupPrepareSender,
         feature_resolver: FeatureResolver,
     ) -> TenantShard {
         assert!(!attached_conf.location.generation.is_none());
@@ -4489,7 +4485,7 @@ impl TenantShard {
             ongoing_timeline_detach: std::sync::Mutex::default(),
             gc_block: Default::default(),
             l0_flush_global_state,
-            basebackup_cache,
+            basebackup_prepare_sender,
             feature_resolver,
         }
     }
@@ -4529,10 +4525,6 @@ impl TenantShard {
         Ok(toml_edit::de::from_str::<LocationConf>(&config)?)
     }
 
-    /// Stores a tenant location config to disk.
-    ///
-    /// NB: make sure to call `ShardIdentity::assert_equal` before persisting a new config, to avoid
-    /// changes to shard parameters that may result in data corruption.
     #[tracing::instrument(skip_all, fields(tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug()))]
     pub(super) async fn persist_tenant_config(
         conf: &'static PageServerConf,
@@ -5422,7 +5414,7 @@ impl TenantShard {
             pagestream_throttle_metrics: self.pagestream_throttle_metrics.clone(),
             l0_compaction_trigger: self.l0_compaction_trigger.clone(),
             l0_flush_global_state: self.l0_flush_global_state.clone(),
-            basebackup_cache: self.basebackup_cache.clone(),
+            basebackup_prepare_sender: self.basebackup_prepare_sender.clone(),
             feature_resolver: self.feature_resolver.clone(),
         }
     }
@@ -6008,7 +6000,7 @@ pub(crate) mod harness {
         ) -> anyhow::Result<Arc<TenantShard>> {
             let walredo_mgr = Arc::new(WalRedoManager::from(TestRedoManager));
 
-            let (basebackup_cache, _) = BasebackupCache::new(Utf8PathBuf::new(), None);
+            let (basebackup_requst_sender, _) = tokio::sync::mpsc::unbounded_channel();
 
             let tenant = Arc::new(TenantShard::new(
                 TenantState::Attaching,
@@ -6016,7 +6008,7 @@ pub(crate) mod harness {
                 AttachedTenantConf::try_from(LocationConf::attached_single(
                     self.tenant_conf.clone(),
                     self.generation,
-                    ShardParameters::default(),
+                    &ShardParameters::default(),
                 ))
                 .unwrap(),
                 self.shard_identity,
@@ -6026,7 +6018,7 @@ pub(crate) mod harness {
                 self.deletion_queue.new_client(),
                 // TODO: ideally we should run all unit tests with both configs
                 L0FlushGlobalState::new(L0FlushConfig::default()),
-                basebackup_cache,
+                basebackup_requst_sender,
                 FeatureResolver::new_disabled(),
             ));
 
@@ -11437,11 +11429,11 @@ mod tests {
         if left != right {
             eprintln!("---LEFT---");
             for left in left.iter() {
-                eprintln!("{left}");
+                eprintln!("{}", left);
             }
             eprintln!("---RIGHT---");
             for right in right.iter() {
-                eprintln!("{right}");
+                eprintln!("{}", right);
             }
             assert_eq!(left, right);
         }

pageserver/src/tenant/config.rs

@@ -12,7 +12,6 @@
 use pageserver_api::models;
 use pageserver_api::shard::{ShardCount, ShardIdentity, ShardNumber, ShardStripeSize};
 use serde::{Deserialize, Serialize};
-use utils::critical;
 use utils::generation::Generation;
 
 #[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -137,7 +136,7 @@ impl LocationConf {
     pub(crate) fn attached_single(
         tenant_conf: pageserver_api::models::TenantConfig,
         generation: Generation,
-        shard_params: models::ShardParameters,
+        shard_params: &models::ShardParameters,
     ) -> Self {
         Self {
             mode: LocationMode::Attached(AttachedLocationConfig {
@@ -172,16 +171,6 @@ impl LocationConf {
             }
         }
 
-        // This should never happen.
-        // TODO: turn this into a proper assertion.
-        if stripe_size != self.shard.stripe_size {
-            critical!(
-                "stripe size mismatch: {} != {}",
-                self.shard.stripe_size,
-                stripe_size,
-            );
-        }
-
         self.shard.stripe_size = stripe_size;
     }
 

pageserver/src/tenant/mgr.rs

@@ -880,9 +880,6 @@ impl TenantManager {
         // phase of writing config and/or waiting for flush, before returning.
         match fast_path_taken {
             Some(FastPathModified::Attached(tenant)) => {
-                tenant
-                    .shard_identity
-                    .assert_equal(new_location_config.shard);
                 TenantShard::persist_tenant_config(
                     self.conf,
                     &tenant_shard_id,
@@ -917,10 +914,7 @@ impl TenantManager {
 
                 return Ok(Some(tenant));
             }
-            Some(FastPathModified::Secondary(secondary_tenant)) => {
-                secondary_tenant
-                    .shard_identity
-                    .assert_equal(new_location_config.shard);
+            Some(FastPathModified::Secondary(_secondary_tenant)) => {
                 TenantShard::persist_tenant_config(
                     self.conf,
                     &tenant_shard_id,
@@ -954,10 +948,6 @@ impl TenantManager {
 
         match slot_guard.get_old_value() {
             Some(TenantSlot::Attached(tenant)) => {
-                tenant
-                    .shard_identity
-                    .assert_equal(new_location_config.shard);
-
                 // The case where we keep a Tenant alive was covered above in the special case
                 // for Attached->Attached transitions in the same generation.  By this point,
                 // if we see an attached tenant we know it will be discarded and should be
@@ -991,13 +981,9 @@ impl TenantManager {
                 // rather than assuming it to be empty.
                 spawn_mode = SpawnMode::Eager;
             }
-            Some(TenantSlot::Secondary(secondary_tenant)) => {
-                secondary_tenant
-                    .shard_identity
-                    .assert_equal(new_location_config.shard);
-
+            Some(TenantSlot::Secondary(state)) => {
                 info!("Shutting down secondary tenant");
-                secondary_tenant.shutdown().await;
+                state.shutdown().await;
             }
             Some(TenantSlot::InProgress(_)) => {
                 // This should never happen: acquire_slot should error out
@@ -2214,7 +2200,7 @@ impl TenantManager {
         selector: ShardSelector,
     ) -> ShardResolveResult {
         let tenants = self.tenants.read().unwrap();
-        let mut want_shard: Option<ShardIndex> = None;
+        let mut want_shard = None;
         let mut any_in_progress = None;
 
         match &*tenants {
@@ -2239,23 +2225,14 @@ impl TenantManager {
                             return ShardResolveResult::Found(tenant.clone());
                         }
                         ShardSelector::Page(key) => {
-                            // Each time we find an attached slot with a different shard count,
-                            // recompute the expected shard number: during shard splits we might
-                            // have multiple shards with the old shard count.
-                            if want_shard.is_none()
-                                || want_shard.unwrap().shard_count != tenant.shard_identity.count
-                            {
-                                want_shard = Some(ShardIndex {
-                                    shard_number: tenant.shard_identity.get_shard_number(&key),
-                                    shard_count: tenant.shard_identity.count,
-                                });
+                            // First slot we see for this tenant, calculate the expected shard number
+                            // for the key: we will use this for checking if this and subsequent
+                            // slots contain the key, rather than recalculating the hash each time.
+                            if want_shard.is_none() {
+                                want_shard = Some(tenant.shard_identity.get_shard_number(&key));
                             }
 
-                            if Some(ShardIndex {
-                                shard_number: tenant.shard_identity.number,
-                                shard_count: tenant.shard_identity.count,
-                            }) == want_shard
-                            {
+                            if Some(tenant.shard_identity.number) == want_shard {
                                 return ShardResolveResult::Found(tenant.clone());
                             }
                         }
@@ -2914,18 +2891,14 @@ mod tests {
     use std::collections::BTreeMap;
     use std::sync::Arc;
 
-    use camino::Utf8PathBuf;
     use storage_broker::BrokerClientChannel;
     use tracing::Instrument;
 
     use super::super::harness::TenantHarness;
     use super::TenantsMap;
-    use crate::{
-        basebackup_cache::BasebackupCache,
-        tenant::{
-            TenantSharedResources,
-            mgr::{BackgroundPurges, TenantManager, TenantSlot},
-        },
+    use crate::tenant::{
+        TenantSharedResources,
+        mgr::{BackgroundPurges, TenantManager, TenantSlot},
     };
 
     #[tokio::test(start_paused = true)]
@@ -2951,7 +2924,9 @@ mod tests {
         // Invoke remove_tenant_from_memory with a cleanup hook that blocks until we manually
         // permit it to proceed: that will stick the tenant in InProgress
 
-        let (basebackup_cache, _) = BasebackupCache::new(Utf8PathBuf::new(), None);
+        let (basebackup_prepare_sender, _) = tokio::sync::mpsc::unbounded_channel::<
+            crate::basebackup_cache::BasebackupPrepareRequest,
+        >();
 
         let tenant_manager = TenantManager {
             tenants: std::sync::RwLock::new(TenantsMap::Open(tenants)),
@@ -2965,7 +2940,7 @@ mod tests {
                 l0_flush_global_state: crate::l0_flush::L0FlushGlobalState::new(
                     h.conf.l0_flush.clone(),
                 ),
-                basebackup_cache,
+                basebackup_prepare_sender,
                 feature_resolver: crate::feature_resolver::FeatureResolver::new_disabled(),
             },
             cancel: tokio_util::sync::CancellationToken::new(),

pageserver/src/tenant/secondary.rs

@@ -101,7 +101,7 @@ pub(crate) struct SecondaryTenant {
     // Secondary mode does not need the full shard identity or the pageserver_api::models::TenantConfig.  However,
     // storing these enables us to report our full LocationConf, enabling convenient reconciliation
     // by the control plane (see [`Self::get_location_conf`])
-    pub(crate) shard_identity: ShardIdentity,
+    shard_identity: ShardIdentity,
     tenant_conf: std::sync::Mutex<pageserver_api::models::TenantConfig>,
 
     // Internal state used by the Downloader.

pageserver/src/tenant/storage_layer/batch_split_writer.rs

@@ -55,11 +55,11 @@ pub struct BatchLayerWriter {
 }
 
 impl BatchLayerWriter {
-    pub fn new(conf: &'static PageServerConf) -> Self {
-        Self {
+    pub async fn new(conf: &'static PageServerConf) -> anyhow::Result<Self> {
+        Ok(Self {
             generated_layer_writers: Vec::new(),
             conf,
-        }
+        })
     }
 
     pub fn add_unfinished_image_writer(
@@ -209,7 +209,6 @@ impl<'a> SplitImageLayerWriter<'a> {
     ) -> anyhow::Result<Self> {
         Ok(Self {
             target_layer_size,
-            // XXX make this lazy like in SplitDeltaLayerWriter?
             inner: ImageLayerWriter::new(
                 conf,
                 timeline_id,
@@ -224,7 +223,7 @@ impl<'a> SplitImageLayerWriter<'a> {
             conf,
             timeline_id,
             tenant_shard_id,
-            batches: BatchLayerWriter::new(conf),
+            batches: BatchLayerWriter::new(conf).await?,
             lsn,
             start_key,
             gate,
@@ -320,7 +319,7 @@ pub struct SplitDeltaLayerWriter<'a> {
 }
 
 impl<'a> SplitDeltaLayerWriter<'a> {
-    pub fn new(
+    pub async fn new(
         conf: &'static PageServerConf,
         timeline_id: TimelineId,
         tenant_shard_id: TenantShardId,
@@ -328,8 +327,8 @@ impl<'a> SplitDeltaLayerWriter<'a> {
         target_layer_size: u64,
         gate: &'a utils::sync::gate::Gate,
         cancel: CancellationToken,
-    ) -> Self {
-        Self {
+    ) -> anyhow::Result<Self> {
+        Ok(Self {
             target_layer_size,
             inner: None,
             conf,
@@ -337,10 +336,10 @@ impl<'a> SplitDeltaLayerWriter<'a> {
             tenant_shard_id,
             lsn_range,
             last_key_written: Key::MIN,
-            batches: BatchLayerWriter::new(conf),
+            batches: BatchLayerWriter::new(conf).await?,
             gate,
             cancel,
-        }
+        })
     }
 
     pub async fn put_value(
@@ -511,7 +510,9 @@ mod tests {
             4 * 1024 * 1024,
             &tline.gate,
             tline.cancel.clone(),
-        );
+        )
+        .await
+        .unwrap();
 
         image_writer
             .put_image(get_key(0), get_img(0), &ctx)
@@ -589,7 +590,9 @@ mod tests {
             4 * 1024 * 1024,
             &tline.gate,
             tline.cancel.clone(),
-        );
+        )
+        .await
+        .unwrap();
         const N: usize = 2000;
         for i in 0..N {
             let i = i as u32;
@@ -689,7 +692,9 @@ mod tests {
             4 * 1024,
             &tline.gate,
             tline.cancel.clone(),
-        );
+        )
+        .await
+        .unwrap();
 
         image_writer
             .put_image(get_key(0), get_img(0), &ctx)
@@ -765,7 +770,9 @@ mod tests {
             4 * 1024 * 1024,
             &tline.gate,
             tline.cancel.clone(),
-        );
+        )
+        .await
+        .unwrap();
 
         for i in 0..N {
             let i = i as u32;

pageserver/src/tenant/tasks.rs

@@ -17,17 +17,14 @@ use tracing::*;
 use utils::backoff::exponential_backoff_duration;
 use utils::completion::Barrier;
 use utils::pausable_failpoint;
-use utils::sync::gate::GateError;
 
 use crate::context::{DownloadBehavior, RequestContext};
 use crate::metrics::{self, BackgroundLoopSemaphoreMetricsRecorder, TENANT_TASK_EVENTS};
 use crate::task_mgr::{self, BACKGROUND_RUNTIME, TOKIO_WORKER_THREADS, TaskKind};
-use crate::tenant::blob_io::WriteBlobError;
 use crate::tenant::throttle::Stats;
 use crate::tenant::timeline::CompactionError;
 use crate::tenant::timeline::compaction::CompactionOutcome;
 use crate::tenant::{TenantShard, TenantState};
-use crate::virtual_file::owned_buffers_io::write::FlushTaskError;
 
 /// Semaphore limiting concurrent background tasks (across all tenants).
 ///
@@ -316,20 +313,7 @@ pub(crate) fn log_compaction_error(
             let timeline = root_cause
                 .downcast_ref::<PageReconstructError>()
                 .is_some_and(|e| e.is_stopping());
-            let buffered_writer_flush_task_canelled = root_cause
-                .downcast_ref::<FlushTaskError>()
-                .is_some_and(|e| e.is_cancel());
-            let write_blob_cancelled = root_cause
-                .downcast_ref::<WriteBlobError>()
-                .is_some_and(|e| e.is_cancel());
-            let gate_closed = root_cause
-                .downcast_ref::<GateError>()
-                .is_some_and(|e| e.is_cancel());
-            let is_stopping = upload_queue
-                || timeline
-                || buffered_writer_flush_task_canelled
-                || write_blob_cancelled
-                || gate_closed;
+            let is_stopping = upload_queue || timeline;
 
             if is_stopping {
                 Level::INFO

pageserver/src/tenant/timeline.rs

@@ -95,12 +95,12 @@ use super::storage_layer::{LayerFringe, LayerVisibilityHint, ReadableLayer};
 use super::tasks::log_compaction_error;
 use super::upload_queue::NotInitialized;
 use super::{
-    AttachedTenantConf, GcError, HeatMapTimeline, MaybeOffloaded,
+    AttachedTenantConf, BasebackupPrepareSender, GcError, HeatMapTimeline, MaybeOffloaded,
     debug_assert_current_span_has_tenant_and_timeline_id,
 };
 use crate::PERF_TRACE_TARGET;
 use crate::aux_file::AuxFileSizeEstimator;
-use crate::basebackup_cache::BasebackupCache;
+use crate::basebackup_cache::BasebackupPrepareRequest;
 use crate::config::PageServerConf;
 use crate::context::{
     DownloadBehavior, PerfInstrumentFutureExt, RequestContext, RequestContextBuilder,
@@ -201,7 +201,7 @@ pub struct TimelineResources {
     pub pagestream_throttle_metrics: Arc<crate::metrics::tenant_throttling::Pagestream>,
     pub l0_compaction_trigger: Arc<Notify>,
     pub l0_flush_global_state: l0_flush::L0FlushGlobalState,
-    pub basebackup_cache: Arc<BasebackupCache>,
+    pub basebackup_prepare_sender: BasebackupPrepareSender,
     pub feature_resolver: FeatureResolver,
 }
 
@@ -448,7 +448,7 @@ pub struct Timeline {
     wait_lsn_log_slow: tokio::sync::Semaphore,
 
     /// A channel to send async requests to prepare a basebackup for the basebackup cache.
-    basebackup_cache: Arc<BasebackupCache>,
+    basebackup_prepare_sender: BasebackupPrepareSender,
 
     feature_resolver: FeatureResolver,
 }
@@ -763,7 +763,7 @@ pub(crate) enum CreateImageLayersError {
     PageReconstructError(#[source] PageReconstructError),
 
     #[error(transparent)]
-    Other(anyhow::Error),
+    Other(#[from] anyhow::Error),
 }
 
 impl From<layer_manager::Shutdown> for CreateImageLayersError {
@@ -2500,37 +2500,6 @@ impl Timeline {
             .unwrap_or(self.conf.default_tenant_conf.basebackup_cache_enabled)
     }
 
-    /// Try to get a basebackup from the on-disk cache.
-    pub(crate) async fn get_cached_basebackup(&self, lsn: Lsn) -> Option<tokio::fs::File> {
-        self.basebackup_cache
-            .get(self.tenant_shard_id.tenant_id, self.timeline_id, lsn)
-            .await
-    }
-
-    /// Convenience method to attempt fetching a basebackup for the timeline if enabled and safe for
-    /// the given request parameters.
-    ///
-    /// TODO: consider moving this onto GrpcPageServiceHandler once the libpq handler is gone.
-    pub async fn get_cached_basebackup_if_enabled(
-        &self,
-        lsn: Option<Lsn>,
-        prev_lsn: Option<Lsn>,
-        full: bool,
-        replica: bool,
-        gzip: bool,
-    ) -> Option<tokio::fs::File> {
-        if !self.is_basebackup_cache_enabled() || !self.basebackup_cache.is_enabled() {
-            return None;
-        }
-        // We have to know which LSN to fetch the basebackup for.
-        let lsn = lsn?;
-        // We only cache gzipped, non-full basebackups for primary computes with automatic prev_lsn.
-        if prev_lsn.is_some() || full || replica || !gzip {
-            return None;
-        }
-        self.get_cached_basebackup(lsn).await
-    }
-
     /// Prepare basebackup for the given LSN and store it in the basebackup cache.
     /// The method is asynchronous and returns immediately.
     /// The actual basebackup preparation is performed in the background
@@ -2552,8 +2521,17 @@ impl Timeline {
             return;
         }
 
-        self.basebackup_cache
-            .send_prepare(self.tenant_shard_id, self.timeline_id, lsn);
+        let res = self
+            .basebackup_prepare_sender
+            .send(BasebackupPrepareRequest {
+                tenant_shard_id: self.tenant_shard_id,
+                timeline_id: self.timeline_id,
+                lsn,
+            });
+        if let Err(e) = res {
+            // May happen during shutdown, it's not critical.
+            info!("Failed to send shutdown checkpoint: {e:#}");
+        }
     }
 }
 
@@ -3110,7 +3088,7 @@ impl Timeline {
 
                 wait_lsn_log_slow: tokio::sync::Semaphore::new(1),
 
-                basebackup_cache: resources.basebackup_cache,
+                basebackup_prepare_sender: resources.basebackup_prepare_sender,
 
                 feature_resolver: resources.feature_resolver,
             };
@@ -4680,16 +4658,6 @@ impl Timeline {
         mut layer_flush_start_rx: tokio::sync::watch::Receiver<(u64, Lsn)>,
         ctx: &RequestContext,
     ) {
-        // Always notify waiters about the flush loop exiting since the loop might stop
-        // when the timeline hasn't been cancelled.
-        let scopeguard_rx = layer_flush_start_rx.clone();
-        scopeguard::defer! {
-            let (flush_counter, _) = *scopeguard_rx.borrow();
-            let _ = self
-                .layer_flush_done_tx
-                .send_replace((flush_counter, Err(FlushLayerError::Cancelled)));
-        }
-
         // Subscribe to L0 delta layer updates, for compaction backpressure.
         let mut watch_l0 = match self
             .layers
@@ -4719,6 +4687,9 @@ impl Timeline {
             let result = loop {
                 if self.cancel.is_cancelled() {
                     info!("dropping out of flush loop for timeline shutdown");
+                    // Note: we do not bother transmitting into [`layer_flush_done_tx`], because
+                    // anyone waiting on that will respect self.cancel as well: they will stop
+                    // waiting at the same time we as drop out of this loop.
                     return;
                 }
 
@@ -5590,7 +5561,7 @@ impl Timeline {
                 self.should_check_if_image_layers_required(lsn)
             };
 
-        let mut batch_image_writer = BatchLayerWriter::new(self.conf);
+        let mut batch_image_writer = BatchLayerWriter::new(self.conf).await?;
 
         let mut all_generated = true;
 
@@ -5694,8 +5665,7 @@ impl Timeline {
                 self.cancel.clone(),
                 ctx,
             )
-            .await
-            .map_err(CreateImageLayersError::Other)?;
+            .await?;
 
             fail_point!("image-layer-writer-fail-before-finish", |_| {
                 Err(CreateImageLayersError::Other(anyhow::anyhow!(
@@ -5790,10 +5760,7 @@ impl Timeline {
             }
         }
 
-        let image_layers = batch_image_writer
-            .finish(self, ctx)
-            .await
-            .map_err(CreateImageLayersError::Other)?;
+        let image_layers = batch_image_writer.finish(self, ctx).await?;
 
         let mut guard = self.layers.write(LayerManagerLockHolder::Compaction).await;
 

pageserver/src/tenant/timeline/compaction.rs

@@ -3531,7 +3531,10 @@ impl Timeline {
             self.get_compaction_target_size(),
             &self.gate,
             self.cancel.clone(),
-        );
+        )
+        .await
+        .context("failed to create delta layer writer")
+        .map_err(CompactionError::Other)?;
 
         #[derive(Default)]
         struct RewritingLayers {
@@ -4327,8 +4330,7 @@ impl TimelineAdaptor {
             self.timeline.cancel.clone(),
             ctx,
         )
-        .await
-        .map_err(CreateImageLayersError::Other)?;
+        .await?;
 
         fail_point!("image-layer-writer-fail-before-finish", |_| {
             Err(CreateImageLayersError::Other(anyhow::anyhow!(
@@ -4337,10 +4339,7 @@ impl TimelineAdaptor {
         });
 
         let keyspace = KeySpace {
-            ranges: self
-                .get_keyspace(key_range, lsn, ctx)
-                .await
-                .map_err(CreateImageLayersError::Other)?,
+            ranges: self.get_keyspace(key_range, lsn, ctx).await?,
         };
         // TODO set proper (stateful) start. The create_image_layer_for_rel_blocks function mostly
         let outcome = self
@@ -4359,13 +4358,9 @@ impl TimelineAdaptor {
             unfinished_image_layer,
         } = outcome
         {
-            let (desc, path) = unfinished_image_layer
-                .finish(ctx)
-                .await
-                .map_err(CreateImageLayersError::Other)?;
+            let (desc, path) = unfinished_image_layer.finish(ctx).await?;
             let image_layer =
-                Layer::finish_creating(self.timeline.conf, &self.timeline, desc, &path)
-                    .map_err(CreateImageLayersError::Other)?;
+                Layer::finish_creating(self.timeline.conf, &self.timeline, desc, &path)?;
             self.new_images.push(image_layer);
         }
 

pageserver/src/tenant/timeline/delete.rs

@@ -241,17 +241,8 @@ impl DeleteTimelineFlow {
                 {
                     Ok(r) => r,
                     Err(DownloadError::NotFound) => {
-                        // Deletion is already complete.
-                        // As we came here, we will need to remove the timeline from the tenant though.
+                        // Deletion is already complete
                         tracing::info!("Timeline already deleted in remote storage");
-                        if let TimelineOrOffloaded::Offloaded(_) = &timeline {
-                            // We only supoprt this for offloaded timelines, as we don't know which state non-offloaded timelines are in.
-                            tracing::info!(
-                                "Timeline with gone index part is offloaded timeline. Removing from tenant."
-                            );
-                            remove_maybe_offloaded_timeline_from_tenant(tenant, &timeline, &guard)
-                                .await?;
-                        }
                         return Ok(());
                     }
                     Err(e) => {

pageserver/src/tenant/timeline/detach_ancestor.rs

@@ -885,7 +885,7 @@ async fn remote_copy(
                 }
                 tracing::info!("Deleting orphan layer file to make way for hard linking");
                 // Delete orphan layer file and try again, to ensure this layer has a well understood source
-                std::fs::remove_file(&adoptee_path)
+                std::fs::remove_file(adopted_path)
                     .map_err(|e| Error::launder(e.into(), Error::Prepare))?;
                 std::fs::hard_link(adopted_path, &adoptee_path)
                     .map_err(|e| Error::launder(e.into(), Error::Prepare))?;

pageserver/src/tenant/timeline/handle.rs

@@ -887,7 +887,7 @@ mod tests {
             .expect("we still have it");
     }
 
-    fn make_relation_key_for_shard(shard: ShardNumber, params: ShardParameters) -> Key {
+    fn make_relation_key_for_shard(shard: ShardNumber, params: &ShardParameters) -> Key {
         rel_block_to_key(
             RelTag {
                 spcnode: 1663,
@@ -917,14 +917,14 @@ mod tests {
         let child0 = Arc::new_cyclic(|myself| StubTimeline {
             gate: Default::default(),
             id: timeline_id,
-            shard: ShardIdentity::from_params(ShardNumber(0), child_params),
+            shard: ShardIdentity::from_params(ShardNumber(0), &child_params),
             per_timeline_state: PerTimelineState::default(),
             myself: myself.clone(),
         });
         let child1 = Arc::new_cyclic(|myself| StubTimeline {
             gate: Default::default(),
             id: timeline_id,
-            shard: ShardIdentity::from_params(ShardNumber(1), child_params),
+            shard: ShardIdentity::from_params(ShardNumber(1), &child_params),
             per_timeline_state: PerTimelineState::default(),
             myself: myself.clone(),
         });
@@ -937,7 +937,7 @@ mod tests {
             let handle = cache
                 .get(
                     timeline_id,
-                    ShardSelector::Page(make_relation_key_for_shard(ShardNumber(i), child_params)),
+                    ShardSelector::Page(make_relation_key_for_shard(ShardNumber(i), &child_params)),
                     &StubManager {
                         shards: vec![parent.clone()],
                     },
@@ -961,7 +961,7 @@ mod tests {
             let handle = cache
                 .get(
                     timeline_id,
-                    ShardSelector::Page(make_relation_key_for_shard(ShardNumber(i), child_params)),
+                    ShardSelector::Page(make_relation_key_for_shard(ShardNumber(i), &child_params)),
                     &StubManager {
                         shards: vec![], // doesn't matter what's in here, the cache is fully loaded
                     },
@@ -978,7 +978,7 @@ mod tests {
         let parent_handle = cache
             .get(
                 timeline_id,
-                ShardSelector::Page(make_relation_key_for_shard(ShardNumber(0), child_params)),
+                ShardSelector::Page(make_relation_key_for_shard(ShardNumber(0), &child_params)),
                 &StubManager {
                     shards: vec![parent.clone()],
                 },
@@ -995,7 +995,7 @@ mod tests {
             let handle = cache
                 .get(
                     timeline_id,
-                    ShardSelector::Page(make_relation_key_for_shard(ShardNumber(i), child_params)),
+                    ShardSelector::Page(make_relation_key_for_shard(ShardNumber(i), &child_params)),
                     &StubManager {
                         shards: vec![child0.clone(), child1.clone()], // <====== this changed compared to previous loop
                     },

pgxn/neon/file_cache.c

@@ -1295,8 +1295,7 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 
 		if (iteration_hits != 0)
 		{
-			/* chunk offset (#
-			   of pages) into the LFC file */
+			/* chunk offset (# of pages) into the LFC file */
 			off_t	first_read_offset = (off_t) entry_offset * lfc_blocks_per_chunk;
 			int		nwrite = iov_last_used - first_block_in_chunk_read;
 			/* offset of first IOV */
@@ -1314,6 +1313,16 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 				lfc_disable("read");
 				return -1;
 			}
+
+			/*
+			 * We successfully read the pages we know were valid when we
+			 * started reading; now mark those pages as read
+			 */
+			for (int i = first_block_in_chunk_read; i < iov_last_used; i++)
+			{
+				if (BITMAP_ISSET(chunk_mask, i))
+					BITMAP_SET(mask, buf_offset + i);
+			}
 		}
 
 		/* Place entry to the head of LRU list */
@@ -1331,15 +1340,6 @@ lfc_readv_select(NRelFileInfo rinfo, ForkNumber forkNum, BlockNumber blkno,
 			{
 				lfc_ctl->time_read += io_time_us;
 				inc_page_cache_read_wait(io_time_us);
-				/*
-				 * We successfully read the pages we know were valid when we
-				 * started reading; now mark those pages as read
-				 */
-				for (int i = first_block_in_chunk_read; i < iov_last_used; i++)
-				{
-					if (BITMAP_ISSET(chunk_mask, i))
-						BITMAP_SET(mask, buf_offset + i);
-				}
 			}
 
 			CriticalAssert(entry->access_count > 0);

postgres.mk

@@ -1,121 +0,0 @@
-# Sub-makefile for compiling PostgreSQL as part of Neon. This is
-# included from the main Makefile, and is not meant to be called
-# directly.
-#
-# CI workflows and Dockerfiles can take advantage of the following
-# properties for caching:
-#
-# - Compiling the targets in this file only builds the PostgreSQL sources
-#   under the vendor/ subdirectory, nothing else from the repository.
-# - All outputs go to POSTGRES_INSTALL_DIR (by default 'pg_install',
-#   see parent Makefile)
-# - intermediate build artifacts go to BUILD_DIR
-#
-#
-# Variables passed from the parent Makefile that control what gets
-# installed and where:
-# - POSTGRES_VERSIONS
-# - POSTGRES_INSTALL_DIR
-# - BUILD_DIR
-#
-# Variables passed from the parent Makefile that affect the build
-# process and the resulting binaries:
-# - PG_CONFIGURE_OPTS
-# - PG_CFLAGS
-# - PG_LDFLAGS
-# - EXTRA_PATH_OVERRIDES
-
-###
-### Main targets
-###
-### These are called from the main Makefile, and can also be called
-### directly from command line
-
-# Compile and install a specific PostgreSQL version
-postgres-install-%: postgres-configure-% \
-		  postgres-headers-install-% # to prevent `make install` conflicts with neon's `postgres-headers`
-
-# Install the PostgreSQL header files into $(POSTGRES_INSTALL_DIR)/<version>/include
-#
-# This is implicitly part of the 'postgres-install-%' target, but this can be handy
-# if you want to install just the headers without building PostgreSQL, e.g. for building
-# extensions.
-postgres-headers-install-%: postgres-configure-%
-	+@echo "Installing PostgreSQL $* headers"
-	$(MAKE) -C $(BUILD_DIR)/$*/src/include MAKELEVEL=0 install
-
-# Run Postgres regression tests
-postgres-check-%: postgres-install-%
-	$(MAKE) -C $(BUILD_DIR)/$* MAKELEVEL=0 check
-
-###
-### Shorthands for the main targets, for convenience
-###
-
-# Same as the above main targets, but for all supported PostgreSQL versions
-# For example, 'make postgres-install' is equivalent to
-# 'make postgres-install-v14 postgres-install-v15 postgres-install-v16 postgres-install-v17'
-all_version_targets=postgres-install postgres-headers-install postgres-check
-.PHONY: $(all_version_targets)
-$(all_version_targets): postgres-%: $(foreach pg_version,$(POSTGRES_VERSIONS),postgres-%-$(pg_version))
-
-.PHONY: postgres
-postgres: postgres-install
-
-.PHONY: postgres-headers
-postgres-headers: postgres-headers-install
-
-# 'postgres-v17' is an alias for 'postgres-install-v17' etc.
-$(foreach pg_version,$(POSTGRES_VERSIONS),postgres-$(pg_version)): postgres-%: postgres-install-%
-
-###
-### Intermediate targets
-###
-### These are not intended to be called directly, but are dependencies for the
-### main targets.
-
-# Run 'configure'
-$(BUILD_DIR)/%/config.status:
-	mkdir -p $(BUILD_DIR)
-	test -e $(BUILD_DIR)/CACHEDIR.TAG || echo "$(CACHEDIR_TAG_CONTENTS)" > $(BUILD_DIR)/CACHEDIR.TAG
-
-	+@echo "Configuring Postgres $* build"
-	@test -s $(ROOT_PROJECT_DIR)/vendor/postgres-$*/configure || { \
-		echo "\nPostgres submodule not found in $(ROOT_PROJECT_DIR)/vendor/postgres-$*/, execute "; \
-		echo "'git submodule update --init --recursive --depth 2 --progress .' in project root.\n"; \
-		exit 1; }
-	mkdir -p $(BUILD_DIR)/$*
-
-	VERSION=$*; \
-	EXTRA_VERSION=$$(cd $(ROOT_PROJECT_DIR)/vendor/postgres-$$VERSION && git rev-parse HEAD); \
-	(cd $(BUILD_DIR)/$$VERSION && \
-	env PATH="$(EXTRA_PATH_OVERRIDES):$$PATH" $(ROOT_PROJECT_DIR)/vendor/postgres-$$VERSION/configure \
-		CFLAGS='$(PG_CFLAGS)' LDFLAGS='$(PG_LDFLAGS)' \
-		$(PG_CONFIGURE_OPTS) --with-extra-version=" ($$EXTRA_VERSION)" \
-		--prefix=$(abspath $(POSTGRES_INSTALL_DIR))/$$VERSION > configure.log)
-
-# nicer alias to run 'configure'.
-#
-# This tries to accomplish this rule:
-#
-# postgres-configure-%: $(BUILD_DIR)/%/config.status
-#
-# XXX: I'm not sure why the above rule doesn't work directly. But this accomplishses
-# the same thing
-$(foreach pg_version,$(POSTGRES_VERSIONS),postgres-configure-$(pg_version)): postgres-configure-%: FORCE $(BUILD_DIR)/%/config.status
-
-# Compile and install PostgreSQL (and a few contrib modules used in tests)
-postgres-install-%: postgres-configure-% \
-		  postgres-headers-install-% # to prevent `make install` conflicts with neon's `postgres-headers-install`
-	+@echo "Compiling PostgreSQL $*"
-	$(MAKE) -C $(BUILD_DIR)/$* MAKELEVEL=0 install
-	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pg_prewarm install
-	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pg_buffercache install
-	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pg_visibility install
-	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pageinspect install
-	$(MAKE) -C $(BUILD_DIR)/$*/contrib/pg_trgm install
-	$(MAKE) -C $(BUILD_DIR)/$*/contrib/amcheck install
-	$(MAKE) -C $(BUILD_DIR)/$*/contrib/test_decoding install
-
-.PHONY: FORCE
-FORCE:

proxy/src/binary/local_proxy.rs

@@ -279,6 +279,7 @@ fn build_config(args: &LocalProxyCliArgs) -> anyhow::Result<&'static ProxyConfig
         },
         proxy_protocol_v2: config::ProxyProtocolV2::Rejected,
         handshake_timeout: Duration::from_secs(10),
+        region: "local".into(),
         wake_compute_retry_config: RetryConfig::parse(RetryConfig::WAKE_COMPUTE_DEFAULT_VALUES)?,
         connect_compute_locks,
         connect_to_compute: compute_config,

proxy/src/binary/pg_sni_router.rs

@@ -26,10 +26,9 @@ use utils::sentry_init::init_sentry;
 
 use crate::context::RequestContext;
 use crate::metrics::{Metrics, ThreadPoolMetrics};
-use crate::pglb::TlsRequired;
 use crate::pqproto::FeStartupPacket;
 use crate::protocol2::ConnectionInfo;
-use crate::proxy::{ErrorSource, copy_bidirectional_client_compute};
+use crate::proxy::{ErrorSource, TlsRequired, copy_bidirectional_client_compute};
 use crate::stream::{PqStream, Stream};
 use crate::util::run_until_cancelled;
 
@@ -237,6 +236,7 @@ pub(super) async fn task_main(
                         extra: None,
                     },
                     crate::metrics::Protocol::SniRouter,
+                    "sni",
                 );
                 handle_client(ctx, dest_suffix, tls_config, compute_tls_config, socket).await
             }

proxy/src/binary/proxy.rs

@@ -123,6 +123,12 @@ struct ProxyCliArgs {
     /// timeout for the TLS handshake
     #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
     handshake_timeout: tokio::time::Duration,
+    /// http endpoint to receive periodic metric updates
+    #[clap(long)]
+    metric_collection_endpoint: Option<String>,
+    /// how often metrics should be sent to a collection endpoint
+    #[clap(long)]
+    metric_collection_interval: Option<String>,
     /// cache for `wake_compute` api method (use `size=0` to disable)
     #[clap(long, default_value = config::CacheOptions::CACHE_DEFAULT_OPTIONS)]
     wake_compute_cache: String,
@@ -149,31 +155,40 @@ struct ProxyCliArgs {
     /// Wake compute rate limiter max number of requests per second.
     #[clap(long, default_values_t = RateBucketInfo::DEFAULT_SET)]
     wake_compute_limit: Vec<RateBucketInfo>,
+    /// Redis rate limiter max number of requests per second.
+    #[clap(long, default_values_t = RateBucketInfo::DEFAULT_REDIS_SET)]
+    redis_rps_limit: Vec<RateBucketInfo>,
     /// Cancellation channel size (max queue size for redis kv client)
     #[clap(long, default_value_t = 1024)]
     cancellation_ch_size: usize,
     /// Cancellation ops batch size for redis
     #[clap(long, default_value_t = 8)]
     cancellation_batch_size: usize,
-    /// redis url for plain authentication
-    #[clap(long, alias("redis-notifications"))]
-    redis_plain: Option<String>,
-    /// what from the available authentications type to use for redis. Supported are "irsa" and "plain".
+    /// cache for `allowed_ips` (use `size=0` to disable)
+    #[clap(long, default_value = config::CacheOptions::CACHE_DEFAULT_OPTIONS)]
+    allowed_ips_cache: String,
+    /// cache for `role_secret` (use `size=0` to disable)
+    #[clap(long, default_value = config::CacheOptions::CACHE_DEFAULT_OPTIONS)]
+    role_secret_cache: String,
+    /// redis url for notifications (if empty, redis_host:port will be used for both notifications and streaming connections)
+    #[clap(long)]
+    redis_notifications: Option<String>,
+    /// what from the available authentications type to use for the regional redis we have. Supported are "irsa" and "plain".
     #[clap(long, default_value = "irsa")]
     redis_auth_type: String,
-    /// redis host for irsa authentication
+    /// redis host for streaming connections (might be different from the notifications host)
     #[clap(long)]
     redis_host: Option<String>,
-    /// redis port for irsa authentication
+    /// redis port for streaming connections (might be different from the notifications host)
     #[clap(long)]
     redis_port: Option<u16>,
-    /// redis cluster name for irsa authentication
+    /// redis cluster name, used in aws elasticache
     #[clap(long)]
     redis_cluster_name: Option<String>,
-    /// redis user_id for irsa authentication
+    /// redis user_id, used in aws elasticache
     #[clap(long)]
     redis_user_id: Option<String>,
-    /// aws region for irsa authentication
+    /// aws region to retrieve credentials
     #[clap(long, default_value_t = String::new())]
     aws_region: String,
     /// cache for `project_info` (use `size=0` to disable)
@@ -185,12 +200,6 @@ struct ProxyCliArgs {
     #[clap(flatten)]
     parquet_upload: ParquetUploadArgs,
 
-    /// http endpoint to receive periodic metric updates
-    #[clap(long)]
-    metric_collection_endpoint: Option<String>,
-    /// how often metrics should be sent to a collection endpoint
-    #[clap(long)]
-    metric_collection_interval: Option<String>,
     /// interval for backup metric collection
     #[clap(long, default_value = "10m", value_parser = humantime::parse_duration)]
     metric_backup_collection_interval: std::time::Duration,
@@ -203,7 +212,6 @@ struct ProxyCliArgs {
     /// Size of each event is no more than 400 bytes, so 2**22 is about 200MB before the compression.
     #[clap(long, default_value = "4194304")]
     metric_backup_collection_chunk_size: usize,
-
     /// Whether to retry the connection to the compute node
     #[clap(long, default_value = config::RetryConfig::CONNECT_TO_COMPUTE_DEFAULT_VALUES)]
     connect_to_compute_retry: String,
@@ -323,7 +331,7 @@ pub async fn run() -> anyhow::Result<()> {
         Either::Right(auth_backend) => info!("Authentication backend: {auth_backend:?}"),
     }
     info!("Using region: {}", args.aws_region);
-    let redis_client = configure_redis(&args).await?;
+    let (regional_redis_client, redis_notifications_client) = configure_redis(&args).await?;
 
     // Check that we can bind to address before further initialization
     info!("Starting http on {}", args.http);
@@ -378,6 +386,13 @@ pub async fn run() -> anyhow::Result<()> {
 
     let cancellation_token = CancellationToken::new();
 
+    let redis_rps_limit = Vec::leak(args.redis_rps_limit.clone());
+    RateBucketInfo::validate(redis_rps_limit)?;
+
+    let redis_kv_client = regional_redis_client
+        .as_ref()
+        .map(|redis_publisher| RedisKVClient::new(redis_publisher.clone(), redis_rps_limit));
+
     let cancellation_handler = Arc::new(CancellationHandler::new(&config.connect_to_compute));
 
     let endpoint_rate_limiter = Arc::new(EndpointRateLimiter::new_with_shards(
@@ -392,7 +407,7 @@ pub async fn run() -> anyhow::Result<()> {
     match auth_backend {
         Either::Left(auth_backend) => {
             if let Some(proxy_listener) = proxy_listener {
-                client_tasks.spawn(crate::pglb::task_main(
+                client_tasks.spawn(crate::proxy::task_main(
                     config,
                     auth_backend,
                     proxy_listener,
@@ -457,7 +472,6 @@ pub async fn run() -> anyhow::Result<()> {
     client_tasks.spawn(crate::context::parquet::worker(
         cancellation_token.clone(),
         args.parquet_upload,
-        args.region,
     ));
 
     // maintenance tasks. these never return unless there's an error
@@ -481,17 +495,32 @@ pub async fn run() -> anyhow::Result<()> {
     #[cfg_attr(not(any(test, feature = "testing")), expect(irrefutable_let_patterns))]
     if let Either::Left(auth::Backend::ControlPlane(api, ())) = &auth_backend {
         if let crate::control_plane::client::ControlPlaneClient::ProxyV1(api) = &**api {
-            if let Some(client) = redis_client {
-                // project info cache and invalidation of that cache.
-                let cache = api.caches.project_info.clone();
-                maintenance_tasks.spawn(notifications::task_main(client.clone(), cache.clone()));
-                maintenance_tasks.spawn(async move { cache.clone().gc_worker().await });
+            match (redis_notifications_client, regional_redis_client.clone()) {
+                (None, None) => {}
+                (client1, client2) => {
+                    let cache = api.caches.project_info.clone();
+                    if let Some(client) = client1 {
+                        maintenance_tasks.spawn(notifications::task_main(
+                            client,
+                            cache.clone(),
+                            args.region.clone(),
+                        ));
+                    }
+                    if let Some(client) = client2 {
+                        maintenance_tasks.spawn(notifications::task_main(
+                            client,
+                            cache.clone(),
+                            args.region.clone(),
+                        ));
+                    }
+                    maintenance_tasks.spawn(async move { cache.clone().gc_worker().await });
+                }
+            }
 
-                // Try to connect to Redis 3 times with 1 + (0..0.1) second interval.
-                // This prevents immediate exit and pod restart,
-                // which can cause hammering of the redis in case of connection issues.
-                // cancellation key management
-                let mut redis_kv_client = RedisKVClient::new(client.clone());
+            // Try to connect to Redis 3 times with 1 + (0..0.1) second interval.
+            // This prevents immediate exit and pod restart,
+            // which can cause hammering of the redis in case of connection issues.
+            if let Some(mut redis_kv_client) = redis_kv_client {
                 for attempt in (0..3).with_position() {
                     match redis_kv_client.try_connect().await {
                         Ok(()) => {
@@ -516,12 +545,14 @@ pub async fn run() -> anyhow::Result<()> {
                         }
                     }
                 }
+            }
 
-                // listen for notifications of new projects/endpoints/branches
+            if let Some(regional_redis_client) = regional_redis_client {
                 let cache = api.caches.endpoints_cache.clone();
+                let con = regional_redis_client;
                 let span = tracing::info_span!("endpoints_cache");
                 maintenance_tasks.spawn(
-                    async move { cache.do_read(client, cancellation_token.clone()).await }
+                    async move { cache.do_read(con, cancellation_token.clone()).await }
                         .instrument(span),
                 );
             }
@@ -650,6 +681,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
         authentication_config,
         proxy_protocol_v2: args.proxy_protocol_v2,
         handshake_timeout: args.handshake_timeout,
+        region: args.region.clone(),
         wake_compute_retry_config: config::RetryConfig::parse(&args.wake_compute_retry)?,
         connect_compute_locks,
         connect_to_compute: compute_config,
@@ -811,18 +843,21 @@ fn build_auth_backend(
 
 async fn configure_redis(
     args: &ProxyCliArgs,
-) -> anyhow::Result<Option<ConnectionWithCredentialsProvider>> {
+) -> anyhow::Result<(
+    Option<ConnectionWithCredentialsProvider>,
+    Option<ConnectionWithCredentialsProvider>,
+)> {
     // TODO: untangle the config args
-    let redis_client = match &*args.redis_auth_type {
-        "plain" => match &args.redis_plain {
+    let regional_redis_client = match (args.redis_auth_type.as_str(), &args.redis_notifications) {
+        ("plain", redis_url) => match redis_url {
             None => {
-                bail!("plain auth requires redis_plain to be set");
+                bail!("plain auth requires redis_notifications to be set");
             }
             Some(url) => {
                 Some(ConnectionWithCredentialsProvider::new_with_static_credentials(url.clone()))
             }
         },
-        "irsa" => match (&args.redis_host, args.redis_port) {
+        ("irsa", _) => match (&args.redis_host, args.redis_port) {
             (Some(host), Some(port)) => Some(
                 ConnectionWithCredentialsProvider::new_with_credentials_provider(
                     host.clone(),
@@ -846,12 +881,18 @@ async fn configure_redis(
                 bail!("redis-host and redis-port must be specified together");
             }
         },
-        auth_type => {
-            bail!("unknown auth type {auth_type:?} given")
+        _ => {
+            bail!("unknown auth type given");
         }
     };
 
-    Ok(redis_client)
+    let redis_notifications_client = if let Some(url) = &args.redis_notifications {
+        Some(ConnectionWithCredentialsProvider::new_with_static_credentials(&**url))
+    } else {
+        regional_redis_client.clone()
+    };
+
+    Ok((regional_redis_client, redis_notifications_client))
 }
 
 #[cfg(test)]

proxy/src/cache/timed_lru.rs

@@ -30,7 +30,7 @@ use super::{Cache, timed_lru};
 ///
 /// * There's an API for immediate invalidation (removal) of a cache entry;
 ///   It's useful in case we know for sure that the entry is no longer correct.
-///   See [`timed_lru::Cached`] for more information.
+///   See [`timed_lru::LookupInfo`] & [`timed_lru::Cached`] for more information.
 ///
 /// * Expired entries are kept in the cache, until they are evicted by the LRU policy,
 ///   or by a successful lookup (i.e. the entry hasn't expired yet).
@@ -54,7 +54,7 @@ pub(crate) struct TimedLru<K, V> {
 impl<K: Hash + Eq, V> Cache for TimedLru<K, V> {
     type Key = K;
     type Value = V;
-    type LookupInfo<Key> = Key;
+    type LookupInfo<Key> = LookupInfo<Key>;
 
     fn invalidate(&self, info: &Self::LookupInfo<K>) {
         self.invalidate_raw(info);
@@ -87,24 +87,30 @@ impl<K: Hash + Eq, V> TimedLru<K, V> {
 
     /// Drop an entry from the cache if it's outdated.
     #[tracing::instrument(level = "debug", fields(cache = self.name), skip_all)]
-    fn invalidate_raw(&self, key: &K) {
+    fn invalidate_raw(&self, info: &LookupInfo<K>) {
+        let now = Instant::now();
+
         // Do costly things before taking the lock.
         let mut cache = self.cache.lock();
-        let entry = match cache.raw_entry_mut().from_key(key) {
+        let raw_entry = match cache.raw_entry_mut().from_key(&info.key) {
             RawEntryMut::Vacant(_) => return,
-            RawEntryMut::Occupied(x) => x.remove(),
+            RawEntryMut::Occupied(x) => x,
         };
+
+        // Remove the entry if it was created prior to lookup timestamp.
+        let entry = raw_entry.get();
+        let (created_at, expires_at) = (entry.created_at, entry.expires_at);
+        let should_remove = created_at <= info.created_at || expires_at <= now;
+
+        if should_remove {
+            raw_entry.remove();
+        }
+
         drop(cache); // drop lock before logging
-
-        let Entry {
-            created_at,
-            expires_at,
-            ..
-        } = entry;
-
         debug!(
-            ?created_at,
-            ?expires_at,
+            created_at = format_args!("{created_at:?}"),
+            expires_at = format_args!("{expires_at:?}"),
+            entry_removed = should_remove,
             "processed a cache entry invalidation event"
         );
     }
@@ -205,10 +211,10 @@ impl<K: Hash + Eq + Clone, V: Clone> TimedLru<K, V> {
     }
 
     pub(crate) fn insert_unit(&self, key: K, value: V) -> (Option<V>, Cached<&Self, ()>) {
-        let (_, old) = self.insert_raw(key.clone(), value);
+        let (created_at, old) = self.insert_raw(key.clone(), value);
 
         let cached = Cached {
-            token: Some((self, key)),
+            token: Some((self, LookupInfo { created_at, key })),
             value: (),
         };
 
@@ -223,9 +229,28 @@ impl<K: Hash + Eq, V: Clone> TimedLru<K, V> {
         K: Borrow<Q> + Clone,
         Q: Hash + Eq + ?Sized,
     {
-        self.get_raw(key, |key, entry| Cached {
-            token: Some((self, key.clone())),
-            value: entry.value.clone(),
+        self.get_raw(key, |key, entry| {
+            let info = LookupInfo {
+                created_at: entry.created_at,
+                key: key.clone(),
+            };
+
+            Cached {
+                token: Some((self, info)),
+                value: entry.value.clone(),
+            }
         })
     }
 }
+
+/// Lookup information for key invalidation.
+pub(crate) struct LookupInfo<K> {
+    /// Time of creation of a cache [`Entry`].
+    /// We use this during invalidation lookups to prevent eviction of a newer
+    /// entry sharing the same key (it might've been inserted by a different
+    /// task after we got the entry we're trying to invalidate now).
+    created_at: Instant,
+
+    /// Search by this key.
+    key: K,
+}

proxy/src/cancellation.rs

@@ -350,7 +350,7 @@ impl CancellationHandler {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CancelClosure {
     socket_addr: SocketAddr,
-    cancel_token: RawCancelToken,
+    pub cancel_token: RawCancelToken,
     hostname: String, // for pg_sni router
     user_info: ComputeUserInfo,
 }

proxy/src/compute/mod.rs

@@ -86,6 +86,14 @@ pub(crate) enum ConnectionError {
 
     #[error("error acquiring resource permit: {0}")]
     TooManyConnectionAttempts(#[from] ApiLockError),
+
+    #[cfg(test)]
+    #[error("retryable: {retryable}, wakeable: {wakeable}, kind: {kind:?}")]
+    TestError {
+        retryable: bool,
+        wakeable: bool,
+        kind: crate::error::ErrorKind,
+    },
 }
 
 impl UserFacingError for ConnectionError {
@@ -96,6 +104,8 @@ impl UserFacingError for ConnectionError {
                 "Failed to acquire permit to connect to the database. Too many database connection attempts are currently ongoing.".to_owned()
             }
             ConnectionError::TlsError(_) => COULD_NOT_CONNECT.to_owned(),
+            #[cfg(test)]
+            ConnectionError::TestError { .. } => self.to_string(),
         }
     }
 }
@@ -106,6 +116,8 @@ impl ReportableError for ConnectionError {
             ConnectionError::TlsError(_) => crate::error::ErrorKind::Compute,
             ConnectionError::WakeComputeError(e) => e.get_error_kind(),
             ConnectionError::TooManyConnectionAttempts(e) => e.get_error_kind(),
+            #[cfg(test)]
+            ConnectionError::TestError { kind, .. } => *kind,
         }
     }
 }
@@ -236,7 +248,7 @@ impl AuthInfo {
         &self,
         ctx: &RequestContext,
         compute: &mut ComputeConnection,
-        user_info: &ComputeUserInfo,
+        user_info: ComputeUserInfo,
     ) -> Result<PostgresSettings, PostgresError> {
         // client config with stubbed connect info.
         // TODO(conrad): should we rewrite this to bypass tokio-postgres2 entirely,
@@ -252,6 +264,19 @@ impl AuthInfo {
             .await?;
         drop(pause);
 
+        // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
+        info!(
+            compute_id = %compute.aux.compute_id,
+            pid = connection.process_id,
+            cold_start_info = ctx.cold_start_info().as_str(),
+            query_id = ctx.get_testodrome_id().as_deref(),
+            sslmode = ?compute.ssl_mode,
+            "connected to compute node at {} ({}) latency={}",
+            compute.hostname,
+            compute.socket_addr,
+            ctx.get_proxy_latency(),
+        );
+
         let RawConnection {
             stream: _,
             parameters,
@@ -260,8 +285,6 @@ impl AuthInfo {
             secret_key,
         } = connection;
 
-        tracing::Span::current().record("pid", tracing::field::display(process_id));
-
         // NB: CancelToken is supposed to hold socket_addr, but we use connect_raw.
         // Yet another reason to rework the connection establishing code.
         let cancel_closure = CancelClosure::new(
@@ -272,7 +295,7 @@ impl AuthInfo {
                 secret_key,
             },
             compute.hostname.to_string(),
-            user_info.clone(),
+            user_info,
         );
 
         Ok(PostgresSettings {
@@ -288,6 +311,7 @@ impl ConnectInfo {
     async fn connect_raw(
         &self,
         config: &ComputeConfig,
+        direct: bool,
     ) -> Result<(SocketAddr, MaybeTlsStream<TcpStream, RustlsStream>), TlsError> {
         let timeout = config.timeout;
 
@@ -330,7 +354,7 @@ impl ConnectInfo {
         match connect_once(&*addrs).await {
             Ok((sockaddr, stream)) => Ok((
                 sockaddr,
-                tls::connect_tls(stream, self.ssl_mode, config, host).await?,
+                tls::connect_tls(stream, self.ssl_mode, config, host, direct).await?,
             )),
             Err(err) => {
                 warn!("couldn't connect to compute node at {host}:{port}: {err}");
@@ -357,7 +381,7 @@ pub struct PostgresSettings {
 
 pub struct ComputeConnection {
     /// Socket connected to a compute node.
-    pub stream: MaybeTlsStream<tokio::net::TcpStream, RustlsStream>,
+    pub stream: MaybeRustlsStream,
     /// Labels for proxy's metrics.
     pub aux: MetricsAuxInfo,
     pub hostname: Host,
@@ -373,23 +397,12 @@ impl ConnectInfo {
         ctx: &RequestContext,
         aux: &MetricsAuxInfo,
         config: &ComputeConfig,
+        direct: bool,
     ) -> Result<ComputeConnection, ConnectionError> {
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let (socket_addr, stream) = self.connect_raw(config).await?;
+        let (socket_addr, stream) = self.connect_raw(config, direct).await?;
         drop(pause);
 
-        tracing::Span::current().record("compute_id", tracing::field::display(&aux.compute_id));
-
-        // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
-        info!(
-            cold_start_info = ctx.cold_start_info().as_str(),
-            "connected to compute node at {} ({socket_addr}) sslmode={:?}, latency={}, query_id={}",
-            self.host,
-            self.ssl_mode,
-            ctx.get_proxy_latency(),
-            ctx.get_testodrome_id().unwrap_or_default(),
-        );
-
         let connection = ComputeConnection {
             stream,
             socket_addr,

proxy/src/compute/tls.rs

@@ -11,8 +11,6 @@ use crate::proxy::retry::CouldRetry;
 
 #[derive(Debug, Error)]
 pub enum TlsError {
-    #[error(transparent)]
-    Dns(#[from] InvalidDnsNameError),
     #[error(transparent)]
     Connection(#[from] std::io::Error),
     #[error("TLS required but not provided")]
@@ -22,7 +20,6 @@ pub enum TlsError {
 impl CouldRetry for TlsError {
     fn could_retry(&self) -> bool {
         match self {
-            TlsError::Dns(_) => false,
             TlsError::Connection(err) => err.could_retry(),
             // perhaps compute didn't realise it supports TLS?
             TlsError::Required => true,
@@ -35,6 +32,7 @@ pub async fn connect_tls<S, T>(
     mode: SslMode,
     tls: &T,
     host: &str,
+    direct: bool,
 ) -> Result<MaybeTlsStream<S, T::Stream>, TlsError>
 where
     S: AsyncRead + AsyncWrite + Unpin + Send,
@@ -49,7 +47,7 @@ where
         SslMode::Prefer | SslMode::Require => {}
     }
 
-    if !request_tls(&mut stream).await? {
+    if !direct && !request_tls(&mut stream).await? {
         if SslMode::Require == mode {
             return Err(TlsError::Required);
         }
@@ -57,7 +55,6 @@ where
         return Ok(MaybeTlsStream::Raw(stream));
     }
 
-    Ok(MaybeTlsStream::Tls(
-        tls.make_tls_connect(host)?.connect(stream).boxed().await?,
-    ))
+    let c = tls.make_tls_connect(host).map_err(std::io::Error::other)?;
+    Ok(MaybeTlsStream::Tls(c.connect(stream).boxed().await?))
 }

proxy/src/config.rs

@@ -22,6 +22,7 @@ pub struct ProxyConfig {
     pub http_config: HttpConfig,
     pub authentication_config: AuthenticationConfig,
     pub proxy_protocol_v2: ProxyProtocolV2,
+    pub region: String,
     pub handshake_timeout: Duration,
     pub wake_compute_retry_config: RetryConfig,
     pub connect_compute_locks: ApiLocks<Host>,

proxy/src/console_redirect_proxy.rs

@@ -11,12 +11,11 @@ use crate::config::{ProxyConfig, ProxyProtocolV2};
 use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
-use crate::pglb::ClientRequestError;
 use crate::pglb::handshake::{HandshakeData, handshake};
 use crate::pglb::passthrough::ProxyPassthrough;
 use crate::protocol2::{ConnectHeader, ConnectionInfo, read_proxy_protocol};
 use crate::proxy::connect_compute::{TcpMechanism, connect_to_compute};
-use crate::proxy::{ErrorSource, finish_client_init};
+use crate::proxy::{ClientRequestError, ErrorSource, prepare_client_connection};
 use crate::util::run_until_cancelled;
 
 pub async fn task_main(
@@ -90,7 +89,12 @@ pub async fn task_main(
                 }
             }
 
-            let ctx = RequestContext::new(session_id, conn_info, crate::metrics::Protocol::Tcp);
+            let ctx = RequestContext::new(
+                session_id,
+                conn_info,
+                crate::metrics::Protocol::Tcp,
+                &config.region,
+            );
 
             let res = handle_client(
                 config,
@@ -218,6 +222,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         ctx,
         &TcpMechanism {
             locks: &config.connect_compute_locks,
+            direct: false,
         },
         &node_info,
         config.wake_compute_retry_config,
@@ -227,13 +232,13 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     .await?;
 
     let pg_settings = auth_info
-        .authenticate(ctx, &mut node, &user_info)
+        .authenticate(ctx, &mut node, user_info)
         .or_else(|e| async { Err(stream.throw_error(e, Some(ctx)).await) })
         .await?;
 
     let session = cancellation_handler.get_key();
 
-    finish_client_init(&pg_settings, *session.key(), &mut stream);
+    prepare_client_connection(&pg_settings, *session.key(), &mut stream);
     let stream = stream.flush_and_into_inner().await?;
 
     let session_id = ctx.session_id();

proxy/src/context/mod.rs

@@ -46,6 +46,7 @@ struct RequestContextInner {
     pub(crate) session_id: Uuid,
     pub(crate) protocol: Protocol,
     first_packet: chrono::DateTime<Utc>,
+    region: &'static str,
     pub(crate) span: Span,
 
     // filled in as they are discovered
@@ -93,6 +94,7 @@ impl Clone for RequestContext {
             session_id: inner.session_id,
             protocol: inner.protocol,
             first_packet: inner.first_packet,
+            region: inner.region,
             span: info_span!("background_task"),
 
             project: inner.project,
@@ -122,7 +124,12 @@ impl Clone for RequestContext {
 }
 
 impl RequestContext {
-    pub fn new(session_id: Uuid, conn_info: ConnectionInfo, protocol: Protocol) -> Self {
+    pub fn new(
+        session_id: Uuid,
+        conn_info: ConnectionInfo,
+        protocol: Protocol,
+        region: &'static str,
+    ) -> Self {
         // TODO: be careful with long lived spans
         let span = info_span!(
             "connect_request",
@@ -138,6 +145,7 @@ impl RequestContext {
             session_id,
             protocol,
             first_packet: Utc::now(),
+            region,
             span,
 
             project: None,
@@ -171,7 +179,7 @@ impl RequestContext {
         let ip = IpAddr::from([127, 0, 0, 1]);
         let addr = SocketAddr::new(ip, 5432);
         let conn_info = ConnectionInfo { addr, extra: None };
-        RequestContext::new(Uuid::now_v7(), conn_info, Protocol::Tcp)
+        RequestContext::new(Uuid::now_v7(), conn_info, Protocol::Tcp, "test")
     }
 
     pub(crate) fn console_application_name(&self) -> String {

proxy/src/context/parquet.rs

@@ -74,7 +74,7 @@ pub(crate) const FAILED_UPLOAD_MAX_RETRIES: u32 = 10;
 
 #[derive(parquet_derive::ParquetRecordWriter)]
 pub(crate) struct RequestData {
-    region: String,
+    region: &'static str,
     protocol: &'static str,
     /// Must be UTC. The derive macro doesn't like the timezones
     timestamp: chrono::NaiveDateTime,
@@ -147,7 +147,7 @@ impl From<&RequestContextInner> for RequestData {
             }),
             jwt_issuer: value.jwt_issuer.clone(),
             protocol: value.protocol.as_str(),
-            region: String::new(),
+            region: value.region,
             error: value.error_kind.as_ref().map(|e| e.to_metric_label()),
             success: value.success,
             cold_start_info: value.cold_start_info.as_str(),
@@ -167,7 +167,6 @@ impl From<&RequestContextInner> for RequestData {
 pub async fn worker(
     cancellation_token: CancellationToken,
     config: ParquetUploadArgs,
-    region: String,
 ) -> anyhow::Result<()> {
     let Some(remote_storage_config) = config.parquet_upload_remote_storage else {
         tracing::warn!("parquet request upload: no s3 bucket configured");
@@ -233,17 +232,12 @@ pub async fn worker(
                 .context("remote storage for disconnect events init")?;
         let parquet_config_disconnect = parquet_config.clone();
         tokio::try_join!(
-            worker_inner(storage, rx, parquet_config, &region),
-            worker_inner(
-                storage_disconnect,
-                rx_disconnect,
-                parquet_config_disconnect,
-                &region
-            )
+            worker_inner(storage, rx, parquet_config),
+            worker_inner(storage_disconnect, rx_disconnect, parquet_config_disconnect)
         )
         .map(|_| ())
     } else {
-        worker_inner(storage, rx, parquet_config, &region).await
+        worker_inner(storage, rx, parquet_config).await
     }
 }
 
@@ -263,7 +257,6 @@ async fn worker_inner(
     storage: GenericRemoteStorage,
     rx: impl Stream<Item = RequestData>,
     config: ParquetConfig,
-    region: &str,
 ) -> anyhow::Result<()> {
     #[cfg(any(test, feature = "testing"))]
     let storage = if config.test_remote_failures > 0 {
@@ -284,8 +277,7 @@ async fn worker_inner(
     let mut last_upload = time::Instant::now();
 
     let mut len = 0;
-    while let Some(mut row) = rx.next().await {
-        region.clone_into(&mut row.region);
+    while let Some(row) = rx.next().await {
         rows.push(row);
         let force = last_upload.elapsed() > config.max_duration;
         if rows.len() == config.rows_per_group || force {
@@ -541,7 +533,7 @@ mod tests {
             auth_method: None,
             jwt_issuer: None,
             protocol: ["tcp", "ws", "http"][rng.gen_range(0..3)],
-            region: String::new(),
+            region: "us-east-1",
             error: None,
             success: rng.r#gen(),
             cold_start_info: "no",
@@ -573,9 +565,7 @@ mod tests {
             .await
             .unwrap();
 
-        worker_inner(storage, rx, config, "us-east-1")
-            .await
-            .unwrap();
+        worker_inner(storage, rx, config).await.unwrap();
 
         let mut files = WalkDir::new(tmpdir.as_std_path())
             .into_iter()

proxy/src/control_plane/client/cplane_proxy_v1.rs

@@ -263,7 +263,12 @@ impl NeonControlPlaneClient {
                 None => SslMode::Disable,
             };
             let host = match body.server_name {
-                Some(host) => host.into(),
+                Some(host) => {
+                    if rustls::pki_types::DnsName::try_from_str(&host).is_err() {
+                        return Err(WakeComputeError::BadComputeAddress(host.into_boxed_str()));
+                    }
+                    host.into()
+                }
                 None => host.into(),
             };
 

proxy/src/control_plane/mod.rs

@@ -77,8 +77,9 @@ impl NodeInfo {
         &self,
         ctx: &RequestContext,
         config: &ComputeConfig,
+        direct: bool,
     ) -> Result<compute::ComputeConnection, compute::ConnectionError> {
-        self.conn_info.connect(ctx, &self.aux, config).await
+        self.conn_info.connect(ctx, &self.aux, config, direct).await
     }
 }
 

proxy/src/pglb/handshake.rs

@@ -8,10 +8,10 @@ use crate::config::TlsConfig;
 use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::Metrics;
-use crate::pglb::TlsRequired;
 use crate::pqproto::{
     BeMessage, CancelKeyData, FeStartupPacket, ProtocolVersion, StartupMessageParams,
 };
+use crate::proxy::TlsRequired;
 use crate::stream::{PqStream, Stream, StreamUpgradeError};
 use crate::tls::PG_ALPN_PROTOCOL;
 

proxy/src/pglb/mod.rs

@@ -2,332 +2,3 @@ pub mod copy_bidirectional;
 pub mod handshake;
 pub mod inprocess;
 pub mod passthrough;
-
-use std::sync::Arc;
-
-use futures::FutureExt;
-use smol_str::ToSmolStr;
-use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite};
-use tokio_util::sync::CancellationToken;
-use tracing::{Instrument, debug, error, info, warn};
-
-use crate::auth;
-use crate::cancellation::{self, CancellationHandler};
-use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig};
-use crate::context::RequestContext;
-use crate::error::{ReportableError, UserFacingError};
-use crate::metrics::{Metrics, NumClientConnectionsGuard};
-pub use crate::pglb::copy_bidirectional::ErrorSource;
-use crate::pglb::handshake::{HandshakeData, HandshakeError, handshake};
-use crate::pglb::passthrough::ProxyPassthrough;
-use crate::protocol2::{ConnectHeader, ConnectionInfo, ConnectionInfoExtra, read_proxy_protocol};
-use crate::proxy::handle_client;
-use crate::rate_limiter::EndpointRateLimiter;
-use crate::stream::Stream;
-use crate::util::run_until_cancelled;
-
-pub const ERR_INSECURE_CONNECTION: &str = "connection is insecure (try using `sslmode=require`)";
-
-#[derive(Error, Debug)]
-#[error("{ERR_INSECURE_CONNECTION}")]
-pub struct TlsRequired;
-
-impl ReportableError for TlsRequired {
-    fn get_error_kind(&self) -> crate::error::ErrorKind {
-        crate::error::ErrorKind::User
-    }
-}
-
-impl UserFacingError for TlsRequired {}
-
-pub async fn task_main(
-    config: &'static ProxyConfig,
-    auth_backend: &'static auth::Backend<'static, ()>,
-    listener: tokio::net::TcpListener,
-    cancellation_token: CancellationToken,
-    cancellation_handler: Arc<CancellationHandler>,
-    endpoint_rate_limiter: Arc<EndpointRateLimiter>,
-) -> anyhow::Result<()> {
-    scopeguard::defer! {
-        info!("proxy has shut down");
-    }
-
-    // When set for the server socket, the keepalive setting
-    // will be inherited by all accepted client sockets.
-    socket2::SockRef::from(&listener).set_keepalive(true)?;
-
-    let connections = tokio_util::task::task_tracker::TaskTracker::new();
-    let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
-
-    while let Some(accept_result) =
-        run_until_cancelled(listener.accept(), &cancellation_token).await
-    {
-        let (socket, peer_addr) = accept_result?;
-
-        let conn_gauge = Metrics::get()
-            .proxy
-            .client_connections
-            .guard(crate::metrics::Protocol::Tcp);
-
-        let session_id = uuid::Uuid::new_v4();
-        let cancellation_handler = Arc::clone(&cancellation_handler);
-        let cancellations = cancellations.clone();
-
-        debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
-        let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
-
-        connections.spawn(async move {
-            let (socket, conn_info) = match config.proxy_protocol_v2 {
-                ProxyProtocolV2::Required => {
-                    match read_proxy_protocol(socket).await {
-                        Err(e) => {
-                            warn!("per-client task finished with an error: {e:#}");
-                            return;
-                        }
-                        // our load balancers will not send any more data. let's just exit immediately
-                        Ok((_socket, ConnectHeader::Local)) => {
-                            debug!("healthcheck received");
-                            return;
-                        }
-                        Ok((socket, ConnectHeader::Proxy(info))) => (socket, info),
-                    }
-                }
-                // ignore the header - it cannot be confused for a postgres or http connection so will
-                // error later.
-                ProxyProtocolV2::Rejected => (
-                    socket,
-                    ConnectionInfo {
-                        addr: peer_addr,
-                        extra: None,
-                    },
-                ),
-            };
-
-            match socket.set_nodelay(true) {
-                Ok(()) => {}
-                Err(e) => {
-                    error!(
-                        "per-client task finished with an error: failed to set socket option: {e:#}"
-                    );
-                    return;
-                }
-            }
-
-            let ctx = RequestContext::new(session_id, conn_info, crate::metrics::Protocol::Tcp);
-
-            let res = handle_connection(
-                config,
-                auth_backend,
-                &ctx,
-                cancellation_handler,
-                socket,
-                ClientMode::Tcp,
-                endpoint_rate_limiter2,
-                conn_gauge,
-                cancellations,
-            )
-            .instrument(ctx.span())
-            .boxed()
-            .await;
-
-            match res {
-                Err(e) => {
-                    ctx.set_error_kind(e.get_error_kind());
-                    warn!(parent: &ctx.span(), "per-client task finished with an error: {e:#}");
-                }
-                Ok(None) => {
-                    ctx.set_success();
-                }
-                Ok(Some(p)) => {
-                    ctx.set_success();
-                    let _disconnect = ctx.log_connect();
-                    match p.proxy_pass().await {
-                        Ok(()) => {}
-                        Err(ErrorSource::Client(e)) => {
-                            warn!(
-                                ?session_id,
-                                "per-client task finished with an IO error from the client: {e:#}"
-                            );
-                        }
-                        Err(ErrorSource::Compute(e)) => {
-                            error!(
-                                ?session_id,
-                                "per-client task finished with an IO error from the compute: {e:#}"
-                            );
-                        }
-                    }
-                }
-            }
-        });
-    }
-
-    connections.close();
-    cancellations.close();
-    drop(listener);
-
-    // Drain connections
-    connections.wait().await;
-    cancellations.wait().await;
-
-    Ok(())
-}
-
-pub(crate) enum ClientMode {
-    Tcp,
-    Websockets { hostname: Option<String> },
-}
-
-/// Abstracts the logic of handling TCP vs WS clients
-impl ClientMode {
-    pub fn allow_cleartext(&self) -> bool {
-        match self {
-            ClientMode::Tcp => false,
-            ClientMode::Websockets { .. } => true,
-        }
-    }
-
-    pub fn hostname<'a, S>(&'a self, s: &'a Stream<S>) -> Option<&'a str> {
-        match self {
-            ClientMode::Tcp => s.sni_hostname(),
-            ClientMode::Websockets { hostname } => hostname.as_deref(),
-        }
-    }
-
-    pub fn handshake_tls<'a>(&self, tls: Option<&'a TlsConfig>) -> Option<&'a TlsConfig> {
-        match self {
-            ClientMode::Tcp => tls,
-            // TLS is None here if using websockets, because the connection is already encrypted.
-            ClientMode::Websockets { .. } => None,
-        }
-    }
-}
-
-#[derive(Debug, Error)]
-// almost all errors should be reported to the user, but there's a few cases where we cannot
-// 1. Cancellation: we are not allowed to tell the client any cancellation statuses for security reasons
-// 2. Handshake: handshake reports errors if it can, otherwise if the handshake fails due to protocol violation,
-//    we cannot be sure the client even understands our error message
-// 3. PrepareClient: The client disconnected, so we can't tell them anyway...
-pub(crate) enum ClientRequestError {
-    #[error("{0}")]
-    Cancellation(#[from] cancellation::CancelError),
-    #[error("{0}")]
-    Handshake(#[from] HandshakeError),
-    #[error("{0}")]
-    HandshakeTimeout(#[from] tokio::time::error::Elapsed),
-    #[error("{0}")]
-    PrepareClient(#[from] std::io::Error),
-    #[error("{0}")]
-    ReportedError(#[from] crate::stream::ReportedError),
-}
-
-impl ReportableError for ClientRequestError {
-    fn get_error_kind(&self) -> crate::error::ErrorKind {
-        match self {
-            ClientRequestError::Cancellation(e) => e.get_error_kind(),
-            ClientRequestError::Handshake(e) => e.get_error_kind(),
-            ClientRequestError::HandshakeTimeout(_) => crate::error::ErrorKind::RateLimit,
-            ClientRequestError::ReportedError(e) => e.get_error_kind(),
-            ClientRequestError::PrepareClient(_) => crate::error::ErrorKind::ClientDisconnect,
-        }
-    }
-}
-
-#[allow(clippy::too_many_arguments)]
-pub(crate) async fn handle_connection<S: AsyncRead + AsyncWrite + Unpin + Send>(
-    config: &'static ProxyConfig,
-    auth_backend: &'static auth::Backend<'static, ()>,
-    ctx: &RequestContext,
-    cancellation_handler: Arc<CancellationHandler>,
-    client: S,
-    mode: ClientMode,
-    endpoint_rate_limiter: Arc<EndpointRateLimiter>,
-    conn_gauge: NumClientConnectionsGuard<'static>,
-    cancellations: tokio_util::task::task_tracker::TaskTracker,
-) -> Result<Option<ProxyPassthrough<S>>, ClientRequestError> {
-    debug!(
-        protocol = %ctx.protocol(),
-        "handling interactive connection from client"
-    );
-
-    let metrics = &Metrics::get().proxy;
-    let proto = ctx.protocol();
-    let request_gauge = metrics.connection_requests.guard(proto);
-
-    let tls = config.tls_config.load();
-    let tls = tls.as_deref();
-
-    let record_handshake_error = !ctx.has_private_peer_addr();
-    let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Client);
-    let do_handshake = handshake(ctx, client, mode.handshake_tls(tls), record_handshake_error);
-
-    let (mut client, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
-        .await??
-    {
-        HandshakeData::Startup(client, params) => (client, params),
-        HandshakeData::Cancel(cancel_key_data) => {
-            // spawn a task to cancel the session, but don't wait for it
-            cancellations.spawn({
-                let cancellation_handler_clone = Arc::clone(&cancellation_handler);
-                let ctx = ctx.clone();
-                let cancel_span = tracing::span!(parent: None, tracing::Level::INFO, "cancel_session", session_id = ?ctx.session_id());
-                cancel_span.follows_from(tracing::Span::current());
-                async move {
-                    cancellation_handler_clone
-                        .cancel_session(
-                            cancel_key_data,
-                            ctx,
-                            config.authentication_config.ip_allowlist_check_enabled,
-                            config.authentication_config.is_vpc_acccess_proxy,
-                            auth_backend.get_api(),
-                        )
-                        .await
-                        .inspect_err(|e | debug!(error = ?e, "cancel_session failed")).ok();
-                }.instrument(cancel_span)
-            });
-
-            return Ok(None);
-        }
-    };
-    drop(pause);
-
-    ctx.set_db_options(params.clone());
-
-    let common_names = tls.map(|tls| &tls.common_names);
-
-    let (node, cancel_on_shutdown) = handle_client(
-        config,
-        auth_backend,
-        ctx,
-        cancellation_handler,
-        &mut client,
-        &mode,
-        endpoint_rate_limiter,
-        common_names,
-        &params,
-    )
-    .await?;
-
-    let client = client.flush_and_into_inner().await?;
-
-    let private_link_id = match ctx.extra() {
-        Some(ConnectionInfoExtra::Aws { vpce_id }) => Some(vpce_id.clone()),
-        Some(ConnectionInfoExtra::Azure { link_id }) => Some(link_id.to_smolstr()),
-        None => None,
-    };
-
-    Ok(Some(ProxyPassthrough {
-        client,
-        compute: node.stream,
-
-        aux: node.aux,
-        private_link_id,
-
-        _cancel_on_shutdown: cancel_on_shutdown,
-
-        _req: request_gauge,
-        _conn: conn_gauge,
-        _db_conn: node.guage,
-    }))
-}

proxy/src/proxy/connect_compute.rs

@@ -1,18 +1,15 @@
-use async_trait::async_trait;
 use tokio::time;
 use tracing::{debug, info, warn};
 
 use crate::compute::{self, COULD_NOT_CONNECT, ComputeConnection};
 use crate::config::{ComputeConfig, RetryConfig};
 use crate::context::RequestContext;
-use crate::control_plane::errors::WakeComputeError;
 use crate::control_plane::locks::ApiLocks;
 use crate::control_plane::{self, NodeInfo};
-use crate::error::ReportableError;
 use crate::metrics::{
     ConnectOutcome, ConnectionFailureKind, Metrics, RetriesMetricGroup, RetryType,
 };
-use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute, retry_after, should_retry};
+use crate::proxy::retry::{ShouldRetryWakeCompute, retry_after, should_retry};
 use crate::proxy::wake_compute::{WakeComputeBackend, wake_compute};
 use crate::types::Host;
 
@@ -35,42 +32,34 @@ pub(crate) fn invalidate_cache(node_info: control_plane::CachedNodeInfo) -> Node
     node_info.invalidate()
 }
 
-#[async_trait]
 pub(crate) trait ConnectMechanism {
     type Connection;
-    type ConnectError: ReportableError;
-    type Error: From<Self::ConnectError>;
     async fn connect_once(
         &self,
         ctx: &RequestContext,
         node_info: &control_plane::CachedNodeInfo,
         config: &ComputeConfig,
-    ) -> Result<Self::Connection, Self::ConnectError>;
+    ) -> Result<Self::Connection, compute::ConnectionError>;
 }
 
 pub(crate) struct TcpMechanism {
     /// connect_to_compute concurrency lock
     pub(crate) locks: &'static ApiLocks<Host>,
+    // whether to negotiate TLS for postgres protocol.
+    pub(crate) direct: bool,
 }
 
-#[async_trait]
 impl ConnectMechanism for TcpMechanism {
     type Connection = ComputeConnection;
-    type ConnectError = compute::ConnectionError;
-    type Error = compute::ConnectionError;
 
-    #[tracing::instrument(skip_all, fields(
-        pid = tracing::field::Empty,
-        compute_id = tracing::field::Empty
-    ))]
     async fn connect_once(
         &self,
         ctx: &RequestContext,
         node_info: &control_plane::CachedNodeInfo,
         config: &ComputeConfig,
-    ) -> Result<ComputeConnection, Self::Error> {
+    ) -> Result<ComputeConnection, compute::ConnectionError> {
         let permit = self.locks.get_permit(&node_info.conn_info.host).await?;
-        permit.release_result(node_info.connect(ctx, config).await)
+        permit.release_result(node_info.connect(ctx, config, self.direct).await)
     }
 }
 
@@ -82,11 +71,7 @@ pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: WakeComputeBacken
     user_info: &B,
     wake_compute_retry_config: RetryConfig,
     compute: &ComputeConfig,
-) -> Result<M::Connection, M::Error>
-where
-    M::ConnectError: CouldRetry + ShouldRetryWakeCompute + std::fmt::Debug,
-    M::Error: From<WakeComputeError>,
-{
+) -> Result<M::Connection, compute::ConnectionError> {
     let mut num_retries = 0;
     let node_info =
         wake_compute(&mut num_retries, ctx, user_info, wake_compute_retry_config).await?;
@@ -112,7 +97,7 @@ where
     let node_info = if !node_info.cached() || !err.should_retry_wake_compute() {
         // If we just recieved this from cplane and didn't get it from cache, we shouldn't retry.
         // Do not need to retrieve a new node_info, just return the old one.
-        if !should_retry(&err, num_retries, compute.retry) {
+        if should_retry(&err, num_retries, compute.retry) {
             Metrics::get().proxy.retries_metric.observe(
                 RetriesMetricGroup {
                     outcome: ConnectOutcome::Failed,
@@ -120,7 +105,7 @@ where
                 },
                 num_retries.into(),
             );
-            return Err(err.into());
+            return Err(err);
         }
         node_info
     } else {
@@ -161,7 +146,7 @@ where
                         },
                         num_retries.into(),
                     );
-                    return Err(e.into());
+                    return Err(e);
                 }
 
                 warn!(error = ?e, num_retries, retriable = true, COULD_NOT_CONNECT);

proxy/src/proxy/mod.rs

@@ -5,64 +5,328 @@ pub(crate) mod connect_compute;
 pub(crate) mod retry;
 pub(crate) mod wake_compute;
 
-use std::collections::HashSet;
-use std::convert::Infallible;
 use std::sync::Arc;
 
+use futures::FutureExt;
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
 use regex::Regex;
 use serde::{Deserialize, Serialize};
-use smol_str::{SmolStr, format_smolstr};
+use smol_str::{SmolStr, ToSmolStr, format_smolstr};
+use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite};
-use tokio::sync::oneshot;
-use tracing::Instrument;
+use tokio_util::sync::CancellationToken;
+use tracing::{Instrument, debug, error, info, warn};
 
-use crate::cache::Cache;
-use crate::cancellation::CancellationHandler;
-use crate::compute::ComputeConnection;
-use crate::config::ProxyConfig;
+use crate::cancellation::{self, CancellationHandler};
+use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig};
 use crate::context::RequestContext;
-use crate::control_plane::client::ControlPlaneClient;
+use crate::error::{ReportableError, UserFacingError};
+use crate::metrics::{Metrics, NumClientConnectionsGuard};
 pub use crate::pglb::copy_bidirectional::{ErrorSource, copy_bidirectional_client_compute};
-use crate::pglb::{ClientMode, ClientRequestError};
+use crate::pglb::handshake::{HandshakeData, HandshakeError, handshake};
+use crate::pglb::passthrough::ProxyPassthrough;
 use crate::pqproto::{BeMessage, CancelKeyData, StartupMessageParams};
+use crate::protocol2::{ConnectHeader, ConnectionInfo, ConnectionInfoExtra, read_proxy_protocol};
 use crate::proxy::connect_compute::{TcpMechanism, connect_to_compute};
-use crate::proxy::retry::ShouldRetryWakeCompute;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::stream::{PqStream, Stream};
 use crate::types::EndpointCacheKey;
+use crate::util::run_until_cancelled;
 use crate::{auth, compute};
 
+const ERR_INSECURE_CONNECTION: &str = "connection is insecure (try using `sslmode=require`)";
+
+#[derive(Error, Debug)]
+#[error("{ERR_INSECURE_CONNECTION}")]
+pub struct TlsRequired;
+
+impl ReportableError for TlsRequired {
+    fn get_error_kind(&self) -> crate::error::ErrorKind {
+        crate::error::ErrorKind::User
+    }
+}
+
+impl UserFacingError for TlsRequired {}
+
+pub async fn task_main(
+    config: &'static ProxyConfig,
+    auth_backend: &'static auth::Backend<'static, ()>,
+    listener: tokio::net::TcpListener,
+    cancellation_token: CancellationToken,
+    cancellation_handler: Arc<CancellationHandler>,
+    endpoint_rate_limiter: Arc<EndpointRateLimiter>,
+) -> anyhow::Result<()> {
+    scopeguard::defer! {
+        info!("proxy has shut down");
+    }
+
+    // When set for the server socket, the keepalive setting
+    // will be inherited by all accepted client sockets.
+    socket2::SockRef::from(&listener).set_keepalive(true)?;
+
+    let connections = tokio_util::task::task_tracker::TaskTracker::new();
+    let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
+
+    while let Some(accept_result) =
+        run_until_cancelled(listener.accept(), &cancellation_token).await
+    {
+        let (socket, peer_addr) = accept_result?;
+
+        let conn_gauge = Metrics::get()
+            .proxy
+            .client_connections
+            .guard(crate::metrics::Protocol::Tcp);
+
+        let session_id = uuid::Uuid::new_v4();
+        let cancellation_handler = Arc::clone(&cancellation_handler);
+        let cancellations = cancellations.clone();
+
+        debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
+        let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
+
+        connections.spawn(async move {
+            let (socket, conn_info) = match config.proxy_protocol_v2 {
+                ProxyProtocolV2::Required => {
+                    match read_proxy_protocol(socket).await {
+                        Err(e) => {
+                            warn!("per-client task finished with an error: {e:#}");
+                            return;
+                        }
+                        // our load balancers will not send any more data. let's just exit immediately
+                        Ok((_socket, ConnectHeader::Local)) => {
+                            debug!("healthcheck received");
+                            return;
+                        }
+                        Ok((socket, ConnectHeader::Proxy(info))) => (socket, info),
+                    }
+                }
+                // ignore the header - it cannot be confused for a postgres or http connection so will
+                // error later.
+                ProxyProtocolV2::Rejected => (
+                    socket,
+                    ConnectionInfo {
+                        addr: peer_addr,
+                        extra: None,
+                    },
+                ),
+            };
+
+            match socket.set_nodelay(true) {
+                Ok(()) => {}
+                Err(e) => {
+                    error!(
+                        "per-client task finished with an error: failed to set socket option: {e:#}"
+                    );
+                    return;
+                }
+            }
+
+            let ctx = RequestContext::new(
+                session_id,
+                conn_info,
+                crate::metrics::Protocol::Tcp,
+                &config.region,
+            );
+
+            let res = handle_client(
+                config,
+                auth_backend,
+                &ctx,
+                cancellation_handler,
+                socket,
+                ClientMode::Tcp,
+                endpoint_rate_limiter2,
+                conn_gauge,
+                cancellations,
+            )
+            .instrument(ctx.span())
+            .boxed()
+            .await;
+
+            match res {
+                Err(e) => {
+                    ctx.set_error_kind(e.get_error_kind());
+                    warn!(parent: &ctx.span(), "per-client task finished with an error: {e:#}");
+                }
+                Ok(None) => {
+                    ctx.set_success();
+                }
+                Ok(Some(p)) => {
+                    ctx.set_success();
+                    let _disconnect = ctx.log_connect();
+                    match p.proxy_pass().await {
+                        Ok(()) => {}
+                        Err(ErrorSource::Client(e)) => {
+                            warn!(
+                                ?session_id,
+                                "per-client task finished with an IO error from the client: {e:#}"
+                            );
+                        }
+                        Err(ErrorSource::Compute(e)) => {
+                            error!(
+                                ?session_id,
+                                "per-client task finished with an IO error from the compute: {e:#}"
+                            );
+                        }
+                    }
+                }
+            }
+        });
+    }
+
+    connections.close();
+    cancellations.close();
+    drop(listener);
+
+    // Drain connections
+    connections.wait().await;
+    cancellations.wait().await;
+
+    Ok(())
+}
+
+pub(crate) enum ClientMode {
+    Tcp,
+    Websockets { hostname: Option<String> },
+}
+
+/// Abstracts the logic of handling TCP vs WS clients
+impl ClientMode {
+    pub(crate) fn allow_cleartext(&self) -> bool {
+        match self {
+            ClientMode::Tcp => false,
+            ClientMode::Websockets { .. } => true,
+        }
+    }
+
+    fn hostname<'a, S>(&'a self, s: &'a Stream<S>) -> Option<&'a str> {
+        match self {
+            ClientMode::Tcp => s.sni_hostname(),
+            ClientMode::Websockets { hostname } => hostname.as_deref(),
+        }
+    }
+
+    fn handshake_tls<'a>(&self, tls: Option<&'a TlsConfig>) -> Option<&'a TlsConfig> {
+        match self {
+            ClientMode::Tcp => tls,
+            // TLS is None here if using websockets, because the connection is already encrypted.
+            ClientMode::Websockets { .. } => None,
+        }
+    }
+}
+
+#[derive(Debug, Error)]
+// almost all errors should be reported to the user, but there's a few cases where we cannot
+// 1. Cancellation: we are not allowed to tell the client any cancellation statuses for security reasons
+// 2. Handshake: handshake reports errors if it can, otherwise if the handshake fails due to protocol violation,
+//    we cannot be sure the client even understands our error message
+// 3. PrepareClient: The client disconnected, so we can't tell them anyway...
+pub(crate) enum ClientRequestError {
+    #[error("{0}")]
+    Cancellation(#[from] cancellation::CancelError),
+    #[error("{0}")]
+    Handshake(#[from] HandshakeError),
+    #[error("{0}")]
+    HandshakeTimeout(#[from] tokio::time::error::Elapsed),
+    #[error("{0}")]
+    PrepareClient(#[from] std::io::Error),
+    #[error("{0}")]
+    ReportedError(#[from] crate::stream::ReportedError),
+}
+
+impl ReportableError for ClientRequestError {
+    fn get_error_kind(&self) -> crate::error::ErrorKind {
+        match self {
+            ClientRequestError::Cancellation(e) => e.get_error_kind(),
+            ClientRequestError::Handshake(e) => e.get_error_kind(),
+            ClientRequestError::HandshakeTimeout(_) => crate::error::ErrorKind::RateLimit,
+            ClientRequestError::ReportedError(e) => e.get_error_kind(),
+            ClientRequestError::PrepareClient(_) => crate::error::ErrorKind::ClientDisconnect,
+        }
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     config: &'static ProxyConfig,
     auth_backend: &'static auth::Backend<'static, ()>,
     ctx: &RequestContext,
     cancellation_handler: Arc<CancellationHandler>,
-    client: &mut PqStream<Stream<S>>,
-    mode: &ClientMode,
+    stream: S,
+    mode: ClientMode,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
-    common_names: Option<&HashSet<String>>,
-    params: &StartupMessageParams,
-) -> Result<(ComputeConnection, oneshot::Sender<Infallible>), ClientRequestError> {
-    let hostname = mode.hostname(client.get_ref());
+    conn_gauge: NumClientConnectionsGuard<'static>,
+    cancellations: tokio_util::task::task_tracker::TaskTracker,
+) -> Result<Option<ProxyPassthrough<S>>, ClientRequestError> {
+    debug!(
+        protocol = %ctx.protocol(),
+        "handling interactive connection from client"
+    );
+
+    let metrics = &Metrics::get().proxy;
+    let proto = ctx.protocol();
+    let request_gauge = metrics.connection_requests.guard(proto);
+
+    let tls = config.tls_config.load();
+    let tls = tls.as_deref();
+
+    let record_handshake_error = !ctx.has_private_peer_addr();
+    let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Client);
+    let do_handshake = handshake(ctx, stream, mode.handshake_tls(tls), record_handshake_error);
+
+    let (mut stream, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
+        .await??
+    {
+        HandshakeData::Startup(stream, params) => (stream, params),
+        HandshakeData::Cancel(cancel_key_data) => {
+            // spawn a task to cancel the session, but don't wait for it
+            cancellations.spawn({
+                let cancellation_handler_clone = Arc::clone(&cancellation_handler);
+                let ctx = ctx.clone();
+                let cancel_span = tracing::span!(parent: None, tracing::Level::INFO, "cancel_session", session_id = ?ctx.session_id());
+                cancel_span.follows_from(tracing::Span::current());
+                async move {
+                    cancellation_handler_clone
+                        .cancel_session(
+                            cancel_key_data,
+                            ctx,
+                            config.authentication_config.ip_allowlist_check_enabled,
+                            config.authentication_config.is_vpc_acccess_proxy,
+                            auth_backend.get_api(),
+                        )
+                        .await
+                        .inspect_err(|e | debug!(error = ?e, "cancel_session failed")).ok();
+                }.instrument(cancel_span)
+            });
+
+            return Ok(None);
+        }
+    };
+    drop(pause);
+
+    ctx.set_db_options(params.clone());
+
+    let hostname = mode.hostname(stream.get_ref());
+
+    let common_names = tls.map(|tls| &tls.common_names);
+
     // Extract credentials which we're going to use for auth.
     let result = auth_backend
         .as_ref()
-        .map(|()| auth::ComputeUserInfoMaybeEndpoint::parse(ctx, params, hostname, common_names))
+        .map(|()| auth::ComputeUserInfoMaybeEndpoint::parse(ctx, &params, hostname, common_names))
         .transpose();
 
     let user_info = match result {
         Ok(user_info) => user_info,
-        Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
+        Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
     };
 
     let user = user_info.get_user().to_owned();
     let user_info = match user_info
         .authenticate(
             ctx,
-            client,
+            &mut stream,
             mode.allow_cleartext(),
             &config.authentication_config,
             endpoint_rate_limiter,
@@ -75,7 +339,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
             let app = params.get("application_name");
             let params_span = tracing::info_span!("", ?user, ?db, ?app);
 
-            return Err(client
+            return Err(stream
                 .throw_error(e, Some(ctx))
                 .instrument(params_span)
                 .await)?;
@@ -88,67 +352,38 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     };
     let params_compat = creds.info.options.get(NeonOptions::PARAMS_COMPAT).is_some();
     let mut auth_info = compute::AuthInfo::with_auth_keys(creds.keys);
-    auth_info.set_startup_params(params, params_compat);
+    auth_info.set_startup_params(&params, params_compat);
 
-    let mut node;
-    let mut attempt = 0;
-    let connect = TcpMechanism {
-        locks: &config.connect_compute_locks,
+    let res = connect_to_compute(
+        ctx,
+        &TcpMechanism {
+            locks: &config.connect_compute_locks,
+            direct: false,
+        },
+        &auth::Backend::ControlPlane(cplane, creds.info.clone()),
+        config.wake_compute_retry_config,
+        &config.connect_to_compute,
+    )
+    .await;
+
+    let mut node = match res {
+        Ok(node) => node,
+        Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
     };
-    let backend = auth::Backend::ControlPlane(cplane, creds.info);
 
-    // NOTE: This is messy, but should hopefully be detangled with PGLB.
-    // We wanted to separate the concerns of **connect** to compute (a PGLB operation),
-    // from **authenticate** to compute (a NeonKeeper operation).
-    //
-    // This unfortunately removed retry handling for one error case where
-    // the compute was cached, and we connected, but the compute cache was actually stale
-    // and is associated with the wrong endpoint. We detect this when the **authentication** fails.
-    // As such, we retry once here if the `authenticate` function fails and the error is valid to retry.
-    let pg_settings = loop {
-        attempt += 1;
-
-        // TODO: callback to pglb
-        let res = connect_to_compute(
-            ctx,
-            &connect,
-            &backend,
-            config.wake_compute_retry_config,
-            &config.connect_to_compute,
-        )
-        .await;
-
-        match res {
-            Ok(n) => node = n,
-            Err(e) => return Err(client.throw_error(e, Some(ctx)).await)?,
-        }
-
-        let auth::Backend::ControlPlane(cplane, user_info) = &backend else {
-            unreachable!("ensured above");
-        };
-
-        let res = auth_info.authenticate(ctx, &mut node, user_info).await;
-        match res {
-            Ok(pg_settings) => break pg_settings,
-            Err(e) if attempt < 2 && e.should_retry_wake_compute() => {
-                tracing::warn!(error = ?e, "retrying wake compute");
-
-                #[allow(irrefutable_let_patterns)]
-                if let ControlPlaneClient::ProxyV1(cplane_proxy_v1) = &**cplane {
-                    let key = user_info.endpoint_cache_key();
-                    cplane_proxy_v1.caches.node_info.invalidate(&key);
-                }
-            }
-            Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
-        }
+    let pg_settings = auth_info.authenticate(ctx, &mut node, creds.info).await;
+    let pg_settings = match pg_settings {
+        Ok(pg_settings) => pg_settings,
+        Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
     };
 
     let session = cancellation_handler.get_key();
 
-    finish_client_init(&pg_settings, *session.key(), client);
+    prepare_client_connection(&pg_settings, *session.key(), &mut stream);
+    let stream = stream.flush_and_into_inner().await?;
 
     let session_id = ctx.session_id();
-    let (cancel_on_shutdown, cancel) = oneshot::channel();
+    let (cancel_on_shutdown, cancel) = tokio::sync::oneshot::channel();
     tokio::spawn(async move {
         session
             .maintain_cancel_key(
@@ -160,32 +395,50 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
             .await;
     });
 
-    Ok((node, cancel_on_shutdown))
+    let private_link_id = match ctx.extra() {
+        Some(ConnectionInfoExtra::Aws { vpce_id }) => Some(vpce_id.clone()),
+        Some(ConnectionInfoExtra::Azure { link_id }) => Some(link_id.to_smolstr()),
+        None => None,
+    };
+
+    Ok(Some(ProxyPassthrough {
+        client: stream,
+        compute: node.stream,
+
+        aux: node.aux,
+        private_link_id,
+
+        _cancel_on_shutdown: cancel_on_shutdown,
+
+        _req: request_gauge,
+        _conn: conn_gauge,
+        _db_conn: node.guage,
+    }))
 }
 
 /// Finish client connection initialization: confirm auth success, send params, etc.
-pub(crate) fn finish_client_init(
+pub(crate) fn prepare_client_connection(
     settings: &compute::PostgresSettings,
     cancel_key_data: CancelKeyData,
-    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    stream: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
 ) {
     // Forward all deferred notices to the client.
     for notice in &settings.delayed_notice {
-        client.write_raw(notice.as_bytes().len(), b'N', |buf| {
+        stream.write_raw(notice.as_bytes().len(), b'N', |buf| {
             buf.extend_from_slice(notice.as_bytes());
         });
     }
 
     // Forward all postgres connection params to the client.
     for (name, value) in &settings.params {
-        client.write_message(BeMessage::ParameterStatus {
+        stream.write_message(BeMessage::ParameterStatus {
             name: name.as_bytes(),
             value: value.as_bytes(),
         });
     }
 
-    client.write_message(BeMessage::BackendKeyData(cancel_key_data));
-    client.write_message(BeMessage::ReadyForQuery);
+    stream.write_message(BeMessage::BackendKeyData(cancel_key_data));
+    stream.write_message(BeMessage::ReadyForQuery);
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
@@ -195,7 +448,7 @@ impl NeonOptions {
     // proxy options:
 
     /// `PARAMS_COMPAT` allows opting in to forwarding all startup parameters from client to compute.
-    pub const PARAMS_COMPAT: &str = "proxy_params_compat";
+    const PARAMS_COMPAT: &str = "proxy_params_compat";
 
     // cplane options:
 

proxy/src/proxy/retry.rs

@@ -1,9 +1,8 @@
-use std::error::Error;
 use std::io;
 
 use tokio::time;
 
-use crate::compute::{self, PostgresError};
+use crate::compute;
 use crate::config::RetryConfig;
 
 pub(crate) trait CouldRetry {
@@ -31,153 +30,32 @@ impl CouldRetry for io::Error {
     }
 }
 
-impl CouldRetry for postgres_client::error::DbError {
-    fn could_retry(&self) -> bool {
-        use postgres_client::error::SqlState;
-        matches!(
-            self.code(),
-            &SqlState::CONNECTION_FAILURE
-                | &SqlState::CONNECTION_EXCEPTION
-                | &SqlState::CONNECTION_DOES_NOT_EXIST
-                | &SqlState::SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION,
-        )
-    }
-}
-impl ShouldRetryWakeCompute for postgres_client::error::DbError {
-    fn should_retry_wake_compute(&self) -> bool {
-        use postgres_client::error::SqlState;
-        // Here are errors that happens after the user successfully authenticated to the database.
-        // TODO: there are pgbouncer errors that should be retried, but they are not listed here.
-        let non_retriable_pg_errors = matches!(
-            self.code(),
-            &SqlState::TOO_MANY_CONNECTIONS
-                | &SqlState::OUT_OF_MEMORY
-                | &SqlState::SYNTAX_ERROR
-                | &SqlState::T_R_SERIALIZATION_FAILURE
-                | &SqlState::INVALID_CATALOG_NAME
-                | &SqlState::INVALID_SCHEMA_NAME
-                | &SqlState::INVALID_PARAMETER_VALUE,
-        );
-        if non_retriable_pg_errors {
-            return false;
-        }
-        // PGBouncer errors that should not trigger a wake_compute retry.
-        if self.code() == &SqlState::PROTOCOL_VIOLATION {
-            // Source for the error message:
-            // https://github.com/pgbouncer/pgbouncer/blob/f15997fe3effe3a94ba8bcc1ea562e6117d1a131/src/client.c#L1070
-            return !self
-                .message()
-                .contains("no more connections allowed (max_client_conn)");
-        }
-        true
-    }
-}
-
-impl CouldRetry for postgres_client::Error {
-    fn could_retry(&self) -> bool {
-        if let Some(io_err) = self.source().and_then(|x| x.downcast_ref()) {
-            io::Error::could_retry(io_err)
-        } else if let Some(db_err) = self.source().and_then(|x| x.downcast_ref()) {
-            postgres_client::error::DbError::could_retry(db_err)
-        } else {
-            false
-        }
-    }
-}
-impl ShouldRetryWakeCompute for postgres_client::Error {
-    fn should_retry_wake_compute(&self) -> bool {
-        if let Some(db_err) = self.source().and_then(|x| x.downcast_ref()) {
-            postgres_client::error::DbError::should_retry_wake_compute(db_err)
-        } else {
-            // likely an IO error. Possible the compute has shutdown and the
-            // cache is stale.
-            true
-        }
-    }
-}
-
 impl CouldRetry for compute::ConnectionError {
     fn could_retry(&self) -> bool {
         match self {
             compute::ConnectionError::TlsError(err) => err.could_retry(),
             compute::ConnectionError::WakeComputeError(err) => err.could_retry(),
             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
+            #[cfg(test)]
+            compute::ConnectionError::TestError { retryable, .. } => *retryable,
         }
     }
 }
+
 impl ShouldRetryWakeCompute for compute::ConnectionError {
     fn should_retry_wake_compute(&self) -> bool {
         match self {
             // the cache entry was not checked for validity
             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
+            #[cfg(test)]
+            compute::ConnectionError::TestError { wakeable, .. } => *wakeable,
             _ => true,
         }
     }
 }
 
-impl ShouldRetryWakeCompute for PostgresError {
-    fn should_retry_wake_compute(&self) -> bool {
-        match self {
-            PostgresError::Postgres(error) => error.should_retry_wake_compute(),
-        }
-    }
-}
-
 pub(crate) fn retry_after(num_retries: u32, config: RetryConfig) -> time::Duration {
     config
         .base_delay
         .mul_f64(config.backoff_factor.powi((num_retries as i32) - 1))
 }
-
-#[cfg(test)]
-mod tests {
-    use postgres_client::error::{DbError, SqlState};
-
-    use super::ShouldRetryWakeCompute;
-
-    #[test]
-    fn should_retry_wake_compute_for_db_error() {
-        // These SQLStates should NOT trigger a wake_compute retry.
-        let non_retry_states = [
-            SqlState::TOO_MANY_CONNECTIONS,
-            SqlState::OUT_OF_MEMORY,
-            SqlState::SYNTAX_ERROR,
-            SqlState::T_R_SERIALIZATION_FAILURE,
-            SqlState::INVALID_CATALOG_NAME,
-            SqlState::INVALID_SCHEMA_NAME,
-            SqlState::INVALID_PARAMETER_VALUE,
-        ];
-        for state in non_retry_states {
-            let err = DbError::new_test_error(state.clone(), "oops".to_string());
-            assert!(
-                !err.should_retry_wake_compute(),
-                "State {state:?} unexpectedly retried"
-            );
-        }
-
-        // Errors coming from pgbouncer should not trigger a wake_compute retry
-        let non_retry_pgbouncer_errors = ["no more connections allowed (max_client_conn)"];
-        for error in non_retry_pgbouncer_errors {
-            let err = DbError::new_test_error(SqlState::PROTOCOL_VIOLATION, error.to_string());
-            assert!(
-                !err.should_retry_wake_compute(),
-                "PGBouncer error {error:?} unexpectedly retried"
-            );
-        }
-
-        // These SQLStates should trigger a wake_compute retry.
-        let retry_states = [
-            SqlState::CONNECTION_FAILURE,
-            SqlState::CONNECTION_EXCEPTION,
-            SqlState::CONNECTION_DOES_NOT_EXIST,
-            SqlState::SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION,
-        ];
-        for state in retry_states {
-            let err = DbError::new_test_error(state.clone(), "oops".to_string());
-            assert!(
-                err.should_retry_wake_compute(),
-                "State {state:?} unexpectedly skipped retry"
-            );
-        }
-    }
-}

proxy/src/proxy/tests/mitm.rs

@@ -14,9 +14,6 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream};
 use tokio_util::codec::{Decoder, Encoder};
 
 use super::*;
-use crate::config::TlsConfig;
-use crate::context::RequestContext;
-use crate::pglb::handshake::{HandshakeData, handshake};
 
 enum Intercept {
     None,

proxy/src/proxy/tests/mod.rs

@@ -3,7 +3,6 @@
 
 mod mitm;
 
-use std::sync::Arc;
 use std::time::Duration;
 
 use anyhow::{Context, bail};
@@ -11,31 +10,27 @@ use async_trait::async_trait;
 use http::StatusCode;
 use postgres_client::config::SslMode;
 use postgres_client::tls::{MakeTlsConnect, NoTls};
+use retry::retry_after;
 use rstest::rstest;
 use rustls::crypto::ring;
 use rustls::pki_types;
-use tokio::io::{AsyncRead, AsyncWrite, DuplexStream};
+use tokio::io::DuplexStream;
 use tracing_test::traced_test;
 
 use super::retry::CouldRetry;
+use super::*;
 use crate::auth::backend::{ComputeUserInfo, MaybeOwned};
-use crate::config::{ComputeConfig, RetryConfig, TlsConfig};
-use crate::context::RequestContext;
+use crate::compute::ConnectionError;
+use crate::config::{ComputeConfig, RetryConfig};
 use crate::control_plane::client::{ControlPlaneClient, TestControlPlaneClient};
 use crate::control_plane::messages::{ControlPlaneErrorMessage, Details, MetricsAuxInfo, Status};
 use crate::control_plane::{self, CachedNodeInfo, NodeInfo, NodeInfoCache};
-use crate::error::{ErrorKind, ReportableError};
-use crate::pglb::ERR_INSECURE_CONNECTION;
-use crate::pglb::handshake::{HandshakeData, handshake};
-use crate::pqproto::BeMessage;
-use crate::proxy::NeonOptions;
-use crate::proxy::connect_compute::{ConnectMechanism, connect_to_compute};
-use crate::proxy::retry::{ShouldRetryWakeCompute, retry_after};
-use crate::stream::{PqStream, Stream};
+use crate::error::ErrorKind;
+use crate::proxy::connect_compute::ConnectMechanism;
 use crate::tls::client_config::compute_client_config_with_certs;
 use crate::tls::server_config::CertResolver;
 use crate::types::{BranchId, EndpointId, ProjectId};
-use crate::{auth, compute, sasl, scram};
+use crate::{sasl, scram};
 
 /// Generate a set of TLS certificates: CA + server.
 fn generate_certs(
@@ -380,7 +375,6 @@ fn connect_compute_total_wait() {
 #[derive(Clone, Copy, Debug)]
 enum ConnectAction {
     Wake,
-    WakeCold,
     WakeFail,
     WakeRetry,
     Connect,
@@ -430,71 +424,36 @@ impl TestConnectMechanism {
 #[derive(Debug)]
 struct TestConnection;
 
-#[derive(Debug)]
-struct TestConnectError {
-    retryable: bool,
-    wakeable: bool,
-    kind: crate::error::ErrorKind,
-}
-
-impl ReportableError for TestConnectError {
-    fn get_error_kind(&self) -> crate::error::ErrorKind {
-        self.kind
-    }
-}
-
-impl std::fmt::Display for TestConnectError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{self:?}")
-    }
-}
-
-impl std::error::Error for TestConnectError {}
-
-impl CouldRetry for TestConnectError {
-    fn could_retry(&self) -> bool {
-        self.retryable
-    }
-}
-impl ShouldRetryWakeCompute for TestConnectError {
-    fn should_retry_wake_compute(&self) -> bool {
-        self.wakeable
-    }
-}
-
-#[async_trait]
 impl ConnectMechanism for TestConnectMechanism {
     type Connection = TestConnection;
-    type ConnectError = TestConnectError;
-    type Error = anyhow::Error;
 
     async fn connect_once(
         &self,
         _ctx: &RequestContext,
         _node_info: &control_plane::CachedNodeInfo,
         _config: &ComputeConfig,
-    ) -> Result<Self::Connection, Self::ConnectError> {
+    ) -> Result<Self::Connection, ConnectionError> {
         let mut counter = self.counter.lock().unwrap();
         let action = self.sequence[*counter];
         *counter += 1;
         match action {
             ConnectAction::Connect => Ok(TestConnection),
-            ConnectAction::Retry => Err(TestConnectError {
+            ConnectAction::Retry => Err(ConnectionError::TestError {
                 retryable: true,
                 wakeable: true,
                 kind: ErrorKind::Compute,
             }),
-            ConnectAction::RetryNoWake => Err(TestConnectError {
+            ConnectAction::RetryNoWake => Err(ConnectionError::TestError {
                 retryable: true,
                 wakeable: false,
                 kind: ErrorKind::Compute,
             }),
-            ConnectAction::Fail => Err(TestConnectError {
+            ConnectAction::Fail => Err(ConnectionError::TestError {
                 retryable: false,
                 wakeable: true,
                 kind: ErrorKind::Compute,
             }),
-            ConnectAction::FailNoWake => Err(TestConnectError {
+            ConnectAction::FailNoWake => Err(ConnectionError::TestError {
                 retryable: false,
                 wakeable: false,
                 kind: ErrorKind::Compute,
@@ -511,9 +470,6 @@ impl TestControlPlaneClient for TestConnectMechanism {
         *counter += 1;
         match action {
             ConnectAction::Wake => Ok(helper_create_cached_node_info(self.cache)),
-            ConnectAction::WakeCold => Ok(CachedNodeInfo::new_uncached(
-                helper_create_uncached_node_info(),
-            )),
             ConnectAction::WakeFail => {
                 let err = control_plane::errors::ControlPlaneError::Message(Box::new(
                     ControlPlaneErrorMessage {
@@ -561,8 +517,8 @@ impl TestControlPlaneClient for TestConnectMechanism {
     }
 }
 
-fn helper_create_uncached_node_info() -> NodeInfo {
-    NodeInfo {
+fn helper_create_cached_node_info(cache: &'static NodeInfoCache) -> CachedNodeInfo {
+    let node = NodeInfo {
         conn_info: compute::ConnectInfo {
             host: "test".into(),
             port: 5432,
@@ -576,11 +532,7 @@ fn helper_create_uncached_node_info() -> NodeInfo {
             compute_id: "compute".into(),
             cold_start_info: crate::control_plane::messages::ColdStartInfo::Warm,
         },
-    }
-}
-
-fn helper_create_cached_node_info(cache: &'static NodeInfoCache) -> CachedNodeInfo {
-    let node = helper_create_uncached_node_info();
+    };
     let (_, node2) = cache.insert_unit("key".into(), Ok(node.clone()));
     node2.map(|()| node)
 }
@@ -756,7 +708,7 @@ async fn fail_no_wake_skips_cache_invalidation() {
     let ctx = RequestContext::test();
     let mech = TestConnectMechanism::new(vec![
         ConnectAction::Wake,
-        ConnectAction::RetryNoWake,
+        ConnectAction::FailNoWake,
         ConnectAction::Connect,
     ]);
     let user = helper_create_connect_info(&mech);
@@ -802,7 +754,7 @@ async fn retry_no_wake_skips_invalidation() {
 
     let ctx = RequestContext::test();
     // Wake → RetryNoWake (retryable + NOT wakeable)
-    let mechanism = TestConnectMechanism::new(vec![Wake, RetryNoWake, Fail]);
+    let mechanism = TestConnectMechanism::new(vec![Wake, RetryNoWake]);
     let user_info = helper_create_connect_info(&mechanism);
     let cfg = config();
 
@@ -816,44 +768,3 @@ async fn retry_no_wake_skips_invalidation() {
         "invalidating stalled compute node info cache entry"
     ));
 }
-
-#[tokio::test]
-#[traced_test]
-async fn retry_no_wake_error_fast() {
-    let _ = env_logger::try_init();
-    use ConnectAction::*;
-
-    let ctx = RequestContext::test();
-    // Wake → FailNoWake (not retryable + NOT wakeable)
-    let mechanism = TestConnectMechanism::new(vec![Wake, FailNoWake]);
-    let user_info = helper_create_connect_info(&mechanism);
-    let cfg = config();
-
-    connect_to_compute(&ctx, &mechanism, &user_info, cfg.retry, &cfg)
-        .await
-        .unwrap_err();
-    mechanism.verify();
-
-    // Because FailNoWake has wakeable=false, we must NOT see invalidate_cache
-    assert!(!logs_contain(
-        "invalidating stalled compute node info cache entry"
-    ));
-}
-
-#[tokio::test]
-#[traced_test]
-async fn retry_cold_wake_skips_invalidation() {
-    let _ = env_logger::try_init();
-    use ConnectAction::*;
-
-    let ctx = RequestContext::test();
-    // WakeCold → FailNoWake (not retryable + NOT wakeable)
-    let mechanism = TestConnectMechanism::new(vec![WakeCold, Retry, Connect]);
-    let user_info = helper_create_connect_info(&mechanism);
-    let cfg = config();
-
-    connect_to_compute(&ctx, &mechanism, &user_info, cfg.retry, &cfg)
-        .await
-        .unwrap();
-    mechanism.verify();
-}

proxy/src/rate_limiter/limiter.rs

@@ -139,6 +139,12 @@ impl RateBucketInfo {
         Self::new(200, Duration::from_secs(600)),
     ];
 
+    // For all the sessions will be cancel key. So this limit is essentially global proxy limit.
+    pub const DEFAULT_REDIS_SET: [Self; 2] = [
+        Self::new(100_000, Duration::from_secs(1)),
+        Self::new(50_000, Duration::from_secs(10)),
+    ];
+
     pub fn rps(&self) -> f64 {
         (self.max_rpi as f64) / self.interval.as_secs_f64()
     }

proxy/src/redis/kv_ops.rs

@@ -5,9 +5,11 @@ use redis::aio::ConnectionLike;
 use redis::{Cmd, FromRedisValue, Pipeline, RedisResult};
 
 use super::connection_with_credentials_provider::ConnectionWithCredentialsProvider;
+use crate::rate_limiter::{GlobalRateLimiter, RateBucketInfo};
 
 pub struct RedisKVClient {
     client: ConnectionWithCredentialsProvider,
+    limiter: GlobalRateLimiter,
 }
 
 #[allow(async_fn_in_trait)]
@@ -28,8 +30,11 @@ impl Queryable for Cmd {
 }
 
 impl RedisKVClient {
-    pub fn new(client: ConnectionWithCredentialsProvider) -> Self {
-        Self { client }
+    pub fn new(client: ConnectionWithCredentialsProvider, info: &'static [RateBucketInfo]) -> Self {
+        Self {
+            client,
+            limiter: GlobalRateLimiter::new(info.into()),
+        }
     }
 
     pub async fn try_connect(&mut self) -> anyhow::Result<()> {
@@ -44,6 +49,11 @@ impl RedisKVClient {
         &mut self,
         q: &impl Queryable,
     ) -> anyhow::Result<T> {
+        if !self.limiter.check() {
+            tracing::info!("Rate limit exceeded. Skipping query");
+            return Err(anyhow::anyhow!("Rate limit exceeded"));
+        }
+
         let e = match q.query(&mut self.client).await {
             Ok(t) => return Ok(t),
             Err(e) => e,

proxy/src/redis/notifications.rs

@@ -141,19 +141,29 @@ where
 
 struct MessageHandler<C: ProjectInfoCache + Send + Sync + 'static> {
     cache: Arc<C>,
+    region_id: String,
 }
 
 impl<C: ProjectInfoCache + Send + Sync + 'static> Clone for MessageHandler<C> {
     fn clone(&self) -> Self {
         Self {
             cache: self.cache.clone(),
+            region_id: self.region_id.clone(),
         }
     }
 }
 
 impl<C: ProjectInfoCache + Send + Sync + 'static> MessageHandler<C> {
-    pub(crate) fn new(cache: Arc<C>) -> Self {
-        Self { cache }
+    pub(crate) fn new(cache: Arc<C>, region_id: String) -> Self {
+        Self { cache, region_id }
+    }
+
+    pub(crate) async fn increment_active_listeners(&self) {
+        self.cache.increment_active_listeners().await;
+    }
+
+    pub(crate) async fn decrement_active_listeners(&self) {
+        self.cache.decrement_active_listeners().await;
     }
 
     #[tracing::instrument(skip(self, msg), fields(session_id = tracing::field::Empty))]
@@ -266,7 +276,7 @@ async fn handle_messages<C: ProjectInfoCache + Send + Sync + 'static>(
         }
         let mut conn = match try_connect(&redis).await {
             Ok(conn) => {
-                handler.cache.increment_active_listeners().await;
+                handler.increment_active_listeners().await;
                 conn
             }
             Err(e) => {
@@ -287,11 +297,11 @@ async fn handle_messages<C: ProjectInfoCache + Send + Sync + 'static>(
                 }
             }
             if cancellation_token.is_cancelled() {
-                handler.cache.decrement_active_listeners().await;
+                handler.decrement_active_listeners().await;
                 return Ok(());
             }
         }
-        handler.cache.decrement_active_listeners().await;
+        handler.decrement_active_listeners().await;
     }
 }
 
@@ -300,11 +310,12 @@ async fn handle_messages<C: ProjectInfoCache + Send + Sync + 'static>(
 pub async fn task_main<C>(
     redis: ConnectionWithCredentialsProvider,
     cache: Arc<C>,
+    region_id: String,
 ) -> anyhow::Result<Infallible>
 where
     C: ProjectInfoCache + Send + Sync + 'static,
 {
-    let handler = MessageHandler::new(cache);
+    let handler = MessageHandler::new(cache, region_id);
     // 6h - 1m.
     // There will be 1 minute overlap between two tasks. But at least we can be sure that no message is lost.
     let mut interval = tokio::time::interval(std::time::Duration::from_secs(6 * 60 * 60 - 60));

proxy/src/serverless/backend.rs

@@ -1,17 +1,12 @@
-use std::io;
-use std::net::{IpAddr, SocketAddr};
 use std::sync::Arc;
 use std::time::Duration;
 
-use async_trait::async_trait;
 use ed25519_dalek::SigningKey;
 use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
 use jose_jwk::jose_b64;
-use postgres_client::config::SslMode;
+use postgres_client::SocketConfig;
+use postgres_client::maybe_tls_stream::MaybeTlsStream;
 use rand::rngs::OsRng;
-use rustls::pki_types::{DnsName, ServerName};
-use tokio::net::{TcpStream, lookup_host};
-use tokio_rustls::TlsConnector;
 use tracing::field::display;
 use tracing::{debug, info};
 
@@ -23,21 +18,19 @@ use super::local_conn_pool::{self, EXT_NAME, EXT_SCHEMA, EXT_VERSION, LocalConnP
 use crate::auth::backend::local::StaticAuthRules;
 use crate::auth::backend::{ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo};
 use crate::auth::{self, AuthError};
+use crate::compute::{self, ComputeConnection};
 use crate::compute_ctl::{
     ComputeCtlError, ExtensionInstallRequest, Privilege, SetRoleGrantsRequest,
 };
-use crate::config::{ComputeConfig, ProxyConfig};
+use crate::config::ProxyConfig;
 use crate::context::RequestContext;
-use crate::control_plane::CachedNodeInfo;
 use crate::control_plane::client::ApiLockError;
 use crate::control_plane::errors::{GetAuthInfoError, WakeComputeError};
-use crate::control_plane::locks::ApiLocks;
 use crate::error::{ErrorKind, ReportableError, UserFacingError};
 use crate::intern::EndpointIdInt;
-use crate::proxy::connect_compute::ConnectMechanism;
-use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute};
+use crate::proxy::connect_compute::TcpMechanism;
 use crate::rate_limiter::EndpointRateLimiter;
-use crate::types::{EndpointId, Host, LOCAL_PROXY_SUFFIX};
+use crate::types::{EndpointId, LOCAL_PROXY_SUFFIX};
 
 pub(crate) struct PoolingBackend {
     pub(crate) http_conn_pool: Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
@@ -157,11 +150,6 @@ impl PoolingBackend {
     // Wake up the destination if needed. Code here is a bit involved because
     // we reuse the code from the usual proxy and we need to prepare few structures
     // that this code expects.
-    #[tracing::instrument(skip_all, fields(
-        pid = tracing::field::Empty,
-        compute_id = tracing::field::Empty,
-        conn_id = tracing::field::Empty,
-    ))]
     pub(crate) async fn connect_to_compute(
         &self,
         ctx: &RequestContext,
@@ -181,30 +169,24 @@ impl PoolingBackend {
             return Ok(client);
         }
         let conn_id = uuid::Uuid::new_v4();
-        tracing::Span::current().record("conn_id", display(conn_id));
         info!(%conn_id, "pool: opening a new connection '{conn_info}'");
         let backend = self.auth_backend.as_ref().map(|()| keys.info);
-        crate::proxy::connect_compute::connect_to_compute(
+        let connection = crate::proxy::connect_compute::connect_to_compute(
             ctx,
-            &TokioMechanism {
-                conn_id,
-                conn_info,
-                pool: self.pool.clone(),
+            &TcpMechanism {
                 locks: &self.config.connect_compute_locks,
-                keys: keys.keys,
+                direct: false,
             },
             &backend,
             self.config.wake_compute_retry_config,
             &self.config.connect_to_compute,
         )
-        .await
+        .await?;
+
+        authenticate(ctx, &self.pool, &conn_info, keys.keys, connection, conn_id).await
     }
 
     // Wake up the destination if needed
-    #[tracing::instrument(skip_all, fields(
-        compute_id = tracing::field::Empty,
-        conn_id = tracing::field::Empty,
-    ))]
     pub(crate) async fn connect_to_local_proxy(
         &self,
         ctx: &RequestContext,
@@ -216,7 +198,6 @@ impl PoolingBackend {
         }
 
         let conn_id = uuid::Uuid::new_v4();
-        tracing::Span::current().record("conn_id", display(conn_id));
         debug!(%conn_id, "pool: opening a new connection '{conn_info}'");
         let backend = self.auth_backend.as_ref().map(|()| ComputeUserInfo {
             user: conn_info.user_info.user.clone(),
@@ -226,19 +207,19 @@ impl PoolingBackend {
             )),
             options: conn_info.user_info.options.clone(),
         });
-        crate::proxy::connect_compute::connect_to_compute(
+        let connection = crate::proxy::connect_compute::connect_to_compute(
             ctx,
-            &HyperMechanism {
-                conn_id,
-                conn_info,
-                pool: self.http_conn_pool.clone(),
+            &TcpMechanism {
                 locks: &self.config.connect_compute_locks,
+                direct: true,
             },
             &backend,
             self.config.wake_compute_retry_config,
             &self.config.connect_to_compute,
         )
-        .await
+        .await?;
+
+        h2handshake(ctx, &self.http_conn_pool, &conn_info, connection, conn_id).await
     }
 
     /// Connect to postgres over localhost.
@@ -248,10 +229,6 @@ impl PoolingBackend {
     /// # Panics
     ///
     /// Panics if called with a non-local_proxy backend.
-    #[tracing::instrument(skip_all, fields(
-        pid = tracing::field::Empty,
-        conn_id = tracing::field::Empty,
-    ))]
     pub(crate) async fn connect_to_local_postgres(
         &self,
         ctx: &RequestContext,
@@ -373,6 +350,8 @@ fn create_random_jwk() -> (SigningKey, jose_jwk::Key) {
 pub(crate) enum HttpConnError {
     #[error("pooled connection closed at inconsistent state")]
     ConnectionClosedAbruptly(#[from] tokio::sync::watch::error::SendError<uuid::Uuid>),
+    #[error("could not connect to compute")]
+    ConnectError(#[from] compute::ConnectionError),
     #[error("could not connect to postgres in compute")]
     PostgresConnectionError(#[from] postgres_client::Error),
     #[error("could not connect to local-proxy in compute")]
@@ -394,8 +373,6 @@ pub(crate) enum HttpConnError {
 
 #[derive(Debug, thiserror::Error)]
 pub(crate) enum LocalProxyConnError {
-    #[error("error with connection to local-proxy")]
-    Io(#[source] std::io::Error),
     #[error("could not establish h2 connection")]
     H2(#[from] hyper::Error),
 }
@@ -403,6 +380,7 @@ pub(crate) enum LocalProxyConnError {
 impl ReportableError for HttpConnError {
     fn get_error_kind(&self) -> ErrorKind {
         match self {
+            HttpConnError::ConnectError(_) => ErrorKind::Compute,
             HttpConnError::ConnectionClosedAbruptly(_) => ErrorKind::Compute,
             HttpConnError::PostgresConnectionError(p) => p.get_error_kind(),
             HttpConnError::LocalProxyConnectionError(_) => ErrorKind::Compute,
@@ -419,6 +397,7 @@ impl ReportableError for HttpConnError {
 impl UserFacingError for HttpConnError {
     fn to_string_client(&self) -> String {
         match self {
+            HttpConnError::ConnectError(p) => p.to_string_client(),
             HttpConnError::ConnectionClosedAbruptly(_) => self.to_string(),
             HttpConnError::PostgresConnectionError(p) => p.to_string(),
             HttpConnError::LocalProxyConnectionError(p) => p.to_string(),
@@ -434,36 +413,9 @@ impl UserFacingError for HttpConnError {
     }
 }
 
-impl CouldRetry for HttpConnError {
-    fn could_retry(&self) -> bool {
-        match self {
-            HttpConnError::PostgresConnectionError(e) => e.could_retry(),
-            HttpConnError::LocalProxyConnectionError(e) => e.could_retry(),
-            HttpConnError::ComputeCtl(_) => false,
-            HttpConnError::ConnectionClosedAbruptly(_) => false,
-            HttpConnError::JwtPayloadError(_) => false,
-            HttpConnError::GetAuthInfo(_) => false,
-            HttpConnError::AuthError(_) => false,
-            HttpConnError::WakeCompute(_) => false,
-            HttpConnError::TooManyConnectionAttempts(_) => false,
-        }
-    }
-}
-impl ShouldRetryWakeCompute for HttpConnError {
-    fn should_retry_wake_compute(&self) -> bool {
-        match self {
-            HttpConnError::PostgresConnectionError(e) => e.should_retry_wake_compute(),
-            // we never checked cache validity
-            HttpConnError::TooManyConnectionAttempts(_) => false,
-            _ => true,
-        }
-    }
-}
-
 impl ReportableError for LocalProxyConnError {
     fn get_error_kind(&self) -> ErrorKind {
         match self {
-            LocalProxyConnError::Io(_) => ErrorKind::Compute,
             LocalProxyConnError::H2(_) => ErrorKind::Compute,
         }
     }
@@ -475,208 +427,106 @@ impl UserFacingError for LocalProxyConnError {
     }
 }
 
-impl CouldRetry for LocalProxyConnError {
-    fn could_retry(&self) -> bool {
-        match self {
-            LocalProxyConnError::Io(_) => false,
-            LocalProxyConnError::H2(_) => false,
-        }
-    }
-}
-impl ShouldRetryWakeCompute for LocalProxyConnError {
-    fn should_retry_wake_compute(&self) -> bool {
-        match self {
-            LocalProxyConnError::Io(_) => false,
-            LocalProxyConnError::H2(_) => false,
-        }
-    }
-}
-
-struct TokioMechanism {
-    pool: Arc<GlobalConnPool<postgres_client::Client, EndpointConnPool<postgres_client::Client>>>,
-    conn_info: ConnInfo,
-    conn_id: uuid::Uuid,
+async fn authenticate(
+    ctx: &RequestContext,
+    pool: &Arc<GlobalConnPool<postgres_client::Client, EndpointConnPool<postgres_client::Client>>>,
+    conn_info: &ConnInfo,
     keys: ComputeCredentialKeys,
-
-    /// connect_to_compute concurrency lock
-    locks: &'static ApiLocks<Host>,
-}
-
-#[async_trait]
-impl ConnectMechanism for TokioMechanism {
-    type Connection = Client<postgres_client::Client>;
-    type ConnectError = HttpConnError;
-    type Error = HttpConnError;
-
-    async fn connect_once(
-        &self,
-        ctx: &RequestContext,
-        node_info: &CachedNodeInfo,
-        compute_config: &ComputeConfig,
-    ) -> Result<Self::Connection, Self::ConnectError> {
-        let permit = self.locks.get_permit(&node_info.conn_info.host).await?;
-
-        let mut config = node_info.conn_info.to_postgres_client_config();
-        let config = config
-            .user(&self.conn_info.user_info.user)
-            .dbname(&self.conn_info.dbname)
-            .connect_timeout(compute_config.timeout);
-
-        if let ComputeCredentialKeys::AuthKeys(auth_keys) = self.keys {
-            config.auth_keys(auth_keys);
-        }
-
-        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let res = config.connect(compute_config).await;
-        drop(pause);
-        let (client, connection) = permit.release_result(res)?;
-
-        tracing::Span::current().record("pid", tracing::field::display(client.get_process_id()));
-        tracing::Span::current().record(
-            "compute_id",
-            tracing::field::display(&node_info.aux.compute_id),
-        );
-
-        if let Some(query_id) = ctx.get_testodrome_id() {
-            info!("latency={}, query_id={}", ctx.get_proxy_latency(), query_id);
-        }
-
-        Ok(poll_client(
-            self.pool.clone(),
-            ctx,
-            self.conn_info.clone(),
-            client,
-            connection,
-            self.conn_id,
-            node_info.aux.clone(),
-        ))
-    }
-}
-
-struct HyperMechanism {
-    pool: Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
-    conn_info: ConnInfo,
+    compute: ComputeConnection,
     conn_id: uuid::Uuid,
+) -> Result<Client<postgres_client::Client>, HttpConnError> {
+    // client config with stubbed connect info.
+    let mut config = postgres_client::Config::new(String::new(), 0);
+    config
+        .user(&conn_info.user_info.user)
+        .dbname(&conn_info.dbname);
 
-    /// connect_to_compute concurrency lock
-    locks: &'static ApiLocks<Host>,
-}
-
-#[async_trait]
-impl ConnectMechanism for HyperMechanism {
-    type Connection = http_conn_pool::Client<Send>;
-    type ConnectError = HttpConnError;
-    type Error = HttpConnError;
-
-    async fn connect_once(
-        &self,
-        ctx: &RequestContext,
-        node_info: &CachedNodeInfo,
-        config: &ComputeConfig,
-    ) -> Result<Self::Connection, Self::ConnectError> {
-        let host_addr = node_info.conn_info.host_addr;
-        let host = &node_info.conn_info.host;
-        let permit = self.locks.get_permit(host).await?;
-
-        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-
-        let tls = if node_info.conn_info.ssl_mode == SslMode::Disable {
-            None
-        } else {
-            Some(&config.tls)
-        };
-
-        let port = node_info.conn_info.port;
-        let res = connect_http2(host_addr, host, port, config.timeout, tls).await;
-        drop(pause);
-        let (client, connection) = permit.release_result(res)?;
-
-        tracing::Span::current().record(
-            "compute_id",
-            tracing::field::display(&node_info.aux.compute_id),
-        );
-
-        if let Some(query_id) = ctx.get_testodrome_id() {
-            info!("latency={}, query_id={}", ctx.get_proxy_latency(), query_id);
-        }
-
-        Ok(poll_http2_client(
-            self.pool.clone(),
-            ctx,
-            &self.conn_info,
-            client,
-            connection,
-            self.conn_id,
-            node_info.aux.clone(),
-        ))
+    if let ComputeCredentialKeys::AuthKeys(auth_keys) = keys {
+        config.auth_keys(auth_keys);
     }
+
+    let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
+    let connection = config.authenticate(compute.stream).await?;
+    drop(pause);
+
+    // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
+    info!(
+        compute_id = %compute.aux.compute_id,
+        pid = connection.process_id,
+        cold_start_info = ctx.cold_start_info().as_str(),
+        query_id = ctx.get_testodrome_id().as_deref(),
+        sslmode = ?compute.ssl_mode,
+        %conn_id,
+        "connected to compute node at {} ({}) latency={}",
+        compute.hostname,
+        compute.socket_addr,
+        ctx.get_proxy_latency(),
+    );
+
+    let (client, connection) = connection.into_managed_conn(
+        SocketConfig {
+            host_addr: Some(compute.socket_addr.ip()),
+            host: postgres_client::config::Host::Tcp(compute.hostname.to_string()),
+            port: compute.socket_addr.port(),
+            connect_timeout: None,
+        },
+        compute.ssl_mode,
+    );
+
+    Ok(poll_client(
+        pool.clone(),
+        ctx,
+        conn_info.clone(),
+        client,
+        connection,
+        conn_id,
+        compute.aux,
+    ))
 }
 
-async fn connect_http2(
-    host_addr: Option<IpAddr>,
-    host: &str,
-    port: u16,
-    timeout: Duration,
-    tls: Option<&Arc<rustls::ClientConfig>>,
-) -> Result<(http_conn_pool::Send, http_conn_pool::Connect), LocalProxyConnError> {
-    let addrs = match host_addr {
-        Some(addr) => vec![SocketAddr::new(addr, port)],
-        None => lookup_host((host, port))
-            .await
-            .map_err(LocalProxyConnError::Io)?
-            .collect(),
-    };
-    let mut last_err = None;
-
-    let mut addrs = addrs.into_iter();
-    let stream = loop {
-        let Some(addr) = addrs.next() else {
-            return Err(last_err.unwrap_or_else(|| {
-                LocalProxyConnError::Io(io::Error::new(
-                    io::ErrorKind::InvalidInput,
-                    "could not resolve any addresses",
-                ))
-            }));
-        };
-
-        match tokio::time::timeout(timeout, TcpStream::connect(addr)).await {
-            Ok(Ok(stream)) => {
-                stream.set_nodelay(true).map_err(LocalProxyConnError::Io)?;
-                break stream;
-            }
-            Ok(Err(e)) => {
-                last_err = Some(LocalProxyConnError::Io(e));
-            }
-            Err(e) => {
-                last_err = Some(LocalProxyConnError::Io(io::Error::new(
-                    io::ErrorKind::TimedOut,
-                    e,
-                )));
-            }
-        }
-    };
-
-    let stream = if let Some(tls) = tls {
-        let host = DnsName::try_from(host)
-            .map_err(io::Error::other)
-            .map_err(LocalProxyConnError::Io)?
-            .to_owned();
-        let stream = TlsConnector::from(tls.clone())
-            .connect(ServerName::DnsName(host), stream)
-            .await
-            .map_err(LocalProxyConnError::Io)?;
-        Box::pin(stream) as AsyncRW
-    } else {
-        Box::pin(stream) as AsyncRW
+async fn h2handshake(
+    ctx: &RequestContext,
+    pool: &Arc<GlobalConnPool<Send, HttpConnPool<Send>>>,
+    conn_info: &ConnInfo,
+    compute: ComputeConnection,
+    conn_id: uuid::Uuid,
+) -> Result<http_conn_pool::Client<Send>, HttpConnError> {
+    let stream = match compute.stream {
+        MaybeTlsStream::Raw(tcp) => Box::pin(tcp) as AsyncRW,
+        MaybeTlsStream::Tls(tls) => Box::into_pin(tls.0) as AsyncRW,
     };
 
+    let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
     let (client, connection) = hyper::client::conn::http2::Builder::new(TokioExecutor::new())
         .timer(TokioTimer::new())
         .keep_alive_interval(Duration::from_secs(20))
         .keep_alive_while_idle(true)
         .keep_alive_timeout(Duration::from_secs(5))
         .handshake(TokioIo::new(stream))
-        .await?;
+        .await
+        .map_err(LocalProxyConnError::H2)?;
+    drop(pause);
 
-    Ok((client, connection))
+    // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
+    info!(
+        compute_id = %compute.aux.compute_id,
+        cold_start_info = ctx.cold_start_info().as_str(),
+        query_id = ctx.get_testodrome_id().as_deref(),
+        sslmode = ?compute.ssl_mode,
+        %conn_id,
+        "connected to compute node at {} ({}) latency={}",
+        compute.hostname,
+        compute.socket_addr,
+        ctx.get_proxy_latency(),
+    );
+
+    Ok(poll_http2_client(
+        pool.clone(),
+        ctx,
+        conn_info,
+        client,
+        connection,
+        conn_id,
+        compute.aux,
+    ))
 }

proxy/src/serverless/conn_pool.rs

@@ -69,7 +69,7 @@ pub(crate) fn poll_client<C: ClientInnerExt>(
     let mut session_id = ctx.session_id();
     let (tx, mut rx) = tokio::sync::watch::channel(session_id);
 
-    let span = info_span!(parent: None, "connection", %conn_id);
+    let span = info_span!(parent: None, "connection", %conn_id, pid=client.get_process_id(), compute_id=%aux.compute_id);
     let cold_start_info = ctx.cold_start_info();
     span.in_scope(|| {
         info!(cold_start_info = cold_start_info.as_str(), %conn_info, %session_id, "new connection");

proxy/src/serverless/conn_pool_lib.rs

@@ -518,15 +518,14 @@ impl<C: ClientInnerExt> GlobalConnPool<C, EndpointConnPool<C>> {
                 info!("pool: cached connection '{conn_info}' is closed, opening a new one");
                 return Ok(None);
             }
-            tracing::Span::current()
-                .record("conn_id", tracing::field::display(client.get_conn_id()));
-            tracing::Span::current().record(
-                "pid",
-                tracing::field::display(client.inner.get_process_id()),
-            );
-            debug!(
+            info!(
+                conn_id = %client.get_conn_id(),
+                pid = client.inner.get_process_id(),
+                compute_id = &*client.aux.compute_id,
                 cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
-                "pool: reusing connection '{conn_info}'"
+                query_id = ctx.get_testodrome_id().as_deref(),
+                "reusing connection: latency={}",
+                ctx.get_proxy_latency(),
             );
 
             match client.get_data() {

proxy/src/serverless/http_conn_pool.rs

@@ -6,7 +6,7 @@ use hyper::client::conn::http2;
 use hyper_util::rt::{TokioExecutor, TokioIo};
 use parking_lot::RwLock;
 use smol_str::ToSmolStr;
-use tracing::{Instrument, debug, error, info, info_span};
+use tracing::{Instrument, error, info, info_span};
 
 use super::AsyncRW;
 use super::backend::HttpConnError;
@@ -115,7 +115,6 @@ impl<C: ClientInnerExt + Clone> Drop for HttpConnPool<C> {
 }
 
 impl<C: ClientInnerExt + Clone> GlobalConnPool<C, HttpConnPool<C>> {
-    #[expect(unused_results)]
     pub(crate) fn get(
         self: &Arc<Self>,
         ctx: &RequestContext,
@@ -132,10 +131,13 @@ impl<C: ClientInnerExt + Clone> GlobalConnPool<C, HttpConnPool<C>> {
             return result;
         };
 
-        tracing::Span::current().record("conn_id", tracing::field::display(client.conn.conn_id));
-        debug!(
+        info!(
+            conn_id = %client.conn.conn_id,
+            compute_id = &*client.conn.aux.compute_id,
             cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
-            "pool: reusing connection '{conn_info}'"
+            query_id = ctx.get_testodrome_id().as_deref(),
+            "reusing connection: latency={}",
+            ctx.get_proxy_latency(),
         );
         ctx.set_cold_start_info(ColdStartInfo::HttpPoolHit);
         ctx.success();
@@ -197,7 +199,7 @@ pub(crate) fn poll_http2_client(
     let conn_gauge = Metrics::get().proxy.db_connections.guard(ctx.protocol());
     let session_id = ctx.session_id();
 
-    let span = info_span!(parent: None, "connection", %conn_id);
+    let span = info_span!(parent: None, "connection", %conn_id, compute_id=%aux.compute_id);
     let cold_start_info = ctx.cold_start_info();
     span.in_scope(|| {
         info!(cold_start_info = cold_start_info.as_str(), %conn_info, %session_id, "new connection");
@@ -229,6 +231,8 @@ pub(crate) fn poll_http2_client(
 
     tokio::spawn(
         async move {
+            info!("new local proxy connection");
+
             let _conn_gauge = conn_gauge;
             let res = connection.await;
             match res {

proxy/src/serverless/local_conn_pool.rs

@@ -30,7 +30,7 @@ use serde_json::value::RawValue;
 use tokio::net::TcpStream;
 use tokio::time::Instant;
 use tokio_util::sync::CancellationToken;
-use tracing::{Instrument, debug, error, info, info_span, warn};
+use tracing::{Instrument, error, info, info_span, warn};
 
 use super::backend::HttpConnError;
 use super::conn_pool_lib::{
@@ -107,15 +107,13 @@ impl<C: ClientInnerExt> LocalConnPool<C> {
                 return Ok(None);
             }
 
-            tracing::Span::current()
-                .record("conn_id", tracing::field::display(client.get_conn_id()));
-            tracing::Span::current().record(
-                "pid",
-                tracing::field::display(client.inner.get_process_id()),
-            );
-            debug!(
+            info!(
+                pid = client.inner.get_process_id(),
+                conn_id = %client.get_conn_id(),
                 cold_start_info = ColdStartInfo::HttpPoolHit.as_str(),
-                "local_pool: reusing connection '{conn_info}'"
+                query_id = ctx.get_testodrome_id().as_deref(),
+                "reusing connection: latency={}",
+                ctx.get_proxy_latency(),
             );
 
             match client.get_data() {

proxy/src/serverless/mod.rs

@@ -417,7 +417,12 @@ async fn request_handler(
     if config.http_config.accept_websockets
         && framed_websockets::upgrade::is_upgrade_request(&request)
     {
-        let ctx = RequestContext::new(session_id, conn_info, crate::metrics::Protocol::Ws);
+        let ctx = RequestContext::new(
+            session_id,
+            conn_info,
+            crate::metrics::Protocol::Ws,
+            &config.region,
+        );
 
         ctx.set_user_agent(
             request
@@ -457,7 +462,12 @@ async fn request_handler(
         // Return the response so the spawned future can continue.
         Ok(response.map(|b| b.map_err(|x| match x {}).boxed()))
     } else if request.uri().path() == "/sql" && *request.method() == Method::POST {
-        let ctx = RequestContext::new(session_id, conn_info, crate::metrics::Protocol::Http);
+        let ctx = RequestContext::new(
+            session_id,
+            conn_info,
+            crate::metrics::Protocol::Http,
+            &config.region,
+        );
         let span = ctx.span();
 
         let testodrome_id = request

proxy/src/serverless/websocket.rs

@@ -17,8 +17,7 @@ use crate::config::ProxyConfig;
 use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::Metrics;
-use crate::pglb::{ClientMode, handle_connection};
-use crate::proxy::ErrorSource;
+use crate::proxy::{ClientMode, ErrorSource, handle_client};
 use crate::rate_limiter::EndpointRateLimiter;
 
 pin_project! {
@@ -143,7 +142,7 @@ pub(crate) async fn serve_websocket(
         .client_connections
         .guard(crate::metrics::Protocol::Ws);
 
-    let res = Box::pin(handle_connection(
+    let res = Box::pin(handle_client(
         config,
         auth_backend,
         &ctx,

proxy/src/tls/postgres_rustls.rs

@@ -60,7 +60,7 @@ mod private {
         }
     }
 
-    pub struct RustlsStream<S>(Box<TlsStream<S>>);
+    pub struct RustlsStream<S>(pub Box<TlsStream<S>>);
 
     impl<S> postgres_client::tls::TlsStream for RustlsStream<S>
     where

rust-toolchain.toml

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

safekeeper/src/handler.rs

@@ -220,7 +220,7 @@ impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
                                     stripe_size: ShardStripeSize(stripe_size),
                                 };
                                 self.shard =
-                                    Some(ShardIdentity::from_params(ShardNumber(number), params));
+                                    Some(ShardIdentity::from_params(ShardNumber(number), &params));
                             }
                             _ => {
                                 return Err(QueryError::Other(anyhow::anyhow!(

safekeeper/src/pull_timeline.rs

@@ -1,6 +1,5 @@
 use std::cmp::min;
 use std::io::{self, ErrorKind};
-use std::ops::RangeInclusive;
 use std::sync::Arc;
 
 use anyhow::{Context, Result, anyhow, bail};
@@ -35,7 +34,7 @@ use crate::control_file::CONTROL_FILE_NAME;
 use crate::state::{EvictionState, TimelinePersistentState};
 use crate::timeline::{Timeline, TimelineError, WalResidentTimeline};
 use crate::timelines_global_map::{create_temp_timeline_dir, validate_temp_timeline};
-use crate::wal_storage::{open_wal_file, wal_file_paths};
+use crate::wal_storage::open_wal_file;
 use crate::{GlobalTimelines, debug_dump, wal_backup};
 
 /// Stream tar archive of timeline to tx.
@@ -96,8 +95,8 @@ pub async fn stream_snapshot(
 
 /// State needed while streaming the snapshot.
 pub struct SnapshotContext {
-    /// The interval of segment numbers. If None, the timeline hasn't had writes yet, so only send the control file
-    pub from_to_segno: Option<RangeInclusive<XLogSegNo>>,
+    pub from_segno: XLogSegNo, // including
+    pub upto_segno: XLogSegNo, // including
     pub term: Term,
     pub last_log_term: Term,
     pub flush_lsn: Lsn,
@@ -175,35 +174,23 @@ pub async fn stream_snapshot_resident_guts(
         .await?;
     pausable_failpoint!("sk-snapshot-after-list-pausable");
 
-    if let Some(from_to_segno) = &bctx.from_to_segno {
-        let tli_dir = tli.get_timeline_dir();
-        info!(
-            "sending {} segments [{:#X}-{:#X}], term={}, last_log_term={}, flush_lsn={}",
-            from_to_segno.end() - from_to_segno.start() + 1,
-            from_to_segno.start(),
-            from_to_segno.end(),
-            bctx.term,
-            bctx.last_log_term,
-            bctx.flush_lsn,
-        );
-        for segno in from_to_segno.clone() {
-            let Some((mut sf, is_partial)) =
-                open_wal_file(&tli_dir, segno, bctx.wal_seg_size).await?
-            else {
-                // File is not found
-                let (wal_file_path, _wal_file_partial_path) =
-                    wal_file_paths(&tli_dir, segno, bctx.wal_seg_size);
-                tracing::warn!("couldn't find WAL segment file {wal_file_path}");
-                bail!("couldn't find WAL segment file {wal_file_path}")
-            };
-            let mut wal_file_name = XLogFileName(PG_TLI, segno, bctx.wal_seg_size);
-            if is_partial {
-                wal_file_name.push_str(".partial");
-            }
-            ar.append_file(&wal_file_name, &mut sf).await?;
+    let tli_dir = tli.get_timeline_dir();
+    info!(
+        "sending {} segments [{:#X}-{:#X}], term={}, last_log_term={}, flush_lsn={}",
+        bctx.upto_segno - bctx.from_segno + 1,
+        bctx.from_segno,
+        bctx.upto_segno,
+        bctx.term,
+        bctx.last_log_term,
+        bctx.flush_lsn,
+    );
+    for segno in bctx.from_segno..=bctx.upto_segno {
+        let (mut sf, is_partial) = open_wal_file(&tli_dir, segno, bctx.wal_seg_size).await?;
+        let mut wal_file_name = XLogFileName(PG_TLI, segno, bctx.wal_seg_size);
+        if is_partial {
+            wal_file_name.push_str(".partial");
         }
-    } else {
-        info!("Not including any segments into the snapshot");
+        ar.append_file(&wal_file_name, &mut sf).await?;
     }
 
     // Do the term check before ar.finish to make archive corrupted in case of
@@ -351,26 +338,19 @@ impl WalResidentTimeline {
         // removed further than `backup_lsn`. Since we're holding shared_state
         // lock and setting `wal_removal_on_hold` later, it guarantees that WAL
         // won't be removed until we're done.
-        let timeline_state = shared_state.sk.state();
         let from_lsn = min(
-            timeline_state.remote_consistent_lsn,
-            timeline_state.backup_lsn,
-        );
-        let flush_lsn = shared_state.sk.flush_lsn();
-        let (send_segments, msg) = if from_lsn == Lsn::INVALID {
-            (false, "snapshot is called on uninitialized timeline")
-        } else {
-            (true, "timeline is initialized")
-        };
-        tracing::info!(
-            remote_consistent_lsn=%timeline_state.remote_consistent_lsn,
-            backup_lsn=%timeline_state.backup_lsn,
-            %flush_lsn,
-            "{msg}"
+            shared_state.sk.state().remote_consistent_lsn,
+            shared_state.sk.state().backup_lsn,
         );
+        if from_lsn == Lsn::INVALID {
+            // this is possible if snapshot is called before handling first
+            // elected message
+            bail!("snapshot is called on uninitialized timeline");
+        }
         let from_segno = from_lsn.segment_number(wal_seg_size);
         let term = shared_state.sk.state().acceptor_state.term;
         let last_log_term = shared_state.sk.last_log_term();
+        let flush_lsn = shared_state.sk.flush_lsn();
         let upto_segno = flush_lsn.segment_number(wal_seg_size);
         // have some limit on max number of segments as a sanity check
         const MAX_ALLOWED_SEGS: u64 = 1000;
@@ -396,9 +376,9 @@ impl WalResidentTimeline {
         drop(shared_state);
 
         let tli_copy = self.wal_residence_guard().await?;
-        let from_to_segno = send_segments.then_some(from_segno..=upto_segno);
         let bctx = SnapshotContext {
-            from_to_segno,
+            from_segno,
+            upto_segno,
             term,
             last_log_term,
             flush_lsn,