Remove notion of ParsedSpec

Signed-off-by: Tristan Partin <tristan@neon.tech>

Cargo.lock

@@ -753,7 +753,6 @@ dependencies = [
  "axum",
  "axum-core",
  "bytes",
- "form_urlencoded",
  "futures-util",
  "headers",
  "http 1.1.0",
@@ -762,8 +761,6 @@ dependencies = [
  "mime",
  "pin-project-lite",
  "serde",
- "serde_html_form",
- "serde_path_to_error",
  "tower 0.5.2",
  "tower-layer",
  "tower-service",
@@ -1281,10 +1278,13 @@ dependencies = [
  "chrono",
  "indexmap 2.9.0",
  "jsonwebtoken",
+ "postgres",
  "regex",
  "remote_storage",
  "serde",
  "serde_json",
+ "tokio-postgres",
+ "url",
  "utils",
 ]
 
@@ -6425,19 +6425,6 @@ dependencies = [
  "syn 2.0.100",
 ]
 
-[[package]]
-name = "serde_html_form"
-version = "0.2.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9d2de91cf02bbc07cde38891769ccd5d4f073d22a40683aa4bc7a95781aaa2c4"
-dependencies = [
- "form_urlencoded",
- "indexmap 2.9.0",
- "itoa",
- "ryu",
- "serde",
-]
-
 [[package]]
 name = "serde_json"
 version = "1.0.125"
@@ -6763,6 +6750,7 @@ dependencies = [
  "chrono",
  "clap",
  "clashmap",
+ "compute_api",
  "control_plane",
  "cron",
  "diesel",

Cargo.toml

@@ -71,7 +71,7 @@ aws-credential-types = "1.2.0"
 aws-sigv4 = { version = "1.2", features = ["sign-http"] }
 aws-types = "1.3"
 axum = { version = "0.8.1", features = ["ws"] }
-axum-extra = { version = "0.10.0", features = ["typed-header", "query"] }
+axum-extra = { version = "0.10.0", features = ["typed-header"] }
 base64 = "0.13.0"
 bincode = "1.3"
 bindgen = "0.71"

compute_tools/src/bin/compute_ctl.rs

@@ -145,28 +145,47 @@ impl Cli {
     }
 }
 
-fn main() -> Result<()> {
+impl Cli {
+    pub fn get_config(&self) -> Result<ComputeConfig> {
+        // First, read the config from the path if provided
+        if let Some(ref config) = self.config {
+            let file = File::open(config)?;
+            return Ok(serde_json::from_reader(&file)?);
+        }
+
+        // If the config wasn't provided in the CLI arguments, then retrieve it from
+        // the control plane
+        match get_config_from_control_plane(
+            self.control_plane_uri.as_ref().unwrap(),
+            &self.compute_id,
+        ) {
+            Ok(config) => Ok(config),
+            Err(e) => {
+                error!(
+                    "cannot get response from control plane: {}\n\
+                    neither spec nor confirmation that compute is in the Empty state was received",
+                    e
+                );
+                Err(e)
+            }
+        }
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
     let cli = Cli::parse();
 
     let scenario = failpoint_support::init();
 
-    // For historical reasons, the main thread that processes the config and launches postgres
-    // is synchronous, but we always have this tokio runtime available and we "enter" it so
-    // that you can use tokio::spawn() and tokio::runtime::Handle::current().block_on(...)
-    // from all parts of compute_ctl.
-    let runtime = tokio::runtime::Builder::new_multi_thread()
-        .enable_all()
-        .build()?;
-    let _rt_guard = runtime.enter();
-
-    runtime.block_on(init())?;
+    init().await?;
 
     // enable core dumping for all child processes
     setrlimit(Resource::CORE, rlimit::INFINITY, rlimit::INFINITY)?;
 
     let connstr = Url::parse(&cli.connstr).context("cannot parse connstr as a URL")?;
 
-    let config = get_config(&cli)?;
+    let config = cli.get_config()?;
 
     let compute_node = ComputeNode::new(
         ComputeNodeParams {
@@ -191,7 +210,7 @@ fn main() -> Result<()> {
         config,
     )?;
 
-    let exit_code = compute_node.run()?;
+    let exit_code = compute_node.run().await?;
 
     scenario.teardown();
 
@@ -213,28 +232,6 @@ async fn init() -> Result<()> {
     Ok(())
 }
 
-fn get_config(cli: &Cli) -> Result<ComputeConfig> {
-    // First, read the config from the path if provided
-    if let Some(ref config) = cli.config {
-        let file = File::open(config)?;
-        return Ok(serde_json::from_reader(&file)?);
-    }
-
-    // If the config wasn't provided in the CLI arguments, then retrieve it from
-    // the control plane
-    match get_config_from_control_plane(cli.control_plane_uri.as_ref().unwrap(), &cli.compute_id) {
-        Ok(config) => Ok(config),
-        Err(e) => {
-            error!(
-                "cannot get response from control plane: {}\n\
-                neither spec nor confirmation that compute is in the Empty state was received",
-                e
-            );
-            Err(e)
-        }
-    }
-}
-
 fn deinit_and_exit(exit_code: Option<i32>) -> ! {
     // Shutdown trace pipeline gracefully, so that it has a chance to send any
     // pending traces before we exit. Shutting down OTEL tracing provider may

compute_tools/src/compute.rs

@@ -15,12 +15,8 @@ use itertools::Itertools;
 use nix::sys::signal::{Signal, kill};
 use nix::unistd::Pid;
 use once_cell::sync::Lazy;
-use postgres;
-use postgres::NoTls;
-use postgres::error::SqlState;
 use remote_storage::{DownloadError, RemotePath};
 use std::collections::{HashMap, HashSet};
-use std::net::SocketAddr;
 use std::os::unix::fs::{PermissionsExt, symlink};
 use std::path::Path;
 use std::process::{Command, Stdio};
@@ -30,9 +26,9 @@ use std::sync::{Arc, Condvar, Mutex, RwLock};
 use std::time::{Duration, Instant};
 use std::{env, fs};
 use tokio::spawn;
+use tokio_postgres::{NoTls, error::SqlState};
 use tracing::{Instrument, debug, error, info, instrument, warn};
 use url::Url;
-use utils::id::{TenantId, TimelineId};
 use utils::lsn::Lsn;
 use utils::measured_stream::MeasuredReader;
 
@@ -144,7 +140,7 @@ pub struct ComputeState {
 
     /// Compute spec. This can be received from the CLI or - more likely -
     /// passed by the control plane with a /configure HTTP request.
-    pub pspec: Option<ParsedSpec>,
+    pub spec: Option<ComputeSpec>,
 
     /// If the spec is passed by a /configure request, 'startup_span' is the
     /// /configure request's tracing span. The main thread enters it when it
@@ -171,7 +167,7 @@ impl ComputeState {
             status: ComputeStatus::Empty,
             last_active: None,
             error: None,
-            pspec: None,
+            spec: None,
             startup_span: None,
             metrics: ComputeMetrics::default(),
             lfc_prewarm_state: LfcPrewarmState::default(),
@@ -203,94 +199,6 @@ impl Default for ComputeState {
     }
 }
 
-#[derive(Clone, Debug)]
-pub struct ParsedSpec {
-    pub spec: ComputeSpec,
-    pub tenant_id: TenantId,
-    pub timeline_id: TimelineId,
-    pub pageserver_connstr: String,
-    pub safekeeper_connstrings: Vec<String>,
-    pub storage_auth_token: Option<String>,
-    pub endpoint_storage_addr: Option<SocketAddr>,
-    pub endpoint_storage_token: Option<String>,
-}
-
-impl TryFrom<ComputeSpec> for ParsedSpec {
-    type Error = String;
-    fn try_from(spec: ComputeSpec) -> Result<Self, String> {
-        // Extract the options from the spec file that are needed to connect to
-        // the storage system.
-        //
-        // For backwards-compatibility, the top-level fields in the spec file
-        // may be empty. In that case, we need to dig them from the GUCs in the
-        // cluster.settings field.
-        let pageserver_connstr = spec
-            .pageserver_connstring
-            .clone()
-            .or_else(|| spec.cluster.settings.find("neon.pageserver_connstring"))
-            .ok_or("pageserver connstr should be provided")?;
-        let safekeeper_connstrings = if spec.safekeeper_connstrings.is_empty() {
-            if matches!(spec.mode, ComputeMode::Primary) {
-                spec.cluster
-                    .settings
-                    .find("neon.safekeepers")
-                    .ok_or("safekeeper connstrings should be provided")?
-                    .split(',')
-                    .map(|str| str.to_string())
-                    .collect()
-            } else {
-                vec![]
-            }
-        } else {
-            spec.safekeeper_connstrings.clone()
-        };
-        let storage_auth_token = spec.storage_auth_token.clone();
-        let tenant_id: TenantId = if let Some(tenant_id) = spec.tenant_id {
-            tenant_id
-        } else {
-            spec.cluster
-                .settings
-                .find("neon.tenant_id")
-                .ok_or("tenant id should be provided")
-                .map(|s| TenantId::from_str(&s))?
-                .or(Err("invalid tenant id"))?
-        };
-        let timeline_id: TimelineId = if let Some(timeline_id) = spec.timeline_id {
-            timeline_id
-        } else {
-            spec.cluster
-                .settings
-                .find("neon.timeline_id")
-                .ok_or("timeline id should be provided")
-                .map(|s| TimelineId::from_str(&s))?
-                .or(Err("invalid timeline id"))?
-        };
-
-        let endpoint_storage_addr: Option<SocketAddr> = spec
-            .endpoint_storage_addr
-            .clone()
-            .or_else(|| spec.cluster.settings.find("neon.endpoint_storage_addr"))
-            .unwrap_or_default()
-            .parse()
-            .ok();
-        let endpoint_storage_token = spec
-            .endpoint_storage_token
-            .clone()
-            .or_else(|| spec.cluster.settings.find("neon.endpoint_storage_token"));
-
-        Ok(ParsedSpec {
-            spec,
-            pageserver_connstr,
-            safekeeper_connstrings,
-            storage_auth_token,
-            tenant_id,
-            timeline_id,
-            endpoint_storage_addr,
-            endpoint_storage_token,
-        })
-    }
-}
-
 /// If we are a VM, returns a [`Command`] that will run in the `neon-postgres`
 /// cgroup. Otherwise returns the default `Command::new(cmd)`
 ///
@@ -368,10 +276,7 @@ impl ComputeNode {
         tokio_conn_conf.options(&options);
 
         let mut new_state = ComputeState::new();
-        if let Some(spec) = config.spec {
-            let pspec = ParsedSpec::try_from(spec).map_err(|msg| anyhow::anyhow!(msg))?;
-            new_state.pspec = Some(pspec);
-        }
+        new_state.spec = config.spec;
 
         Ok(ComputeNode {
             params,
@@ -386,10 +291,10 @@ impl ComputeNode {
 
     /// Top-level control flow of compute_ctl. Returns a process exit code we should
     /// exit with.
-    pub fn run(self) -> Result<Option<i32>> {
+    pub async fn run(self) -> Result<Option<i32>> {
         let this = Arc::new(self);
 
-        let cli_spec = this.state.lock().unwrap().pspec.clone();
+        let cli_spec = this.state.lock().unwrap().spec.clone();
 
         // If this is a pooled VM, prewarm before starting HTTP server and becoming
         // available for binding. Prewarming helps Postgres start quicker later,
@@ -425,7 +330,7 @@ impl ComputeNode {
 
         // If we got a spec from the CLI already, use that. Otherwise wait for the
         // control plane to pass it to us with a /configure HTTP request
-        let pspec = if let Some(cli_spec) = cli_spec {
+        let spec = if let Some(cli_spec) = cli_spec {
             cli_spec
         } else {
             this.wait_spec()?
@@ -438,11 +343,11 @@ impl ComputeNode {
         let mut vm_monitor = None;
         let mut pg_process: Option<PostgresHandle> = None;
 
-        match this.start_compute(&mut pg_process) {
+        match this.start_compute(&mut pg_process).await {
             Ok(()) => {
                 // Success! Launch remaining services (just vm-monitor currently)
                 vm_monitor =
-                    Some(this.start_vm_monitor(pspec.spec.disable_lfc_resizing.unwrap_or(false)));
+                    Some(this.start_vm_monitor(spec.disable_lfc_resizing.unwrap_or(false)));
             }
             Err(err) => {
                 // Something went wrong with the startup. Log it and expose the error to
@@ -487,7 +392,7 @@ impl ComputeNode {
         }
 
         // Reap the postgres process
-        delay_exit |= this.cleanup_after_postgres_exit()?;
+        delay_exit |= this.cleanup_after_postgres_exit().await?;
 
         // If launch failed, keep serving HTTP requests for a while, so the cloud
         // control plane can get the actual error.
@@ -498,7 +403,7 @@ impl ComputeNode {
         Ok(exit_code)
     }
 
-    pub fn wait_spec(&self) -> Result<ParsedSpec> {
+    pub fn wait_spec(&self) -> Result<ComputeSpec> {
         info!("no compute spec provided, waiting");
         let mut state = self.state.lock().unwrap();
         while state.status != ComputeStatus::ConfigurationPending {
@@ -506,7 +411,7 @@ impl ComputeNode {
         }
 
         info!("got spec, continue configuration");
-        let spec = state.pspec.as_ref().unwrap().clone();
+        let spec = state.spec.as_ref().unwrap().clone();
 
         // Record for how long we slept waiting for the spec.
         let now = Utc::now();
@@ -539,7 +444,7 @@ impl ComputeNode {
     ///
     /// Note that this is in the critical path of a compute cold start. Keep this fast.
     /// Try to do things concurrently, to hide the latencies.
-    fn start_compute(self: &Arc<Self>, pg_handle: &mut Option<PostgresHandle>) -> Result<()> {
+    async fn start_compute(self: &Arc<Self>, pg_handle: &mut Option<PostgresHandle>) -> Result<()> {
         let compute_state: ComputeState;
 
         let start_compute_span;
@@ -574,18 +479,17 @@ impl ComputeNode {
             compute_state = state_guard.clone()
         }
 
-        let pspec = compute_state.pspec.as_ref().expect("spec must be set");
+        let spec = compute_state.spec.as_ref().expect("spec must be set");
         info!(
-            "starting compute for project {}, operation {}, tenant {}, timeline {}, project {}, branch {}, endpoint {}, features {:?}, spec.remote_extensions {:?}",
-            pspec.spec.cluster.cluster_id.as_deref().unwrap_or("None"),
-            pspec.spec.operation_uuid.as_deref().unwrap_or("None"),
-            pspec.tenant_id,
-            pspec.timeline_id,
-            pspec.spec.project_id.as_deref().unwrap_or("None"),
-            pspec.spec.branch_id.as_deref().unwrap_or("None"),
-            pspec.spec.endpoint_id.as_deref().unwrap_or("None"),
-            pspec.spec.features,
-            pspec.spec.remote_extensions,
+            "starting compute for operation {}, tenant {}, timeline {}, project {}, branch {}, endpoint {}, features {:?}, spec.remote_extensions {:?}",
+            spec.operation_uuid.as_deref().unwrap_or("None"),
+            spec.tenant_id,
+            spec.timeline_id,
+            spec.project_id,
+            spec.branch_id,
+            spec.endpoint_id,
+            spec.features,
+            spec.remote_extensions,
         );
 
         ////// PRE-STARTUP PHASE: things that need to be finished before we start the Postgres process
@@ -606,8 +510,8 @@ 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());
+        if spec.remote_extensions.is_some() {
+            let (this, spec) = (self.clone(), spec.clone());
             pre_tasks.spawn(async move {
                 this.download_preload_extensions(&spec)
                     .in_current_span()
@@ -618,13 +522,11 @@ impl ComputeNode {
         // Prepare pgdata directory. This downloads the basebackup, among other things.
         {
             let (this, cs) = (self.clone(), compute_state.clone());
-            pre_tasks.spawn_blocking_child(move || this.prepare_pgdata(&cs));
+            pre_tasks.spawn(async move { this.prepare_pgdata(&cs).await });
         }
 
         // Resize swap to the desired size if the compute spec says so
-        if let (Some(size_bytes), true) =
-            (pspec.spec.swap_size_bytes, self.params.resize_swap_on_bind)
-        {
+        if let (Some(size_bytes), true) = (spec.swap_size_bytes, self.params.resize_swap_on_bind) {
             pre_tasks.spawn_blocking_child(move || {
                 // To avoid 'swapoff' hitting postgres startup, we need to run resize-swap to completion
                 // *before* starting postgres.
@@ -642,7 +544,7 @@ impl ComputeNode {
 
         // Set disk quota if the compute spec says so
         if let (Some(disk_quota_bytes), Some(disk_quota_fs_mountpoint)) = (
-            pspec.spec.disk_quota_bytes,
+            spec.disk_quota_bytes,
             self.params.set_disk_quota_for_fs.as_ref(),
         ) {
             let disk_quota_fs_mountpoint = disk_quota_fs_mountpoint.clone();
@@ -657,7 +559,7 @@ impl ComputeNode {
         }
 
         // tune pgbouncer
-        if let Some(pgbouncer_settings) = &pspec.spec.pgbouncer_settings {
+        if let Some(pgbouncer_settings) = &spec.pgbouncer_settings {
             info!("tuning pgbouncer");
 
             let pgbouncer_settings = pgbouncer_settings.clone();
@@ -675,7 +577,7 @@ impl ComputeNode {
         }
 
         // configure local_proxy
-        if let Some(local_proxy) = &pspec.spec.local_proxy_config {
+        if let Some(local_proxy) = &spec.local_proxy_config {
             info!("configuring local_proxy");
 
             // Spawn a background task to do the configuration,
@@ -693,7 +595,7 @@ impl ComputeNode {
         }
 
         // Configure and start rsyslog for compliance audit logging
-        match pspec.spec.audit_log_level {
+        match spec.audit_log_level {
             ComputeAudit::Hipaa | ComputeAudit::Extended | ComputeAudit::Full => {
                 let remote_endpoint =
                     std::env::var("AUDIT_LOGGING_ENDPOINT").unwrap_or("".to_string());
@@ -704,16 +606,10 @@ impl ComputeNode {
                 let log_directory_path = Path::new(&self.params.pgdata).join("log");
                 let log_directory_path = log_directory_path.to_string_lossy().to_string();
 
-                // Add project_id,endpoint_id to identify the logs.
-                //
-                // These ids are passed from cplane,
-                let endpoint_id = pspec.spec.endpoint_id.as_deref().unwrap_or("");
-                let project_id = pspec.spec.project_id.as_deref().unwrap_or("");
-
                 configure_audit_rsyslog(
                     log_directory_path.clone(),
-                    endpoint_id,
-                    project_id,
+                    &spec.endpoint_id,
+                    &spec.project_id,
                     &remote_endpoint,
                 )?;
 
@@ -724,7 +620,7 @@ impl ComputeNode {
         }
 
         // Configure and start rsyslog for Postgres logs export
-        let conf = PostgresLogsRsyslogConfig::new(pspec.spec.logs_export_host.as_deref());
+        let conf = PostgresLogsRsyslogConfig::new(spec.logs_export_host.as_deref());
         configure_postgres_logs_export(conf)?;
 
         // Launch remaining service threads
@@ -732,21 +628,20 @@ impl ComputeNode {
         let _configurator_handle = launch_configurator(self);
 
         // Wait for all the pre-tasks to finish before starting postgres
-        let rt = tokio::runtime::Handle::current();
-        while let Some(res) = rt.block_on(pre_tasks.join_next()) {
+        while let Some(res) = pre_tasks.join_next().await {
             res??;
         }
 
         ////// START POSTGRES
         let start_time = Utc::now();
-        let pg_process = self.start_postgres(pspec.storage_auth_token.clone())?;
+        let pg_process = self.start_postgres(spec.storage_auth_token.clone())?;
         let postmaster_pid = pg_process.pid();
         *pg_handle = Some(pg_process);
 
         // If this is a primary endpoint, perform some post-startup configuration before
         // opening it up for the world.
         let config_time = Utc::now();
-        if pspec.spec.mode == ComputeMode::Primary {
+        if spec.mode == ComputeMode::Primary {
             self.configure_as_primary(&compute_state)?;
 
             let conf = self.get_tokio_conn_conf(None);
@@ -785,8 +680,9 @@ impl ComputeNode {
         self.spawn_extension_stats_task();
 
         if pspec.spec.autoprewarm {
-            self.prewarm_lfc(None);
+            self.prewarm_lfc();
         }
+
         Ok(())
     }
 
@@ -862,14 +758,14 @@ impl ComputeNode {
         }
     }
 
-    fn cleanup_after_postgres_exit(&self) -> Result<bool> {
+    async fn cleanup_after_postgres_exit(&self) -> Result<bool> {
         // Maybe sync safekeepers again, to speed up next startup
         let compute_state = self.state.lock().unwrap().clone();
-        let pspec = compute_state.pspec.as_ref().expect("spec must be set");
-        if matches!(pspec.spec.mode, compute_api::spec::ComputeMode::Primary) {
+        let spec = compute_state.spec.as_ref().expect("spec must be set");
+        if matches!(spec.mode, compute_api::spec::ComputeMode::Primary) {
             info!("syncing safekeepers on shutdown");
-            let storage_auth_token = pspec.storage_auth_token.clone();
-            let lsn = self.sync_safekeepers(storage_auth_token)?;
+            let storage_auth_token = spec.storage_auth_token.clone();
+            let lsn = self.sync_safekeepers(storage_auth_token).await?;
             info!("synced safekeepers at lsn {lsn}");
         }
 
@@ -892,13 +788,12 @@ impl ComputeNode {
 
     /// Check that compute node has corresponding feature enabled.
     pub fn has_feature(&self, feature: ComputeFeature) -> bool {
-        let state = self.state.lock().unwrap();
-
-        if let Some(s) = state.pspec.as_ref() {
-            s.spec.features.contains(&feature)
-        } else {
-            false
-        }
+        self.state
+            .lock()
+            .unwrap()
+            .spec
+            .as_ref()
+            .is_some_and(|spec| spec.features.contains(&feature))
     }
 
     pub fn set_status(&self, status: ComputeStatus) {
@@ -915,13 +810,15 @@ impl ComputeNode {
         self.state.lock().unwrap().status
     }
 
-    pub fn get_timeline_id(&self) -> Option<TimelineId> {
+    pub fn get_timeline_id(&self) -> TimelineId {
         self.state
             .lock()
             .unwrap()
-            .pspec
+            .spec
             .as_ref()
-            .map(|s| s.timeline_id)
+            .unwrap()
+            .timeline_id
+            .clone()
     }
 
     // Remove `pgdata` directory and create it again with right permissions.
@@ -940,11 +837,10 @@ impl ComputeNode {
     // unarchive it to `pgdata` directory overriding all its previous content.
     #[instrument(skip_all, fields(%lsn))]
     fn try_get_basebackup(&self, compute_state: &ComputeState, lsn: Lsn) -> Result<()> {
-        let spec = compute_state.pspec.as_ref().expect("spec must be set");
+        let spec = compute_state.spec.as_ref().expect("spec must be set");
         let start_time = Instant::now();
 
-        let shard0_connstr = spec.pageserver_connstr.split(',').next().unwrap();
-        let mut config = postgres::Config::from_str(shard0_connstr)?;
+        let mut config = postgres::Config::from(&spec.pageservers[0]);
 
         // Use the storage auth token from the config file, if given.
         // Note: this overrides any password set in the connection string.
@@ -956,20 +852,17 @@ impl ComputeNode {
         }
 
         config.application_name("compute_ctl");
-        if let Some(spec) = &compute_state.pspec {
-            config.options(&format!(
-                "-c neon.compute_mode={}",
-                spec.spec.mode.to_type_str()
-            ));
+        if let Some(spec) = &compute_state.spec {
+            config.options(&format!("-c neon.compute_mode={}", spec.mode.to_type_str()));
         }
 
         // Connect to pageserver
-        let mut client = config.connect(NoTls)?;
+        let mut client = config.connect(postgres::NoTls)?;
         let pageserver_connect_micros = start_time.elapsed().as_micros() as u64;
 
         let basebackup_cmd = match lsn {
             Lsn(0) => {
-                if spec.spec.mode != ComputeMode::Primary {
+                if spec.mode != ComputeMode::Primary {
                     format!(
                         "basebackup {} {} --gzip --replica",
                         spec.tenant_id, spec.timeline_id
@@ -979,7 +872,7 @@ impl ComputeNode {
                 }
             }
             _ => {
-                if spec.spec.mode != ComputeMode::Primary {
+                if spec.mode != ComputeMode::Primary {
                     format!(
                         "basebackup {} {} {} --gzip --replica",
                         spec.tenant_id, spec.timeline_id, lsn
@@ -1055,35 +948,34 @@ impl ComputeNode {
         compute_state: &ComputeState,
     ) -> Result<Option<Lsn>> {
         // Construct a connection config for each safekeeper
-        let pspec: ParsedSpec = compute_state
-            .pspec
+        let spec = compute_state
+            .spec
             .as_ref()
             .expect("spec must be set")
             .clone();
-        let sk_connstrs: Vec<String> = pspec.safekeeper_connstrings.clone();
-        let sk_configs = sk_connstrs.into_iter().map(|connstr| {
-            // Format connstr
-            let id = connstr.clone();
-            let connstr = format!("postgresql://no_user@{}", connstr);
-            let options = format!(
-                "-c timeline_id={} tenant_id={}",
-                pspec.timeline_id, pspec.tenant_id
-            );
+        let safekeepers = spec
+            .safekeepers
+            .iter()
+            .map(|s| {
+                let mut config = tokio_postgres::Config::from(s);
 
-            // Construct client
-            let mut config = tokio_postgres::Config::from_str(&connstr).unwrap();
-            config.options(&options);
-            if let Some(storage_auth_token) = pspec.storage_auth_token.clone() {
-                config.password(storage_auth_token);
-            }
+                config.user("no_user");
+                config.options(&format!(
+                    "-c timeline_id={} tenant_id={}",
+                    spec.timeline_id, spec.tenant_id
+                ));
+                if let Some(storage_auth_token) = &spec.storage_auth_token {
+                    config.password(storage_auth_token);
+                }
 
-            (id, config)
-        });
+                (format!("{}:{}", s.host, s.port), config)
+            })
+            .collect::<Vec<(String, tokio_postgres::Config)>>();
 
         // Create task set to query all safekeepers
         let mut tasks = FuturesUnordered::new();
-        let quorum = sk_configs.len() / 2 + 1;
-        for (id, config) in sk_configs {
+        let quorum = safekeepers.len() / 2 + 1;
+        for (id, config) in safekeepers {
             let timeout = tokio::time::Duration::from_millis(100);
             let task = tokio::time::timeout(timeout, ping_safekeeper(id, config));
             tasks.push(tokio::spawn(task));
@@ -1128,11 +1020,13 @@ impl ComputeNode {
     // Fast path for sync_safekeepers. If they're already synced we get the lsn
     // in one roundtrip. If not, we should do a full sync_safekeepers.
     #[instrument(skip_all)]
-    pub fn check_safekeepers_synced(&self, compute_state: &ComputeState) -> Result<Option<Lsn>> {
+    pub async fn check_safekeepers_synced(
+        &self,
+        compute_state: &ComputeState,
+    ) -> Result<Option<Lsn>> {
         let start_time = Utc::now();
 
-        let rt = tokio::runtime::Handle::current();
-        let result = rt.block_on(self.check_safekeepers_synced_async(compute_state));
+        let result = self.check_safekeepers_synced_async(compute_state).await;
 
         // Record runtime
         self.state.lock().unwrap().metrics.sync_sk_check_ms = Utc::now()
@@ -1146,7 +1040,7 @@ impl ComputeNode {
     // Run `postgres` in a special mode with `--sync-safekeepers` argument
     // and return the reported LSN back to the caller.
     #[instrument(skip_all)]
-    pub fn sync_safekeepers(&self, storage_auth_token: Option<String>) -> Result<Lsn> {
+    pub async fn sync_safekeepers(&self, storage_auth_token: Option<String>) -> Result<Lsn> {
         let start_time = Utc::now();
 
         let mut sync_handle = maybe_cgexec(&self.params.pgbin)
@@ -1178,8 +1072,8 @@ impl ComputeNode {
         SYNC_SAFEKEEPERS_PID.store(0, Ordering::SeqCst);
 
         // Process has exited, so we can join the logs thread.
-        let _ = tokio::runtime::Handle::current()
-            .block_on(logs_handle)
+        let _ = logs_handle
+            .await
             .map_err(|e| tracing::error!("log task panicked: {:?}", e));
 
         if !sync_output.status.success() {
@@ -1205,9 +1099,8 @@ impl ComputeNode {
     /// Do all the preparations like PGDATA directory creation, configuration,
     /// safekeepers sync, basebackup, etc.
     #[instrument(skip_all)]
-    pub fn prepare_pgdata(&self, compute_state: &ComputeState) -> Result<()> {
-        let pspec = compute_state.pspec.as_ref().expect("spec must be set");
-        let spec = &pspec.spec;
+    pub async fn prepare_pgdata(&self, compute_state: &ComputeState) -> Result<()> {
+        let spec = compute_state.spec.as_ref().expect("spec must be set");
         let pgdata_path = Path::new(&self.params.pgdata);
 
         let tls_config = self.tls_config(&pspec.spec);
@@ -1216,7 +1109,7 @@ impl ComputeNode {
         self.create_pgdata()?;
         config::write_postgres_conf(
             pgdata_path,
-            &pspec.spec,
+            spec,
             self.params.internal_http_port,
             tls_config,
         )?;
@@ -1227,11 +1120,13 @@ impl ComputeNode {
         let lsn = match spec.mode {
             ComputeMode::Primary => {
                 info!("checking if safekeepers are synced");
-                let lsn = if let Ok(Some(lsn)) = self.check_safekeepers_synced(compute_state) {
+                let lsn = if let Ok(Some(lsn)) = self.check_safekeepers_synced(compute_state).await
+                {
                     lsn
                 } else {
                     info!("starting safekeepers syncing");
-                    self.sync_safekeepers(pspec.storage_auth_token.clone())
+                    self.sync_safekeepers(spec.storage_auth_token.clone())
+                        .await
                         .with_context(|| "failed to sync safekeepers")?
                 };
                 info!("safekeepers synced at LSN {}", lsn);
@@ -1248,13 +1143,13 @@ impl ComputeNode {
         };
 
         info!(
-            "getting basebackup@{} from pageserver {}",
-            lsn, &pspec.pageserver_connstr
+            "getting basebackup@{} from pageserver {}:{}",
+            lsn, spec.pageservers[0].host, spec.pageservers[0].port
         );
         self.get_basebackup(compute_state, lsn).with_context(|| {
             format!(
-                "failed to get basebackup@{} from pageserver {}",
-                lsn, &pspec.pageserver_connstr
+                "failed to get basebackup@{} from pageserver {}:{}",
+                lsn, spec.pageservers[0].host, spec.pageservers[0].port
             )
         })?;
 
@@ -1536,10 +1431,9 @@ impl ComputeNode {
         let conf = Arc::new(conf);
         let spec = Arc::new(
             compute_state
-                .pspec
+                .spec
                 .as_ref()
                 .expect("spec must be set")
-                .spec
                 .clone(),
         );
 
@@ -1608,7 +1502,7 @@ impl ComputeNode {
     /// as it's used to reconfigure a previously started and configured Postgres node.
     #[instrument(skip_all)]
     pub fn reconfigure(&self) -> Result<()> {
-        let spec = self.state.lock().unwrap().pspec.clone().unwrap().spec;
+        let spec = self.state.lock().unwrap().spec.as_ref().unwrap().clone();
 
         let tls_config = self.tls_config(&spec);
 
@@ -1690,10 +1584,10 @@ impl ComputeNode {
 
     #[instrument(skip_all)]
     pub fn configure_as_primary(&self, compute_state: &ComputeState) -> Result<()> {
-        let pspec = compute_state.pspec.as_ref().expect("spec must be set");
+        let spec = compute_state.spec.as_ref().expect("spec must be set");
 
-        assert!(pspec.spec.mode == ComputeMode::Primary);
-        if !pspec.spec.skip_pg_catalog_updates {
+        assert!(spec.mode == ComputeMode::Primary);
+        if !spec.skip_pg_catalog_updates {
             let pgdata_path = Path::new(&self.params.pgdata);
             // temporarily reset max_cluster_size in config
             // to avoid the possibility of hitting the limit, while we are applying config:
@@ -2189,24 +2083,23 @@ LIMIT 100",
     /// the pageserver connection strings has changed.
     ///
     /// The operation will time out after a specified duration.
-    pub fn wait_timeout_while_pageserver_connstr_unchanged(&self, duration: Duration) {
+    pub fn wait_timeout_while_pageservers_unchanged(&self, duration: Duration) {
         let state = self.state.lock().unwrap();
-        let old_pageserver_connstr = state
-            .pspec
+        let old_pageservers = state
+            .spec
             .as_ref()
             .expect("spec must be set")
-            .pageserver_connstr
+            .pageservers
             .clone();
         let mut unchanged = true;
         let _ = self
             .state_changed
             .wait_timeout_while(state, duration, |s| {
-                let pageserver_connstr = &s
-                    .pspec
-                    .as_ref()
-                    .expect("spec must be set")
-                    .pageserver_connstr;
-                unchanged = pageserver_connstr == &old_pageserver_connstr;
+                let current_pageservers = &s.spec.as_ref().expect("spec must be set").pageservers;
+                unchanged = current_pageservers
+                    .iter()
+                    .zip(&old_pageservers)
+                    .all(|(c, o)| c == o);
                 unchanged
             })
             .unwrap();

compute_tools/src/compute_prewarm.rs

@@ -3,6 +3,7 @@ use anyhow::{Context, Result, bail};
 use async_compression::tokio::bufread::{ZstdDecoder, ZstdEncoder};
 use compute_api::responses::LfcOffloadState;
 use compute_api::responses::LfcPrewarmState;
+use compute_api::spec::ComputeSpec;
 use http::StatusCode;
 use reqwest::Client;
 use std::sync::Arc;
@@ -25,29 +26,30 @@ struct EndpointStoragePair {
 }
 
 const KEY: &str = "lfc_state";
-impl EndpointStoragePair {
-    /// endpoint_id is set to None while prewarming from other endpoint, see replica promotion
-    /// If not None, takes precedence over pspec.spec.endpoint_id
-    fn from_spec_and_endpoint(
-        pspec: &crate::compute::ParsedSpec,
-        endpoint_id: Option<String>,
-    ) -> Result<Self> {
-        let endpoint_id = endpoint_id.as_ref().or(pspec.spec.endpoint_id.as_ref());
-        let Some(ref endpoint_id) = endpoint_id else {
-            bail!("pspec.endpoint_id missing, other endpoint_id not provided")
+impl TryFrom<&ComputeSpec> for EndpointStoragePair {
+    type Error = anyhow::Error;
+    fn try_from(spec: &ComputeSpec) -> Result<Self, Self::Error> {
+        let Some(ref addr) = spec.endpoint_storage_addr else {
+            bail!("spec.endpoint_storage_addr missing")
         };
-        let Some(ref base_uri) = pspec.endpoint_storage_addr else {
-            bail!("pspec.endpoint_storage_addr missing")
-        };
-        let tenant_id = pspec.tenant_id;
-        let timeline_id = pspec.timeline_id;
 
-        let url = format!("http://{base_uri}/{tenant_id}/{timeline_id}/{endpoint_id}/{KEY}");
-        let Some(ref token) = pspec.endpoint_storage_token else {
-            bail!("pspec.endpoint_storage_token missing")
+        let url = format!(
+            "http://{addr}/{tenant_id}/{timeline_id}/{endpoint_id}/{key}",
+            addr = addr,
+            tenant_id = spec.tenant_id,
+            timeline_id = spec.timeline_id,
+            endpoint_id = spec.endpoint_id,
+            key = KEY
+        );
+
+        let Some(ref token) = spec.endpoint_storage_token else {
+            bail!("spec.endpoint_storage_token missing")
         };
-        let token = token.clone();
-        Ok(EndpointStoragePair { url, token })
+
+        Ok(EndpointStoragePair {
+            url,
+            token: token.clone(),
+        })
     }
 }
 
@@ -89,7 +91,7 @@ impl ComputeNode {
     }
 
     /// Returns false if there is a prewarm request ongoing, true otherwise
-    pub fn prewarm_lfc(self: &Arc<Self>, from_endpoint: Option<String>) -> bool {
+    pub fn prewarm_lfc(self: &Arc<Self>) -> bool {
         crate::metrics::LFC_PREWARM_REQUESTS.inc();
         {
             let state = &mut self.state.lock().unwrap().lfc_prewarm_state;
@@ -102,7 +104,7 @@ impl ComputeNode {
 
         let cloned = self.clone();
         spawn(async move {
-            let Err(err) = cloned.prewarm_impl(from_endpoint).await else {
+            let Err(err) = cloned.prewarm_impl().await else {
                 cloned.state.lock().unwrap().lfc_prewarm_state = LfcPrewarmState::Completed;
                 return;
             };
@@ -114,14 +116,13 @@ impl ComputeNode {
         true
     }
 
-    /// from_endpoint: None for endpoint managed by this compute_ctl
-    fn endpoint_storage_pair(&self, from_endpoint: Option<String>) -> Result<EndpointStoragePair> {
+    fn endpoint_storage_pair(&self) -> Result<EndpointStoragePair> {
         let state = self.state.lock().unwrap();
-        EndpointStoragePair::from_spec_and_endpoint(state.pspec.as_ref().unwrap(), from_endpoint)
+        state.spec.as_ref().unwrap().try_into()
     }
 
-    async fn prewarm_impl(&self, from_endpoint: Option<String>) -> Result<()> {
-        let EndpointStoragePair { url, token } = self.endpoint_storage_pair(from_endpoint)?;
+    async fn prewarm_impl(&self) -> Result<()> {
+        let EndpointStoragePair { url, token } = self.endpoint_storage_pair()?;
         info!(%url, "requesting LFC state from endpoint storage");
 
         let request = Client::new().get(&url).bearer_auth(token);
@@ -179,7 +180,7 @@ impl ComputeNode {
     }
 
     async fn offload_lfc_impl(&self) -> Result<()> {
-        let EndpointStoragePair { url, token } = self.endpoint_storage_pair(None)?;
+        let EndpointStoragePair { url, token } = self.endpoint_storage_pair()?;
         info!(%url, "requesting LFC state from postgres");
 
         let mut compressed = Vec::new();

compute_tools/src/config.rs

@@ -56,13 +56,24 @@ pub fn write_postgres_conf(
 
     // Add options for connecting to storage
     writeln!(file, "# Neon storage settings")?;
-    if let Some(s) = &spec.pageserver_connstring {
-        writeln!(file, "neon.pageserver_connstring={}", escape_conf_value(s))?;
+    if !spec.pageservers.is_empty() {
+        writeln!(
+            file,
+            "neon.pageserver_connstring={}",
+            escape_conf_value(
+                &spec
+                    .pageservers
+                    .iter()
+                    .map(|p| format!("host={} port={}", p.host, p.port))
+                    .collect::<Vec<_>>()
+                    .join(",")
+            )
+        )?;
     }
     if let Some(stripe_size) = spec.shard_stripe_size {
         writeln!(file, "neon.stripe_size={stripe_size}")?;
     }
-    if !spec.safekeeper_connstrings.is_empty() {
+    if !spec.safekeepers.is_empty() {
         let mut neon_safekeepers_value = String::new();
         tracing::info!(
             "safekeepers_connstrings is not zero, gen: {:?}",
@@ -72,32 +83,45 @@ pub fn write_postgres_conf(
         if let Some(generation) = spec.safekeepers_generation {
             write!(neon_safekeepers_value, "g#{}:", generation)?;
         }
-        neon_safekeepers_value.push_str(&spec.safekeeper_connstrings.join(","));
+        neon_safekeepers_value.push_str(
+            &spec
+                .safekeepers
+                .iter()
+                .map(|s| format!("{}:{}", s.host.to_string(), s.port))
+                .collect::<Vec<_>>()
+                .join(","),
+        );
         writeln!(
             file,
             "neon.safekeepers={}",
             escape_conf_value(&neon_safekeepers_value)
         )?;
     }
-    if let Some(s) = &spec.tenant_id {
-        writeln!(file, "neon.tenant_id={}", escape_conf_value(&s.to_string()))?;
-    }
-    if let Some(s) = &spec.timeline_id {
-        writeln!(
-            file,
-            "neon.timeline_id={}",
-            escape_conf_value(&s.to_string())
-        )?;
-    }
-    if let Some(s) = &spec.project_id {
-        writeln!(file, "neon.project_id={}", escape_conf_value(s))?;
-    }
-    if let Some(s) = &spec.branch_id {
-        writeln!(file, "neon.branch_id={}", escape_conf_value(s))?;
-    }
-    if let Some(s) = &spec.endpoint_id {
-        writeln!(file, "neon.endpoint_id={}", escape_conf_value(s))?;
-    }
+    writeln!(
+        file,
+        "neon.tenant_id={}",
+        escape_conf_value(&spec.tenant_id.to_string())
+    )?;
+    writeln!(
+        file,
+        "neon.timeline_id={}",
+        escape_conf_value(&spec.timeline_id.to_string())
+    )?;
+    writeln!(
+        file,
+        "neon.project_id={}",
+        escape_conf_value(&spec.project_id)
+    )?;
+    writeln!(
+        file,
+        "neon.branch_id={}",
+        escape_conf_value(&spec.branch_id)
+    )?;
+    writeln!(
+        file,
+        "neon.endpoint_id={}",
+        escape_conf_value(&spec.endpoint_id)
+    )?;
 
     // tls
     if let Some(tls_config) = tls_config {

compute_tools/src/http/routes/configure.rs

@@ -8,7 +8,7 @@ use http::StatusCode;
 use tokio::task;
 use tracing::info;
 
-use crate::compute::{ComputeNode, ParsedSpec};
+use crate::compute::ComputeNode;
 use crate::http::JsonResponse;
 use crate::http::extract::Json;
 
@@ -22,11 +22,6 @@ 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) {
-        Ok(p) => p,
-        Err(e) => return JsonResponse::error(StatusCode::BAD_REQUEST, e),
-    };
-
     // XXX: wrap state update under lock in a code block. Otherwise, we will try
     // to `Send` `mut state` into the spawned thread bellow, which will cause
     // the following rustc error:
@@ -43,7 +38,7 @@ pub(in crate::http) async fn configure(
         // configure request for tracing purposes.
         state.startup_span = Some(tracing::Span::current());
 
-        state.pspec = Some(pspec);
+        state.spec = Some(request.spec.clone());
         state.set_status(ComputeStatus::ConfigurationPending, &compute.state_changed);
         drop(state);
     }

compute_tools/src/http/routes/extension_server.rs

@@ -31,8 +31,7 @@ pub(in crate::http) async fn download_extension(
 
     let ext = {
         let state = compute.state.lock().unwrap();
-        let pspec = state.pspec.as_ref().unwrap();
-        let spec = &pspec.spec;
+        let spec = &state.spec.as_ref().unwrap();
 
         let remote_extensions = match spec.remote_extensions.as_ref() {
             Some(r) => r,

compute_tools/src/http/routes/lfc.rs

@@ -2,7 +2,6 @@ use crate::compute_prewarm::LfcPrewarmStateWithProgress;
 use crate::http::JsonResponse;
 use axum::response::{IntoResponse, Response};
 use axum::{Json, http::StatusCode};
-use axum_extra::extract::OptionalQuery;
 use compute_api::responses::LfcOffloadState;
 type Compute = axum::extract::State<std::sync::Arc<crate::compute::ComputeNode>>;
 
@@ -17,16 +16,8 @@ pub(in crate::http) async fn offload_state(compute: Compute) -> Json<LfcOffloadS
     Json(compute.lfc_offload_state())
 }
 
-#[derive(serde::Deserialize)]
-pub struct PrewarmQuery {
-    pub from_endpoint: String,
-}
-
-pub(in crate::http) async fn prewarm(
-    compute: Compute,
-    OptionalQuery(query): OptionalQuery<PrewarmQuery>,
-) -> Response {
-    if compute.prewarm_lfc(query.map(|q| q.from_endpoint)) {
+pub(in crate::http) async fn prewarm(compute: Compute) -> Response {
+    if compute.prewarm_lfc() {
         StatusCode::ACCEPTED.into_response()
     } else {
         JsonResponse::error(

compute_tools/src/http/routes/mod.rs

@@ -21,14 +21,8 @@ impl From<&ComputeState> for ComputeStatusResponse {
     fn from(state: &ComputeState) -> Self {
         ComputeStatusResponse {
             start_time: state.start_time,
-            tenant: state
-                .pspec
-                .as_ref()
-                .map(|pspec| pspec.tenant_id.to_string()),
-            timeline: state
-                .pspec
-                .as_ref()
-                .map(|pspec| pspec.timeline_id.to_string()),
+            tenant: state.spec.as_ref().map(|spec| spec.tenant_id.to_string()),
+            timeline: state.spec.as_ref().map(|spec| spec.timeline_id.to_string()),
             status: state.status,
             last_active: state.last_active,
             error: state.error.clone(),

compute_tools/src/lsn_lease.rs

@@ -18,8 +18,8 @@ use crate::compute::ComputeNode;
 pub fn launch_lsn_lease_bg_task_for_static(compute: &Arc<ComputeNode>) {
     let (tenant_id, timeline_id, lsn) = {
         let state = compute.state.lock().unwrap();
-        let spec = state.pspec.as_ref().expect("Spec must be set");
-        match spec.spec.mode {
+        let spec = state.spec.as_ref().expect("Spec must be set");
+        match spec.mode {
             ComputeMode::Static(lsn) => (spec.tenant_id, spec.timeline_id, lsn),
             _ => return,
         }
@@ -58,7 +58,7 @@ fn lsn_lease_bg_task(
             "Request succeeded, sleeping for {} seconds",
             sleep_duration.as_secs()
         );
-        compute.wait_timeout_while_pageserver_connstr_unchanged(sleep_duration);
+        compute.wait_timeout_while_pageservers_unchanged(sleep_duration);
     }
 }
 
@@ -79,18 +79,11 @@ fn acquire_lsn_lease_with_retry(
         let configs = {
             let state = compute.state.lock().unwrap();
 
-            let spec = state.pspec.as_ref().expect("spec must be set");
+            let spec = state.spec.as_ref().expect("spec must be set");
 
-            let conn_strings = spec.pageserver_connstr.split(',');
-
-            conn_strings
-                .map(|connstr| {
-                    let mut config = postgres::Config::from_str(connstr).expect("Invalid connstr");
-                    if let Some(storage_auth_token) = &spec.storage_auth_token {
-                        config.password(storage_auth_token.clone());
-                    }
-                    config
-                })
+            spec.pageservers
+                .iter()
+                .map(|p| postgres::Config::from(p))
                 .collect::<Vec<_>>()
         };
 
@@ -105,7 +98,7 @@ fn acquire_lsn_lease_with_retry(
             Err(e) => {
                 warn!("Failed to acquire lsn lease: {e} (attempt {attempts})");
 
-                compute.wait_timeout_while_pageserver_connstr_unchanged(Duration::from_millis(
+                compute.wait_timeout_while_pageservers_unchanged(Duration::from_millis(
                     retry_period_ms as u64,
                 ));
                 retry_period_ms *= 1.5;

compute_tools/src/spec_apply.rs

@@ -4,7 +4,7 @@ use std::future::Future;
 use std::iter::{empty, once};
 use std::sync::Arc;
 
-use anyhow::{Context, Result};
+use anyhow::Result;
 use compute_api::responses::ComputeStatus;
 use compute_api::spec::{ComputeAudit, ComputeSpec, Database, PgIdent, Role};
 use futures::future::join_all;
@@ -74,7 +74,7 @@ impl ComputeNode {
             let mut drop_subscriptions_done = false;
 
             if spec.drop_subscriptions_before_start {
-                let timeline_id = self.get_timeline_id().context("timeline_id must be set")?;
+                let timeline_id = self.get_timeline_id();
 
                 info!("Checking if drop subscription operation was already performed for timeline_id: {}", timeline_id);
 

compute_tools/src/sync_sk.rs

@@ -37,7 +37,7 @@ pub async fn ping_safekeeper(
 
     // Parse result
     info!("done with {}", id);
-    if let postgres::SimpleQueryMessage::Row(row) = &result[0] {
+    if let tokio_postgres::SimpleQueryMessage::Row(row) = &result[0] {
         use std::str::FromStr;
         let response = TimelineStatusResponse::Ok(TimelineStatusOkResponse {
             flush_lsn: Lsn::from_str(row.get("flush_lsn").unwrap())?,

control_plane/src/bin/neon_local.rs

@@ -1493,7 +1493,10 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
                 let conf = env.get_pageserver_conf(pageserver_id).unwrap();
                 let parsed = parse_host_port(&conf.listen_pg_addr).expect("Bad config");
                 (
-                    vec![(parsed.0, parsed.1.unwrap_or(5432))],
+                    vec![compute_api::spec::Pageserver {
+                        host: parsed.0,
+                        port: parsed.1.unwrap_or(5432),
+                    }],
                     // If caller is telling us what pageserver to use, this is not a tenant which is
                     // full managed by storage controller, therefore not sharded.
                     DEFAULT_STRIPE_SIZE,
@@ -1516,11 +1519,11 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
                                 .await?;
                         }
 
-                        anyhow::Ok((
-                            Host::parse(&shard.listen_pg_addr)
+                        anyhow::Ok(compute_api::spec::Pageserver {
+                            host: Host::parse(&shard.listen_pg_addr)
                                 .expect("Storage controller reported bad hostname"),
-                            shard.listen_pg_port,
-                        ))
+                            port: shard.listen_pg_port,
+                        })
                     }),
                 )
                 .await?;
@@ -1576,10 +1579,10 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
                 .with_context(|| format!("postgres endpoint {endpoint_id} is not found"))?;
             let pageservers = if let Some(ps_id) = args.endpoint_pageserver_id {
                 let pageserver = PageServerNode::from_env(env, env.get_pageserver_conf(ps_id)?);
-                vec![(
-                    pageserver.pg_connection_config.host().clone(),
-                    pageserver.pg_connection_config.port(),
-                )]
+                vec![compute_api::spec::Pageserver {
+                    host: pageserver.pg_connection_config.host().clone(),
+                    port: pageserver.pg_connection_config.port(),
+                }]
             } else {
                 let storage_controller = StorageController::from_env(env);
                 storage_controller
@@ -1587,12 +1590,10 @@ async fn handle_endpoint(subcmd: &EndpointCmd, env: &local_env::LocalEnv) -> Res
                     .await?
                     .shards
                     .into_iter()
-                    .map(|shard| {
-                        (
-                            Host::parse(&shard.listen_pg_addr)
-                                .expect("Storage controller reported malformed host"),
-                            shard.listen_pg_port,
-                        )
+                    .map(|shard| compute_api::spec::Pageserver {
+                        host: Host::parse(&shard.listen_pg_addr)
+                            .expect("Storage controller reported malformed host"),
+                        port: shard.listen_pg_port,
                     })
                     .collect::<Vec<_>>()
             };

control_plane/src/endpoint.rs

@@ -52,8 +52,8 @@ use compute_api::responses::{
     ComputeConfig, ComputeCtlConfig, ComputeStatus, ComputeStatusResponse, TlsConfig,
 };
 use compute_api::spec::{
-    Cluster, ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, Database, PgIdent,
-    RemoteExtSpec, Role,
+    Cluster, ComputeAudit, ComputeFeature, ComputeMode, ComputeSpec, Database, Pageserver, PgIdent,
+    RemoteExtSpec, Role, Safekeeper,
 };
 use jsonwebtoken::jwk::{
     AlgorithmParameters, CommonParameters, EllipticCurve, Jwk, JwkSet, KeyAlgorithm, KeyOperations,
@@ -606,29 +606,25 @@ impl Endpoint {
         }
     }
 
-    fn build_pageserver_connstr(pageservers: &[(Host, u16)]) -> String {
-        pageservers
-            .iter()
-            .map(|(host, port)| format!("postgresql://no_user@{host}:{port}"))
-            .collect::<Vec<_>>()
-            .join(",")
-    }
+    fn safekeepers_from_nodes(&self, ids: Vec<NodeId>) -> Result<Vec<Safekeeper>> {
+        let mut s = Vec::new();
 
-    /// Map safekeepers ids to the actual connection strings.
-    fn build_safekeepers_connstrs(&self, sk_ids: Vec<NodeId>) -> Result<Vec<String>> {
-        let mut safekeeper_connstrings = Vec::new();
         if self.mode == ComputeMode::Primary {
-            for sk_id in sk_ids {
+            for id in ids {
                 let sk = self
                     .env
                     .safekeepers
                     .iter()
-                    .find(|node| node.id == sk_id)
-                    .ok_or_else(|| anyhow!("safekeeper {sk_id} does not exist"))?;
-                safekeeper_connstrings.push(format!("127.0.0.1:{}", sk.get_compute_port()));
+                    .find(|node| node.id == id)
+                    .ok_or_else(|| anyhow!("safekeeper {id} does not exist"))?;
+                s.push(Safekeeper {
+                    host: Host::parse("127.0.0.1")?,
+                    port: sk.get_compute_port(),
+                });
             }
         }
-        Ok(safekeeper_connstrings)
+
+        Ok(s)
     }
 
     /// Generate a JWT with the correct claims.
@@ -654,7 +650,7 @@ impl Endpoint {
         endpoint_storage_addr: String,
         safekeepers_generation: Option<SafekeeperGeneration>,
         safekeepers: Vec<NodeId>,
-        pageservers: Vec<(Host, u16)>,
+        pageservers: Vec<Pageserver>,
         remote_ext_base_url: Option<&String>,
         shard_stripe_size: usize,
         create_test_user: bool,
@@ -672,11 +668,6 @@ impl Endpoint {
             std::fs::remove_dir_all(self.pgdata())?;
         }
 
-        let pageserver_connstring = Self::build_pageserver_connstr(&pageservers);
-        assert!(!pageserver_connstring.is_empty());
-
-        let safekeeper_connstrings = self.build_safekeepers_connstrs(safekeepers)?;
-
         // check for file remote_extensions_spec.json
         // if it is present, read it and pass to compute_ctl
         let remote_extensions_spec_path = self.endpoint_path().join("remote_extensions_spec.json");
@@ -727,15 +718,34 @@ impl Endpoint {
                     postgresql_conf: Some(postgresql_conf.clone()),
                 },
                 delta_operations: None,
-                tenant_id: Some(self.tenant_id),
-                timeline_id: Some(self.timeline_id),
-                project_id: None,
-                branch_id: None,
-                endpoint_id: Some(self.endpoint_id.clone()),
+                tenant_id: self.tenant_id.clone(),
+                timeline_id: self.timeline_id.clone(),
+                project_id: self.tenant_id.to_string(),
+                branch_id: self.timeline_id.to_string(),
+                endpoint_id: self.endpoint_id.clone(),
                 mode: self.mode,
-                pageserver_connstring: Some(pageserver_connstring),
+                pageservers,
+                safekeepers: {
+                    let mut s = Vec::new();
+
+                    if self.mode == ComputeMode::Primary {
+                        for id in safekeepers {
+                            let sk = self
+                                .env
+                                .safekeepers
+                                .iter()
+                                .find(|node| node.id == id)
+                                .ok_or_else(|| anyhow!("safekeeper {id} does not exist"))?;
+                            s.push(Safekeeper {
+                                host: Host::parse("127.0.0.1")?,
+                                port: sk.get_compute_port(),
+                            });
+                        }
+                    }
+
+                    s
+                },
                 safekeepers_generation: safekeepers_generation.map(|g| g.into_inner()),
-                safekeeper_connstrings,
                 storage_auth_token: auth_token.clone(),
                 remote_extensions,
                 pgbouncer_settings: None,
@@ -939,7 +949,7 @@ impl Endpoint {
 
     pub async fn reconfigure(
         &self,
-        mut pageservers: Vec<(Host, u16)>,
+        pageservers: Vec<Pageserver>,
         stripe_size: Option<ShardStripeSize>,
         safekeepers: Option<Vec<NodeId>>,
     ) -> Result<()> {
@@ -958,30 +968,24 @@ impl Endpoint {
         if pageservers.is_empty() {
             let storage_controller = StorageController::from_env(&self.env);
             let locate_result = storage_controller.tenant_locate(self.tenant_id).await?;
-            pageservers = locate_result
+            spec.pageservers = locate_result
                 .shards
                 .into_iter()
-                .map(|shard| {
-                    (
-                        Host::parse(&shard.listen_pg_addr)
-                            .expect("Storage controller reported bad hostname"),
-                        shard.listen_pg_port,
-                    )
+                .map(|shard| Pageserver {
+                    host: Host::parse(&shard.listen_pg_addr)
+                        .expect("Storage controller reported bad hostname"),
+                    port: shard.listen_pg_port,
                 })
                 .collect::<Vec<_>>();
         }
 
-        let pageserver_connstr = Self::build_pageserver_connstr(&pageservers);
-        assert!(!pageserver_connstr.is_empty());
-        spec.pageserver_connstring = Some(pageserver_connstr);
         if stripe_size.is_some() {
             spec.shard_stripe_size = stripe_size.map(|s| s.0 as usize);
         }
 
         // If safekeepers are not specified, don't change them.
         if let Some(safekeepers) = safekeepers {
-            let safekeeper_connstrings = self.build_safekeepers_connstrs(safekeepers)?;
-            spec.safekeeper_connstrings = safekeeper_connstrings;
+            spec.safekeepers = self.safekeepers_from_nodes(safekeepers)?;
         }
 
         let client = reqwest::Client::builder()

libs/compute_api/Cargo.toml

@@ -9,8 +9,11 @@ anyhow.workspace = true
 chrono.workspace = true
 indexmap.workspace = true
 jsonwebtoken.workspace = true
+postgres.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+tokio-postgres.workspace = true
+url.workspace = true
 regex.workspace = true
 
 utils = { path = "../utils" }

libs/compute_api/src/spec.rs

@@ -9,7 +9,8 @@ use indexmap::IndexMap;
 use regex::Regex;
 use remote_storage::RemotePath;
 use serde::{Deserialize, Serialize};
-use utils::id::{TenantId, TimelineId};
+use url::Host;
+use utils::id::{BranchId, EndpointId, ProjectId, TenantId, TimelineId};
 use utils::lsn::Lsn;
 
 use crate::responses::TlsConfig;
@@ -21,13 +22,77 @@ pub type PgIdent = String;
 /// String type alias representing Postgres extension version
 pub type ExtVersion = String;
 
+/// Pageserver settings.
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+pub struct Pageserver {
+    /// Hostname of the pageserver.
+    pub host: Host,
+
+    /// Port that the safekeeper listens on.
+    pub port: u16,
+}
+
+impl From<&Pageserver> for postgres::Config {
+    fn from(ps: &Pageserver) -> Self {
+        let mut config = postgres::Config::new();
+
+        config.host(&ps.host.to_string());
+        config.port(ps.port);
+
+        config
+    }
+}
+
+impl From<&Pageserver> for tokio_postgres::Config {
+    fn from(ps: &Pageserver) -> Self {
+        let mut config = tokio_postgres::Config::new();
+
+        config.host(&ps.host.to_string());
+        config.port(ps.port);
+
+        config
+    }
+}
+
+/// Safekeeper settings.
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+pub struct Safekeeper {
+    /// Hostname of the safekeeper.
+    pub host: Host,
+
+    /// Port that the safekeeper listens on.
+    pub port: u16,
+}
+
+impl From<&Safekeeper> for postgres::Config {
+    fn from(sk: &Safekeeper) -> Self {
+        let mut config = postgres::Config::new();
+
+        config.host(&sk.host.to_string());
+        config.port(sk.port);
+
+        config
+    }
+}
+
+impl From<&Safekeeper> for tokio_postgres::Config {
+    fn from(sk: &Safekeeper) -> Self {
+        let mut config = tokio_postgres::Config::new();
+
+        config.host(&sk.host.to_string());
+        config.port(sk.port);
+
+        config
+    }
+}
+
 fn default_reconfigure_concurrency() -> usize {
     1
 }
 
 /// Cluster spec or configuration represented as an optional number of
 /// delta operations + final cluster state description.
-#[derive(Clone, Debug, Default, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct ComputeSpec {
     pub format_version: f32,
 
@@ -90,25 +155,13 @@ pub struct ComputeSpec {
 
     // Information needed to connect to the storage layer.
     //
-    // `tenant_id`, `timeline_id` and `pageserver_connstring` are always needed.
-    //
     // Depending on `mode`, this can be a primary read-write node, a read-only
     // replica, or a read-only node pinned at an older LSN.
     // `safekeeper_connstrings` must be set for a primary.
-    //
-    // For backwards compatibility, the control plane may leave out all of
-    // these, and instead set the "neon.tenant_id", "neon.timeline_id",
-    // etc. GUCs in cluster.settings. TODO: Once the control plane has been
-    // updated to fill these fields, we can make these non optional.
-    pub tenant_id: Option<TenantId>,
-    pub timeline_id: Option<TimelineId>,
-    pub pageserver_connstring: Option<String>,
+    pub pageservers: Vec<Pageserver>,
 
-    // More neon ids that we expose to the compute_ctl
-    // and to postgres as neon extension GUCs.
-    pub project_id: Option<String>,
-    pub branch_id: Option<String>,
-    pub endpoint_id: Option<String>,
+    #[serde(default)]
+    pub safekeepers_generation: Option<u32>,
 
     /// Safekeeper membership config generation. It is put in
     /// neon.safekeepers GUC and serves two purposes:
@@ -120,9 +173,18 @@ pub struct ComputeSpec {
     /// Note: it could be SafekeeperGeneration, but this needs linking
     /// compute_ctl with postgres_ffi.
     #[serde(default)]
-    pub safekeepers_generation: Option<u32>,
-    #[serde(default)]
-    pub safekeeper_connstrings: Vec<String>,
+    pub safekeepers: Vec<Safekeeper>,
+
+    /// The Neon tenant ID. Exposed to Postgres as `neon.tenant_id`.
+    pub tenant_id: TenantId,
+    /// The Neon timeline ID. Exposed to Postgres as `neon.timeline_id`.
+    pub timeline_id: TimelineId,
+    /// The Neon project ID. Exposed to Postgres as `neon.project_id`.
+    pub project_id: ProjectId,
+    /// The Neon branch ID. Exposed to Postgres as `neon.branch_id`.
+    pub branch_id: BranchId,
+    /// The Neon endpoint ID. Exposed to Postgres as `neon.endpoint_id`.
+    pub endpoint_id: EndpointId,
 
     #[serde(default)]
     pub mode: ComputeMode,

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

@@ -1,3 +1,5 @@
+use std::io;
+
 use tokio::net::TcpStream;
 
 use crate::client::SocketConfig;
@@ -6,7 +8,7 @@ use crate::tls::MakeTlsConnect;
 use crate::{Error, cancel_query_raw, connect_socket};
 
 pub(crate) async fn cancel_query<T>(
-    config: SocketConfig,
+    config: Option<SocketConfig>,
     ssl_mode: SslMode,
     tls: T,
     process_id: i32,
@@ -15,6 +17,16 @@ pub(crate) async fn cancel_query<T>(
 where
     T: MakeTlsConnect<TcpStream>,
 {
+    let config = match config {
+        Some(config) => config,
+        None => {
+            return Err(Error::connect(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                "unknown host",
+            )));
+        }
+    };
+
     let hostname = match &config.host {
         Host::Tcp(host) => &**host,
     };

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

@@ -9,16 +9,9 @@ use crate::{Error, cancel_query, cancel_query_raw};
 
 /// The capability to request cancellation of in-progress queries on a
 /// connection.
-#[derive(Clone)]
+#[derive(Clone, Serialize, Deserialize)]
 pub struct CancelToken {
-    pub socket_config: SocketConfig,
-    pub raw: RawCancelToken,
-}
-
-/// The capability to request cancellation of in-progress queries on a
-/// connection.
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct RawCancelToken {
+    pub socket_config: Option<SocketConfig>,
     pub ssl_mode: SslMode,
     pub process_id: i32,
     pub secret_key: i32,
@@ -43,16 +36,14 @@ impl CancelToken {
     {
         cancel_query::cancel_query(
             self.socket_config.clone(),
-            self.raw.ssl_mode,
+            self.ssl_mode,
             tls,
-            self.raw.process_id,
-            self.raw.secret_key,
+            self.process_id,
+            self.secret_key,
         )
         .await
     }
-}
 
-impl RawCancelToken {
     /// Like `cancel_query`, but uses a stream which is already connected to the server rather than opening a new
     /// connection itself.
     pub async fn cancel_query_raw<S, T>(&self, stream: S, tls: T) -> Result<(), Error>

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

@@ -12,7 +12,6 @@ use postgres_protocol2::message::frontend;
 use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
 
-use crate::cancel_token::RawCancelToken;
 use crate::codec::{BackendMessages, FrontendMessage};
 use crate::config::{Host, SslMode};
 use crate::query::RowStream;
@@ -332,12 +331,10 @@ impl Client {
     /// connection associated with this client.
     pub fn cancel_token(&self) -> CancelToken {
         CancelToken {
-            socket_config: self.socket_config.clone(),
-            raw: RawCancelToken {
-                ssl_mode: self.ssl_mode,
-                process_id: self.process_id,
-                secret_key: self.secret_key,
-            },
+            socket_config: Some(self.socket_config.clone()),
+            ssl_mode: self.ssl_mode,
+            process_id: self.process_id,
+            secret_key: self.secret_key,
         }
     }
 

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

@@ -3,7 +3,7 @@
 
 use postgres_protocol2::message::backend::ReadyForQueryBody;
 
-pub use crate::cancel_token::{CancelToken, RawCancelToken};
+pub use crate::cancel_token::CancelToken;
 pub use crate::client::{Client, SocketConfig};
 pub use crate::config::Config;
 pub use crate::connect_raw::RawConnection;

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

@@ -14,7 +14,7 @@ pub(crate) mod private {
 
 /// Channel binding information returned from a TLS handshake.
 pub struct ChannelBinding {
-    pub tls_server_end_point: Option<Vec<u8>>,
+    pub(crate) tls_server_end_point: Option<Vec<u8>>,
 }
 
 impl ChannelBinding {

libs/remote_storage/src/azure_blob.rs

@@ -10,7 +10,7 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 use std::{env, io};
 
-use anyhow::{Context, Result, anyhow};
+use anyhow::{Context, Result};
 use azure_core::request_options::{IfMatchCondition, MaxResults, Metadata, Range};
 use azure_core::{Continuable, HttpClient, RetryOptions, TransportOptions};
 use azure_storage::StorageCredentials;
@@ -37,7 +37,6 @@ use crate::metrics::{AttemptOutcome, RequestKind, start_measuring_requests};
 use crate::{
     ConcurrencyLimiter, Download, DownloadError, DownloadKind, DownloadOpts, Listing, ListingMode,
     ListingObject, RemotePath, RemoteStorage, StorageMetadata, TimeTravelError, TimeoutOrCancel,
-    Version, VersionKind,
 };
 
 pub struct AzureBlobStorage {
@@ -406,39 +405,6 @@ impl AzureBlobStorage {
     pub fn container_name(&self) -> &str {
         &self.container_name
     }
-
-    async fn list_versions_with_permit(
-        &self,
-        _permit: &tokio::sync::SemaphorePermit<'_>,
-        prefix: Option<&RemotePath>,
-        mode: ListingMode,
-        max_keys: Option<NonZeroU32>,
-        cancel: &CancellationToken,
-    ) -> Result<crate::VersionListing, DownloadError> {
-        let customize_builder = |mut builder: ListBlobsBuilder| {
-            builder = builder.include_versions(true);
-            // We do not return this info back to `VersionListing` yet.
-            builder = builder.include_deleted(true);
-            builder
-        };
-        let kind = RequestKind::ListVersions;
-
-        let mut stream = std::pin::pin!(self.list_streaming_for_fn(
-            prefix,
-            mode,
-            max_keys,
-            cancel,
-            kind,
-            customize_builder
-        ));
-        let mut combined: crate::VersionListing =
-            stream.next().await.expect("At least one item required")?;
-        while let Some(list) = stream.next().await {
-            let list = list?;
-            combined.versions.extend(list.versions.into_iter());
-        }
-        Ok(combined)
-    }
 }
 
 trait ListingCollector {
@@ -522,10 +488,27 @@ impl RemoteStorage for AzureBlobStorage {
         max_keys: Option<NonZeroU32>,
         cancel: &CancellationToken,
     ) -> std::result::Result<crate::VersionListing, DownloadError> {
+        let customize_builder = |mut builder: ListBlobsBuilder| {
+            builder = builder.include_versions(true);
+            builder
+        };
         let kind = RequestKind::ListVersions;
-        let permit = self.permit(kind, cancel).await?;
-        self.list_versions_with_permit(&permit, prefix, mode, max_keys, cancel)
-            .await
+
+        let mut stream = std::pin::pin!(self.list_streaming_for_fn(
+            prefix,
+            mode,
+            max_keys,
+            cancel,
+            kind,
+            customize_builder
+        ));
+        let mut combined: crate::VersionListing =
+            stream.next().await.expect("At least one item required")?;
+        while let Some(list) = stream.next().await {
+            let list = list?;
+            combined.versions.extend(list.versions.into_iter());
+        }
+        Ok(combined)
     }
 
     async fn head_object(
@@ -820,158 +803,14 @@ impl RemoteStorage for AzureBlobStorage {
 
     async fn time_travel_recover(
         &self,
-        prefix: Option<&RemotePath>,
-        timestamp: SystemTime,
-        done_if_after: SystemTime,
-        cancel: &CancellationToken,
+        _prefix: Option<&RemotePath>,
+        _timestamp: SystemTime,
+        _done_if_after: SystemTime,
+        _cancel: &CancellationToken,
     ) -> Result<(), TimeTravelError> {
-        let msg = "PLEASE NOTE: Azure Blob storage time-travel recovery may not work as expected "
-            .to_string()
-            + "for some specific files. If a file gets deleted but then overwritten and we want to recover "
-            + "to the time during the file was not present, this functionality will recover the file. Only "
-            + "use the functionality for services that can tolerate this. For example, recovering a state of the "
-            + "pageserver tenants.";
-        tracing::error!("{}", msg);
-
-        let kind = RequestKind::TimeTravel;
-        let permit = self.permit(kind, cancel).await?;
-
-        let mode = ListingMode::NoDelimiter;
-        let version_listing = self
-            .list_versions_with_permit(&permit, prefix, mode, None, cancel)
-            .await
-            .map_err(|err| match err {
-                DownloadError::Other(e) => TimeTravelError::Other(e),
-                DownloadError::Cancelled => TimeTravelError::Cancelled,
-                other => TimeTravelError::Other(other.into()),
-            })?;
-        let versions_and_deletes = version_listing.versions;
-
-        tracing::info!(
-            "Built list for time travel with {} versions and deletions",
-            versions_and_deletes.len()
-        );
-
-        // Work on the list of references instead of the objects directly,
-        // otherwise we get lifetime errors in the sort_by_key call below.
-        let mut versions_and_deletes = versions_and_deletes.iter().collect::<Vec<_>>();
-
-        versions_and_deletes.sort_by_key(|vd| (&vd.key, &vd.last_modified));
-
-        let mut vds_for_key = HashMap::<_, Vec<_>>::new();
-
-        for vd in &versions_and_deletes {
-            let Version { key, .. } = &vd;
-            let version_id = vd.version_id().map(|v| v.0.as_str());
-            if version_id == Some("null") {
-                return Err(TimeTravelError::Other(anyhow!(
-                    "Received ListVersions response for key={key} with version_id='null', \
-                        indicating either disabled versioning, or legacy objects with null version id values"
-                )));
-            }
-            tracing::trace!("Parsing version key={key} kind={:?}", vd.kind);
-
-            vds_for_key.entry(key).or_default().push(vd);
-        }
-
-        let warn_threshold = 3;
-        let max_retries = 10;
-        let is_permanent = |e: &_| matches!(e, TimeTravelError::Cancelled);
-
-        for (key, versions) in vds_for_key {
-            let last_vd = versions.last().unwrap();
-            let key = self.relative_path_to_name(key);
-            if last_vd.last_modified > done_if_after {
-                tracing::debug!("Key {key} has version later than done_if_after, skipping");
-                continue;
-            }
-            // the version we want to restore to.
-            let version_to_restore_to =
-                match versions.binary_search_by_key(&timestamp, |tpl| tpl.last_modified) {
-                    Ok(v) => v,
-                    Err(e) => e,
-                };
-            if version_to_restore_to == versions.len() {
-                tracing::debug!("Key {key} has no changes since timestamp, skipping");
-                continue;
-            }
-            let mut do_delete = false;
-            if version_to_restore_to == 0 {
-                // All versions more recent, so the key didn't exist at the specified time point.
-                tracing::debug!(
-                    "All {} versions more recent for {key}, deleting",
-                    versions.len()
-                );
-                do_delete = true;
-            } else {
-                match &versions[version_to_restore_to - 1] {
-                    Version {
-                        kind: VersionKind::Version(version_id),
-                        ..
-                    } => {
-                        let source_url = format!(
-                            "{}/{}?versionid={}",
-                            self.client
-                                .url()
-                                .map_err(|e| TimeTravelError::Other(anyhow!("{e}")))?,
-                            key,
-                            version_id.0
-                        );
-                        tracing::debug!(
-                            "Promoting old version {} for {key} at {}...",
-                            version_id.0,
-                            source_url
-                        );
-                        backoff::retry(
-                            || async {
-                                let blob_client = self.client.blob_client(key.clone());
-                                let op = blob_client.copy(Url::from_str(&source_url).unwrap());
-                                tokio::select! {
-                                    res = op => res.map_err(|e| TimeTravelError::Other(e.into())),
-                                    _ = cancel.cancelled() => Err(TimeTravelError::Cancelled),
-                                }
-                            },
-                            is_permanent,
-                            warn_threshold,
-                            max_retries,
-                            "copying object version for time_travel_recover",
-                            cancel,
-                        )
-                        .await
-                        .ok_or_else(|| TimeTravelError::Cancelled)
-                        .and_then(|x| x)?;
-                        tracing::info!(?version_id, %key, "Copied old version in Azure blob storage");
-                    }
-                    Version {
-                        kind: VersionKind::DeletionMarker,
-                        ..
-                    } => {
-                        do_delete = true;
-                    }
-                }
-            };
-            if do_delete {
-                if matches!(last_vd.kind, VersionKind::DeletionMarker) {
-                    // Key has since been deleted (but there was some history), no need to do anything
-                    tracing::debug!("Key {key} already deleted, skipping.");
-                } else {
-                    tracing::debug!("Deleting {key}...");
-
-                    self.delete(&RemotePath::from_string(&key).unwrap(), cancel)
-                        .await
-                        .map_err(|e| {
-                            // delete_oid0 will use TimeoutOrCancel
-                            if TimeoutOrCancel::caused_by_cancel(&e) {
-                                TimeTravelError::Cancelled
-                            } else {
-                                TimeTravelError::Other(e)
-                            }
-                        })?;
-                }
-            }
-        }
-
-        Ok(())
+        // TODO use Azure point in time recovery feature for this
+        // https://learn.microsoft.com/en-us/azure/storage/blobs/point-in-time-restore-overview
+        Err(TimeTravelError::Unimplemented)
     }
 }
 

libs/remote_storage/src/s3_bucket.rs

@@ -1022,7 +1022,6 @@ impl RemoteStorage for S3Bucket {
             let Version { key, .. } = &vd;
             let version_id = vd.version_id().map(|v| v.0.as_str());
             if version_id == Some("null") {
-                // TODO: check the behavior of using the SDK on a non-versioned container
                 return Err(TimeTravelError::Other(anyhow!(
                     "Received ListVersions response for key={key} with version_id='null', \
                     indicating either disabled versioning, or legacy objects with null version id values"

libs/utils/src/id.rs

@@ -295,7 +295,11 @@ pub struct TenantId(Id);
 
 id_newtype!(TenantId);
 
-/// If needed, reuse small string from proxy/src/types.rc
+/// Type representing a project ID.
+pub type ProjectId = String;
+/// Type representing a branch ID.
+pub type BranchId = String;
+/// Type representing an endpoint ID.
 pub type EndpointId = String;
 
 // A pair uniquely identifying Neon instance.

pageserver/src/bin/pageserver.rs

@@ -573,8 +573,7 @@ fn start_pageserver(
         tokio::sync::mpsc::unbounded_channel();
     let deletion_queue_client = deletion_queue.new_client();
     let background_purges = mgr::BackgroundPurges::default();
-
-    let tenant_manager = mgr::init(
+    let tenant_manager = BACKGROUND_RUNTIME.block_on(mgr::init_tenant_mgr(
         conf,
         background_purges.clone(),
         TenantSharedResources {
@@ -585,10 +584,10 @@ fn start_pageserver(
             basebackup_prepare_sender,
             feature_resolver,
         },
+        order,
         shutdown_pageserver.clone(),
-    );
+    ))?;
     let tenant_manager = Arc::new(tenant_manager);
-    BACKGROUND_RUNTIME.block_on(mgr::init_tenant_mgr(tenant_manager.clone(), order))?;
 
     let basebackup_cache = BasebackupCache::spawn(
         BACKGROUND_RUNTIME.handle(),

pageserver/src/feature_resolver.rs

@@ -1,6 +1,5 @@
 use std::{collections::HashMap, sync::Arc, time::Duration};
 
-use pageserver_api::config::NodeMetadata;
 use posthog_client_lite::{
     CaptureEvent, FeatureResolverBackgroundLoop, PostHogClientConfig, PostHogEvaluationError,
     PostHogFlagFilterPropertyValue,
@@ -87,35 +86,7 @@ impl FeatureResolver {
                         }
                     }
                 }
-                // TODO: move this to a background task so that we don't block startup in case of slow disk
-                let metadata_path = conf.metadata_path();
-                match std::fs::read_to_string(&metadata_path) {
-                    Ok(metadata_str) => match serde_json::from_str::<NodeMetadata>(&metadata_str) {
-                        Ok(metadata) => {
-                            properties.insert(
-                                "hostname".to_string(),
-                                PostHogFlagFilterPropertyValue::String(metadata.http_host),
-                            );
-                            if let Some(cplane_region) = metadata.other.get("region_id") {
-                                if let Some(cplane_region) = cplane_region.as_str() {
-                                    // This region contains the cell number
-                                    properties.insert(
-                                        "neon_region".to_string(),
-                                        PostHogFlagFilterPropertyValue::String(
-                                            cplane_region.to_string(),
-                                        ),
-                                    );
-                                }
-                            }
-                        }
-                        Err(e) => {
-                            tracing::warn!("Failed to parse metadata.json: {}", e);
-                        }
-                    },
-                    Err(e) => {
-                        tracing::warn!("Failed to read metadata.json: {}", e);
-                    }
-                }
+                // TODO: add pageserver URL.
                 Arc::new(properties)
             };
             let fake_tenants = {

pageserver/src/tenant/mgr.rs

@@ -12,6 +12,7 @@ use anyhow::Context;
 use camino::{Utf8DirEntry, Utf8Path, Utf8PathBuf};
 use futures::StreamExt;
 use itertools::Itertools;
+use once_cell::sync::Lazy;
 use pageserver_api::key::Key;
 use pageserver_api::models::{DetachBehavior, LocationConfigMode};
 use pageserver_api::shard::{
@@ -102,7 +103,7 @@ pub(crate) enum TenantsMap {
     /// [`init_tenant_mgr`] is not done yet.
     Initializing,
     /// [`init_tenant_mgr`] is done, all on-disk tenants have been loaded.
-    /// New tenants can be added using [`TenantManager::tenant_map_acquire_slot`].
+    /// New tenants can be added using [`tenant_map_acquire_slot`].
     Open(BTreeMap<TenantShardId, TenantSlot>),
     /// The pageserver has entered shutdown mode via [`TenantManager::shutdown`].
     /// Existing tenants are still accessible, but no new tenants can be created.
@@ -283,6 +284,9 @@ impl BackgroundPurges {
     }
 }
 
+static TENANTS: Lazy<std::sync::RwLock<TenantsMap>> =
+    Lazy::new(|| std::sync::RwLock::new(TenantsMap::Initializing));
+
 /// Responsible for storing and mutating the collection of all tenants
 /// that this pageserver has state for.
 ///
@@ -293,7 +297,10 @@ impl BackgroundPurges {
 /// and attached modes concurrently.
 pub struct TenantManager {
     conf: &'static PageServerConf,
-    tenants: std::sync::RwLock<TenantsMap>,
+    // TODO: currently this is a &'static pointing to TENANTs.  When we finish refactoring
+    // out of that static variable, the TenantManager can own this.
+    // See https://github.com/neondatabase/neon/issues/5796
+    tenants: &'static std::sync::RwLock<TenantsMap>,
     resources: TenantSharedResources,
 
     // Long-running operations that happen outside of a [`Tenant`] lifetime should respect this token.
@@ -472,43 +479,21 @@ pub(crate) enum DeleteTenantError {
     Other(#[from] anyhow::Error),
 }
 
-/// Initialize repositories at `Initializing` state.
-pub fn init(
-    conf: &'static PageServerConf,
-    background_purges: BackgroundPurges,
-    resources: TenantSharedResources,
-    cancel: CancellationToken,
-) -> TenantManager {
-    TenantManager {
-        conf,
-        tenants: std::sync::RwLock::new(TenantsMap::Initializing),
-        resources,
-        cancel,
-        background_purges,
-    }
-}
-
-/// Transition repositories from `Initializing` state to `Open` state with locally available timelines.
+/// Initialize repositories with locally available timelines.
 /// Timelines that are only partially available locally (remote storage has more data than this pageserver)
 /// are scheduled for download and added to the tenant once download is completed.
 #[instrument(skip_all)]
 pub async fn init_tenant_mgr(
-    tenant_manager: Arc<TenantManager>,
+    conf: &'static PageServerConf,
+    background_purges: BackgroundPurges,
+    resources: TenantSharedResources,
     init_order: InitializationOrder,
-) -> anyhow::Result<()> {
-    debug_assert!(matches!(
-        *tenant_manager.tenants.read().unwrap(),
-        TenantsMap::Initializing
-    ));
+    cancel: CancellationToken,
+) -> anyhow::Result<TenantManager> {
     let mut tenants = BTreeMap::new();
 
     let ctx = RequestContext::todo_child(TaskKind::Startup, DownloadBehavior::Warn);
 
-    let conf = tenant_manager.conf;
-    let resources = &tenant_manager.resources;
-    let cancel = &tenant_manager.cancel;
-    let background_purges = &tenant_manager.background_purges;
-
     // Initialize dynamic limits that depend on system resources
     let system_memory =
         sysinfo::System::new_with_specifics(sysinfo::RefreshKind::new().with_memory())
@@ -527,7 +512,7 @@ pub async fn init_tenant_mgr(
     let tenant_configs = init_load_tenant_configs(conf).await;
 
     // Determine which tenants are to be secondary or attached, and in which generation
-    let tenant_modes = init_load_generations(conf, &tenant_configs, resources, cancel).await?;
+    let tenant_modes = init_load_generations(conf, &tenant_configs, &resources, &cancel).await?;
 
     tracing::info!(
         "Attaching {} tenants at startup, warming up {} at a time",
@@ -684,10 +669,18 @@ pub async fn init_tenant_mgr(
 
     info!("Processed {} local tenants at startup", tenants.len());
 
-    let mut tenant_map = tenant_manager.tenants.write().unwrap();
-    *tenant_map = TenantsMap::Open(tenants);
+    let mut tenants_map = TENANTS.write().unwrap();
+    assert!(matches!(&*tenants_map, &TenantsMap::Initializing));
 
-    Ok(())
+    *tenants_map = TenantsMap::Open(tenants);
+
+    Ok(TenantManager {
+        conf,
+        tenants: &TENANTS,
+        resources,
+        cancel: CancellationToken::new(),
+        background_purges,
+    })
 }
 
 /// Wrapper for Tenant::spawn that checks invariants before running
@@ -726,6 +719,142 @@ fn tenant_spawn(
     )
 }
 
+async fn shutdown_all_tenants0(tenants: &std::sync::RwLock<TenantsMap>) {
+    let mut join_set = JoinSet::new();
+
+    #[cfg(all(debug_assertions, not(test)))]
+    {
+        // Check that our metrics properly tracked the size of the tenants map.  This is a convenient location to check,
+        // as it happens implicitly at the end of tests etc.
+        let m = tenants.read().unwrap();
+        debug_assert_eq!(METRICS.slots_total(), m.len() as u64);
+    }
+
+    // Atomically, 1. create the shutdown tasks and 2. prevent creation of new tenants.
+    let (total_in_progress, total_attached) = {
+        let mut m = tenants.write().unwrap();
+        match &mut *m {
+            TenantsMap::Initializing => {
+                *m = TenantsMap::ShuttingDown(BTreeMap::default());
+                info!("tenants map is empty");
+                return;
+            }
+            TenantsMap::Open(tenants) => {
+                let mut shutdown_state = BTreeMap::new();
+                let mut total_in_progress = 0;
+                let mut total_attached = 0;
+
+                for (tenant_shard_id, v) in std::mem::take(tenants).into_iter() {
+                    match v {
+                        TenantSlot::Attached(t) => {
+                            shutdown_state.insert(tenant_shard_id, TenantSlot::Attached(t.clone()));
+                            join_set.spawn(
+                                async move {
+                                    let res = {
+                                        let (_guard, shutdown_progress) = completion::channel();
+                                        t.shutdown(shutdown_progress, ShutdownMode::FreezeAndFlush).await
+                                    };
+
+                                    if let Err(other_progress) = res {
+                                        // join the another shutdown in progress
+                                        other_progress.wait().await;
+                                    }
+
+                                    // we cannot afford per tenant logging here, because if s3 is degraded, we are
+                                    // going to log too many lines
+                                    debug!("tenant successfully stopped");
+                                }
+                                .instrument(info_span!("shutdown", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug())),
+                            );
+
+                            total_attached += 1;
+                        }
+                        TenantSlot::Secondary(state) => {
+                            // We don't need to wait for this individually per-tenant: the
+                            // downloader task will be waited on eventually, this cancel
+                            // is just to encourage it to drop out if it is doing work
+                            // for this tenant right now.
+                            state.cancel.cancel();
+
+                            shutdown_state.insert(tenant_shard_id, TenantSlot::Secondary(state));
+                        }
+                        TenantSlot::InProgress(notify) => {
+                            // InProgress tenants are not visible in TenantsMap::ShuttingDown: we will
+                            // wait for their notifications to fire in this function.
+                            join_set.spawn(async move {
+                                notify.wait().await;
+                            });
+
+                            total_in_progress += 1;
+                        }
+                    }
+                }
+                *m = TenantsMap::ShuttingDown(shutdown_state);
+                (total_in_progress, total_attached)
+            }
+            TenantsMap::ShuttingDown(_) => {
+                error!(
+                    "already shutting down, this function isn't supposed to be called more than once"
+                );
+                return;
+            }
+        }
+    };
+
+    let started_at = std::time::Instant::now();
+
+    info!(
+        "Waiting for {} InProgress tenants and {} Attached tenants to shut down",
+        total_in_progress, total_attached
+    );
+
+    let total = join_set.len();
+    let mut panicked = 0;
+    let mut buffering = true;
+    const BUFFER_FOR: std::time::Duration = std::time::Duration::from_millis(500);
+    let mut buffered = std::pin::pin!(tokio::time::sleep(BUFFER_FOR));
+
+    while !join_set.is_empty() {
+        tokio::select! {
+            Some(joined) = join_set.join_next() => {
+                match joined {
+                    Ok(()) => {},
+                    Err(join_error) if join_error.is_cancelled() => {
+                        unreachable!("we are not cancelling any of the tasks");
+                    }
+                    Err(join_error) if join_error.is_panic() => {
+                        // cannot really do anything, as this panic is likely a bug
+                        panicked += 1;
+                    }
+                    Err(join_error) => {
+                        warn!("unknown kind of JoinError: {join_error}");
+                    }
+                }
+                if !buffering {
+                    // buffer so that every 500ms since the first update (or starting) we'll log
+                    // how far away we are; this is because we will get SIGKILL'd at 10s, and we
+                    // are not able to log *then*.
+                    buffering = true;
+                    buffered.as_mut().reset(tokio::time::Instant::now() + BUFFER_FOR);
+                }
+            },
+            _ = &mut buffered, if buffering => {
+                buffering = false;
+                info!(remaining = join_set.len(), total, elapsed_ms = started_at.elapsed().as_millis(), "waiting for tenants to shutdown");
+            }
+        }
+    }
+
+    if panicked > 0 {
+        warn!(
+            panicked,
+            total, "observed panicks while shutting down tenants"
+        );
+    }
+
+    // caller will log how long we took
+}
+
 #[derive(thiserror::Error, Debug)]
 pub(crate) enum UpsertLocationError {
     #[error("Bad config request: {0}")]
@@ -927,8 +1056,7 @@ impl TenantManager {
         // the tenant is inaccessible to the outside world while we are doing this, but that is sensible:
         // the state is ill-defined while we're in transition.  Transitions are async, but fast: we do
         // not do significant I/O, and shutdowns should be prompt via cancellation tokens.
-        let mut slot_guard = self
-            .tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)
+        let mut slot_guard = tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)
             .map_err(|e| match e {
                 TenantSlotError::NotFound(_) => {
                     unreachable!("Called with mode Any")
@@ -1095,75 +1223,6 @@ impl TenantManager {
         }
     }
 
-    fn tenant_map_acquire_slot(
-        &self,
-        tenant_shard_id: &TenantShardId,
-        mode: TenantSlotAcquireMode,
-    ) -> Result<SlotGuard, TenantSlotError> {
-        use TenantSlotAcquireMode::*;
-        METRICS.tenant_slot_writes.inc();
-
-        let mut locked = self.tenants.write().unwrap();
-        let span = tracing::info_span!("acquire_slot", tenant_id=%tenant_shard_id.tenant_id, shard_id = %tenant_shard_id.shard_slug());
-        let _guard = span.enter();
-
-        let m = match &mut *locked {
-            TenantsMap::Initializing => return Err(TenantMapError::StillInitializing.into()),
-            TenantsMap::ShuttingDown(_) => return Err(TenantMapError::ShuttingDown.into()),
-            TenantsMap::Open(m) => m,
-        };
-
-        use std::collections::btree_map::Entry;
-
-        let entry = m.entry(*tenant_shard_id);
-
-        match entry {
-            Entry::Vacant(v) => match mode {
-                MustExist => {
-                    tracing::debug!("Vacant && MustExist: return NotFound");
-                    Err(TenantSlotError::NotFound(*tenant_shard_id))
-                }
-                _ => {
-                    let (completion, barrier) = utils::completion::channel();
-                    let inserting = TenantSlot::InProgress(barrier);
-                    METRICS.slot_inserted(&inserting);
-                    v.insert(inserting);
-                    tracing::debug!("Vacant, inserted InProgress");
-                    Ok(SlotGuard::new(
-                        *tenant_shard_id,
-                        None,
-                        completion,
-                        &self.tenants,
-                    ))
-                }
-            },
-            Entry::Occupied(mut o) => {
-                // Apply mode-driven checks
-                match (o.get(), mode) {
-                    (TenantSlot::InProgress(_), _) => {
-                        tracing::debug!("Occupied, failing for InProgress");
-                        Err(TenantSlotError::InProgress)
-                    }
-                    _ => {
-                        // Happy case: the slot was not in any state that violated our mode
-                        let (completion, barrier) = utils::completion::channel();
-                        let in_progress = TenantSlot::InProgress(barrier);
-                        METRICS.slot_inserted(&in_progress);
-                        let old_value = o.insert(in_progress);
-                        METRICS.slot_removed(&old_value);
-                        tracing::debug!("Occupied, replaced with InProgress");
-                        Ok(SlotGuard::new(
-                            *tenant_shard_id,
-                            Some(old_value),
-                            completion,
-                            &self.tenants,
-                        ))
-                    }
-                }
-            }
-        }
-    }
-
     /// Resetting a tenant is equivalent to detaching it, then attaching it again with the same
     /// LocationConf that was last used to attach it.  Optionally, the local file cache may be
     /// dropped before re-attaching.
@@ -1180,8 +1239,7 @@ impl TenantManager {
         drop_cache: bool,
         ctx: &RequestContext,
     ) -> anyhow::Result<()> {
-        let mut slot_guard =
-            self.tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)?;
+        let mut slot_guard = tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)?;
         let Some(old_slot) = slot_guard.get_old_value() else {
             anyhow::bail!("Tenant not found when trying to reset");
         };
@@ -1330,8 +1388,7 @@ impl TenantManager {
             Ok(())
         }
 
-        let slot_guard =
-            self.tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)?;
+        let slot_guard = tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)?;
         match &slot_guard.old_value {
             Some(TenantSlot::Attached(tenant)) => {
                 // Legacy deletion flow: the tenant remains attached, goes to Stopping state, and
@@ -1482,7 +1539,7 @@ impl TenantManager {
         // Phase 2: Put the parent shard to InProgress and grab a reference to the parent Tenant
         drop(tenant);
         let mut parent_slot_guard =
-            self.tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)?;
+            tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::Any)?;
         let parent = match parent_slot_guard.get_old_value() {
             Some(TenantSlot::Attached(t)) => t,
             Some(TenantSlot::Secondary(_)) => anyhow::bail!("Tenant location in secondary mode"),
@@ -1786,145 +1843,7 @@ impl TenantManager {
     pub(crate) async fn shutdown(&self) {
         self.cancel.cancel();
 
-        self.shutdown_all_tenants0().await
-    }
-
-    async fn shutdown_all_tenants0(&self) {
-        let mut join_set = JoinSet::new();
-
-        #[cfg(all(debug_assertions, not(test)))]
-        {
-            // Check that our metrics properly tracked the size of the tenants map.  This is a convenient location to check,
-            // as it happens implicitly at the end of tests etc.
-            let m = self.tenants.read().unwrap();
-            debug_assert_eq!(METRICS.slots_total(), m.len() as u64);
-        }
-
-        // Atomically, 1. create the shutdown tasks and 2. prevent creation of new tenants.
-        let (total_in_progress, total_attached) = {
-            let mut m = self.tenants.write().unwrap();
-            match &mut *m {
-                TenantsMap::Initializing => {
-                    *m = TenantsMap::ShuttingDown(BTreeMap::default());
-                    info!("tenants map is empty");
-                    return;
-                }
-                TenantsMap::Open(tenants) => {
-                    let mut shutdown_state = BTreeMap::new();
-                    let mut total_in_progress = 0;
-                    let mut total_attached = 0;
-
-                    for (tenant_shard_id, v) in std::mem::take(tenants).into_iter() {
-                        match v {
-                            TenantSlot::Attached(t) => {
-                                shutdown_state
-                                    .insert(tenant_shard_id, TenantSlot::Attached(t.clone()));
-                                join_set.spawn(
-                                    async move {
-                                        let res = {
-                                            let (_guard, shutdown_progress) = completion::channel();
-                                            t.shutdown(shutdown_progress, ShutdownMode::FreezeAndFlush).await
-                                        };
-
-                                        if let Err(other_progress) = res {
-                                            // join the another shutdown in progress
-                                            other_progress.wait().await;
-                                        }
-
-                                        // we cannot afford per tenant logging here, because if s3 is degraded, we are
-                                        // going to log too many lines
-                                        debug!("tenant successfully stopped");
-                                    }
-                                    .instrument(info_span!("shutdown", tenant_id=%tenant_shard_id.tenant_id, shard_id=%tenant_shard_id.shard_slug())),
-                                );
-
-                                total_attached += 1;
-                            }
-                            TenantSlot::Secondary(state) => {
-                                // We don't need to wait for this individually per-tenant: the
-                                // downloader task will be waited on eventually, this cancel
-                                // is just to encourage it to drop out if it is doing work
-                                // for this tenant right now.
-                                state.cancel.cancel();
-
-                                shutdown_state
-                                    .insert(tenant_shard_id, TenantSlot::Secondary(state));
-                            }
-                            TenantSlot::InProgress(notify) => {
-                                // InProgress tenants are not visible in TenantsMap::ShuttingDown: we will
-                                // wait for their notifications to fire in this function.
-                                join_set.spawn(async move {
-                                    notify.wait().await;
-                                });
-
-                                total_in_progress += 1;
-                            }
-                        }
-                    }
-                    *m = TenantsMap::ShuttingDown(shutdown_state);
-                    (total_in_progress, total_attached)
-                }
-                TenantsMap::ShuttingDown(_) => {
-                    error!(
-                        "already shutting down, this function isn't supposed to be called more than once"
-                    );
-                    return;
-                }
-            }
-        };
-
-        let started_at = std::time::Instant::now();
-
-        info!(
-            "Waiting for {} InProgress tenants and {} Attached tenants to shut down",
-            total_in_progress, total_attached
-        );
-
-        let total = join_set.len();
-        let mut panicked = 0;
-        let mut buffering = true;
-        const BUFFER_FOR: std::time::Duration = std::time::Duration::from_millis(500);
-        let mut buffered = std::pin::pin!(tokio::time::sleep(BUFFER_FOR));
-
-        while !join_set.is_empty() {
-            tokio::select! {
-                Some(joined) = join_set.join_next() => {
-                    match joined {
-                        Ok(()) => {},
-                        Err(join_error) if join_error.is_cancelled() => {
-                            unreachable!("we are not cancelling any of the tasks");
-                        }
-                        Err(join_error) if join_error.is_panic() => {
-                            // cannot really do anything, as this panic is likely a bug
-                            panicked += 1;
-                        }
-                        Err(join_error) => {
-                            warn!("unknown kind of JoinError: {join_error}");
-                        }
-                    }
-                    if !buffering {
-                        // buffer so that every 500ms since the first update (or starting) we'll log
-                        // how far away we are; this is because we will get SIGKILL'd at 10s, and we
-                        // are not able to log *then*.
-                        buffering = true;
-                        buffered.as_mut().reset(tokio::time::Instant::now() + BUFFER_FOR);
-                    }
-                },
-                _ = &mut buffered, if buffering => {
-                    buffering = false;
-                    info!(remaining = join_set.len(), total, elapsed_ms = started_at.elapsed().as_millis(), "waiting for tenants to shutdown");
-                }
-            }
-        }
-
-        if panicked > 0 {
-            warn!(
-                panicked,
-                total, "observed panicks while shutting down tenants"
-            );
-        }
-
-        // caller will log how long we took
+        shutdown_all_tenants0(self.tenants).await
     }
 
     /// Detaches a tenant, and removes its local files asynchronously.
@@ -1970,12 +1889,12 @@ impl TenantManager {
                 .map(Some)
         };
 
-        let mut removal_result = self
-            .remove_tenant_from_memory(
-                tenant_shard_id,
-                tenant_dir_rename_operation(tenant_shard_id),
-            )
-            .await;
+        let mut removal_result = remove_tenant_from_memory(
+            self.tenants,
+            tenant_shard_id,
+            tenant_dir_rename_operation(tenant_shard_id),
+        )
+        .await;
 
         // If the tenant was not found, it was likely already removed. Attempt to remove the tenant
         // directory on disk anyway. For example, during shard splits, we shut down and remove the
@@ -2029,16 +1948,17 @@ impl TenantManager {
     ) -> Result<HashSet<TimelineId>, detach_ancestor::Error> {
         use detach_ancestor::Error;
 
-        let slot_guard = self
-            .tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::MustExist)
-            .map_err(|e| {
-                use TenantSlotError::*;
+        let slot_guard =
+            tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::MustExist).map_err(
+                |e| {
+                    use TenantSlotError::*;
 
-                match e {
-                    MapState(TenantMapError::ShuttingDown) => Error::ShuttingDown,
-                    NotFound(_) | InProgress | MapState(_) => Error::DetachReparent(e.into()),
-                }
-            })?;
+                    match e {
+                        MapState(TenantMapError::ShuttingDown) => Error::ShuttingDown,
+                        NotFound(_) | InProgress | MapState(_) => Error::DetachReparent(e.into()),
+                    }
+                },
+            )?;
 
         let tenant = {
             let old_slot = slot_guard
@@ -2371,80 +2291,6 @@ impl TenantManager {
             other => ApiError::InternalServerError(anyhow::anyhow!(other)),
         })
     }
-
-    /// Stops and removes the tenant from memory, if it's not [`TenantState::Stopping`] already, bails otherwise.
-    /// Allows to remove other tenant resources manually, via `tenant_cleanup`.
-    /// If the cleanup fails, tenant will stay in memory in [`TenantState::Broken`] state, and another removal
-    async fn remove_tenant_from_memory<V, F>(
-        &self,
-        tenant_shard_id: TenantShardId,
-        tenant_cleanup: F,
-    ) -> Result<V, TenantStateError>
-    where
-        F: std::future::Future<Output = anyhow::Result<V>>,
-    {
-        let mut slot_guard =
-            self.tenant_map_acquire_slot(&tenant_shard_id, TenantSlotAcquireMode::MustExist)?;
-
-        // allow pageserver shutdown to await for our completion
-        let (_guard, progress) = completion::channel();
-
-        // The SlotGuard allows us to manipulate the Tenant object without fear of some
-        // concurrent API request doing something else for the same tenant ID.
-        let attached_tenant = match slot_guard.get_old_value() {
-            Some(TenantSlot::Attached(tenant)) => {
-                // whenever we remove a tenant from memory, we don't want to flush and wait for upload
-                let shutdown_mode = ShutdownMode::Hard;
-
-                // shutdown is sure to transition tenant to stopping, and wait for all tasks to complete, so
-                // that we can continue safely to cleanup.
-                match tenant.shutdown(progress, shutdown_mode).await {
-                    Ok(()) => {}
-                    Err(_other) => {
-                        // if pageserver shutdown or other detach/ignore is already ongoing, we don't want to
-                        // wait for it but return an error right away because these are distinct requests.
-                        slot_guard.revert();
-                        return Err(TenantStateError::IsStopping(tenant_shard_id));
-                    }
-                }
-                Some(tenant)
-            }
-            Some(TenantSlot::Secondary(secondary_state)) => {
-                tracing::info!("Shutting down in secondary mode");
-                secondary_state.shutdown().await;
-                None
-            }
-            Some(TenantSlot::InProgress(_)) => {
-                // Acquiring a slot guarantees its old value was not InProgress
-                unreachable!();
-            }
-            None => None,
-        };
-
-        match tenant_cleanup
-            .await
-            .with_context(|| format!("Failed to run cleanup for tenant {tenant_shard_id}"))
-        {
-            Ok(hook_value) => {
-                // Success: drop the old TenantSlot::Attached.
-                slot_guard
-                    .drop_old_value()
-                    .expect("We just called shutdown");
-
-                Ok(hook_value)
-            }
-            Err(e) => {
-                // If we had a Tenant, set it to Broken and put it back in the TenantsMap
-                if let Some(attached_tenant) = attached_tenant {
-                    attached_tenant.set_broken(e.to_string()).await;
-                }
-                // Leave the broken tenant in the map
-                slot_guard.revert();
-
-                Err(TenantStateError::Other(e))
-            }
-        }
-    }
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -2609,7 +2455,7 @@ pub(crate) enum TenantMapError {
 /// this tenant to retry later, or wait for the InProgress state to end.
 ///
 /// This structure enforces the important invariant that we do not have overlapping
-/// tasks that will try to use local storage for a the same tenant ID: we enforce that
+/// tasks that will try use local storage for a the same tenant ID: we enforce that
 /// the previous contents of a slot have been shut down before the slot can be
 /// left empty or used for something else
 ///
@@ -2622,7 +2468,7 @@ pub(crate) enum TenantMapError {
 /// The `old_value` may be dropped before the SlotGuard is dropped, by calling
 /// `drop_old_value`.  It is an error to call this without shutting down
 /// the conents of `old_value`.
-pub(crate) struct SlotGuard<'a> {
+pub(crate) struct SlotGuard {
     tenant_shard_id: TenantShardId,
     old_value: Option<TenantSlot>,
     upserted: bool,
@@ -2630,23 +2476,19 @@ pub(crate) struct SlotGuard<'a> {
     /// [`TenantSlot::InProgress`] carries the corresponding Barrier: it will
     /// release any waiters as soon as this SlotGuard is dropped.
     completion: utils::completion::Completion,
-
-    tenants: &'a std::sync::RwLock<TenantsMap>,
 }
 
-impl<'a> SlotGuard<'a> {
+impl SlotGuard {
     fn new(
         tenant_shard_id: TenantShardId,
         old_value: Option<TenantSlot>,
         completion: utils::completion::Completion,
-        tenants: &'a std::sync::RwLock<TenantsMap>,
     ) -> Self {
         Self {
             tenant_shard_id,
             old_value,
             upserted: false,
             completion,
-            tenants,
         }
     }
 
@@ -2670,8 +2512,8 @@ impl<'a> SlotGuard<'a> {
             ));
         }
 
-        let replaced: Option<TenantSlot> = {
-            let mut locked = self.tenants.write().unwrap();
+        let replaced = {
+            let mut locked = TENANTS.write().unwrap();
 
             if let TenantSlot::InProgress(_) = new_value {
                 // It is never expected to try and upsert InProgress via this path: it should
@@ -2779,7 +2621,7 @@ impl<'a> SlotGuard<'a> {
     }
 }
 
-impl<'a> Drop for SlotGuard<'a> {
+impl Drop for SlotGuard {
     fn drop(&mut self) {
         if self.upserted {
             return;
@@ -2787,7 +2629,7 @@ impl<'a> Drop for SlotGuard<'a> {
         // Our old value is already shutdown, or it never existed: it is safe
         // for us to fully release the TenantSlot back into an empty state
 
-        let mut locked = self.tenants.write().unwrap();
+        let mut locked = TENANTS.write().unwrap();
 
         let m = match &mut *locked {
             TenantsMap::Initializing => {
@@ -2869,6 +2711,151 @@ enum TenantSlotAcquireMode {
     MustExist,
 }
 
+fn tenant_map_acquire_slot(
+    tenant_shard_id: &TenantShardId,
+    mode: TenantSlotAcquireMode,
+) -> Result<SlotGuard, TenantSlotError> {
+    tenant_map_acquire_slot_impl(tenant_shard_id, &TENANTS, mode)
+}
+
+fn tenant_map_acquire_slot_impl(
+    tenant_shard_id: &TenantShardId,
+    tenants: &std::sync::RwLock<TenantsMap>,
+    mode: TenantSlotAcquireMode,
+) -> Result<SlotGuard, TenantSlotError> {
+    use TenantSlotAcquireMode::*;
+    METRICS.tenant_slot_writes.inc();
+
+    let mut locked = tenants.write().unwrap();
+    let span = tracing::info_span!("acquire_slot", tenant_id=%tenant_shard_id.tenant_id, shard_id = %tenant_shard_id.shard_slug());
+    let _guard = span.enter();
+
+    let m = match &mut *locked {
+        TenantsMap::Initializing => return Err(TenantMapError::StillInitializing.into()),
+        TenantsMap::ShuttingDown(_) => return Err(TenantMapError::ShuttingDown.into()),
+        TenantsMap::Open(m) => m,
+    };
+
+    use std::collections::btree_map::Entry;
+
+    let entry = m.entry(*tenant_shard_id);
+
+    match entry {
+        Entry::Vacant(v) => match mode {
+            MustExist => {
+                tracing::debug!("Vacant && MustExist: return NotFound");
+                Err(TenantSlotError::NotFound(*tenant_shard_id))
+            }
+            _ => {
+                let (completion, barrier) = utils::completion::channel();
+                let inserting = TenantSlot::InProgress(barrier);
+                METRICS.slot_inserted(&inserting);
+                v.insert(inserting);
+                tracing::debug!("Vacant, inserted InProgress");
+                Ok(SlotGuard::new(*tenant_shard_id, None, completion))
+            }
+        },
+        Entry::Occupied(mut o) => {
+            // Apply mode-driven checks
+            match (o.get(), mode) {
+                (TenantSlot::InProgress(_), _) => {
+                    tracing::debug!("Occupied, failing for InProgress");
+                    Err(TenantSlotError::InProgress)
+                }
+                _ => {
+                    // Happy case: the slot was not in any state that violated our mode
+                    let (completion, barrier) = utils::completion::channel();
+                    let in_progress = TenantSlot::InProgress(barrier);
+                    METRICS.slot_inserted(&in_progress);
+                    let old_value = o.insert(in_progress);
+                    METRICS.slot_removed(&old_value);
+                    tracing::debug!("Occupied, replaced with InProgress");
+                    Ok(SlotGuard::new(
+                        *tenant_shard_id,
+                        Some(old_value),
+                        completion,
+                    ))
+                }
+            }
+        }
+    }
+}
+
+/// Stops and removes the tenant from memory, if it's not [`TenantState::Stopping`] already, bails otherwise.
+/// Allows to remove other tenant resources manually, via `tenant_cleanup`.
+/// If the cleanup fails, tenant will stay in memory in [`TenantState::Broken`] state, and another removal
+/// operation would be needed to remove it.
+async fn remove_tenant_from_memory<V, F>(
+    tenants: &std::sync::RwLock<TenantsMap>,
+    tenant_shard_id: TenantShardId,
+    tenant_cleanup: F,
+) -> Result<V, TenantStateError>
+where
+    F: std::future::Future<Output = anyhow::Result<V>>,
+{
+    let mut slot_guard =
+        tenant_map_acquire_slot_impl(&tenant_shard_id, tenants, TenantSlotAcquireMode::MustExist)?;
+
+    // allow pageserver shutdown to await for our completion
+    let (_guard, progress) = completion::channel();
+
+    // The SlotGuard allows us to manipulate the Tenant object without fear of some
+    // concurrent API request doing something else for the same tenant ID.
+    let attached_tenant = match slot_guard.get_old_value() {
+        Some(TenantSlot::Attached(tenant)) => {
+            // whenever we remove a tenant from memory, we don't want to flush and wait for upload
+            let shutdown_mode = ShutdownMode::Hard;
+
+            // shutdown is sure to transition tenant to stopping, and wait for all tasks to complete, so
+            // that we can continue safely to cleanup.
+            match tenant.shutdown(progress, shutdown_mode).await {
+                Ok(()) => {}
+                Err(_other) => {
+                    // if pageserver shutdown or other detach/ignore is already ongoing, we don't want to
+                    // wait for it but return an error right away because these are distinct requests.
+                    slot_guard.revert();
+                    return Err(TenantStateError::IsStopping(tenant_shard_id));
+                }
+            }
+            Some(tenant)
+        }
+        Some(TenantSlot::Secondary(secondary_state)) => {
+            tracing::info!("Shutting down in secondary mode");
+            secondary_state.shutdown().await;
+            None
+        }
+        Some(TenantSlot::InProgress(_)) => {
+            // Acquiring a slot guarantees its old value was not InProgress
+            unreachable!();
+        }
+        None => None,
+    };
+
+    match tenant_cleanup
+        .await
+        .with_context(|| format!("Failed to run cleanup for tenant {tenant_shard_id}"))
+    {
+        Ok(hook_value) => {
+            // Success: drop the old TenantSlot::Attached.
+            slot_guard
+                .drop_old_value()
+                .expect("We just called shutdown");
+
+            Ok(hook_value)
+        }
+        Err(e) => {
+            // If we had a Tenant, set it to Broken and put it back in the TenantsMap
+            if let Some(attached_tenant) = attached_tenant {
+                attached_tenant.set_broken(e.to_string()).await;
+            }
+            // Leave the broken tenant in the map
+            slot_guard.revert();
+
+            Err(TenantStateError::Other(e))
+        }
+    }
+}
+
 use http_utils::error::ApiError;
 use pageserver_api::models::TimelineGcRequest;
 
@@ -2879,15 +2866,11 @@ mod tests {
     use std::collections::BTreeMap;
     use std::sync::Arc;
 
-    use storage_broker::BrokerClientChannel;
     use tracing::Instrument;
 
     use super::super::harness::TenantHarness;
     use super::TenantsMap;
-    use crate::tenant::{
-        TenantSharedResources,
-        mgr::{BackgroundPurges, TenantManager, TenantSlot},
-    };
+    use crate::tenant::mgr::TenantSlot;
 
     #[tokio::test(start_paused = true)]
     async fn shutdown_awaits_in_progress_tenant() {
@@ -2908,47 +2891,23 @@ mod tests {
         let _e = span.enter();
 
         let tenants = BTreeMap::from([(id, TenantSlot::Attached(t.clone()))]);
+        let tenants = Arc::new(std::sync::RwLock::new(TenantsMap::Open(tenants)));
 
         // Invoke remove_tenant_from_memory with a cleanup hook that blocks until we manually
         // permit it to proceed: that will stick the tenant in InProgress
 
-        let (basebackup_prepare_sender, _) = tokio::sync::mpsc::unbounded_channel::<
-            crate::basebackup_cache::BasebackupPrepareRequest,
-        >();
-
-        let tenant_manager = TenantManager {
-            tenants: std::sync::RwLock::new(TenantsMap::Open(tenants)),
-            conf: h.conf,
-            resources: TenantSharedResources {
-                broker_client: BrokerClientChannel::connect_lazy("foobar.com")
-                    .await
-                    .unwrap(),
-                remote_storage: h.remote_storage.clone(),
-                deletion_queue_client: h.deletion_queue.new_client(),
-                l0_flush_global_state: crate::l0_flush::L0FlushGlobalState::new(
-                    h.conf.l0_flush.clone(),
-                ),
-                basebackup_prepare_sender,
-                feature_resolver: crate::feature_resolver::FeatureResolver::new_disabled(),
-            },
-            cancel: tokio_util::sync::CancellationToken::new(),
-            background_purges: BackgroundPurges::default(),
-        };
-
-        let tenant_manager = Arc::new(tenant_manager);
-
         let (until_cleanup_completed, can_complete_cleanup) = utils::completion::channel();
         let (until_cleanup_started, cleanup_started) = utils::completion::channel();
         let mut remove_tenant_from_memory_task = {
-            let tenant_manager = tenant_manager.clone();
             let jh = tokio::spawn({
+                let tenants = tenants.clone();
                 async move {
                     let cleanup = async move {
                         drop(until_cleanup_started);
                         can_complete_cleanup.wait().await;
                         anyhow::Ok(())
                     };
-                    tenant_manager.remove_tenant_from_memory(id, cleanup).await
+                    super::remove_tenant_from_memory(&tenants, id, cleanup).await
                 }
                 .instrument(h.span())
             });
@@ -2961,11 +2920,9 @@ mod tests {
         let mut shutdown_task = {
             let (until_shutdown_started, shutdown_started) = utils::completion::channel();
 
-            let tenant_manager = tenant_manager.clone();
-
             let shutdown_task = tokio::spawn(async move {
                 drop(until_shutdown_started);
-                tenant_manager.shutdown_all_tenants0().await;
+                super::shutdown_all_tenants0(&tenants).await;
             });
 
             shutdown_started.wait().await;

pgxn/neon/communicator.c

@@ -1092,15 +1092,13 @@ communicator_prefetch_register_bufferv(BufferTag tag, neon_request_lsns *frlsns,
 		   MyPState->ring_last <= ring_index);
 }
 
-/* Internal version. Returns the ring index of the last block (result of this function is used only
-*  when nblocks==1)
-*/
+/* internal version. Returns the ring index */
 static uint64
 prefetch_register_bufferv(BufferTag tag, neon_request_lsns *frlsns,
 						  BlockNumber nblocks, const bits8 *mask,
 						  bool is_prefetch)
 {
-	uint64		last_ring_index;
+	uint64		min_ring_index;
 	PrefetchRequest hashkey;
 #ifdef USE_ASSERT_CHECKING
 	bool		any_hits = false;
@@ -1124,12 +1122,13 @@ Retry:
 		MyPState->ring_unused - MyPState->ring_receive;
 	MyNeonCounters->getpage_prefetches_buffered =
 		MyPState->n_responses_buffered;
-	last_ring_index = UINT64_MAX;
 
+	min_ring_index = UINT64_MAX;
 	for (int i = 0; i < nblocks; i++)
 	{
 		PrefetchRequest *slot = NULL;
 		PrfHashEntry *entry = NULL;
+		uint64		ring_index;
 		neon_request_lsns *lsns;
 
 		if (PointerIsValid(mask) && BITMAP_ISSET(mask, i))
@@ -1153,12 +1152,12 @@ Retry:
 		if (entry != NULL)
 		{
 			slot = entry->slot;
-			last_ring_index = slot->my_ring_index;
-			Assert(slot == GetPrfSlot(last_ring_index));
+			ring_index = slot->my_ring_index;
+			Assert(slot == GetPrfSlot(ring_index));
 
 			Assert(slot->status != PRFS_UNUSED);
-			Assert(MyPState->ring_last <= last_ring_index &&
-				   last_ring_index < MyPState->ring_unused);
+			Assert(MyPState->ring_last <= ring_index &&
+				   ring_index < MyPState->ring_unused);
 			Assert(BufferTagsEqual(&slot->buftag, &hashkey.buftag));
 
 			/*
@@ -1170,9 +1169,9 @@ Retry:
 				if (!neon_prefetch_response_usable(lsns, slot))
 				{
 					/* Wait for the old request to finish and discard it */
-					if (!prefetch_wait_for(last_ring_index))
+					if (!prefetch_wait_for(ring_index))
 						goto Retry;
-					prefetch_set_unused(last_ring_index);
+					prefetch_set_unused(ring_index);
 					entry = NULL;
 					slot = NULL;
 					pgBufferUsage.prefetch.expired += 1;
@@ -1189,12 +1188,13 @@ Retry:
 				 */
 				if (slot->status == PRFS_TAG_REMAINS)
 				{
-					prefetch_set_unused(last_ring_index);
+					prefetch_set_unused(ring_index);
 					entry = NULL;
 					slot = NULL;
 				}
 				else
 				{
+					min_ring_index = Min(min_ring_index, ring_index);
 					/* The buffered request is good enough, return that index */
 					if (is_prefetch)
 						pgBufferUsage.prefetch.duplicates++;
@@ -1283,12 +1283,12 @@ Retry:
 		 * The next buffer pointed to by `ring_unused` is now definitely empty, so
 		 * we can insert the new request to it.
 		 */
-		last_ring_index = MyPState->ring_unused;
+		ring_index = MyPState->ring_unused;
 
-		Assert(MyPState->ring_last <= last_ring_index &&
-			   last_ring_index <= MyPState->ring_unused);
+		Assert(MyPState->ring_last <= ring_index &&
+			   ring_index <= MyPState->ring_unused);
 
-		slot = GetPrfSlotNoCheck(last_ring_index);
+		slot = GetPrfSlotNoCheck(ring_index);
 
 		Assert(slot->status == PRFS_UNUSED);
 
@@ -1298,9 +1298,11 @@ Retry:
 		 */
 		slot->buftag = hashkey.buftag;
 		slot->shard_no = get_shard_number(&tag);
-		slot->my_ring_index = last_ring_index;
+		slot->my_ring_index = ring_index;
 		slot->flags = 0;
 
+		min_ring_index = Min(min_ring_index, ring_index);
+
 		if (is_prefetch)
 			MyNeonCounters->getpage_prefetch_requests_total++;
 		else
@@ -1313,12 +1315,11 @@ Retry:
 		MyPState->ring_unused - MyPState->ring_receive;
 
 	Assert(any_hits);
-	Assert(last_ring_index != UINT64_MAX);
 
-	Assert(GetPrfSlot(last_ring_index)->status == PRFS_REQUESTED ||
-		   GetPrfSlot(last_ring_index)->status == PRFS_RECEIVED);
-	Assert(MyPState->ring_last <= last_ring_index &&
-		   last_ring_index < MyPState->ring_unused);
+	Assert(GetPrfSlot(min_ring_index)->status == PRFS_REQUESTED ||
+		   GetPrfSlot(min_ring_index)->status == PRFS_RECEIVED);
+	Assert(MyPState->ring_last <= min_ring_index &&
+		   min_ring_index < MyPState->ring_unused);
 
 	if (flush_every_n_requests > 0 &&
 		MyPState->ring_unused - MyPState->ring_flush >= flush_every_n_requests)
@@ -1334,7 +1335,7 @@ Retry:
 		MyPState->ring_flush = MyPState->ring_unused;
 	}
 
-	return last_ring_index;
+	return min_ring_index;
 }
 
 static bool

pgxn/neon/walproposer.c

@@ -1135,7 +1135,7 @@ VotesCollectedMset(WalProposer *wp, MemberSet *mset, Safekeeper **msk, StringInf
 				wp->propTermStartLsn = sk->voteResponse.flushLsn;
 				wp->donor = sk;
 			}
-			wp->truncateLsn = Max(sk->voteResponse.truncateLsn, wp->truncateLsn);
+			wp->truncateLsn = Max(wp->safekeeper[i].voteResponse.truncateLsn, wp->truncateLsn);
 
 			if (n_votes > 0)
 				appendStringInfoString(s, ", ");

proxy/src/auth/backend/classic.rs

@@ -7,13 +7,13 @@ use crate::auth::{self, AuthFlow};
 use crate::config::AuthenticationConfig;
 use crate::context::RequestContext;
 use crate::control_plane::AuthSecret;
-use crate::stream::{PqFeStream, Stream};
+use crate::stream::{PqStream, Stream};
 use crate::{compute, sasl};
 
 pub(super) async fn authenticate(
     ctx: &RequestContext,
     creds: ComputeUserInfo,
-    client: &mut PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     config: &'static AuthenticationConfig,
     secret: AuthSecret,
 ) -> auth::Result<ComputeCredentials> {

proxy/src/auth/backend/console_redirect.rs

@@ -14,10 +14,10 @@ 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::wake_compute::WakeComputeBackend;
-use crate::stream::PqFeStream;
+use crate::stream::PqStream;
 use crate::types::RoleName;
 use crate::{auth, compute, waiters};
 
@@ -96,7 +96,7 @@ impl ConsoleRedirectBackend {
         &self,
         ctx: &RequestContext,
         auth_config: &'static AuthenticationConfig,
-        client: &mut PqFeStream<impl AsyncRead + AsyncWrite + Unpin>,
+        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     ) -> auth::Result<(ConsoleRedirectNodeInfo, AuthInfo, ComputeUserInfo)> {
         authenticate(ctx, auth_config, &self.console_uri, client)
             .await
@@ -109,7 +109,7 @@ impl ConsoleRedirectBackend {
 pub struct ConsoleRedirectNodeInfo(pub(super) NodeInfo);
 
 #[async_trait]
-impl WakeComputeBackend for ConsoleRedirectNodeInfo {
+impl ComputeConnectBackend for ConsoleRedirectNodeInfo {
     async fn wake_compute(
         &self,
         _ctx: &RequestContext,
@@ -122,7 +122,7 @@ async fn authenticate(
     ctx: &RequestContext,
     auth_config: &'static AuthenticationConfig,
     link_uri: &reqwest::Url,
-    client: &mut PqFeStream<impl AsyncRead + AsyncWrite + Unpin>,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
 ) -> auth::Result<(NodeInfo, AuthInfo, ComputeUserInfo)> {
     ctx.set_auth_method(crate::context::AuthMethod::ConsoleRedirect);
 

proxy/src/auth/backend/hacks.rs

@@ -17,7 +17,7 @@ use crate::stream::{self, Stream};
 pub(crate) async fn authenticate_cleartext(
     ctx: &RequestContext,
     info: ComputeUserInfo,
-    client: &mut stream::PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     secret: AuthSecret,
     config: &'static AuthenticationConfig,
 ) -> auth::Result<ComputeCredentials> {
@@ -61,7 +61,7 @@ pub(crate) async fn authenticate_cleartext(
 pub(crate) async fn password_hack_no_authentication(
     ctx: &RequestContext,
     info: ComputeUserInfoNoEndpoint,
-    client: &mut stream::PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
 ) -> auth::Result<(ComputeUserInfo, Vec<u8>)> {
     debug!("project not specified, resorting to the password hack auth flow");
     ctx.set_auth_method(crate::context::AuthMethod::Cleartext);

proxy/src/auth/backend/mod.rs

@@ -14,21 +14,20 @@ use serde::{Deserialize, Serialize};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tracing::{debug, info};
 
-use crate::auth::{self, ComputeUserInfoMaybeEndpoint, validate_password_and_exchange};
+use crate::auth::{self, AuthError, ComputeUserInfoMaybeEndpoint, validate_password_and_exchange};
 use crate::cache::Cached;
 use crate::config::AuthenticationConfig;
 use crate::context::RequestContext;
 use crate::control_plane::client::ControlPlaneClient;
 use crate::control_plane::errors::GetAuthInfoError;
-use crate::control_plane::messages::EndpointRateLimitConfig;
 use crate::control_plane::{
     self, AccessBlockerFlags, AuthSecret, CachedNodeInfo, ControlPlaneApi, EndpointAccessControl,
     RoleAccessControl,
 };
 use crate::intern::EndpointIdInt;
+use crate::pglb::connect_compute::ComputeConnectBackend;
 use crate::pqproto::BeMessage;
 use crate::proxy::NeonOptions;
-use crate::proxy::wake_compute::WakeComputeBackend;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::stream::Stream;
 use crate::types::{EndpointCacheKey, EndpointId, RoleName};
@@ -201,7 +200,7 @@ async fn auth_quirks(
     ctx: &RequestContext,
     api: &impl control_plane::ControlPlaneApi,
     user_info: ComputeUserInfoMaybeEndpoint,
-    client: &mut stream::PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     allow_cleartext: bool,
     config: &'static AuthenticationConfig,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
@@ -231,8 +230,11 @@ async fn auth_quirks(
         config.is_vpc_acccess_proxy,
     )?;
 
-    access_controls.connection_attempt_rate_limit(ctx, &info.endpoint, &endpoint_rate_limiter)?;
-
+    let endpoint = EndpointIdInt::from(&info.endpoint);
+    let rate_limit_config = None;
+    if !endpoint_rate_limiter.check(endpoint, rate_limit_config, 1) {
+        return Err(AuthError::too_many_connections());
+    }
     let role_access = api
         .get_role_access_control(ctx, &info.endpoint, &info.user)
         .await?;
@@ -267,7 +269,7 @@ async fn authenticate_with_secret(
     ctx: &RequestContext,
     secret: AuthSecret,
     info: ComputeUserInfo,
-    client: &mut stream::PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     unauthenticated_password: Option<Vec<u8>>,
     allow_cleartext: bool,
     config: &'static AuthenticationConfig,
@@ -318,7 +320,7 @@ impl<'a> Backend<'a, ComputeUserInfoMaybeEndpoint> {
     pub(crate) async fn authenticate(
         self,
         ctx: &RequestContext,
-        client: &mut stream::PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+        client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
         allow_cleartext: bool,
         config: &'static AuthenticationConfig,
         endpoint_rate_limiter: Arc<EndpointRateLimiter>,
@@ -399,20 +401,19 @@ impl Backend<'_, ComputeUserInfo> {
                 allowed_ips: Arc::new(vec![]),
                 allowed_vpce: Arc::new(vec![]),
                 flags: AccessBlockerFlags::default(),
-                rate_limits: EndpointRateLimitConfig::default(),
             }),
         }
     }
 }
 
 #[async_trait::async_trait]
-impl WakeComputeBackend for Backend<'_, ComputeUserInfo> {
+impl ComputeConnectBackend for Backend<'_, ComputeCredentials> {
     async fn wake_compute(
         &self,
         ctx: &RequestContext,
     ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError> {
         match self {
-            Self::ControlPlane(api, info) => api.wake_compute(ctx, info).await,
+            Self::ControlPlane(api, creds) => api.wake_compute(ctx, &creds.info).await,
             Self::Local(local) => Ok(Cached::new_uncached(local.node_info.clone())),
         }
     }
@@ -438,7 +439,6 @@ mod tests {
     use crate::auth::{ComputeUserInfoMaybeEndpoint, IpPattern};
     use crate::config::AuthenticationConfig;
     use crate::context::RequestContext;
-    use crate::control_plane::messages::EndpointRateLimitConfig;
     use crate::control_plane::{
         self, AccessBlockerFlags, CachedNodeInfo, EndpointAccessControl, RoleAccessControl,
     };
@@ -446,7 +446,7 @@ mod tests {
     use crate::rate_limiter::EndpointRateLimiter;
     use crate::scram::ServerSecret;
     use crate::scram::threadpool::ThreadPool;
-    use crate::stream::{PqFeStream, Stream};
+    use crate::stream::{PqStream, Stream};
 
     struct Auth {
         ips: Vec<IpPattern>,
@@ -477,7 +477,6 @@ mod tests {
                 allowed_ips: Arc::new(self.ips.clone()),
                 allowed_vpce: Arc::new(self.vpc_endpoint_ids.clone()),
                 flags: self.access_blocker_flags,
-                rate_limits: EndpointRateLimitConfig::default(),
             })
         }
 
@@ -522,7 +521,7 @@ mod tests {
     #[tokio::test]
     async fn auth_quirks_scram() {
         let (mut client, server) = tokio::io::duplex(1024);
-        let mut stream = PqFeStream::new_skip_handshake(Stream::from_raw(server));
+        let mut stream = PqStream::new_skip_handshake(Stream::from_raw(server));
 
         let ctx = RequestContext::test();
         let api = Auth {
@@ -604,7 +603,7 @@ mod tests {
     #[tokio::test]
     async fn auth_quirks_cleartext() {
         let (mut client, server) = tokio::io::duplex(1024);
-        let mut stream = PqFeStream::new_skip_handshake(Stream::from_raw(server));
+        let mut stream = PqStream::new_skip_handshake(Stream::from_raw(server));
 
         let ctx = RequestContext::test();
         let api = Auth {
@@ -658,7 +657,7 @@ mod tests {
     #[tokio::test]
     async fn auth_quirks_password_hack() {
         let (mut client, server) = tokio::io::duplex(1024);
-        let mut stream = PqFeStream::new_skip_handshake(Stream::from_raw(server));
+        let mut stream = PqStream::new_skip_handshake(Stream::from_raw(server));
 
         let ctx = RequestContext::test();
         let api = Auth {

proxy/src/auth/flow.rs

@@ -15,7 +15,7 @@ use crate::pqproto::{BeAuthenticationSaslMessage, BeMessage};
 use crate::sasl;
 use crate::scram::threadpool::ThreadPool;
 use crate::scram::{self};
-use crate::stream::{PqFeStream, Stream};
+use crate::stream::{PqStream, Stream};
 use crate::tls::TlsServerEndPoint;
 
 /// Use [SCRAM](crate::scram)-based auth in [`AuthFlow`].
@@ -53,7 +53,7 @@ pub(crate) struct CleartextPassword {
 #[must_use]
 pub(crate) struct AuthFlow<'a, S, State> {
     /// The underlying stream which implements libpq's protocol.
-    stream: &'a mut PqFeStream<Stream<S>>,
+    stream: &'a mut PqStream<Stream<S>>,
     /// State might contain ancillary data.
     state: State,
     tls_server_end_point: TlsServerEndPoint,
@@ -62,7 +62,7 @@ pub(crate) struct AuthFlow<'a, S, State> {
 /// Initial state of the stream wrapper.
 impl<'a, S: AsyncRead + AsyncWrite + Unpin, M> AuthFlow<'a, S, M> {
     /// Create a new wrapper for client authentication.
-    pub(crate) fn new(stream: &'a mut PqFeStream<Stream<S>>, method: M) -> Self {
+    pub(crate) fn new(stream: &'a mut PqStream<Stream<S>>, method: M) -> Self {
         let tls_server_end_point = stream.get_ref().tls_server_end_point();
 
         Self {

proxy/src/batch.rs

@@ -1,145 +0,0 @@
-//! Batch processing system based on intrusive linked lists.
-//!
-//! Enqueuing a batch job requires no allocations, with
-//! direct support for cancelling jobs early.
-use std::collections::BTreeMap;
-use std::pin::pin;
-use std::sync::Mutex;
-
-use futures::future::Either;
-use scopeguard::ScopeGuard;
-use tokio::sync::oneshot::error::TryRecvError;
-
-use crate::ext::LockExt;
-
-pub trait QueueProcessing: Send + 'static {
-    type Req: Send + 'static;
-    type Res: Send;
-
-    /// Get the desired batch size.
-    fn batch_size(&self, queue_size: usize) -> usize;
-
-    /// This applies a full batch of events.
-    /// Must respond with a full batch of replies.
-    ///
-    /// If this apply can error, it's expected that errors be forwarded to each Self::Res.
-    ///
-    /// Batching does not need to happen atomically.
-    fn apply(&mut self, req: Vec<Self::Req>) -> impl Future<Output = Vec<Self::Res>> + Send;
-}
-
-pub struct BatchQueue<P: QueueProcessing> {
-    processor: tokio::sync::Mutex<P>,
-    inner: Mutex<BatchQueueInner<P>>,
-}
-
-struct BatchJob<P: QueueProcessing> {
-    req: P::Req,
-    res: tokio::sync::oneshot::Sender<P::Res>,
-}
-
-impl<P: QueueProcessing> BatchQueue<P> {
-    pub fn new(p: P) -> Self {
-        Self {
-            processor: tokio::sync::Mutex::new(p),
-            inner: Mutex::new(BatchQueueInner {
-                version: 0,
-                queue: BTreeMap::new(),
-            }),
-        }
-    }
-
-    pub async fn call(&self, req: P::Req) -> P::Res {
-        let (id, mut rx) = self.inner.lock_propagate_poison().register_job(req);
-        let guard = scopeguard::guard(id, move |id| {
-            if self.inner.lock_propagate_poison().queue.remove(&id).is_some() {
-                tracing::debug!("batched task cancelled before completion");
-            }
-        });
-
-        let resp = loop {
-            // try become the leader, or try wait for success.
-            let mut processor = match futures::future::select(rx, pin!(self.processor.lock())).await
-            {
-                // we got the resp.
-                Either::Left((resp, _)) => break resp.ok(),
-                // we are the leader.
-                Either::Right((p, rx_)) => {
-                    rx = rx_;
-                    p
-                }
-            };
-
-            let (reqs, resps) = self.inner.lock_propagate_poison().get_batch(&processor);
-
-            // apply a batch.
-            let values = processor.apply(reqs).await;
-
-            // send response values.
-            for (tx, value) in std::iter::zip(resps, values) {
-                // sender hung up but that's fine.
-                drop(tx.send(value));
-            }
-
-            match rx.try_recv() {
-                Ok(resp) => break Some(resp),
-                Err(TryRecvError::Closed) => break None,
-                // edge case - there was a race condition where
-                // we became the leader but were not in the batch.
-                //
-                // Example:
-                // thread 1: register job id=1
-                // thread 2: register job id=2
-                // thread 2: processor.lock().await
-                // thread 1: processor.lock().await
-                // thread 2: becomes leader, batch_size=1, jobs=[1].
-                Err(TryRecvError::Empty) => {}
-            }
-        };
-
-        // already removed.
-        ScopeGuard::into_inner(guard);
-
-        resp.expect("no response found. batch processer should not panic")
-    }
-}
-
-struct BatchQueueInner<P: QueueProcessing> {
-    version: u64,
-    queue: BTreeMap<u64, BatchJob<P>>,
-}
-
-impl<P: QueueProcessing> BatchQueueInner<P> {
-    fn register_job(&mut self, req: P::Req) -> (u64, tokio::sync::oneshot::Receiver<P::Res>) {
-        let (tx, rx) = tokio::sync::oneshot::channel();
-
-        let id = self.version;
-
-        // Overflow concern:
-        // This is a u64, and we might enqueue 2^16 tasks per second.
-        // This gives us 2^48 seconds (9 million years).
-        // Even if this does overflow, it will not break, but some
-        // jobs with the higher version might never get prioritised.
-        self.version += 1;
-
-        self.queue.insert(id, BatchJob { req, res: tx });
-
-        (id, rx)
-    }
-
-    fn get_batch(&mut self, p: &P) -> (Vec<P::Req>, Vec<tokio::sync::oneshot::Sender<P::Res>>) {
-        let batch_size = p.batch_size(self.queue.len());
-        let mut reqs = Vec::with_capacity(batch_size);
-        let mut resps = Vec::with_capacity(batch_size);
-
-        while reqs.len() < batch_size {
-            let Some((_, job)) = self.queue.pop_first() else {
-                break;
-            };
-            reqs.push(job.req);
-            resps.push(job.res);
-        }
-
-        (reqs, resps)
-    }
-}

proxy/src/binary/local_proxy.rs

@@ -201,7 +201,7 @@ pub async fn run() -> anyhow::Result<()> {
         auth_backend,
         http_listener,
         shutdown.clone(),
-        Arc::new(CancellationHandler::new(&config.connect_to_compute)),
+        Arc::new(CancellationHandler::new(&config.connect_to_compute, None)),
         endpoint_rate_limiter,
     );
 

proxy/src/binary/pg_sni_router.rs

@@ -28,9 +28,10 @@ use crate::context::RequestContext;
 use crate::metrics::{Metrics, ThreadPoolMetrics};
 use crate::pqproto::FeStartupPacket;
 use crate::protocol2::ConnectionInfo;
-use crate::proxy::{ErrorSource, TlsRequired, copy_bidirectional_client_compute};
-use crate::stream::{PqFeStream, Stream};
-use crate::util::run_until_cancelled;
+use crate::proxy::{
+    ErrorSource, TlsRequired, copy_bidirectional_client_compute, run_until_cancelled,
+};
+use crate::stream::{PqStream, Stream};
 
 project_git_version!(GIT_VERSION);
 
@@ -262,7 +263,7 @@ async fn ssl_handshake<S: AsyncRead + AsyncWrite + Unpin>(
     raw_stream: S,
     tls_config: Arc<rustls::ServerConfig>,
 ) -> anyhow::Result<TlsStream<S>> {
-    let (mut stream, msg) = PqFeStream::parse_startup(Stream::from_raw(raw_stream)).await?;
+    let (mut stream, msg) = PqStream::parse_startup(Stream::from_raw(raw_stream)).await?;
     match msg {
         FeStartupPacket::SslRequest { direct: None } => {
             let raw = stream.accept_tls().await?;

proxy/src/binary/proxy.rs

@@ -21,8 +21,7 @@ use utils::{project_build_tag, project_git_version};
 
 use crate::auth::backend::jwt::JwkCache;
 use crate::auth::backend::{ConsoleRedirectBackend, MaybeOwned};
-use crate::batch::BatchQueue;
-use crate::cancellation::{CancellationHandler, CancellationProcessor};
+use crate::cancellation::{CancellationHandler, handle_cancel_messages};
 use crate::config::{
     self, AuthenticationConfig, CacheOptions, ComputeConfig, HttpConfig, ProjectInfoCacheOptions,
     ProxyConfig, ProxyProtocolV2, remote_storage_from_toml,
@@ -391,7 +390,13 @@ pub async fn run() -> anyhow::Result<()> {
         .as_ref()
         .map(|redis_publisher| RedisKVClient::new(redis_publisher.clone(), redis_rps_limit));
 
-    let cancellation_handler = Arc::new(CancellationHandler::new(&config.connect_to_compute));
+    // channel size should be higher than redis client limit to avoid blocking
+    let cancel_ch_size = args.cancellation_ch_size;
+    let (tx_cancel, rx_cancel) = tokio::sync::mpsc::channel(cancel_ch_size);
+    let cancellation_handler = Arc::new(CancellationHandler::new(
+        &config.connect_to_compute,
+        Some(tx_cancel),
+    ));
 
     let endpoint_rate_limiter = Arc::new(EndpointRateLimiter::new_with_shards(
         RateBucketInfo::to_leaky_bucket(&args.endpoint_rps_limit)
@@ -518,10 +523,14 @@ pub async fn run() -> anyhow::Result<()> {
             if let Some(mut redis_kv_client) = redis_kv_client {
                 maintenance_tasks.spawn(async move {
                     redis_kv_client.try_connect().await?;
-                    cancellation_handler.init_tx(BatchQueue::new(CancellationProcessor {
-                        client: redis_kv_client,
-                        batch_size: args.cancellation_batch_size,
-                    }));
+                    handle_cancel_messages(
+                        &mut redis_kv_client,
+                        rx_cancel,
+                        args.cancellation_batch_size,
+                    )
+                    .await?;
+
+                    drop(redis_kv_client);
 
                     // `handle_cancel_messages` was terminated due to the tx_cancel
                     // being dropped. this is not worthy of an error, and this task can only return `Err`,

proxy/src/cache/project_info.rs

@@ -364,7 +364,6 @@ mod tests {
     use std::sync::Arc;
 
     use super::*;
-    use crate::control_plane::messages::EndpointRateLimitConfig;
     use crate::control_plane::{AccessBlockerFlags, AuthSecret};
     use crate::scram::ServerSecret;
     use crate::types::ProjectId;
@@ -400,7 +399,6 @@ mod tests {
                 allowed_ips: allowed_ips.clone(),
                 allowed_vpce: Arc::new(vec![]),
                 flags: AccessBlockerFlags::default(),
-                rate_limits: EndpointRateLimitConfig::default(),
             },
             RoleAccessControl {
                 secret: secret1.clone(),
@@ -416,7 +414,6 @@ mod tests {
                 allowed_ips: allowed_ips.clone(),
                 allowed_vpce: Arc::new(vec![]),
                 flags: AccessBlockerFlags::default(),
-                rate_limits: EndpointRateLimitConfig::default(),
             },
             RoleAccessControl {
                 secret: secret2.clone(),
@@ -442,7 +439,6 @@ mod tests {
                 allowed_ips: allowed_ips.clone(),
                 allowed_vpce: Arc::new(vec![]),
                 flags: AccessBlockerFlags::default(),
-                rate_limits: EndpointRateLimitConfig::default(),
             },
             RoleAccessControl {
                 secret: secret3.clone(),

proxy/src/cancellation.rs

@@ -1,23 +1,19 @@
-use std::convert::Infallible;
 use std::net::{IpAddr, SocketAddr};
-use std::sync::{Arc, OnceLock};
-use std::time::Duration;
+use std::sync::Arc;
 
-use anyhow::anyhow;
-use futures::FutureExt;
+use anyhow::{Context, anyhow};
 use ipnet::{IpNet, Ipv4Net, Ipv6Net};
-use postgres_client::RawCancelToken;
+use postgres_client::CancelToken;
 use postgres_client::tls::MakeTlsConnect;
 use redis::{Cmd, FromRedisValue, Value};
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio::net::TcpStream;
-use tokio::time::timeout;
-use tracing::{debug, error, info};
+use tokio::sync::{mpsc, oneshot};
+use tracing::{debug, error, info, warn};
 
 use crate::auth::AuthError;
 use crate::auth::backend::ComputeUserInfo;
-use crate::batch::{BatchQueue, QueueProcessing};
 use crate::config::ComputeConfig;
 use crate::context::RequestContext;
 use crate::control_plane::ControlPlaneApi;
@@ -31,36 +27,46 @@ use crate::redis::kv_ops::RedisKVClient;
 
 type IpSubnetKey = IpNet;
 
-const CANCEL_KEY_TTL: std::time::Duration = std::time::Duration::from_secs(600);
-const CANCEL_KEY_REFRESH: std::time::Duration = std::time::Duration::from_secs(570);
+const CANCEL_KEY_TTL: i64 = 1_209_600; // 2 weeks cancellation key expire time
 
 // Message types for sending through mpsc channel
 pub enum CancelKeyOp {
     StoreCancelKey {
-        key: CancelKeyData,
-        value: Box<str>,
-        expire: std::time::Duration,
+        key: String,
+        field: String,
+        value: String,
+        resp_tx: Option<oneshot::Sender<anyhow::Result<()>>>,
+        _guard: CancelChannelSizeGuard<'static>,
+        expire: i64, // TTL for key
     },
     GetCancelData {
-        key: CancelKeyData,
+        key: String,
+        resp_tx: oneshot::Sender<anyhow::Result<Vec<(String, String)>>>,
+        _guard: CancelChannelSizeGuard<'static>,
+    },
+    RemoveCancelKey {
+        key: String,
+        field: String,
+        resp_tx: Option<oneshot::Sender<anyhow::Result<()>>>,
+        _guard: CancelChannelSizeGuard<'static>,
     },
 }
 
 pub struct Pipeline {
     inner: redis::Pipeline,
-    replies: usize,
+    replies: Vec<CancelReplyOp>,
 }
 
 impl Pipeline {
     fn with_capacity(n: usize) -> Self {
         Self {
             inner: redis::Pipeline::with_capacity(n),
-            replies: 0,
+            replies: Vec::with_capacity(n),
         }
     }
 
-    async fn execute(self, client: &mut RedisKVClient) -> Vec<anyhow::Result<Value>> {
-        let responses = self.replies;
+    async fn execute(&mut self, client: &mut RedisKVClient) {
+        let responses = self.replies.len();
         let batch_size = self.inner.len();
 
         match client.query(&self.inner).await {
@@ -70,73 +76,176 @@ impl Pipeline {
                     batch_size,
                     responses, "successfully completed cancellation jobs",
                 );
-                values.into_iter().map(Ok).collect()
+                for (value, reply) in std::iter::zip(values, self.replies.drain(..)) {
+                    reply.send_value(value);
+                }
             }
             Ok(value) => {
                 error!(batch_size, ?value, "unexpected redis return value");
-                std::iter::repeat_with(|| Err(anyhow!("incorrect response type from redis")))
-                    .take(responses)
-                    .collect()
+                for reply in self.replies.drain(..) {
+                    reply.send_err(anyhow!("incorrect response type from redis"));
+                }
             }
             Err(err) => {
-                std::iter::repeat_with(|| Err(anyhow!("could not send cmd to redis: {err}")))
-                    .take(responses)
-                    .collect()
+                for reply in self.replies.drain(..) {
+                    reply.send_err(anyhow!("could not send cmd to redis: {err}"));
+                }
             }
         }
+
+        self.inner.clear();
+        self.replies.clear();
     }
 
-    fn add_command_with_reply(&mut self, cmd: Cmd) {
+    fn add_command_with_reply(&mut self, cmd: Cmd, reply: CancelReplyOp) {
         self.inner.add_command(cmd);
-        self.replies += 1;
+        self.replies.push(reply);
     }
 
     fn add_command_no_reply(&mut self, cmd: Cmd) {
         self.inner.add_command(cmd).ignore();
     }
+
+    fn add_command(&mut self, cmd: Cmd, reply: Option<CancelReplyOp>) {
+        match reply {
+            Some(reply) => self.add_command_with_reply(cmd, reply),
+            None => self.add_command_no_reply(cmd),
+        }
+    }
 }
 
 impl CancelKeyOp {
-    fn register(&self, pipe: &mut Pipeline) {
+    fn register(self, pipe: &mut Pipeline) {
         #[allow(clippy::used_underscore_binding)]
         match self {
-            CancelKeyOp::StoreCancelKey { key, value, expire } => {
-                let key = KeyPrefix::Cancel(*key).build_redis_key();
-                pipe.add_command_with_reply(Cmd::hset(&key, "data", &**value));
-                pipe.add_command_no_reply(Cmd::expire(&key, expire.as_secs() as i64));
+            CancelKeyOp::StoreCancelKey {
+                key,
+                field,
+                value,
+                resp_tx,
+                _guard,
+                expire,
+            } => {
+                let reply =
+                    resp_tx.map(|resp_tx| CancelReplyOp::StoreCancelKey { resp_tx, _guard });
+                pipe.add_command(Cmd::hset(&key, field, value), reply);
+                pipe.add_command_no_reply(Cmd::expire(key, expire));
             }
-            CancelKeyOp::GetCancelData { key } => {
-                let key = KeyPrefix::Cancel(*key).build_redis_key();
-                pipe.add_command_with_reply(Cmd::hget(key, "data"));
+            CancelKeyOp::GetCancelData {
+                key,
+                resp_tx,
+                _guard,
+            } => {
+                let reply = CancelReplyOp::GetCancelData { resp_tx, _guard };
+                pipe.add_command_with_reply(Cmd::hgetall(key), reply);
+            }
+            CancelKeyOp::RemoveCancelKey {
+                key,
+                field,
+                resp_tx,
+                _guard,
+            } => {
+                let reply =
+                    resp_tx.map(|resp_tx| CancelReplyOp::RemoveCancelKey { resp_tx, _guard });
+                pipe.add_command(Cmd::hdel(key, field), reply);
             }
         }
     }
 }
 
-pub struct CancellationProcessor {
-    pub client: RedisKVClient,
-    pub batch_size: usize,
+// Message types for sending through mpsc channel
+pub enum CancelReplyOp {
+    StoreCancelKey {
+        resp_tx: oneshot::Sender<anyhow::Result<()>>,
+        _guard: CancelChannelSizeGuard<'static>,
+    },
+    GetCancelData {
+        resp_tx: oneshot::Sender<anyhow::Result<Vec<(String, String)>>>,
+        _guard: CancelChannelSizeGuard<'static>,
+    },
+    RemoveCancelKey {
+        resp_tx: oneshot::Sender<anyhow::Result<()>>,
+        _guard: CancelChannelSizeGuard<'static>,
+    },
 }
 
-impl QueueProcessing for CancellationProcessor {
-    type Req = (CancelChannelSizeGuard<'static>, CancelKeyOp);
-    type Res = anyhow::Result<redis::Value>;
-
-    fn batch_size(&self, _queue_size: usize) -> usize {
-        self.batch_size
+impl CancelReplyOp {
+    fn send_err(self, e: anyhow::Error) {
+        match self {
+            CancelReplyOp::StoreCancelKey { resp_tx, _guard } => {
+                resp_tx
+                    .send(Err(e))
+                    .inspect_err(|_| tracing::debug!("could not send reply"))
+                    .ok();
+            }
+            CancelReplyOp::GetCancelData { resp_tx, _guard } => {
+                resp_tx
+                    .send(Err(e))
+                    .inspect_err(|_| tracing::debug!("could not send reply"))
+                    .ok();
+            }
+            CancelReplyOp::RemoveCancelKey { resp_tx, _guard } => {
+                resp_tx
+                    .send(Err(e))
+                    .inspect_err(|_| tracing::debug!("could not send reply"))
+                    .ok();
+            }
+        }
     }
 
-    async fn apply(&mut self, batch: Vec<Self::Req>) -> Vec<Self::Res> {
-        let mut pipeline = Pipeline::with_capacity(batch.len());
+    fn send_value(self, v: redis::Value) {
+        match self {
+            CancelReplyOp::StoreCancelKey { resp_tx, _guard } => {
+                let send =
+                    FromRedisValue::from_owned_redis_value(v).context("could not parse value");
+                resp_tx
+                    .send(send)
+                    .inspect_err(|_| tracing::debug!("could not send reply"))
+                    .ok();
+            }
+            CancelReplyOp::GetCancelData { resp_tx, _guard } => {
+                let send =
+                    FromRedisValue::from_owned_redis_value(v).context("could not parse value");
+                resp_tx
+                    .send(send)
+                    .inspect_err(|_| tracing::debug!("could not send reply"))
+                    .ok();
+            }
+            CancelReplyOp::RemoveCancelKey { resp_tx, _guard } => {
+                let send =
+                    FromRedisValue::from_owned_redis_value(v).context("could not parse value");
+                resp_tx
+                    .send(send)
+                    .inspect_err(|_| tracing::debug!("could not send reply"))
+                    .ok();
+            }
+        }
+    }
+}
+
+// Running as a separate task to accept messages through the rx channel
+pub async fn handle_cancel_messages(
+    client: &mut RedisKVClient,
+    mut rx: mpsc::Receiver<CancelKeyOp>,
+    batch_size: usize,
+) -> anyhow::Result<()> {
+    let mut batch = Vec::with_capacity(batch_size);
+    let mut pipeline = Pipeline::with_capacity(batch_size);
+
+    loop {
+        if rx.recv_many(&mut batch, batch_size).await == 0 {
+            warn!("shutting down cancellation queue");
+            break Ok(());
+        }
 
         let batch_size = batch.len();
         debug!(batch_size, "running cancellation jobs");
 
-        for (_, op) in &batch {
-            op.register(&mut pipeline);
+        for msg in batch.drain(..) {
+            msg.register(&mut pipeline);
         }
 
-        pipeline.execute(&mut self.client).await
+        pipeline.execute(client).await;
     }
 }
 
@@ -147,7 +256,7 @@ pub struct CancellationHandler {
     compute_config: &'static ComputeConfig,
     // rate limiter of cancellation requests
     limiter: Arc<std::sync::Mutex<LeakyBucketRateLimiter<IpSubnetKey>>>,
-    tx: OnceLock<BatchQueue<CancellationProcessor>>, // send messages to the redis KV client task
+    tx: Option<mpsc::Sender<CancelKeyOp>>, // send messages to the redis KV client task
 }
 
 #[derive(Debug, Error)]
@@ -187,10 +296,13 @@ impl ReportableError for CancelError {
 }
 
 impl CancellationHandler {
-    pub fn new(compute_config: &'static ComputeConfig) -> Self {
+    pub fn new(
+        compute_config: &'static ComputeConfig,
+        tx: Option<mpsc::Sender<CancelKeyOp>>,
+    ) -> Self {
         Self {
             compute_config,
-            tx: OnceLock::new(),
+            tx,
             limiter: Arc::new(std::sync::Mutex::new(
                 LeakyBucketRateLimiter::<IpSubnetKey>::new_with_shards(
                     LeakyBucketRateLimiter::<IpSubnetKey>::DEFAULT,
@@ -200,14 +312,7 @@ impl CancellationHandler {
         }
     }
 
-    pub fn init_tx(&self, queue: BatchQueue<CancellationProcessor>) {
-        self.tx
-            .set(queue)
-            .map_err(|_| {})
-            .expect("cancellation queue should be registered once");
-    }
-
-    pub(crate) fn get_key(self: Arc<Self>) -> Session {
+    pub(crate) fn get_key(self: &Arc<Self>) -> Session {
         // we intentionally generate a random "backend pid" and "secret key" here.
         // we use the corresponding u64 as an identifier for the
         // actual endpoint+pid+secret for postgres/pgbouncer.
@@ -217,10 +322,14 @@ impl CancellationHandler {
 
         let key: CancelKeyData = rand::random();
 
+        let prefix_key: KeyPrefix = KeyPrefix::Cancel(key);
+        let redis_key = prefix_key.build_redis_key();
+
         debug!("registered new query cancellation key {key}");
         Session {
             key,
-            cancellation_handler: self,
+            redis_key,
+            cancellation_handler: Arc::clone(self),
         }
     }
 
@@ -228,43 +337,62 @@ impl CancellationHandler {
         &self,
         key: CancelKeyData,
     ) -> Result<Option<CancelClosure>, CancelError> {
-        let guard = Metrics::get()
-            .proxy
-            .cancel_channel_size
-            .guard(RedisMsgKind::HGet);
-        let op = CancelKeyOp::GetCancelData { key };
+        let prefix_key: KeyPrefix = KeyPrefix::Cancel(key);
+        let redis_key = prefix_key.build_redis_key();
 
-        let Some(tx) = self.tx.get() else {
+        let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
+        let op = CancelKeyOp::GetCancelData {
+            key: redis_key,
+            resp_tx,
+            _guard: Metrics::get()
+                .proxy
+                .cancel_channel_size
+                .guard(RedisMsgKind::HGetAll),
+        };
+
+        let Some(tx) = &self.tx else {
             tracing::warn!("cancellation handler is not available");
             return Err(CancelError::InternalError);
         };
 
-        const TIMEOUT: Duration = Duration::from_secs(5);
-        let result = timeout(TIMEOUT, tx.call((guard, op)))
-            .await
-            .map_err(|_| {
-                tracing::warn!("timed out waiting to receive GetCancelData response");
-                CancelError::RateLimit
-            })?
+        tx.try_send(op)
             .map_err(|e| {
-                tracing::warn!("failed to receive GetCancelData response: {e}");
-                CancelError::InternalError
-            })?;
+                tracing::warn!("failed to send GetCancelData for {key}: {e}");
+            })
+            .map_err(|()| CancelError::InternalError)?;
 
-        let cancel_state_str = String::from_owned_redis_value(result).map_err(|e| {
+        let result = resp_rx.await.map_err(|e| {
             tracing::warn!("failed to receive GetCancelData response: {e}");
             CancelError::InternalError
         })?;
 
-        let cancel_closure: CancelClosure =
-            serde_json::from_str(&cancel_state_str).map_err(|e| {
-                tracing::warn!("failed to deserialize cancel state: {e}");
-                CancelError::InternalError
-            })?;
+        let cancel_state_str: Option<String> = match result {
+            Ok(mut state) => {
+                if state.len() == 1 {
+                    Some(state.remove(0).1)
+                } else {
+                    tracing::warn!("unexpected number of entries in cancel state: {state:?}");
+                    return Err(CancelError::InternalError);
+                }
+            }
+            Err(e) => {
+                tracing::warn!("failed to receive cancel state from redis: {e}");
+                return Err(CancelError::InternalError);
+            }
+        };
 
-        Ok(Some(cancel_closure))
+        let cancel_state: Option<CancelClosure> = match cancel_state_str {
+            Some(state) => {
+                let cancel_closure: CancelClosure = serde_json::from_str(&state).map_err(|e| {
+                    tracing::warn!("failed to deserialize cancel state: {e}");
+                    CancelError::InternalError
+                })?;
+                Some(cancel_closure)
+            }
+            None => None,
+        };
+        Ok(cancel_state)
     }
-
     /// Try to cancel a running query for the corresponding connection.
     /// If the cancellation key is not found, it will be published to Redis.
     /// check_allowed - if true, check if the IP is allowed to cancel the query.
@@ -339,10 +467,10 @@ impl CancellationHandler {
 /// This should've been a [`std::future::Future`], but
 /// it's impossible to name a type of an unboxed future
 /// (we'd need something like `#![feature(type_alias_impl_trait)]`).
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize)]
 pub struct CancelClosure {
     socket_addr: SocketAddr,
-    cancel_token: RawCancelToken,
+    cancel_token: CancelToken,
     hostname: String, // for pg_sni router
     user_info: ComputeUserInfo,
 }
@@ -350,7 +478,7 @@ pub struct CancelClosure {
 impl CancelClosure {
     pub(crate) fn new(
         socket_addr: SocketAddr,
-        cancel_token: RawCancelToken,
+        cancel_token: CancelToken,
         hostname: String,
         user_info: ComputeUserInfo,
     ) -> Self {
@@ -363,7 +491,7 @@ impl CancelClosure {
     }
     /// Cancels the query running on user's compute node.
     pub(crate) async fn try_cancel_query(
-        &self,
+        self,
         compute_config: &ComputeConfig,
     ) -> Result<(), CancelError> {
         let socket = TcpStream::connect(self.socket_addr).await?;
@@ -384,6 +512,7 @@ impl CancelClosure {
 pub(crate) struct Session {
     /// The user-facing key identifying this session.
     key: CancelKeyData,
+    redis_key: String,
     cancellation_handler: Arc<CancellationHandler>,
 }
 
@@ -392,66 +521,60 @@ impl Session {
         &self.key
     }
 
-    /// Ensure the cancel key is continously refreshed,
-    /// but stop when the channel is dropped.
-    pub(crate) async fn maintain_cancel_key(
+    // Send the store key op to the cancellation handler and set TTL for the key
+    pub(crate) fn write_cancel_key(
         &self,
-        session_id: uuid::Uuid,
-        cancel: tokio::sync::oneshot::Receiver<Infallible>,
-        cancel_closure: &CancelClosure,
-        compute_config: &ComputeConfig,
-    ) {
-        futures::future::select(
-            std::pin::pin!(self.maintain_redis_cancel_key(cancel_closure)),
-            cancel,
-        )
-        .await;
-
-        if let Err(err) = cancel_closure
-            .try_cancel_query(compute_config)
-            .boxed()
-            .await
-        {
-            tracing::warn!(
-                ?session_id,
-                ?err,
-                "could not cancel the query in the database"
-            );
-        }
-    }
-
-    // Ensure the cancel key is continously refreshed.
-    async fn maintain_redis_cancel_key(&self, cancel_closure: &CancelClosure) -> ! {
-        let Some(tx) = self.cancellation_handler.tx.get() else {
+        cancel_closure: CancelClosure,
+    ) -> Result<(), CancelError> {
+        let Some(tx) = &self.cancellation_handler.tx else {
             tracing::warn!("cancellation handler is not available");
-            // don't exit, as we only want to exit if cancelled externally.
-            std::future::pending().await
+            return Err(CancelError::InternalError);
         };
 
-        let closure_json = serde_json::to_string(&cancel_closure)
-            .expect("serialising to json string should not fail")
-            .into_boxed_str();
+        let closure_json = serde_json::to_string(&cancel_closure).map_err(|e| {
+            tracing::warn!("failed to serialize cancel closure: {e}");
+            CancelError::InternalError
+        })?;
 
-        loop {
-            let guard = Metrics::get()
+        let op = CancelKeyOp::StoreCancelKey {
+            key: self.redis_key.clone(),
+            field: "data".to_string(),
+            value: closure_json,
+            resp_tx: None,
+            _guard: Metrics::get()
                 .proxy
                 .cancel_channel_size
-                .guard(RedisMsgKind::HSet);
-            let op = CancelKeyOp::StoreCancelKey {
-                key: self.key,
-                value: closure_json.clone(),
-                expire: CANCEL_KEY_TTL,
-            };
+                .guard(RedisMsgKind::HSet),
+            expire: CANCEL_KEY_TTL,
+        };
 
-            tracing::debug!(
-                src=%self.key,
-                dest=?cancel_closure.cancel_token,
-                "registering cancellation key"
-            );
+        let _ = tx.try_send(op).map_err(|e| {
+            let key = self.key;
+            tracing::warn!("failed to send StoreCancelKey for {key}: {e}");
+        });
+        Ok(())
+    }
 
-            if tx.call((guard, op)).await.is_ok() {
-                tokio::time::sleep(CANCEL_KEY_REFRESH).await;
-            }
-        }
+    pub(crate) fn remove_cancel_key(&self) -> Result<(), CancelError> {
+        let Some(tx) = &self.cancellation_handler.tx else {
+            tracing::warn!("cancellation handler is not available");
+            return Err(CancelError::InternalError);
+        };
+
+        let op = CancelKeyOp::RemoveCancelKey {
+            key: self.redis_key.clone(),
+            field: "data".to_string(),
+            resp_tx: None,
+            _guard: Metrics::get()
+                .proxy
+                .cancel_channel_size
+                .guard(RedisMsgKind::HDel),
+        };
+
+        let _ = tx.try_send(op).map_err(|e| {
+            let key = self.key;
+            tracing::warn!("failed to send RemoveCancelKey for {key}: {e}");
+        });
+        Ok(())
     }
 }

proxy/src/compute/authenticate.rs

@@ -1,146 +0,0 @@
-use bytes::BufMut;
-use postgres_client::tls::{ChannelBinding, TlsStream};
-use postgres_protocol::authentication::sasl;
-use postgres_protocol::authentication::sasl::{SCRAM_SHA_256, SCRAM_SHA_256_PLUS};
-use tokio::io::{AsyncRead, AsyncWrite};
-
-use super::{Auth, MaybeRustlsStream};
-use crate::compute::RustlsStream;
-use crate::pqproto::{
-    AUTH_OK, AUTH_SASL, AUTH_SASL_CONT, AUTH_SASL_FINAL, FE_PASSWORD_MESSAGE, StartupMessageParams,
-};
-use crate::stream::{PostgresError, PqBeStream};
-
-pub async fn authenticate<S>(
-    stream: MaybeRustlsStream<S>,
-    auth: Option<&Auth>,
-    params: &StartupMessageParams,
-) -> Result<PqBeStream<MaybeRustlsStream<S>>, PostgresError>
-where
-    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
-    RustlsStream<S>: TlsStream + Unpin,
-{
-    let mut stream = PqBeStream::new(stream, params);
-    stream.flush().await?;
-
-    let channel_binding = stream.get_ref().channel_binding();
-
-    // TODO: rather than checking for SASL, maybe we can just assume it.
-    // With SCRAM_SHA_256 if we're not using TLS,
-    // and SCRAM_SHA_256_PLUS if we are using TLS.
-
-    let (channel_binding, mechanism) = match stream.read_auth_message().await? {
-        (AUTH_OK, _) => return Ok(stream),
-        (AUTH_SASL, mechanisms) => {
-            let mut has_scram = false;
-            let mut has_scram_plus = false;
-            for mechanism in mechanisms.split(|&b| b == b'\0') {
-                match mechanism {
-                    b"SCRAM-SHA-256" => has_scram = true,
-                    b"SCRAM-SHA-256-PLUS" => has_scram_plus = true,
-                    _ => {}
-                }
-            }
-
-            match (channel_binding, has_scram, has_scram_plus) {
-                (cb, true, false) => {
-                    if cb.tls_server_end_point.is_some() {
-                        // I don't think this can happen in our setup, but I would like to monitor it.
-                        tracing::warn!(
-                            "TLS is enabled, but compute doesn't support SCRAM-SHA-256-PLUS."
-                        );
-                    }
-                    (sasl::ChannelBinding::unrequested(), SCRAM_SHA_256)
-                }
-                (
-                    ChannelBinding {
-                        tls_server_end_point: None,
-                    },
-                    true,
-                    _,
-                ) => (sasl::ChannelBinding::unsupported(), SCRAM_SHA_256),
-                (
-                    ChannelBinding {
-                        tls_server_end_point: Some(h),
-                    },
-                    _,
-                    true,
-                ) => (
-                    sasl::ChannelBinding::tls_server_end_point(h),
-                    SCRAM_SHA_256_PLUS,
-                ),
-                (_, false, _) => {
-                    tracing::error!(
-                        "compute responded with unsupported auth mechanisms: {}",
-                        String::from_utf8_lossy(mechanisms)
-                    );
-                    return Err(PostgresError::InvalidAuthMessage);
-                }
-            }
-        }
-        (tag, msg) => {
-            tracing::error!(
-                "compute responded with unexpected auth message with tag[{tag}]: {}",
-                String::from_utf8_lossy(msg)
-            );
-            return Err(PostgresError::InvalidAuthMessage);
-        }
-    };
-
-    let mut scram = match auth {
-        // We only touch passwords when it comes to console-redirect.
-        Some(Auth::Password(pw)) => sasl::ScramSha256::new(pw, channel_binding),
-        Some(Auth::Scram(keys)) => sasl::ScramSha256::new_with_keys(**keys, channel_binding),
-        None => {
-            // local_proxy does not set credentials, since it relies on trust and expects an OK message above
-            tracing::error!("compute requested SASL auth, but there are no credentials available",);
-            return Err(PostgresError::InvalidAuthMessage);
-        }
-    };
-
-    stream.write_raw(0, FE_PASSWORD_MESSAGE.0, |buf| {
-        buf.put_slice(mechanism.as_bytes());
-        buf.put_u8(b'\0');
-
-        let data = scram.message();
-        buf.put_u32(data.len() as u32);
-        buf.put_slice(data);
-    });
-    stream.flush().await?;
-
-    loop {
-        // wait for SASLContinue or SASLFinal.
-        match stream.read_auth_message().await? {
-            (AUTH_SASL_CONT, data) => scram.update(data).await?,
-            (AUTH_SASL_FINAL, data) => {
-                scram.finish(data)?;
-                break;
-            }
-            (tag, msg) => {
-                tracing::error!(
-                    "compute responded with unexpected auth message with tag[{tag}]: {}",
-                    String::from_utf8_lossy(msg)
-                );
-                return Err(PostgresError::InvalidAuthMessage);
-            }
-        }
-
-        stream.write_raw(0, FE_PASSWORD_MESSAGE.0, |buf| {
-            buf.put_slice(scram.message());
-        });
-        stream.flush().await?;
-    }
-
-    match stream.read_auth_message().await? {
-        (AUTH_OK, _) => {}
-        (tag, msg) => {
-            tracing::error!(
-                "compute responded with unexpected auth message with tag[{tag}]: {}",
-                String::from_utf8_lossy(msg)
-            );
-            return Err(PostgresError::InvalidAuthMessage);
-        }
-    }
-
-    Ok(stream)
-}

proxy/src/compute/mod.rs

@@ -1,4 +1,3 @@
-mod authenticate;
 mod tls;
 
 use std::fmt::Debug;
@@ -10,12 +9,15 @@ 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_protocol::message::backend::NoticeResponseBody;
 use thiserror::Error;
 use tokio::net::{TcpStream, lookup_host};
 use tracing::{debug, error, info, warn};
 
-use crate::auth::backend::ComputeCredentialKeys;
+use crate::auth::backend::{ComputeCredentialKeys, ComputeUserInfo};
 use crate::auth::parse_endpoint_param;
+use crate::cancellation::CancelClosure;
 use crate::compute::tls::TlsError;
 use crate::config::ComputeConfig;
 use crate::context::RequestContext;
@@ -26,7 +28,6 @@ use crate::error::{ReportableError, UserFacingError};
 use crate::metrics::{Metrics, NumDbConnectionsGuard};
 use crate::pqproto::StartupMessageParams;
 use crate::proxy::neon_option;
-use crate::stream::{PostgresError, PqBeStream};
 use crate::types::Host;
 
 pub const COULD_NOT_CONNECT: &str = "Couldn't connect to compute node";
@@ -36,7 +37,7 @@ pub(crate) enum ConnectionError {
     /// This error doesn't seem to reveal any secrets; for instance,
     /// `postgres_client::error::Kind` doesn't contain ip addresses and such.
     #[error("{COULD_NOT_CONNECT}: {0}")]
-    Postgres(#[from] PostgresError),
+    Postgres(#[from] postgres_client::Error),
 
     #[error("{COULD_NOT_CONNECT}: {0}")]
     TlsError(#[from] TlsError),
@@ -53,21 +54,20 @@ impl UserFacingError for ConnectionError {
         match self {
             // This helps us drop irrelevant library-specific prefixes.
             // TODO: propagate severity level and other parameters.
-            ConnectionError::Postgres(PostgresError::Error(err)) => {
-                let (_code, msg) = err.parse();
-                let msg = String::from_utf8_lossy(msg);
+            ConnectionError::Postgres(err) => match err.as_db_error() {
+                Some(err) => {
+                    let msg = err.message();
 
-                if msg.starts_with("unsupported startup parameter: ")
-                    || msg.starts_with("unsupported startup parameter in options: ")
-                {
-                    format!(
-                        "{msg}. Please use unpooled connection or remove this parameter from the startup package. More details: https://neon.tech/docs/connect/connection-errors#unsupported-startup-parameter"
-                    )
-                } else {
-                    msg.into_owned()
+                    if msg.starts_with("unsupported startup parameter: ")
+                        || msg.starts_with("unsupported startup parameter in options: ")
+                    {
+                        format!("{msg}. Please use unpooled connection or remove this parameter from the startup package. More details: https://neon.tech/docs/connect/connection-errors#unsupported-startup-parameter")
+                    } else {
+                        msg.to_owned()
+                    }
                 }
-            }
-            ConnectionError::Postgres(err) => err.to_string(),
+                None => err.to_string(),
+            },
             ConnectionError::WakeComputeError(err) => err.to_string_client(),
             ConnectionError::TooManyConnectionAttempts(_) => {
                 "Failed to acquire permit to connect to the database. Too many database connection attempts are currently ongoing.".to_owned()
@@ -80,12 +80,10 @@ impl UserFacingError for ConnectionError {
 impl ReportableError for ConnectionError {
     fn get_error_kind(&self) -> crate::error::ErrorKind {
         match self {
-            ConnectionError::Postgres(PostgresError::Io(_)) => crate::error::ErrorKind::Compute,
-            ConnectionError::Postgres(
-                PostgresError::Error(_)
-                | PostgresError::InvalidAuthMessage
-                | PostgresError::Unexpected(_),
-            ) => crate::error::ErrorKind::Postgres,
+            ConnectionError::Postgres(e) if e.as_db_error().is_some() => {
+                crate::error::ErrorKind::Postgres
+            }
+            ConnectionError::Postgres(_) => crate::error::ErrorKind::Compute,
             ConnectionError::TlsError(_) => crate::error::ErrorKind::Compute,
             ConnectionError::WakeComputeError(e) => e.get_error_kind(),
             ConnectionError::TooManyConnectionAttempts(e) => e.get_error_kind(),
@@ -138,11 +136,11 @@ impl AuthInfo {
         }
     }
 
-    pub(crate) fn with_auth_keys(keys: ComputeCredentialKeys) -> Self {
+    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)))
+                    Some(Auth::Scram(Box::new(*auth_keys)))
                 }
                 ComputeCredentialKeys::JwtPayload(_) | ComputeCredentialKeys::None => None,
             },
@@ -164,6 +162,18 @@ impl ConnectInfo {
 }
 
 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(
         &mut self,
@@ -203,7 +213,7 @@ impl ConnectInfo {
     async fn connect_raw(
         &self,
         config: &ComputeConfig,
-    ) -> Result<(SocketAddr, MaybeRustlsStream<TcpStream>), TlsError> {
+    ) -> Result<(SocketAddr, MaybeTlsStream<TcpStream, RustlsStream>), TlsError> {
         let timeout = config.timeout;
 
         // wrap TcpStream::connect with timeout
@@ -255,19 +265,21 @@ impl ConnectInfo {
     }
 }
 
-pub type RustlsStream<S> = <ComputeConfig as MakeTlsConnect<S>>::Stream;
-pub type MaybeRustlsStream<S> = MaybeTlsStream<S, RustlsStream<S>>;
+type RustlsStream = <ComputeConfig as MakeTlsConnect<tokio::net::TcpStream>>::Stream;
 
-pub struct PostgresConnection {
+pub(crate) struct PostgresConnection {
     /// Socket connected to a compute node.
-    pub stream: PqBeStream<MaybeRustlsStream<TcpStream>>,
+    pub(crate) stream: MaybeTlsStream<tokio::net::TcpStream, RustlsStream>,
+    /// PostgreSQL connection parameters.
+    pub(crate) params: std::collections::HashMap<String, String>,
+    /// Query cancellation token.
+    pub(crate) cancel_closure: CancelClosure,
+    /// Labels for proxy's metrics.
+    pub(crate) aux: MetricsAuxInfo,
+    /// Notices received from compute after authenticating
+    pub(crate) delayed_notice: Vec<NoticeResponseBody>,
 
-    pub socket_addr: SocketAddr,
-    pub hostname: String,
-    pub ssl_mode: SslMode,
-    pub aux: MetricsAuxInfo,
-
-    pub guage: NumDbConnectionsGuard<'static>,
+    _guage: NumDbConnectionsGuard<'static>,
 }
 
 impl ConnectInfo {
@@ -275,18 +287,31 @@ impl ConnectInfo {
     pub(crate) async fn connect(
         &self,
         ctx: &RequestContext,
-        aux: &MetricsAuxInfo,
+        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 stream =
-            authenticate::authenticate(stream, auth.auth.as_ref(), &auth.server_params).await?;
+        let connection = tmp_config.connect_raw(stream, NoTls).await?;
         drop(pause);
 
-        // tracing::Span::current().record("pid", tracing::field::display(process_id));
+        let RawConnection {
+            stream,
+            parameters,
+            delayed_notice,
+            process_id,
+            secret_key,
+        } = connection;
+
+        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();
 
         // TODO: lots of useful info but maybe we can move it elsewhere (eg traces?)
         info!(
@@ -298,13 +323,27 @@ impl ConnectInfo {
             ctx.get_testodrome_id().unwrap_or_default(),
         );
 
+        // NB: CancelToken is supposed to hold socket_addr, but we use connect_raw.
+        // Yet another reason to rework the connection establishing code.
+        let cancel_closure = CancelClosure::new(
+            socket_addr,
+            CancelToken {
+                socket_config: None,
+                ssl_mode: self.ssl_mode,
+                process_id,
+                secret_key,
+            },
+            self.host.to_string(),
+            user_info,
+        );
+
         let connection = PostgresConnection {
             stream,
-            socket_addr,
-            hostname: self.host.to_string(),
-            ssl_mode: self.ssl_mode,
-            aux: aux.clone(),
-            guage: Metrics::get().proxy.db_connections.guard(ctx.protocol()),
+            params: parameters,
+            delayed_notice,
+            cancel_closure,
+            aux,
+            _guage: Metrics::get().proxy.db_connections.guard(ctx.protocol()),
         };
 
         Ok(connection)

proxy/src/console_redirect_proxy.rs

@@ -11,12 +11,13 @@ 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::{ClientRequestError, ErrorSource, prepare_client_connection};
-use crate::util::run_until_cancelled;
+use crate::proxy::{
+    ClientRequestError, ErrorSource, prepare_client_connection, run_until_cancelled,
+};
 
 pub async fn task_main(
     config: &'static ProxyConfig,
@@ -120,7 +121,7 @@ pub async fn task_main(
                 Ok(Some(p)) => {
                     ctx.set_success();
                     let _disconnect = ctx.log_connect();
-                    match p.proxy_pass().await {
+                    match p.proxy_pass(&config.connect_to_compute).await {
                         Ok(()) => {}
                         Err(ErrorSource::Client(e)) => {
                             error!(
@@ -177,7 +178,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Client);
     let do_handshake = handshake(ctx, stream, tls, record_handshake_error);
 
-    let (mut client, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
+    let (mut stream, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
         .await??
     {
         HandshakeData::Startup(stream, params) => (stream, params),
@@ -210,17 +211,18 @@ 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
-        .authenticate(ctx, &config.authentication_config, &mut client)
+        .authenticate(ctx, &config.authentication_config, &mut stream)
         .await
     {
         Ok(auth_result) => auth_result,
-        Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
+        Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
     };
     auth_info.set_startup_params(&params, true);
 
-    let mut node = connect_to_compute(
+    let node = connect_to_compute(
         ctx,
         &TcpMechanism {
+            user_info,
             auth: auth_info,
             locks: &config.connect_compute_locks,
         },
@@ -228,41 +230,25 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         config.wake_compute_retry_config,
         &config.connect_to_compute,
     )
-    .or_else(|e| async { Err(client.throw_error(e, Some(ctx)).await) })
+    .or_else(|e| async { Err(stream.throw_error(e, Some(ctx)).await) })
     .await?;
 
-    let session = cancellation_handler.get_key();
+    let cancellation_handler_clone = Arc::clone(&cancellation_handler);
+    let session = cancellation_handler_clone.get_key();
 
-    let cancel_closure =
-        prepare_client_connection(&mut node, session.key(), &mut client, user_info).await?;
+    session.write_cancel_key(node.cancel_closure.clone())?;
 
-    let session_id = ctx.session_id();
-    let (cancel_on_shutdown, cancel) = tokio::sync::oneshot::channel();
-    tokio::spawn(async move {
-        session
-            .maintain_cancel_key(
-                session_id,
-                cancel,
-                &cancel_closure,
-                &config.connect_to_compute,
-            )
-            .await;
-    });
-
-    let client = client.flush_and_into_inner().await?;
-    let compute = node.stream.flush_and_into_inner().await?;
+    prepare_client_connection(&node, *session.key(), &mut stream);
+    let stream = stream.flush_and_into_inner().await?;
 
     Ok(Some(ProxyPassthrough {
-        client,
-        compute,
-
-        aux: node.aux,
+        client: stream,
+        aux: node.aux.clone(),
         private_link_id: None,
-
-        _cancel_on_shutdown: cancel_on_shutdown,
-
+        compute: node,
+        session_id: ctx.session_id(),
+        cancel: session,
         _req: request_gauge,
         _conn: conn_gauge,
-        _db_conn: node.guage,
     }))
 }

proxy/src/control_plane/client/cplane_proxy_v1.rs

@@ -146,7 +146,6 @@ impl NeonControlPlaneClient {
                     public_access_blocked: block_public_connections,
                     vpc_access_blocked: block_vpc_connections,
                 },
-                rate_limits: body.rate_limits,
             })
         }
         .inspect_err(|e| tracing::debug!(error = ?e))
@@ -313,7 +312,6 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
             allowed_ips: Arc::new(auth_info.allowed_ips),
             allowed_vpce: Arc::new(auth_info.allowed_vpc_endpoint_ids),
             flags: auth_info.access_blocker_flags,
-            rate_limits: auth_info.rate_limits,
         };
         let role_control = RoleAccessControl {
             secret: auth_info.secret,
@@ -359,7 +357,6 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
             allowed_ips: Arc::new(auth_info.allowed_ips),
             allowed_vpce: Arc::new(auth_info.allowed_vpc_endpoint_ids),
             flags: auth_info.access_blocker_flags,
-            rate_limits: auth_info.rate_limits,
         };
         let role_control = RoleAccessControl {
             secret: auth_info.secret,

proxy/src/control_plane/client/mock.rs

@@ -20,7 +20,7 @@ use crate::context::RequestContext;
 use crate::control_plane::errors::{
     ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError,
 };
-use crate::control_plane::messages::{EndpointRateLimitConfig, MetricsAuxInfo};
+use crate::control_plane::messages::MetricsAuxInfo;
 use crate::control_plane::{
     AccessBlockerFlags, AuthInfo, AuthSecret, CachedNodeInfo, EndpointAccessControl, NodeInfo,
     RoleAccessControl,
@@ -130,7 +130,6 @@ impl MockControlPlane {
             project_id: None,
             account_id: None,
             access_blocker_flags: AccessBlockerFlags::default(),
-            rate_limits: EndpointRateLimitConfig::default(),
         })
     }
 
@@ -234,7 +233,6 @@ impl super::ControlPlaneApi for MockControlPlane {
             allowed_ips: Arc::new(info.allowed_ips),
             allowed_vpce: Arc::new(info.allowed_vpc_endpoint_ids),
             flags: info.access_blocker_flags,
-            rate_limits: info.rate_limits,
         })
     }
 

proxy/src/control_plane/client/mod.rs

@@ -10,7 +10,6 @@ use clashmap::ClashMap;
 use tokio::time::Instant;
 use tracing::{debug, info};
 
-use super::{EndpointAccessControl, RoleAccessControl};
 use crate::auth::backend::ComputeUserInfo;
 use crate::auth::backend::jwt::{AuthRule, FetchAuthRules, FetchAuthRulesError};
 use crate::cache::endpoints::EndpointsCache;
@@ -23,6 +22,8 @@ use crate::metrics::ApiLockMetrics;
 use crate::rate_limiter::{DynamicLimiter, Outcome, RateLimiterConfig, Token};
 use crate::types::EndpointId;
 
+use super::{EndpointAccessControl, RoleAccessControl};
+
 #[non_exhaustive]
 #[derive(Clone)]
 pub enum ControlPlaneClient {

proxy/src/control_plane/messages.rs

@@ -227,35 +227,12 @@ pub(crate) struct UserFacingMessage {
 #[derive(Deserialize)]
 pub(crate) struct GetEndpointAccessControl {
     pub(crate) role_secret: Box<str>,
-
-    pub(crate) project_id: Option<ProjectIdInt>,
-    pub(crate) account_id: Option<AccountIdInt>,
-
     pub(crate) allowed_ips: Option<Vec<IpPattern>>,
     pub(crate) allowed_vpc_endpoint_ids: Option<Vec<String>>,
+    pub(crate) project_id: Option<ProjectIdInt>,
+    pub(crate) account_id: Option<AccountIdInt>,
     pub(crate) block_public_connections: Option<bool>,
     pub(crate) block_vpc_connections: Option<bool>,
-
-    #[serde(default)]
-    pub(crate) rate_limits: EndpointRateLimitConfig,
-}
-
-#[derive(Copy, Clone, Deserialize, Default)]
-pub struct EndpointRateLimitConfig {
-    pub connection_attempts: ConnectionAttemptsLimit,
-}
-
-#[derive(Copy, Clone, Deserialize, Default)]
-pub struct ConnectionAttemptsLimit {
-    pub tcp: Option<LeakyBucketSetting>,
-    pub ws: Option<LeakyBucketSetting>,
-    pub http: Option<LeakyBucketSetting>,
-}
-
-#[derive(Copy, Clone, Deserialize)]
-pub struct LeakyBucketSetting {
-    pub rps: f64,
-    pub burst: f64,
 }
 
 /// Response which holds compute node's `host:port` pair.

proxy/src/control_plane/mod.rs

@@ -11,8 +11,6 @@ pub(crate) mod errors;
 
 use std::sync::Arc;
 
-use messages::EndpointRateLimitConfig;
-
 use crate::auth::backend::ComputeUserInfo;
 use crate::auth::backend::jwt::AuthRule;
 use crate::auth::{AuthError, IpPattern, check_peer_addr_is_in_list};
@@ -20,9 +18,8 @@ use crate::cache::{Cached, TimedLru};
 use crate::config::ComputeConfig;
 use crate::context::RequestContext;
 use crate::control_plane::messages::{ControlPlaneErrorMessage, MetricsAuxInfo};
-use crate::intern::{AccountIdInt, EndpointIdInt, ProjectIdInt};
+use crate::intern::{AccountIdInt, ProjectIdInt};
 use crate::protocol2::ConnectionInfoExtra;
-use crate::rate_limiter::{EndpointRateLimiter, LeakyBucketConfig};
 use crate::types::{EndpointCacheKey, EndpointId, RoleName};
 use crate::{compute, scram};
 
@@ -59,8 +56,6 @@ pub(crate) struct AuthInfo {
     pub(crate) account_id: Option<AccountIdInt>,
     /// Are public connections or VPC connections blocked?
     pub(crate) access_blocker_flags: AccessBlockerFlags,
-    /// The rate limits for this endpoint.
-    pub(crate) rate_limits: EndpointRateLimitConfig,
 }
 
 /// Info for establishing a connection to a compute node.
@@ -78,8 +73,11 @@ impl NodeInfo {
         ctx: &RequestContext,
         auth: &compute::AuthInfo,
         config: &ComputeConfig,
+        user_info: ComputeUserInfo,
     ) -> Result<compute::PostgresConnection, compute::ConnectionError> {
-        self.conn_info.connect(ctx, &self.aux, auth, config).await
+        self.conn_info
+            .connect(ctx, self.aux.clone(), auth, config, user_info)
+            .await
     }
 }
 
@@ -103,8 +101,6 @@ pub struct EndpointAccessControl {
     pub allowed_ips: Arc<Vec<IpPattern>>,
     pub allowed_vpce: Arc<Vec<String>>,
     pub flags: AccessBlockerFlags,
-
-    pub rate_limits: EndpointRateLimitConfig,
 }
 
 impl EndpointAccessControl {
@@ -143,36 +139,6 @@ impl EndpointAccessControl {
 
         Ok(())
     }
-
-    pub fn connection_attempt_rate_limit(
-        &self,
-        ctx: &RequestContext,
-        endpoint: &EndpointId,
-        rate_limiter: &EndpointRateLimiter,
-    ) -> Result<(), AuthError> {
-        let endpoint = EndpointIdInt::from(endpoint);
-
-        let limits = &self.rate_limits.connection_attempts;
-        let config = match ctx.protocol() {
-            crate::metrics::Protocol::Http => limits.http,
-            crate::metrics::Protocol::Ws => limits.ws,
-            crate::metrics::Protocol::Tcp => limits.tcp,
-            crate::metrics::Protocol::SniRouter => return Ok(()),
-        };
-        let config = config.and_then(|config| {
-            if config.rps <= 0.0 || config.burst <= 0.0 {
-                return None;
-            }
-
-            Some(LeakyBucketConfig::new(config.rps, config.burst))
-        });
-
-        if !rate_limiter.check(endpoint, config, 1) {
-            return Err(AuthError::too_many_connections());
-        }
-
-        Ok(())
-    }
 }
 
 /// This will allocate per each call, but the http requests alone

proxy/src/lib.rs

@@ -75,7 +75,6 @@
 pub mod binary;
 
 mod auth;
-mod batch;
 mod cache;
 mod cancellation;
 mod compute;
@@ -107,5 +106,4 @@ mod tls;
 mod types;
 mod url;
 mod usage_metrics;
-mod util;
 mod waiters;

proxy/src/pglb/connect_compute.rs

@@ -2,24 +2,25 @@ 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 crate::config::{ComputeConfig, RetryConfig};
 use crate::context::RequestContext;
 use crate::control_plane::errors::WakeComputeError;
 use crate::control_plane::locks::ApiLocks;
-use crate::control_plane::{self, NodeInfo};
+use crate::control_plane::{self, CachedNodeInfo, NodeInfo};
 use crate::error::ReportableError;
 use crate::metrics::{
     ConnectOutcome, ConnectionFailureKind, Metrics, RetriesMetricGroup, RetryType,
 };
 use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute, retry_after, should_retry};
-use crate::proxy::wake_compute::{WakeComputeBackend, wake_compute};
+use crate::proxy::wake_compute::wake_compute;
 use crate::types::Host;
 
 /// If we couldn't connect, a cached connection info might be to blame
 /// (e.g. the compute node's address might've changed at the wrong time).
 /// Invalidate the cache entry (if any) to prevent subsequent errors.
-#[tracing::instrument(skip_all)]
+#[tracing::instrument(name = "invalidate_cache", skip_all)]
 pub(crate) fn invalidate_cache(node_info: control_plane::CachedNodeInfo) -> NodeInfo {
     let is_cached = node_info.cached();
     if is_cached {
@@ -48,10 +49,19 @@ pub(crate) trait ConnectMechanism {
     ) -> Result<Self::Connection, Self::ConnectError>;
 }
 
+#[async_trait]
+pub(crate) trait ComputeConnectBackend {
+    async fn wake_compute(
+        &self,
+        ctx: &RequestContext,
+    ) -> Result<CachedNodeInfo, control_plane::errors::WakeComputeError>;
+}
+
 pub(crate) struct TcpMechanism {
     pub(crate) auth: AuthInfo,
     /// connect_to_compute concurrency lock
     pub(crate) locks: &'static ApiLocks<Host>,
+    pub(crate) user_info: ComputeUserInfo,
 }
 
 #[async_trait]
@@ -71,13 +81,17 @@ impl ConnectMechanism for TcpMechanism {
         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).await)
+        permit.release_result(
+            node_info
+                .connect(ctx, &self.auth, config, self.user_info.clone())
+                .await,
+        )
     }
 }
 
 /// Try to connect to the compute node, retrying if necessary.
 #[tracing::instrument(skip_all)]
-pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: WakeComputeBackend>(
+pub(crate) async fn connect_to_compute<M: ConnectMechanism, B: ComputeConnectBackend>(
     ctx: &RequestContext,
     mechanism: &M,
     user_info: &B,

proxy/src/pglb/handshake.rs

@@ -12,7 +12,7 @@ use crate::pqproto::{
     BeMessage, CancelKeyData, FeStartupPacket, ProtocolVersion, StartupMessageParams,
 };
 use crate::proxy::TlsRequired;
-use crate::stream::{PqFeStream, Stream, StreamUpgradeError};
+use crate::stream::{PqStream, Stream, StreamUpgradeError};
 use crate::tls::PG_ALPN_PROTOCOL;
 
 #[derive(Error, Debug)]
@@ -49,7 +49,7 @@ impl ReportableError for HandshakeError {
 }
 
 pub(crate) enum HandshakeData<S> {
-    Startup(PqFeStream<Stream<S>>, StartupMessageParams),
+    Startup(PqStream<Stream<S>>, StartupMessageParams),
     Cancel(CancelKeyData),
 }
 
@@ -70,7 +70,7 @@ pub(crate) async fn handshake<S: AsyncRead + AsyncWrite + Unpin + Send>(
     const PG_PROTOCOL_EARLIEST: ProtocolVersion = ProtocolVersion::new(3, 0);
     const PG_PROTOCOL_LATEST: ProtocolVersion = ProtocolVersion::new(3, 0);
 
-    let (mut stream, mut msg) = PqFeStream::parse_startup(Stream::from_raw(stream)).await?;
+    let (mut stream, mut msg) = PqStream::parse_startup(Stream::from_raw(stream)).await?;
     loop {
         match msg {
             FeStartupPacket::SslRequest { direct } => match stream.get_ref() {
@@ -152,7 +152,7 @@ pub(crate) async fn handshake<S: AsyncRead + AsyncWrite + Unpin + Send>(
                             tls: tls_stream,
                             tls_server_end_point,
                         };
-                        (stream, msg) = PqFeStream::parse_startup(tls).await?;
+                        (stream, msg) = PqStream::parse_startup(tls).await?;
                     } else {
                         if direct.is_some() {
                             // client sent us a ClientHello already, we can't do anything with it.

proxy/src/pglb/mod.rs

@@ -1,3 +1,4 @@
+pub mod connect_compute;
 pub mod copy_bidirectional;
 pub mod handshake;
 pub mod inprocess;

proxy/src/pglb/passthrough.rs

@@ -1,18 +1,15 @@
-use std::convert::Infallible;
-
+use futures::FutureExt;
 use smol_str::SmolStr;
 use tokio::io::{AsyncRead, AsyncWrite};
-use tokio::net::TcpStream;
 use tracing::debug;
 use utils::measured_stream::MeasuredStream;
 
 use super::copy_bidirectional::ErrorSource;
-use crate::compute::MaybeRustlsStream;
+use crate::cancellation;
+use crate::compute::PostgresConnection;
+use crate::config::ComputeConfig;
 use crate::control_plane::messages::MetricsAuxInfo;
-use crate::metrics::{
-    Direction, Metrics, NumClientConnectionsGuard, NumConnectionRequestsGuard,
-    NumDbConnectionsGuard,
-};
+use crate::metrics::{Direction, Metrics, NumClientConnectionsGuard, NumConnectionRequestsGuard};
 use crate::stream::Stream;
 use crate::usage_metrics::{Ids, MetricCounterRecorder, USAGE_METRICS};
 
@@ -67,19 +64,40 @@ pub(crate) async fn proxy_pass(
 
 pub(crate) struct ProxyPassthrough<S> {
     pub(crate) client: Stream<S>,
-    pub(crate) compute: MaybeRustlsStream<TcpStream>,
+    pub(crate) compute: PostgresConnection,
     pub(crate) aux: MetricsAuxInfo,
+    pub(crate) session_id: uuid::Uuid,
     pub(crate) private_link_id: Option<SmolStr>,
-
-    pub(crate) _cancel_on_shutdown: tokio::sync::oneshot::Sender<Infallible>,
+    pub(crate) cancel: cancellation::Session,
 
     pub(crate) _req: NumConnectionRequestsGuard<'static>,
     pub(crate) _conn: NumClientConnectionsGuard<'static>,
-    pub(crate) _db_conn: NumDbConnectionsGuard<'static>,
 }
 
 impl<S: AsyncRead + AsyncWrite + Unpin> ProxyPassthrough<S> {
-    pub(crate) async fn proxy_pass(self) -> Result<(), ErrorSource> {
-        proxy_pass(self.client, self.compute, self.aux, self.private_link_id).await
+    pub(crate) async fn proxy_pass(
+        self,
+        compute_config: &ComputeConfig,
+    ) -> Result<(), ErrorSource> {
+        let res = proxy_pass(
+            self.client,
+            self.compute.stream,
+            self.aux,
+            self.private_link_id,
+        )
+        .await;
+        if let Err(err) = self
+            .compute
+            .cancel_closure
+            .try_cancel_query(compute_config)
+            .boxed()
+            .await
+        {
+            tracing::warn!(session_id = ?self.session_id, ?err, "could not cancel the query in the database");
+        }
+
+        drop(self.cancel.remove_cancel_key()); // we don't need a result. If the queue is full, we just log the error
+
+        res
     }
 }

proxy/src/pqproto.rs

@@ -11,86 +11,11 @@ use rand::distributions::{Distribution, Standard};
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use zerocopy::{FromBytes, Immutable, IntoBytes, big_endian};
 
-#[derive(Debug, Copy, Clone, PartialEq)]
-pub struct ErrorCode(pub [u8; 5]);
+pub type ErrorCode = [u8; 5];
 
-#[derive(Debug, Copy, Clone, PartialEq)]
-pub struct FeTag(pub u8);
+pub const FE_PASSWORD_MESSAGE: u8 = b'p';
 
-#[derive(Debug, Copy, Clone, PartialEq)]
-pub struct BeTag(pub u8);
-
-#[derive(Debug, Copy, Clone, PartialEq)]
-pub struct AuthTag(pub i32);
-
-pub const FE_PASSWORD_MESSAGE: FeTag = FeTag(b'p');
-
-pub const BE_AUTH_MESSAGE: BeTag = BeTag(b'R');
-pub const BE_ERR_MESSAGE: BeTag = BeTag(b'E');
-pub const BE_KEY_MESSAGE: BeTag = BeTag(b'K');
-pub const BE_READY_MESSAGE: BeTag = BeTag(b'Z');
-pub const BE_NEGOTIATE_MESSAGE: BeTag = BeTag(b'v');
-
-pub const AUTH_OK: AuthTag = AuthTag(0);
-pub const AUTH_SASL: AuthTag = AuthTag(10);
-pub const AUTH_SASL_CONT: AuthTag = AuthTag(11);
-pub const AUTH_SASL_FINAL: AuthTag = AuthTag(12);
-
-pub const SQLSTATE_INTERNAL_ERROR: ErrorCode = ErrorCode(*b"XX000");
-pub const CONNECTION_EXCEPTION: ErrorCode = ErrorCode(*b"08000");
-pub const CONNECTION_DOES_NOT_EXIST: ErrorCode = ErrorCode(*b"08003");
-pub const CONNECTION_FAILURE: ErrorCode = ErrorCode(*b"08006");
-pub const SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION: ErrorCode = ErrorCode(*b"08001");
-pub const PROTOCOL_VIOLATION: ErrorCode = ErrorCode(*b"08P01");
-pub const INVALID_PARAMETER_VALUE: ErrorCode = ErrorCode(*b"22023");
-pub const INVALID_CATALOG_NAME: ErrorCode = ErrorCode(*b"3D000");
-pub const INVALID_SCHEMA_NAME: ErrorCode = ErrorCode(*b"3F000");
-pub const T_R_SERIALIZATION_FAILURE: ErrorCode = ErrorCode(*b"40001");
-pub const SYNTAX_ERROR: ErrorCode = ErrorCode(*b"42601");
-pub const OUT_OF_MEMORY: ErrorCode = ErrorCode(*b"53200");
-pub const TOO_MANY_CONNECTIONS: ErrorCode = ErrorCode(*b"53300");
-
-impl fmt::Display for AuthTag {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match self.0 {
-            0 => f.write_str("Ok"),
-            2 => f.write_str("KerberosV5"),
-            3 => f.write_str("CleartextPassword"),
-            5 => f.write_str("MD5Password"),
-            7 => f.write_str("GSS"),
-            8 => f.write_str("GSSContinue"),
-            9 => f.write_str("SSPI"),
-            10 => f.write_str("SASL"),
-            11 => f.write_str("SASLContinue"),
-            12 => f.write_str("SASLFinal"),
-            x => write!(f, "{x}"),
-        }
-    }
-}
-
-impl fmt::Display for BeTag {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match *self {
-            BE_AUTH_MESSAGE => f.write_str("Authentication"),
-            BE_KEY_MESSAGE => f.write_str("BackendKeyData"),
-            BE_ERR_MESSAGE => f.write_str("ErrorResponse"),
-            BE_READY_MESSAGE => f.write_str("ReadyForQuery"),
-            BE_NEGOTIATE_MESSAGE => f.write_str("NegotiateProtocolVersion"),
-            BeTag(b'S') => f.write_str("ParameterStatus"),
-            BeTag(b'N') => f.write_str("NoticeMessage"),
-            BeTag(x) => write!(f, "{:?}", char::from(x)),
-        }
-    }
-}
-
-impl fmt::Display for FeTag {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match *self {
-            FE_PASSWORD_MESSAGE => f.write_str("Password"),
-            FeTag(x) => write!(f, "{:?}", char::from(x)),
-        }
-    }
-}
+pub const SQLSTATE_INTERNAL_ERROR: [u8; 5] = *b"XX000";
 
 /// The protocol version number.
 ///
@@ -355,10 +280,6 @@ impl WriteBuf {
         Self(Cursor::new(Vec::new()))
     }
 
-    pub const fn len(&self) -> usize {
-        self.0.get_ref().len()
-    }
-
     /// Use a heuristic to determine if we should shrink the write buffer.
     #[inline]
     fn should_shrink(&self) -> bool {
@@ -392,19 +313,6 @@ impl WriteBuf {
         self.0.set_position(0);
     }
 
-    /// Write a startup message.
-    pub fn startup(&mut self, params: &StartupMessageParams) {
-        self.0.get_mut().extend_from_slice(
-            StartupHeader {
-                len: big_endian::U32::new(params.params.len() as u32 + 9),
-                version: ProtocolVersion::new(3, 0),
-            }
-            .as_bytes(),
-        );
-        self.0.get_mut().extend_from_slice(params.params.as_bytes());
-        self.0.get_mut().push(0);
-    }
-
     /// Write a raw message to the internal buffer.
     ///
     /// The size_hint value is only a hint for reserving space. It's ok if it's incorrect, since
@@ -445,7 +353,7 @@ impl WriteBuf {
 
             // Code: error_code
             buf.put_u8(b'C');
-            buf.put_slice(&error_code.0);
+            buf.put_slice(&error_code);
             buf.put_u8(0);
 
             // Message: msg
@@ -560,11 +468,11 @@ pub enum BeMessage<'a> {
     AuthenticationOk,
     AuthenticationSasl(BeAuthenticationSaslMessage<'a>),
     AuthenticationCleartextPassword,
+    BackendKeyData(CancelKeyData),
     ParameterStatus {
         name: &'a [u8],
         value: &'a [u8],
     },
-    #[cfg(test)]
     ReadyForQuery,
     NoticeResponse(&'a str),
     NegotiateProtocolVersion {
@@ -586,17 +494,17 @@ impl BeMessage<'_> {
         match self {
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-AUTHENTICATIONCLEARTEXTPASSWORD>
             BeMessage::AuthenticationOk => {
-                buf.write_raw(1, BE_AUTH_MESSAGE.0, |buf| buf.put_i32(0));
+                buf.write_raw(1, b'R', |buf| buf.put_i32(0));
             }
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-AUTHENTICATIONCLEARTEXTPASSWORD>
             BeMessage::AuthenticationCleartextPassword => {
-                buf.write_raw(1, BE_AUTH_MESSAGE.0, |buf| buf.put_i32(3));
+                buf.write_raw(1, b'R', |buf| buf.put_i32(3));
             }
 
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-AUTHENTICATIONSASL>
             BeMessage::AuthenticationSasl(BeAuthenticationSaslMessage::Methods(methods)) => {
                 let len: usize = methods.iter().map(|m| m.len() + 1).sum();
-                buf.write_raw(len + 2, BE_AUTH_MESSAGE.0, |buf| {
+                buf.write_raw(len + 2, b'R', |buf| {
                     buf.put_i32(10); // Specifies that SASL auth method is used.
                     for method in methods {
                         buf.put_slice(method.as_bytes());
@@ -607,19 +515,24 @@ impl BeMessage<'_> {
             }
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-AUTHENTICATIONSASL>
             BeMessage::AuthenticationSasl(BeAuthenticationSaslMessage::Continue(extra)) => {
-                buf.write_raw(extra.len() + 1, BE_AUTH_MESSAGE.0, |buf| {
+                buf.write_raw(extra.len() + 1, b'R', |buf| {
                     buf.put_i32(11); // Continue SASL auth.
                     buf.put_slice(extra);
                 });
             }
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-AUTHENTICATIONSASL>
             BeMessage::AuthenticationSasl(BeAuthenticationSaslMessage::Final(extra)) => {
-                buf.write_raw(extra.len() + 1, BE_AUTH_MESSAGE.0, |buf| {
+                buf.write_raw(extra.len() + 1, b'R', |buf| {
                     buf.put_i32(12); // Send final SASL message.
                     buf.put_slice(extra);
                 });
             }
 
+            // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-BACKENDKEYDATA>
+            BeMessage::BackendKeyData(key_data) => {
+                buf.write_raw(8, b'K', |buf| buf.put_slice(key_data.as_bytes()));
+            }
+
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-NOTICERESPONSE>
             // <https://www.postgresql.org/docs/current/protocol-error-fields.html>
             BeMessage::NoticeResponse(msg) => {
@@ -651,16 +564,15 @@ impl BeMessage<'_> {
                 });
             }
 
-            #[cfg(test)]
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-NEGOTIATEPROTOCOLVERSION>
             BeMessage::ReadyForQuery => {
-                buf.write_raw(1, BE_READY_MESSAGE.0, |buf| buf.put_u8(b'I'));
+                buf.write_raw(1, b'Z', |buf| buf.put_u8(b'I'));
             }
 
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-NEGOTIATEPROTOCOLVERSION>
             BeMessage::NegotiateProtocolVersion { version, options } => {
                 let len: usize = options.iter().map(|o| o.len() + 1).sum();
-                buf.write_raw(8 + len, BE_NEGOTIATE_MESSAGE.0, |buf| {
+                buf.write_raw(8 + len, b'v', |buf| {
                     buf.put_slice(version.as_bytes());
                     buf.put_u32(options.len() as u32);
                     for option in options {

proxy/src/proxy/mod.rs

@@ -1,16 +1,13 @@
 #[cfg(test)]
 mod tests;
 
-pub(crate) mod connect_compute;
 pub(crate) mod retry;
 pub(crate) mod wake_compute;
-
 use std::sync::Arc;
 
 use futures::FutureExt;
 use itertools::Itertools;
 use once_cell::sync::OnceCell;
-use postgres_client::RawCancelToken;
 use regex::Regex;
 use serde::{Deserialize, Serialize};
 use smol_str::{SmolStr, ToSmolStr, format_smolstr};
@@ -19,22 +16,20 @@ use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, debug, error, info, warn};
 
-use crate::auth::backend::ComputeUserInfo;
-use crate::cancellation::{self, CancelClosure, CancellationHandler};
+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::{CancelKeyData, StartupMessageParams};
+use crate::pqproto::{BeMessage, CancelKeyData, StartupMessageParams};
 use crate::protocol2::{ConnectHeader, ConnectionInfo, ConnectionInfoExtra, read_proxy_protocol};
-use crate::proxy::connect_compute::{TcpMechanism, connect_to_compute};
 use crate::rate_limiter::EndpointRateLimiter;
-use crate::stream::{PostgresError, PqFeStream, Stream};
+use crate::stream::{PqStream, Stream};
 use crate::types::EndpointCacheKey;
-use crate::util::run_until_cancelled;
 use crate::{auth, compute};
 
 const ERR_INSECURE_CONNECTION: &str = "connection is insecure (try using `sslmode=require`)";
@@ -51,6 +46,21 @@ impl ReportableError for TlsRequired {
 
 impl UserFacingError for TlsRequired {}
 
+pub async fn run_until_cancelled<F: std::future::Future>(
+    f: F,
+    cancellation_token: &CancellationToken,
+) -> Option<F::Output> {
+    match futures::future::select(
+        std::pin::pin!(f),
+        std::pin::pin!(cancellation_token.cancelled()),
+    )
+    .await
+    {
+        futures::future::Either::Left((f, _)) => Some(f),
+        futures::future::Either::Right(((), _)) => None,
+    }
+}
+
 pub async fn task_main(
     config: &'static ProxyConfig,
     auth_backend: &'static auth::Backend<'static, ()>,
@@ -157,7 +167,7 @@ pub async fn task_main(
                 Ok(Some(p)) => {
                     ctx.set_success();
                     let _disconnect = ctx.log_connect();
-                    match p.proxy_pass().await {
+                    match p.proxy_pass(&config.connect_to_compute).await {
                         Ok(()) => {}
                         Err(ErrorSource::Client(e)) => {
                             warn!(
@@ -255,7 +265,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     auth_backend: &'static auth::Backend<'static, ()>,
     ctx: &RequestContext,
     cancellation_handler: Arc<CancellationHandler>,
-    client: S,
+    stream: S,
     mode: ClientMode,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     conn_gauge: NumClientConnectionsGuard<'static>,
@@ -275,9 +285,9 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
 
     let record_handshake_error = !ctx.has_private_peer_addr();
     let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Client);
-    let do_handshake = handshake(ctx, client, mode.handshake_tls(tls), record_handshake_error);
+    let do_handshake = handshake(ctx, stream, mode.handshake_tls(tls), record_handshake_error);
 
-    let (mut client, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
+    let (mut stream, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
         .await??
     {
         HandshakeData::Startup(stream, params) => (stream, params),
@@ -309,7 +319,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
 
     ctx.set_db_options(params.clone());
 
-    let hostname = mode.hostname(client.get_ref());
+    let hostname = mode.hostname(stream.get_ref());
 
     let common_names = tls.map(|tls| &tls.common_names);
 
@@ -321,14 +331,14 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
 
     let user_info = match result {
         Ok(user_info) => user_info,
-        Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
+        Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
     };
 
     let user = user_info.get_user().to_owned();
     let user_info = match user_info
         .authenticate(
             ctx,
-            &mut client,
+            &mut stream,
             mode.allow_cleartext(),
             &config.authentication_config,
             endpoint_rate_limiter,
@@ -341,58 +351,45 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
             let app = params.get("application_name");
             let params_span = tracing::info_span!("", ?user, ?db, ?app);
 
-            return Err(client
+            return Err(stream
                 .throw_error(e, Some(ctx))
                 .instrument(params_span)
                 .await)?;
         }
     };
 
-    let (cplane, creds) = match user_info {
-        auth::Backend::ControlPlane(cplane, creds) => (cplane, creds),
+    let creds = match &user_info {
+        auth::Backend::ControlPlane(_, creds) => creds,
         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);
+    let mut auth_info = compute::AuthInfo::with_auth_keys(&creds.keys);
     auth_info.set_startup_params(&params, params_compat);
 
     let res = connect_to_compute(
         ctx,
         &TcpMechanism {
+            user_info: creds.info.clone(),
             auth: auth_info,
             locks: &config.connect_compute_locks,
         },
-        &auth::Backend::ControlPlane(cplane, creds.info.clone()),
+        &user_info,
         config.wake_compute_retry_config,
         &config.connect_to_compute,
     )
     .await;
 
-    let mut node = match res {
+    let node = match res {
         Ok(node) => node,
-        Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
+        Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
     };
 
-    let session = cancellation_handler.get_key();
+    let cancellation_handler_clone = Arc::clone(&cancellation_handler);
+    let session = cancellation_handler_clone.get_key();
 
-    let cancel_closure =
-        prepare_client_connection(&mut node, session.key(), &mut client, creds.info).await?;
-
-    let session_id = ctx.session_id();
-    let (cancel_on_shutdown, cancel) = tokio::sync::oneshot::channel();
-    tokio::spawn(async move {
-        session
-            .maintain_cancel_key(
-                session_id,
-                cancel,
-                &cancel_closure,
-                &config.connect_to_compute,
-            )
-            .await;
-    });
-
-    let client = client.flush_and_into_inner().await?;
-    let compute = node.stream.flush_and_into_inner().await?;
+    session.write_cancel_key(node.cancel_closure.clone())?;
+    prepare_client_connection(&node, *session.key(), &mut stream);
+    let stream = stream.flush_and_into_inner().await?;
 
     let private_link_id = match ctx.extra() {
         Some(ConnectionInfoExtra::Aws { vpce_id }) => Some(vpce_id.clone()),
@@ -401,75 +398,40 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     };
 
     Ok(Some(ProxyPassthrough {
-        client,
-        compute,
-
-        aux: node.aux,
+        client: stream,
+        aux: node.aux.clone(),
         private_link_id,
-
-        _cancel_on_shutdown: cancel_on_shutdown,
-
+        compute: node,
+        session_id: ctx.session_id(),
+        cancel: session,
         _req: request_gauge,
         _conn: conn_gauge,
-        _db_conn: node.guage,
     }))
 }
 
 /// Finish client connection initialization: confirm auth success, send params, etc.
-pub(crate) async fn prepare_client_connection(
-    node: &mut compute::PostgresConnection,
-    key_data: &CancelKeyData,
-    stream: &mut PqFeStream<impl AsyncRead + AsyncWrite + Unpin>,
-    user_info: ComputeUserInfo,
-) -> Result<CancelClosure, std::io::Error> {
-    use zerocopy::{FromBytes, IntoBytes};
-
-    use crate::pqproto::{BE_KEY_MESSAGE, BE_READY_MESSAGE};
-
-    let mut process_id = 0;
-    let mut secret_key = 0;
-
-    loop {
-        match node.stream.read_raw_be(1024).await {
-            // parse backend keys, and substitute our own.
-            Ok((tag @ BE_KEY_MESSAGE, msg)) => {
-                stream.write_raw(8, tag, |b| b.extend_from_slice(key_data.as_bytes()));
-
-                let key_data = CancelKeyData::read_from_bytes(msg)
-                    .map_err(|_| std::io::Error::other("invalid msg len"))?;
-
-                process_id = (key_data.0.get() >> 32) as i32;
-                secret_key = (key_data.0.get() & 0xffff_ffff) as i32;
-            }
-            // ready for query, we're done :)
-            Ok((tag @ BE_READY_MESSAGE, msg)) => {
-                stream.write_raw(msg.len(), tag, |b| b.extend_from_slice(msg.as_bytes()));
-                break;
-            }
-            // either a notice or a parameter status.
-            Ok((tag, msg)) => {
-                stream.write_raw(msg.len(), tag, |b| b.extend_from_slice(msg.as_bytes()));
-            }
-            Err(PostgresError::Io(io)) => return Err(io),
-            Err(PostgresError::Error(e)) => return Err(std::io::Error::other(e)),
-            Err(_) => unreachable!("read_raw_be only returns IO or BackendError types"),
-        }
-
-        if stream.write_buf_len() > 512 {
-            stream.flush().await?;
-        }
+pub(crate) fn prepare_client_connection(
+    node: &compute::PostgresConnection,
+    cancel_key_data: CancelKeyData,
+    stream: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+) {
+    // Forward all deferred notices to the client.
+    for notice in &node.delayed_notice {
+        stream.write_raw(notice.as_bytes().len(), b'N', |buf| {
+            buf.extend_from_slice(notice.as_bytes());
+        });
     }
 
-    Ok(CancelClosure::new(
-        node.socket_addr,
-        RawCancelToken {
-            ssl_mode: node.ssl_mode,
-            process_id,
-            secret_key,
-        },
-        node.hostname.clone(),
-        user_info,
-    ))
+    // Forward all postgres connection params to the client.
+    for (name, value) in &node.params {
+        stream.write_message(BeMessage::ParameterStatus {
+            name: name.as_bytes(),
+            value: value.as_bytes(),
+        });
+    }
+
+    stream.write_message(BeMessage::BackendKeyData(cancel_key_data));
+    stream.write_message(BeMessage::ReadyForQuery);
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]

proxy/src/proxy/retry.rs

@@ -1,12 +1,10 @@
 use std::error::Error;
 use std::io;
 
-use bstr::ByteSlice;
 use tokio::time;
 
 use crate::compute;
 use crate::config::RetryConfig;
-use crate::stream::{BackendError, PostgresError};
 
 pub(crate) trait CouldRetry {
     /// Returns true if the error could be retried
@@ -98,55 +96,10 @@ impl ShouldRetryWakeCompute for postgres_client::Error {
     }
 }
 
-impl CouldRetry for BackendError {
-    fn could_retry(&self) -> bool {
-        let (code, _message) = self.parse();
-        matches!(
-            code,
-            crate::pqproto::CONNECTION_FAILURE
-                | crate::pqproto::CONNECTION_EXCEPTION
-                | crate::pqproto::CONNECTION_DOES_NOT_EXIST
-                | crate::pqproto::SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION,
-        )
-    }
-}
-
-impl ShouldRetryWakeCompute for BackendError {
-    fn should_retry_wake_compute(&self) -> bool {
-        let (code, message) = self.parse();
-
-        // Here are errors that happens after the user successfully authenticated to the database.
-        let non_retriable_pg_errors = matches!(
-            code,
-            crate::pqproto::TOO_MANY_CONNECTIONS
-                | crate::pqproto::OUT_OF_MEMORY
-                | crate::pqproto::SYNTAX_ERROR
-                | crate::pqproto::T_R_SERIALIZATION_FAILURE
-                | crate::pqproto::INVALID_CATALOG_NAME
-                | crate::pqproto::INVALID_SCHEMA_NAME
-                | crate::pqproto::INVALID_PARAMETER_VALUE,
-        );
-        if non_retriable_pg_errors {
-            return false;
-        }
-
-        // PGBouncer errors that should not trigger a wake_compute retry.
-        if code == crate::pqproto::PROTOCOL_VIOLATION {
-            // Source for the error message:
-            // https://github.com/pgbouncer/pgbouncer/blob/f15997fe3effe3a94ba8bcc1ea562e6117d1a131/src/client.c#L1070
-            return message.contains_str("no more connections allowed (max_client_conn)");
-        }
-        true
-    }
-}
-
 impl CouldRetry for compute::ConnectionError {
     fn could_retry(&self) -> bool {
         match self {
-            compute::ConnectionError::Postgres(PostgresError::Error(err)) => err.could_retry(),
-            compute::ConnectionError::Postgres(PostgresError::Io(err)) => err.could_retry(),
-            compute::ConnectionError::Postgres(PostgresError::Unexpected(_)) => false,
-            compute::ConnectionError::Postgres(PostgresError::InvalidAuthMessage) => false,
+            compute::ConnectionError::Postgres(err) => err.could_retry(),
             compute::ConnectionError::TlsError(err) => err.could_retry(),
             compute::ConnectionError::WakeComputeError(err) => err.could_retry(),
             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
@@ -156,12 +109,7 @@ impl CouldRetry for compute::ConnectionError {
 impl ShouldRetryWakeCompute for compute::ConnectionError {
     fn should_retry_wake_compute(&self) -> bool {
         match self {
-            compute::ConnectionError::Postgres(PostgresError::Error(err)) => {
-                err.should_retry_wake_compute()
-            }
-            compute::ConnectionError::Postgres(PostgresError::Io(_)) => true,
-            compute::ConnectionError::Postgres(PostgresError::Unexpected(_)) => false,
-            compute::ConnectionError::Postgres(PostgresError::InvalidAuthMessage) => false,
+            compute::ConnectionError::Postgres(err) => err.should_retry_wake_compute(),
             // the cache entry was not checked for validity
             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
             _ => true,

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::SslMode;
+use postgres_client::config::{AuthKeys, ScramKeys, SslMode};
 use postgres_client::tls::{MakeTlsConnect, NoTls};
 use retry::{ShouldRetryWakeCompute, retry_after};
 use rstest::rstest;
@@ -19,14 +19,15 @@ use tracing_test::traced_test;
 
 use super::retry::CouldRetry;
 use super::*;
-use crate::auth::backend::{ComputeUserInfo, MaybeOwned};
+use crate::auth::backend::{
+    ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, MaybeOwned,
+};
 use crate::config::{ComputeConfig, RetryConfig};
 use crate::control_plane::client::{ControlPlaneClient, TestControlPlaneClient};
 use crate::control_plane::messages::{ControlPlaneErrorMessage, Details, MetricsAuxInfo, Status};
 use crate::control_plane::{self, CachedNodeInfo, NodeInfo, NodeInfoCache};
 use crate::error::ErrorKind;
-use crate::pqproto::BeMessage;
-use crate::proxy::connect_compute::ConnectMechanism;
+use crate::pglb::connect_compute::ConnectMechanism;
 use crate::tls::client_config::compute_client_config_with_certs;
 use crate::tls::server_config::CertResolver;
 use crate::types::{BranchId, EndpointId, ProjectId};
@@ -123,7 +124,7 @@ fn generate_tls_config<'a>(
 trait TestAuth: Sized {
     async fn authenticate<S: AsyncRead + AsyncWrite + Unpin + Send>(
         self,
-        stream: &mut PqFeStream<Stream<S>>,
+        stream: &mut PqStream<Stream<S>>,
     ) -> anyhow::Result<()> {
         stream.write_message(BeMessage::AuthenticationOk);
         Ok(())
@@ -152,7 +153,7 @@ impl Scram {
 impl TestAuth for Scram {
     async fn authenticate<S: AsyncRead + AsyncWrite + Unpin + Send>(
         self,
-        stream: &mut PqFeStream<Stream<S>>,
+        stream: &mut PqStream<Stream<S>>,
     ) -> anyhow::Result<()> {
         let outcome = auth::AuthFlow::new(stream, auth::Scram(&self.0, &RequestContext::test()))
             .authenticate()
@@ -574,13 +575,19 @@ fn helper_create_cached_node_info(cache: &'static NodeInfoCache) -> CachedNodeIn
 
 fn helper_create_connect_info(
     mechanism: &TestConnectMechanism,
-) -> auth::Backend<'static, ComputeUserInfo> {
+) -> auth::Backend<'static, ComputeCredentials> {
     auth::Backend::ControlPlane(
         MaybeOwned::Owned(ControlPlaneClient::Test(Box::new(mechanism.clone()))),
-        ComputeUserInfo {
-            endpoint: "endpoint".into(),
-            user: "user".into(),
-            options: NeonOptions::parse_options_raw(""),
+        ComputeCredentials {
+            info: ComputeUserInfo {
+                endpoint: "endpoint".into(),
+                user: "user".into(),
+                options: NeonOptions::parse_options_raw(""),
+            },
+            keys: ComputeCredentialKeys::AuthKeys(AuthKeys::ScramSha256(ScramKeys {
+                client_key: [0; 32],
+                server_key: [0; 32],
+            })),
         },
     )
 }

proxy/src/proxy/wake_compute.rs

@@ -1,4 +1,3 @@
-use async_trait::async_trait;
 use tracing::{error, info};
 
 use crate::config::RetryConfig;
@@ -9,6 +8,7 @@ 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.
@@ -23,12 +23,7 @@ macro_rules! log_wake_compute_error {
     };
 }
 
-#[async_trait]
-pub(crate) trait WakeComputeBackend {
-    async fn wake_compute(&self, ctx: &RequestContext) -> Result<CachedNodeInfo, WakeComputeError>;
-}
-
-pub(crate) async fn wake_compute<B: WakeComputeBackend>(
+pub(crate) async fn wake_compute<B: ComputeConnectBackend>(
     num_retries: &mut u32,
     ctx: &RequestContext,
     api: &B,

proxy/src/rate_limiter/leaky_bucket.rs

@@ -69,8 +69,9 @@ pub struct LeakyBucketConfig {
     pub max: f64,
 }
 
+#[cfg(test)]
 impl LeakyBucketConfig {
-    pub fn new(rps: f64, max: f64) -> Self {
+    pub(crate) fn new(rps: f64, max: f64) -> Self {
         assert!(rps > 0.0, "rps must be positive");
         assert!(max > 0.0, "max must be positive");
         Self { rps, max }

proxy/src/rate_limiter/limiter.rs

@@ -12,10 +12,11 @@ use rand::{Rng, SeedableRng};
 use tokio::time::{Duration, Instant};
 use tracing::info;
 
-use super::LeakyBucketConfig;
 use crate::ext::LockExt;
 use crate::intern::EndpointIdInt;
 
+use super::LeakyBucketConfig;
+
 pub struct GlobalRateLimiter {
     data: Vec<RateBucket>,
     info: Vec<RateBucketInfo>,

proxy/src/redis/keys.rs

@@ -1,4 +1,8 @@
-use crate::pqproto::CancelKeyData;
+use std::io::ErrorKind;
+
+use anyhow::Ok;
+
+use crate::pqproto::{CancelKeyData, id_to_cancel_key};
 
 pub mod keyspace {
     pub const CANCEL_PREFIX: &str = "cancel";
@@ -19,12 +23,39 @@ impl KeyPrefix {
             }
         }
     }
+
+    #[allow(dead_code)]
+    pub(crate) fn as_str(&self) -> &'static str {
+        match self {
+            KeyPrefix::Cancel(_) => keyspace::CANCEL_PREFIX,
+        }
+    }
+}
+
+#[allow(dead_code)]
+pub(crate) fn parse_redis_key(key: &str) -> anyhow::Result<KeyPrefix> {
+    let (prefix, key_str) = key.split_once(':').ok_or_else(|| {
+        anyhow::anyhow!(std::io::Error::new(
+            ErrorKind::InvalidData,
+            "missing prefix"
+        ))
+    })?;
+
+    match prefix {
+        keyspace::CANCEL_PREFIX => {
+            let id = u64::from_str_radix(key_str, 16)?;
+
+            Ok(KeyPrefix::Cancel(id_to_cancel_key(id)))
+        }
+        _ => Err(anyhow::anyhow!(std::io::Error::new(
+            ErrorKind::InvalidData,
+            "unknown prefix"
+        ))),
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::pqproto::id_to_cancel_key;
-
     use super::*;
 
     #[test]
@@ -34,4 +65,16 @@ mod tests {
         let redis_key = cancel_key.build_redis_key();
         assert_eq!(redis_key, "cancel:30390000d431");
     }
+
+    #[test]
+    fn test_parse_redis_key() {
+        let redis_key = "cancel:30390000d431";
+        let key: KeyPrefix = parse_redis_key(redis_key).expect("Failed to parse key");
+
+        let ref_key = id_to_cancel_key(12345 << 32 | 54321);
+
+        assert_eq!(key.as_str(), KeyPrefix::Cancel(ref_key).as_str());
+        let KeyPrefix::Cancel(cancel_key) = key;
+        assert_eq!(ref_key, cancel_key);
+    }
 }

proxy/src/redis/kv_ops.rs

@@ -1,6 +1,3 @@
-use std::time::Duration;
-
-use futures::FutureExt;
 use redis::aio::ConnectionLike;
 use redis::{Cmd, FromRedisValue, Pipeline, RedisResult};
 
@@ -38,11 +35,14 @@ impl RedisKVClient {
     }
 
     pub async fn try_connect(&mut self) -> anyhow::Result<()> {
-        self.client
-            .connect()
-            .boxed()
-            .await
-            .inspect_err(|e| tracing::error!("failed to connect to redis: {e}"))
+        match self.client.connect().await {
+            Ok(()) => {}
+            Err(e) => {
+                tracing::error!("failed to connect to redis: {e}");
+                return Err(e);
+            }
+        }
+        Ok(())
     }
 
     pub(crate) async fn query<T: FromRedisValue>(
@@ -54,25 +54,15 @@ impl RedisKVClient {
             return Err(anyhow::anyhow!("Rate limit exceeded"));
         }
 
-        let e = match q.query(&mut self.client).await {
+        match q.query(&mut self.client).await {
             Ok(t) => return Ok(t),
-            Err(e) => e,
-        };
-
-        tracing::error!("failed to run query: {e}");
-        match e.retry_method() {
-            redis::RetryMethod::Reconnect => {
-                tracing::info!("Redis client is disconnected. Reconnecting...");
-                self.try_connect().await?;
+            Err(e) => {
+                tracing::error!("failed to run query: {e}");
             }
-            redis::RetryMethod::RetryImmediately => {}
-            redis::RetryMethod::WaitAndRetry => {
-                // somewhat arbitrary.
-                tokio::time::sleep(Duration::from_millis(100)).await;
-            }
-            _ => Err(e)?,
         }
 
+        tracing::info!("Redis client is disconnected. Reconnecting...");
+        self.try_connect().await?;
         Ok(q.query(&mut self.client).await?)
     }
 }

proxy/src/sasl/stream.rs

@@ -7,7 +7,7 @@ use tokio::io::{AsyncRead, AsyncWrite};
 use super::{Mechanism, Step};
 use crate::context::RequestContext;
 use crate::pqproto::{BeAuthenticationSaslMessage, BeMessage};
-use crate::stream::PqFeStream;
+use crate::stream::PqStream;
 
 /// SASL authentication outcome.
 /// It's much easier to match on those two variants
@@ -22,7 +22,7 @@ pub(crate) enum Outcome<R> {
 
 pub async fn authenticate<S, F, M>(
     ctx: &RequestContext,
-    stream: &mut PqFeStream<S>,
+    stream: &mut PqStream<S>,
     mechanism: F,
 ) -> super::Result<Outcome<M::Output>>
 where

proxy/src/serverless/backend.rs

@@ -21,7 +21,7 @@ use super::conn_pool_lib::{Client, ConnInfo, EndpointConnPool, GlobalConnPool};
 use super::http_conn_pool::{self, HttpConnPool, Send, poll_http2_client};
 use super::local_conn_pool::{self, EXT_NAME, EXT_SCHEMA, EXT_VERSION, LocalConnPool};
 use crate::auth::backend::local::StaticAuthRules;
-use crate::auth::backend::{ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo};
+use crate::auth::backend::{ComputeCredentials, ComputeUserInfo};
 use crate::auth::{self, AuthError};
 use crate::compute_ctl::{
     ComputeCtlError, ExtensionInstallRequest, Privilege, SetRoleGrantsRequest,
@@ -34,7 +34,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::proxy::connect_compute::ConnectMechanism;
+use crate::pglb::connect_compute::ConnectMechanism;
 use crate::proxy::retry::{CouldRetry, ShouldRetryWakeCompute};
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::types::{EndpointId, Host, LOCAL_PROXY_SUFFIX};
@@ -68,20 +68,17 @@ impl PoolingBackend {
             self.config.authentication_config.is_vpc_acccess_proxy,
         )?;
 
-        access_control.connection_attempt_rate_limit(
-            ctx,
-            &user_info.endpoint,
-            &self.endpoint_rate_limiter,
-        )?;
-
+        let ep = EndpointIdInt::from(&user_info.endpoint);
+        let rate_limit_config = None;
+        if !self.endpoint_rate_limiter.check(ep, rate_limit_config, 1) {
+            return Err(AuthError::too_many_connections());
+        }
         let role_access = backend.get_role_secret(ctx).await?;
         let Some(secret) = role_access.secret else {
             // If we don't have an authentication secret, for the http flow we can just return an error.
             info!("authentication info not found");
             return Err(AuthError::password_failed(&*user_info.user));
         };
-
-        let ep = EndpointIdInt::from(&user_info.endpoint);
         let auth_outcome = crate::auth::validate_password_and_exchange(
             &self.config.authentication_config.thread_pool,
             ep,
@@ -183,15 +180,14 @@ impl PoolingBackend {
         let conn_id = uuid::Uuid::new_v4();
         tracing::Span::current().record("conn_id", display(conn_id));
         info!(%conn_id, "pool: opening a new connection '{conn_info}'");
-        let backend = self.auth_backend.as_ref().map(|()| keys.info);
-        crate::proxy::connect_compute::connect_to_compute(
+        let backend = self.auth_backend.as_ref().map(|()| keys);
+        crate::pglb::connect_compute::connect_to_compute(
             ctx,
             &TokioMechanism {
                 conn_id,
                 conn_info,
                 pool: self.pool.clone(),
                 locks: &self.config.connect_compute_locks,
-                keys: keys.keys,
             },
             &backend,
             self.config.wake_compute_retry_config,
@@ -218,15 +214,18 @@ impl PoolingBackend {
         let conn_id = uuid::Uuid::new_v4();
         tracing::Span::current().record("conn_id", display(conn_id));
         debug!(%conn_id, "pool: opening a new connection '{conn_info}'");
-        let backend = self.auth_backend.as_ref().map(|()| ComputeUserInfo {
-            user: conn_info.user_info.user.clone(),
-            endpoint: EndpointId::from(format!(
-                "{}{LOCAL_PROXY_SUFFIX}",
-                conn_info.user_info.endpoint.normalize()
-            )),
-            options: conn_info.user_info.options.clone(),
+        let backend = self.auth_backend.as_ref().map(|()| ComputeCredentials {
+            info: ComputeUserInfo {
+                user: conn_info.user_info.user.clone(),
+                endpoint: EndpointId::from(format!(
+                    "{}{LOCAL_PROXY_SUFFIX}",
+                    conn_info.user_info.endpoint.normalize()
+                )),
+                options: conn_info.user_info.options.clone(),
+            },
+            keys: crate::auth::backend::ComputeCredentialKeys::None,
         });
-        crate::proxy::connect_compute::connect_to_compute(
+        crate::pglb::connect_compute::connect_to_compute(
             ctx,
             &HyperMechanism {
                 conn_id,
@@ -496,7 +495,6 @@ struct TokioMechanism {
     pool: Arc<GlobalConnPool<postgres_client::Client, EndpointConnPool<postgres_client::Client>>>,
     conn_info: ConnInfo,
     conn_id: uuid::Uuid,
-    keys: ComputeCredentialKeys,
 
     /// connect_to_compute concurrency lock
     locks: &'static ApiLocks<Host>,
@@ -522,10 +520,6 @@ impl ConnectMechanism for TokioMechanism {
             .dbname(&self.conn_info.dbname)
             .connect_timeout(compute_config.timeout);
 
-        if let ComputeCredentialKeys::AuthKeys(auth_keys) = self.keys {
-            config.auth_keys(auth_keys);
-        }
-
         let pause = ctx.latency_timer_pause(crate::metrics::Waiting::Compute);
         let res = config.connect(compute_config).await;
         drop(pause);

proxy/src/serverless/mod.rs

@@ -50,10 +50,10 @@ use crate::context::RequestContext;
 use crate::ext::TaskExt;
 use crate::metrics::Metrics;
 use crate::protocol2::{ConnectHeader, ConnectionInfo, read_proxy_protocol};
+use crate::proxy::run_until_cancelled;
 use crate::rate_limiter::EndpointRateLimiter;
 use crate::serverless::backend::PoolingBackend;
 use crate::serverless::http_util::{api_error_into_response, json_response};
-use crate::util::run_until_cancelled;
 
 pub(crate) const SERVERLESS_DRIVER_SNI: &str = "api";
 pub(crate) const AUTH_BROKER_SNI: &str = "apiauth";

proxy/src/serverless/sql_over_http.rs

@@ -41,11 +41,10 @@ use crate::error::{ErrorKind, ReportableError, UserFacingError};
 use crate::http::{ReadBodyError, read_body_with_limit};
 use crate::metrics::{HttpDirection, Metrics, SniGroup, SniKind};
 use crate::pqproto::StartupMessageParams;
-use crate::proxy::NeonOptions;
+use crate::proxy::{NeonOptions, run_until_cancelled};
 use crate::serverless::backend::HttpConnError;
 use crate::types::{DbName, RoleName};
 use crate::usage_metrics::{MetricCounter, MetricCounterRecorder};
-use crate::util::run_until_cancelled;
 
 #[derive(serde::Deserialize)]
 #[serde(rename_all = "camelCase")]

proxy/src/serverless/websocket.rs

@@ -167,7 +167,7 @@ pub(crate) async fn serve_websocket(
         Ok(Some(p)) => {
             ctx.set_success();
             ctx.log_connect();
-            match p.proxy_pass().await {
+            match p.proxy_pass(&config.connect_to_compute).await {
                 Ok(()) => Ok(()),
                 Err(ErrorSource::Client(err)) => Err(err).context("client"),
                 Err(ErrorSource::Compute(err)) => Err(err).context("compute"),

proxy/src/stream.rs

@@ -0,0 +1,351 @@
+use std::pin::Pin;
+use std::sync::Arc;
+use std::{io, task};
+
+use rustls::ServerConfig;
+use thiserror::Error;
+use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf};
+use tokio_rustls::server::TlsStream;
+
+use crate::error::{ErrorKind, ReportableError, UserFacingError};
+use crate::metrics::Metrics;
+use crate::pqproto::{
+    BeMessage, FE_PASSWORD_MESSAGE, FeStartupPacket, SQLSTATE_INTERNAL_ERROR, WriteBuf,
+    read_message, read_startup,
+};
+use crate::tls::TlsServerEndPoint;
+
+/// Stream wrapper which implements libpq's protocol.
+///
+/// NOTE: This object deliberately doesn't implement [`AsyncRead`]
+/// or [`AsyncWrite`] to prevent subtle errors (e.g. trying
+/// to pass random malformed bytes through the connection).
+pub struct PqStream<S> {
+    stream: S,
+    read: Vec<u8>,
+    write: WriteBuf,
+}
+
+impl<S> PqStream<S> {
+    pub fn get_ref(&self) -> &S {
+        &self.stream
+    }
+
+    /// Construct a new libpq protocol wrapper over a stream without the first startup message.
+    #[cfg(test)]
+    pub fn new_skip_handshake(stream: S) -> Self {
+        Self {
+            stream,
+            read: Vec::new(),
+            write: WriteBuf::new(),
+        }
+    }
+}
+
+impl<S: AsyncRead + AsyncWrite + Unpin> PqStream<S> {
+    /// Construct a new libpq protocol wrapper and read the first startup message.
+    ///
+    /// This is not cancel safe.
+    pub async fn parse_startup(mut stream: S) -> io::Result<(Self, FeStartupPacket)> {
+        let startup = read_startup(&mut stream).await?;
+        Ok((
+            Self {
+                stream,
+                read: Vec::new(),
+                write: WriteBuf::new(),
+            },
+            startup,
+        ))
+    }
+
+    /// Tell the client that encryption is not supported.
+    ///
+    /// This is not cancel safe
+    pub async fn reject_encryption(&mut self) -> io::Result<FeStartupPacket> {
+        // N for No.
+        self.write.encryption(b'N');
+        self.flush().await?;
+        read_startup(&mut self.stream).await
+    }
+}
+
+impl<S: AsyncRead + Unpin> PqStream<S> {
+    /// Read a raw postgres packet, which will respect the max length requested.
+    /// This is not cancel safe.
+    async fn read_raw_expect(&mut self, tag: u8, max: u32) -> io::Result<&mut [u8]> {
+        let (actual_tag, msg) = read_message(&mut self.stream, &mut self.read, max).await?;
+        if actual_tag != tag {
+            return Err(io::Error::other(format!(
+                "incorrect message tag, expected {:?}, got {:?}",
+                tag as char, actual_tag as char,
+            )));
+        }
+        Ok(msg)
+    }
+
+    /// Read a postgres password message, which will respect the max length requested.
+    /// This is not cancel safe.
+    pub async fn read_password_message(&mut self) -> io::Result<&mut [u8]> {
+        // passwords are usually pretty short
+        // and SASL SCRAM messages are no longer than 256 bytes in my testing
+        // (a few hashes and random bytes, encoded into base64).
+        const MAX_PASSWORD_LENGTH: u32 = 512;
+        self.read_raw_expect(FE_PASSWORD_MESSAGE, MAX_PASSWORD_LENGTH)
+            .await
+    }
+}
+
+#[derive(Debug)]
+pub struct ReportedError {
+    source: anyhow::Error,
+    error_kind: ErrorKind,
+}
+
+impl ReportedError {
+    pub fn new(e: (impl UserFacingError + Into<anyhow::Error>)) -> Self {
+        let error_kind = e.get_error_kind();
+        Self {
+            source: e.into(),
+            error_kind,
+        }
+    }
+}
+
+impl std::fmt::Display for ReportedError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.source.fmt(f)
+    }
+}
+
+impl std::error::Error for ReportedError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        self.source.source()
+    }
+}
+
+impl ReportableError for ReportedError {
+    fn get_error_kind(&self) -> ErrorKind {
+        self.error_kind
+    }
+}
+
+impl<S: AsyncWrite + Unpin> PqStream<S> {
+    /// Tell the client that we are willing to accept SSL.
+    /// This is not cancel safe
+    pub async fn accept_tls(mut self) -> io::Result<S> {
+        // S for SSL.
+        self.write.encryption(b'S');
+        self.flush().await?;
+        Ok(self.stream)
+    }
+
+    /// Assert that we are using direct TLS.
+    pub fn accept_direct_tls(self) -> S {
+        self.stream
+    }
+
+    /// Write a raw message to the internal buffer.
+    pub fn write_raw(&mut self, size_hint: usize, tag: u8, f: impl FnOnce(&mut Vec<u8>)) {
+        self.write.write_raw(size_hint, tag, f);
+    }
+
+    /// Write the message into an internal buffer
+    pub fn write_message(&mut self, message: BeMessage<'_>) {
+        message.write_message(&mut self.write);
+    }
+
+    /// Flush the output buffer into the underlying stream.
+    ///
+    /// This is cancel safe.
+    pub async fn flush(&mut self) -> io::Result<()> {
+        self.stream.write_all_buf(&mut self.write).await?;
+        self.write.reset();
+
+        self.stream.flush().await?;
+
+        Ok(())
+    }
+
+    /// Flush the output buffer into the underlying stream.
+    ///
+    /// This is cancel safe.
+    pub async fn flush_and_into_inner(mut self) -> io::Result<S> {
+        self.flush().await?;
+        Ok(self.stream)
+    }
+
+    /// Write the error message to the client, then re-throw it.
+    ///
+    /// Trait [`UserFacingError`] acts as an allowlist for error types.
+    /// If `ctx` is provided and has testodrome_id set, error messages will be prefixed according to error kind.
+    pub(crate) async fn throw_error<E>(
+        &mut self,
+        error: E,
+        ctx: Option<&crate::context::RequestContext>,
+    ) -> ReportedError
+    where
+        E: UserFacingError + Into<anyhow::Error>,
+    {
+        let error_kind = error.get_error_kind();
+        let msg = error.to_string_client();
+
+        if error_kind != ErrorKind::RateLimit && error_kind != ErrorKind::User {
+            tracing::info!(
+                kind = error_kind.to_metric_label(),
+                msg,
+                "forwarding error to user"
+            );
+        }
+
+        let probe_msg;
+        let mut msg = &*msg;
+        if let Some(ctx) = ctx {
+            if ctx.get_testodrome_id().is_some() {
+                let tag = match error_kind {
+                    ErrorKind::User => "client",
+                    ErrorKind::ClientDisconnect => "client",
+                    ErrorKind::RateLimit => "proxy",
+                    ErrorKind::ServiceRateLimit => "proxy",
+                    ErrorKind::Quota => "proxy",
+                    ErrorKind::Service => "proxy",
+                    ErrorKind::ControlPlane => "controlplane",
+                    ErrorKind::Postgres => "other",
+                    ErrorKind::Compute => "compute",
+                };
+                probe_msg = typed_json::json!({
+                    "tag": tag,
+                    "msg": msg,
+                    "cold_start_info": ctx.cold_start_info(),
+                })
+                .to_string();
+                msg = &probe_msg;
+            }
+        }
+
+        // TODO: either preserve the error code from postgres, or assign error codes to proxy errors.
+        self.write.write_error(msg, SQLSTATE_INTERNAL_ERROR);
+
+        self.flush()
+            .await
+            .unwrap_or_else(|e| tracing::debug!("write_message failed: {e}"));
+
+        ReportedError::new(error)
+    }
+}
+
+/// Wrapper for upgrading raw streams into secure streams.
+pub enum Stream<S> {
+    /// We always begin with a raw stream,
+    /// which may then be upgraded into a secure stream.
+    Raw { raw: S },
+    Tls {
+        /// We box [`TlsStream`] since it can be quite large.
+        tls: Box<TlsStream<S>>,
+        /// Channel binding parameter
+        tls_server_end_point: TlsServerEndPoint,
+    },
+}
+
+impl<S: Unpin> Unpin for Stream<S> {}
+
+impl<S> Stream<S> {
+    /// Construct a new instance from a raw stream.
+    pub fn from_raw(raw: S) -> Self {
+        Self::Raw { raw }
+    }
+
+    /// Return SNI hostname when it's available.
+    pub fn sni_hostname(&self) -> Option<&str> {
+        match self {
+            Stream::Raw { .. } => None,
+            Stream::Tls { tls, .. } => tls.get_ref().1.server_name(),
+        }
+    }
+
+    pub(crate) fn tls_server_end_point(&self) -> TlsServerEndPoint {
+        match self {
+            Stream::Raw { .. } => TlsServerEndPoint::Undefined,
+            Stream::Tls {
+                tls_server_end_point,
+                ..
+            } => *tls_server_end_point,
+        }
+    }
+}
+
+#[derive(Debug, Error)]
+#[error("Can't upgrade TLS stream")]
+pub enum StreamUpgradeError {
+    #[error("Bad state reached: can't upgrade TLS stream")]
+    AlreadyTls,
+
+    #[error("Can't upgrade stream: IO error: {0}")]
+    Io(#[from] io::Error),
+}
+
+impl<S: AsyncRead + AsyncWrite + Unpin> Stream<S> {
+    /// If possible, upgrade raw stream into a secure TLS-based stream.
+    pub async fn upgrade(
+        self,
+        cfg: Arc<ServerConfig>,
+        record_handshake_error: bool,
+    ) -> Result<TlsStream<S>, StreamUpgradeError> {
+        match self {
+            Stream::Raw { raw } => Ok(tokio_rustls::TlsAcceptor::from(cfg)
+                .accept(raw)
+                .await
+                .inspect_err(|_| {
+                    if record_handshake_error {
+                        Metrics::get().proxy.tls_handshake_failures.inc();
+                    }
+                })?),
+            Stream::Tls { .. } => Err(StreamUpgradeError::AlreadyTls),
+        }
+    }
+}
+
+impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for Stream<S> {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+        buf: &mut ReadBuf<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        match &mut *self {
+            Self::Raw { raw } => Pin::new(raw).poll_read(context, buf),
+            Self::Tls { tls, .. } => Pin::new(tls).poll_read(context, buf),
+        }
+    }
+}
+
+impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for Stream<S> {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+        buf: &[u8],
+    ) -> task::Poll<io::Result<usize>> {
+        match &mut *self {
+            Self::Raw { raw } => Pin::new(raw).poll_write(context, buf),
+            Self::Tls { tls, .. } => Pin::new(tls).poll_write(context, buf),
+        }
+    }
+
+    fn poll_flush(
+        mut self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        match &mut *self {
+            Self::Raw { raw } => Pin::new(raw).poll_flush(context),
+            Self::Tls { tls, .. } => Pin::new(tls).poll_flush(context),
+        }
+    }
+
+    fn poll_shutdown(
+        mut self: Pin<&mut Self>,
+        context: &mut task::Context<'_>,
+    ) -> task::Poll<io::Result<()>> {
+        match &mut *self {
+            Self::Raw { raw } => Pin::new(raw).poll_shutdown(context),
+            Self::Tls { tls, .. } => Pin::new(tls).poll_shutdown(context),
+        }
+    }
+}

proxy/src/stream/mod.rs

@@ -1,168 +0,0 @@
-mod pq_backend;
-mod pq_frontend;
-
-use std::pin::Pin;
-use std::sync::Arc;
-use std::{io, task};
-
-pub use pq_backend::{BackendError, PostgresError, PqBeStream};
-pub use pq_frontend::PqFeStream;
-use rustls::ServerConfig;
-use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
-use tokio_rustls::server::TlsStream;
-
-use crate::error::{ErrorKind, ReportableError, UserFacingError};
-use crate::metrics::Metrics;
-use crate::tls::TlsServerEndPoint;
-
-#[derive(Debug)]
-pub struct ReportedError {
-    source: anyhow::Error,
-    error_kind: ErrorKind,
-}
-
-impl ReportedError {
-    pub fn new(e: (impl UserFacingError + Into<anyhow::Error>)) -> Self {
-        let error_kind = e.get_error_kind();
-        Self {
-            source: e.into(),
-            error_kind,
-        }
-    }
-}
-
-impl std::fmt::Display for ReportedError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.source.fmt(f)
-    }
-}
-
-impl std::error::Error for ReportedError {
-    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
-        self.source.source()
-    }
-}
-
-impl ReportableError for ReportedError {
-    fn get_error_kind(&self) -> ErrorKind {
-        self.error_kind
-    }
-}
-
-/// Wrapper for upgrading raw streams into secure streams.
-pub enum Stream<S> {
-    /// We always begin with a raw stream,
-    /// which may then be upgraded into a secure stream.
-    Raw { raw: S },
-    Tls {
-        /// We box [`TlsStream`] since it can be quite large.
-        tls: Box<TlsStream<S>>,
-        /// Channel binding parameter
-        tls_server_end_point: TlsServerEndPoint,
-    },
-}
-
-impl<S: Unpin> Unpin for Stream<S> {}
-
-impl<S> Stream<S> {
-    /// Construct a new instance from a raw stream.
-    pub fn from_raw(raw: S) -> Self {
-        Self::Raw { raw }
-    }
-
-    /// Return SNI hostname when it's available.
-    pub fn sni_hostname(&self) -> Option<&str> {
-        match self {
-            Stream::Raw { .. } => None,
-            Stream::Tls { tls, .. } => tls.get_ref().1.server_name(),
-        }
-    }
-
-    pub(crate) fn tls_server_end_point(&self) -> TlsServerEndPoint {
-        match self {
-            Stream::Raw { .. } => TlsServerEndPoint::Undefined,
-            Stream::Tls {
-                tls_server_end_point,
-                ..
-            } => *tls_server_end_point,
-        }
-    }
-}
-
-#[derive(Debug, Error)]
-#[error("Can't upgrade TLS stream")]
-pub enum StreamUpgradeError {
-    #[error("Bad state reached: can't upgrade TLS stream")]
-    AlreadyTls,
-
-    #[error("Can't upgrade stream: IO error: {0}")]
-    Io(#[from] io::Error),
-}
-
-impl<S: AsyncRead + AsyncWrite + Unpin> Stream<S> {
-    /// If possible, upgrade raw stream into a secure TLS-based stream.
-    pub async fn upgrade(
-        self,
-        cfg: Arc<ServerConfig>,
-        record_handshake_error: bool,
-    ) -> Result<TlsStream<S>, StreamUpgradeError> {
-        match self {
-            Stream::Raw { raw } => Ok(tokio_rustls::TlsAcceptor::from(cfg)
-                .accept(raw)
-                .await
-                .inspect_err(|_| {
-                    if record_handshake_error {
-                        Metrics::get().proxy.tls_handshake_failures.inc();
-                    }
-                })?),
-            Stream::Tls { .. } => Err(StreamUpgradeError::AlreadyTls),
-        }
-    }
-}
-
-impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for Stream<S> {
-    fn poll_read(
-        mut self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-        buf: &mut ReadBuf<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        match &mut *self {
-            Self::Raw { raw } => Pin::new(raw).poll_read(context, buf),
-            Self::Tls { tls, .. } => Pin::new(tls).poll_read(context, buf),
-        }
-    }
-}
-
-impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for Stream<S> {
-    fn poll_write(
-        mut self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-        buf: &[u8],
-    ) -> task::Poll<io::Result<usize>> {
-        match &mut *self {
-            Self::Raw { raw } => Pin::new(raw).poll_write(context, buf),
-            Self::Tls { tls, .. } => Pin::new(tls).poll_write(context, buf),
-        }
-    }
-
-    fn poll_flush(
-        mut self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        match &mut *self {
-            Self::Raw { raw } => Pin::new(raw).poll_flush(context),
-            Self::Tls { tls, .. } => Pin::new(tls).poll_flush(context),
-        }
-    }
-
-    fn poll_shutdown(
-        mut self: Pin<&mut Self>,
-        context: &mut task::Context<'_>,
-    ) -> task::Poll<io::Result<()>> {
-        match &mut *self {
-            Self::Raw { raw } => Pin::new(raw).poll_shutdown(context),
-            Self::Tls { tls, .. } => Pin::new(tls).poll_shutdown(context),
-        }
-    }
-}

proxy/src/stream/pq_backend.rs

@@ -1,165 +0,0 @@
-//! Postgres connection from backend, proxy is the frontend.
-
-use std::io;
-
-use bytes::Bytes;
-use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
-
-use crate::pqproto::{
-    AuthTag, BE_AUTH_MESSAGE, BE_ERR_MESSAGE, BeTag, ErrorCode, SQLSTATE_INTERNAL_ERROR,
-    StartupMessageParams, WriteBuf, read_message,
-};
-
-/// Stream wrapper which implements libpq's protocol.
-pub struct PqBeStream<S> {
-    stream: S,
-    read: Vec<u8>,
-    write: WriteBuf,
-}
-
-impl<S> PqBeStream<S> {
-    pub fn get_ref(&self) -> &S {
-        &self.stream
-    }
-
-    /// Construct a new libpq protocol wrapper and write the first startup message.
-    pub fn new(stream: S, params: &StartupMessageParams) -> Self {
-        let mut write = WriteBuf::new();
-        write.startup(params);
-        Self {
-            stream,
-            read: Vec::new(),
-            write,
-        }
-    }
-}
-
-impl<S: AsyncRead + Unpin> PqBeStream<S> {
-    /// Read a raw postgres packet from the backend, which will respect the max length requested,
-    /// as well as handling postgres error messages.
-    ///
-    /// This is not cancel safe.
-    pub async fn read_raw_be(&mut self, max: u32) -> Result<(BeTag, &mut [u8]), PostgresError> {
-        let (tag, msg) = read_message(&mut self.stream, &mut self.read, max).await?;
-        match BeTag(tag) {
-            BE_ERR_MESSAGE => Err(PostgresError::Error(BackendError {
-                data: msg.to_vec().into(),
-            })),
-            tag => Ok((tag, msg)),
-        }
-    }
-
-    /// Read a raw postgres packet, which will respect the max length requested.
-    /// This is not cancel safe.
-    async fn read_raw_be_expect(
-        &mut self,
-        tag: BeTag,
-        max: u32,
-    ) -> Result<&mut [u8], PostgresError> {
-        let (actual_tag, msg) = self.read_raw_be(max).await?;
-        if actual_tag != tag {
-            return Err(PostgresError::Unexpected(UnexpectedMessage {
-                expected: tag,
-                tag: actual_tag,
-                data: msg.to_vec().into(),
-            }));
-        }
-        Ok(msg)
-    }
-
-    /// Read a postgres backend auth message.
-    /// This is not cancel safe.
-    pub async fn read_auth_message(&mut self) -> Result<(AuthTag, &mut [u8]), PostgresError> {
-        const MAX_AUTH_LENGTH: u32 = 512;
-
-        self.read_raw_be_expect(BE_AUTH_MESSAGE, MAX_AUTH_LENGTH)
-            .await?
-            .split_first_chunk_mut()
-            .map(|(tag, msg)| (AuthTag(i32::from_be_bytes(*tag)), msg))
-            .ok_or(PostgresError::InvalidAuthMessage)
-    }
-}
-
-impl<S: AsyncWrite + Unpin> PqBeStream<S> {
-    /// Write a raw message to the internal buffer.
-    pub fn write_raw(&mut self, size_hint: usize, tag: u8, f: impl FnOnce(&mut Vec<u8>)) {
-        self.write.write_raw(size_hint, tag, f);
-    }
-
-    /// Flush the output buffer into the underlying stream.
-    ///
-    /// This is cancel safe.
-    pub async fn flush(&mut self) -> io::Result<()> {
-        self.stream.write_all_buf(&mut self.write).await?;
-        self.write.reset();
-
-        self.stream.flush().await?;
-
-        Ok(())
-    }
-
-    /// Flush the output buffer into the underlying stream.
-    ///
-    /// This is cancel safe.
-    pub async fn flush_and_into_inner(mut self) -> io::Result<S> {
-        self.flush().await?;
-        Ok(self.stream)
-    }
-}
-
-#[derive(Debug, Error)]
-pub enum PostgresError {
-    #[error("postgres responded with error {0}")]
-    Error(#[from] BackendError),
-    #[error("postgres responded with an unexpected message: {0}")]
-    Unexpected(#[from] UnexpectedMessage),
-    #[error("postgres responded with an invalid authentication message")]
-    InvalidAuthMessage,
-    #[error("IO error from compute: {0}")]
-    Io(#[from] io::Error),
-}
-
-#[derive(Debug, Error)]
-#[error("expected {expected}, got {tag} with data {data:?}")]
-pub struct UnexpectedMessage {
-    expected: BeTag,
-    tag: BeTag,
-    data: Bytes,
-}
-
-pub struct BackendError {
-    data: Bytes,
-}
-
-impl BackendError {
-    pub fn parse(&self) -> (ErrorCode, &[u8]) {
-        let mut code = &[] as &[u8];
-        let mut message = &[] as &[u8];
-
-        for param in self.data.split(|b| *b == 0) {
-            match param {
-                [b'M', rest @ ..] => message = rest,
-                [b'C', rest @ ..] => code = rest,
-                _ => {}
-            }
-        }
-
-        let code = code.try_into().map_or(SQLSTATE_INTERNAL_ERROR, ErrorCode);
-
-        (code, message)
-    }
-}
-
-impl std::fmt::Debug for BackendError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{self}")
-    }
-}
-
-impl std::fmt::Display for BackendError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{:?}", &self.data)
-    }
-}
-impl std::error::Error for BackendError {}

proxy/src/stream/pq_frontend.rs

@@ -1,197 +0,0 @@
-//! Postgres connection from frontend, proxy is the backend.
-
-use std::io;
-
-use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
-
-use crate::error::{ErrorKind, UserFacingError};
-use crate::pqproto::{
-    BeMessage, BeTag, FE_PASSWORD_MESSAGE, FeStartupPacket, FeTag, SQLSTATE_INTERNAL_ERROR,
-    WriteBuf, read_message, read_startup,
-};
-use crate::stream::ReportedError;
-
-/// Stream wrapper which implements libpq's protocol.
-pub struct PqFeStream<S> {
-    stream: S,
-    read: Vec<u8>,
-    write: WriteBuf,
-}
-
-impl<S> PqFeStream<S> {
-    pub fn get_ref(&self) -> &S {
-        &self.stream
-    }
-
-    /// Construct a new libpq protocol wrapper over a stream without the first startup message.
-    #[cfg(test)]
-    pub fn new_skip_handshake(stream: S) -> Self {
-        Self {
-            stream,
-            read: Vec::new(),
-            write: WriteBuf::new(),
-        }
-    }
-
-    pub fn write_buf_len(&self) -> usize {
-        self.write.len()
-    }
-}
-
-impl<S: AsyncRead + AsyncWrite + Unpin> PqFeStream<S> {
-    /// Construct a new libpq protocol wrapper and read the first startup message.
-    ///
-    /// This is not cancel safe.
-    pub async fn parse_startup(mut stream: S) -> io::Result<(Self, FeStartupPacket)> {
-        let startup = read_startup(&mut stream).await?;
-        Ok((
-            Self {
-                stream,
-                read: Vec::new(),
-                write: WriteBuf::new(),
-            },
-            startup,
-        ))
-    }
-
-    /// Tell the client that encryption is not supported.
-    ///
-    /// This is not cancel safe
-    pub async fn reject_encryption(&mut self) -> io::Result<FeStartupPacket> {
-        // N for No.
-        self.write.encryption(b'N');
-        self.flush().await?;
-        read_startup(&mut self.stream).await
-    }
-}
-
-impl<S: AsyncRead + Unpin> PqFeStream<S> {
-    /// Read a raw postgres packet, which will respect the max length requested.
-    /// This is not cancel safe.
-    async fn read_raw_expect(&mut self, tag: FeTag, max: u32) -> io::Result<&mut [u8]> {
-        let (actual_tag, msg) = read_message(&mut self.stream, &mut self.read, max).await?;
-        let actual_tag = FeTag(actual_tag);
-        if actual_tag != tag {
-            return Err(io::Error::other(format!(
-                "incorrect message tag, expected {tag}, got {actual_tag}",
-            )));
-        }
-        Ok(msg)
-    }
-
-    /// Read a postgres password message, which will respect the max length requested.
-    /// This is not cancel safe.
-    pub async fn read_password_message(&mut self) -> io::Result<&mut [u8]> {
-        // passwords are usually pretty short
-        // and SASL SCRAM messages are no longer than 256 bytes in my testing
-        // (a few hashes and random bytes, encoded into base64).
-        const MAX_PASSWORD_LENGTH: u32 = 512;
-        self.read_raw_expect(FE_PASSWORD_MESSAGE, MAX_PASSWORD_LENGTH)
-            .await
-    }
-}
-
-impl<S: AsyncWrite + Unpin> PqFeStream<S> {
-    /// Tell the client that we are willing to accept SSL.
-    /// This is not cancel safe
-    pub async fn accept_tls(mut self) -> io::Result<S> {
-        // S for SSL.
-        self.write.encryption(b'S');
-        self.flush().await?;
-        Ok(self.stream)
-    }
-
-    /// Assert that we are using direct TLS.
-    pub fn accept_direct_tls(self) -> S {
-        self.stream
-    }
-
-    /// Write a raw message to the internal buffer.
-    pub fn write_raw(&mut self, size_hint: usize, tag: BeTag, f: impl FnOnce(&mut Vec<u8>)) {
-        self.write.write_raw(size_hint, tag.0, f);
-    }
-
-    /// Write the message into an internal buffer
-    pub fn write_message(&mut self, message: BeMessage<'_>) {
-        message.write_message(&mut self.write);
-    }
-
-    /// Flush the output buffer into the underlying stream.
-    ///
-    /// This is cancel safe.
-    pub async fn flush(&mut self) -> io::Result<()> {
-        self.stream.write_all_buf(&mut self.write).await?;
-        self.write.reset();
-
-        self.stream.flush().await?;
-
-        Ok(())
-    }
-
-    /// Flush the output buffer into the underlying stream.
-    ///
-    /// This is cancel safe.
-    pub async fn flush_and_into_inner(mut self) -> io::Result<S> {
-        self.flush().await?;
-        Ok(self.stream)
-    }
-
-    /// Write the error message to the client, then re-throw it.
-    ///
-    /// Trait [`UserFacingError`] acts as an allowlist for error types.
-    /// If `ctx` is provided and has testodrome_id set, error messages will be prefixed according to error kind.
-    pub(crate) async fn throw_error<E>(
-        &mut self,
-        error: E,
-        ctx: Option<&crate::context::RequestContext>,
-    ) -> ReportedError
-    where
-        E: UserFacingError + Into<anyhow::Error>,
-    {
-        let error_kind = error.get_error_kind();
-        let msg = error.to_string_client();
-
-        if error_kind != ErrorKind::RateLimit && error_kind != ErrorKind::User {
-            tracing::info!(
-                kind = error_kind.to_metric_label(),
-                %error,
-                msg,
-                "forwarding error to user"
-            );
-        }
-
-        let probe_msg;
-        let mut msg = &*msg;
-        if let Some(ctx) = ctx {
-            if ctx.get_testodrome_id().is_some() {
-                let tag = match error_kind {
-                    ErrorKind::User => "client",
-                    ErrorKind::ClientDisconnect => "client",
-                    ErrorKind::RateLimit => "proxy",
-                    ErrorKind::ServiceRateLimit => "proxy",
-                    ErrorKind::Quota => "proxy",
-                    ErrorKind::Service => "proxy",
-                    ErrorKind::ControlPlane => "controlplane",
-                    ErrorKind::Postgres => "other",
-                    ErrorKind::Compute => "compute",
-                };
-                probe_msg = typed_json::json!({
-                    "tag": tag,
-                    "msg": msg,
-                    "cold_start_info": ctx.cold_start_info(),
-                })
-                .to_string();
-                msg = &probe_msg;
-            }
-        }
-
-        // TODO: either preserve the error code from postgres, or assign error codes to proxy errors.
-        self.write.write_error(msg, SQLSTATE_INTERNAL_ERROR);
-
-        self.flush()
-            .await
-            .unwrap_or_else(|e| tracing::debug!("write_message failed: {e}"));
-
-        ReportedError::new(error)
-    }
-}

proxy/src/util.rs

@@ -1,14 +0,0 @@
-use std::pin::pin;
-
-use futures::future::{Either, select};
-use tokio_util::sync::CancellationToken;
-
-pub async fn run_until_cancelled<F: Future>(
-    f: F,
-    cancellation_token: &CancellationToken,
-) -> Option<F::Output> {
-    match select(pin!(f), pin!(cancellation_token.cancelled())).await {
-        Either::Left((f, _)) => Some(f),
-        Either::Right(((), _)) => None,
-    }
-}

storage_controller/Cargo.toml

@@ -65,8 +65,9 @@ diesel-async = { version = "0.5.2", features = ["postgres", "bb8", "async-connec
 diesel_migrations = { version = "2.2.0" }
 scoped-futures = "0.1.4"
 
+compute_api = { path = "../libs/compute_api/" }
 http-utils = { path = "../libs/http-utils/" }
 utils = { path = "../libs/utils/" }
 metrics = { path = "../libs/metrics/" }
 control_plane = { path = "../control_plane" }
-workspace_hack = { version = "0.1", path = "../workspace_hack" }
+workspace_hack = { version = "0.1", path = "../workspace_hack" }

storage_controller/src/compute_hook.rs

@@ -428,7 +428,10 @@ impl ComputeHook {
                     .expect("Unknown pageserver");
                 let (pg_host, pg_port) = parse_host_port(&ps_conf.listen_pg_addr)
                     .expect("Unable to parse listen_pg_addr");
-                (pg_host, pg_port.unwrap_or(5432))
+                compute_api::spec::Pageserver {
+                    host: pg_host,
+                    port: pg_port.unwrap_or(5432),
+                }
             })
             .collect::<Vec<_>>();
 

test_runner/fixtures/endpoint/http.py

@@ -69,17 +69,15 @@ class EndpointHttpClient(requests.Session):
         json: dict[str, str] = res.json()
         return json
 
-    def prewarm_lfc(self, from_endpoint_id: str | None = None):
-        url: str = f"http://localhost:{self.external_port}/lfc/prewarm"
-        params = {"from_endpoint": from_endpoint_id} if from_endpoint_id else dict()
-        self.post(url, params=params).raise_for_status()
+    def prewarm_lfc(self):
+        self.post(f"http://localhost:{self.external_port}/lfc/prewarm").raise_for_status()
 
         def prewarmed():
             json = self.prewarm_lfc_status()
             status, err = json["status"], json.get("error")
             assert status == "completed", f"{status}, error {err}"
 
-        wait_until(prewarmed, timeout=60)
+        wait_until(prewarmed)
 
     def offload_lfc(self):
         url = f"http://localhost:{self.external_port}/lfc/offload"

test_runner/fixtures/neon_api.py

@@ -129,18 +129,6 @@ class NeonAPI:
 
         return cast("dict[str, Any]", resp.json())
 
-    def get_project_limits(self, project_id: str) -> dict[str, Any]:
-        resp = self.__request(
-            "GET",
-            f"/projects/{project_id}/limits",
-            headers={
-                "Accept": "application/json",
-                "Content-Type": "application/json",
-            },
-        )
-
-        return cast("dict[str, Any]", resp.json())
-
     def delete_project(
         self,
         project_id: str,

test_runner/fixtures/neon_fixtures.py

@@ -4046,16 +4046,6 @@ def static_proxy(
         "CREATE TABLE neon_control_plane.endpoints (endpoint_id VARCHAR(255) PRIMARY KEY, allowed_ips VARCHAR(255))"
     )
 
-    vanilla_pg.stop()
-    vanilla_pg.edit_hba(
-        [
-            "local all all              trust",
-            "host  all all 127.0.0.1/32 scram-sha-256",
-            "host  all all ::1/128      scram-sha-256",
-        ]
-    )
-    vanilla_pg.start()
-
     proxy_port = port_distributor.get_port()
     mgmt_port = port_distributor.get_port()
     http_port = port_distributor.get_port()

test_runner/random_ops/test_random_ops.py

@@ -45,8 +45,6 @@ class NeonEndpoint:
         if self.branch.connect_env:
             self.connect_env = self.branch.connect_env.copy()
             self.connect_env["PGHOST"] = self.host
-        if self.type == "read_only":
-            self.project.read_only_endpoints_total += 1
 
     def delete(self):
         self.project.delete_endpoint(self.id)
@@ -230,13 +228,8 @@ class NeonProject:
         self.benchmarks: dict[str, subprocess.Popen[Any]] = {}
         self.restore_num: int = 0
         self.restart_pgbench_on_console_errors: bool = False
-        self.limits: dict[str, Any] = self.get_limits()["limits"]
-        self.read_only_endpoints_total: int = 0
 
-    def get_limits(self) -> dict[str, Any]:
-        return self.neon_api.get_project_limits(self.id)
-
-    def delete(self) -> None:
+    def delete(self):
         self.neon_api.delete_project(self.id)
 
     def create_branch(self, parent_id: str | None = None) -> NeonBranch | None:
@@ -289,7 +282,6 @@ class NeonProject:
         self.neon_api.delete_endpoint(self.id, endpoint_id)
         self.endpoints[endpoint_id].branch.endpoints.pop(endpoint_id)
         self.endpoints.pop(endpoint_id)
-        self.read_only_endpoints_total -= 1
         self.wait()
 
     def start_benchmark(self, target: str, clients: int = 10) -> subprocess.Popen[Any]:
@@ -377,64 +369,49 @@ def setup_class(
         print(f"::warning::Retried on 524 error {neon_api.retries524} times")
     if neon_api.retries4xx > 0:
         print(f"::warning::Retried on 4xx error {neon_api.retries4xx} times")
-    log.info("Removing the project %s", project.id)
+    log.info("Removing the project")
     project.delete()
 
 
-def do_action(project: NeonProject, action: str) -> bool:
+def do_action(project: NeonProject, action: str) -> None:
     """
     Runs the action
     """
     log.info("Action: %s", action)
     if action == "new_branch":
         log.info("Trying to create a new branch")
-        if 0 <= project.limits["max_branches"] <= len(project.branches):
-            log.info(
-                "Maximum branch limit exceeded (%s of %s)",
-                len(project.branches),
-                project.limits["max_branches"],
-            )
-            return False
         parent = project.branches[
             random.choice(list(set(project.branches.keys()) - project.reset_branches))
         ]
         log.info("Parent: %s", parent)
         child = parent.create_child_branch()
         if child is None:
-            return False
+            return
         log.info("Created branch %s", child)
         child.start_benchmark()
     elif action == "delete_branch":
         if project.leaf_branches:
-            target: NeonBranch = random.choice(list(project.leaf_branches.values()))
+            target = random.choice(list(project.leaf_branches.values()))
             log.info("Trying to delete branch %s", target)
             target.delete()
         else:
             log.info("Leaf branches not found, skipping")
-            return False
     elif action == "new_ro_endpoint":
-        if 0 <= project.limits["max_read_only_endpoints"] <= project.read_only_endpoints_total:
-            log.info(
-                "Maximum read only endpoint limit exceeded (%s of %s)",
-                project.read_only_endpoints_total,
-                project.limits["max_read_only_endpoints"],
-            )
-            return False
         ep = random.choice(
             [br for br in project.branches.values() if br.id not in project.reset_branches]
         ).create_ro_endpoint()
         log.info("Created the RO endpoint with id %s branch: %s", ep.id, ep.branch.id)
         ep.start_benchmark()
     elif action == "delete_ro_endpoint":
-        if project.read_only_endpoints_total == 0:
-            log.info("no read_only endpoints present, skipping")
-            return False
         ro_endpoints: list[NeonEndpoint] = [
             endpoint for endpoint in project.endpoints.values() if endpoint.type == "read_only"
         ]
-        target_ep: NeonEndpoint = random.choice(ro_endpoints)
-        target_ep.delete()
-        log.info("endpoint %s deleted", target_ep.id)
+        if ro_endpoints:
+            target_ep: NeonEndpoint = random.choice(ro_endpoints)
+            target_ep.delete()
+            log.info("endpoint %s deleted", target_ep.id)
+        else:
+            log.info("no read_only endpoints present, skipping")
     elif action == "restore_random_time":
         if project.leaf_branches:
             br: NeonBranch = random.choice(list(project.leaf_branches.values()))
@@ -442,10 +419,8 @@ def do_action(project: NeonProject, action: str) -> bool:
             br.restore_random_time()
         else:
             log.info("No leaf branches found")
-            return False
     else:
         raise ValueError(f"The action {action} is unknown")
-    return True
 
 
 @pytest.mark.timeout(7200)
@@ -482,9 +457,8 @@ def test_api_random(
     pg_bin.run(["pgbench", "-i", "-I", "dtGvp", "-s100"], env=project.main_branch.connect_env)
     for _ in range(num_operations):
         log.info("Starting action #%s", _ + 1)
-        while not do_action(
+        do_action(
             project, random.choices([a[0] for a in ACTIONS], weights=[w[1] for w in ACTIONS])[0]
-        ):
-            log.info("Retrying...")
+        )
         project.check_all_benchmarks()
     assert True

test_runner/regress/test_lfc_prewarm.py

@@ -188,8 +188,7 @@ def test_lfc_prewarm_under_workload(neon_simple_env: NeonEnv, query: LfcQueryMet
             pg_cur.execute("select pg_reload_conf()")
 
             if query is LfcQueryMethod.COMPUTE_CTL:
-                # Same thing as prewarm_lfc(), testing other method
-                http_client.prewarm_lfc(endpoint.endpoint_id)
+                http_client.prewarm_lfc()
             else:
                 pg_cur.execute("select prewarm_local_cache(%s)", (lfc_state,))
 

test_runner/regress/test_pageserver_restarts_under_workload.py

@@ -16,7 +16,7 @@ if TYPE_CHECKING:
 
 # Test restarting page server, while safekeeper and compute node keep
 # running.
-def test_pageserver_restarts_under_workload(neon_simple_env: NeonEnv, pg_bin: PgBin):
+def test_pageserver_restarts_under_worload(neon_simple_env: NeonEnv, pg_bin: PgBin):
     env = neon_simple_env
     env.create_branch("test_pageserver_restarts")
     endpoint = env.endpoints.create_start("test_pageserver_restarts")
@@ -28,11 +28,7 @@ def test_pageserver_restarts_under_workload(neon_simple_env: NeonEnv, pg_bin: Pg
         pg_bin.run_capture(["pgbench", "-i", "-I", "dtGvp", f"-s{scale}", connstr])
         pg_bin.run_capture(["pgbench", f"-T{n_restarts}", connstr])
 
-    thread = threading.Thread(
-        target=run_pgbench,
-        args=(endpoint.connstr(options="-cstatement_timeout=360s"),),
-        daemon=True,
-    )
+    thread = threading.Thread(target=run_pgbench, args=(endpoint.connstr(),), daemon=True)
     thread.start()
 
     for _ in range(n_restarts):

test_runner/regress/test_proxy_allowed_ips.py

@@ -19,15 +19,11 @@ TABLE_NAME = "neon_control_plane.endpoints"
 async def test_proxy_psql_allowed_ips(static_proxy: NeonProxy, vanilla_pg: VanillaPostgres):
     # Shouldn't be able to connect to this project
     vanilla_pg.safe_psql(
-        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('private-project', '8.8.8.8')",
-        user="proxy",
-        password="password",
+        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('private-project', '8.8.8.8')"
     )
     # Should be able to connect to this project
     vanilla_pg.safe_psql(
-        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('generic-project', '::1,127.0.0.1')",
-        user="proxy",
-        password="password",
+        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('generic-project', '::1,127.0.0.1')"
     )
 
     def check_cannot_connect(**kwargs):
@@ -64,9 +60,7 @@ async def test_proxy_http_allowed_ips(static_proxy: NeonProxy, vanilla_pg: Vanil
 
     # Shouldn't be able to connect to this project
     vanilla_pg.safe_psql(
-        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('proxy', '8.8.8.8')",
-        user="proxy",
-        password="password",
+        f"INSERT INTO {TABLE_NAME} (endpoint_id, allowed_ips) VALUES ('proxy', '8.8.8.8')"
     )
 
     def query(status: int, query: str, *args):
@@ -81,8 +75,6 @@ async def test_proxy_http_allowed_ips(static_proxy: NeonProxy, vanilla_pg: Vanil
     query(400, "select 1;")  # ip address is not allowed
     # Should be able to connect to this project
     vanilla_pg.safe_psql(
-        f"UPDATE {TABLE_NAME} SET allowed_ips = '8.8.8.8,127.0.0.1' WHERE endpoint_id = 'proxy'",
-        user="proxy",
-        password="password",
+        f"UPDATE {TABLE_NAME} SET allowed_ips = '8.8.8.8,127.0.0.1' WHERE endpoint_id = 'proxy'"
     )
     query(200, "select 1;")  # should work now