Proxy release 2025-06-03 06:01 UTC

Proxy release 2025-03-06

Cargo.lock

@@ -1445,7 +1445,6 @@ dependencies = [
  "regex",
  "reqwest",
  "safekeeper_api",
- "safekeeper_client",
  "scopeguard",
  "serde",
  "serde_json",
@@ -4238,7 +4237,6 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "async-trait",
- "bytes",
  "camino",
  "clap",
  "futures",
@@ -4471,6 +4469,7 @@ dependencies = [
  "pageserver_api",
  "postgres_ffi",
  "prost 0.13.5",
+ "smallvec",
  "thiserror 1.0.69",
  "tonic 0.13.1",
  "tonic-build",

Dockerfile

@@ -110,19 +110,6 @@ RUN set -e \
 	# System postgres for use with client libraries (e.g. in storage controller)
         postgresql-15 \
         openssl \
-        unzip \
-        curl \
-    && ARCH=$(uname -m) \
-    && if [ "$ARCH" = "x86_64" ]; then \
-        curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"; \
-    elif [ "$ARCH" = "aarch64" ]; then \
-        curl "https://awscli.amazonaws.com/awscli-exe-linux-aarch64.zip" -o "awscliv2.zip"; \
-    else \
-        echo "Unsupported architecture: $ARCH" && exit 1; \
-    fi \
-    && unzip awscliv2.zip \
-    && ./aws/install \
-    && rm -rf aws awscliv2.zip \
     && rm -f /etc/apt/apt.conf.d/80-retries \
     && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* \
     && useradd -d /data neon \

compute/compute-node.Dockerfile

@@ -602,7 +602,7 @@ RUN case "${PG_VERSION:?}" in \
         ;; \
     esac && \
 	wget https://github.com/knizhnik/online_advisor/archive/refs/tags/1.0.tar.gz -O online_advisor.tar.gz && \
-    echo "37dcadf8f7cc8d6cc1f8831276ee245b44f1b0274f09e511e47a67738ba9ed0f online_advisor.tar.gz" | sha256sum --check && \
+    echo "059b7d9e5a90013a58bdd22e9505b88406ce05790675eb2d8434e5b215652d54 online_advisor.tar.gz" | sha256sum --check && \
     mkdir online_advisor-src && cd online_advisor-src && tar xzf ../online_advisor.tar.gz --strip-components=1 -C .
 
 FROM pg-build AS online_advisor-build
@@ -1180,14 +1180,14 @@ RUN cd exts/rag && \
 RUN cd exts/rag_bge_small_en_v15 && \
     sed -i 's/pgrx = "0.14.1"/pgrx = { version = "0.14.1", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     ORT_LIB_LOCATION=/ext-src/onnxruntime-src/build/Linux \
-        REMOTE_ONNX_URL=http://pg-ext-s3-gateway.pg-ext-s3-gateway.svc.cluster.local/pgrag-data/bge_small_en_v15.onnx \
+        REMOTE_ONNX_URL=http://pg-ext-s3-gateway/pgrag-data/bge_small_en_v15.onnx \
         cargo pgrx install --release --features remote_onnx && \
     echo "trusted = true" >> /usr/local/pgsql/share/extension/rag_bge_small_en_v15.control
 
 RUN cd exts/rag_jina_reranker_v1_tiny_en && \
     sed -i 's/pgrx = "0.14.1"/pgrx = { version = "0.14.1", features = [ "unsafe-postgres" ] }/g' Cargo.toml && \
     ORT_LIB_LOCATION=/ext-src/onnxruntime-src/build/Linux \
-        REMOTE_ONNX_URL=http://pg-ext-s3-gateway.pg-ext-s3-gateway.svc.cluster.local/pgrag-data/jina_reranker_v1_tiny_en.onnx \
+        REMOTE_ONNX_URL=http://pg-ext-s3-gateway/pgrag-data/jina_reranker_v1_tiny_en.onnx \
         cargo pgrx install --release --features remote_onnx && \
     echo "trusted = true" >> /usr/local/pgsql/share/extension/rag_jina_reranker_v1_tiny_en.control
 
@@ -1886,7 +1886,6 @@ COPY compute/patches/pg_repack.patch /ext-src
 RUN cd /ext-src/pg_repack-src && patch -p1 </ext-src/pg_repack.patch && rm -f /ext-src/pg_repack.patch
 
 COPY --chmod=755 docker-compose/run-tests.sh /run-tests.sh
-RUN echo /usr/local/pgsql/lib > /etc/ld.so.conf.d/00-neon.conf && /sbin/ldconfig
 RUN apt-get update && apt-get install -y libtap-parser-sourcehandler-pgtap-perl jq \
    && apt clean && rm -rf /ext-src/*.tar.gz /ext-src/*.patch /var/lib/apt/lists/*
 ENV PATH=/usr/local/pgsql/bin:$PATH

compute/manifest.yaml

@@ -1,121 +0,0 @@
-pg_settings:
-  # Common settings for primaries and replicas of all versions.
-  common:
-    # Check for client disconnection every 1 minute. By default, Postgres will detect the
-    # loss of the connection only at the next interaction with the socket, when it waits
-    # for, receives or sends data, so it will likely waste resources till the end of the
-    # query execution. There should be no drawbacks in setting this for everyone, so enable
-    # it by default. If anyone will complain, we can allow editing it.
-    # https://www.postgresql.org/docs/16/runtime-config-connection.html#GUC-CLIENT-CONNECTION-CHECK-INTERVAL
-    client_connection_check_interval: "60000" # 1 minute
-    # ---- IO ---- 
-    effective_io_concurrency: "20"
-    maintenance_io_concurrency: "100"
-    fsync: "off"
-    hot_standby: "off"
-    # We allow users to change this if needed, but by default we
-    # just don't want to see long-lasting idle transactions, as they
-    # prevent activity monitor from suspending projects.
-    idle_in_transaction_session_timeout: "300000" # 5 minutes
-    listen_addresses: "*"
-    # --- LOGGING ---- helps investigations
-    log_connections: "on"
-    log_disconnections: "on"
-    # 1GB, unit is KB
-    log_temp_files: "1048576"
-    # Disable dumping customer data to logs, both to increase data privacy
-    # and to reduce the amount the logs.
-    log_error_verbosity: "terse"
-    log_min_error_statement: "panic"
-    max_connections: "100"
-    # --- WAL ---
-    # - flush lag is the max amount of WAL that has been generated but not yet stored
-    # to disk in the page server. A smaller value means less delay after a pageserver
-    # restart, but if you set it too small you might again need to slow down writes if the
-    # pageserver cannot flush incoming WAL to disk fast enough. This must be larger
-    # than the pageserver's checkpoint interval, currently 1 GB! Otherwise you get a
-    # a deadlock where the compute node refuses to generate more WAL before the
-    # old WAL has been uploaded to S3, but the pageserver is waiting for more WAL
-    # to be generated before it is uploaded to S3.
-    max_replication_flush_lag: "10GB"
-    max_replication_slots: "10"
-    # Backpressure configuration:
-    # - write lag is the max amount of WAL that has been generated by Postgres but not yet
-    # processed by the page server. Making this smaller reduces the worst case latency
-    # of a GetPage request, if you request a page that was recently modified. On the other
-    # hand, if this is too small, the compute node might need to wait on a write if there is a
-    # hiccup in the network or page server so that the page server has temporarily fallen
-    # behind.
-    #
-    # Previously it was set to 500 MB, but it caused compute being unresponsive under load
-    # https://github.com/neondatabase/neon/issues/2028
-    max_replication_write_lag: "500MB"
-    max_wal_senders: "10"
-    # A Postgres checkpoint is cheap in storage, as doesn't involve any significant amount
-    # of real I/O. Only the SLRU buffers and some other small files are flushed to disk.
-    # However, as long as we have full_page_writes=on, page updates after a checkpoint
-    # include full-page images which bloats the WAL. So may want to bump max_wal_size to
-    # reduce the WAL bloating, but at the same it will increase pg_wal directory size on
-    # compute and can lead to out of disk error on k8s nodes.
-    max_wal_size: "1024"
-    wal_keep_size: "0"
-    wal_level: "replica"
-    # Reduce amount of WAL generated by default.
-    wal_log_hints: "off"
-    # - without wal_sender_timeout set we don't get feedback messages,
-    # required for backpressure.
-    wal_sender_timeout: "10000"
-    # We have some experimental extensions, which we don't want users to install unconsciously.
-    # To install them, users would need to set the `neon.allow_unstable_extensions` setting.
-    # There are two of them currently:
-    # - `pgrag` - https://github.com/neondatabase-labs/pgrag - extension is actually called just `rag`,
-    #                                                          and two dependencies:
-    #                                                          - `rag_bge_small_en_v15`
-    #                                                          - `rag_jina_reranker_v1_tiny_en`
-    # - `pg_mooncake` - https://github.com/Mooncake-Labs/pg_mooncake/  
-    neon.unstable_extensions: "rag,rag_bge_small_en_v15,rag_jina_reranker_v1_tiny_en,pg_mooncake,anon"
-    neon.protocol_version: "3"
-    password_encryption: "scram-sha-256"
-    # This is important to prevent Postgres from trying to perform
-    # a local WAL redo after backend crash. It should exit and let
-    # the systemd or k8s to do a fresh startup with compute_ctl.
-    restart_after_crash: "off"
-    # By default 3. We have the following persistent connections in the VM:
-    # * compute_activity_monitor (from compute_ctl)
-    # * postgres-exporter (metrics collector; it has 2 connections)
-    # * sql_exporter (metrics collector; we have 2 instances [1 for us & users; 1 for autoscaling])
-    # * vm-monitor (to query & change file cache size)
-    # i.e. total of 6. Let's reserve 7, so there's still at least one left over.
-    superuser_reserved_connections: "7"
-    synchronous_standby_names: "walproposer"
-
-  replica:
-    hot_standby: "on"
-
-  per_version:
-    17:
-      common:
-        # PostgreSQL 17 has a new IO system called "read stream", which can combine IOs up to some
-        # size. It still has some issues with readahead, though, so we default to disabled/
-        # "no combining of IOs" to make sure we get the maximum prefetch depth.
-        # See also: https://github.com/neondatabase/neon/pull/9860
-        io_combine_limit: "1"
-      replica:
-        # prefetching of blocks referenced in WAL doesn't make sense for us
-        # Neon hot standby ignores pages that are not in the shared_buffers
-        recovery_prefetch: "off"
-    16:
-      common:
-      replica:
-        # prefetching of blocks referenced in WAL doesn't make sense for us
-        # Neon hot standby ignores pages that are not in the shared_buffers
-        recovery_prefetch: "off"
-    15:
-      common:
-      replica:
-        # prefetching of blocks referenced in WAL doesn't make sense for us
-        # Neon hot standby ignores pages that are not in the shared_buffers
-        recovery_prefetch: "off"
-    14:
-      common:
-      replica:

compute_tools/src/bin/compute_ctl.rs

@@ -40,7 +40,7 @@ use std::sync::mpsc;
 use std::thread;
 use std::time::Duration;
 
-use anyhow::{Context, Result, bail};
+use anyhow::{Context, Result};
 use clap::Parser;
 use compute_api::responses::ComputeConfig;
 use compute_tools::compute::{
@@ -57,14 +57,14 @@ use tracing::{error, info};
 use url::Url;
 use utils::failpoint_support;
 
-#[derive(Debug, Parser)]
+#[derive(Parser)]
 #[command(rename_all = "kebab-case")]
 struct Cli {
     #[arg(short = 'b', long, default_value = "postgres", env = "POSTGRES_PATH")]
     pub pgbin: String,
 
     /// The base URL for the remote extension storage proxy gateway.
-    #[arg(short = 'r', long, value_parser = Self::parse_remote_ext_base_url)]
+    #[arg(short = 'r', long)]
     pub remote_ext_base_url: Option<Url>,
 
     /// The port to bind the external listening HTTP server to. Clients running
@@ -126,25 +126,6 @@ struct Cli {
     pub installed_extensions_collection_interval: u64,
 }
 
-impl Cli {
-    /// Parse a URL from an argument. By default, this isn't necessary, but we
-    /// want to do some sanity checking.
-    fn parse_remote_ext_base_url(value: &str) -> Result<Url> {
-        // Remove extra trailing slashes, and add one. We use Url::join() later
-        // when downloading remote extensions. If the base URL is something like
-        // http://example.com/pg-ext-s3-gateway, and join() is called with
-        // something like "xyz", the resulting URL is http://example.com/xyz.
-        let value = value.trim_end_matches('/').to_owned() + "/";
-        let url = Url::parse(&value)?;
-
-        if url.query_pairs().count() != 0 {
-            bail!("parameters detected in remote extensions base URL")
-        }
-
-        Ok(url)
-    }
-}
-
 fn main() -> Result<()> {
     let cli = Cli::parse();
 
@@ -271,8 +252,7 @@ fn handle_exit_signal(sig: i32) {
 
 #[cfg(test)]
 mod test {
-    use clap::{CommandFactory, Parser};
-    use url::Url;
+    use clap::CommandFactory;
 
     use super::Cli;
 
@@ -280,43 +260,4 @@ mod test {
     fn verify_cli() {
         Cli::command().debug_assert()
     }
-
-    #[test]
-    fn verify_remote_ext_base_url() {
-        let cli = Cli::parse_from([
-            "compute_ctl",
-            "--pgdata=test",
-            "--connstr=test",
-            "--compute-id=test",
-            "--remote-ext-base-url",
-            "https://example.com/subpath",
-        ]);
-        assert_eq!(
-            cli.remote_ext_base_url.unwrap(),
-            Url::parse("https://example.com/subpath/").unwrap()
-        );
-
-        let cli = Cli::parse_from([
-            "compute_ctl",
-            "--pgdata=test",
-            "--connstr=test",
-            "--compute-id=test",
-            "--remote-ext-base-url",
-            "https://example.com//",
-        ]);
-        assert_eq!(
-            cli.remote_ext_base_url.unwrap(),
-            Url::parse("https://example.com").unwrap()
-        );
-
-        Cli::try_parse_from([
-            "compute_ctl",
-            "--pgdata=test",
-            "--connstr=test",
-            "--compute-id=test",
-            "--remote-ext-base-url",
-            "https://example.com?hello=world",
-        ])
-        .expect_err("URL parameters are not allowed");
-    }
 }

compute_tools/src/compute.rs

@@ -3,7 +3,7 @@ use chrono::{DateTime, Utc};
 use compute_api::privilege::Privilege;
 use compute_api::responses::{
     ComputeConfig, ComputeCtlConfig, ComputeMetrics, ComputeStatus, LfcOffloadState,
-    LfcPrewarmState, TlsConfig,
+    LfcPrewarmState,
 };
 use compute_api::spec::{
     ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, ExtVersion, PgIdent,
@@ -396,7 +396,7 @@ impl ComputeNode {
         // because QEMU will already have its memory allocated from the host, and
         // the necessary binaries will already be cached.
         if cli_spec.is_none() {
-            this.prewarm_postgres_vm_memory()?;
+            this.prewarm_postgres()?;
         }
 
         // Set the up metric with Empty status before starting the HTTP server.
@@ -603,8 +603,6 @@ impl ComputeNode {
             });
         }
 
-        let tls_config = self.tls_config(&pspec.spec);
-
         // If there are any remote extensions in shared_preload_libraries, start downloading them
         if pspec.spec.remote_extensions.is_some() {
             let (this, spec) = (self.clone(), pspec.spec.clone());
@@ -661,7 +659,7 @@ impl ComputeNode {
             info!("tuning pgbouncer");
 
             let pgbouncer_settings = pgbouncer_settings.clone();
-            let tls_config = tls_config.clone();
+            let tls_config = self.compute_ctl_config.tls.clone();
 
             // Spawn a background task to do the tuning,
             // so that we don't block the main thread that starts Postgres.
@@ -680,10 +678,7 @@ impl ComputeNode {
 
             // Spawn a background task to do the configuration,
             // so that we don't block the main thread that starts Postgres.
-
-            let mut local_proxy = local_proxy.clone();
-            local_proxy.tls = tls_config.clone();
-
+            let local_proxy = local_proxy.clone();
             let _handle = tokio::spawn(async move {
                 if let Err(err) = local_proxy::configure(&local_proxy) {
                     error!("error while configuring local_proxy: {err:?}");
@@ -784,7 +779,7 @@ impl ComputeNode {
         // Spawn the extension stats background task
         self.spawn_extension_stats_task();
 
-        if pspec.spec.autoprewarm {
+        if pspec.spec.prewarm_lfc_on_startup {
             self.prewarm_lfc();
         }
         Ok(())
@@ -1210,15 +1205,13 @@ impl ComputeNode {
         let spec = &pspec.spec;
         let pgdata_path = Path::new(&self.params.pgdata);
 
-        let tls_config = self.tls_config(&pspec.spec);
-
         // Remove/create an empty pgdata directory and put configuration there.
         self.create_pgdata()?;
         config::write_postgres_conf(
             pgdata_path,
             &pspec.spec,
             self.params.internal_http_port,
-            tls_config,
+            &self.compute_ctl_config.tls,
         )?;
 
         // Syncing safekeepers is only safe with primary nodes: if a primary
@@ -1314,8 +1307,8 @@ impl ComputeNode {
     }
 
     /// Start and stop a postgres process to warm up the VM for startup.
-    pub fn prewarm_postgres_vm_memory(&self) -> Result<()> {
-        info!("prewarming VM memory");
+    pub fn prewarm_postgres(&self) -> Result<()> {
+        info!("prewarming");
 
         // Create pgdata
         let pgdata = &format!("{}.warmup", self.params.pgdata);
@@ -1357,7 +1350,7 @@ impl ComputeNode {
         kill(pm_pid, Signal::SIGQUIT)?;
         info!("sent SIGQUIT signal");
         pg.wait()?;
-        info!("done prewarming vm memory");
+        info!("done prewarming");
 
         // clean up
         let _ok = fs::remove_dir_all(pgdata);
@@ -1543,22 +1536,14 @@ impl ComputeNode {
                 .clone(),
         );
 
-        let mut tls_config = None::<TlsConfig>;
-        if spec.features.contains(&ComputeFeature::TlsExperimental) {
-            tls_config = self.compute_ctl_config.tls.clone();
-        }
-
         let max_concurrent_connections = self.max_service_connections(compute_state, &spec);
 
         // Merge-apply spec & changes to PostgreSQL state.
         self.apply_spec_sql(spec.clone(), conf.clone(), max_concurrent_connections)?;
 
         if let Some(local_proxy) = &spec.clone().local_proxy_config {
-            let mut local_proxy = local_proxy.clone();
-            local_proxy.tls = tls_config.clone();
-
             info!("configuring local_proxy");
-            local_proxy::configure(&local_proxy).context("apply_config local_proxy")?;
+            local_proxy::configure(local_proxy).context("apply_config local_proxy")?;
         }
 
         // Run migrations separately to not hold up cold starts
@@ -1610,13 +1595,11 @@ impl ComputeNode {
     pub fn reconfigure(&self) -> Result<()> {
         let spec = self.state.lock().unwrap().pspec.clone().unwrap().spec;
 
-        let tls_config = self.tls_config(&spec);
-
         if let Some(ref pgbouncer_settings) = spec.pgbouncer_settings {
             info!("tuning pgbouncer");
 
             let pgbouncer_settings = pgbouncer_settings.clone();
-            let tls_config = tls_config.clone();
+            let tls_config = self.compute_ctl_config.tls.clone();
 
             // Spawn a background task to do the tuning,
             // so that we don't block the main thread that starts Postgres.
@@ -1634,7 +1617,7 @@ impl ComputeNode {
             // Spawn a background task to do the configuration,
             // so that we don't block the main thread that starts Postgres.
             let mut local_proxy = local_proxy.clone();
-            local_proxy.tls = tls_config.clone();
+            local_proxy.tls = self.compute_ctl_config.tls.clone();
             tokio::spawn(async move {
                 if let Err(err) = local_proxy::configure(&local_proxy) {
                     error!("error while configuring local_proxy: {err:?}");
@@ -1652,7 +1635,7 @@ impl ComputeNode {
             pgdata_path,
             &spec,
             self.params.internal_http_port,
-            tls_config,
+            &self.compute_ctl_config.tls,
         )?;
 
         if !spec.skip_pg_catalog_updates {
@@ -1772,14 +1755,6 @@ impl ComputeNode {
         }
     }
 
-    pub fn tls_config(&self, spec: &ComputeSpec) -> &Option<TlsConfig> {
-        if spec.features.contains(&ComputeFeature::TlsExperimental) {
-            &self.compute_ctl_config.tls
-        } else {
-            &None::<TlsConfig>
-        }
-    }
-
     /// Update the `last_active` in the shared state, but ensure that it's a more recent one.
     pub fn update_last_active(&self, last_active: Option<DateTime<Utc>>) {
         let mut state = self.state.lock().unwrap();

compute_tools/src/extension_server.rs

@@ -166,7 +166,7 @@ pub async fn download_extension(
 
     // TODO add retry logic
     let download_buffer =
-        match download_extension_tar(remote_ext_base_url, &ext_path.to_string()).await {
+        match download_extension_tar(remote_ext_base_url.as_str(), &ext_path.to_string()).await {
             Ok(buffer) => buffer,
             Err(error_message) => {
                 return Err(anyhow::anyhow!(
@@ -271,14 +271,10 @@ pub fn create_control_files(remote_extensions: &RemoteExtSpec, pgbin: &str) {
 }
 
 // Do request to extension storage proxy, e.g.,
-// curl http://pg-ext-s3-gateway.pg-ext-s3-gateway.svc.cluster.local/latest/v15/extensions/anon.tar.zst
+// curl http://pg-ext-s3-gateway/latest/v15/extensions/anon.tar.zst
 // using HTTP GET and return the response body as bytes.
-async fn download_extension_tar(remote_ext_base_url: &Url, ext_path: &str) -> Result<Bytes> {
-    let uri = remote_ext_base_url.join(ext_path).with_context(|| {
-        format!(
-            "failed to create the remote extension URI for {ext_path} using {remote_ext_base_url}"
-        )
-    })?;
+async fn download_extension_tar(remote_ext_base_url: &str, ext_path: &str) -> Result<Bytes> {
+    let uri = format!("{}/{}", remote_ext_base_url, ext_path);
     let filename = Path::new(ext_path)
         .file_name()
         .unwrap_or_else(|| std::ffi::OsStr::new("unknown"))
@@ -288,7 +284,7 @@ async fn download_extension_tar(remote_ext_base_url: &Url, ext_path: &str) -> Re
 
     info!("Downloading extension file '{}' from uri {}", filename, uri);
 
-    match do_extension_server_request(uri).await {
+    match do_extension_server_request(&uri).await {
         Ok(resp) => {
             info!("Successfully downloaded remote extension data {}", ext_path);
             REMOTE_EXT_REQUESTS_TOTAL
@@ -307,7 +303,7 @@ async fn download_extension_tar(remote_ext_base_url: &Url, ext_path: &str) -> Re
 
 // Do a single remote extensions server request.
 // Return result or (error message + stringified status code) in case of any failures.
-async fn do_extension_server_request(uri: Url) -> Result<Bytes, (String, String)> {
+async fn do_extension_server_request(uri: &str) -> Result<Bytes, (String, String)> {
     let resp = reqwest::get(uri).await.map_err(|e| {
         (
             format!(

compute_tools/src/http/mod.rs

@@ -48,9 +48,11 @@ impl JsonResponse {
 
     /// Create an error response related to the compute being in an invalid state
     pub(self) fn invalid_status(status: ComputeStatus) -> Response {
-        Self::error(
+        Self::create_response(
             StatusCode::PRECONDITION_FAILED,
-            format!("invalid compute status: {status}"),
+            &GenericAPIError {
+                error: format!("invalid compute status: {status}"),
+            },
         )
     }
 }

compute_tools/src/http/routes/configure.rs

@@ -22,7 +22,7 @@ pub(in crate::http) async fn configure(
     State(compute): State<Arc<ComputeNode>>,
     request: Json<ConfigurationRequest>,
 ) -> Response {
-    let pspec = match ParsedSpec::try_from(request.0.spec) {
+    let pspec = match ParsedSpec::try_from(request.spec.clone()) {
         Ok(p) => p,
         Err(e) => return JsonResponse::error(StatusCode::BAD_REQUEST, e),
     };

compute_tools/src/monitor.rs

@@ -13,12 +13,6 @@ use crate::metrics::{PG_CURR_DOWNTIME_MS, PG_TOTAL_DOWNTIME_MS};
 
 const MONITOR_CHECK_INTERVAL: Duration = Duration::from_millis(500);
 
-/// Struct to store runtime state of the compute monitor thread.
-/// In theory, this could be a part of `Compute`, but i)
-/// this state is expected to be accessed only by single thread,
-/// so we don't need to care about locking; ii) `Compute` is
-/// already quite big. Thus, it seems to be a good idea to keep
-/// all the activity/health monitoring parts here.
 struct ComputeMonitor {
     compute: Arc<ComputeNode>,
 
@@ -76,36 +70,12 @@ impl ComputeMonitor {
         )
     }
 
-    /// Check if compute is in some terminal or soon-to-be-terminal
-    /// state, then return `true`, signalling the caller that it
-    /// should exit gracefully. Otherwise, return `false`.
-    fn check_interrupts(&mut self) -> bool {
-        let compute_status = self.compute.get_status();
-        if matches!(
-            compute_status,
-            ComputeStatus::Terminated | ComputeStatus::TerminationPending | ComputeStatus::Failed
-        ) {
-            info!(
-                "compute is in {} status, stopping compute monitor",
-                compute_status
-            );
-            return true;
-        }
-
-        false
-    }
-
     /// Spin in a loop and figure out the last activity time in the Postgres.
-    /// Then update it in the shared state. This function currently never
-    /// errors out explicitly, but there is a graceful termination path.
-    /// Every time we receive an error trying to check Postgres, we use
-    /// [`ComputeMonitor::check_interrupts()`] because it could be that
-    /// compute is being terminated already, then we can exit gracefully
-    /// to not produce errors' noise in the log.
+    /// Then update it in the shared state. This function never errors out.
     /// NB: the only expected panic is at `Mutex` unwrap(), all other errors
     /// should be handled gracefully.
     #[instrument(skip_all)]
-    pub fn run(&mut self) -> anyhow::Result<()> {
+    pub fn run(&mut self) {
         // Suppose that `connstr` doesn't change
         let connstr = self.compute.params.connstr.clone();
         let conf = self
@@ -123,10 +93,6 @@ impl ComputeMonitor {
         info!("starting compute monitor for {}", connstr);
 
         loop {
-            if self.check_interrupts() {
-                break;
-            }
-
             match &mut client {
                 Ok(cli) => {
                     if cli.is_closed() {
@@ -134,10 +100,6 @@ impl ComputeMonitor {
                             downtime_info = self.downtime_info(),
                             "connection to Postgres is closed, trying to reconnect"
                         );
-                        if self.check_interrupts() {
-                            break;
-                        }
-
                         self.report_down();
 
                         // Connection is closed, reconnect and try again.
@@ -149,19 +111,15 @@ impl ComputeMonitor {
                                 self.compute.update_last_active(self.last_active);
                             }
                             Err(e) => {
-                                error!(
-                                    downtime_info = self.downtime_info(),
-                                    "could not check Postgres: {}", e
-                                );
-                                if self.check_interrupts() {
-                                    break;
-                                }
-
                                 // Although we have many places where we can return errors in `check()`,
                                 // normally it shouldn't happen. I.e., we will likely return error if
                                 // connection got broken, query timed out, Postgres returned invalid data, etc.
                                 // In all such cases it's suspicious, so let's report this as downtime.
                                 self.report_down();
+                                error!(
+                                    downtime_info = self.downtime_info(),
+                                    "could not check Postgres: {}", e
+                                );
 
                                 // Reconnect to Postgres just in case. During tests, I noticed
                                 // that queries in `check()` can fail with `connection closed`,
@@ -178,10 +136,6 @@ impl ComputeMonitor {
                         downtime_info = self.downtime_info(),
                         "could not connect to Postgres: {}, retrying", e
                     );
-                    if self.check_interrupts() {
-                        break;
-                    }
-
                     self.report_down();
 
                     // Establish a new connection and try again.
@@ -193,9 +147,6 @@ impl ComputeMonitor {
             self.last_checked = Utc::now();
             thread::sleep(MONITOR_CHECK_INTERVAL);
         }
-
-        // Graceful termination path
-        Ok(())
     }
 
     #[instrument(skip_all)]
@@ -478,10 +429,7 @@ pub fn launch_monitor(compute: &Arc<ComputeNode>) -> thread::JoinHandle<()> {
         .spawn(move || {
             let span = span!(Level::INFO, "compute_monitor");
             let _enter = span.enter();
-            match monitor.run() {
-                Ok(_) => info!("compute monitor thread terminated gracefully"),
-                Err(err) => error!("compute monitor thread terminated abnormally {:?}", err),
-            }
+            monitor.run();
         })
         .expect("cannot launch compute monitor thread")
 }

compute_tools/tests/pg_helpers_tests.rs

@@ -30,7 +30,7 @@ mod pg_helpers_tests {
             r#"fsync = off
 wal_level = logical
 hot_standby = on
-autoprewarm = off
+prewarm_lfc_on_startup = off
 neon.safekeepers = '127.0.0.1:6502,127.0.0.1:6503,127.0.0.1:6501'
 wal_log_hints = on
 log_connections = on

control_plane/Cargo.toml

@@ -36,7 +36,6 @@ pageserver_api.workspace = true
 pageserver_client.workspace = true
 postgres_backend.workspace = true
 safekeeper_api.workspace = true
-safekeeper_client.workspace = true
 postgres_connection.workspace = true
 storage_broker.workspace = true
 http-utils.workspace = true

control_plane/src/bin/neon_local.rs

@@ -45,7 +45,7 @@ use pageserver_api::models::{
 use pageserver_api::shard::{DEFAULT_STRIPE_SIZE, ShardCount, ShardStripeSize, TenantShardId};
 use postgres_backend::AuthType;
 use postgres_connection::parse_host_port;
-use safekeeper_api::membership::{SafekeeperGeneration, SafekeeperId};
+use safekeeper_api::membership::SafekeeperGeneration;
 use safekeeper_api::{
     DEFAULT_HTTP_LISTEN_PORT as DEFAULT_SAFEKEEPER_HTTP_PORT,
     DEFAULT_PG_LISTEN_PORT as DEFAULT_SAFEKEEPER_PG_PORT,
@@ -1255,45 +1255,6 @@ async fn handle_timeline(cmd: &TimelineCmd, env: &mut local_env::LocalEnv) -> Re
             pageserver
                 .timeline_import(tenant_id, timeline_id, base, pg_wal, args.pg_version)
                 .await?;
-            if env.storage_controller.timelines_onto_safekeepers {
-                println!("Creating timeline on safekeeper ...");
-                let timeline_info = pageserver
-                    .timeline_info(
-                        TenantShardId::unsharded(tenant_id),
-                        timeline_id,
-                        pageserver_client::mgmt_api::ForceAwaitLogicalSize::No,
-                    )
-                    .await?;
-                let default_sk = SafekeeperNode::from_env(env, env.safekeepers.first().unwrap());
-                let default_host = default_sk
-                    .conf
-                    .listen_addr
-                    .clone()
-                    .unwrap_or_else(|| "localhost".to_string());
-                let mconf = safekeeper_api::membership::Configuration {
-                    generation: SafekeeperGeneration::new(1),
-                    members: safekeeper_api::membership::MemberSet {
-                        m: vec![SafekeeperId {
-                            host: default_host,
-                            id: default_sk.conf.id,
-                            pg_port: default_sk.conf.pg_port,
-                        }],
-                    },
-                    new_members: None,
-                };
-                let pg_version = args.pg_version * 10000;
-                let req = safekeeper_api::models::TimelineCreateRequest {
-                    tenant_id,
-                    timeline_id,
-                    mconf,
-                    pg_version,
-                    system_id: None,
-                    wal_seg_size: None,
-                    start_lsn: timeline_info.last_record_lsn,
-                    commit_lsn: None,
-                };
-                default_sk.create_timeline(&req).await?;
-            }
             env.register_branch_mapping(branch_name.to_string(), tenant_id, timeline_id)?;
             println!("Done");
         }

control_plane/src/endpoint.rs

@@ -747,7 +747,7 @@ impl Endpoint {
                 logs_export_host: None::<String>,
                 endpoint_storage_addr: Some(endpoint_storage_addr),
                 endpoint_storage_token: Some(endpoint_storage_token),
-                autoprewarm: false,
+                prewarm_lfc_on_startup: false,
             };
 
             // this strange code is needed to support respec() in tests

control_plane/src/pageserver.rs

@@ -513,6 +513,11 @@ impl PageServerNode {
                 .map(|x| x.parse::<bool>())
                 .transpose()
                 .context("Failed to parse 'timeline_offloading' as bool")?,
+            wal_receiver_protocol_override: settings
+                .remove("wal_receiver_protocol_override")
+                .map(serde_json::from_str)
+                .transpose()
+                .context("parse `wal_receiver_protocol_override` from json")?,
             rel_size_v2_enabled: settings
                 .remove("rel_size_v2_enabled")
                 .map(|x| x.parse::<bool>())
@@ -635,16 +640,4 @@ impl PageServerNode {
 
         Ok(())
     }
-    pub async fn timeline_info(
-        &self,
-        tenant_shard_id: TenantShardId,
-        timeline_id: TimelineId,
-        force_await_logical_size: mgmt_api::ForceAwaitLogicalSize,
-    ) -> anyhow::Result<TimelineInfo> {
-        let timeline_info = self
-            .http_client
-            .timeline_info(tenant_shard_id, timeline_id, force_await_logical_size)
-            .await?;
-        Ok(timeline_info)
-    }
 }

control_plane/src/safekeeper.rs

@@ -6,6 +6,7 @@
 //!   .neon/safekeepers/<safekeeper id>
 //! ```
 use std::error::Error as _;
+use std::future::Future;
 use std::io::Write;
 use std::path::PathBuf;
 use std::time::Duration;
@@ -13,9 +14,9 @@ use std::{io, result};
 
 use anyhow::Context;
 use camino::Utf8PathBuf;
+use http_utils::error::HttpErrorBody;
 use postgres_connection::PgConnectionConfig;
-use safekeeper_api::models::TimelineCreateRequest;
-use safekeeper_client::mgmt_api;
+use reqwest::{IntoUrl, Method};
 use thiserror::Error;
 use utils::auth::{Claims, Scope};
 use utils::id::NodeId;
@@ -34,14 +35,25 @@ pub enum SafekeeperHttpError {
 
 type Result<T> = result::Result<T, SafekeeperHttpError>;
 
-fn err_from_client_err(err: mgmt_api::Error) -> SafekeeperHttpError {
-    use mgmt_api::Error::*;
-    match err {
-        ApiError(_, str) => SafekeeperHttpError::Response(str),
-        Cancelled => SafekeeperHttpError::Response("Cancelled".to_owned()),
-        ReceiveBody(err) => SafekeeperHttpError::Transport(err),
-        ReceiveErrorBody(err) => SafekeeperHttpError::Response(err),
-        Timeout(str) => SafekeeperHttpError::Response(format!("timeout: {str}")),
+pub(crate) trait ResponseErrorMessageExt: Sized {
+    fn error_from_body(self) -> impl Future<Output = Result<Self>> + Send;
+}
+
+impl ResponseErrorMessageExt for reqwest::Response {
+    async fn error_from_body(self) -> Result<Self> {
+        let status = self.status();
+        if !(status.is_client_error() || status.is_server_error()) {
+            return Ok(self);
+        }
+
+        // reqwest does not export its error construction utility functions, so let's craft the message ourselves
+        let url = self.url().to_owned();
+        Err(SafekeeperHttpError::Response(
+            match self.json::<HttpErrorBody>().await {
+                Ok(err_body) => format!("Error: {}", err_body.msg),
+                Err(_) => format!("Http error ({}) at {}.", status.as_u16(), url),
+            },
+        ))
     }
 }
 
@@ -58,8 +70,9 @@ pub struct SafekeeperNode {
 
     pub pg_connection_config: PgConnectionConfig,
     pub env: LocalEnv,
-    pub http_client: mgmt_api::Client,
+    pub http_client: reqwest::Client,
     pub listen_addr: String,
+    pub http_base_url: String,
 }
 
 impl SafekeeperNode {
@@ -69,14 +82,13 @@ impl SafekeeperNode {
         } else {
             "127.0.0.1".to_string()
         };
-        let jwt = None;
-        let http_base_url = format!("http://{}:{}", listen_addr, conf.http_port);
         SafekeeperNode {
             id: conf.id,
             conf: conf.clone(),
             pg_connection_config: Self::safekeeper_connection_config(&listen_addr, conf.pg_port),
             env: env.clone(),
-            http_client: mgmt_api::Client::new(env.create_http_client(), http_base_url, jwt),
+            http_client: env.create_http_client(),
+            http_base_url: format!("http://{}:{}/v1", listen_addr, conf.http_port),
             listen_addr,
         }
     }
@@ -266,19 +278,20 @@ impl SafekeeperNode {
         )
     }
 
-    pub async fn check_status(&self) -> Result<()> {
-        self.http_client
-            .status()
-            .await
-            .map_err(err_from_client_err)?;
-        Ok(())
+    fn http_request<U: IntoUrl>(&self, method: Method, url: U) -> reqwest::RequestBuilder {
+        // TODO: authentication
+        //if self.env.auth_type == AuthType::NeonJWT {
+        //    builder = builder.bearer_auth(&self.env.safekeeper_auth_token)
+        //}
+        self.http_client.request(method, url)
     }
 
-    pub async fn create_timeline(&self, req: &TimelineCreateRequest) -> Result<()> {
-        self.http_client
-            .create_timeline(req)
-            .await
-            .map_err(err_from_client_err)?;
+    pub async fn check_status(&self) -> Result<()> {
+        self.http_request(Method::GET, format!("{}/{}", self.http_base_url, "status"))
+            .send()
+            .await?
+            .error_from_body()
+            .await?;
         Ok(())
     }
 }

control_plane/storcon_cli/src/main.rs

@@ -61,16 +61,10 @@ enum Command {
         #[arg(long)]
         scheduling: Option<NodeSchedulingPolicy>,
     },
-    // Set a node status as deleted.
     NodeDelete {
         #[arg(long)]
         node_id: NodeId,
     },
-    /// Delete a tombstone of node from the storage controller.
-    NodeDeleteTombstone {
-        #[arg(long)]
-        node_id: NodeId,
-    },
     /// Modify a tenant's policies in the storage controller
     TenantPolicy {
         #[arg(long)]
@@ -88,8 +82,6 @@ enum Command {
     },
     /// List nodes known to the storage controller
     Nodes {},
-    /// List soft deleted nodes known to the storage controller
-    NodeTombstones {},
     /// List tenants known to the storage controller
     Tenants {
         /// If this field is set, it will list the tenants on a specific node
@@ -908,39 +900,6 @@ async fn main() -> anyhow::Result<()> {
                 .dispatch::<(), ()>(Method::DELETE, format!("control/v1/node/{node_id}"), None)
                 .await?;
         }
-        Command::NodeDeleteTombstone { node_id } => {
-            storcon_client
-                .dispatch::<(), ()>(
-                    Method::DELETE,
-                    format!("debug/v1/tombstone/{node_id}"),
-                    None,
-                )
-                .await?;
-        }
-        Command::NodeTombstones {} => {
-            let mut resp = storcon_client
-                .dispatch::<(), Vec<NodeDescribeResponse>>(
-                    Method::GET,
-                    "debug/v1/tombstone".to_string(),
-                    None,
-                )
-                .await?;
-
-            resp.sort_by(|a, b| a.listen_http_addr.cmp(&b.listen_http_addr));
-
-            let mut table = comfy_table::Table::new();
-            table.set_header(["Id", "Hostname", "AZ", "Scheduling", "Availability"]);
-            for node in resp {
-                table.add_row([
-                    format!("{}", node.id),
-                    node.listen_http_addr,
-                    node.availability_zone_id,
-                    format!("{:?}", node.scheduling),
-                    format!("{:?}", node.availability),
-                ]);
-            }
-            println!("{table}");
-        }
         Command::TenantSetTimeBasedEviction {
             tenant_id,
             period,

docker-compose/compute_wrapper/shell/compute.sh

@@ -1,18 +1,18 @@
-#!/usr/bin/env bash
+#!/bin/bash
 set -eux
 
 # Generate a random tenant or timeline ID
 #
 # Takes a variable name as argument. The result is stored in that variable.
 generate_id() {
-    local -n resvar=${1}
-    printf -v resvar '%08x%08x%08x%08x' ${SRANDOM} ${SRANDOM} ${SRANDOM} ${SRANDOM}
+    local -n resvar=$1
+    printf -v resvar '%08x%08x%08x%08x' $SRANDOM $SRANDOM $SRANDOM $SRANDOM
 }
 
 PG_VERSION=${PG_VERSION:-14}
 
-readonly CONFIG_FILE_ORG=/var/db/postgres/configs/config.json
-readonly CONFIG_FILE=/tmp/config.json
+CONFIG_FILE_ORG=/var/db/postgres/configs/config.json
+CONFIG_FILE=/tmp/config.json
 
 # Test that the first library path that the dynamic loader looks in is the path
 # that we use for custom compiled software
@@ -20,17 +20,17 @@ first_path="$(ldconfig --verbose 2>/dev/null \
     | grep --invert-match ^$'\t' \
     | cut --delimiter=: --fields=1 \
     | head --lines=1)"
-test "${first_path}" = '/usr/local/lib'
+test "$first_path" == '/usr/local/lib'
 
 echo "Waiting pageserver become ready."
 while ! nc -z pageserver 6400; do
-     sleep 1
+     sleep 1;
 done
 echo "Page server is ready."
 
-cp "${CONFIG_FILE_ORG}" "${CONFIG_FILE}"
+cp ${CONFIG_FILE_ORG} ${CONFIG_FILE}
 
- if [[ -n "${TENANT_ID:-}" && -n "${TIMELINE_ID:-}" ]]; then
+ if [ -n "${TENANT_ID:-}" ] && [ -n "${TIMELINE_ID:-}" ]; then
    tenant_id=${TENANT_ID}
    timeline_id=${TIMELINE_ID}
 else
@@ -41,7 +41,7 @@ else
        "http://pageserver:9898/v1/tenant"
   )
   tenant_id=$(curl "${PARAMS[@]}" | jq -r .[0].id)
-  if [[ -z "${tenant_id}" || "${tenant_id}" = null ]]; then
+  if [ -z "${tenant_id}" ] || [ "${tenant_id}" = null ]; then
     echo "Create a tenant"
     generate_id tenant_id
     PARAMS=(
@@ -51,7 +51,7 @@ else
         "http://pageserver:9898/v1/tenant/${tenant_id}/location_config"
     )
     result=$(curl "${PARAMS[@]}")
-    printf '%s\n' "${result}" | jq .
+    echo $result | jq .
   fi
 
   echo "Check if a timeline present"
@@ -61,7 +61,7 @@ else
        "http://pageserver:9898/v1/tenant/${tenant_id}/timeline"
   )
   timeline_id=$(curl "${PARAMS[@]}" | jq -r .[0].timeline_id)
-  if [[ -z "${timeline_id}" || "${timeline_id}" = null ]]; then
+  if [ -z "${timeline_id}" ] || [ "${timeline_id}" = null ]; then
     generate_id timeline_id
     PARAMS=(
         -sbf
@@ -71,7 +71,7 @@ else
         "http://pageserver:9898/v1/tenant/${tenant_id}/timeline/"
     )
     result=$(curl "${PARAMS[@]}")
-    printf '%s\n' "${result}" | jq .
+    echo $result | jq .
   fi
 fi
 
@@ -82,10 +82,10 @@ else
 fi
 echo "Adding pgx_ulid"
 shared_libraries=$(jq -r '.spec.cluster.settings[] | select(.name=="shared_preload_libraries").value' ${CONFIG_FILE})
-sed -i "s|${shared_libraries}|${shared_libraries},${ulid_extension}|" ${CONFIG_FILE}
+sed -i "s/${shared_libraries}/${shared_libraries},${ulid_extension}/" ${CONFIG_FILE}
 echo "Overwrite tenant id and timeline id in spec file"
-sed -i "s|TENANT_ID|${tenant_id}|" ${CONFIG_FILE}
-sed -i "s|TIMELINE_ID|${timeline_id}|" ${CONFIG_FILE}
+sed -i "s/TENANT_ID/${tenant_id}/" ${CONFIG_FILE}
+sed -i "s/TIMELINE_ID/${timeline_id}/" ${CONFIG_FILE}
 
 cat ${CONFIG_FILE}
 
@@ -93,5 +93,5 @@ echo "Start compute node"
 /usr/local/bin/compute_ctl --pgdata /var/db/postgres/compute \
      -C "postgresql://cloud_admin@localhost:55433/postgres"  \
      -b /usr/local/bin/postgres                              \
-     --compute-id "compute-${RANDOM}"                          \
-     --config "${CONFIG_FILE}"
+     --compute-id "compute-$RANDOM"                          \
+     --config "$CONFIG_FILE"

libs/compute_api/src/spec.rs

@@ -178,9 +178,9 @@ pub struct ComputeSpec {
     /// JWT for authorizing requests to endpoint storage service
     pub endpoint_storage_token: Option<String>,
 
-    /// Download LFC state from endpoint_storage and pass it to Postgres on startup
+    /// If true, download LFC state from endpoint_storage and pass it to Postgres on startup
     #[serde(default)]
-    pub autoprewarm: bool,
+    pub prewarm_lfc_on_startup: bool,
 }
 
 /// Feature flag to signal `compute_ctl` to enable certain experimental functionality.
@@ -192,9 +192,6 @@ pub enum ComputeFeature {
     /// track short-lived connections as user activity.
     ActivityMonitorExperimental,
 
-    /// Enable TLS functionality.
-    TlsExperimental,
-
     /// This is a special feature flag that is used to represent unknown feature flags.
     /// Basically all unknown to enum flags are represented as this one. See unit test
     /// `parse_unknown_features()` for more details.
@@ -253,44 +250,34 @@ impl RemoteExtSpec {
         }
 
         match self.extension_data.get(real_ext_name) {
-            Some(_ext_data) => Ok((
-                real_ext_name.to_string(),
-                Self::build_remote_path(build_tag, pg_major_version, real_ext_name)?,
-            )),
+            Some(_ext_data) => {
+                // We have decided to use the Go naming convention due to Kubernetes.
+
+                let arch = match std::env::consts::ARCH {
+                    "x86_64" => "amd64",
+                    "aarch64" => "arm64",
+                    arch => arch,
+                };
+
+                // Construct the path to the extension archive
+                // BUILD_TAG/PG_MAJOR_VERSION/extensions/EXTENSION_NAME.tar.zst
+                //
+                // Keep it in sync with path generation in
+                // https://github.com/neondatabase/build-custom-extensions/tree/main
+                let archive_path_str = format!(
+                    "{build_tag}/{arch}/{pg_major_version}/extensions/{real_ext_name}.tar.zst"
+                );
+                Ok((
+                    real_ext_name.to_string(),
+                    RemotePath::from_string(&archive_path_str)?,
+                ))
+            }
             None => Err(anyhow::anyhow!(
                 "real_ext_name {} is not found",
                 real_ext_name
             )),
         }
     }
-
-    /// Get the architecture-specific portion of the remote extension path. We
-    /// use the Go naming convention due to Kubernetes.
-    fn get_arch() -> &'static str {
-        match std::env::consts::ARCH {
-            "x86_64" => "amd64",
-            "aarch64" => "arm64",
-            arch => arch,
-        }
-    }
-
-    /// Build a [`RemotePath`] for an extension.
-    fn build_remote_path(
-        build_tag: &str,
-        pg_major_version: &str,
-        ext_name: &str,
-    ) -> anyhow::Result<RemotePath> {
-        let arch = Self::get_arch();
-
-        // Construct the path to the extension archive
-        // BUILD_TAG/PG_MAJOR_VERSION/extensions/EXTENSION_NAME.tar.zst
-        //
-        // Keep it in sync with path generation in
-        // https://github.com/neondatabase/build-custom-extensions/tree/main
-        RemotePath::from_string(&format!(
-            "{build_tag}/{arch}/{pg_major_version}/extensions/{ext_name}.tar.zst"
-        ))
-    }
 }
 
 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Deserialize, Serialize)]
@@ -531,37 +518,6 @@ mod tests {
             .expect("Library should be found");
     }
 
-    #[test]
-    fn remote_extension_path() {
-        let rspec: RemoteExtSpec = serde_json::from_value(serde_json::json!({
-            "public_extensions": ["ext"],
-            "custom_extensions": [],
-            "library_index": {
-                "extlib": "ext",
-            },
-            "extension_data": {
-                "ext": {
-                    "control_data": {
-                        "ext.control": ""
-                    },
-                    "archive_path": ""
-                }
-            },
-        }))
-        .unwrap();
-
-        let (_ext_name, ext_path) = rspec
-            .get_ext("ext", false, "latest", "v17")
-            .expect("Extension should be found");
-        // Starting with a forward slash would have consequences for the
-        // Url::join() that occurs when downloading a remote extension.
-        assert!(!ext_path.to_string().starts_with("/"));
-        assert_eq!(
-            ext_path,
-            RemoteExtSpec::build_remote_path("latest", "v17", "ext").unwrap()
-        );
-    }
-
     #[test]
     fn parse_spec_file() {
         let file = File::open("tests/cluster_spec.json").unwrap();

libs/compute_api/tests/cluster_spec.json

@@ -85,7 +85,7 @@
                 "vartype": "bool"
             },
             {
-                "name": "autoprewarm",
+                "name": "prewarm_lfc_on_startup",
                 "value": "off",
                 "vartype": "bool"
             },

libs/pageserver_api/src/config.rs

@@ -20,6 +20,7 @@ use postgres_backend::AuthType;
 use remote_storage::RemoteStorageConfig;
 use serde_with::serde_as;
 use utils::logging::LogFormat;
+use utils::postgres_client::PostgresClientProtocol;
 
 use crate::models::{ImageCompressionAlgorithm, LsnLease};
 
@@ -180,7 +181,6 @@ pub struct ConfigToml {
     pub virtual_file_io_engine: Option<crate::models::virtual_file::IoEngineKind>,
     pub ingest_batch_size: u64,
     pub max_vectored_read_bytes: MaxVectoredReadBytes,
-    pub max_get_vectored_keys: MaxGetVectoredKeys,
     pub image_compression: ImageCompressionAlgorithm,
     pub timeline_offloading: bool,
     pub ephemeral_bytes_per_memory_kb: usize,
@@ -188,6 +188,7 @@ pub struct ConfigToml {
     pub virtual_file_io_mode: Option<crate::models::virtual_file::IoMode>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub no_sync: Option<bool>,
+    pub wal_receiver_protocol: PostgresClientProtocol,
     pub page_service_pipelining: PageServicePipeliningConfig,
     pub get_vectored_concurrent_io: GetVectoredConcurrentIo,
     pub enable_read_path_debugging: Option<bool>,
@@ -228,7 +229,7 @@ pub enum PageServicePipeliningConfig {
 }
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 pub struct PageServicePipeliningConfigPipelined {
-    /// Failed config parsing and validation if larger than `max_get_vectored_keys`.
+    /// Causes runtime errors if larger than max get_vectored batch size.
     pub max_batch_size: NonZeroUsize,
     pub execution: PageServiceProtocolPipelinedExecutionStrategy,
     // The default below is such that new versions of the software can start
@@ -328,8 +329,6 @@ pub struct TimelineImportConfig {
     pub import_job_concurrency: NonZeroUsize,
     pub import_job_soft_size_limit: NonZeroUsize,
     pub import_job_checkpoint_threshold: NonZeroUsize,
-    /// Max size of the remote storage partial read done by any job
-    pub import_job_max_byte_range_size: NonZeroUsize,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
@@ -404,16 +403,6 @@ impl Default for EvictionOrder {
 #[serde(transparent)]
 pub struct MaxVectoredReadBytes(pub NonZeroUsize);
 
-#[derive(Copy, Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
-#[serde(transparent)]
-pub struct MaxGetVectoredKeys(NonZeroUsize);
-
-impl MaxGetVectoredKeys {
-    pub fn get(&self) -> usize {
-        self.0.get()
-    }
-}
-
 /// Tenant-level configuration values, used for various purposes.
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
 #[serde(default)]
@@ -525,6 +514,8 @@ pub struct TenantConfigToml {
     /// (either this flag or the pageserver-global one need to be set)
     pub timeline_offloading: bool,
 
+    pub wal_receiver_protocol_override: Option<PostgresClientProtocol>,
+
     /// Enable rel_size_v2 for this tenant. Once enabled, the tenant will persist this information into
     /// `index_part.json`, and it cannot be reversed.
     pub rel_size_v2_enabled: bool,
@@ -596,8 +587,6 @@ pub mod defaults {
     /// That is, slightly above 128 kB.
     pub const DEFAULT_MAX_VECTORED_READ_BYTES: usize = 130 * 1024; // 130 KiB
 
-    pub const DEFAULT_MAX_GET_VECTORED_KEYS: usize = 32;
-
     pub const DEFAULT_IMAGE_COMPRESSION: ImageCompressionAlgorithm =
         ImageCompressionAlgorithm::Zstd { level: Some(1) };
 
@@ -605,6 +594,9 @@ pub mod defaults {
 
     pub const DEFAULT_IO_BUFFER_ALIGNMENT: usize = 512;
 
+    pub const DEFAULT_WAL_RECEIVER_PROTOCOL: utils::postgres_client::PostgresClientProtocol =
+        utils::postgres_client::PostgresClientProtocol::Vanilla;
+
     pub const DEFAULT_SSL_KEY_FILE: &str = "server.key";
     pub const DEFAULT_SSL_CERT_FILE: &str = "server.crt";
 }
@@ -693,9 +685,6 @@ impl Default for ConfigToml {
             max_vectored_read_bytes: (MaxVectoredReadBytes(
                 NonZeroUsize::new(DEFAULT_MAX_VECTORED_READ_BYTES).unwrap(),
             )),
-            max_get_vectored_keys: (MaxGetVectoredKeys(
-                NonZeroUsize::new(DEFAULT_MAX_GET_VECTORED_KEYS).unwrap(),
-            )),
             image_compression: (DEFAULT_IMAGE_COMPRESSION),
             timeline_offloading: true,
             ephemeral_bytes_per_memory_kb: (DEFAULT_EPHEMERAL_BYTES_PER_MEMORY_KB),
@@ -703,6 +692,7 @@ impl Default for ConfigToml {
             virtual_file_io_mode: None,
             tenant_config: TenantConfigToml::default(),
             no_sync: None,
+            wal_receiver_protocol: DEFAULT_WAL_RECEIVER_PROTOCOL,
             page_service_pipelining: PageServicePipeliningConfig::Pipelined(
                 PageServicePipeliningConfigPipelined {
                     max_batch_size: NonZeroUsize::new(32).unwrap(),
@@ -726,7 +716,6 @@ impl Default for ConfigToml {
                 import_job_concurrency: NonZeroUsize::new(32).unwrap(),
                 import_job_soft_size_limit: NonZeroUsize::new(256 * 1024 * 1024).unwrap(),
                 import_job_checkpoint_threshold: NonZeroUsize::new(32).unwrap(),
-                import_job_max_byte_range_size: NonZeroUsize::new(4 * 1024 * 1024).unwrap(),
             },
             basebackup_cache_config: None,
             posthog_config: None,
@@ -847,6 +836,7 @@ impl Default for TenantConfigToml {
             lsn_lease_length: LsnLease::DEFAULT_LENGTH,
             lsn_lease_length_for_ts: LsnLease::DEFAULT_LENGTH_FOR_TS,
             timeline_offloading: true,
+            wal_receiver_protocol_override: None,
             rel_size_v2_enabled: false,
             gc_compaction_enabled: DEFAULT_GC_COMPACTION_ENABLED,
             gc_compaction_verification: DEFAULT_GC_COMPACTION_VERIFICATION,

libs/pageserver_api/src/controller_api.rs

@@ -344,35 +344,6 @@ impl Default for ShardSchedulingPolicy {
     }
 }
 
-#[derive(Serialize, Deserialize, Clone, Copy, Eq, PartialEq, Debug)]
-pub enum NodeLifecycle {
-    Active,
-    Deleted,
-}
-
-impl FromStr for NodeLifecycle {
-    type Err = anyhow::Error;
-
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "active" => Ok(Self::Active),
-            "deleted" => Ok(Self::Deleted),
-            _ => Err(anyhow::anyhow!("Unknown node lifecycle '{s}'")),
-        }
-    }
-}
-
-impl From<NodeLifecycle> for String {
-    fn from(value: NodeLifecycle) -> String {
-        use NodeLifecycle::*;
-        match value {
-            Active => "active",
-            Deleted => "deleted",
-        }
-        .to_string()
-    }
-}
-
 #[derive(Serialize, Deserialize, Clone, Copy, Eq, PartialEq, Debug)]
 pub enum NodeSchedulingPolicy {
     Active,

libs/pageserver_api/src/models.rs

@@ -20,6 +20,7 @@ use serde_with::serde_as;
 pub use utilization::PageserverUtilization;
 use utils::id::{NodeId, TenantId, TimelineId};
 use utils::lsn::Lsn;
+use utils::postgres_client::PostgresClientProtocol;
 use utils::{completion, serde_system_time};
 
 use crate::config::Ratio;
@@ -621,6 +622,8 @@ pub struct TenantConfigPatch {
     #[serde(skip_serializing_if = "FieldPatch::is_noop")]
     pub timeline_offloading: FieldPatch<bool>,
     #[serde(skip_serializing_if = "FieldPatch::is_noop")]
+    pub wal_receiver_protocol_override: FieldPatch<PostgresClientProtocol>,
+    #[serde(skip_serializing_if = "FieldPatch::is_noop")]
     pub rel_size_v2_enabled: FieldPatch<bool>,
     #[serde(skip_serializing_if = "FieldPatch::is_noop")]
     pub gc_compaction_enabled: FieldPatch<bool>,
@@ -745,6 +748,9 @@ pub struct TenantConfig {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub timeline_offloading: Option<bool>,
 
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub wal_receiver_protocol_override: Option<PostgresClientProtocol>,
+
     #[serde(skip_serializing_if = "Option::is_none")]
     pub rel_size_v2_enabled: Option<bool>,
 
@@ -806,6 +812,7 @@ impl TenantConfig {
             mut lsn_lease_length,
             mut lsn_lease_length_for_ts,
             mut timeline_offloading,
+            mut wal_receiver_protocol_override,
             mut rel_size_v2_enabled,
             mut gc_compaction_enabled,
             mut gc_compaction_verification,
@@ -898,6 +905,9 @@ impl TenantConfig {
             .map(|v| humantime::parse_duration(&v))?
             .apply(&mut lsn_lease_length_for_ts);
         patch.timeline_offloading.apply(&mut timeline_offloading);
+        patch
+            .wal_receiver_protocol_override
+            .apply(&mut wal_receiver_protocol_override);
         patch.rel_size_v2_enabled.apply(&mut rel_size_v2_enabled);
         patch
             .gc_compaction_enabled
@@ -950,6 +960,7 @@ impl TenantConfig {
             lsn_lease_length,
             lsn_lease_length_for_ts,
             timeline_offloading,
+            wal_receiver_protocol_override,
             rel_size_v2_enabled,
             gc_compaction_enabled,
             gc_compaction_verification,
@@ -1047,6 +1058,9 @@ impl TenantConfig {
             timeline_offloading: self
                 .timeline_offloading
                 .unwrap_or(global_conf.timeline_offloading),
+            wal_receiver_protocol_override: self
+                .wal_receiver_protocol_override
+                .or(global_conf.wal_receiver_protocol_override),
             rel_size_v2_enabled: self
                 .rel_size_v2_enabled
                 .unwrap_or(global_conf.rel_size_v2_enabled),

libs/posthog_client_lite/src/background_loop.rs

@@ -6,7 +6,7 @@ use arc_swap::ArcSwap;
 use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, info_span};
 
-use crate::{CaptureEvent, FeatureStore, PostHogClient, PostHogClientConfig};
+use crate::{FeatureStore, PostHogClient, PostHogClientConfig};
 
 /// A background loop that fetches feature flags from PostHog and updates the feature store.
 pub struct FeatureResolverBackgroundLoop {
@@ -24,16 +24,9 @@ impl FeatureResolverBackgroundLoop {
         }
     }
 
-    pub fn spawn(
-        self: Arc<Self>,
-        handle: &tokio::runtime::Handle,
-        refresh_period: Duration,
-        fake_tenants: Vec<CaptureEvent>,
-    ) {
+    pub fn spawn(self: Arc<Self>, handle: &tokio::runtime::Handle, refresh_period: Duration) {
         let this = self.clone();
         let cancel = self.cancel.clone();
-
-        // Main loop of updating the feature flags.
         handle.spawn(
             async move {
                 tracing::info!("Starting PostHog feature resolver");
@@ -63,22 +56,6 @@ impl FeatureResolverBackgroundLoop {
             }
             .instrument(info_span!("posthog_feature_resolver")),
         );
-
-        // Report fake tenants to PostHog so that we have the combination of all the properties in the UI.
-        // Do one report per pageserver restart.
-        let this = self.clone();
-        handle.spawn(
-            async move {
-                tracing::info!("Starting PostHog feature reporter");
-                for tenant in &fake_tenants {
-                    tracing::info!("Reporting fake tenant: {:?}", tenant);
-                }
-                if let Err(e) = this.posthog_client.capture_event_batch(&fake_tenants).await {
-                    tracing::warn!("Cannot report fake tenants: {}", e);
-                }
-            }
-            .instrument(info_span!("posthog_feature_reporter")),
-        );
     }
 
     pub fn feature_store(&self) -> Arc<FeatureStore> {

libs/posthog_client_lite/src/lib.rs

@@ -22,16 +22,6 @@ pub enum PostHogEvaluationError {
     Internal(String),
 }
 
-impl PostHogEvaluationError {
-    pub fn as_variant_str(&self) -> &'static str {
-        match self {
-            PostHogEvaluationError::NotAvailable(_) => "not_available",
-            PostHogEvaluationError::NoConditionGroupMatched => "no_condition_group_matched",
-            PostHogEvaluationError::Internal(_) => "internal",
-        }
-    }
-}
-
 #[derive(Deserialize)]
 pub struct LocalEvaluationResponse {
     pub flags: Vec<LocalEvaluationFlag>,
@@ -64,7 +54,7 @@ pub struct LocalEvaluationFlagFilterProperty {
     operator: String,
 }
 
-#[derive(Debug, Serialize, Deserialize, Clone)]
+#[derive(Debug, Serialize, Deserialize)]
 #[serde(untagged)]
 pub enum PostHogFlagFilterPropertyValue {
     String(String),
@@ -507,13 +497,6 @@ pub struct PostHogClient {
     client: reqwest::Client,
 }
 
-#[derive(Serialize, Debug)]
-pub struct CaptureEvent {
-    pub event: String,
-    pub distinct_id: String,
-    pub properties: serde_json::Value,
-}
-
 impl PostHogClient {
     pub fn new(config: PostHogClientConfig) -> Self {
         let client = reqwest::Client::new();
@@ -577,12 +560,12 @@ impl PostHogClient {
         &self,
         event: &str,
         distinct_id: &str,
-        properties: &serde_json::Value,
+        properties: &HashMap<String, PostHogFlagFilterPropertyValue>,
     ) -> anyhow::Result<()> {
         // PUBLIC_URL/capture/
+        // with bearer token of self.client_api_key
         let url = format!("{}/capture/", self.config.public_api_url);
-        let response = self
-            .client
+        self.client
             .post(url)
             .body(serde_json::to_string(&json!({
                 "api_key": self.config.client_api_key,
@@ -592,39 +575,6 @@ impl PostHogClient {
             }))?)
             .send()
             .await?;
-        let status = response.status();
-        let body = response.text().await?;
-        if !status.is_success() {
-            return Err(anyhow::anyhow!(
-                "Failed to capture events: {}, {}",
-                status,
-                body
-            ));
-        }
-        Ok(())
-    }
-
-    pub async fn capture_event_batch(&self, events: &[CaptureEvent]) -> anyhow::Result<()> {
-        // PUBLIC_URL/batch/
-        let url = format!("{}/batch/", self.config.public_api_url);
-        let response = self
-            .client
-            .post(url)
-            .body(serde_json::to_string(&json!({
-                "api_key": self.config.client_api_key,
-                "batch": events,
-            }))?)
-            .send()
-            .await?;
-        let status = response.status();
-        let body = response.text().await?;
-        if !status.is_success() {
-            return Err(anyhow::anyhow!(
-                "Failed to capture events: {}, {}",
-                status,
-                body
-            ));
-        }
         Ok(())
     }
 }

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

@@ -10,7 +10,7 @@ use crate::{Error, cancel_query_raw, connect_socket};
 pub(crate) async fn cancel_query<T>(
     config: Option<SocketConfig>,
     ssl_mode: SslMode,
-    tls: T,
+    mut tls: T,
     process_id: i32,
     secret_key: i32,
 ) -> Result<(), Error>

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

@@ -17,6 +17,7 @@ use crate::{Client, Connection, Error};
 
 /// TLS configuration.
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[non_exhaustive]
 pub enum SslMode {
     /// Do not use TLS.
     Disable,
@@ -230,7 +231,7 @@ impl Config {
     /// Requires the `runtime` Cargo feature (enabled by default).
     pub async fn connect<T>(
         &self,
-        tls: &T,
+        tls: T,
     ) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
     where
         T: MakeTlsConnect<TcpStream>,

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

@@ -13,7 +13,7 @@ use crate::tls::{MakeTlsConnect, TlsConnect};
 use crate::{Client, Config, Connection, Error, RawConnection};
 
 pub async fn connect<T>(
-    tls: &T,
+    mut tls: T,
     config: &Config,
 ) -> Result<(Client, Connection<TcpStream, T::Stream>), Error>
 where

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

@@ -47,7 +47,7 @@ pub trait MakeTlsConnect<S> {
     /// Creates a new `TlsConnect`or.
     ///
     /// The domain name is provided for certificate verification and SNI.
-    fn make_tls_connect(&self, domain: &str) -> Result<Self::TlsConnect, Self::Error>;
+    fn make_tls_connect(&mut self, domain: &str) -> Result<Self::TlsConnect, Self::Error>;
 }
 
 /// An asynchronous function wrapping a stream in a TLS session.
@@ -85,7 +85,7 @@ impl<S> MakeTlsConnect<S> for NoTls {
     type TlsConnect = NoTls;
     type Error = NoTlsError;
 
-    fn make_tls_connect(&self, _: &str) -> Result<NoTls, NoTlsError> {
+    fn make_tls_connect(&mut self, _: &str) -> Result<NoTls, NoTlsError> {
         Ok(NoTls)
     }
 }

libs/safekeeper_api/src/models.rs

@@ -13,7 +13,7 @@ use utils::pageserver_feedback::PageserverFeedback;
 use crate::membership::Configuration;
 use crate::{ServerInfo, Term};
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize)]
 pub struct SafekeeperStatus {
     pub id: NodeId,
 }

libs/walproposer/src/api_bindings.rs

@@ -439,7 +439,6 @@ pub fn empty_shmem() -> crate::bindings::WalproposerShmemState {
         currentClusterSize: crate::bindings::pg_atomic_uint64 { value: 0 },
         shard_ps_feedback: [empty_feedback; 128],
         num_shards: 0,
-        replica_promote: false,
         min_ps_feedback: empty_feedback,
     }
 }

pageserver/page_api/Cargo.toml

@@ -9,6 +9,7 @@ bytes.workspace = true
 pageserver_api.workspace = true
 postgres_ffi.workspace = true
 prost.workspace = true
+smallvec.workspace = true
 thiserror.workspace = true
 tonic.workspace = true
 utils.workspace = true

pageserver/page_api/src/model.rs

@@ -9,16 +9,12 @@
 //! - Use more precise datatypes, e.g. Lsn and uints shorter than 32 bits.
 //!
 //! - Validate protocol invariants, via try_from() and try_into().
-//!
-//! Validation only happens on the receiver side, i.e. when converting from Protobuf to domain
-//! types. This is where it matters -- the Protobuf types are less strict than the domain types, and
-//! receivers should expect all sorts of junk from senders. This also allows the sender to use e.g.
-//! stream combinators without dealing with errors, and avoids validating the same message twice.
 
 use std::fmt::Display;
 
 use bytes::Bytes;
 use postgres_ffi::Oid;
+use smallvec::SmallVec;
 // 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;
@@ -76,35 +72,47 @@ impl Display for ReadLsn {
     }
 }
 
+impl ReadLsn {
+    /// Validates the ReadLsn.
+    pub fn validate(&self) -> Result<(), ProtocolError> {
+        if self.request_lsn == Lsn::INVALID {
+            return Err(ProtocolError::invalid("request_lsn", self.request_lsn));
+        }
+        if self.not_modified_since_lsn > Some(self.request_lsn) {
+            return Err(ProtocolError::invalid(
+                "not_modified_since_lsn",
+                self.not_modified_since_lsn,
+            ));
+        }
+        Ok(())
+    }
+}
+
 impl TryFrom<proto::ReadLsn> for ReadLsn {
     type Error = ProtocolError;
 
     fn try_from(pb: proto::ReadLsn) -> Result<Self, Self::Error> {
-        if pb.request_lsn == 0 {
-            return Err(ProtocolError::invalid("request_lsn", pb.request_lsn));
-        }
-        if pb.not_modified_since_lsn > pb.request_lsn {
-            return Err(ProtocolError::invalid(
-                "not_modified_since_lsn",
-                pb.not_modified_since_lsn,
-            ));
-        }
-        Ok(Self {
+        let read_lsn = Self {
             request_lsn: Lsn(pb.request_lsn),
             not_modified_since_lsn: match pb.not_modified_since_lsn {
                 0 => None,
                 lsn => Some(Lsn(lsn)),
             },
-        })
+        };
+        read_lsn.validate()?;
+        Ok(read_lsn)
     }
 }
 
-impl From<ReadLsn> for proto::ReadLsn {
-    fn from(read_lsn: ReadLsn) -> Self {
-        Self {
+impl TryFrom<ReadLsn> for proto::ReadLsn {
+    type Error = ProtocolError;
+
+    fn try_from(read_lsn: ReadLsn) -> Result<Self, Self::Error> {
+        read_lsn.validate()?;
+        Ok(Self {
             request_lsn: read_lsn.request_lsn.0,
             not_modified_since_lsn: read_lsn.not_modified_since_lsn.unwrap_or_default().0,
-        }
+        })
     }
 }
 
@@ -159,15 +167,6 @@ impl TryFrom<proto::CheckRelExistsRequest> for CheckRelExistsRequest {
     }
 }
 
-impl From<CheckRelExistsRequest> for proto::CheckRelExistsRequest {
-    fn from(request: CheckRelExistsRequest) -> Self {
-        Self {
-            read_lsn: Some(request.read_lsn.into()),
-            rel: Some(request.rel.into()),
-        }
-    }
-}
-
 pub type CheckRelExistsResponse = bool;
 
 impl From<proto::CheckRelExistsResponse> for CheckRelExistsResponse {
@@ -205,12 +204,14 @@ impl TryFrom<proto::GetBaseBackupRequest> for GetBaseBackupRequest {
     }
 }
 
-impl From<GetBaseBackupRequest> for proto::GetBaseBackupRequest {
-    fn from(request: GetBaseBackupRequest) -> Self {
-        Self {
-            read_lsn: Some(request.read_lsn.into()),
+impl TryFrom<GetBaseBackupRequest> for proto::GetBaseBackupRequest {
+    type Error = ProtocolError;
+
+    fn try_from(request: GetBaseBackupRequest) -> Result<Self, Self::Error> {
+        Ok(Self {
+            read_lsn: Some(request.read_lsn.try_into()?),
             replica: request.replica,
-        }
+        })
     }
 }
 
@@ -227,9 +228,14 @@ impl TryFrom<proto::GetBaseBackupResponseChunk> for GetBaseBackupResponseChunk {
     }
 }
 
-impl From<GetBaseBackupResponseChunk> for proto::GetBaseBackupResponseChunk {
-    fn from(chunk: GetBaseBackupResponseChunk) -> Self {
-        Self { chunk }
+impl TryFrom<GetBaseBackupResponseChunk> for proto::GetBaseBackupResponseChunk {
+    type Error = ProtocolError;
+
+    fn try_from(chunk: GetBaseBackupResponseChunk) -> Result<Self, Self::Error> {
+        if chunk.is_empty() {
+            return Err(ProtocolError::Missing("chunk"));
+        }
+        Ok(Self { chunk })
     }
 }
 
@@ -254,12 +260,14 @@ impl TryFrom<proto::GetDbSizeRequest> for GetDbSizeRequest {
     }
 }
 
-impl From<GetDbSizeRequest> for proto::GetDbSizeRequest {
-    fn from(request: GetDbSizeRequest) -> Self {
-        Self {
-            read_lsn: Some(request.read_lsn.into()),
+impl TryFrom<GetDbSizeRequest> for proto::GetDbSizeRequest {
+    type Error = ProtocolError;
+
+    fn try_from(request: GetDbSizeRequest) -> Result<Self, Self::Error> {
+        Ok(Self {
+            read_lsn: Some(request.read_lsn.try_into()?),
             db_oid: request.db_oid,
-        }
+        })
     }
 }
 
@@ -294,7 +302,7 @@ pub struct GetPageRequest {
     /// Multiple pages will be executed as a single batch by the Pageserver, amortizing layer access
     /// costs and parallelizing them. This may increase the latency of any individual request, but
     /// improves the overall latency and throughput of the batch as a whole.
-    pub block_numbers: Vec<u32>,
+    pub block_numbers: SmallVec<[u32; 1]>,
 }
 
 impl TryFrom<proto::GetPageRequest> for GetPageRequest {
@@ -312,20 +320,25 @@ impl TryFrom<proto::GetPageRequest> for GetPageRequest {
                 .ok_or(ProtocolError::Missing("read_lsn"))?
                 .try_into()?,
             rel: pb.rel.ok_or(ProtocolError::Missing("rel"))?.try_into()?,
-            block_numbers: pb.block_number,
+            block_numbers: pb.block_number.into(),
         })
     }
 }
 
-impl From<GetPageRequest> for proto::GetPageRequest {
-    fn from(request: GetPageRequest) -> Self {
-        Self {
+impl TryFrom<GetPageRequest> for proto::GetPageRequest {
+    type Error = ProtocolError;
+
+    fn try_from(request: GetPageRequest) -> Result<Self, Self::Error> {
+        if request.block_numbers.is_empty() {
+            return Err(ProtocolError::Missing("block_number"));
+        }
+        Ok(Self {
             request_id: request.request_id,
             request_class: request.request_class.into(),
-            read_lsn: Some(request.read_lsn.into()),
+            read_lsn: Some(request.read_lsn.try_into()?),
             rel: Some(request.rel.into()),
-            block_number: request.block_numbers,
-        }
+            block_number: request.block_numbers.into_vec(),
+        })
     }
 }
 
@@ -397,7 +410,7 @@ pub struct GetPageResponse {
     /// A string describing the status, if any.
     pub reason: Option<String>,
     /// The 8KB page images, in the same order as the request. Empty if status != OK.
-    pub page_images: Vec<Bytes>,
+    pub page_images: SmallVec<[Bytes; 1]>,
 }
 
 impl From<proto::GetPageResponse> for GetPageResponse {
@@ -406,7 +419,7 @@ impl From<proto::GetPageResponse> for GetPageResponse {
             request_id: pb.request_id,
             status_code: pb.status_code.into(),
             reason: Some(pb.reason).filter(|r| !r.is_empty()),
-            page_images: pb.page_image,
+            page_images: pb.page_image.into(),
         }
     }
 }
@@ -417,7 +430,7 @@ impl From<GetPageResponse> for proto::GetPageResponse {
             request_id: response.request_id,
             status_code: response.status_code.into(),
             reason: response.reason.unwrap_or_default(),
-            page_image: response.page_images,
+            page_image: response.page_images.into_vec(),
         }
     }
 }
@@ -506,12 +519,14 @@ impl TryFrom<proto::GetRelSizeRequest> for GetRelSizeRequest {
     }
 }
 
-impl From<GetRelSizeRequest> for proto::GetRelSizeRequest {
-    fn from(request: GetRelSizeRequest) -> Self {
-        Self {
-            read_lsn: Some(request.read_lsn.into()),
+impl TryFrom<GetRelSizeRequest> for proto::GetRelSizeRequest {
+    type Error = ProtocolError;
+
+    fn try_from(request: GetRelSizeRequest) -> Result<Self, Self::Error> {
+        Ok(Self {
+            read_lsn: Some(request.read_lsn.try_into()?),
             rel: Some(request.rel.into()),
-        }
+        })
     }
 }
 
@@ -554,13 +569,15 @@ impl TryFrom<proto::GetSlruSegmentRequest> for GetSlruSegmentRequest {
     }
 }
 
-impl From<GetSlruSegmentRequest> for proto::GetSlruSegmentRequest {
-    fn from(request: GetSlruSegmentRequest) -> Self {
-        Self {
-            read_lsn: Some(request.read_lsn.into()),
+impl TryFrom<GetSlruSegmentRequest> for proto::GetSlruSegmentRequest {
+    type Error = ProtocolError;
+
+    fn try_from(request: GetSlruSegmentRequest) -> Result<Self, Self::Error> {
+        Ok(Self {
+            read_lsn: Some(request.read_lsn.try_into()?),
             kind: request.kind as u32,
             segno: request.segno,
-        }
+        })
     }
 }
 
@@ -577,9 +594,15 @@ impl TryFrom<proto::GetSlruSegmentResponse> for GetSlruSegmentResponse {
     }
 }
 
-impl From<GetSlruSegmentResponse> for proto::GetSlruSegmentResponse {
-    fn from(segment: GetSlruSegmentResponse) -> Self {
-        Self { segment }
+impl TryFrom<GetSlruSegmentResponse> for proto::GetSlruSegmentResponse {
+    type Error = ProtocolError;
+
+    fn try_from(segment: GetSlruSegmentResponse) -> Result<Self, Self::Error> {
+        // TODO: can a segment legitimately be empty?
+        if segment.is_empty() {
+            return Err(ProtocolError::Missing("segment"));
+        }
+        Ok(Self { segment })
     }
 }
 

pageserver/pagebench/Cargo.toml

@@ -9,7 +9,6 @@ license.workspace = true
 [dependencies]
 anyhow.workspace = true
 async-trait.workspace = true
-bytes.workspace = true
 camino.workspace = true
 clap.workspace = true
 futures.workspace = true

pageserver/pagebench/src/cmd/getpage_latest_lsn.rs

@@ -1,4 +1,4 @@
-use std::collections::{HashMap, HashSet, VecDeque};
+use std::collections::{HashSet, VecDeque};
 use std::future::Future;
 use std::num::NonZeroUsize;
 use std::pin::Pin;
@@ -8,12 +8,12 @@ use std::time::{Duration, Instant};
 
 use anyhow::Context;
 use async_trait::async_trait;
-use bytes::Bytes;
 use camino::Utf8PathBuf;
 use pageserver_api::key::Key;
 use pageserver_api::keyspace::KeySpaceAccum;
-use pageserver_api::models::{PagestreamGetPageRequest, PagestreamRequest};
-use pageserver_api::reltag::RelTag;
+use pageserver_api::models::{
+    PagestreamGetPageRequest, PagestreamGetPageResponse, PagestreamRequest,
+};
 use pageserver_api::shard::TenantShardId;
 use pageserver_page_api::proto;
 use rand::prelude::*;
@@ -77,16 +77,6 @@ pub(crate) struct Args {
     #[clap(long, default_value = "1")]
     queue_depth: NonZeroUsize,
 
-    /// Batch size of contiguous pages generated by each client. This is equivalent to how Postgres
-    /// will request page batches (e.g. prefetches or vectored reads). A batch counts as 1 RPS and
-    /// 1 queue depth.
-    ///
-    /// The libpq protocol does not support client-side batching, and will submit batches as many
-    /// individual requests, in the hope that the server will batch them. Each batch still counts as
-    /// 1 RPS and 1 queue depth.
-    #[clap(long, default_value = "1")]
-    batch_size: NonZeroUsize,
-
     #[clap(long)]
     only_relnode: Option<u32>,
 
@@ -402,16 +392,7 @@ async fn run_worker(
     shared_state.start_work_barrier.wait().await;
     let client_start = Instant::now();
     let mut ticks_processed = 0;
-    let mut req_id = 0;
-    let batch_size: usize = args.batch_size.into();
-
-    // Track inflight requests by request ID and start time. This times the request duration, and
-    // ensures responses match requests. We don't expect responses back in any particular order.
-    //
-    // NB: this does not check that all requests received a response, because we don't wait for the
-    // inflight requests to complete when the duration elapses.
-    let mut inflight: HashMap<u64, Instant> = HashMap::new();
-
+    let mut inflight = VecDeque::new();
     while !cancel.is_cancelled() {
         // Detect if a request took longer than the RPS rate
         if let Some(period) = &rps_period {
@@ -427,72 +408,36 @@ async fn run_worker(
         }
 
         while inflight.len() < args.queue_depth.get() {
-            req_id += 1;
             let start = Instant::now();
-            let (req_lsn, mod_lsn, rel, blks) = {
-                /// Converts a compact i128 key to a relation tag and block number.
-                fn key_to_block(key: i128) -> (RelTag, u32) {
-                    let key = Key::from_i128(key);
-                    assert!(key.is_rel_block_key());
-                    key.to_rel_block()
-                        .expect("we filter non-rel-block keys out above")
-                }
-
-                // Pick a random page from a random relation.
+            let req = {
                 let mut rng = rand::thread_rng();
                 let r = &ranges[weights.sample(&mut rng)];
                 let key: i128 = rng.gen_range(r.start..r.end);
-                let (rel_tag, block_no) = key_to_block(key);
-
-                let mut blks = VecDeque::with_capacity(batch_size);
-                blks.push_back(block_no);
-
-                // If requested, populate a batch of sequential pages. This is how Postgres will
-                // request page batches (e.g. prefetches). If we hit the end of the relation, we
-                // grow the batch towards the start too.
-                for i in 1..batch_size {
-                    let (r, b) = key_to_block(key + i as i128);
-                    if r != rel_tag {
-                        break; // went outside relation
-                    }
-                    blks.push_back(b)
+                let key = Key::from_i128(key);
+                assert!(key.is_rel_block_key());
+                let (rel_tag, block_no) = key
+                    .to_rel_block()
+                    .expect("we filter non-rel-block keys out above");
+                PagestreamGetPageRequest {
+                    hdr: PagestreamRequest {
+                        reqid: 0,
+                        request_lsn: if rng.gen_bool(args.req_latest_probability) {
+                            Lsn::MAX
+                        } else {
+                            r.timeline_lsn
+                        },
+                        not_modified_since: r.timeline_lsn,
+                    },
+                    rel: rel_tag,
+                    blkno: block_no,
                 }
-
-                if blks.len() < batch_size {
-                    // Grow batch backwards if needed.
-                    for i in 1..batch_size {
-                        let (r, b) = key_to_block(key - i as i128);
-                        if r != rel_tag {
-                            break; // went outside relation
-                        }
-                        blks.push_front(b)
-                    }
-                }
-
-                // We assume that the entire batch can fit within the relation.
-                assert_eq!(blks.len(), batch_size, "incomplete batch");
-
-                let req_lsn = if rng.gen_bool(args.req_latest_probability) {
-                    Lsn::MAX
-                } else {
-                    r.timeline_lsn
-                };
-                (req_lsn, r.timeline_lsn, rel_tag, blks.into())
             };
-            client
-                .send_get_page(req_id, req_lsn, mod_lsn, rel, blks)
-                .await
-                .unwrap();
-            let old = inflight.insert(req_id, start);
-            assert!(old.is_none(), "duplicate request ID {req_id}");
+            client.send_get_page(req).await.unwrap();
+            inflight.push_back(start);
         }
 
-        let (req_id, pages) = client.recv_get_page().await.unwrap();
-        assert_eq!(pages.len(), batch_size, "unexpected page count");
-        assert!(pages.iter().all(|p| !p.is_empty()), "empty page");
-        let start = inflight
-            .remove(&req_id)
-            .expect("response for unknown request ID");
+        let start = inflight.pop_front().unwrap();
+        client.recv_get_page().await.unwrap();
         let end = Instant::now();
         shared_state.live_stats.request_done();
         ticks_processed += 1;
@@ -522,24 +467,15 @@ async fn run_worker(
 #[async_trait]
 trait Client: Send {
     /// Sends an asynchronous GetPage request to the pageserver.
-    async fn send_get_page(
-        &mut self,
-        req_id: u64,
-        req_lsn: Lsn,
-        mod_lsn: Lsn,
-        rel: RelTag,
-        blks: Vec<u32>,
-    ) -> anyhow::Result<()>;
+    async fn send_get_page(&mut self, req: PagestreamGetPageRequest) -> anyhow::Result<()>;
 
     /// Receives the next GetPage response from the pageserver.
-    async fn recv_get_page(&mut self) -> anyhow::Result<(u64, Vec<Bytes>)>;
+    async fn recv_get_page(&mut self) -> anyhow::Result<PagestreamGetPageResponse>;
 }
 
 /// A libpq-based Pageserver client.
 struct LibpqClient {
     inner: pageserver_client::page_service::PagestreamClient,
-    // Track sent batches, so we know how many responses to expect.
-    batch_sizes: VecDeque<usize>,
 }
 
 impl LibpqClient {
@@ -548,55 +484,18 @@ impl LibpqClient {
             .await?
             .pagestream(ttid.tenant_id, ttid.timeline_id)
             .await?;
-        Ok(Self {
-            inner,
-            batch_sizes: VecDeque::new(),
-        })
+        Ok(Self { inner })
     }
 }
 
 #[async_trait]
 impl Client for LibpqClient {
-    async fn send_get_page(
-        &mut self,
-        req_id: u64,
-        req_lsn: Lsn,
-        mod_lsn: Lsn,
-        rel: RelTag,
-        blks: Vec<u32>,
-    ) -> anyhow::Result<()> {
-        // libpq doesn't support client-side batches, so we send a bunch of individual requests
-        // instead in the hope that the server will batch them for us. We use the same request ID
-        // for all, because we'll return a single batch response.
-        self.batch_sizes.push_back(blks.len());
-        for blkno in blks {
-            let req = PagestreamGetPageRequest {
-                hdr: PagestreamRequest {
-                    reqid: req_id,
-                    request_lsn: req_lsn,
-                    not_modified_since: mod_lsn,
-                },
-                rel,
-                blkno,
-            };
-            self.inner.getpage_send(req).await?;
-        }
-        Ok(())
+    async fn send_get_page(&mut self, req: PagestreamGetPageRequest) -> anyhow::Result<()> {
+        self.inner.getpage_send(req).await
     }
 
-    async fn recv_get_page(&mut self) -> anyhow::Result<(u64, Vec<Bytes>)> {
-        let batch_size = self.batch_sizes.pop_front().unwrap();
-        let mut batch = Vec::with_capacity(batch_size);
-        let mut req_id = None;
-        for _ in 0..batch_size {
-            let resp = self.inner.getpage_recv().await?;
-            if req_id.is_none() {
-                req_id = Some(resp.req.hdr.reqid);
-            }
-            assert_eq!(req_id, Some(resp.req.hdr.reqid), "request ID mismatch");
-            batch.push(resp.page);
-        }
-        Ok((req_id.unwrap(), batch))
+    async fn recv_get_page(&mut self) -> anyhow::Result<PagestreamGetPageResponse> {
+        self.inner.getpage_recv().await
     }
 }
 
@@ -633,35 +532,31 @@ impl GrpcClient {
 
 #[async_trait]
 impl Client for GrpcClient {
-    async fn send_get_page(
-        &mut self,
-        req_id: u64,
-        req_lsn: Lsn,
-        mod_lsn: Lsn,
-        rel: RelTag,
-        blks: Vec<u32>,
-    ) -> anyhow::Result<()> {
+    async fn send_get_page(&mut self, req: PagestreamGetPageRequest) -> anyhow::Result<()> {
         let req = proto::GetPageRequest {
-            request_id: req_id,
+            request_id: 0,
             request_class: proto::GetPageClass::Normal as i32,
             read_lsn: Some(proto::ReadLsn {
-                request_lsn: req_lsn.0,
-                not_modified_since_lsn: mod_lsn.0,
+                request_lsn: req.hdr.request_lsn.0,
+                not_modified_since_lsn: req.hdr.not_modified_since.0,
             }),
-            rel: Some(rel.into()),
-            block_number: blks,
+            rel: Some(req.rel.into()),
+            block_number: vec![req.blkno],
         };
         self.req_tx.send(req).await?;
         Ok(())
     }
 
-    async fn recv_get_page(&mut self) -> anyhow::Result<(u64, Vec<Bytes>)> {
+    async fn recv_get_page(&mut self) -> anyhow::Result<PagestreamGetPageResponse> {
         let resp = self.resp_rx.message().await?.unwrap();
         anyhow::ensure!(
             resp.status_code == proto::GetPageStatusCode::Ok as i32,
             "unexpected status code: {}",
             resp.status_code
         );
-        Ok((resp.request_id, resp.page_image))
+        Ok(PagestreamGetPageResponse {
+            page: resp.page_image[0].clone(),
+            req: PagestreamGetPageRequest::default(), // dummy
+        })
     }
 }

pageserver/src/basebackup.rs

@@ -371,7 +371,7 @@ where
                 .await?
                 .partition(
                     self.timeline.get_shard_identity(),
-                    self.timeline.conf.max_get_vectored_keys.get() as u64 * BLCKSZ as u64,
+                    Timeline::MAX_GET_VECTORED_KEYS * BLCKSZ as u64,
                 );
 
             let mut slru_builder = SlruSegmentsBuilder::new(&mut self.ar);

pageserver/src/bin/pageserver.rs

@@ -23,7 +23,6 @@ use pageserver::deletion_queue::DeletionQueue;
 use pageserver::disk_usage_eviction_task::{self, launch_disk_usage_global_eviction_task};
 use pageserver::feature_resolver::FeatureResolver;
 use pageserver::metrics::{STARTUP_DURATION, STARTUP_IS_LOADING};
-use pageserver::page_service::GrpcPageServiceHandler;
 use pageserver::task_mgr::{
     BACKGROUND_RUNTIME, COMPUTE_REQUEST_RUNTIME, MGMT_REQUEST_RUNTIME, WALRECEIVER_RUNTIME,
 };
@@ -159,6 +158,7 @@ fn main() -> anyhow::Result<()> {
     // (maybe we should automate this with a visitor?).
     info!(?conf.virtual_file_io_engine, "starting with virtual_file IO engine");
     info!(?conf.virtual_file_io_mode, "starting with virtual_file IO mode");
+    info!(?conf.wal_receiver_protocol, "starting with WAL receiver protocol");
     info!(?conf.validate_wal_contiguity, "starting with WAL contiguity validation");
     info!(?conf.page_service_pipelining, "starting with page service pipelining config");
     info!(?conf.get_vectored_concurrent_io, "starting with get_vectored IO concurrency config");
@@ -815,11 +815,10 @@ fn start_pageserver(
     // necessary?
     let mut page_service_grpc = None;
     if let Some(grpc_listener) = grpc_listener {
-        page_service_grpc = Some(GrpcPageServiceHandler::spawn(
+        page_service_grpc = Some(page_service::spawn_grpc(
             tenant_manager.clone(),
             grpc_auth,
             otel_guard.as_ref().map(|g| g.dispatch.clone()),
-            conf.get_vectored_concurrent_io,
             grpc_listener,
         )?);
     }

pageserver/src/config.rs

@@ -14,10 +14,7 @@ use std::time::Duration;
 use anyhow::{Context, bail, ensure};
 use camino::{Utf8Path, Utf8PathBuf};
 use once_cell::sync::OnceCell;
-use pageserver_api::config::{
-    DiskUsageEvictionTaskConfig, MaxGetVectoredKeys, MaxVectoredReadBytes,
-    PageServicePipeliningConfig, PageServicePipeliningConfigPipelined, PostHogConfig,
-};
+use pageserver_api::config::{DiskUsageEvictionTaskConfig, MaxVectoredReadBytes, PostHogConfig};
 use pageserver_api::models::ImageCompressionAlgorithm;
 use pageserver_api::shard::TenantShardId;
 use pem::Pem;
@@ -27,6 +24,7 @@ use reqwest::Url;
 use storage_broker::Uri;
 use utils::id::{NodeId, TimelineId};
 use utils::logging::{LogFormat, SecretString};
+use utils::postgres_client::PostgresClientProtocol;
 
 use crate::tenant::storage_layer::inmemory_layer::IndexEntry;
 use crate::tenant::{TENANTS_SEGMENT_NAME, TIMELINES_SEGMENT_NAME};
@@ -187,9 +185,6 @@ pub struct PageServerConf {
 
     pub max_vectored_read_bytes: MaxVectoredReadBytes,
 
-    /// Maximum number of keys to be read in a single get_vectored call.
-    pub max_get_vectored_keys: MaxGetVectoredKeys,
-
     pub image_compression: ImageCompressionAlgorithm,
 
     /// Whether to offload archived timelines automatically
@@ -210,6 +205,8 @@ pub struct PageServerConf {
     /// Optionally disable disk syncs (unsafe!)
     pub no_sync: bool,
 
+    pub wal_receiver_protocol: PostgresClientProtocol,
+
     pub page_service_pipelining: pageserver_api::config::PageServicePipeliningConfig,
 
     pub get_vectored_concurrent_io: pageserver_api::config::GetVectoredConcurrentIo,
@@ -407,7 +404,6 @@ impl PageServerConf {
             secondary_download_concurrency,
             ingest_batch_size,
             max_vectored_read_bytes,
-            max_get_vectored_keys,
             image_compression,
             timeline_offloading,
             ephemeral_bytes_per_memory_kb,
@@ -418,6 +414,7 @@ impl PageServerConf {
             virtual_file_io_engine,
             tenant_config,
             no_sync,
+            wal_receiver_protocol,
             page_service_pipelining,
             get_vectored_concurrent_io,
             enable_read_path_debugging,
@@ -473,13 +470,13 @@ impl PageServerConf {
             secondary_download_concurrency,
             ingest_batch_size,
             max_vectored_read_bytes,
-            max_get_vectored_keys,
             image_compression,
             timeline_offloading,
             ephemeral_bytes_per_memory_kb,
             import_pgdata_upcall_api,
             import_pgdata_upcall_api_token: import_pgdata_upcall_api_token.map(SecretString::from),
             import_pgdata_aws_endpoint_url,
+            wal_receiver_protocol,
             page_service_pipelining,
             get_vectored_concurrent_io,
             tracing,
@@ -601,19 +598,6 @@ impl PageServerConf {
                 )
             })?;
 
-        if let PageServicePipeliningConfig::Pipelined(PageServicePipeliningConfigPipelined {
-            max_batch_size,
-            ..
-        }) = conf.page_service_pipelining
-        {
-            if max_batch_size.get() > conf.max_get_vectored_keys.get() {
-                return Err(anyhow::anyhow!(
-                    "`max_batch_size` ({max_batch_size}) must be less than or equal to `max_get_vectored_keys` ({})",
-                    conf.max_get_vectored_keys.get()
-                ));
-            }
-        };
-
         Ok(conf)
     }
 
@@ -701,7 +685,6 @@ impl ConfigurableSemaphore {
 mod tests {
 
     use camino::Utf8PathBuf;
-    use rstest::rstest;
     use utils::id::NodeId;
 
     use super::PageServerConf;
@@ -741,28 +724,4 @@ mod tests {
         PageServerConf::parse_and_validate(NodeId(0), config_toml, &workdir)
             .expect_err("parse_and_validate should fail for endpoint without scheme");
     }
-
-    #[rstest]
-    #[case(32, 32, true)]
-    #[case(64, 32, false)]
-    #[case(64, 64, true)]
-    #[case(128, 128, true)]
-    fn test_config_max_batch_size_is_valid(
-        #[case] max_batch_size: usize,
-        #[case] max_get_vectored_keys: usize,
-        #[case] is_valid: bool,
-    ) {
-        let input = format!(
-            r#"
-            control_plane_api = "http://localhost:6666"
-            max_get_vectored_keys = {max_get_vectored_keys}
-            page_service_pipelining = {{ mode="pipelined", execution="concurrent-futures", max_batch_size={max_batch_size}, batching="uniform-lsn" }}
-        "#,
-        );
-        let config_toml = toml_edit::de::from_str::<pageserver_api::config::ConfigToml>(&input)
-            .expect("config has valid fields");
-        let workdir = Utf8PathBuf::from("/nonexistent");
-        let result = PageServerConf::parse_and_validate(NodeId(0), config_toml, &workdir);
-        assert_eq!(result.is_ok(), is_valid);
-    }
 }

pageserver/src/feature_resolver.rs

@@ -1,28 +1,21 @@
 use std::{collections::HashMap, sync::Arc, time::Duration};
 
 use posthog_client_lite::{
-    CaptureEvent, FeatureResolverBackgroundLoop, PostHogClientConfig, PostHogEvaluationError,
-    PostHogFlagFilterPropertyValue,
+    FeatureResolverBackgroundLoop, PostHogClientConfig, PostHogEvaluationError,
 };
-use remote_storage::RemoteStorageKind;
-use serde_json::json;
 use tokio_util::sync::CancellationToken;
 use utils::id::TenantId;
 
-use crate::{config::PageServerConf, metrics::FEATURE_FLAG_EVALUATION};
+use crate::config::PageServerConf;
 
 #[derive(Clone)]
 pub struct FeatureResolver {
     inner: Option<Arc<FeatureResolverBackgroundLoop>>,
-    internal_properties: Option<Arc<HashMap<String, PostHogFlagFilterPropertyValue>>>,
 }
 
 impl FeatureResolver {
     pub fn new_disabled() -> Self {
-        Self {
-            inner: None,
-            internal_properties: None,
-        }
+        Self { inner: None }
     }
 
     pub fn spawn(
@@ -43,114 +36,14 @@ impl FeatureResolver {
                 shutdown_pageserver,
             );
             let inner = Arc::new(inner);
-
-            // The properties shared by all tenants on this pageserver.
-            let internal_properties = {
-                let mut properties = HashMap::new();
-                properties.insert(
-                    "pageserver_id".to_string(),
-                    PostHogFlagFilterPropertyValue::String(conf.id.to_string()),
-                );
-                if let Some(availability_zone) = &conf.availability_zone {
-                    properties.insert(
-                        "availability_zone".to_string(),
-                        PostHogFlagFilterPropertyValue::String(availability_zone.clone()),
-                    );
-                }
-                // Infer region based on the remote storage config.
-                if let Some(remote_storage) = &conf.remote_storage_config {
-                    match &remote_storage.storage {
-                        RemoteStorageKind::AwsS3(config) => {
-                            properties.insert(
-                                "region".to_string(),
-                                PostHogFlagFilterPropertyValue::String(format!(
-                                    "aws-{}",
-                                    config.bucket_region
-                                )),
-                            );
-                        }
-                        RemoteStorageKind::AzureContainer(config) => {
-                            properties.insert(
-                                "region".to_string(),
-                                PostHogFlagFilterPropertyValue::String(format!(
-                                    "azure-{}",
-                                    config.container_region
-                                )),
-                            );
-                        }
-                        RemoteStorageKind::LocalFs { .. } => {
-                            properties.insert(
-                                "region".to_string(),
-                                PostHogFlagFilterPropertyValue::String("local".to_string()),
-                            );
-                        }
-                    }
-                }
-                // TODO: add pageserver URL.
-                Arc::new(properties)
-            };
-            let fake_tenants = {
-                let mut tenants = Vec::new();
-                for i in 0..10 {
-                    let distinct_id = format!(
-                        "fake_tenant_{}_{}_{}",
-                        conf.availability_zone.as_deref().unwrap_or_default(),
-                        conf.id,
-                        i
-                    );
-                    let properties = Self::collect_properties_inner(
-                        distinct_id.clone(),
-                        Some(&internal_properties),
-                    );
-                    tenants.push(CaptureEvent {
-                        event: "initial_tenant_report".to_string(),
-                        distinct_id,
-                        properties: json!({ "$set": properties }), // use `$set` to set the person properties instead of the event properties
-                    });
-                }
-                tenants
-            };
-            // TODO: make refresh period configurable
-            inner
-                .clone()
-                .spawn(handle, Duration::from_secs(60), fake_tenants);
-            Ok(FeatureResolver {
-                inner: Some(inner),
-                internal_properties: Some(internal_properties),
-            })
+            // TODO: make this configurable
+            inner.clone().spawn(handle, Duration::from_secs(60));
+            Ok(FeatureResolver { inner: Some(inner) })
         } else {
-            Ok(FeatureResolver {
-                inner: None,
-                internal_properties: None,
-            })
+            Ok(FeatureResolver { inner: None })
         }
     }
 
-    fn collect_properties_inner(
-        tenant_id: String,
-        internal_properties: Option<&HashMap<String, PostHogFlagFilterPropertyValue>>,
-    ) -> HashMap<String, PostHogFlagFilterPropertyValue> {
-        let mut properties = HashMap::new();
-        if let Some(internal_properties) = internal_properties {
-            for (key, value) in internal_properties.iter() {
-                properties.insert(key.clone(), value.clone());
-            }
-        }
-        properties.insert(
-            "tenant_id".to_string(),
-            PostHogFlagFilterPropertyValue::String(tenant_id),
-        );
-        properties
-    }
-
-    /// Collect all properties availble for the feature flag evaluation.
-    pub(crate) fn collect_properties(
-        &self,
-        tenant_id: TenantId,
-    ) -> HashMap<String, PostHogFlagFilterPropertyValue> {
-        Self::collect_properties_inner(tenant_id.to_string(), self.internal_properties.as_deref())
-    }
-
     /// Evaluate a multivariate feature flag. Currently, we do not support any properties.
     ///
     /// Error handling: the caller should inspect the error and decide the behavior when a feature flag
@@ -162,24 +55,11 @@ impl FeatureResolver {
         tenant_id: TenantId,
     ) -> Result<String, PostHogEvaluationError> {
         if let Some(inner) = &self.inner {
-            let res = inner.feature_store().evaluate_multivariate(
+            inner.feature_store().evaluate_multivariate(
                 flag_key,
                 &tenant_id.to_string(),
-                &self.collect_properties(tenant_id),
-            );
-            match &res {
-                Ok(value) => {
-                    FEATURE_FLAG_EVALUATION
-                        .with_label_values(&[flag_key, "ok", value])
-                        .inc();
-                }
-                Err(e) => {
-                    FEATURE_FLAG_EVALUATION
-                        .with_label_values(&[flag_key, "error", e.as_variant_str()])
-                        .inc();
-                }
-            }
-            res
+                &HashMap::new(),
+            )
         } else {
             Err(PostHogEvaluationError::NotAvailable(
                 "PostHog integration is not enabled".to_string(),
@@ -200,24 +80,11 @@ impl FeatureResolver {
         tenant_id: TenantId,
     ) -> Result<(), PostHogEvaluationError> {
         if let Some(inner) = &self.inner {
-            let res = inner.feature_store().evaluate_boolean(
+            inner.feature_store().evaluate_boolean(
                 flag_key,
                 &tenant_id.to_string(),
-                &self.collect_properties(tenant_id),
-            );
-            match &res {
-                Ok(()) => {
-                    FEATURE_FLAG_EVALUATION
-                        .with_label_values(&[flag_key, "ok", "true"])
-                        .inc();
-                }
-                Err(e) => {
-                    FEATURE_FLAG_EVALUATION
-                        .with_label_values(&[flag_key, "error", e.as_variant_str()])
-                        .inc();
-                }
-            }
-            res
+                &HashMap::new(),
+            )
         } else {
             Err(PostHogEvaluationError::NotAvailable(
                 "PostHog integration is not enabled".to_string(),

pageserver/src/http/routes.rs

@@ -43,7 +43,6 @@ use pageserver_api::models::{
 use pageserver_api::shard::{ShardCount, TenantShardId};
 use remote_storage::{DownloadError, GenericRemoteStorage, TimeTravelError};
 use scopeguard::defer;
-use serde_json::json;
 use tenant_size_model::svg::SvgBranchKind;
 use tenant_size_model::{SizeResult, StorageModel};
 use tokio::time::Instant;
@@ -3680,24 +3679,23 @@ async fn tenant_evaluate_feature_flag(
         let tenant = state
             .tenant_manager
             .get_attached_tenant_shard(tenant_shard_id)?;
-        let properties = tenant.feature_resolver.collect_properties(tenant_shard_id.tenant_id);
         if as_type == "boolean" {
             let result = tenant.feature_resolver.evaluate_boolean(&flag, tenant_shard_id.tenant_id);
             let result = result.map(|_| true).map_err(|e| e.to_string());
-            json_response(StatusCode::OK, json!({ "result": result, "properties": properties }))
+            json_response(StatusCode::OK, result)
         } else if as_type == "multivariate" {
             let result = tenant.feature_resolver.evaluate_multivariate(&flag, tenant_shard_id.tenant_id).map_err(|e| e.to_string());
-            json_response(StatusCode::OK, json!({ "result": result, "properties": properties }))
+            json_response(StatusCode::OK, result)
         } else {
             // Auto infer the type of the feature flag.
             let is_boolean = tenant.feature_resolver.is_feature_flag_boolean(&flag).map_err(|e| ApiError::InternalServerError(anyhow::anyhow!("{e}")))?;
             if is_boolean {
                 let result = tenant.feature_resolver.evaluate_boolean(&flag, tenant_shard_id.tenant_id);
                 let result = result.map(|_| true).map_err(|e| e.to_string());
-                json_response(StatusCode::OK, json!({ "result": result, "properties": properties }))
+                json_response(StatusCode::OK, result)
             } else {
                 let result = tenant.feature_resolver.evaluate_multivariate(&flag, tenant_shard_id.tenant_id).map_err(|e| e.to_string());
-                json_response(StatusCode::OK, json!({ "result": result, "properties": properties }))
+                json_response(StatusCode::OK, result)
             }
         }
     }

pageserver/src/metrics.rs

@@ -15,7 +15,6 @@ use metrics::{
     register_int_gauge, register_int_gauge_vec, register_uint_gauge, register_uint_gauge_vec,
 };
 use once_cell::sync::Lazy;
-use pageserver_api::config::defaults::DEFAULT_MAX_GET_VECTORED_KEYS;
 use pageserver_api::config::{
     PageServicePipeliningConfig, PageServicePipeliningConfigPipelined,
     PageServiceProtocolPipelinedBatchingStrategy, PageServiceProtocolPipelinedExecutionStrategy,
@@ -33,6 +32,7 @@ use crate::config::PageServerConf;
 use crate::context::{PageContentKind, RequestContext};
 use crate::pgdatadir_mapping::DatadirModificationStats;
 use crate::task_mgr::TaskKind;
+use crate::tenant::Timeline;
 use crate::tenant::layer_map::LayerMap;
 use crate::tenant::mgr::TenantSlot;
 use crate::tenant::storage_layer::{InMemoryLayer, PersistentLayerDesc};
@@ -446,15 +446,6 @@ static PAGE_CACHE_ERRORS: Lazy<IntCounterVec> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
-pub(crate) static FEATURE_FLAG_EVALUATION: Lazy<CounterVec> = Lazy::new(|| {
-    register_counter_vec!(
-        "pageserver_feature_flag_evaluation",
-        "Number of times a feature flag is evaluated",
-        &["flag_key", "status", "value"],
-    )
-    .unwrap()
-});
-
 #[derive(IntoStaticStr)]
 #[strum(serialize_all = "kebab_case")]
 pub(crate) enum PageCacheErrorKind {
@@ -1948,7 +1939,7 @@ static SMGR_QUERY_TIME_GLOBAL: Lazy<HistogramVec> = Lazy::new(|| {
 });
 
 static PAGE_SERVICE_BATCH_SIZE_BUCKETS_GLOBAL: Lazy<Vec<f64>> = Lazy::new(|| {
-    (1..=u32::try_from(DEFAULT_MAX_GET_VECTORED_KEYS).unwrap())
+    (1..=u32::try_from(Timeline::MAX_GET_VECTORED_KEYS).unwrap())
         .map(|v| v.into())
         .collect()
 });
@@ -1966,7 +1957,7 @@ static PAGE_SERVICE_BATCH_SIZE_BUCKETS_PER_TIMELINE: Lazy<Vec<f64>> = Lazy::new(
     let mut buckets = Vec::new();
     for i in 0.. {
         let bucket = 1 << i;
-        if bucket > u32::try_from(DEFAULT_MAX_GET_VECTORED_KEYS).unwrap() {
+        if bucket > u32::try_from(Timeline::MAX_GET_VECTORED_KEYS).unwrap() {
             break;
         }
         buckets.push(bucket.into());
@@ -2855,6 +2846,7 @@ pub(crate) struct WalIngestMetrics {
     pub(crate) records_received: IntCounter,
     pub(crate) records_observed: IntCounter,
     pub(crate) records_committed: IntCounter,
+    pub(crate) records_filtered: IntCounter,
     pub(crate) values_committed_metadata_images: IntCounter,
     pub(crate) values_committed_metadata_deltas: IntCounter,
     pub(crate) values_committed_data_images: IntCounter,
@@ -2910,6 +2902,11 @@ pub(crate) static WAL_INGEST: Lazy<WalIngestMetrics> = Lazy::new(|| {
         "Number of WAL records which resulted in writes to pageserver storage"
     )
     .expect("failed to define a metric"),
+    records_filtered: register_int_counter!(
+        "pageserver_wal_ingest_records_filtered",
+        "Number of WAL records filtered out due to sharding"
+    )
+    .expect("failed to define a metric"),
     values_committed_metadata_images: values_committed.with_label_values(&["metadata", "image"]),
     values_committed_metadata_deltas: values_committed.with_label_values(&["metadata", "delta"]),
     values_committed_data_images: values_committed.with_label_values(&["data", "image"]),

pageserver/src/page_service.rs

@@ -169,6 +169,96 @@ pub fn spawn(
     Listener { cancel, task }
 }
 
+/// Spawns a gRPC server for the page service.
+///
+/// TODO: move this onto GrpcPageServiceHandler::spawn().
+/// TODO: this doesn't support TLS. We need TLS reloading via ReloadingCertificateResolver, so we
+/// need to reimplement the TCP+TLS accept loop ourselves.
+pub fn spawn_grpc(
+    tenant_manager: Arc<TenantManager>,
+    auth: Option<Arc<SwappableJwtAuth>>,
+    perf_trace_dispatch: Option<Dispatch>,
+    listener: std::net::TcpListener,
+) -> anyhow::Result<CancellableTask> {
+    let cancel = CancellationToken::new();
+    let ctx = RequestContextBuilder::new(TaskKind::PageRequestHandler)
+        .download_behavior(DownloadBehavior::Download)
+        .perf_span_dispatch(perf_trace_dispatch)
+        .detached_child();
+    let gate = Gate::default();
+
+    // Set up the TCP socket. We take a preconfigured TcpListener to bind the
+    // port early during startup.
+    let incoming = {
+        let _runtime = COMPUTE_REQUEST_RUNTIME.enter(); // required by TcpListener::from_std
+        listener.set_nonblocking(true)?;
+        tonic::transport::server::TcpIncoming::from(tokio::net::TcpListener::from_std(listener)?)
+            .with_nodelay(Some(GRPC_TCP_NODELAY))
+            .with_keepalive(Some(GRPC_TCP_KEEPALIVE_TIME))
+    };
+
+    // Set up the gRPC server.
+    //
+    // TODO: consider tuning window sizes.
+    let mut server = tonic::transport::Server::builder()
+        .http2_keepalive_interval(Some(GRPC_HTTP2_KEEPALIVE_INTERVAL))
+        .http2_keepalive_timeout(Some(GRPC_HTTP2_KEEPALIVE_TIMEOUT))
+        .max_concurrent_streams(Some(GRPC_MAX_CONCURRENT_STREAMS));
+
+    // Main page service stack. Uses a mix of Tonic interceptors and Tower layers:
+    //
+    // * Interceptors: can inspect and modify the gRPC request. Sync code only, runs before service.
+    //
+    // * Layers: allow async code, can run code after the service response. However, only has access
+    //   to the raw HTTP request/response, not the gRPC types.
+    let page_service_handler = GrpcPageServiceHandler {
+        tenant_manager,
+        ctx,
+    };
+
+    let observability_layer = ObservabilityLayer;
+    let mut tenant_interceptor = TenantMetadataInterceptor;
+    let mut auth_interceptor = TenantAuthInterceptor::new(auth);
+
+    let page_service = tower::ServiceBuilder::new()
+        // Create tracing span and record request start time.
+        .layer(observability_layer)
+        // Intercept gRPC requests.
+        .layer(tonic::service::InterceptorLayer::new(move |mut req| {
+            // Extract tenant metadata.
+            req = tenant_interceptor.call(req)?;
+            // Authenticate tenant JWT token.
+            req = auth_interceptor.call(req)?;
+            Ok(req)
+        }))
+        .service(proto::PageServiceServer::new(page_service_handler));
+    let server = server.add_service(page_service);
+
+    // Reflection service for use with e.g. grpcurl.
+    let reflection_service = tonic_reflection::server::Builder::configure()
+        .register_encoded_file_descriptor_set(proto::FILE_DESCRIPTOR_SET)
+        .build_v1()?;
+    let server = server.add_service(reflection_service);
+
+    // Spawn server task.
+    let task_cancel = cancel.clone();
+    let task = COMPUTE_REQUEST_RUNTIME.spawn(task_mgr::exit_on_panic_or_error(
+        "grpc listener",
+        async move {
+            let result = server
+                .serve_with_incoming_shutdown(incoming, task_cancel.cancelled())
+                .await;
+            if result.is_ok() {
+                // TODO: revisit shutdown logic once page service is implemented.
+                gate.close().await;
+            }
+            result
+        },
+    ));
+
+    Ok(CancellableTask { task, cancel })
+}
+
 impl Listener {
     pub async fn stop_accepting(self) -> Connections {
         self.cancel.cancel();
@@ -407,6 +497,10 @@ async fn page_service_conn_main(
 }
 
 /// Page service connection handler.
+///
+/// TODO: for gRPC, this will be shared by all requests from all connections.
+/// Decompose it into global state and per-connection/request state, and make
+/// libpq-specific options (e.g. pipelining) separate.
 struct PageServerHandler {
     auth: Option<Arc<SwappableJwtAuth>>,
     claims: Option<Claims>,
@@ -671,7 +765,7 @@ impl PageStreamError {
             request_id,
             status_code,
             reason: Some(status.message().to_string()),
-            page_images: Vec::new(),
+            page_images: SmallVec::new(),
         }
         .into())
     }
@@ -3268,106 +3362,9 @@ where
 pub struct GrpcPageServiceHandler {
     tenant_manager: Arc<TenantManager>,
     ctx: RequestContext,
-    gate_guard: GateGuard,
-    get_vectored_concurrent_io: GetVectoredConcurrentIo,
 }
 
 impl GrpcPageServiceHandler {
-    /// Spawns a gRPC server for the page service.
-    ///
-    /// TODO: this doesn't support TLS. We need TLS reloading via ReloadingCertificateResolver, so we
-    /// need to reimplement the TCP+TLS accept loop ourselves.
-    pub fn spawn(
-        tenant_manager: Arc<TenantManager>,
-        auth: Option<Arc<SwappableJwtAuth>>,
-        perf_trace_dispatch: Option<Dispatch>,
-        get_vectored_concurrent_io: GetVectoredConcurrentIo,
-        listener: std::net::TcpListener,
-    ) -> anyhow::Result<CancellableTask> {
-        let cancel = CancellationToken::new();
-        let ctx = RequestContextBuilder::new(TaskKind::PageRequestHandler)
-            .download_behavior(DownloadBehavior::Download)
-            .perf_span_dispatch(perf_trace_dispatch)
-            .detached_child();
-        let gate = Gate::default();
-
-        // Set up the TCP socket. We take a preconfigured TcpListener to bind the
-        // port early during startup.
-        let incoming = {
-            let _runtime = COMPUTE_REQUEST_RUNTIME.enter(); // required by TcpListener::from_std
-            listener.set_nonblocking(true)?;
-            tonic::transport::server::TcpIncoming::from(tokio::net::TcpListener::from_std(
-                listener,
-            )?)
-            .with_nodelay(Some(GRPC_TCP_NODELAY))
-            .with_keepalive(Some(GRPC_TCP_KEEPALIVE_TIME))
-        };
-
-        // Set up the gRPC server.
-        //
-        // TODO: consider tuning window sizes.
-        let mut server = tonic::transport::Server::builder()
-            .http2_keepalive_interval(Some(GRPC_HTTP2_KEEPALIVE_INTERVAL))
-            .http2_keepalive_timeout(Some(GRPC_HTTP2_KEEPALIVE_TIMEOUT))
-            .max_concurrent_streams(Some(GRPC_MAX_CONCURRENT_STREAMS));
-
-        // Main page service stack. Uses a mix of Tonic interceptors and Tower layers:
-        //
-        // * Interceptors: can inspect and modify the gRPC request. Sync code only, runs before service.
-        //
-        // * Layers: allow async code, can run code after the service response. However, only has access
-        //   to the raw HTTP request/response, not the gRPC types.
-        let page_service_handler = GrpcPageServiceHandler {
-            tenant_manager,
-            ctx,
-            gate_guard: gate.enter().expect("gate was just created"),
-            get_vectored_concurrent_io,
-        };
-
-        let observability_layer = ObservabilityLayer;
-        let mut tenant_interceptor = TenantMetadataInterceptor;
-        let mut auth_interceptor = TenantAuthInterceptor::new(auth);
-
-        let page_service = tower::ServiceBuilder::new()
-            // Create tracing span and record request start time.
-            .layer(observability_layer)
-            // Intercept gRPC requests.
-            .layer(tonic::service::InterceptorLayer::new(move |mut req| {
-                // Extract tenant metadata.
-                req = tenant_interceptor.call(req)?;
-                // Authenticate tenant JWT token.
-                req = auth_interceptor.call(req)?;
-                Ok(req)
-            }))
-            // Run the page service.
-            .service(proto::PageServiceServer::new(page_service_handler));
-        let server = server.add_service(page_service);
-
-        // Reflection service for use with e.g. grpcurl.
-        let reflection_service = tonic_reflection::server::Builder::configure()
-            .register_encoded_file_descriptor_set(proto::FILE_DESCRIPTOR_SET)
-            .build_v1()?;
-        let server = server.add_service(reflection_service);
-
-        // Spawn server task.
-        let task_cancel = cancel.clone();
-        let task = COMPUTE_REQUEST_RUNTIME.spawn(task_mgr::exit_on_panic_or_error(
-            "grpc listener",
-            async move {
-                let result = server
-                    .serve_with_incoming_shutdown(incoming, task_cancel.cancelled())
-                    .await;
-                if result.is_ok() {
-                    // TODO: revisit shutdown logic once page service is implemented.
-                    gate.close().await;
-                }
-                result
-            },
-        ));
-
-        Ok(CancellableTask { task, cancel })
-    }
-
     /// Errors if the request is executed on a non-zero shard. Only shard 0 has a complete view of
     /// relations and their sizes, as well as SLRU segments and similar data.
     #[allow(clippy::result_large_err)]
@@ -3524,7 +3521,7 @@ impl GrpcPageServiceHandler {
             request_id: req.request_id,
             status_code: page_api::GetPageStatusCode::Ok,
             reason: None,
-            page_images: Vec::with_capacity(results.len()),
+            page_images: SmallVec::with_capacity(results.len()),
         };
 
         for result in results {
@@ -3663,7 +3660,7 @@ impl proto::PageService for GrpcPageServiceHandler {
                     if chunk.is_empty() {
                         break;
                     }
-                    yield proto::GetBaseBackupResponseChunk::from(chunk.clone().freeze());
+                    yield proto::GetBaseBackupResponseChunk::try_from(chunk.clone().freeze())?;
                     chunk.clear();
                 }
             }
@@ -3724,14 +3721,6 @@ impl proto::PageService for GrpcPageServiceHandler {
             .get(ttid.tenant_id, ttid.timeline_id, shard_selector)
             .await?;
 
-        // Spawn an IoConcurrency sidecar, if enabled.
-        let Ok(gate_guard) = self.gate_guard.try_clone() else {
-            return Err(tonic::Status::unavailable("shutting down"));
-        };
-        let io_concurrency =
-            IoConcurrency::spawn_from_conf(self.get_vectored_concurrent_io, gate_guard);
-
-        // Spawn a task to handle the GetPageRequest stream.
         let span = Span::current();
         let ctx = self.ctx.attached_child();
         let mut reqs = req.into_inner();
@@ -3742,7 +3731,8 @@ impl proto::PageService for GrpcPageServiceHandler {
                 .await?
                 .downgrade();
             while let Some(req) = reqs.message().await? {
-                yield Self::get_page(&ctx, &timeline, req, io_concurrency.clone())
+                // TODO: implement IoConcurrency sidecar.
+                yield Self::get_page(&ctx, &timeline, req, IoConcurrency::Sequential)
                     .instrument(span.clone()) // propagate request span
                     .await?
             }
@@ -3816,7 +3806,7 @@ impl proto::PageService for GrpcPageServiceHandler {
         let resp =
             PageServerHandler::handle_get_slru_segment_request(&timeline, &req, &ctx).await?;
         let resp: page_api::GetSlruSegmentResponse = resp.segment;
-        Ok(tonic::Response::new(resp.into()))
+        Ok(tonic::Response::new(resp.try_into()?))
     }
 }
 

pageserver/src/pgdatadir_mapping.rs

@@ -431,10 +431,10 @@ impl Timeline {
                         GetVectoredError::InvalidLsn(e) => {
                             Err(anyhow::anyhow!("invalid LSN: {e:?}").into())
                         }
-                        // NB: this should never happen in practice because we limit batch size to be smaller than max_get_vectored_keys
+                        // NB: this should never happen in practice because we limit MAX_GET_VECTORED_KEYS
                         // TODO: we can prevent this error class by moving this check into the type system
-                        GetVectoredError::Oversized(err, max) => {
-                            Err(anyhow::anyhow!("batching oversized: {err} > {max}").into())
+                        GetVectoredError::Oversized(err) => {
+                            Err(anyhow::anyhow!("batching oversized: {err:?}").into())
                         }
                     };
 
@@ -719,7 +719,7 @@ impl Timeline {
 
         let batches = keyspace.partition(
             self.get_shard_identity(),
-            self.conf.max_get_vectored_keys.get() as u64 * BLCKSZ as u64,
+            Timeline::MAX_GET_VECTORED_KEYS * BLCKSZ as u64,
         );
 
         let io_concurrency = IoConcurrency::spawn_from_conf(
@@ -959,7 +959,7 @@ impl Timeline {
 
             let batches = keyspace.partition(
                 self.get_shard_identity(),
-                self.conf.max_get_vectored_keys.get() as u64 * BLCKSZ as u64,
+                Timeline::MAX_GET_VECTORED_KEYS * BLCKSZ as u64,
             );
 
             let io_concurrency = IoConcurrency::spawn_from_conf(

pageserver/src/tenant.rs

@@ -7197,7 +7197,7 @@ mod tests {
             let end = desc
                 .key_range
                 .start
-                .add(tenant.conf.max_get_vectored_keys.get() as u32);
+                .add(Timeline::MAX_GET_VECTORED_KEYS.try_into().unwrap());
             reads.push(KeySpace {
                 ranges: vec![start..end],
             });
@@ -11260,11 +11260,11 @@ mod tests {
                 let mut keyspaces_at_lsn: HashMap<Lsn, KeySpaceRandomAccum> = HashMap::default();
                 let mut used_keys: HashSet<Key> = HashSet::default();
 
-                while used_keys.len() < tenant.conf.max_get_vectored_keys.get() {
+                while used_keys.len() < Timeline::MAX_GET_VECTORED_KEYS as usize {
                     let selected_lsn = interesting_lsns.choose(&mut random).expect("not empty");
                     let mut selected_key = start_key.add(random.gen_range(0..KEY_DIMENSION_SIZE));
 
-                    while used_keys.len() < tenant.conf.max_get_vectored_keys.get() {
+                    while used_keys.len() < Timeline::MAX_GET_VECTORED_KEYS as usize {
                         if used_keys.contains(&selected_key)
                             || selected_key >= start_key.add(KEY_DIMENSION_SIZE)
                         {

pageserver/src/tenant/timeline.rs

@@ -817,8 +817,8 @@ pub(crate) enum GetVectoredError {
     #[error("timeline shutting down")]
     Cancelled,
 
-    #[error("requested too many keys: {0} > {1}")]
-    Oversized(u64, u64),
+    #[error("requested too many keys: {0} > {}", Timeline::MAX_GET_VECTORED_KEYS)]
+    Oversized(u64),
 
     #[error("requested at invalid LSN: {0}")]
     InvalidLsn(Lsn),
@@ -1019,7 +1019,7 @@ impl From<GetVectoredError> for PageReconstructError {
         match e {
             GetVectoredError::Cancelled => PageReconstructError::Cancelled,
             GetVectoredError::InvalidLsn(_) => PageReconstructError::Other(anyhow!("Invalid LSN")),
-            err @ GetVectoredError::Oversized(_, _) => PageReconstructError::Other(err.into()),
+            err @ GetVectoredError::Oversized(_) => PageReconstructError::Other(err.into()),
             GetVectoredError::MissingKey(err) => PageReconstructError::MissingKey(err),
             GetVectoredError::GetReadyAncestorError(err) => PageReconstructError::from(err),
             GetVectoredError::Other(err) => PageReconstructError::Other(err),
@@ -1199,6 +1199,7 @@ impl Timeline {
         }
     }
 
+    pub(crate) const MAX_GET_VECTORED_KEYS: u64 = 32;
     pub(crate) const LAYERS_VISITED_WARN_THRESHOLD: u32 = 100;
 
     /// Look up multiple page versions at a given LSN
@@ -1213,12 +1214,9 @@ impl Timeline {
     ) -> Result<BTreeMap<Key, Result<Bytes, PageReconstructError>>, GetVectoredError> {
         let total_keyspace = query.total_keyspace();
 
-        let key_count = total_keyspace.total_raw_size();
-        if key_count > self.conf.max_get_vectored_keys.get() {
-            return Err(GetVectoredError::Oversized(
-                key_count as u64,
-                self.conf.max_get_vectored_keys.get() as u64,
-            ));
+        let key_count = total_keyspace.total_raw_size().try_into().unwrap();
+        if key_count > Timeline::MAX_GET_VECTORED_KEYS {
+            return Err(GetVectoredError::Oversized(key_count));
         }
 
         for range in &total_keyspace.ranges {
@@ -2506,13 +2504,6 @@ impl Timeline {
             // Preparing basebackup doesn't make sense for shards other than shard zero.
             return;
         }
-        if !self.is_active() {
-            // May happen during initial timeline creation.
-            // Such timeline is not in the global timeline map yet,
-            // so basebackup cache will not be able to find it.
-            // TODO(diko): We can prepare such timelines in finish_creation().
-            return;
-        }
 
         let res = self
             .basebackup_prepare_sender
@@ -2852,6 +2843,21 @@ impl Timeline {
             )
     }
 
+    /// Resolve the effective WAL receiver protocol to use for this tenant.
+    ///
+    /// Priority order is:
+    /// 1. Tenant config override
+    /// 2. Default value for tenant config override
+    /// 3. Pageserver config override
+    /// 4. Pageserver config default
+    pub fn resolve_wal_receiver_protocol(&self) -> PostgresClientProtocol {
+        let tenant_conf = self.tenant_conf.load().tenant_conf.clone();
+        tenant_conf
+            .wal_receiver_protocol_override
+            .or(self.conf.default_tenant_conf.wal_receiver_protocol_override)
+            .unwrap_or(self.conf.wal_receiver_protocol)
+    }
+
     pub(super) fn tenant_conf_updated(&self, new_conf: &AttachedTenantConf) {
         // NB: Most tenant conf options are read by background loops, so,
         // changes will automatically be picked up.
@@ -3207,16 +3213,10 @@ impl Timeline {
             guard.is_none(),
             "multiple launches / re-launches of WAL receiver are not supported"
         );
-
-        let protocol = PostgresClientProtocol::Interpreted {
-            format: utils::postgres_client::InterpretedFormat::Protobuf,
-            compression: Some(utils::postgres_client::Compression::Zstd { level: 1 }),
-        };
-
         *guard = Some(WalReceiver::start(
             Arc::clone(self),
             WalReceiverConf {
-                protocol,
+                protocol: self.resolve_wal_receiver_protocol(),
                 wal_connect_timeout,
                 lagging_wal_timeout,
                 max_lsn_wal_lag,
@@ -5270,7 +5270,7 @@ impl Timeline {
                 key = key.next();
 
                 // Maybe flush `key_rest_accum`
-                if key_request_accum.raw_size() >= self.conf.max_get_vectored_keys.get() as u64
+                if key_request_accum.raw_size() >= Timeline::MAX_GET_VECTORED_KEYS
                     || (last_key_in_range && key_request_accum.raw_size() > 0)
                 {
                     let query =

pageserver/src/tenant/timeline/import_pgdata.rs

@@ -201,8 +201,8 @@ async fn prepare_import(
         .await;
         match res {
             Ok(_) => break,
-            Err(_err) => {
-                info!("indefinitely waiting for pgdata to finish");
+            Err(err) => {
+                info!(?err, "indefinitely waiting for pgdata to finish");
                 if tokio::time::timeout(std::time::Duration::from_secs(10), cancel.cancelled())
                     .await
                     .is_ok()

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

@@ -100,7 +100,6 @@ async fn run_v1(
                         .unwrap(),
                     import_job_concurrency: base.import_job_concurrency,
                     import_job_checkpoint_threshold: base.import_job_checkpoint_threshold,
-                    import_job_max_byte_range_size: base.import_job_max_byte_range_size,
                 }
             }
             None => timeline.conf.timeline_import_config.clone(),
@@ -442,7 +441,6 @@ impl Plan {
 
         let mut last_completed_job_idx = start_after_job_idx.unwrap_or(0);
         let checkpoint_every: usize = import_config.import_job_checkpoint_threshold.into();
-        let max_byte_range_size: usize = import_config.import_job_max_byte_range_size.into();
 
         // Run import jobs concurrently up to the limit specified by the pageserver configuration.
         // Note that we process completed futures in the oreder of insertion. This will be the
@@ -458,7 +456,7 @@ impl Plan {
 
                     work.push_back(tokio::task::spawn(async move {
                         let _permit = permit;
-                        let res = job.run(job_timeline, max_byte_range_size, &ctx).await;
+                        let res = job.run(job_timeline, &ctx).await;
                         (job_idx, res)
                     }));
                 },
@@ -473,8 +471,6 @@ impl Plan {
                             last_completed_job_idx = job_idx;
 
                             if last_completed_job_idx % checkpoint_every == 0 {
-                                tracing::info!(last_completed_job_idx, jobs=%jobs_in_plan, "Checkpointing import status");
-
                                 let progress = ShardImportProgressV1 {
                                     jobs: jobs_in_plan,
                                     completed: last_completed_job_idx,
@@ -496,6 +492,8 @@ impl Plan {
                                     anyhow::anyhow!("Shut down while putting timeline import status")
                                 })?;
                             }
+
+                            tracing::info!(last_completed_job_idx, jobs=%jobs_in_plan, "Checkpointing import status");
                         },
                         Some(Err(_)) => {
                             anyhow::bail!(
@@ -681,7 +679,6 @@ trait ImportTask {
     async fn doit(
         self,
         layer_writer: &mut ImageLayerWriter,
-        max_byte_range_size: usize,
         ctx: &RequestContext,
     ) -> anyhow::Result<usize>;
 }
@@ -718,7 +715,6 @@ impl ImportTask for ImportSingleKeyTask {
     async fn doit(
         self,
         layer_writer: &mut ImageLayerWriter,
-        _max_byte_range_size: usize,
         ctx: &RequestContext,
     ) -> anyhow::Result<usize> {
         layer_writer.put_image(self.key, self.buf, ctx).await?;
@@ -772,9 +768,10 @@ impl ImportTask for ImportRelBlocksTask {
     async fn doit(
         self,
         layer_writer: &mut ImageLayerWriter,
-        max_byte_range_size: usize,
         ctx: &RequestContext,
     ) -> anyhow::Result<usize> {
+        const MAX_BYTE_RANGE_SIZE: usize = 4 * 1024 * 1024;
+
         debug!("Importing relation file");
 
         let (rel_tag, start_blk) = self.key_range.start.to_rel_block()?;
@@ -799,7 +796,7 @@ impl ImportTask for ImportRelBlocksTask {
                 assert_eq!(key.len(), 1);
                 assert!(!acc.is_empty());
                 assert!(acc_end > acc_start);
-                if acc_end == start && end - acc_start <= max_byte_range_size {
+                if acc_end == start && end - acc_start <= MAX_BYTE_RANGE_SIZE {
                     acc.push(key.pop().unwrap());
                     Ok((acc, acc_start, end))
                 } else {
@@ -863,7 +860,6 @@ impl ImportTask for ImportSlruBlocksTask {
     async fn doit(
         self,
         layer_writer: &mut ImageLayerWriter,
-        _max_byte_range_size: usize,
         ctx: &RequestContext,
     ) -> anyhow::Result<usize> {
         debug!("Importing SLRU segment file {}", self.path);
@@ -910,13 +906,12 @@ impl ImportTask for AnyImportTask {
     async fn doit(
         self,
         layer_writer: &mut ImageLayerWriter,
-        max_byte_range_size: usize,
         ctx: &RequestContext,
     ) -> anyhow::Result<usize> {
         match self {
-            Self::SingleKey(t) => t.doit(layer_writer, max_byte_range_size, ctx).await,
-            Self::RelBlocks(t) => t.doit(layer_writer, max_byte_range_size, ctx).await,
-            Self::SlruBlocks(t) => t.doit(layer_writer, max_byte_range_size, ctx).await,
+            Self::SingleKey(t) => t.doit(layer_writer, ctx).await,
+            Self::RelBlocks(t) => t.doit(layer_writer, ctx).await,
+            Self::SlruBlocks(t) => t.doit(layer_writer, ctx).await,
         }
     }
 }
@@ -957,12 +952,7 @@ impl ChunkProcessingJob {
         }
     }
 
-    async fn run(
-        self,
-        timeline: Arc<Timeline>,
-        max_byte_range_size: usize,
-        ctx: &RequestContext,
-    ) -> anyhow::Result<()> {
+    async fn run(self, timeline: Arc<Timeline>, ctx: &RequestContext) -> anyhow::Result<()> {
         let mut writer = ImageLayerWriter::new(
             timeline.conf,
             timeline.timeline_id,
@@ -977,7 +967,7 @@ impl ChunkProcessingJob {
 
         let mut nimages = 0;
         for task in self.tasks {
-            nimages += task.doit(&mut writer, max_byte_range_size, ctx).await?;
+            nimages += task.doit(&mut writer, ctx).await?;
         }
 
         let resident_layer = if nimages > 0 {

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

@@ -32,7 +32,9 @@ use utils::backoff::{
 };
 use utils::id::{NodeId, TenantTimelineId};
 use utils::lsn::Lsn;
-use utils::postgres_client::{ConnectionConfigArgs, wal_stream_connection_config};
+use utils::postgres_client::{
+    ConnectionConfigArgs, PostgresClientProtocol, wal_stream_connection_config,
+};
 
 use super::walreceiver_connection::{WalConnectionStatus, WalReceiverError};
 use super::{TaskEvent, TaskHandle, TaskStateUpdate, WalReceiverConf};
@@ -989,12 +991,19 @@ impl ConnectionManagerState {
                     return None; // no connection string, ignore sk
                 }
 
-                let shard_identity = self.timeline.get_shard_identity();
-                let (shard_number, shard_count, shard_stripe_size) = (
-                    Some(shard_identity.number.0),
-                    Some(shard_identity.count.0),
-                    Some(shard_identity.stripe_size.0),
-                );
+                let (shard_number, shard_count, shard_stripe_size) = match self.conf.protocol {
+                    PostgresClientProtocol::Vanilla => {
+                        (None, None, None)
+                    },
+                    PostgresClientProtocol::Interpreted { .. } => {
+                        let shard_identity = self.timeline.get_shard_identity();
+                        (
+                            Some(shard_identity.number.0),
+                            Some(shard_identity.count.0),
+                            Some(shard_identity.stripe_size.0),
+                        )
+                    }
+                };
 
                 let connection_conf_args = ConnectionConfigArgs {
                     protocol: self.conf.protocol,
@@ -1111,8 +1120,8 @@ impl ReconnectReason {
 
 #[cfg(test)]
 mod tests {
+    use pageserver_api::config::defaults::DEFAULT_WAL_RECEIVER_PROTOCOL;
     use url::Host;
-    use utils::postgres_client::PostgresClientProtocol;
 
     use super::*;
     use crate::tenant::harness::{TIMELINE_ID, TenantHarness};
@@ -1543,11 +1552,6 @@ mod tests {
             .await
             .expect("Failed to create an empty timeline for dummy wal connection manager");
 
-        let protocol = PostgresClientProtocol::Interpreted {
-            format: utils::postgres_client::InterpretedFormat::Protobuf,
-            compression: Some(utils::postgres_client::Compression::Zstd { level: 1 }),
-        };
-
         ConnectionManagerState {
             id: TenantTimelineId {
                 tenant_id: harness.tenant_shard_id.tenant_id,
@@ -1556,7 +1560,7 @@ mod tests {
             timeline,
             cancel: CancellationToken::new(),
             conf: WalReceiverConf {
-                protocol,
+                protocol: DEFAULT_WAL_RECEIVER_PROTOCOL,
                 wal_connect_timeout: Duration::from_secs(1),
                 lagging_wal_timeout: Duration::from_secs(1),
                 max_lsn_wal_lag: NonZeroU64::new(1024 * 1024).unwrap(),

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

@@ -15,7 +15,7 @@ use postgres_backend::is_expected_io_error;
 use postgres_connection::PgConnectionConfig;
 use postgres_ffi::WAL_SEGMENT_SIZE;
 use postgres_ffi::v14::xlog_utils::normalize_lsn;
-use postgres_ffi::waldecoder::WalDecodeError;
+use postgres_ffi::waldecoder::{WalDecodeError, WalStreamDecoder};
 use postgres_protocol::message::backend::ReplicationMessage;
 use postgres_types::PgLsn;
 use tokio::sync::watch;
@@ -31,7 +31,7 @@ use utils::lsn::Lsn;
 use utils::pageserver_feedback::PageserverFeedback;
 use utils::postgres_client::PostgresClientProtocol;
 use utils::sync::gate::GateError;
-use wal_decoder::models::{FlushUncommittedRecords, InterpretedWalRecords};
+use wal_decoder::models::{FlushUncommittedRecords, InterpretedWalRecord, InterpretedWalRecords};
 use wal_decoder::wire_format::FromWireFormat;
 
 use super::TaskStateUpdate;
@@ -275,6 +275,8 @@ pub(super) async fn handle_walreceiver_connection(
     let copy_stream = replication_client.copy_both_simple(&query).await?;
     let mut physical_stream = pin!(ReplicationStream::new(copy_stream));
 
+    let mut waldecoder = WalStreamDecoder::new(startpoint, timeline.pg_version);
+
     let mut walingest = WalIngest::new(timeline.as_ref(), startpoint, &ctx)
         .await
         .map_err(|e| match e.kind {
@@ -282,16 +284,14 @@ pub(super) async fn handle_walreceiver_connection(
             _ => WalReceiverError::Other(e.into()),
         })?;
 
-    let (format, compression) = match protocol {
+    let shard = vec![*timeline.get_shard_identity()];
+
+    let interpreted_proto_config = match protocol {
+        PostgresClientProtocol::Vanilla => None,
         PostgresClientProtocol::Interpreted {
             format,
             compression,
-        } => (format, compression),
-        PostgresClientProtocol::Vanilla => {
-            return Err(WalReceiverError::Other(anyhow!(
-                "Vanilla WAL receiver protocol is no longer supported for ingest"
-            )));
-        }
+        } => Some((format, compression)),
     };
 
     let mut expected_wal_start = startpoint;
@@ -313,6 +313,16 @@ pub(super) async fn handle_walreceiver_connection(
         // Update the connection status before processing the message. If the message processing
         // fails (e.g. in walingest), we still want to know latests LSNs from the safekeeper.
         match &replication_message {
+            ReplicationMessage::XLogData(xlog_data) => {
+                connection_status.latest_connection_update = now;
+                connection_status.commit_lsn = Some(Lsn::from(xlog_data.wal_end()));
+                connection_status.streaming_lsn = Some(Lsn::from(
+                    xlog_data.wal_start() + xlog_data.data().len() as u64,
+                ));
+                if !xlog_data.data().is_empty() {
+                    connection_status.latest_wal_update = now;
+                }
+            }
             ReplicationMessage::PrimaryKeepAlive(keepalive) => {
                 connection_status.latest_connection_update = now;
                 connection_status.commit_lsn = Some(Lsn::from(keepalive.wal_end()));
@@ -343,6 +353,7 @@ pub(super) async fn handle_walreceiver_connection(
                 // were interpreted.
                 let streaming_lsn = Lsn::from(raw.streaming_lsn());
 
+                let (format, compression) = interpreted_proto_config.unwrap();
                 let batch = InterpretedWalRecords::from_wire(raw.data(), format, compression)
                     .await
                     .with_context(|| {
@@ -498,6 +509,138 @@ pub(super) async fn handle_walreceiver_connection(
                 Some(streaming_lsn)
             }
 
+            ReplicationMessage::XLogData(xlog_data) => {
+                async fn commit(
+                    modification: &mut DatadirModification<'_>,
+                    uncommitted: &mut u64,
+                    filtered: &mut u64,
+                    ctx: &RequestContext,
+                ) -> anyhow::Result<()> {
+                    let stats = modification.stats();
+                    modification.commit(ctx).await?;
+                    WAL_INGEST
+                        .records_committed
+                        .inc_by(*uncommitted - *filtered);
+                    WAL_INGEST.inc_values_committed(&stats);
+                    *uncommitted = 0;
+                    *filtered = 0;
+                    Ok(())
+                }
+
+                // Pass the WAL data to the decoder, and see if we can decode
+                // more records as a result.
+                let data = xlog_data.data();
+                let startlsn = Lsn::from(xlog_data.wal_start());
+                let endlsn = startlsn + data.len() as u64;
+
+                trace!("received XLogData between {startlsn} and {endlsn}");
+
+                WAL_INGEST.bytes_received.inc_by(data.len() as u64);
+                waldecoder.feed_bytes(data);
+
+                {
+                    let mut modification = timeline.begin_modification(startlsn);
+                    let mut uncommitted_records = 0;
+                    let mut filtered_records = 0;
+
+                    while let Some((next_record_lsn, recdata)) = waldecoder.poll_decode()? {
+                        // It is important to deal with the aligned records as lsn in getPage@LSN is
+                        // aligned and can be several bytes bigger. Without this alignment we are
+                        // at risk of hitting a deadlock.
+                        if !next_record_lsn.is_aligned() {
+                            return Err(WalReceiverError::Other(anyhow!("LSN not aligned")));
+                        }
+
+                        // Deserialize and interpret WAL record
+                        let interpreted = InterpretedWalRecord::from_bytes_filtered(
+                            recdata,
+                            &shard,
+                            next_record_lsn,
+                            modification.tline.pg_version,
+                        )?
+                        .remove(timeline.get_shard_identity())
+                        .unwrap();
+
+                        if matches!(interpreted.flush_uncommitted, FlushUncommittedRecords::Yes)
+                            && uncommitted_records > 0
+                        {
+                            // Special case: legacy PG database creations operate by reading pages from a 'template' database:
+                            // these are the only kinds of WAL record that require reading data blocks while ingesting.  Ensure
+                            // all earlier writes of data blocks are visible by committing any modification in flight.
+                            commit(
+                                &mut modification,
+                                &mut uncommitted_records,
+                                &mut filtered_records,
+                                &ctx,
+                            )
+                            .await?;
+                        }
+
+                        // Ingest the records without immediately committing them.
+                        timeline.metrics.wal_records_received.inc();
+                        let ingested = walingest
+                            .ingest_record(interpreted, &mut modification, &ctx)
+                            .await
+                            .with_context(|| {
+                                format!("could not ingest record at {next_record_lsn}")
+                            })
+                            .inspect_err(|err| {
+                                // TODO: we can't differentiate cancellation errors with
+                                // anyhow::Error, so just ignore it if we're cancelled.
+                                if !cancellation.is_cancelled() && !timeline.is_stopping() {
+                                    critical!("{err:?}")
+                                }
+                            })?;
+                        if !ingested {
+                            tracing::debug!("ingest: filtered out record @ LSN {next_record_lsn}");
+                            WAL_INGEST.records_filtered.inc();
+                            filtered_records += 1;
+                        }
+
+                        // FIXME: this cannot be made pausable_failpoint without fixing the
+                        // failpoint library; in tests, the added amount of debugging will cause us
+                        // to timeout the tests.
+                        fail_point!("walreceiver-after-ingest");
+
+                        last_rec_lsn = next_record_lsn;
+
+                        // Commit every ingest_batch_size records. Even if we filtered out
+                        // all records, we still need to call commit to advance the LSN.
+                        uncommitted_records += 1;
+                        if uncommitted_records >= ingest_batch_size
+                            || modification.approx_pending_bytes()
+                                > DatadirModification::MAX_PENDING_BYTES
+                        {
+                            commit(
+                                &mut modification,
+                                &mut uncommitted_records,
+                                &mut filtered_records,
+                                &ctx,
+                            )
+                            .await?;
+                        }
+                    }
+
+                    // Commit the remaining records.
+                    if uncommitted_records > 0 {
+                        commit(
+                            &mut modification,
+                            &mut uncommitted_records,
+                            &mut filtered_records,
+                            &ctx,
+                        )
+                        .await?;
+                    }
+                }
+
+                if !caught_up && endlsn >= end_of_wal {
+                    info!("caught up at LSN {endlsn}");
+                    caught_up = true;
+                }
+
+                Some(endlsn)
+            }
+
             ReplicationMessage::PrimaryKeepAlive(keepalive) => {
                 let wal_end = keepalive.wal_end();
                 let timestamp = keepalive.timestamp();

pgxn/neon/neon.c

@@ -16,7 +16,6 @@
 #if PG_MAJORVERSION_NUM >= 15
 #include "access/xlogrecovery.h"
 #endif
-#include "executor/instrument.h"
 #include "replication/logical.h"
 #include "replication/logicallauncher.h"
 #include "replication/slot.h"
@@ -34,7 +33,6 @@
 #include "file_cache.h"
 #include "neon.h"
 #include "neon_lwlsncache.h"
-#include "neon_perf_counters.h"
 #include "control_plane_connector.h"
 #include "logical_replication_monitor.h"
 #include "unstable_extensions.h"
@@ -48,13 +46,6 @@ void		_PG_init(void);
 
 
 static int  running_xacts_overflow_policy;
-static bool monitor_query_exec_time = false;
-
-static ExecutorStart_hook_type prev_ExecutorStart = NULL;
-static ExecutorEnd_hook_type prev_ExecutorEnd = NULL;
-
-static void neon_ExecutorStart(QueryDesc *queryDesc, int eflags);
-static void neon_ExecutorEnd(QueryDesc *queryDesc);
 
 #if PG_MAJORVERSION_NUM >= 16
 static shmem_startup_hook_type prev_shmem_startup_hook;
@@ -479,16 +470,6 @@ _PG_init(void)
 							0,
 							NULL, NULL, NULL);
 
-	DefineCustomBoolVariable(
-							"neon.monitor_query_exec_time",
-							"Collect infortmation about query execution time",
-							NULL,
-							&monitor_query_exec_time,
-							false,
-							PGC_USERSET,
-							0,
-							NULL, NULL, NULL);
-
 	DefineCustomBoolVariable(
 							"neon.allow_replica_misconfig",
 							"Allow replica startup when some critical GUCs have smaller value than on primary node",
@@ -527,11 +508,6 @@ _PG_init(void)
 	EmitWarningsOnPlaceholders("neon");
 
 	ReportSearchPath();
-
-	prev_ExecutorStart = ExecutorStart_hook;
-	ExecutorStart_hook = neon_ExecutorStart;
-	prev_ExecutorEnd = ExecutorEnd_hook;
-	ExecutorEnd_hook = neon_ExecutorEnd;
 }
 
 PG_FUNCTION_INFO_V1(pg_cluster_size);
@@ -605,55 +581,3 @@ neon_shmem_startup_hook(void)
 #endif
 }
 #endif
-
-/*
- * ExecutorStart hook: start up tracking if needed
- */
-static void
-neon_ExecutorStart(QueryDesc *queryDesc, int eflags)
-{
-	if (prev_ExecutorStart)
-		prev_ExecutorStart(queryDesc, eflags);
-	else
-		standard_ExecutorStart(queryDesc, eflags);
-
-	if (monitor_query_exec_time)
-	{
-		/*
-		 * Set up to track total elapsed time in ExecutorRun.  Make sure the
-		 * space is allocated in the per-query context so it will go away at
-		 * ExecutorEnd.
-		 */
-		if (queryDesc->totaltime == NULL)
-		{
-			MemoryContext oldcxt;
-
-			oldcxt = MemoryContextSwitchTo(queryDesc->estate->es_query_cxt);
-			queryDesc->totaltime = InstrAlloc(1, INSTRUMENT_TIMER, false);
-			MemoryContextSwitchTo(oldcxt);
-		}
-	}
-}
-
-/*
- * ExecutorEnd hook: store results if needed
- */
-static void
-neon_ExecutorEnd(QueryDesc *queryDesc)
-{
-	if (monitor_query_exec_time && queryDesc->totaltime)
-	{
-		/*
-		 * Make sure stats accumulation is done.  (Note: it's okay if several
-		 * levels of hook all do this.)
-		 */
-		InstrEndLoop(queryDesc->totaltime);
-
-		inc_query_time(queryDesc->totaltime->total*1000000); /* convert to usec */
-	}
-
-	if (prev_ExecutorEnd)
-		prev_ExecutorEnd(queryDesc);
-	else
-		standard_ExecutorEnd(queryDesc);
-}

pgxn/neon/neon_perf_counters.c

@@ -71,27 +71,6 @@ inc_iohist(IOHistogram hist, uint64 latency_us)
 	hist->wait_us_count++;
 }
 
-static inline void
-inc_qthist(QTHistogram hist, uint64 elapsed_us)
-{
-	int			lo = 0;
-	int			hi = NUM_QT_BUCKETS - 1;
-
-	/* Find the right bucket with binary search */
-	while (lo < hi)
-	{
-		int			mid = (lo + hi) / 2;
-
-		if (elapsed_us < qt_bucket_thresholds[mid])
-			hi = mid;
-		else
-			lo = mid + 1;
-	}
-	hist->elapsed_us_bucket[lo]++;
-	hist->elapsed_us_sum += elapsed_us;
-	hist->elapsed_us_count++;
-}
-
 /*
  * Count a GetPage wait operation.
  */
@@ -119,13 +98,6 @@ inc_page_cache_write_wait(uint64 latency)
 	inc_iohist(&MyNeonCounters->file_cache_write_hist, latency);
 }
 
-
-void
-inc_query_time(uint64 elapsed)
-{
-	inc_qthist(&MyNeonCounters->query_time_hist, elapsed);
-}
-
 /*
  * Support functions for the views, neon_backend_perf_counters and
  * neon_perf_counters.
@@ -140,11 +112,11 @@ typedef struct
 } metric_t;
 
 static int
-io_histogram_to_metrics(IOHistogram histogram,
-						metric_t *metrics,
-						const char *count,
-						const char *sum,
-						const char *bucket)
+histogram_to_metrics(IOHistogram histogram,
+					 metric_t *metrics,
+					 const char *count,
+					 const char *sum,
+					 const char *bucket)
 {
 	int		i = 0;
 	uint64	bucket_accum = 0;
@@ -173,44 +145,10 @@ io_histogram_to_metrics(IOHistogram histogram,
 	return i;
 }
 
-static int
-qt_histogram_to_metrics(QTHistogram histogram,
-						metric_t *metrics,
-						const char *count,
-						const char *sum,
-						const char *bucket)
-{
-	int		i = 0;
-	uint64	bucket_accum = 0;
-
-	metrics[i].name = count;
-	metrics[i].is_bucket = false;
-	metrics[i].value = (double) histogram->elapsed_us_count;
-	i++;
-	metrics[i].name = sum;
-	metrics[i].is_bucket = false;
-	metrics[i].value = (double) histogram->elapsed_us_sum / 1000000.0;
-	i++;
-	for (int bucketno = 0; bucketno < NUM_QT_BUCKETS; bucketno++)
-	{
-		uint64		threshold = qt_bucket_thresholds[bucketno];
-
-		bucket_accum += histogram->elapsed_us_bucket[bucketno];
-
-		metrics[i].name = bucket;
-		metrics[i].is_bucket = true;
-		metrics[i].bucket_le = (threshold == UINT64_MAX) ? INFINITY : ((double) threshold) / 1000000.0;
-		metrics[i].value = (double) bucket_accum;
-		i++;
-	}
-
-	return i;
-}
-
 static metric_t *
 neon_perf_counters_to_metrics(neon_per_backend_counters *counters)
 {
-#define NUM_METRICS ((2 + NUM_IO_WAIT_BUCKETS) * 3 + (2 + NUM_QT_BUCKETS) + 12)
+#define NUM_METRICS ((2 + NUM_IO_WAIT_BUCKETS) * 3 + 12)
 	metric_t   *metrics = palloc((NUM_METRICS + 1) * sizeof(metric_t));
 	int			i = 0;
 
@@ -221,10 +159,10 @@ neon_perf_counters_to_metrics(neon_per_backend_counters *counters)
 		i++; \
 	} while (false)
 
-	i += io_histogram_to_metrics(&counters->getpage_hist, &metrics[i],
-								 "getpage_wait_seconds_count",
-								 "getpage_wait_seconds_sum",
-								 "getpage_wait_seconds_bucket");
+	i += histogram_to_metrics(&counters->getpage_hist, &metrics[i],
+							  "getpage_wait_seconds_count",
+							  "getpage_wait_seconds_sum",
+							  "getpage_wait_seconds_bucket");
 
 	APPEND_METRIC(getpage_prefetch_requests_total);
 	APPEND_METRIC(getpage_sync_requests_total);
@@ -240,19 +178,14 @@ neon_perf_counters_to_metrics(neon_per_backend_counters *counters)
 
 	APPEND_METRIC(file_cache_hits_total);
 
-	i += io_histogram_to_metrics(&counters->file_cache_read_hist, &metrics[i],
-								 "file_cache_read_wait_seconds_count",
-								 "file_cache_read_wait_seconds_sum",
-								 "file_cache_read_wait_seconds_bucket");
-	i += io_histogram_to_metrics(&counters->file_cache_write_hist, &metrics[i],
-								 "file_cache_write_wait_seconds_count",
-								 "file_cache_write_wait_seconds_sum",
-								 "file_cache_write_wait_seconds_bucket");
-
-	i += qt_histogram_to_metrics(&counters->query_time_hist, &metrics[i],
-								 "query_time_seconds_count",
-								 "query_time_seconds_sum",
-								 "query_time_seconds_bucket");
+	i += histogram_to_metrics(&counters->file_cache_read_hist, &metrics[i],
+							  "file_cache_read_wait_seconds_count",
+							  "file_cache_read_wait_seconds_sum",
+							  "file_cache_read_wait_seconds_bucket");
+	i += histogram_to_metrics(&counters->file_cache_write_hist, &metrics[i],
+							  "file_cache_write_wait_seconds_count",
+							  "file_cache_write_wait_seconds_sum",
+							  "file_cache_write_wait_seconds_bucket");
 
 	Assert(i == NUM_METRICS);
 
@@ -324,7 +257,7 @@ neon_get_backend_perf_counters(PG_FUNCTION_ARGS)
 }
 
 static inline void
-io_histogram_merge_into(IOHistogram into, IOHistogram from)
+histogram_merge_into(IOHistogram into, IOHistogram from)
 {
 	into->wait_us_count += from->wait_us_count;
 	into->wait_us_sum += from->wait_us_sum;
@@ -332,15 +265,6 @@ io_histogram_merge_into(IOHistogram into, IOHistogram from)
 		into->wait_us_bucket[bucketno] += from->wait_us_bucket[bucketno];
 }
 
-static inline void
-qt_histogram_merge_into(QTHistogram into, QTHistogram from)
-{
-	into->elapsed_us_count += from->elapsed_us_count;
-	into->elapsed_us_sum += from->elapsed_us_sum;
-	for (int bucketno = 0; bucketno < NUM_QT_BUCKETS; bucketno++)
-		into->elapsed_us_bucket[bucketno] += from->elapsed_us_bucket[bucketno];
-}
-
 PG_FUNCTION_INFO_V1(neon_get_perf_counters);
 Datum
 neon_get_perf_counters(PG_FUNCTION_ARGS)
@@ -359,7 +283,7 @@ neon_get_perf_counters(PG_FUNCTION_ARGS)
 	{
 		neon_per_backend_counters *counters = &neon_per_backend_counters_shared[procno];
 
-		io_histogram_merge_into(&totals.getpage_hist, &counters->getpage_hist);
+		histogram_merge_into(&totals.getpage_hist, &counters->getpage_hist);
 		totals.getpage_prefetch_requests_total += counters->getpage_prefetch_requests_total;
 		totals.getpage_sync_requests_total += counters->getpage_sync_requests_total;
 		totals.getpage_prefetch_misses_total += counters->getpage_prefetch_misses_total;
@@ -370,13 +294,13 @@ neon_get_perf_counters(PG_FUNCTION_ARGS)
 		totals.pageserver_open_requests += counters->pageserver_open_requests;
 		totals.getpage_prefetches_buffered += counters->getpage_prefetches_buffered;
 		totals.file_cache_hits_total += counters->file_cache_hits_total;
+		histogram_merge_into(&totals.file_cache_read_hist, &counters->file_cache_read_hist);
+		histogram_merge_into(&totals.file_cache_write_hist, &counters->file_cache_write_hist);
+
 		totals.compute_getpage_stuck_requests_total += counters->compute_getpage_stuck_requests_total;
 		totals.compute_getpage_max_inflight_stuck_time_ms = Max(
 			totals.compute_getpage_max_inflight_stuck_time_ms,
 			counters->compute_getpage_max_inflight_stuck_time_ms);
-		io_histogram_merge_into(&totals.file_cache_read_hist, &counters->file_cache_read_hist);
-		io_histogram_merge_into(&totals.file_cache_write_hist, &counters->file_cache_write_hist);
-		qt_histogram_merge_into(&totals.query_time_hist, &counters->query_time_hist);
 	}
 
 	metrics = neon_perf_counters_to_metrics(&totals);

pgxn/neon/neon_perf_counters.h

@@ -36,28 +36,6 @@ typedef struct IOHistogramData
 
 typedef IOHistogramData *IOHistogram;
 
-static const uint64 qt_bucket_thresholds[] = {
-	       2,        3,        6,        10,  /* 0 us   - 10 us */
-	      20,       30,       60,       100,  /* 10 us  - 100 us */
-	     200,      300,      600,	   1000,  /* 100 us - 1 ms */
-	    2000,     3000,     6000,     10000,  /* 1 ms   - 10 ms */
-	   20000,    30000,    60000,    100000,  /* 10 ms  - 100 ms */
-	  200000,   300000,   600000,   1000000,  /* 100 ms - 1 s */
-	 2000000,  3000000,  6000000,  10000000,  /* 1 s - 10 s */
-	20000000, 30000000, 60000000, 100000000,  /* 10 s - 100 s */
-	UINT64_MAX,
-};
-#define NUM_QT_BUCKETS (lengthof(qt_bucket_thresholds))
-
-typedef struct QTHistogramData
-{
-	uint64		elapsed_us_count;
-	uint64		elapsed_us_sum;
-	uint64		elapsed_us_bucket[NUM_QT_BUCKETS];
-} QTHistogramData;
-
-typedef QTHistogramData *QTHistogram;
-
 typedef struct
 {
 	/*
@@ -149,11 +127,6 @@ typedef struct
 	/* LFC I/O time buckets */
 	IOHistogramData file_cache_read_hist;
 	IOHistogramData file_cache_write_hist;
-
-	/*
-	 * Histogram of query execution time.
-	 */
-	QTHistogramData query_time_hist;
 } neon_per_backend_counters;
 
 /* Pointer to the shared memory array of neon_per_backend_counters structs */
@@ -176,7 +149,6 @@ extern neon_per_backend_counters *neon_per_backend_counters_shared;
 extern void inc_getpage_wait(uint64 latency);
 extern void inc_page_cache_read_wait(uint64 latency);
 extern void inc_page_cache_write_wait(uint64 latency);
-extern void inc_query_time(uint64 elapsed);
 
 extern Size NeonPerfCountersShmemSize(void);
 extern void NeonPerfCountersShmemInit(void);

pgxn/neon/neon_pgversioncompat.c

@@ -5,7 +5,6 @@
 
 #include "funcapi.h"
 #include "miscadmin.h"
-#include "access/xlog.h"
 #include "utils/tuplestore.h"
 
 #include "neon_pgversioncompat.h"
@@ -42,12 +41,5 @@ InitMaterializedSRF(FunctionCallInfo fcinfo, bits32 flags)
 	rsinfo->setDesc = stored_tupdesc;
 	MemoryContextSwitchTo(old_context);
 }
-
-TimeLineID GetWALInsertionTimeLine(void)
-{
-	return ThisTimeLineID + 1;
-}
-
-
 #endif
 

pgxn/neon/neon_pgversioncompat.h

@@ -162,7 +162,6 @@ InitBufferTag(BufferTag *tag, const RelFileNode *rnode,
 
 #if PG_MAJORVERSION_NUM < 15
 extern void InitMaterializedSRF(FunctionCallInfo fcinfo, bits32 flags);
-extern TimeLineID GetWALInsertionTimeLine(void);
 #endif
 
 #endif							/* NEON_PGVERSIONCOMPAT_H */

pgxn/neon/neon_walreader.c

@@ -69,7 +69,6 @@ struct NeonWALReader
 	WALSegmentContext segcxt;
 	WALOpenSegment seg;
 	int			wre_errno;
-	TimeLineID	local_active_tlid;
 	/* Explains failure to read, static for simplicity. */
 	char		err_msg[NEON_WALREADER_ERR_MSG_LEN];
 
@@ -107,7 +106,7 @@ struct NeonWALReader
 
 /* palloc and initialize NeonWALReader */
 NeonWALReader *
-NeonWALReaderAllocate(int wal_segment_size, XLogRecPtr available_lsn, char *log_prefix, TimeLineID tlid)
+NeonWALReaderAllocate(int wal_segment_size, XLogRecPtr available_lsn, char *log_prefix)
 {
 	NeonWALReader *reader;
 
@@ -119,7 +118,6 @@ NeonWALReaderAllocate(int wal_segment_size, XLogRecPtr available_lsn, char *log_
 		MemoryContextAllocZero(TopMemoryContext, sizeof(NeonWALReader));
 
 	reader->available_lsn = available_lsn;
-	reader->local_active_tlid = tlid;
 	reader->seg.ws_file = -1;
 	reader->seg.ws_segno = 0;
 	reader->seg.ws_tli = 0;
@@ -579,17 +577,6 @@ NeonWALReaderIsRemConnEstablished(NeonWALReader *state)
 	return state->rem_state == RS_ESTABLISHED;
 }
 
-/*
- * Whether remote connection is established. Once this is done, until successful
- * local read or error socket is stable and user can update socket events
- * instead of readding it each time.
- */
-TimeLineID
-NeonWALReaderLocalActiveTimeLineID(NeonWALReader *state)
-{
-	return state->local_active_tlid;
-}
-
 /*
  * Returns events user should wait on connection socket or 0 if remote
  * connection is not active.

pgxn/neon/neon_walreader.h

@@ -19,10 +19,9 @@ typedef enum
 	NEON_WALREAD_ERROR,
 } NeonWALReadResult;
 
-extern NeonWALReader *NeonWALReaderAllocate(int wal_segment_size, XLogRecPtr available_lsn, char *log_prefix, TimeLineID tlid);
+extern NeonWALReader *NeonWALReaderAllocate(int wal_segment_size, XLogRecPtr available_lsn, char *log_prefix);
 extern void NeonWALReaderFree(NeonWALReader *state);
 extern void NeonWALReaderResetRemote(NeonWALReader *state);
-extern TimeLineID NeonWALReaderLocalActiveTimeLineID(NeonWALReader *state);
 extern NeonWALReadResult NeonWALRead(NeonWALReader *state, char *buf, XLogRecPtr startptr, Size count, TimeLineID tli);
 extern pgsocket NeonWALReaderSocket(NeonWALReader *state);
 extern uint32 NeonWALReaderEvents(NeonWALReader *state);

pgxn/neon/walproposer.c

@@ -98,7 +98,6 @@ WalProposerCreate(WalProposerConfig *config, walproposer_api api)
 	wp = palloc0(sizeof(WalProposer));
 	wp->config = config;
 	wp->api = api;
-	wp->localTimeLineID = config->pgTimeline;
 	wp->state = WPS_COLLECTING_TERMS;
 	wp->mconf.generation = INVALID_GENERATION;
 	wp->mconf.members.len = 0;
@@ -120,10 +119,6 @@ WalProposerCreate(WalProposerConfig *config, walproposer_api api)
 		{
 			wp_log(FATAL, "failed to parse neon.safekeepers generation number: %m");
 		}
-		if (*endptr != ':')
-		{
-			wp_log(FATAL, "failed to parse neon.safekeepers: no colon after generation");
-		}
 		/* Skip past : to the first hostname. */
 		host = endptr + 1;
 	}
@@ -1385,7 +1380,7 @@ ProcessPropStartPos(WalProposer *wp)
 	 * we must bail out, as clog and other non rel data is inconsistent.
 	 */
 	walprop_shared = wp->api.get_shmem_state(wp);
-	if (!wp->config->syncSafekeepers && !walprop_shared->replica_promote)
+	if (!wp->config->syncSafekeepers)
 	{
 		/*
 		 * Basebackup LSN always points to the beginning of the record (not

pgxn/neon/walproposer.h

@@ -391,7 +391,6 @@ typedef struct WalproposerShmemState
 	/* last feedback from each shard */
 	PageserverFeedback shard_ps_feedback[MAX_SHARDS];
 	int			num_shards;
-	bool		replica_promote;
 
 	/* aggregated feedback with min LSNs across shards */
 	PageserverFeedback min_ps_feedback;
@@ -807,9 +806,6 @@ typedef struct WalProposer
 	/* Safekeepers walproposer is connecting to. */
 	Safekeeper	safekeeper[MAX_SAFEKEEPERS];
 
-	/* Current local TimeLineId in use */
-	TimeLineID	localTimeLineID;
-
 	/* WAL has been generated up to this point */
 	XLogRecPtr	availableLsn;
 

pgxn/neon/walproposer_pg.c

@@ -35,7 +35,6 @@
 #include "storage/proc.h"
 #include "storage/ipc.h"
 #include "storage/lwlock.h"
-#include "storage/pg_shmem.h"
 #include "storage/shmem.h"
 #include "storage/spin.h"
 #include "tcop/tcopprot.h"
@@ -160,19 +159,12 @@ WalProposerMain(Datum main_arg)
 {
 	WalProposer *wp;
 
-	if (*wal_acceptors_list == '\0')
-	{
-		wpg_log(WARNING, "Safekeepers list is empty");
-		return;
-	}
-
 	init_walprop_config(false);
 	walprop_pg_init_bgworker();
 	am_walproposer = true;
 	walprop_pg_load_libpqwalreceiver();
 
 	wp = WalProposerCreate(&walprop_config, walprop_pg);
-	wp->localTimeLineID = GetWALInsertionTimeLine();
 	wp->last_reconnect_attempt = walprop_pg_get_current_timestamp(wp);
 
 	walprop_pg_init_walsender();
@@ -280,30 +272,6 @@ split_safekeepers_list(char *safekeepers_list, char *safekeepers[])
 	return n_safekeepers;
 }
 
-static char *split_off_safekeepers_generation(char *safekeepers_list, uint32 *generation)
-{
-	char	   *endptr;
-
-	if (strncmp(safekeepers_list, "g#", 2) != 0)
-	{
-		return safekeepers_list;
-	}
-	else
-	{
-		errno = 0;
-		*generation = strtoul(safekeepers_list + 2, &endptr, 10);
-		if (errno != 0)
-		{
-			wp_log(FATAL, "failed to parse neon.safekeepers generation number: %m");
-		}
-		if (*endptr != ':')
-		{
-			wp_log(FATAL, "failed to parse neon.safekeepers: no colon after generation");
-		}
-		return endptr + 1;
-	}
-}
-
 /*
  * Accept two coma-separated strings with list of safekeeper host:port addresses.
  * Split them into arrays and return false if two sets do not match, ignoring the order.
@@ -315,16 +283,6 @@ safekeepers_cmp(char *old, char *new)
 	char	   *safekeepers_new[MAX_SAFEKEEPERS];
 	int			len_old = 0;
 	int			len_new = 0;
-	uint32		gen_old = INVALID_GENERATION;
-	uint32		gen_new = INVALID_GENERATION;
-
-	old = split_off_safekeepers_generation(old, &gen_old);
-	new = split_off_safekeepers_generation(new, &gen_new);
-
-	if (gen_old != gen_new)
-	{
-		return false;
-	}
 
 	len_old = split_safekeepers_list(old, safekeepers_old);
 	len_new = split_safekeepers_list(new, safekeepers_new);
@@ -358,9 +316,6 @@ assign_neon_safekeepers(const char *newval, void *extra)
 	char	   *newval_copy;
 	char	   *oldval;
 
-	if (newval && *newval != '\0' && UsedShmemSegAddr && walprop_shared && RecoveryInProgress())
-		walprop_shared->replica_promote = true;
-
 	if (!am_walproposer)
 		return;
 
@@ -551,15 +506,16 @@ BackpressureThrottlingTime(void)
 
 /*
  * Register a background worker proposing WAL to wal acceptors.
- * We start walproposer bgworker even for replicas in order to support possible replica promotion.
- * When pg_promote() function is called, then walproposer bgworker registered with BgWorkerStart_RecoveryFinished
- * is automatically launched when promotion is completed.
  */
 static void
 walprop_register_bgworker(void)
 {
 	BackgroundWorker bgw;
 
+	/* If no wal acceptors are specified, don't start the background worker. */
+	if (*wal_acceptors_list == '\0')
+		return;
+
 	memset(&bgw, 0, sizeof(bgw));
 	bgw.bgw_flags = BGWORKER_SHMEM_ACCESS;
 	bgw.bgw_start_time = BgWorkerStart_RecoveryFinished;
@@ -1336,7 +1292,9 @@ StartProposerReplication(WalProposer *wp, StartReplicationCmd *cmd)
 
 #if PG_VERSION_NUM < 150000
 	if (ThisTimeLineID == 0)
-		ThisTimeLineID = 1;
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("IDENTIFY_SYSTEM has not been run before START_REPLICATION")));
 #endif
 
 	/*
@@ -1550,7 +1508,7 @@ walprop_pg_wal_reader_allocate(Safekeeper *sk)
 
 	snprintf(log_prefix, sizeof(log_prefix), WP_LOG_PREFIX "sk %s:%s nwr: ", sk->host, sk->port);
 	Assert(!sk->xlogreader);
-	sk->xlogreader = NeonWALReaderAllocate(wal_segment_size, sk->wp->propTermStartLsn, log_prefix, sk->wp->localTimeLineID);
+	sk->xlogreader = NeonWALReaderAllocate(wal_segment_size, sk->wp->propTermStartLsn, log_prefix);
 	if (sk->xlogreader == NULL)
 		wpg_log(FATAL, "failed to allocate xlog reader");
 }
@@ -1564,7 +1522,7 @@ walprop_pg_wal_read(Safekeeper *sk, char *buf, XLogRecPtr startptr, Size count,
 					  buf,
 					  startptr,
 					  count,
-					  sk->wp->localTimeLineID);
+					  walprop_pg_get_timeline_id());
 
 	if (res == NEON_WALREAD_SUCCESS)
 	{

pgxn/neon/walsender_hooks.c

@@ -111,7 +111,7 @@ NeonWALPageRead(
 											readBuf,
 											targetPagePtr,
 											count,
-											NeonWALReaderLocalActiveTimeLineID(wal_reader));
+											walprop_pg_get_timeline_id());
 
 		if (res == NEON_WALREAD_SUCCESS)
 		{
@@ -202,7 +202,7 @@ NeonOnDemandXLogReaderRoutines(XLogReaderRoutine *xlr)
 		{
 			elog(ERROR, "unable to start walsender when basebackupLsn is 0");
 		}
-		wal_reader = NeonWALReaderAllocate(wal_segment_size, basebackupLsn, "[walsender] ", 1);
+		wal_reader = NeonWALReaderAllocate(wal_segment_size, basebackupLsn, "[walsender] ");
 	}
 	xlr->page_read = NeonWALPageRead;
 	xlr->segment_open = NeonWALReadSegmentOpen;

proxy/src/auth/backend/classic.rs

@@ -18,6 +18,11 @@ pub(super) async fn authenticate(
     secret: AuthSecret,
 ) -> auth::Result<ComputeCredentials> {
     let scram_keys = match secret {
+        #[cfg(any(test, feature = "testing"))]
+        AuthSecret::Md5(_) => {
+            debug!("auth endpoint chooses MD5");
+            return Err(auth::AuthError::MalformedPassword("MD5 not supported"));
+        }
         AuthSecret::Scram(secret) => {
             debug!("auth endpoint chooses SCRAM");
 

proxy/src/auth/backend/console_redirect.rs

@@ -6,17 +6,18 @@ use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::{info, info_span};
 
+use super::ComputeCredentialKeys;
+use crate::auth::IpPattern;
 use crate::auth::backend::ComputeUserInfo;
 use crate::cache::Cached;
-use crate::compute::AuthInfo;
 use crate::config::AuthenticationConfig;
 use crate::context::RequestContext;
 use crate::control_plane::client::cplane_proxy_v1;
 use crate::control_plane::{self, CachedNodeInfo, NodeInfo};
 use crate::error::{ReportableError, UserFacingError};
-use crate::pglb::connect_compute::ComputeConnectBackend;
 use crate::pqproto::BeMessage;
 use crate::proxy::NeonOptions;
+use crate::proxy::connect_compute::ComputeConnectBackend;
 use crate::stream::PqStream;
 use crate::types::RoleName;
 use crate::{auth, compute, waiters};
@@ -97,11 +98,15 @@ impl ConsoleRedirectBackend {
         ctx: &RequestContext,
         auth_config: &'static AuthenticationConfig,
         client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-    ) -> auth::Result<(ConsoleRedirectNodeInfo, AuthInfo, ComputeUserInfo)> {
+    ) -> auth::Result<(
+        ConsoleRedirectNodeInfo,
+        ComputeUserInfo,
+        Option<Vec<IpPattern>>,
+    )> {
         authenticate(ctx, auth_config, &self.console_uri, client)
             .await
-            .map(|(node_info, auth_info, user_info)| {
-                (ConsoleRedirectNodeInfo(node_info), auth_info, user_info)
+            .map(|(node_info, user_info, ip_allowlist)| {
+                (ConsoleRedirectNodeInfo(node_info), user_info, ip_allowlist)
             })
     }
 }
@@ -116,6 +121,10 @@ impl ComputeConnectBackend for ConsoleRedirectNodeInfo {
     ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError> {
         Ok(Cached::new_uncached(self.0.clone()))
     }
+
+    fn get_keys(&self) -> &ComputeCredentialKeys {
+        &ComputeCredentialKeys::None
+    }
 }
 
 async fn authenticate(
@@ -123,7 +132,7 @@ async fn authenticate(
     auth_config: &'static AuthenticationConfig,
     link_uri: &reqwest::Url,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-) -> auth::Result<(NodeInfo, AuthInfo, ComputeUserInfo)> {
+) -> auth::Result<(NodeInfo, ComputeUserInfo, Option<Vec<IpPattern>>)> {
     ctx.set_auth_method(crate::context::AuthMethod::ConsoleRedirect);
 
     // registering waiter can fail if we get unlucky with rng.
@@ -183,24 +192,10 @@ async fn authenticate(
 
     client.write_message(BeMessage::NoticeResponse("Connecting to database."));
 
-    // Backwards compatibility. pg_sni_proxy uses "--" in domain names
-    // while direct connections do not. Once we migrate to pg_sni_proxy
-    // everywhere, we can remove this.
-    let ssl_mode = if db_info.host.contains("--") {
-        // we need TLS connection with SNI info to properly route it
-        SslMode::Require
-    } else {
-        SslMode::Disable
-    };
-
-    let conn_info = compute::ConnectInfo {
-        host: db_info.host.into(),
-        port: db_info.port,
-        ssl_mode,
-        host_addr: None,
-    };
-    let auth_info =
-        AuthInfo::for_console_redirect(&db_info.dbname, &db_info.user, db_info.password.as_deref());
+    // This config should be self-contained, because we won't
+    // take username or dbname from client's startup message.
+    let mut config = compute::ConnCfg::new(db_info.host.to_string(), db_info.port);
+    config.dbname(&db_info.dbname).user(&db_info.user);
 
     let user: RoleName = db_info.user.into();
     let user_info = ComputeUserInfo {
@@ -214,12 +209,26 @@ async fn authenticate(
     ctx.set_project(db_info.aux.clone());
     info!("woken up a compute node");
 
+    // Backwards compatibility. pg_sni_proxy uses "--" in domain names
+    // while direct connections do not. Once we migrate to pg_sni_proxy
+    // everywhere, we can remove this.
+    if db_info.host.contains("--") {
+        // we need TLS connection with SNI info to properly route it
+        config.ssl_mode(SslMode::Require);
+    } else {
+        config.ssl_mode(SslMode::Disable);
+    }
+
+    if let Some(password) = db_info.password {
+        config.password(password.as_ref());
+    }
+
     Ok((
         NodeInfo {
-            conn_info,
+            config,
             aux: db_info.aux,
         },
-        auth_info,
         user_info,
+        db_info.allowed_ips,
     ))
 }

proxy/src/auth/backend/local.rs

@@ -1,12 +1,11 @@
 use std::net::SocketAddr;
 
 use arc_swap::ArcSwapOption;
-use postgres_client::config::SslMode;
 use tokio::sync::Semaphore;
 
 use super::jwt::{AuthRule, FetchAuthRules};
 use crate::auth::backend::jwt::FetchAuthRulesError;
-use crate::compute::ConnectInfo;
+use crate::compute::ConnCfg;
 use crate::compute_ctl::ComputeCtlApi;
 use crate::context::RequestContext;
 use crate::control_plane::NodeInfo;
@@ -30,12 +29,7 @@ impl LocalBackend {
                 api: http::Endpoint::new(compute_ctl, http::new_client()),
             },
             node_info: NodeInfo {
-                conn_info: ConnectInfo {
-                    host_addr: Some(postgres_addr.ip()),
-                    host: postgres_addr.ip().to_string().into(),
-                    port: postgres_addr.port(),
-                    ssl_mode: SslMode::Disable,
-                },
+                config: ConnCfg::new(postgres_addr.ip().to_string(), postgres_addr.port()),
                 // TODO(conrad): make this better reflect compute info rather than endpoint info.
                 aux: MetricsAuxInfo {
                     endpoint_id: EndpointIdTag::get_interner().get_or_intern("local"),

proxy/src/auth/backend/mod.rs

@@ -25,9 +25,9 @@ use crate::control_plane::{
     RoleAccessControl,
 };
 use crate::intern::EndpointIdInt;
-use crate::pglb::connect_compute::ComputeConnectBackend;
 use crate::pqproto::BeMessage;
 use crate::proxy::NeonOptions;
+use crate::proxy::connect_compute::ComputeConnectBackend;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::stream::Stream;
 use crate::types::{EndpointCacheKey, EndpointId, RoleName};
@@ -168,6 +168,8 @@ impl ComputeUserInfo {
 
 #[cfg_attr(test, derive(Debug))]
 pub(crate) enum ComputeCredentialKeys {
+    #[cfg(any(test, feature = "testing"))]
+    Password(Vec<u8>),
     AuthKeys(AuthKeys),
     JwtPayload(Vec<u8>),
     None,
@@ -417,6 +419,13 @@ impl ComputeConnectBackend for Backend<'_, ComputeCredentials> {
             Self::Local(local) => Ok(Cached::new_uncached(local.node_info.clone())),
         }
     }
+
+    fn get_keys(&self) -> &ComputeCredentialKeys {
+        match self {
+            Self::ControlPlane(_, creds) => &creds.keys,
+            Self::Local(_) => &ComputeCredentialKeys::None,
+        }
+    }
 }
 
 #[cfg(test)]

proxy/src/auth/flow.rs

@@ -169,6 +169,13 @@ pub(crate) async fn validate_password_and_exchange(
     secret: AuthSecret,
 ) -> super::Result<sasl::Outcome<ComputeCredentialKeys>> {
     match secret {
+        #[cfg(any(test, feature = "testing"))]
+        AuthSecret::Md5(_) => {
+            // test only
+            Ok(sasl::Outcome::Success(ComputeCredentialKeys::Password(
+                password.to_owned(),
+            )))
+        }
         // perform scram authentication as both client and server to validate the keys
         AuthSecret::Scram(scram_secret) => {
             let outcome = crate::scram::exchange(pool, endpoint, &scram_secret, password).await?;

proxy/src/binary/proxy.rs

@@ -221,7 +221,8 @@ struct ProxyCliArgs {
     is_private_access_proxy: bool,
 
     /// Configure whether all incoming requests have a Proxy Protocol V2 packet.
-    #[clap(value_enum, long, default_value_t = ProxyProtocolV2::Rejected)]
+    // TODO(conradludgate): switch default to rejected or required once we've updated all deployments
+    #[clap(value_enum, long, default_value_t = ProxyProtocolV2::Supported)]
     proxy_protocol_v2: ProxyProtocolV2,
 
     /// Time the proxy waits for the webauth session to be confirmed by the control plane.

proxy/src/cancellation.rs

@@ -24,6 +24,7 @@ use crate::pqproto::CancelKeyData;
 use crate::rate_limiter::LeakyBucketRateLimiter;
 use crate::redis::keys::KeyPrefix;
 use crate::redis::kv_ops::RedisKVClient;
+use crate::tls::postgres_rustls::MakeRustlsConnect;
 
 type IpSubnetKey = IpNet;
 
@@ -496,8 +497,10 @@ impl CancelClosure {
     ) -> Result<(), CancelError> {
         let socket = TcpStream::connect(self.socket_addr).await?;
 
-        let tls = <_ as MakeTlsConnect<tokio::net::TcpStream>>::make_tls_connect(
-            compute_config,
+        let mut mk_tls =
+            crate::tls::postgres_rustls::MakeRustlsConnect::new(compute_config.tls.clone());
+        let tls = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::make_tls_connect(
+            &mut mk_tls,
             &self.hostname,
         )
         .map_err(|e| CancelError::IO(std::io::Error::other(e.to_string())))?;

proxy/src/compute.rs

@@ -1,24 +1,21 @@
-mod tls;
-
 use std::fmt::Debug;
 use std::io;
-use std::net::{IpAddr, SocketAddr};
+use std::net::SocketAddr;
+use std::time::Duration;
 
 use futures::{FutureExt, TryFutureExt};
 use itertools::Itertools;
-use postgres_client::config::{AuthKeys, SslMode};
-use postgres_client::maybe_tls_stream::MaybeTlsStream;
 use postgres_client::tls::MakeTlsConnect;
-use postgres_client::{CancelToken, NoTls, RawConnection};
+use postgres_client::{CancelToken, RawConnection};
 use postgres_protocol::message::backend::NoticeResponseBody;
+use rustls::pki_types::InvalidDnsNameError;
 use thiserror::Error;
 use tokio::net::{TcpStream, lookup_host};
 use tracing::{debug, error, info, warn};
 
-use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
+use crate::auth::backend::ComputeUserInfo;
 use crate::auth::parse_endpoint_param;
 use crate::cancellation::CancelClosure;
-use crate::compute::tls::TlsError;
 use crate::config::ComputeConfig;
 use crate::context::RequestContext;
 use crate::control_plane::client::ApiLockError;
@@ -28,6 +25,7 @@ use crate::error::{ReportableError, UserFacingError};
 use crate::metrics::{Metrics, NumDbConnectionsGuard};
 use crate::pqproto::StartupMessageParams;
 use crate::proxy::neon_option;
+use crate::tls::postgres_rustls::MakeRustlsConnect;
 use crate::types::Host;
 
 pub const COULD_NOT_CONNECT: &str = "Couldn't connect to compute node";
@@ -40,7 +38,10 @@ pub(crate) enum ConnectionError {
     Postgres(#[from] postgres_client::Error),
 
     #[error("{COULD_NOT_CONNECT}: {0}")]
-    TlsError(#[from] TlsError),
+    CouldNotConnect(#[from] io::Error),
+
+    #[error("{COULD_NOT_CONNECT}: {0}")]
+    TlsError(#[from] InvalidDnsNameError),
 
     #[error("{COULD_NOT_CONNECT}: {0}")]
     WakeComputeError(#[from] WakeComputeError),
@@ -72,7 +73,7 @@ impl UserFacingError for ConnectionError {
             ConnectionError::TooManyConnectionAttempts(_) => {
                 "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(),
+            _ => COULD_NOT_CONNECT.to_owned(),
         }
     }
 }
@@ -84,6 +85,7 @@ impl ReportableError for ConnectionError {
                 crate::error::ErrorKind::Postgres
             }
             ConnectionError::Postgres(_) => crate::error::ErrorKind::Compute,
+            ConnectionError::CouldNotConnect(_) => crate::error::ErrorKind::Compute,
             ConnectionError::TlsError(_) => crate::error::ErrorKind::Compute,
             ConnectionError::WakeComputeError(e) => e.get_error_kind(),
             ConnectionError::TooManyConnectionAttempts(e) => e.get_error_kind(),
@@ -94,85 +96,34 @@ impl ReportableError for ConnectionError {
 /// A pair of `ClientKey` & `ServerKey` for `SCRAM-SHA-256`.
 pub(crate) type ScramKeys = postgres_client::config::ScramKeys<32>;
 
+/// A config for establishing a connection to compute node.
+/// Eventually, `postgres_client` will be replaced with something better.
+/// Newtype allows us to implement methods on top of it.
 #[derive(Clone)]
-pub enum Auth {
-    /// Only used during console-redirect.
-    Password(Vec<u8>),
-    /// Used by sql-over-http, ws, tcp.
-    Scram(Box<ScramKeys>),
-}
-
-/// A config for authenticating to the compute node.
-pub(crate) struct AuthInfo {
-    /// None for local-proxy, as we use trust-based localhost auth.
-    /// Some for sql-over-http, ws, tcp, and in most cases for console-redirect.
-    /// Might be None for console-redirect, but that's only a consequence of testing environments ATM.
-    auth: Option<Auth>,
-    server_params: StartupMessageParams,
-
-    /// Console redirect sets user and database, we shouldn't re-use those from the params.
-    skip_db_user: bool,
-}
-
-/// Contains only the data needed to establish a secure connection to compute.
-#[derive(Clone)]
-pub struct ConnectInfo {
-    pub host_addr: Option<IpAddr>,
-    pub host: Host,
-    pub port: u16,
-    pub ssl_mode: SslMode,
-}
+pub(crate) struct ConnCfg(Box<postgres_client::Config>);
 
 /// Creation and initialization routines.
-impl AuthInfo {
-    pub(crate) fn for_console_redirect(db: &str, user: &str, pw: Option<&str>) -> Self {
-        let mut server_params = StartupMessageParams::default();
-        server_params.insert("database", db);
-        server_params.insert("user", user);
-        Self {
-            auth: pw.map(|pw| Auth::Password(pw.as_bytes().to_owned())),
-            server_params,
-            skip_db_user: true,
+impl ConnCfg {
+    pub(crate) fn new(host: String, port: u16) -> Self {
+        Self(Box::new(postgres_client::Config::new(host, port)))
+    }
+
+    /// Reuse password or auth keys from the other config.
+    pub(crate) fn reuse_password(&mut self, other: Self) {
+        if let Some(password) = other.get_password() {
+            self.password(password);
+        }
+
+        if let Some(keys) = other.get_auth_keys() {
+            self.auth_keys(keys);
         }
     }
 
-    pub(crate) fn with_auth_keys(keys: &ComputeCredentialKeys) -> Self {
-        Self {
-            auth: match keys {
-                ComputeCredentialKeys::AuthKeys(AuthKeys::ScramSha256(auth_keys)) => {
-                    Some(Auth::Scram(Box::new(*auth_keys)))
-                }
-                ComputeCredentialKeys::JwtPayload(_) | ComputeCredentialKeys::None => None,
-            },
-            server_params: StartupMessageParams::default(),
-            skip_db_user: false,
+    pub(crate) fn get_host(&self) -> Host {
+        match self.0.get_host() {
+            postgres_client::config::Host::Tcp(s) => s.into(),
         }
     }
-}
-
-impl ConnectInfo {
-    pub fn to_postgres_client_config(&self) -> postgres_client::Config {
-        let mut config = postgres_client::Config::new(self.host.to_string(), self.port);
-        config.ssl_mode(self.ssl_mode);
-        if let Some(host_addr) = self.host_addr {
-            config.set_host_addr(host_addr);
-        }
-        config
-    }
-}
-
-impl AuthInfo {
-    fn enrich(&self, mut config: postgres_client::Config) -> postgres_client::Config {
-        match &self.auth {
-            Some(Auth::Scram(keys)) => config.auth_keys(AuthKeys::ScramSha256(**keys)),
-            Some(Auth::Password(pw)) => config.password(pw),
-            None => &mut config,
-        };
-        for (k, v) in self.server_params.iter() {
-            config.set_param(k, v);
-        }
-        config
-    }
 
     /// Apply startup message params to the connection config.
     pub(crate) fn set_startup_params(
@@ -181,26 +132,27 @@ impl AuthInfo {
         arbitrary_params: bool,
     ) {
         if !arbitrary_params {
-            self.server_params.insert("client_encoding", "UTF8");
+            self.set_param("client_encoding", "UTF8");
         }
         for (k, v) in params.iter() {
             match k {
                 // Only set `user` if it's not present in the config.
                 // Console redirect auth flow takes username from the console's response.
-                "user" | "database" if self.skip_db_user => {}
+                "user" if self.user_is_set() => {}
+                "database" if self.db_is_set() => {}
                 "options" => {
                     if let Some(options) = filtered_options(v) {
-                        self.server_params.insert(k, &options);
+                        self.set_param(k, &options);
                     }
                 }
                 "user" | "database" | "application_name" | "replication" => {
-                    self.server_params.insert(k, v);
+                    self.set_param(k, v);
                 }
 
                 // if we allow arbitrary params, then we forward them through.
                 // this is a flag for a period of backwards compatibility
                 k if arbitrary_params => {
-                    self.server_params.insert(k, v);
+                    self.set_param(k, v);
                 }
                 _ => {}
             }
@@ -208,13 +160,25 @@ impl AuthInfo {
     }
 }
 
-impl ConnectInfo {
-    /// Establish a raw TCP+TLS connection to the compute node.
-    async fn connect_raw(
-        &self,
-        config: &ComputeConfig,
-    ) -> Result<(SocketAddr, MaybeTlsStream<TcpStream, RustlsStream>), TlsError> {
-        let timeout = config.timeout;
+impl std::ops::Deref for ConnCfg {
+    type Target = postgres_client::Config;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+/// For now, let's make it easier to setup the config.
+impl std::ops::DerefMut for ConnCfg {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl ConnCfg {
+    /// Establish a raw TCP connection to the compute node.
+    async fn connect_raw(&self, timeout: Duration) -> io::Result<(SocketAddr, TcpStream, &str)> {
+        use postgres_client::config::Host;
 
         // wrap TcpStream::connect with timeout
         let connect_with_timeout = |addrs| {
@@ -244,32 +208,34 @@ impl ConnectInfo {
         // We can't reuse connection establishing logic from `postgres_client` here,
         // because it has no means for extracting the underlying socket which we
         // require for our business.
-        let port = self.port;
-        let host = &*self.host;
+        let port = self.0.get_port();
+        let host = self.0.get_host();
 
-        let addrs = match self.host_addr {
+        let host = match host {
+            Host::Tcp(host) => host.as_str(),
+        };
+
+        let addrs = match self.0.get_host_addr() {
             Some(addr) => vec![SocketAddr::new(addr, port)],
             None => lookup_host((host, port)).await?.collect(),
         };
 
         match connect_once(&*addrs).await {
-            Ok((sockaddr, stream)) => Ok((
-                sockaddr,
-                tls::connect_tls(stream, self.ssl_mode, config, host).await?,
-            )),
+            Ok((sockaddr, stream)) => Ok((sockaddr, stream, host)),
             Err(err) => {
                 warn!("couldn't connect to compute node at {host}:{port}: {err}");
-                Err(TlsError::Connection(err))
+                Err(err)
             }
         }
     }
 }
 
-type RustlsStream = <ComputeConfig as MakeTlsConnect<tokio::net::TcpStream>>::Stream;
+type RustlsStream = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::Stream;
 
 pub(crate) struct PostgresConnection {
     /// Socket connected to a compute node.
-    pub(crate) stream: MaybeTlsStream<tokio::net::TcpStream, RustlsStream>,
+    pub(crate) stream:
+        postgres_client::maybe_tls_stream::MaybeTlsStream<tokio::net::TcpStream, RustlsStream>,
     /// PostgreSQL connection parameters.
     pub(crate) params: std::collections::HashMap<String, String>,
     /// Query cancellation token.
@@ -282,23 +248,28 @@ pub(crate) struct PostgresConnection {
     _guage: NumDbConnectionsGuard<'static>,
 }
 
-impl ConnectInfo {
+impl ConnCfg {
     /// Connect to a corresponding compute node.
     pub(crate) async fn connect(
         &self,
         ctx: &RequestContext,
         aux: MetricsAuxInfo,
-        auth: &AuthInfo,
         config: &ComputeConfig,
         user_info: ComputeUserInfo,
     ) -> Result<PostgresConnection, ConnectionError> {
-        let mut tmp_config = auth.enrich(self.to_postgres_client_config());
-        // we setup SSL early in `ConnectInfo::connect_raw`.
-        tmp_config.ssl_mode(SslMode::Disable);
-
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let (socket_addr, stream) = self.connect_raw(config).await?;
-        let connection = tmp_config.connect_raw(stream, NoTls).await?;
+        let (socket_addr, stream, host) = self.connect_raw(config.timeout).await?;
+        drop(pause);
+
+        let mut mk_tls = crate::tls::postgres_rustls::MakeRustlsConnect::new(config.tls.clone());
+        let tls = <MakeRustlsConnect as MakeTlsConnect<tokio::net::TcpStream>>::make_tls_connect(
+            &mut mk_tls,
+            host,
+        )?;
+
+        // connect_raw() will not use TLS if sslmode is "disable"
+        let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
+        let connection = self.0.connect_raw(stream, tls).await?;
         drop(pause);
 
         let RawConnection {
@@ -311,14 +282,13 @@ impl ConnectInfo {
 
         tracing::Span::current().record("pid", tracing::field::display(process_id));
         tracing::Span::current().record("compute_id", tracing::field::display(&aux.compute_id));
-        let MaybeTlsStream::Raw(stream) = stream.into_inner();
+        let stream = stream.into_inner();
 
         // 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,
+            "connected to compute node at {host} ({socket_addr}) sslmode={:?}, latency={}, query_id={}",
+            self.0.get_ssl_mode(),
             ctx.get_proxy_latency(),
             ctx.get_testodrome_id().unwrap_or_default(),
         );
@@ -329,11 +299,11 @@ impl ConnectInfo {
             socket_addr,
             CancelToken {
                 socket_config: None,
-                ssl_mode: self.ssl_mode,
+                ssl_mode: self.0.get_ssl_mode(),
                 process_id,
                 secret_key,
             },
-            self.host.to_string(),
+            host.to_string(),
             user_info,
         );
 

proxy/src/compute/tls.rs

@@ -1,63 +0,0 @@
-use futures::FutureExt;
-use postgres_client::config::SslMode;
-use postgres_client::maybe_tls_stream::MaybeTlsStream;
-use postgres_client::tls::{MakeTlsConnect, TlsConnect};
-use rustls::pki_types::InvalidDnsNameError;
-use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite};
-
-use crate::pqproto::request_tls;
-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")]
-    Required,
-}
-
-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,
-        }
-    }
-}
-
-pub async fn connect_tls<S, T>(
-    mut stream: S,
-    mode: SslMode,
-    tls: &T,
-    host: &str,
-) -> Result<MaybeTlsStream<S, T::Stream>, TlsError>
-where
-    S: AsyncRead + AsyncWrite + Unpin + Send,
-    T: MakeTlsConnect<
-            S,
-            Error = InvalidDnsNameError,
-            TlsConnect: TlsConnect<S, Error = std::io::Error, Future: Send>,
-        >,
-{
-    match mode {
-        SslMode::Disable => return Ok(MaybeTlsStream::Raw(stream)),
-        SslMode::Prefer | SslMode::Require => {}
-    }
-
-    if !request_tls(&mut stream).await? {
-        if SslMode::Require == mode {
-            return Err(TlsError::Required);
-        }
-
-        return Ok(MaybeTlsStream::Raw(stream));
-    }
-
-    Ok(MaybeTlsStream::Tls(
-        tls.make_tls_connect(host)?.connect(stream).boxed().await?,
-    ))
-}

proxy/src/config.rs

@@ -39,6 +39,8 @@ pub struct ComputeConfig {
 pub enum ProxyProtocolV2 {
     /// Connection will error if PROXY protocol v2 header is missing
     Required,
+    /// Connection will parse PROXY protocol v2 header, but accept the connection if it's missing.
+    Supported,
     /// Connection will error if PROXY protocol v2 header is provided
     Rejected,
 }

proxy/src/console_redirect_proxy.rs

@@ -11,10 +11,10 @@ use crate::config::{ProxyConfig, ProxyProtocolV2};
 use crate::context::RequestContext;
 use crate::error::ReportableError;
 use crate::metrics::{Metrics, NumClientConnectionsGuard};
-use crate::pglb::connect_compute::{TcpMechanism, connect_to_compute};
-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::handshake::{HandshakeData, handshake};
+use crate::proxy::passthrough::ProxyPassthrough;
 use crate::proxy::{
     ClientRequestError, ErrorSource, prepare_client_connection, run_until_cancelled,
 };
@@ -54,24 +54,30 @@ pub async fn task_main(
         debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
 
         connections.spawn(async move {
-            let (socket, conn_info) = match config.proxy_protocol_v2 {
-                ProxyProtocolV2::Required => {
-                    match read_proxy_protocol(socket).await {
-                        Err(e) => {
-                            error!("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),
-                    }
+            let (socket, peer_addr) = match read_proxy_protocol(socket).await {
+                Err(e) => {
+                    error!("per-client task finished with an error: {e:#}");
+                    return;
                 }
-                // ignore the header - it cannot be confused for a postgres or http connection so will
-                // error later.
-                ProxyProtocolV2::Rejected => (
+                // 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::Missing))
+                    if config.proxy_protocol_v2 == ProxyProtocolV2::Required =>
+                {
+                    error!("missing required proxy protocol header");
+                    return;
+                }
+                Ok((_socket, ConnectHeader::Proxy(_)))
+                    if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected =>
+                {
+                    error!("proxy protocol header not supported");
+                    return;
+                }
+                Ok((socket, ConnectHeader::Proxy(info))) => (socket, info),
+                Ok((socket, ConnectHeader::Missing)) => (
                     socket,
                     ConnectionInfo {
                         addr: peer_addr,
@@ -80,7 +86,7 @@ pub async fn task_main(
                 ),
             };
 
-            match socket.set_nodelay(true) {
+            match socket.inner.set_nodelay(true) {
                 Ok(()) => {}
                 Err(e) => {
                     error!(
@@ -92,7 +98,7 @@ pub async fn task_main(
 
             let ctx = RequestContext::new(
                 session_id,
-                conn_info,
+                peer_addr,
                 crate::metrics::Protocol::Tcp,
                 &config.region,
             );
@@ -210,20 +216,20 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
 
     ctx.set_db_options(params.clone());
 
-    let (node_info, mut auth_info, user_info) = match backend
+    let (node_info, user_info, _ip_allowlist) = match backend
         .authenticate(ctx, &config.authentication_config, &mut stream)
         .await
     {
         Ok(auth_result) => auth_result,
         Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
     };
-    auth_info.set_startup_params(&params, true);
 
     let node = connect_to_compute(
         ctx,
         &TcpMechanism {
             user_info,
-            auth: auth_info,
+            params_compat: true,
+            params: &params,
             locks: &config.connect_compute_locks,
         },
         &node_info,

proxy/src/control_plane/client/cplane_proxy_v1.rs

@@ -261,18 +261,24 @@ impl NeonControlPlaneClient {
                 Some(_) => SslMode::Require,
                 None => SslMode::Disable,
             };
-            let host = match body.server_name {
-                Some(host) => host.into(),
-                None => host.into(),
+            let host_name = match body.server_name {
+                Some(host) => host,
+                None => host.to_owned(),
             };
 
+            // Don't set anything but host and port! This config will be cached.
+            // We'll set username and such later using the startup message.
+            // TODO: add more type safety (in progress).
+            let mut config = compute::ConnCfg::new(host_name, port);
+
+            if let Some(addr) = host_addr {
+                config.set_host_addr(addr);
+            }
+
+            config.ssl_mode(ssl_mode);
+
             let node = NodeInfo {
-                conn_info: compute::ConnectInfo {
-                    host_addr,
-                    host,
-                    port,
-                    ssl_mode,
-                },
+                config,
                 aux: body.aux,
             };
 

proxy/src/control_plane/client/mock.rs

@@ -6,7 +6,6 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use futures::TryFutureExt;
-use postgres_client::config::SslMode;
 use thiserror::Error;
 use tokio_postgres::Client;
 use tracing::{Instrument, error, info, info_span, warn};
@@ -15,7 +14,6 @@ use crate::auth::IpPattern;
 use crate::auth::backend::ComputeUserInfo;
 use crate::auth::backend::jwt::AuthRule;
 use crate::cache::Cached;
-use crate::compute::ConnectInfo;
 use crate::context::RequestContext;
 use crate::control_plane::errors::{
     ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError,
@@ -26,9 +24,9 @@ use crate::control_plane::{
     RoleAccessControl,
 };
 use crate::intern::RoleNameInt;
-use crate::scram;
 use crate::types::{BranchId, EndpointId, ProjectId, RoleName};
 use crate::url::ApiUrl;
+use crate::{compute, scram};
 
 #[derive(Debug, Error)]
 enum MockApiError {
@@ -89,7 +87,8 @@ impl MockControlPlane {
             .await?
             {
                 info!("got a secret: {entry}"); // safe since it's not a prod scenario
-                scram::ServerSecret::parse(&entry).map(AuthSecret::Scram)
+                let secret = scram::ServerSecret::parse(&entry).map(AuthSecret::Scram);
+                secret.or_else(|| parse_md5(&entry).map(AuthSecret::Md5))
             } else {
                 warn!("user '{role}' does not exist");
                 None
@@ -171,23 +170,25 @@ impl MockControlPlane {
 
     async fn do_wake_compute(&self) -> Result<NodeInfo, WakeComputeError> {
         let port = self.endpoint.port().unwrap_or(5432);
-        let conn_info = match self.endpoint.host_str() {
-            None => ConnectInfo {
-                host_addr: Some(IpAddr::V4(Ipv4Addr::LOCALHOST)),
-                host: "localhost".into(),
-                port,
-                ssl_mode: SslMode::Disable,
-            },
-            Some(host) => ConnectInfo {
-                host_addr: IpAddr::from_str(host).ok(),
-                host: host.into(),
-                port,
-                ssl_mode: SslMode::Disable,
-            },
+        let mut config = match self.endpoint.host_str() {
+            None => {
+                let mut config = compute::ConnCfg::new("localhost".to_string(), port);
+                config.set_host_addr(IpAddr::V4(Ipv4Addr::LOCALHOST));
+                config
+            }
+            Some(host) => {
+                let mut config = compute::ConnCfg::new(host.to_string(), port);
+                if let Ok(addr) = IpAddr::from_str(host) {
+                    config.set_host_addr(addr);
+                }
+                config
+            }
         };
 
+        config.ssl_mode(postgres_client::config::SslMode::Disable);
+
         let node = NodeInfo {
-            conn_info,
+            config,
             aux: MetricsAuxInfo {
                 endpoint_id: (&EndpointId::from("endpoint")).into(),
                 project_id: (&ProjectId::from("project")).into(),
@@ -265,3 +266,12 @@ impl super::ControlPlaneApi for MockControlPlane {
         self.do_wake_compute().map_ok(Cached::new_uncached).await
     }
 }
+
+fn parse_md5(input: &str) -> Option<[u8; 16]> {
+    let text = input.strip_prefix("md5")?;
+
+    let mut bytes = [0u8; 16];
+    hex::decode_to_slice(text, &mut bytes).ok()?;
+
+    Some(bytes)
+}

proxy/src/control_plane/mod.rs

@@ -11,8 +11,8 @@ pub(crate) mod errors;
 
 use std::sync::Arc;
 
-use crate::auth::backend::ComputeUserInfo;
 use crate::auth::backend::jwt::AuthRule;
+use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
 use crate::auth::{AuthError, IpPattern, check_peer_addr_is_in_list};
 use crate::cache::{Cached, TimedLru};
 use crate::config::ComputeConfig;
@@ -39,6 +39,10 @@ pub mod mgmt;
 /// Auth secret which is managed by the cloud.
 #[derive(Clone, Eq, PartialEq, Debug)]
 pub(crate) enum AuthSecret {
+    #[cfg(any(test, feature = "testing"))]
+    /// Md5 hash of user's password.
+    Md5([u8; 16]),
+
     /// [SCRAM](crate::scram) authentication info.
     Scram(scram::ServerSecret),
 }
@@ -59,9 +63,13 @@ pub(crate) struct AuthInfo {
 }
 
 /// Info for establishing a connection to a compute node.
+/// This is what we get after auth succeeded, but not before!
 #[derive(Clone)]
 pub(crate) struct NodeInfo {
-    pub(crate) conn_info: compute::ConnectInfo,
+    /// Compute node connection params.
+    /// It's sad that we have to clone this, but this will improve
+    /// once we migrate to a bespoke connection logic.
+    pub(crate) config: compute::ConnCfg,
 
     /// Labels for proxy's metrics.
     pub(crate) aux: MetricsAuxInfo,
@@ -71,14 +79,26 @@ impl NodeInfo {
     pub(crate) async fn connect(
         &self,
         ctx: &RequestContext,
-        auth: &compute::AuthInfo,
         config: &ComputeConfig,
         user_info: ComputeUserInfo,
     ) -> Result<compute::PostgresConnection, compute::ConnectionError> {
-        self.conn_info
-            .connect(ctx, self.aux.clone(), auth, config, user_info)
+        self.config
+            .connect(ctx, self.aux.clone(), config, user_info)
             .await
     }
+
+    pub(crate) fn reuse_settings(&mut self, other: Self) {
+        self.config.reuse_password(other.config);
+    }
+
+    pub(crate) fn set_keys(&mut self, keys: &ComputeCredentialKeys) {
+        match keys {
+            #[cfg(any(test, feature = "testing"))]
+            ComputeCredentialKeys::Password(password) => self.config.password(password),
+            ComputeCredentialKeys::AuthKeys(auth_keys) => self.config.auth_keys(*auth_keys),
+            ComputeCredentialKeys::JwtPayload(_) | ComputeCredentialKeys::None => &mut self.config,
+        };
+    }
 }
 
 #[derive(Copy, Clone, Default)]

proxy/src/pglb/mod.rs

@@ -1,5 +1 @@
-pub mod connect_compute;
-pub mod copy_bidirectional;
-pub mod handshake;
 pub mod inprocess;
-pub mod passthrough;

proxy/src/pqproto.rs

@@ -8,7 +8,7 @@ use std::io::{self, Cursor};
 use bytes::{Buf, BufMut};
 use itertools::Itertools;
 use rand::distributions::{Distribution, Standard};
-use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+use tokio::io::{AsyncRead, AsyncReadExt};
 use zerocopy::{FromBytes, Immutable, IntoBytes, big_endian};
 
 pub type ErrorCode = [u8; 5];
@@ -53,28 +53,6 @@ impl fmt::Debug for ProtocolVersion {
     }
 }
 
-/// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L118>
-const MAX_STARTUP_PACKET_LENGTH: usize = 10000;
-const RESERVED_INVALID_MAJOR_VERSION: u16 = 1234;
-/// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L132>
-const CANCEL_REQUEST_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5678);
-/// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L166>
-const NEGOTIATE_SSL_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5679);
-/// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L167>
-const NEGOTIATE_GSS_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5680);
-
-/// This first reads the startup message header, is 8 bytes.
-/// The first 4 bytes is a big-endian message length, and the next 4 bytes is a version number.
-///
-/// The length value is inclusive of the header. For example,
-/// an empty message will always have length 8.
-#[derive(Clone, Copy, FromBytes, IntoBytes, Immutable)]
-#[repr(C)]
-struct StartupHeader {
-    len: big_endian::U32,
-    version: ProtocolVersion,
-}
-
 /// read the type from the stream using zerocopy.
 ///
 /// not cancel safe.
@@ -88,38 +66,32 @@ macro_rules! read {
     }};
 }
 
-/// Returns true if TLS is supported.
-///
-/// This is not cancel safe.
-pub async fn request_tls<S>(stream: &mut S) -> io::Result<bool>
-where
-    S: AsyncRead + AsyncWrite + Unpin,
-{
-    let payload = StartupHeader {
-        len: 8.into(),
-        version: NEGOTIATE_SSL_CODE,
-    };
-    stream.write_all(payload.as_bytes()).await?;
-    stream.flush().await?;
-
-    // we expect back either `S` or `N` as a single byte.
-    let mut res = *b"0";
-    stream.read_exact(&mut res).await?;
-
-    debug_assert!(
-        res == *b"S" || res == *b"N",
-        "unexpected SSL negotiation response: {}",
-        char::from(res[0]),
-    );
-
-    // S for SSL.
-    Ok(res == *b"S")
-}
-
 pub async fn read_startup<S>(stream: &mut S) -> io::Result<FeStartupPacket>
 where
     S: AsyncRead + Unpin,
 {
+    /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L118>
+    const MAX_STARTUP_PACKET_LENGTH: usize = 10000;
+    const RESERVED_INVALID_MAJOR_VERSION: u16 = 1234;
+    /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L132>
+    const CANCEL_REQUEST_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5678);
+    /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L166>
+    const NEGOTIATE_SSL_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5679);
+    /// <https://github.com/postgres/postgres/blob/ca481d3c9ab7bf69ff0c8d71ad3951d407f6a33c/src/include/libpq/pqcomm.h#L167>
+    const NEGOTIATE_GSS_CODE: ProtocolVersion = ProtocolVersion::new(1234, 5680);
+
+    /// This first reads the startup message header, is 8 bytes.
+    /// The first 4 bytes is a big-endian message length, and the next 4 bytes is a version number.
+    ///
+    /// The length value is inclusive of the header. For example,
+    /// an empty message will always have length 8.
+    #[derive(Clone, Copy, FromBytes, IntoBytes, Immutable)]
+    #[repr(C)]
+    struct StartupHeader {
+        len: big_endian::U32,
+        version: ProtocolVersion,
+    }
+
     let header = read!(stream => StartupHeader);
 
     // <https://github.com/postgres/postgres/blob/04bcf9e19a4261fe9c7df37c777592c2e10c32a7/src/backend/tcop/backend_startup.c#L378-L382>
@@ -592,9 +564,10 @@ mod tests {
     use tokio::io::{AsyncWriteExt, duplex};
     use zerocopy::IntoBytes;
 
-    use super::ProtocolVersion;
     use crate::pqproto::{FeStartupPacket, read_message, read_startup};
 
+    use super::ProtocolVersion;
+
     #[tokio::test]
     async fn reject_large_startup() {
         // we're going to define a v3.0 startup message with far too many parameters.

proxy/src/protocol2.rs

@@ -4,13 +4,60 @@
 use core::fmt;
 use std::io;
 use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
+use std::pin::Pin;
+use std::task::{Context, Poll};
 
-use bytes::Buf;
+use bytes::{Buf, Bytes, BytesMut};
+use pin_project_lite::pin_project;
 use smol_str::SmolStr;
 use strum_macros::FromRepr;
-use tokio::io::{AsyncRead, AsyncReadExt};
+use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, ReadBuf};
 use zerocopy::{FromBytes, Immutable, KnownLayout, Unaligned, network_endian};
 
+pin_project! {
+    /// A chained [`AsyncRead`] with [`AsyncWrite`] passthrough
+    pub(crate) struct ChainRW<T> {
+        #[pin]
+        pub(crate) inner: T,
+        buf: BytesMut,
+    }
+}
+
+impl<T: AsyncWrite> AsyncWrite for ChainRW<T> {
+    #[inline]
+    fn poll_write(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, io::Error>> {
+        self.project().inner.poll_write(cx, buf)
+    }
+
+    #[inline]
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        self.project().inner.poll_flush(cx)
+    }
+
+    #[inline]
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        self.project().inner.poll_shutdown(cx)
+    }
+
+    #[inline]
+    fn poll_write_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        bufs: &[io::IoSlice<'_>],
+    ) -> Poll<Result<usize, io::Error>> {
+        self.project().inner.poll_write_vectored(cx, bufs)
+    }
+
+    #[inline]
+    fn is_write_vectored(&self) -> bool {
+        self.inner.is_write_vectored()
+    }
+}
+
 /// Proxy Protocol Version 2 Header
 const SIGNATURE: [u8; 12] = [
     0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D, 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A,
@@ -32,6 +79,7 @@ pub struct ConnectionInfo {
 
 #[derive(PartialEq, Eq, Clone, Debug)]
 pub enum ConnectHeader {
+    Missing,
     Local,
     Proxy(ConnectionInfo),
 }
@@ -58,24 +106,47 @@ pub enum ConnectionInfoExtra {
 
 pub(crate) async fn read_proxy_protocol<T: AsyncRead + Unpin>(
     mut read: T,
-) -> std::io::Result<(T, ConnectHeader)> {
-    let mut header = [0; size_of::<ProxyProtocolV2Header>()];
-    read.read_exact(&mut header).await?;
-    let header: ProxyProtocolV2Header = zerocopy::transmute!(header);
-    if header.signature != SIGNATURE {
-        return Err(std::io::Error::other("invalid proxy protocol header"));
+) -> std::io::Result<(ChainRW<T>, ConnectHeader)> {
+    let mut buf = BytesMut::with_capacity(128);
+    let header = loop {
+        let bytes_read = read.read_buf(&mut buf).await?;
+
+        // exit for bad header signature
+        let len = usize::min(buf.len(), SIGNATURE.len());
+        if buf[..len] != SIGNATURE[..len] {
+            return Ok((ChainRW { inner: read, buf }, ConnectHeader::Missing));
+        }
+
+        // if no more bytes available then exit
+        if bytes_read == 0 {
+            return Ok((ChainRW { inner: read, buf }, ConnectHeader::Missing));
+        }
+
+        // check if we have enough bytes to continue
+        if let Some(header) = buf.try_get::<ProxyProtocolV2Header>() {
+            break header;
+        }
+    };
+
+    let remaining_length = usize::from(header.len.get());
+
+    while buf.len() < remaining_length {
+        if read.read_buf(&mut buf).await? == 0 {
+            return Err(io::Error::new(
+                io::ErrorKind::UnexpectedEof,
+                "stream closed while waiting for proxy protocol addresses",
+            ));
+        }
     }
+    let payload = buf.split_to(remaining_length);
 
-    let mut payload = vec![0; usize::from(header.len.get())];
-    read.read_exact(&mut payload).await?;
-
-    let res = process_proxy_payload(header, &payload)?;
-    Ok((read, res))
+    let res = process_proxy_payload(header, payload)?;
+    Ok((ChainRW { inner: read, buf }, res))
 }
 
 fn process_proxy_payload(
     header: ProxyProtocolV2Header,
-    mut payload: &[u8],
+    mut payload: BytesMut,
 ) -> std::io::Result<ConnectHeader> {
     match header.version_and_command {
         // the connection was established on purpose by the proxy
@@ -91,12 +162,13 @@ fn process_proxy_payload(
         PROXY_V2 => {}
         // other values are unassigned and must not be emitted by senders. Receivers
         // must drop connections presenting unexpected values here.
-        _ => {
-            return Err(io::Error::other(format!(
+        #[rustfmt::skip] // https://github.com/rust-lang/rustfmt/issues/6384
+        _ => return Err(io::Error::other(
+            format!(
                 "invalid proxy protocol command 0x{:02X}. expected local (0x20) or proxy (0x21)",
                 header.version_and_command
-            )));
-        }
+            ),
+        )),
     }
 
     let size_err =
@@ -134,7 +206,7 @@ fn process_proxy_payload(
                 }
                 let subtype = tlv.value.get_u8();
                 match Pp2AwsType::from_repr(subtype) {
-                    Some(Pp2AwsType::VpceId) => match std::str::from_utf8(tlv.value) {
+                    Some(Pp2AwsType::VpceId) => match std::str::from_utf8(&tlv.value) {
                         Ok(s) => {
                             extra = Some(ConnectionInfoExtra::Aws { vpce_id: s.into() });
                         }
@@ -210,28 +282,65 @@ enum Pp2AzureType {
     PrivateEndpointLinkId = 0x01,
 }
 
-#[derive(Debug)]
-struct Tlv<'a> {
-    kind: u8,
-    value: &'a [u8],
+impl<T: AsyncRead> AsyncRead for ChainRW<T> {
+    #[inline]
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        if self.buf.is_empty() {
+            self.project().inner.poll_read(cx, buf)
+        } else {
+            self.read_from_buf(buf)
+        }
+    }
 }
 
-fn read_tlv<'a>(b: &mut &'a [u8]) -> Option<Tlv<'a>> {
+impl<T: AsyncRead> ChainRW<T> {
+    #[cold]
+    fn read_from_buf(self: Pin<&mut Self>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
+        debug_assert!(!self.buf.is_empty());
+        let this = self.project();
+
+        let write = usize::min(this.buf.len(), buf.remaining());
+        let slice = this.buf.split_to(write).freeze();
+        buf.put_slice(&slice);
+
+        // reset the allocation so it can be freed
+        if this.buf.is_empty() {
+            *this.buf = BytesMut::new();
+        }
+
+        Poll::Ready(Ok(()))
+    }
+}
+
+#[derive(Debug)]
+struct Tlv {
+    kind: u8,
+    value: Bytes,
+}
+
+fn read_tlv(b: &mut BytesMut) -> Option<Tlv> {
     let tlv_header = b.try_get::<TlvHeader>()?;
     let len = usize::from(tlv_header.len.get());
+    if b.len() < len {
+        return None;
+    }
     Some(Tlv {
         kind: tlv_header.kind,
-        value: b.split_off(..len)?,
+        value: b.split_to(len).freeze(),
     })
 }
 
 trait BufExt: Sized {
     fn try_get<T: FromBytes>(&mut self) -> Option<T>;
 }
-impl BufExt for &[u8] {
+impl BufExt for BytesMut {
     fn try_get<T: FromBytes>(&mut self) -> Option<T> {
-        let (res, rest) = T::read_from_prefix(self).ok()?;
-        *self = rest;
+        let (res, _) = T::read_from_prefix(self).ok()?;
+        self.advance(size_of::<T>());
         Some(res)
     }
 }
@@ -372,19 +481,27 @@ mod tests {
     }
 
     #[tokio::test]
-    #[should_panic = "invalid proxy protocol header"]
     async fn test_invalid() {
         let data = [0x55; 256];
 
-        read_proxy_protocol(data.as_slice()).await.unwrap();
+        let (mut read, info) = read_proxy_protocol(data.as_slice()).await.unwrap();
+
+        let mut bytes = vec![];
+        read.read_to_end(&mut bytes).await.unwrap();
+        assert_eq!(bytes, data);
+        assert_eq!(info, ConnectHeader::Missing);
     }
 
     #[tokio::test]
-    #[should_panic = "early eof"]
     async fn test_short() {
         let data = [0x55; 10];
 
-        read_proxy_protocol(data.as_slice()).await.unwrap();
+        let (mut read, info) = read_proxy_protocol(data.as_slice()).await.unwrap();
+
+        let mut bytes = vec![];
+        read.read_to_end(&mut bytes).await.unwrap();
+        assert_eq!(bytes, data);
+        assert_eq!(info, ConnectHeader::Missing);
     }
 
     #[tokio::test]

proxy/src/proxy/connect_compute.rs

@@ -2,8 +2,9 @@ use async_trait::async_trait;
 use tokio::time;
 use tracing::{debug, info, warn};
 
-use crate::auth::backend::ComputeUserInfo;
-use crate::compute::{self, AuthInfo, COULD_NOT_CONNECT, PostgresConnection};
+use super::retry::ShouldRetryWakeCompute;
+use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
+use crate::compute::{self, COULD_NOT_CONNECT, PostgresConnection};
 use crate::config::{ComputeConfig, RetryConfig};
 use crate::context::RequestContext;
 use crate::control_plane::errors::WakeComputeError;
@@ -13,7 +14,8 @@ use crate::error::ReportableError;
 use crate::metrics::{
     ConnectOutcome, ConnectionFailureKind, Metrics, RetriesMetricGroup, RetryType,
 };
-use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute, retry_after, should_retry};
+use crate::pqproto::StartupMessageParams;
+use crate::proxy::retry::{CouldRetry, retry_after, should_retry};
 use crate::proxy::wake_compute::wake_compute;
 use crate::types::Host;
 
@@ -47,6 +49,8 @@ pub(crate) trait ConnectMechanism {
         node_info: &control_plane::CachedNodeInfo,
         config: &ComputeConfig,
     ) -> Result<Self::Connection, Self::ConnectError>;
+
+    fn update_connect_config(&self, conf: &mut compute::ConnCfg);
 }
 
 #[async_trait]
@@ -55,17 +59,24 @@ pub(crate) trait ComputeConnectBackend {
         &self,
         ctx: &RequestContext,
     ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError>;
+
+    fn get_keys(&self) -> &ComputeCredentialKeys;
 }
 
-pub(crate) struct TcpMechanism {
-    pub(crate) auth: AuthInfo,
+pub(crate) struct TcpMechanism<'a> {
+    pub(crate) params_compat: bool,
+
+    /// KV-dictionary with PostgreSQL connection params.
+    pub(crate) params: &'a StartupMessageParams,
+
     /// connect_to_compute concurrency lock
     pub(crate) locks: &'static ApiLocks<Host>,
+
     pub(crate) user_info: ComputeUserInfo,
 }
 
 #[async_trait]
-impl ConnectMechanism for TcpMechanism {
+impl ConnectMechanism for TcpMechanism<'_> {
     type Connection = PostgresConnection;
     type ConnectError = compute::ConnectionError;
     type Error = compute::ConnectionError;
@@ -80,12 +91,13 @@ impl ConnectMechanism for TcpMechanism {
         node_info: &control_plane::CachedNodeInfo,
         config: &ComputeConfig,
     ) -> Result<PostgresConnection, Self::Error> {
-        let permit = self.locks.get_permit(&node_info.conn_info.host).await?;
-        permit.release_result(
-            node_info
-                .connect(ctx, &self.auth, config, self.user_info.clone())
-                .await,
-        )
+        let host = node_info.config.get_host();
+        let permit = self.locks.get_permit(&host).await?;
+        permit.release_result(node_info.connect(ctx, config, self.user_info.clone()).await)
+    }
+
+    fn update_connect_config(&self, config: &mut compute::ConnCfg) {
+        config.set_startup_params(self.params, self.params_compat);
     }
 }
 
@@ -103,9 +115,12 @@ where
     M::Error: From<WakeComputeError>,
 {
     let mut num_retries = 0;
-    let node_info =
+    let mut node_info =
         wake_compute(&mut num_retries, ctx, user_info, wake_compute_retry_config).await?;
 
+    node_info.set_keys(user_info.get_keys());
+    mechanism.update_connect_config(&mut node_info.config);
+
     // try once
     let err = match mechanism.connect_once(ctx, &node_info, compute).await {
         Ok(res) => {
@@ -141,9 +156,14 @@ where
     } else {
         // if we failed to connect, it's likely that the compute node was suspended, wake a new compute node
         debug!("compute node's state has likely changed; requesting a wake-up");
-        invalidate_cache(node_info);
+        let old_node_info = invalidate_cache(node_info);
         // TODO: increment num_retries?
-        wake_compute(&mut num_retries, ctx, user_info, wake_compute_retry_config).await?
+        let mut node_info =
+            wake_compute(&mut num_retries, ctx, user_info, wake_compute_retry_config).await?;
+        node_info.reuse_settings(old_node_info);
+
+        mechanism.update_connect_config(&mut node_info.config);
+        node_info
     };
 
     // now that we have a new node, try connect to it repeatedly.

proxy/src/proxy/mod.rs

@@ -1,10 +1,15 @@
 #[cfg(test)]
 mod tests;
 
+pub(crate) mod connect_compute;
+mod copy_bidirectional;
+pub(crate) mod handshake;
+pub(crate) mod passthrough;
 pub(crate) mod retry;
 pub(crate) mod wake_compute;
 use std::sync::Arc;
 
+pub use copy_bidirectional::{ErrorSource, copy_bidirectional_client_compute};
 use futures::FutureExt;
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
@@ -16,17 +21,16 @@ use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, debug, error, info, warn};
 
+use self::connect_compute::{TcpMechanism, connect_to_compute};
+use self::passthrough::ProxyPassthrough;
 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};
-use crate::pglb::connect_compute::{TcpMechanism, connect_to_compute};
-pub use crate::pglb::copy_bidirectional::{ErrorSource, copy_bidirectional_client_compute};
-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::handshake::{HandshakeData, handshake};
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::stream::{PqStream, Stream};
 use crate::types::EndpointCacheKey;
@@ -98,24 +102,30 @@ pub async fn task_main(
         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),
-                    }
+            let (socket, conn_info) = match read_proxy_protocol(socket).await {
+                Err(e) => {
+                    warn!("per-client task finished with an error: {e:#}");
+                    return;
                 }
-                // ignore the header - it cannot be confused for a postgres or http connection so will
-                // error later.
-                ProxyProtocolV2::Rejected => (
+                // 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::Missing))
+                    if config.proxy_protocol_v2 == ProxyProtocolV2::Required =>
+                {
+                    warn!("missing required proxy protocol header");
+                    return;
+                }
+                Ok((_socket, ConnectHeader::Proxy(_)))
+                    if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected =>
+                {
+                    warn!("proxy protocol header not supported");
+                    return;
+                }
+                Ok((socket, ConnectHeader::Proxy(info))) => (socket, info),
+                Ok((socket, ConnectHeader::Missing)) => (
                     socket,
                     ConnectionInfo {
                         addr: peer_addr,
@@ -124,7 +134,7 @@ pub async fn task_main(
                 ),
             };
 
-            match socket.set_nodelay(true) {
+            match socket.inner.set_nodelay(true) {
                 Ok(()) => {}
                 Err(e) => {
                     error!(
@@ -238,7 +248,7 @@ pub(crate) enum ClientRequestError {
     #[error("{0}")]
     Cancellation(#[from] cancellation::CancelError),
     #[error("{0}")]
-    Handshake(#[from] HandshakeError),
+    Handshake(#[from] handshake::HandshakeError),
     #[error("{0}")]
     HandshakeTimeout(#[from] tokio::time::error::Elapsed),
     #[error("{0}")]
@@ -358,19 +368,21 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         }
     };
 
-    let creds = match &user_info {
-        auth::Backend::ControlPlane(_, creds) => creds,
+    let compute_user_info = match &user_info {
+        auth::Backend::ControlPlane(_, info) => &info.info,
         auth::Backend::Local(_) => unreachable!("local proxy does not run tcp proxy service"),
     };
-    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);
+    let params_compat = compute_user_info
+        .options
+        .get(NeonOptions::PARAMS_COMPAT)
+        .is_some();
 
     let res = connect_to_compute(
         ctx,
         &TcpMechanism {
-            user_info: creds.info.clone(),
-            auth: auth_info,
+            user_info: compute_user_info.clone(),
+            params_compat,
+            params: &params,
             locks: &config.connect_compute_locks,
         },
         &user_info,

proxy/src/proxy/passthrough.rs

@@ -53,7 +53,7 @@ pub(crate) async fn proxy_pass(
 
     // Starting from here we only proxy the client's traffic.
     debug!("performing the proxy pass...");
-    let _ = crate::pglb::copy_bidirectional::copy_bidirectional_client_compute(
+    let _ = crate::proxy::copy_bidirectional::copy_bidirectional_client_compute(
         &mut client,
         &mut compute,
     )

proxy/src/proxy/retry.rs

@@ -100,9 +100,9 @@ impl CouldRetry for compute::ConnectionError {
     fn could_retry(&self) -> bool {
         match self {
             compute::ConnectionError::Postgres(err) => err.could_retry(),
-            compute::ConnectionError::TlsError(err) => err.could_retry(),
+            compute::ConnectionError::CouldNotConnect(err) => err.could_retry(),
             compute::ConnectionError::WakeComputeError(err) => err.could_retry(),
-            compute::ConnectionError::TooManyConnectionAttempts(_) => false,
+            _ => false,
         }
     }
 }

proxy/src/proxy/tests/mod.rs

@@ -8,7 +8,7 @@ use std::time::Duration;
 use anyhow::{Context, bail};
 use async_trait::async_trait;
 use http::StatusCode;
-use postgres_client::config::{AuthKeys, ScramKeys, SslMode};
+use postgres_client::config::SslMode;
 use postgres_client::tls::{MakeTlsConnect, NoTls};
 use retry::{ShouldRetryWakeCompute, retry_after};
 use rstest::rstest;
@@ -17,6 +17,7 @@ use rustls::pki_types;
 use tokio::io::DuplexStream;
 use tracing_test::traced_test;
 
+use super::connect_compute::ConnectMechanism;
 use super::retry::CouldRetry;
 use super::*;
 use crate::auth::backend::{
@@ -27,8 +28,8 @@ 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;
-use crate::pglb::connect_compute::ConnectMechanism;
 use crate::tls::client_config::compute_client_config_with_certs;
+use crate::tls::postgres_rustls::MakeRustlsConnect;
 use crate::tls::server_config::CertResolver;
 use crate::types::{BranchId, EndpointId, ProjectId};
 use crate::{sasl, scram};
@@ -71,14 +72,13 @@ struct ClientConfig<'a> {
     hostname: &'a str,
 }
 
-type TlsConnect<S> = <ComputeConfig as MakeTlsConnect<S>>::TlsConnect;
+type TlsConnect<S> = <MakeRustlsConnect as MakeTlsConnect<S>>::TlsConnect;
 
 impl ClientConfig<'_> {
     fn make_tls_connect(self) -> anyhow::Result<TlsConnect<DuplexStream>> {
-        Ok(crate::tls::postgres_rustls::make_tls_connect(
-            &self.config,
-            self.hostname,
-        )?)
+        let mut mk = MakeRustlsConnect::new(self.config);
+        let tls = MakeTlsConnect::<DuplexStream>::make_tls_connect(&mut mk, self.hostname)?;
+        Ok(tls)
     }
 }
 
@@ -173,6 +173,7 @@ async fn dummy_proxy(
     tls: Option<TlsConfig>,
     auth: impl TestAuth + Send,
 ) -> anyhow::Result<()> {
+    let (client, _) = read_proxy_protocol(client).await?;
     let mut stream = match handshake(&RequestContext::test(), client, tls.as_ref(), false).await? {
         HandshakeData::Startup(stream, _) => stream,
         HandshakeData::Cancel(_) => bail!("cancellation not supported"),
@@ -497,6 +498,8 @@ impl ConnectMechanism for TestConnectMechanism {
             x => panic!("expecting action {x:?}, connect is called instead"),
         }
     }
+
+    fn update_connect_config(&self, _conf: &mut compute::ConnCfg) {}
 }
 
 impl TestControlPlaneClient for TestConnectMechanism {
@@ -555,12 +558,7 @@ impl TestControlPlaneClient for TestConnectMechanism {
 
 fn helper_create_cached_node_info(cache: &'static NodeInfoCache) -> CachedNodeInfo {
     let node = NodeInfo {
-        conn_info: compute::ConnectInfo {
-            host: "test".into(),
-            port: 5432,
-            ssl_mode: SslMode::Disable,
-            host_addr: None,
-        },
+        config: compute::ConnCfg::new("test".to_owned(), 5432),
         aux: MetricsAuxInfo {
             endpoint_id: (&EndpointId::from("endpoint")).into(),
             project_id: (&ProjectId::from("project")).into(),
@@ -584,10 +582,7 @@ fn helper_create_connect_info(
                 user: "user".into(),
                 options: NeonOptions::parse_options_raw(""),
             },
-            keys: ComputeCredentialKeys::AuthKeys(AuthKeys::ScramSha256(ScramKeys {
-                client_key: [0; 32],
-                server_key: [0; 32],
-            })),
+            keys: ComputeCredentialKeys::Password("password".into()),
         },
     )
 }

proxy/src/proxy/wake_compute.rs

@@ -1,5 +1,6 @@
 use tracing::{error, info};
 
+use super::connect_compute::ComputeConnectBackend;
 use crate::config::RetryConfig;
 use crate::context::RequestContext;
 use crate::control_plane::CachedNodeInfo;
@@ -8,7 +9,6 @@ use crate::error::ReportableError;
 use crate::metrics::{
     ConnectOutcome, ConnectionFailuresBreakdownGroup, Metrics, RetriesMetricGroup, RetryType,
 };
-use crate::pglb::connect_compute::ComputeConnectBackend;
 use crate::proxy::retry::{retry_after, should_retry};
 
 // Use macro to retain original callsite.

proxy/src/serverless/backend.rs

@@ -23,6 +23,7 @@ use super::local_conn_pool::{self, EXT_NAME, EXT_SCHEMA, EXT_VERSION, LocalConnP
 use crate::auth::backend::local::StaticAuthRules;
 use crate::auth::backend::{ComputeCredentials, ComputeUserInfo};
 use crate::auth::{self, AuthError};
+use crate::compute;
 use crate::compute_ctl::{
     ComputeCtlError, ExtensionInstallRequest, Privilege, SetRoleGrantsRequest,
 };
@@ -34,7 +35,7 @@ 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::pglb::connect_compute::ConnectMechanism;
+use crate::proxy::connect_compute::ConnectMechanism;
 use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute};
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::types::{EndpointId, Host, LOCAL_PROXY_SUFFIX};
@@ -181,7 +182,7 @@ impl PoolingBackend {
         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);
-        crate::pglb::connect_compute::connect_to_compute(
+        crate::proxy::connect_compute::connect_to_compute(
             ctx,
             &TokioMechanism {
                 conn_id,
@@ -225,7 +226,7 @@ impl PoolingBackend {
             },
             keys: crate::auth::backend::ComputeCredentialKeys::None,
         });
-        crate::pglb::connect_compute::connect_to_compute(
+        crate::proxy::connect_compute::connect_to_compute(
             ctx,
             &HyperMechanism {
                 conn_id,
@@ -304,13 +305,12 @@ impl PoolingBackend {
         tracing::Span::current().record("conn_id", display(conn_id));
         info!(%conn_id, "local_pool: opening a new connection '{conn_info}'");
 
+        let mut node_info = local_backend.node_info.clone();
+
         let (key, jwk) = create_random_jwk();
 
-        let mut config = local_backend
-            .node_info
-            .conn_info
-            .to_postgres_client_config();
-        config
+        let config = node_info
+            .config
             .user(&conn_info.user_info.user)
             .dbname(&conn_info.dbname)
             .set_param(
@@ -322,7 +322,7 @@ impl PoolingBackend {
             );
 
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let (client, connection) = config.connect(&postgres_client::NoTls).await?;
+        let (client, connection) = config.connect(postgres_client::NoTls).await?;
         drop(pause);
 
         let pid = client.get_process_id();
@@ -336,7 +336,7 @@ impl PoolingBackend {
             connection,
             key,
             conn_id,
-            local_backend.node_info.aux.clone(),
+            node_info.aux.clone(),
         );
 
         {
@@ -512,16 +512,19 @@ impl ConnectMechanism for TokioMechanism {
         node_info: &CachedNodeInfo,
         compute_config: &ComputeConfig,
     ) -> Result<Self::Connection, Self::ConnectError> {
-        let permit = self.locks.get_permit(&node_info.conn_info.host).await?;
+        let host = node_info.config.get_host();
+        let permit = self.locks.get_permit(&host).await?;
 
-        let mut config = node_info.conn_info.to_postgres_client_config();
+        let mut config = (*node_info.config).clone();
         let config = config
             .user(&self.conn_info.user_info.user)
             .dbname(&self.conn_info.dbname)
             .connect_timeout(compute_config.timeout);
 
+        let mk_tls =
+            crate::tls::postgres_rustls::MakeRustlsConnect::new(compute_config.tls.clone());
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
-        let res = config.connect(compute_config).await;
+        let res = config.connect(mk_tls).await;
         drop(pause);
         let (client, connection) = permit.release_result(res)?;
 
@@ -545,6 +548,8 @@ impl ConnectMechanism for TokioMechanism {
             node_info.aux.clone(),
         ))
     }
+
+    fn update_connect_config(&self, _config: &mut compute::ConnCfg) {}
 }
 
 struct HyperMechanism {
@@ -568,20 +573,20 @@ impl ConnectMechanism for HyperMechanism {
         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 host_addr = node_info.config.get_host_addr();
+        let host = node_info.config.get_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 {
+        let tls = if node_info.config.get_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;
+        let port = node_info.config.get_port();
+        let res = connect_http2(host_addr, &host, port, config.timeout, tls).await;
         drop(pause);
         let (client, connection) = permit.release_result(res)?;
 
@@ -604,6 +609,8 @@ impl ConnectMechanism for HyperMechanism {
             node_info.aux.clone(),
         ))
     }
+
+    fn update_connect_config(&self, _config: &mut compute::ConnCfg) {}
 }
 
 async fn connect_http2(

proxy/src/serverless/conn_pool.rs

@@ -23,12 +23,12 @@ use super::conn_pool_lib::{
     Client, ClientDataEnum, ClientInnerCommon, ClientInnerExt, ConnInfo, EndpointConnPool,
     GlobalConnPool,
 };
-use crate::config::ComputeConfig;
 use crate::context::RequestContext;
 use crate::control_plane::messages::MetricsAuxInfo;
 use crate::metrics::Metrics;
+use crate::tls::postgres_rustls::MakeRustlsConnect;
 
-type TlsStream = <ComputeConfig as MakeTlsConnect<TcpStream>>::Stream;
+type TlsStream = <MakeRustlsConnect as MakeTlsConnect<TcpStream>>::Stream;
 
 #[derive(Debug, Clone)]
 pub(crate) struct ConnInfoWithAuth {

proxy/src/serverless/mod.rs

@@ -49,7 +49,7 @@ use crate::config::{ProxyConfig, ProxyProtocolV2};
 use crate::context::RequestContext;
 use crate::ext::TaskExt;
 use crate::metrics::Metrics;
-use crate::protocol2::{ConnectHeader, ConnectionInfo, read_proxy_protocol};
+use crate::protocol2::{ChainRW, ConnectHeader, ConnectionInfo, read_proxy_protocol};
 use crate::proxy::run_until_cancelled;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::serverless::backend::PoolingBackend;
@@ -207,12 +207,12 @@ pub(crate) type AsyncRW = Pin<Box<dyn AsyncReadWrite>>;
 
 #[async_trait]
 trait MaybeTlsAcceptor: Send + Sync + 'static {
-    async fn accept(&self, conn: TcpStream) -> std::io::Result<AsyncRW>;
+    async fn accept(&self, conn: ChainRW<TcpStream>) -> std::io::Result<AsyncRW>;
 }
 
 #[async_trait]
 impl MaybeTlsAcceptor for &'static ArcSwapOption<crate::config::TlsConfig> {
-    async fn accept(&self, conn: TcpStream) -> std::io::Result<AsyncRW> {
+    async fn accept(&self, conn: ChainRW<TcpStream>) -> std::io::Result<AsyncRW> {
         match &*self.load() {
             Some(config) => Ok(Box::pin(
                 TlsAcceptor::from(config.http_config.clone())
@@ -235,30 +235,33 @@ async fn connection_startup(
     peer_addr: SocketAddr,
 ) -> Option<(AsyncRW, ConnectionInfo)> {
     // handle PROXY protocol
-    let (conn, conn_info) = match config.proxy_protocol_v2 {
-        ProxyProtocolV2::Required => {
-            match read_proxy_protocol(conn).await {
-                Err(e) => {
-                    warn!("per-client task finished with an error: {e:#}");
-                    return None;
-                }
-                // our load balancers will not send any more data. let's just exit immediately
-                Ok((_conn, ConnectHeader::Local)) => {
-                    tracing::debug!("healthcheck received");
-                    return None;
-                }
-                Ok((conn, ConnectHeader::Proxy(info))) => (conn, info),
-            }
+    let (conn, peer) = match read_proxy_protocol(conn).await {
+        Ok(c) => c,
+        Err(e) => {
+            tracing::warn!(?session_id, %peer_addr, "failed to accept TCP connection: invalid PROXY protocol V2 header: {e:#}");
+            return None;
         }
-        // ignore the header - it cannot be confused for a postgres or http connection so will
-        // error later.
-        ProxyProtocolV2::Rejected => (
-            conn,
-            ConnectionInfo {
-                addr: peer_addr,
-                extra: None,
-            },
-        ),
+    };
+
+    let conn_info = match peer {
+        // our load balancers will not send any more data. let's just exit immediately
+        ConnectHeader::Local => {
+            tracing::debug!("healthcheck received");
+            return None;
+        }
+        ConnectHeader::Missing if config.proxy_protocol_v2 == ProxyProtocolV2::Required => {
+            tracing::warn!("missing required proxy protocol header");
+            return None;
+        }
+        ConnectHeader::Proxy(_) if config.proxy_protocol_v2 == ProxyProtocolV2::Rejected => {
+            tracing::warn!("proxy protocol header not supported");
+            return None;
+        }
+        ConnectHeader::Proxy(info) => info,
+        ConnectHeader::Missing => ConnectionInfo {
+            addr: peer_addr,
+            extra: None,
+        },
     };
 
     let has_private_peer_addr = match conn_info.addr.ip() {

proxy/src/tls/postgres_rustls.rs

@@ -2,11 +2,10 @@ use std::convert::TryFrom;
 use std::sync::Arc;
 
 use postgres_client::tls::MakeTlsConnect;
-use rustls::pki_types::{InvalidDnsNameError, ServerName};
+use rustls::ClientConfig;
+use rustls::pki_types::ServerName;
 use tokio::io::{AsyncRead, AsyncWrite};
 
-use crate::config::ComputeConfig;
-
 mod private {
     use std::future::Future;
     use std::io;
@@ -124,27 +123,36 @@ mod private {
     }
 }
 
-impl<S> MakeTlsConnect<S> for ComputeConfig
+/// A `MakeTlsConnect` implementation using `rustls`.
+///
+/// That way you can connect to PostgreSQL using `rustls` as the TLS stack.
+#[derive(Clone)]
+pub struct MakeRustlsConnect {
+    pub config: Arc<ClientConfig>,
+}
+
+impl MakeRustlsConnect {
+    /// Creates a new `MakeRustlsConnect` from the provided `ClientConfig`.
+    #[must_use]
+    pub fn new(config: Arc<ClientConfig>) -> Self {
+        Self { config }
+    }
+}
+
+impl<S> MakeTlsConnect<S> for MakeRustlsConnect
 where
     S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
 {
     type Stream = private::RustlsStream<S>;
     type TlsConnect = private::RustlsConnect;
-    type Error = InvalidDnsNameError;
+    type Error = rustls::pki_types::InvalidDnsNameError;
 
-    fn make_tls_connect(&self, hostname: &str) -> Result<Self::TlsConnect, Self::Error> {
-        make_tls_connect(&self.tls, hostname)
+    fn make_tls_connect(&mut self, hostname: &str) -> Result<Self::TlsConnect, Self::Error> {
+        ServerName::try_from(hostname).map(|dns_name| {
+            private::RustlsConnect(private::RustlsConnectData {
+                hostname: dns_name.to_owned(),
+                connector: Arc::clone(&self.config).into(),
+            })
+        })
     }
 }
-
-pub fn make_tls_connect(
-    tls: &Arc<rustls::ClientConfig>,
-    hostname: &str,
-) -> Result<private::RustlsConnect, InvalidDnsNameError> {
-    ServerName::try_from(hostname).map(|dns_name| {
-        private::RustlsConnect(private::RustlsConnectData {
-            hostname: dns_name.to_owned(),
-            connector: tls.clone().into(),
-        })
-    })
-}

safekeeper/client/src/mgmt_api.rs

@@ -8,8 +8,8 @@ use std::error::Error as _;
 use http_utils::error::HttpErrorBody;
 use reqwest::{IntoUrl, Method, StatusCode};
 use safekeeper_api::models::{
-    self, PullTimelineRequest, PullTimelineResponse, SafekeeperStatus, SafekeeperUtilization,
-    TimelineCreateRequest, TimelineStatus,
+    self, PullTimelineRequest, PullTimelineResponse, SafekeeperUtilization, TimelineCreateRequest,
+    TimelineStatus,
 };
 use utils::id::{NodeId, TenantId, TimelineId};
 use utils::logging::SecretString;
@@ -183,12 +183,6 @@ impl Client {
         self.get(&uri).await
     }
 
-    pub async fn status(&self) -> Result<SafekeeperStatus> {
-        let uri = format!("{}/v1/status", self.mgmt_api_endpoint);
-        let resp = self.get(&uri).await?;
-        resp.json().await.map_err(Error::ReceiveBody)
-    }
-
     pub async fn utilization(&self) -> Result<SafekeeperUtilization> {
         let uri = format!("{}/v1/utilization", self.mgmt_api_endpoint);
         let resp = self.get(&uri).await?;

storage_controller/migrations/2025-06-01-201442_add_lifecycle_to_nodes/down.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes DROP COLUMN lifecycle;

storage_controller/migrations/2025-06-01-201442_add_lifecycle_to_nodes/up.sql

@@ -1 +0,0 @@
-ALTER TABLE nodes ADD COLUMN lifecycle VARCHAR NOT NULL DEFAULT 'active';

storage_controller/src/http.rs

@@ -907,42 +907,6 @@ async fn handle_node_delete(req: Request<Body>) -> Result<Response<Body>, ApiErr
     json_response(StatusCode::OK, state.service.node_delete(node_id).await?)
 }
 
-async fn handle_tombstone_list(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let req = match maybe_forward(req).await {
-        ForwardOutcome::Forwarded(res) => {
-            return res;
-        }
-        ForwardOutcome::NotForwarded(req) => req,
-    };
-
-    let state = get_state(&req);
-    let mut nodes = state.service.tombstone_list().await?;
-    nodes.sort_by_key(|n| n.get_id());
-    let api_nodes = nodes.into_iter().map(|n| n.describe()).collect::<Vec<_>>();
-
-    json_response(StatusCode::OK, api_nodes)
-}
-
-async fn handle_tombstone_delete(req: Request<Body>) -> Result<Response<Body>, ApiError> {
-    check_permissions(&req, Scope::Admin)?;
-
-    let req = match maybe_forward(req).await {
-        ForwardOutcome::Forwarded(res) => {
-            return res;
-        }
-        ForwardOutcome::NotForwarded(req) => req,
-    };
-
-    let state = get_state(&req);
-    let node_id: NodeId = parse_request_param(&req, "node_id")?;
-    json_response(
-        StatusCode::OK,
-        state.service.tombstone_delete(node_id).await?,
-    )
-}
-
 async fn handle_node_configure(req: Request<Body>) -> Result<Response<Body>, ApiError> {
     check_permissions(&req, Scope::Admin)?;
 
@@ -2098,20 +2062,6 @@ pub fn make_router(
         .post("/debug/v1/node/:node_id/drop", |r| {
             named_request_span(r, handle_node_drop, RequestName("debug_v1_node_drop"))
         })
-        .delete("/debug/v1/tombstone/:node_id", |r| {
-            named_request_span(
-                r,
-                handle_tombstone_delete,
-                RequestName("debug_v1_tombstone_delete"),
-            )
-        })
-        .get("/debug/v1/tombstone", |r| {
-            named_request_span(
-                r,
-                handle_tombstone_list,
-                RequestName("debug_v1_tombstone_list"),
-            )
-        })
         .post("/debug/v1/tenant/:tenant_id/import", |r| {
             named_request_span(
                 r,

storage_controller/src/node.rs

@@ -2,7 +2,7 @@ use std::str::FromStr;
 use std::time::Duration;
 
 use pageserver_api::controller_api::{
-    AvailabilityZone, NodeAvailability, NodeDescribeResponse, NodeLifecycle, NodeRegisterRequest,
+    AvailabilityZone, NodeAvailability, NodeDescribeResponse, NodeRegisterRequest,
     NodeSchedulingPolicy, TenantLocateResponseShard,
 };
 use pageserver_api::shard::TenantShardId;
@@ -29,7 +29,6 @@ pub(crate) struct Node {
 
     availability: NodeAvailability,
     scheduling: NodeSchedulingPolicy,
-    lifecycle: NodeLifecycle,
 
     listen_http_addr: String,
     listen_http_port: u16,
@@ -229,7 +228,6 @@ impl Node {
             listen_pg_addr,
             listen_pg_port,
             scheduling: NodeSchedulingPolicy::Active,
-            lifecycle: NodeLifecycle::Active,
             availability: NodeAvailability::Offline,
             availability_zone_id,
             use_https,
@@ -241,7 +239,6 @@ impl Node {
         NodePersistence {
             node_id: self.id.0 as i64,
             scheduling_policy: self.scheduling.into(),
-            lifecycle: self.lifecycle.into(),
             listen_http_addr: self.listen_http_addr.clone(),
             listen_http_port: self.listen_http_port as i32,
             listen_https_port: self.listen_https_port.map(|x| x as i32),
@@ -266,7 +263,6 @@ impl Node {
             availability: NodeAvailability::Offline,
             scheduling: NodeSchedulingPolicy::from_str(&np.scheduling_policy)
                 .expect("Bad scheduling policy in DB"),
-            lifecycle: NodeLifecycle::from_str(&np.lifecycle).expect("Bad lifecycle in DB"),
             listen_http_addr: np.listen_http_addr,
             listen_http_port: np.listen_http_port as u16,
             listen_https_port: np.listen_https_port.map(|x| x as u16),

storage_controller/src/persistence.rs

@@ -19,7 +19,7 @@ use futures::FutureExt;
 use futures::future::BoxFuture;
 use itertools::Itertools;
 use pageserver_api::controller_api::{
-    AvailabilityZone, MetadataHealthRecord, NodeLifecycle, NodeSchedulingPolicy, PlacementPolicy,
+    AvailabilityZone, MetadataHealthRecord, NodeSchedulingPolicy, PlacementPolicy,
     SafekeeperDescribeResponse, ShardSchedulingPolicy, SkSchedulingPolicy,
 };
 use pageserver_api::models::{ShardImportStatus, TenantConfig};
@@ -102,7 +102,6 @@ pub(crate) enum DatabaseOperation {
     UpdateNode,
     DeleteNode,
     ListNodes,
-    ListTombstones,
     BeginShardSplit,
     CompleteShardSplit,
     AbortShardSplit,
@@ -358,8 +357,6 @@ impl Persistence {
     }
 
     /// When a node is first registered, persist it before using it for anything
-    /// If the provided node_id already exists, it will be error.
-    /// The common case is when a node marked for deletion wants to register.
     pub(crate) async fn insert_node(&self, node: &Node) -> DatabaseResult<()> {
         let np = &node.to_persistent();
         self.with_measured_conn(DatabaseOperation::InsertNode, move |conn| {
@@ -376,41 +373,19 @@ impl Persistence {
 
     /// At startup, populate the list of nodes which our shards may be placed on
     pub(crate) async fn list_nodes(&self) -> DatabaseResult<Vec<NodePersistence>> {
-        use crate::schema::nodes::dsl::*;
-
-        let result: Vec<NodePersistence> = self
+        let nodes: Vec<NodePersistence> = self
             .with_measured_conn(DatabaseOperation::ListNodes, move |conn| {
                 Box::pin(async move {
                     Ok(crate::schema::nodes::table
-                        .filter(lifecycle.ne(String::from(NodeLifecycle::Deleted)))
                         .load::<NodePersistence>(conn)
                         .await?)
                 })
             })
             .await?;
 
-        tracing::info!("list_nodes: loaded {} nodes", result.len());
+        tracing::info!("list_nodes: loaded {} nodes", nodes.len());
 
-        Ok(result)
-    }
-
-    pub(crate) async fn list_tombstones(&self) -> DatabaseResult<Vec<NodePersistence>> {
-        use crate::schema::nodes::dsl::*;
-
-        let result: Vec<NodePersistence> = self
-            .with_measured_conn(DatabaseOperation::ListTombstones, move |conn| {
-                Box::pin(async move {
-                    Ok(crate::schema::nodes::table
-                        .filter(lifecycle.eq(String::from(NodeLifecycle::Deleted)))
-                        .load::<NodePersistence>(conn)
-                        .await?)
-                })
-            })
-            .await?;
-
-        tracing::info!("list_tombstones: loaded {} nodes", result.len());
-
-        Ok(result)
+        Ok(nodes)
     }
 
     pub(crate) async fn update_node<V>(
@@ -429,7 +404,6 @@ impl Persistence {
                 Box::pin(async move {
                     let updated = diesel::update(nodes)
                         .filter(node_id.eq(input_node_id.0 as i64))
-                        .filter(lifecycle.ne(String::from(NodeLifecycle::Deleted)))
                         .set(values)
                         .execute(conn)
                         .await?;
@@ -473,57 +447,6 @@ impl Persistence {
         .await
     }
 
-    /// Tombstone is a special state where the node is not deleted from the database,
-    /// but it is not available for usage.
-    /// The main reason for it is to prevent the flaky node to register.
-    pub(crate) async fn set_tombstone(&self, del_node_id: NodeId) -> DatabaseResult<()> {
-        use crate::schema::nodes::dsl::*;
-        self.update_node(
-            del_node_id,
-            lifecycle.eq(String::from(NodeLifecycle::Deleted)),
-        )
-        .await
-    }
-
-    pub(crate) async fn delete_node(&self, del_node_id: NodeId) -> DatabaseResult<()> {
-        use crate::schema::nodes::dsl::*;
-        self.with_measured_conn(DatabaseOperation::DeleteNode, move |conn| {
-            Box::pin(async move {
-                // You can hard delete a node only if it has a tombstone.
-                // So we need to check if the node has lifecycle set to deleted.
-                let node_to_delete = nodes
-                    .filter(node_id.eq(del_node_id.0 as i64))
-                    .first::<NodePersistence>(conn)
-                    .await
-                    .optional()?;
-
-                if let Some(np) = node_to_delete {
-                    let lc = NodeLifecycle::from_str(&np.lifecycle).map_err(|e| {
-                        DatabaseError::Logical(format!(
-                            "Node {} has invalid lifecycle: {}",
-                            del_node_id, e
-                        ))
-                    })?;
-
-                    if lc != NodeLifecycle::Deleted {
-                        return Err(DatabaseError::Logical(format!(
-                            "Node {} was not soft deleted before, cannot hard delete it",
-                            del_node_id
-                        )));
-                    }
-
-                    diesel::delete(nodes)
-                        .filter(node_id.eq(del_node_id.0 as i64))
-                        .execute(conn)
-                        .await?;
-                }
-
-                Ok(())
-            })
-        })
-        .await
-    }
-
     /// At startup, load the high level state for shards, such as their config + policy.  This will
     /// be enriched at runtime with state discovered on pageservers.
     ///
@@ -620,6 +543,21 @@ impl Persistence {
         .await
     }
 
+    pub(crate) async fn delete_node(&self, del_node_id: NodeId) -> DatabaseResult<()> {
+        use crate::schema::nodes::dsl::*;
+        self.with_measured_conn(DatabaseOperation::DeleteNode, move |conn| {
+            Box::pin(async move {
+                diesel::delete(nodes)
+                    .filter(node_id.eq(del_node_id.0 as i64))
+                    .execute(conn)
+                    .await?;
+
+                Ok(())
+            })
+        })
+        .await
+    }
+
     /// When a tenant invokes the /re-attach API, this function is responsible for doing an efficient
     /// batched increment of the generations of all tenants whose generation_pageserver is equal to
     /// the node that called /re-attach.
@@ -633,20 +571,6 @@ impl Persistence {
         let updated = self
             .with_measured_conn(DatabaseOperation::ReAttach, move |conn| {
                 Box::pin(async move {
-                    // Check if the node is not marked as deleted
-                    let deleted_node: i64 = nodes
-                        .filter(node_id.eq(input_node_id.0 as i64))
-                        .filter(lifecycle.eq(String::from(NodeLifecycle::Deleted)))
-                        .count()
-                        .get_result(conn)
-                        .await?;
-                    if deleted_node > 0 {
-                        return Err(DatabaseError::Logical(format!(
-                            "Node {} is marked as deleted, re-attach is not allowed",
-                            input_node_id
-                        )));
-                    }
-
                     let rows_updated = diesel::update(tenant_shards)
                         .filter(generation_pageserver.eq(input_node_id.0 as i64))
                         .set(generation.eq(generation + 1))
@@ -2124,7 +2048,6 @@ pub(crate) struct NodePersistence {
     pub(crate) listen_pg_port: i32,
     pub(crate) availability_zone_id: String,
     pub(crate) listen_https_port: Option<i32>,
-    pub(crate) lifecycle: String,
 }
 
 /// Tenant metadata health status that are stored durably.

storage_controller/src/schema.rs

@@ -33,7 +33,6 @@ diesel::table! {
         listen_pg_port -> Int4,
         availability_zone_id -> Varchar,
         listen_https_port -> Nullable<Int4>,
-        lifecycle -> Varchar,
     }
 }
 

storage_controller/src/service.rs

@@ -166,7 +166,6 @@ enum NodeOperations {
     Register,
     Configure,
     Delete,
-    DeleteTombstone,
 }
 
 /// The leadership status for the storage controller process.
@@ -6910,7 +6909,7 @@ impl Service {
     /// detaching or deleting it on pageservers.  We do not try and re-schedule any
     /// tenants that were on this node.
     pub(crate) async fn node_drop(&self, node_id: NodeId) -> Result<(), ApiError> {
-        self.persistence.set_tombstone(node_id).await?;
+        self.persistence.delete_node(node_id).await?;
 
         let mut locked = self.inner.write().unwrap();
 
@@ -7034,10 +7033,9 @@ impl Service {
         // That is safe because in Service::spawn we only use generation_pageserver if it refers to a node
         // that exists.
 
-        // 2. Actually delete the node from in-memory state and set tombstone to the database
-        // for preventing the node to register again.
+        // 2. Actually delete the node from the database and from in-memory state
         tracing::info!("Deleting node from database");
-        self.persistence.set_tombstone(node_id).await?;
+        self.persistence.delete_node(node_id).await?;
 
         Ok(())
     }
@@ -7056,35 +7054,6 @@ impl Service {
         Ok(nodes)
     }
 
-    pub(crate) async fn tombstone_list(&self) -> Result<Vec<Node>, ApiError> {
-        self.persistence
-            .list_tombstones()
-            .await?
-            .into_iter()
-            .map(|np| Node::from_persistent(np, false))
-            .collect::<Result<Vec<_>, _>>()
-            .map_err(ApiError::InternalServerError)
-    }
-
-    pub(crate) async fn tombstone_delete(&self, node_id: NodeId) -> Result<(), ApiError> {
-        let _node_lock = trace_exclusive_lock(
-            &self.node_op_locks,
-            node_id,
-            NodeOperations::DeleteTombstone,
-        )
-        .await;
-
-        if matches!(self.get_node(node_id).await, Err(ApiError::NotFound(_))) {
-            self.persistence.delete_node(node_id).await?;
-            Ok(())
-        } else {
-            Err(ApiError::Conflict(format!(
-                "Node {} is in use, consider using tombstone API first",
-                node_id
-            )))
-        }
-    }
-
     pub(crate) async fn get_node(&self, node_id: NodeId) -> Result<Node, ApiError> {
         self.inner
             .read()
@@ -7255,25 +7224,7 @@ impl Service {
         };
 
         match registration_status {
-            RegistrationStatus::New => {
-                self.persistence.insert_node(&new_node).await.map_err(|e| {
-                    if matches!(
-                        e,
-                        crate::persistence::DatabaseError::Query(
-                            diesel::result::Error::DatabaseError(
-                                diesel::result::DatabaseErrorKind::UniqueViolation,
-                                _,
-                            )
-                        )
-                    ) {
-                        // The node can be deleted by tombstone API, and not show up in the list of nodes.
-                        // If you see this error, check tombstones first.
-                        ApiError::Conflict(format!("Node {} is already exists", new_node.get_id()))
-                    } else {
-                        ApiError::from(e)
-                    }
-                })?;
-            }
+            RegistrationStatus::New => self.persistence.insert_node(&new_node).await?,
             RegistrationStatus::NeedUpdate => {
                 self.persistence
                     .update_node_on_registration(

test_runner/fixtures/neon_fixtures.py

@@ -357,6 +357,31 @@ class PgProtocol:
         return TimelineId(cast("str", self.safe_psql("show neon.timeline_id")[0][0]))
 
 
+class PageserverWalReceiverProtocol(StrEnum):
+    VANILLA = "vanilla"
+    INTERPRETED = "interpreted"
+
+    @staticmethod
+    def to_config_key_value(proto) -> tuple[str, dict[str, Any]]:
+        if proto == PageserverWalReceiverProtocol.VANILLA:
+            return (
+                "wal_receiver_protocol",
+                {
+                    "type": "vanilla",
+                },
+            )
+        elif proto == PageserverWalReceiverProtocol.INTERPRETED:
+            return (
+                "wal_receiver_protocol",
+                {
+                    "type": "interpreted",
+                    "args": {"format": "protobuf", "compression": {"zstd": {"level": 1}}},
+                },
+            )
+        else:
+            raise ValueError(f"Unknown protocol type: {proto}")
+
+
 @dataclass
 class PageserverTracingConfig:
     sampling_ratio: tuple[int, int]
@@ -398,7 +423,6 @@ class PageserverImportConfig:
             "import_job_concurrency": self.import_job_concurrency,
             "import_job_soft_size_limit": self.import_job_soft_size_limit,
             "import_job_checkpoint_threshold": self.import_job_checkpoint_threshold,
-            "import_job_max_byte_range_size": 4 * 1024 * 1024,  # Pageserver default
         }
         return ("timeline_import_config", value)
 
@@ -450,6 +474,7 @@ class NeonEnvBuilder:
         safekeeper_extra_opts: list[str] | None = None,
         storage_controller_port_override: int | None = None,
         pageserver_virtual_file_io_mode: str | None = None,
+        pageserver_wal_receiver_protocol: PageserverWalReceiverProtocol | None = None,
         pageserver_get_vectored_concurrent_io: str | None = None,
         pageserver_tracing_config: PageserverTracingConfig | None = None,
         pageserver_import_config: PageserverImportConfig | None = None,
@@ -526,6 +551,11 @@ class NeonEnvBuilder:
 
         self.pageserver_virtual_file_io_mode = pageserver_virtual_file_io_mode
 
+        if pageserver_wal_receiver_protocol is not None:
+            self.pageserver_wal_receiver_protocol = pageserver_wal_receiver_protocol
+        else:
+            self.pageserver_wal_receiver_protocol = PageserverWalReceiverProtocol.INTERPRETED
+
         assert test_name.startswith("test_"), (
             "Unexpectedly instantiated from outside a test function"
         )
@@ -1171,6 +1201,7 @@ class NeonEnv:
 
         self.pageserver_virtual_file_io_engine = config.pageserver_virtual_file_io_engine
         self.pageserver_virtual_file_io_mode = config.pageserver_virtual_file_io_mode
+        self.pageserver_wal_receiver_protocol = config.pageserver_wal_receiver_protocol
         self.pageserver_get_vectored_concurrent_io = config.pageserver_get_vectored_concurrent_io
         self.pageserver_tracing_config = config.pageserver_tracing_config
         if config.pageserver_import_config is None:
@@ -1302,6 +1333,13 @@ class NeonEnv:
                         for key, value in override.items():
                             ps_cfg[key] = value
 
+            if self.pageserver_wal_receiver_protocol is not None:
+                key, value = PageserverWalReceiverProtocol.to_config_key_value(
+                    self.pageserver_wal_receiver_protocol
+                )
+                if key not in ps_cfg:
+                    ps_cfg[key] = value
+
             if self.pageserver_tracing_config is not None:
                 key, value = self.pageserver_tracing_config.to_config_key_value()
 
@@ -2054,14 +2092,6 @@ class NeonStorageController(MetricsGetter, LogUtils):
             headers=self.headers(TokenScope.ADMIN),
         )
 
-    def tombstone_delete(self, node_id):
-        log.info(f"tombstone_delete({node_id})")
-        self.request(
-            "DELETE",
-            f"{self.api}/debug/v1/tombstone/{node_id}",
-            headers=self.headers(TokenScope.ADMIN),
-        )
-
     def node_drain(self, node_id):
         log.info(f"node_drain({node_id})")
         self.request(
@@ -2118,14 +2148,6 @@ class NeonStorageController(MetricsGetter, LogUtils):
         )
         return response.json()
 
-    def tombstone_list(self):
-        response = self.request(
-            "GET",
-            f"{self.api}/debug/v1/tombstone",
-            headers=self.headers(TokenScope.ADMIN),
-        )
-        return response.json()
-
     def tenant_shard_dump(self):
         """
         Debug listing API: dumps the internal map of tenant shards
@@ -4688,7 +4710,7 @@ class EndpointFactory:
         origin: Endpoint,
         endpoint_id: str | None = None,
         config_lines: list[str] | None = None,
-    ) -> Endpoint:
+    ):
         branch_name = origin.branch_name
         assert origin in self.endpoints
         assert branch_name is not None
@@ -4707,7 +4729,7 @@ class EndpointFactory:
         origin: Endpoint,
         endpoint_id: str | None = None,
         config_lines: list[str] | None = None,
-    ) -> Endpoint:
+    ):
         branch_name = origin.branch_name
         assert origin in self.endpoints
         assert branch_name is not None

test_runner/performance/test_sharded_ingest.py

@@ -15,10 +15,19 @@ from fixtures.neon_fixtures import (
 
 @pytest.mark.timeout(1200)
 @pytest.mark.parametrize("shard_count", [1, 8, 32])
+@pytest.mark.parametrize(
+    "wal_receiver_protocol",
+    [
+        "vanilla",
+        "interpreted-bincode-compressed",
+        "interpreted-protobuf-compressed",
+    ],
+)
 def test_sharded_ingest(
     neon_env_builder: NeonEnvBuilder,
     zenbenchmark: NeonBenchmarker,
     shard_count: int,
+    wal_receiver_protocol: str,
 ):
     """
     Benchmarks sharded ingestion throughput, by ingesting a large amount of WAL into a Safekeeper
@@ -30,6 +39,36 @@ def test_sharded_ingest(
     neon_env_builder.num_pageservers = shard_count
     env = neon_env_builder.init_configs()
 
+    for ps in env.pageservers:
+        if wal_receiver_protocol == "vanilla":
+            ps.patch_config_toml_nonrecursive(
+                {
+                    "wal_receiver_protocol": {
+                        "type": "vanilla",
+                    }
+                }
+            )
+        elif wal_receiver_protocol == "interpreted-bincode-compressed":
+            ps.patch_config_toml_nonrecursive(
+                {
+                    "wal_receiver_protocol": {
+                        "type": "interpreted",
+                        "args": {"format": "bincode", "compression": {"zstd": {"level": 1}}},
+                    }
+                }
+            )
+        elif wal_receiver_protocol == "interpreted-protobuf-compressed":
+            ps.patch_config_toml_nonrecursive(
+                {
+                    "wal_receiver_protocol": {
+                        "type": "interpreted",
+                        "args": {"format": "protobuf", "compression": {"zstd": {"level": 1}}},
+                    }
+                }
+            )
+        else:
+            raise AssertionError("Test must use explicit wal receiver protocol config")
+
     env.start()
 
     # Create a sharded tenant and timeline, and migrate it to the respective pageservers. Ensure

test_runner/regress/test_attach_tenant_config.py

@@ -182,6 +182,10 @@ def test_fully_custom_config(positive_env: NeonEnv):
         "lsn_lease_length": "1m",
         "lsn_lease_length_for_ts": "5s",
         "timeline_offloading": False,
+        "wal_receiver_protocol_override": {
+            "type": "interpreted",
+            "args": {"format": "bincode", "compression": {"zstd": {"level": 1}}},
+        },
         "rel_size_v2_enabled": True,
         "relsize_snapshot_cache_capacity": 10000,
         "gc_compaction_enabled": True,