re-impl auth startup to compute

Cargo.lock

@@ -903,6 +903,12 @@ version = "0.13.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
 
+[[package]]
+name = "base64"
+version = "0.20.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ea22880d78093b0cbe17c89f64a7d457941e65759157ec6cb31a31d652b05e5"
+
 [[package]]
 name = "base64"
 version = "0.21.7"
@@ -1294,7 +1300,7 @@ dependencies = [
  "aws-smithy-types",
  "axum",
  "axum-extra",
- "base64 0.22.1",
+ "base64 0.13.1",
  "bytes",
  "camino",
  "cfg-if",
@@ -1420,7 +1426,7 @@ name = "control_plane"
 version = "0.1.0"
 dependencies = [
  "anyhow",
- "base64 0.22.1",
+ "base64 0.13.1",
  "camino",
  "clap",
  "comfy-table",
@@ -4812,7 +4818,7 @@ dependencies = [
 name = "postgres-protocol2"
 version = "0.1.0"
 dependencies = [
- "base64 0.22.1",
+ "base64 0.20.0",
  "byteorder",
  "bytes",
  "fallible-iterator",
@@ -5184,7 +5190,7 @@ dependencies = [
  "aws-config",
  "aws-sdk-iam",
  "aws-sigv4",
- "base64 0.22.1",
+ "base64 0.13.1",
  "bstr",
  "bytes",
  "camino",
@@ -6488,17 +6494,15 @@ dependencies = [
 
 [[package]]
 name = "serde_with"
-version = "3.12.0"
+version = "2.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d6b6f7f2fcb69f747921f79f3926bd1e203fce4fef62c268dd3abfb6d86029aa"
+checksum = "07ff71d2c147a7b57362cead5e22f772cd52f6ab31cfcd9edcd7f6aeb2a0afbe"
 dependencies = [
- "base64 0.22.1",
+ "base64 0.13.1",
  "chrono",
  "hex",
  "indexmap 1.9.3",
- "indexmap 2.9.0",
  "serde",
- "serde_derive",
  "serde_json",
  "serde_with_macros",
  "time",
@@ -6506,9 +6510,9 @@ dependencies = [
 
 [[package]]
 name = "serde_with_macros"
-version = "3.12.0"
+version = "2.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8d00caa5193a3c8362ac2b73be6b9e768aa5a4b2f721d8f4b339600c3cb51f8e"
+checksum = "881b6f881b17d13214e5d494c939ebab463d01264ce1811e9d4ac3a882e7695f"
 dependencies = [
  "darling",
  "proc-macro2",
@@ -8579,6 +8583,7 @@ dependencies = [
  "anyhow",
  "axum",
  "axum-core",
+ "base64 0.13.1",
  "base64 0.21.7",
  "base64ct",
  "bytes",

Cargo.toml

@@ -72,7 +72,7 @@ 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"] }
-base64 = "0.22"
+base64 = "0.13.0"
 bincode = "1.3"
 bindgen = "0.71"
 bit_field = "0.10.2"
@@ -171,7 +171,7 @@ sentry = { version = "0.37", default-features = false, features = ["backtrace",
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1"
 serde_path_to_error = "0.1"
-serde_with = { version = "3", features = [ "base64" ] }
+serde_with = { version = "2.0", features = [ "base64" ] }
 serde_assert = "0.5.0"
 sha2 = "0.10.2"
 signal-hook = "0.3"

control_plane/src/endpoint.rs

@@ -45,8 +45,6 @@ use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use anyhow::{Context, Result, anyhow, bail};
-use base64::Engine;
-use base64::prelude::BASE64_URL_SAFE_NO_PAD;
 use compute_api::requests::{
     COMPUTE_AUDIENCE, ComputeClaims, ComputeClaimsScope, ConfigurationRequest,
 };
@@ -166,7 +164,7 @@ impl ComputeControlPlane {
                     public_key_use: Some(PublicKeyUse::Signature),
                     key_operations: Some(vec![KeyOperations::Verify]),
                     key_algorithm: Some(KeyAlgorithm::EdDSA),
-                    key_id: Some(BASE64_URL_SAFE_NO_PAD.encode(key_hash)),
+                    key_id: Some(base64::encode_config(key_hash, base64::URL_SAFE_NO_PAD)),
                     x509_url: None::<String>,
                     x509_chain: None::<Vec<String>>,
                     x509_sha1_fingerprint: None::<String>,
@@ -175,7 +173,7 @@ impl ComputeControlPlane {
                 algorithm: AlgorithmParameters::OctetKeyPair(OctetKeyPairParameters {
                     key_type: OctetKeyPairType::OctetKeyPair,
                     curve: EllipticCurve::Ed25519,
-                    x: BASE64_URL_SAFE_NO_PAD.encode(public_key),
+                    x: base64::encode_config(public_key, base64::URL_SAFE_NO_PAD),
                 }),
             }],
         })

libs/pageserver_api/src/upcall_api.rs

@@ -9,7 +9,7 @@ use utils::id::{NodeId, TimelineId};
 
 use crate::controller_api::NodeRegisterRequest;
 use crate::models::{LocationConfigMode, ShardImportStatus};
-use crate::shard::{ShardStripeSize, TenantShardId};
+use crate::shard::TenantShardId;
 
 /// Upcall message sent by the pageserver to the configured `control_plane_api` on
 /// startup.
@@ -36,10 +36,6 @@ pub struct ReAttachResponseTenant {
     /// Default value only for backward compat: this field should be set
     #[serde(default = "default_mode")]
     pub mode: LocationConfigMode,
-
-    // Default value only for backward compat: this field should be set
-    #[serde(default = "ShardStripeSize::default")]
-    pub stripe_size: ShardStripeSize,
 }
 #[derive(Serialize, Deserialize)]
 pub struct ReAttachResponse {

libs/posthog_client_lite/src/background_loop.rs

@@ -55,16 +55,9 @@ impl FeatureResolverBackgroundLoop {
                             continue;
                         }
                     };
-                    let project_id = this.posthog_client.config.project_id.parse::<u64>().ok();
-                    match FeatureStore::new_with_flags(resp.flags, project_id) {
-                        Ok(feature_store) => {
-                            this.feature_store.store(Arc::new(feature_store));
-                            tracing::info!("Feature flag updated");
-                        }
-                        Err(e) => {
-                            tracing::warn!("Cannot process feature flag spec: {}", e);
-                        }
-                    }
+                    let feature_store = FeatureStore::new_with_flags(resp.flags);
+                    this.feature_store.store(Arc::new(feature_store));
+                    tracing::info!("Feature flag updated");
                 }
                 tracing::info!("PostHog feature resolver stopped");
             }

libs/posthog_client_lite/src/lib.rs

@@ -39,9 +39,6 @@ pub struct LocalEvaluationResponse {
 
 #[derive(Deserialize)]
 pub struct LocalEvaluationFlag {
-    #[allow(dead_code)]
-    id: u64,
-    team_id: u64,
     key: String,
     filters: LocalEvaluationFlagFilters,
     active: bool,
@@ -110,32 +107,17 @@ impl FeatureStore {
         }
     }
 
-    pub fn new_with_flags(
-        flags: Vec<LocalEvaluationFlag>,
-        project_id: Option<u64>,
-    ) -> Result<Self, &'static str> {
+    pub fn new_with_flags(flags: Vec<LocalEvaluationFlag>) -> Self {
         let mut store = Self::new();
-        store.set_flags(flags, project_id)?;
-        Ok(store)
+        store.set_flags(flags);
+        store
     }
 
-    pub fn set_flags(
-        &mut self,
-        flags: Vec<LocalEvaluationFlag>,
-        project_id: Option<u64>,
-    ) -> Result<(), &'static str> {
+    pub fn set_flags(&mut self, flags: Vec<LocalEvaluationFlag>) {
         self.flags.clear();
         for flag in flags {
-            if let Some(project_id) = project_id {
-                if flag.team_id != project_id {
-                    return Err(
-                        "Retrieved a spec with different project id, wrong config? Discarding the feature flags.",
-                    );
-                }
-            }
             self.flags.insert(flag.key.clone(), flag);
         }
-        Ok(())
     }
 
     /// Generate a consistent hash for a user ID (e.g., tenant ID).
@@ -552,13 +534,6 @@ impl PostHogClient {
         })
     }
 
-    /// Check if the server API key is a feature flag secure API key. This key can only be
-    /// used to fetch the feature flag specs and can only be used on a undocumented API
-    /// endpoint.
-    fn is_feature_flag_secure_api_key(&self) -> bool {
-        self.config.server_api_key.starts_with("phs_")
-    }
-
     /// Fetch the feature flag specs from the server.
     ///
     /// This is unfortunately an undocumented API at:
@@ -572,22 +547,10 @@ impl PostHogClient {
     ) -> anyhow::Result<LocalEvaluationResponse> {
         // BASE_URL/api/projects/:project_id/feature_flags/local_evaluation
         // with bearer token of self.server_api_key
-        // OR
-        // BASE_URL/api/feature_flag/local_evaluation/
-        // with bearer token of feature flag specific self.server_api_key
-        let url = if self.is_feature_flag_secure_api_key() {
-            // The new feature local evaluation secure API token
-            format!(
-                "{}/api/feature_flag/local_evaluation",
-                self.config.private_api_url
-            )
-        } else {
-            // The old personal API token
-            format!(
-                "{}/api/projects/{}/feature_flags/local_evaluation",
-                self.config.private_api_url, self.config.project_id
-            )
-        };
+        let url = format!(
+            "{}/api/projects/{}/feature_flags/local_evaluation",
+            self.config.private_api_url, self.config.project_id
+        );
         let response = self
             .client
             .get(url)
@@ -840,7 +803,7 @@ mod tests {
     fn evaluate_multivariate() {
         let mut store = FeatureStore::new();
         let response: LocalEvaluationResponse = serde_json::from_str(data()).unwrap();
-        store.set_flags(response.flags, None).unwrap();
+        store.set_flags(response.flags);
 
         // This lacks the required properties and cannot be evaluated.
         let variant =
@@ -910,7 +873,7 @@ mod tests {
 
         let mut store = FeatureStore::new();
         let response: LocalEvaluationResponse = serde_json::from_str(data()).unwrap();
-        store.set_flags(response.flags, None).unwrap();
+        store.set_flags(response.flags);
 
         // This lacks the required properties and cannot be evaluated.
         let variant = store.evaluate_boolean_inner("boolean-flag", 1.00, &HashMap::new());
@@ -966,7 +929,7 @@ mod tests {
 
         let mut store = FeatureStore::new();
         let response: LocalEvaluationResponse = serde_json::from_str(data()).unwrap();
-        store.set_flags(response.flags, None).unwrap();
+        store.set_flags(response.flags);
 
         // This lacks the required properties and cannot be evaluated.
         let variant =

libs/proxy/postgres-protocol2/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2024"
 license = "MIT/Apache-2.0"
 
 [dependencies]
-base64.workspace = true
+base64 = "0.20"
 byteorder.workspace = true
 bytes.workspace = true
 fallible-iterator.workspace = true

libs/proxy/postgres-protocol2/src/authentication/sasl.rs

@@ -3,8 +3,6 @@
 use std::fmt::Write;
 use std::{io, iter, mem, str};
 
-use base64::Engine as _;
-use base64::prelude::BASE64_STANDARD;
 use hmac::{Hmac, Mac};
 use rand::{self, Rng};
 use sha2::digest::FixedOutput;
@@ -228,7 +226,7 @@ impl ScramSha256 {
 
         let (client_key, server_key) = match password {
             Credentials::Password(password) => {
-                let salt = match BASE64_STANDARD.decode(parsed.salt) {
+                let salt = match base64::decode(parsed.salt) {
                     Ok(salt) => salt,
                     Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
                 };
@@ -257,7 +255,7 @@ impl ScramSha256 {
         let mut cbind_input = vec![];
         cbind_input.extend(channel_binding.gs2_header().as_bytes());
         cbind_input.extend(channel_binding.cbind_data());
-        let cbind_input = BASE64_STANDARD.encode(&cbind_input);
+        let cbind_input = base64::encode(&cbind_input);
 
         self.message.clear();
         write!(&mut self.message, "c={},r={}", cbind_input, parsed.nonce).unwrap();
@@ -274,12 +272,7 @@ impl ScramSha256 {
             *proof ^= signature;
         }
 
-        write!(
-            &mut self.message,
-            ",p={}",
-            BASE64_STANDARD.encode(client_proof)
-        )
-        .unwrap();
+        write!(&mut self.message, ",p={}", base64::encode(client_proof)).unwrap();
 
         self.state = State::Finish {
             server_key,
@@ -313,7 +306,7 @@ impl ScramSha256 {
             ServerFinalMessage::Verifier(verifier) => verifier,
         };
 
-        let verifier = match BASE64_STANDARD.decode(verifier) {
+        let verifier = match base64::decode(verifier) {
             Ok(verifier) => verifier,
             Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
         };

libs/proxy/postgres-protocol2/src/password/mod.rs

@@ -6,8 +6,6 @@
 //! side. This is good because it ensures the cleartext password won't
 //! end up in logs pg_stat displays, etc.
 
-use base64::Engine as _;
-use base64::prelude::BASE64_STANDARD;
 use hmac::{Hmac, Mac};
 use rand::RngCore;
 use sha2::digest::FixedOutput;
@@ -85,8 +83,8 @@ pub(crate) async fn scram_sha_256_salt(
     format!(
         "SCRAM-SHA-256${}:{}${}:{}",
         SCRAM_DEFAULT_ITERATIONS,
-        BASE64_STANDARD.encode(salt),
-        BASE64_STANDARD.encode(stored_key),
-        BASE64_STANDARD.encode(server_key)
+        base64::encode(salt),
+        base64::encode(stored_key),
+        base64::encode(server_key)
     )
 }

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

@@ -1,5 +1,3 @@
-use std::io;
-
 use tokio::net::TcpStream;
 
 use crate::client::SocketConfig;
@@ -8,7 +6,7 @@ use crate::tls::MakeTlsConnect;
 use crate::{Error, cancel_query_raw, connect_socket};
 
 pub(crate) async fn cancel_query<T>(
-    config: Option<SocketConfig>,
+    config: SocketConfig,
     ssl_mode: SslMode,
     tls: T,
     process_id: i32,
@@ -17,16 +15,6 @@ 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,9 +9,16 @@ use crate::{Error, cancel_query, cancel_query_raw};
 
 /// The capability to request cancellation of in-progress queries on a
 /// connection.
-#[derive(Clone, Serialize, Deserialize)]
+#[derive(Clone)]
 pub struct CancelToken {
-    pub socket_config: Option<SocketConfig>,
+    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 ssl_mode: SslMode,
     pub process_id: i32,
     pub secret_key: i32,
@@ -36,14 +43,16 @@ impl CancelToken {
     {
         cancel_query::cancel_query(
             self.socket_config.clone(),
-            self.ssl_mode,
+            self.raw.ssl_mode,
             tls,
-            self.process_id,
-            self.secret_key,
+            self.raw.process_id,
+            self.raw.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,6 +12,7 @@ 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;
@@ -331,10 +332,12 @@ impl Client {
     /// connection associated with this client.
     pub fn cancel_token(&self) -> CancelToken {
         CancelToken {
-            socket_config: Some(self.socket_config.clone()),
-            ssl_mode: self.ssl_mode,
-            process_id: self.process_id,
-            secret_key: self.secret_key,
+            socket_config: self.socket_config.clone(),
+            raw: RawCancelToken {
+                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;
+pub use crate::cancel_token::{CancelToken, RawCancelToken};
 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(crate) tls_server_end_point: Option<Vec<u8>>,
+    pub tls_server_end_point: Option<Vec<u8>>,
 }
 
 impl ChannelBinding {

libs/remote_storage/src/azure_blob.rs

@@ -824,7 +824,6 @@ impl RemoteStorage for AzureBlobStorage {
         timestamp: SystemTime,
         done_if_after: SystemTime,
         cancel: &CancellationToken,
-        _complexity_limit: Option<NonZeroU32>,
     ) -> Result<(), TimeTravelError> {
         let msg = "PLEASE NOTE: Azure Blob storage time-travel recovery may not work as expected "
             .to_string()

libs/remote_storage/src/lib.rs

@@ -440,7 +440,6 @@ pub trait RemoteStorage: Send + Sync + 'static {
         timestamp: SystemTime,
         done_if_after: SystemTime,
         cancel: &CancellationToken,
-        complexity_limit: Option<NonZeroU32>,
     ) -> Result<(), TimeTravelError>;
 }
 
@@ -652,23 +651,22 @@ impl<Other: RemoteStorage> GenericRemoteStorage<Arc<Other>> {
         timestamp: SystemTime,
         done_if_after: SystemTime,
         cancel: &CancellationToken,
-        complexity_limit: Option<NonZeroU32>,
     ) -> Result<(), TimeTravelError> {
         match self {
             Self::LocalFs(s) => {
-                s.time_travel_recover(prefix, timestamp, done_if_after, cancel, complexity_limit)
+                s.time_travel_recover(prefix, timestamp, done_if_after, cancel)
                     .await
             }
             Self::AwsS3(s) => {
-                s.time_travel_recover(prefix, timestamp, done_if_after, cancel, complexity_limit)
+                s.time_travel_recover(prefix, timestamp, done_if_after, cancel)
                     .await
             }
             Self::AzureBlob(s) => {
-                s.time_travel_recover(prefix, timestamp, done_if_after, cancel, complexity_limit)
+                s.time_travel_recover(prefix, timestamp, done_if_after, cancel)
                     .await
             }
             Self::Unreliable(s) => {
-                s.time_travel_recover(prefix, timestamp, done_if_after, cancel, complexity_limit)
+                s.time_travel_recover(prefix, timestamp, done_if_after, cancel)
                     .await
             }
         }

libs/remote_storage/src/local_fs.rs

@@ -610,7 +610,6 @@ impl RemoteStorage for LocalFs {
         _timestamp: SystemTime,
         _done_if_after: SystemTime,
         _cancel: &CancellationToken,
-        _complexity_limit: Option<NonZeroU32>,
     ) -> Result<(), TimeTravelError> {
         Err(TimeTravelError::Unimplemented)
     }

libs/remote_storage/src/s3_bucket.rs

@@ -981,16 +981,22 @@ impl RemoteStorage for S3Bucket {
         timestamp: SystemTime,
         done_if_after: SystemTime,
         cancel: &CancellationToken,
-        complexity_limit: Option<NonZeroU32>,
     ) -> Result<(), TimeTravelError> {
         let kind = RequestKind::TimeTravel;
         let permit = self.permit(kind, cancel).await?;
 
         tracing::trace!("Target time: {timestamp:?}, done_if_after {done_if_after:?}");
 
+        // Limit the number of versions deletions, mostly so that we don't
+        // keep requesting forever if the list is too long, as we'd put the
+        // list in RAM.
+        // Building a list of 100k entries that reaches the limit roughly takes
+        // 40 seconds, and roughly corresponds to tenants of 2 TiB physical size.
+        const COMPLEXITY_LIMIT: Option<NonZeroU32> = NonZeroU32::new(100_000);
+
         let mode = ListingMode::NoDelimiter;
         let version_listing = self
-            .list_versions_with_permit(&permit, prefix, mode, complexity_limit, cancel)
+            .list_versions_with_permit(&permit, prefix, mode, COMPLEXITY_LIMIT, cancel)
             .await
             .map_err(|err| match err {
                 DownloadError::Other(e) => TimeTravelError::Other(e),

libs/remote_storage/src/simulate_failures.rs

@@ -240,12 +240,11 @@ impl RemoteStorage for UnreliableWrapper {
         timestamp: SystemTime,
         done_if_after: SystemTime,
         cancel: &CancellationToken,
-        complexity_limit: Option<NonZeroU32>,
     ) -> Result<(), TimeTravelError> {
         self.attempt(RemoteOp::TimeTravelRecover(prefix.map(|p| p.to_owned())))
             .map_err(TimeTravelError::Other)?;
         self.inner
-            .time_travel_recover(prefix, timestamp, done_if_after, cancel, complexity_limit)
+            .time_travel_recover(prefix, timestamp, done_if_after, cancel)
             .await
     }
 }

libs/remote_storage/tests/test_real_s3.rs

@@ -157,7 +157,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow:
     // No changes after recovery to t2 (no-op)
     let t_final = time_point().await;
     ctx.client
-        .time_travel_recover(None, t2, t_final, &cancel, None)
+        .time_travel_recover(None, t2, t_final, &cancel)
         .await?;
     let t2_files_recovered = list_files(&ctx.client, &cancel).await?;
     println!("after recovery to t2: {t2_files_recovered:?}");
@@ -173,7 +173,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow:
     // after recovery to t1: path1 is back, path2 has the old content
     let t_final = time_point().await;
     ctx.client
-        .time_travel_recover(None, t1, t_final, &cancel, None)
+        .time_travel_recover(None, t1, t_final, &cancel)
         .await?;
     let t1_files_recovered = list_files(&ctx.client, &cancel).await?;
     println!("after recovery to t1: {t1_files_recovered:?}");
@@ -189,7 +189,7 @@ async fn s3_time_travel_recovery_works(ctx: &mut MaybeEnabledStorage) -> anyhow:
     // after recovery to t0: everything is gone except for path1
     let t_final = time_point().await;
     ctx.client
-        .time_travel_recover(None, t0, t_final, &cancel, None)
+        .time_travel_recover(None, t0, t_final, &cancel)
         .await?;
     let t0_files_recovered = list_files(&ctx.client, &cancel).await?;
     println!("after recovery to t0: {t0_files_recovered:?}");

pageserver/ctl/src/main.rs

@@ -176,11 +176,9 @@ async fn main() -> anyhow::Result<()> {
             let config = RemoteStorageConfig::from_toml_str(&cmd.config_toml_str)?;
             let storage = remote_storage::GenericRemoteStorage::from_config(&config).await;
             let cancel = CancellationToken::new();
-            // Complexity limit: as we are running this command locally, we should have a lot of memory available, and we do not
-            // need to limit the number of versions we are going to delete.
             storage
                 .unwrap()
-                .time_travel_recover(Some(&prefix), timestamp, done_if_after, &cancel, None)
+                .time_travel_recover(Some(&prefix), timestamp, done_if_after, &cancel)
                 .await?;
         }
         Commands::Key(dkc) => dkc.execute(),

pageserver/src/basebackup_cache.rs

@@ -1,6 +1,5 @@
 use std::{collections::HashMap, sync::Arc};
 
-use anyhow::Context;
 use async_compression::tokio::write::GzipEncoder;
 use camino::{Utf8Path, Utf8PathBuf};
 use metrics::core::{AtomicU64, GenericCounter};
@@ -168,17 +167,14 @@ impl BasebackupCache {
             .join(Self::entry_filename(tenant_id, timeline_id, lsn))
     }
 
-    fn tmp_dir(&self) -> Utf8PathBuf {
-        self.data_dir.join("tmp")
-    }
-
     fn entry_tmp_path(
         &self,
         tenant_id: TenantId,
         timeline_id: TimelineId,
         lsn: Lsn,
     ) -> Utf8PathBuf {
-        self.tmp_dir()
+        self.data_dir
+            .join("tmp")
             .join(Self::entry_filename(tenant_id, timeline_id, lsn))
     }
 
@@ -198,18 +194,15 @@ impl BasebackupCache {
         Some((tenant_id, timeline_id, lsn))
     }
 
-    // Recreate the tmp directory to clear all files in it.
-    async fn clean_tmp_dir(&self) -> anyhow::Result<()> {
-        let tmp_dir = self.tmp_dir();
-        if tmp_dir.exists() {
-            tokio::fs::remove_dir_all(&tmp_dir).await?;
-        }
-        tokio::fs::create_dir_all(&tmp_dir).await?;
-        Ok(())
-    }
-
     async fn cleanup(&self) -> anyhow::Result<()> {
-        self.clean_tmp_dir().await?;
+        // Cleanup tmp directory.
+        let tmp_dir = self.data_dir.join("tmp");
+        let mut tmp_dir = tokio::fs::read_dir(&tmp_dir).await?;
+        while let Some(dir_entry) = tmp_dir.next_entry().await? {
+            if let Err(e) = tokio::fs::remove_file(dir_entry.path()).await {
+                tracing::warn!("Failed to remove basebackup cache tmp file: {:#}", e);
+            }
+        }
 
         // Remove outdated entries.
         let entries_old = self.entries.lock().unwrap().clone();
@@ -248,14 +241,16 @@ impl BasebackupCache {
     }
 
     async fn on_startup(&self) -> anyhow::Result<()> {
-        // Create data_dir if it does not exist.
-        tokio::fs::create_dir_all(&self.data_dir)
+        // Create data_dir and tmp directory if they do not exist.
+        tokio::fs::create_dir_all(&self.data_dir.join("tmp"))
             .await
-            .context("Failed to create basebackup cache data directory")?;
-
-        self.clean_tmp_dir()
-            .await
-            .context("Failed to clean tmp directory")?;
+            .map_err(|e| {
+                anyhow::anyhow!(
+                    "Failed to create basebackup cache data_dir {:?}: {:?}",
+                    self.data_dir,
+                    e
+                )
+            })?;
 
         // Read existing entries from the data_dir and add them to in-memory state.
         let mut entries = HashMap::new();
@@ -413,19 +408,6 @@ impl BasebackupCache {
             .tenant_manager
             .get_attached_tenant_shard(tenant_shard_id)?;
 
-        let feature_flag = tenant
-            .feature_resolver
-            .evaluate_boolean("enable-basebackup-cache", tenant_shard_id.tenant_id);
-
-        if feature_flag.is_err() {
-            tracing::info!(
-                tenant_id = %tenant_shard_id.tenant_id,
-                "Basebackup cache is disabled for tenant by feature flag, skipping basebackup",
-            );
-            self.prepare_skip_count.inc();
-            return Ok(());
-        }
-
         let tenant_state = tenant.current_state();
         if tenant_state != TenantState::Active {
             anyhow::bail!(
@@ -469,11 +451,6 @@ impl BasebackupCache {
         }
 
         // Move the tmp file to the final location atomically.
-        // The tmp file is fsynced, so it's guaranteed that we will not have a partial file
-        // in the main directory.
-        // It's not necessary to fsync the inode after renaming, because the worst case is that
-        // the rename operation will be rolled back on the disk failure, the entry will disappear
-        // from the main directory, and the entry access will cause a cache miss.
         let entry_path = self.entry_path(tenant_shard_id.tenant_id, timeline_id, req_lsn);
         tokio::fs::rename(&entry_tmp_path, &entry_path).await?;
 
@@ -491,17 +468,16 @@ impl BasebackupCache {
     }
 
     /// Prepares a basebackup in a temporary file.
-    /// Guarantees that the tmp file is fsynced before returning.
     async fn prepare_basebackup_tmp(
         &self,
-        entry_tmp_path: &Utf8Path,
+        emptry_tmp_path: &Utf8Path,
         timeline: &Arc<Timeline>,
         req_lsn: Lsn,
     ) -> anyhow::Result<()> {
         let ctx = RequestContext::new(TaskKind::BasebackupCache, DownloadBehavior::Download);
         let ctx = ctx.with_scope_timeline(timeline);
 
-        let file = tokio::fs::File::create(entry_tmp_path).await?;
+        let file = tokio::fs::File::create(emptry_tmp_path).await?;
         let mut writer = BufWriter::new(file);
 
         let mut encoder = GzipEncoder::with_quality(

pageserver/src/http/routes.rs

@@ -73,7 +73,6 @@ use crate::tenant::remote_timeline_client::{
 use crate::tenant::secondary::SecondaryController;
 use crate::tenant::size::ModelInputs;
 use crate::tenant::storage_layer::{IoConcurrency, LayerAccessStatsReset, LayerName};
-use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
 use crate::tenant::timeline::offload::{OffloadError, offload_timeline};
 use crate::tenant::timeline::{
     CompactFlags, CompactOptions, CompactRequest, CompactionError, MarkInvisibleRequest, Timeline,
@@ -1452,10 +1451,7 @@ async fn timeline_layer_scan_disposable_keys(
     let ctx = RequestContext::new(TaskKind::MgmtRequest, DownloadBehavior::Download)
         .with_scope_timeline(&timeline);
 
-    let guard = timeline
-        .layers
-        .read(LayerManagerLockHolder::GetLayerMapInfo)
-        .await;
+    let guard = timeline.layers.read().await;
     let Some(layer) = guard.try_get_from_key(&layer_name.clone().into()) else {
         return Err(ApiError::NotFound(
             anyhow::anyhow!("Layer {tenant_shard_id}/{timeline_id}/{layer_name} not found").into(),

pageserver/src/tenant.rs

@@ -51,7 +51,6 @@ use secondary::heatmap::{HeatMapTenant, HeatMapTimeline};
 use storage_broker::BrokerClientChannel;
 use timeline::compaction::{CompactionOutcome, GcCompactionQueue};
 use timeline::import_pgdata::ImportingTimeline;
-use timeline::layer_manager::LayerManagerLockHolder;
 use timeline::offload::{OffloadError, offload_timeline};
 use timeline::{
     CompactFlags, CompactOptions, CompactionError, PreviousHeatmap, ShutdownMode, import_pgdata,
@@ -1316,7 +1315,7 @@ impl TenantShard {
                         ancestor.is_some()
                             || timeline
                                 .layers
-                                .read(LayerManagerLockHolder::LoadLayerMap)
+                                .read()
                                 .await
                                 .layer_map()
                                 .expect(
@@ -2644,7 +2643,7 @@ impl TenantShard {
         }
         let layer_names = tline
             .layers
-            .read(LayerManagerLockHolder::Testing)
+            .read()
             .await
             .layer_map()
             .unwrap()
@@ -3159,12 +3158,7 @@ impl TenantShard {
 
         for timeline in &compact {
             // Collect L0 counts. Can't await while holding lock above.
-            if let Ok(lm) = timeline
-                .layers
-                .read(LayerManagerLockHolder::Compaction)
-                .await
-                .layer_map()
-            {
+            if let Ok(lm) = timeline.layers.read().await.layer_map() {
                 l0_counts.insert(timeline.timeline_id, lm.level0_deltas().len());
             }
         }
@@ -4906,7 +4900,7 @@ impl TenantShard {
         }
         let layer_names = tline
             .layers
-            .read(LayerManagerLockHolder::Testing)
+            .read()
             .await
             .layer_map()
             .unwrap()
@@ -6976,7 +6970,7 @@ mod tests {
             .await?;
         make_some_layers(tline.as_ref(), Lsn(0x20), &ctx).await?;
 
-        let layer_map = tline.layers.read(LayerManagerLockHolder::Testing).await;
+        let layer_map = tline.layers.read().await;
         let level0_deltas = layer_map
             .layer_map()?
             .level0_deltas()
@@ -7212,7 +7206,7 @@ mod tests {
         let lsn = Lsn(0x10);
         let inserted = bulk_insert_compact_gc(&tenant, &tline, &ctx, lsn, 50, 10000).await?;
 
-        let guard = tline.layers.read(LayerManagerLockHolder::Testing).await;
+        let guard = tline.layers.read().await;
         let lm = guard.layer_map()?;
 
         lm.dump(true, &ctx).await?;
@@ -8240,23 +8234,12 @@ mod tests {
             tline.freeze_and_flush().await?; // force create a delta layer
         }
 
-        let before_num_l0_delta_files = tline
-            .layers
-            .read(LayerManagerLockHolder::Testing)
-            .await
-            .layer_map()?
-            .level0_deltas()
-            .len();
+        let before_num_l0_delta_files =
+            tline.layers.read().await.layer_map()?.level0_deltas().len();
 
         tline.compact(&cancel, EnumSet::default(), &ctx).await?;
 
-        let after_num_l0_delta_files = tline
-            .layers
-            .read(LayerManagerLockHolder::Testing)
-            .await
-            .layer_map()?
-            .level0_deltas()
-            .len();
+        let after_num_l0_delta_files = tline.layers.read().await.layer_map()?.level0_deltas().len();
 
         assert!(
             after_num_l0_delta_files < before_num_l0_delta_files,

pageserver/src/tenant/config.rs

@@ -61,8 +61,8 @@ pub(crate) struct LocationConf {
     /// The detailed shard identity.  This structure is already scoped within
     /// a TenantShardId, but we need the full ShardIdentity to enable calculating
     /// key->shard mappings.
-    // TODO(vlad): Remove this default once all configs have a shard identity on disk.
     #[serde(default = "ShardIdentity::unsharded")]
+    #[serde(skip_serializing_if = "ShardIdentity::is_unsharded")]
     pub(crate) shard: ShardIdentity,
 
     /// The pan-cluster tenant configuration, the same on all locations
@@ -149,12 +149,7 @@ impl LocationConf {
     /// For use when attaching/re-attaching: update the generation stored in this
     /// structure.  If we were in a secondary state, promote to attached (posession
     /// of a fresh generation implies this).
-    pub(crate) fn attach_in_generation(
-        &mut self,
-        mode: AttachmentMode,
-        generation: Generation,
-        stripe_size: ShardStripeSize,
-    ) {
+    pub(crate) fn attach_in_generation(&mut self, mode: AttachmentMode, generation: Generation) {
         match &mut self.mode {
             LocationMode::Attached(attach_conf) => {
                 attach_conf.generation = generation;
@@ -168,8 +163,6 @@ impl LocationConf {
                 })
             }
         }
-
-        self.shard.stripe_size = stripe_size;
     }
 
     pub(crate) fn try_from(conf: &'_ models::LocationConfig) -> anyhow::Result<Self> {

pageserver/src/tenant/mgr.rs

@@ -51,7 +51,6 @@ use crate::tenant::config::{
 use crate::tenant::span::debug_assert_current_span_has_tenant_id;
 use crate::tenant::storage_layer::inmemory_layer;
 use crate::tenant::timeline::ShutdownMode;
-use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
 use crate::tenant::{
     AttachedTenantConf, GcError, LoadConfigError, SpawnMode, TenantShard, TenantState,
 };
@@ -129,7 +128,7 @@ pub(crate) enum ShardSelector {
 ///
 /// This represents the subset of a LocationConfig that we receive during re-attach.
 pub(crate) enum TenantStartupMode {
-    Attached((AttachmentMode, Generation, ShardStripeSize)),
+    Attached((AttachmentMode, Generation)),
     Secondary,
 }
 
@@ -143,21 +142,15 @@ impl TenantStartupMode {
         match (rart.mode, rart.r#gen) {
             (LocationConfigMode::Detached, _) => None,
             (LocationConfigMode::Secondary, _) => Some(Self::Secondary),
-            (LocationConfigMode::AttachedMulti, Some(g)) => Some(Self::Attached((
-                AttachmentMode::Multi,
-                Generation::new(g),
-                rart.stripe_size,
-            ))),
-            (LocationConfigMode::AttachedSingle, Some(g)) => Some(Self::Attached((
-                AttachmentMode::Single,
-                Generation::new(g),
-                rart.stripe_size,
-            ))),
-            (LocationConfigMode::AttachedStale, Some(g)) => Some(Self::Attached((
-                AttachmentMode::Stale,
-                Generation::new(g),
-                rart.stripe_size,
-            ))),
+            (LocationConfigMode::AttachedMulti, Some(g)) => {
+                Some(Self::Attached((AttachmentMode::Multi, Generation::new(g))))
+            }
+            (LocationConfigMode::AttachedSingle, Some(g)) => {
+                Some(Self::Attached((AttachmentMode::Single, Generation::new(g))))
+            }
+            (LocationConfigMode::AttachedStale, Some(g)) => {
+                Some(Self::Attached((AttachmentMode::Stale, Generation::new(g))))
+            }
             _ => {
                 tracing::warn!(
                     "Received invalid re-attach state for tenant {}: {rart:?}",
@@ -325,11 +318,9 @@ fn emergency_generations(
             Some((
                 *tid,
                 match &lc.mode {
-                    LocationMode::Attached(alc) => TenantStartupMode::Attached((
-                        alc.attach_mode,
-                        alc.generation,
-                        ShardStripeSize::default(),
-                    )),
+                    LocationMode::Attached(alc) => {
+                        TenantStartupMode::Attached((alc.attach_mode, alc.generation))
+                    }
                     LocationMode::Secondary(_) => TenantStartupMode::Secondary,
                 },
             ))
@@ -373,7 +364,7 @@ async fn init_load_generations(
         .iter()
         .flat_map(|(id, start_mode)| {
             match start_mode {
-                TenantStartupMode::Attached((_mode, generation, _stripe_size)) => Some(generation),
+                TenantStartupMode::Attached((_mode, generation)) => Some(generation),
                 TenantStartupMode::Secondary => None,
             }
             .map(|gen_| (*id, *gen_))
@@ -593,7 +584,7 @@ pub async fn init_tenant_mgr(
                         location_conf.mode = LocationMode::Secondary(DEFAULT_SECONDARY_CONF);
                     }
                 }
-                Some(TenantStartupMode::Attached((attach_mode, generation, stripe_size))) => {
+                Some(TenantStartupMode::Attached((attach_mode, generation))) => {
                     let old_gen_higher = match &location_conf.mode {
                         LocationMode::Attached(AttachedLocationConfig {
                             generation: old_generation,
@@ -617,7 +608,7 @@ pub async fn init_tenant_mgr(
                         // local disk content: demote to secondary rather than detaching.
                         location_conf.mode = LocationMode::Secondary(DEFAULT_SECONDARY_CONF);
                     } else {
-                        location_conf.attach_in_generation(*attach_mode, *generation, *stripe_size);
+                        location_conf.attach_in_generation(*attach_mode, *generation);
                     }
                 }
             }
@@ -1667,10 +1658,7 @@ impl TenantManager {
             let parent_timelines = timelines.keys().cloned().collect::<Vec<_>>();
             for timeline in timelines.values() {
                 tracing::info!(timeline_id=%timeline.timeline_id, "Loading list of layers to hardlink");
-                let layers = timeline
-                    .layers
-                    .read(LayerManagerLockHolder::GetLayerMapInfo)
-                    .await;
+                let layers = timeline.layers.read().await;
 
                 for layer in layers.likely_resident_layers() {
                     let relative_path = layer

pageserver/src/tenant/remote_timeline_client/upload.rs

@@ -1,7 +1,6 @@
 //! Helper functions to upload files to remote storage with a RemoteStorage
 
 use std::io::{ErrorKind, SeekFrom};
-use std::num::NonZeroU32;
 use std::time::SystemTime;
 
 use anyhow::{Context, bail};
@@ -229,25 +228,11 @@ pub(crate) async fn time_travel_recover_tenant(
         let timelines_path = super::remote_timelines_path(tenant_shard_id);
         prefixes.push(timelines_path);
     }
-
-    // Limit the number of versions deletions, mostly so that we don't
-    // keep requesting forever if the list is too long, as we'd put the
-    // list in RAM.
-    // Building a list of 100k entries that reaches the limit roughly takes
-    // 40 seconds, and roughly corresponds to tenants of 2 TiB physical size.
-    const COMPLEXITY_LIMIT: Option<NonZeroU32> = NonZeroU32::new(100_000);
-
     for prefix in &prefixes {
         backoff::retry(
             || async {
                 storage
-                    .time_travel_recover(
-                        Some(prefix),
-                        timestamp,
-                        done_if_after,
-                        cancel,
-                        COMPLEXITY_LIMIT,
-                    )
+                    .time_travel_recover(Some(prefix), timestamp, done_if_after, cancel)
                     .await
             },
             |e| !matches!(e, TimeTravelError::Other(_)),

pageserver/src/tenant/storage_layer/delta_layer.rs

@@ -1635,7 +1635,6 @@ pub(crate) mod test {
     use crate::tenant::disk_btree::tests::TestDisk;
     use crate::tenant::harness::{TIMELINE_ID, TenantHarness};
     use crate::tenant::storage_layer::{Layer, ResidentLayer};
-    use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
     use crate::tenant::{TenantShard, Timeline};
 
     /// Construct an index for a fictional delta layer and and then
@@ -2003,7 +2002,7 @@ pub(crate) mod test {
 
         let initdb_layer = timeline
             .layers
-            .read(crate::tenant::timeline::layer_manager::LayerManagerLockHolder::Testing)
+            .read()
             .await
             .likely_resident_layers()
             .next()
@@ -2079,7 +2078,7 @@ pub(crate) mod test {
 
         let new_layer = timeline
             .layers
-            .read(LayerManagerLockHolder::Testing)
+            .read()
             .await
             .likely_resident_layers()
             .find(|&x| x != &initdb_layer)

pageserver/src/tenant/storage_layer/layer/tests.rs

@@ -10,7 +10,6 @@ use super::*;
 use crate::context::DownloadBehavior;
 use crate::tenant::harness::{TenantHarness, test_img};
 use crate::tenant::storage_layer::{IoConcurrency, LayerVisibilityHint};
-use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
 
 /// Used in tests to advance a future to wanted await point, and not futher.
 const ADVANCE: std::time::Duration = std::time::Duration::from_secs(3600);
@@ -60,7 +59,7 @@ async fn smoke_test() {
     // there to avoid the timeline being illegally empty
     let (layer, dummy_layer) = {
         let mut layers = {
-            let layers = timeline.layers.read(LayerManagerLockHolder::Testing).await;
+            let layers = timeline.layers.read().await;
             layers.likely_resident_layers().cloned().collect::<Vec<_>>()
         };
 
@@ -216,7 +215,7 @@ async fn smoke_test() {
 
     // Simulate GC removing our test layer.
     {
-        let mut g = timeline.layers.write(LayerManagerLockHolder::Testing).await;
+        let mut g = timeline.layers.write().await;
 
         let layers = &[layer];
         g.open_mut().unwrap().finish_gc_timeline(layers);
@@ -262,7 +261,7 @@ async fn evict_and_wait_on_wanted_deleted() {
 
     let layer = {
         let mut layers = {
-            let layers = timeline.layers.read(LayerManagerLockHolder::Testing).await;
+            let layers = timeline.layers.read().await;
             layers.likely_resident_layers().cloned().collect::<Vec<_>>()
         };
 
@@ -306,7 +305,7 @@ async fn evict_and_wait_on_wanted_deleted() {
     // assert that once we remove the `layer` from the layer map and drop our reference,
     // the deletion of the layer in remote_storage happens.
     {
-        let mut layers = timeline.layers.write(LayerManagerLockHolder::Testing).await;
+        let mut layers = timeline.layers.write().await;
         layers.open_mut().unwrap().finish_gc_timeline(&[layer]);
     }
 
@@ -348,7 +347,7 @@ fn read_wins_pending_eviction() {
 
         let layer = {
             let mut layers = {
-                let layers = timeline.layers.read(LayerManagerLockHolder::Testing).await;
+                let layers = timeline.layers.read().await;
                 layers.likely_resident_layers().cloned().collect::<Vec<_>>()
             };
 
@@ -481,7 +480,7 @@ fn multiple_pending_evictions_scenario(name: &'static str, in_order: bool) {
 
         let layer = {
             let mut layers = {
-                let layers = timeline.layers.read(LayerManagerLockHolder::Testing).await;
+                let layers = timeline.layers.read().await;
                 layers.likely_resident_layers().cloned().collect::<Vec<_>>()
             };
 
@@ -656,7 +655,7 @@ async fn cancelled_get_or_maybe_download_does_not_cancel_eviction() {
 
     let layer = {
         let mut layers = {
-            let layers = timeline.layers.read(LayerManagerLockHolder::Testing).await;
+            let layers = timeline.layers.read().await;
             layers.likely_resident_layers().cloned().collect::<Vec<_>>()
         };
 
@@ -742,7 +741,7 @@ async fn evict_and_wait_does_not_wait_for_download() {
 
     let layer = {
         let mut layers = {
-            let layers = timeline.layers.read(LayerManagerLockHolder::Testing).await;
+            let layers = timeline.layers.read().await;
             layers.likely_resident_layers().cloned().collect::<Vec<_>>()
         };
 
@@ -863,7 +862,7 @@ async fn eviction_cancellation_on_drop() {
 
     let (evicted_layer, not_evicted) = {
         let mut layers = {
-            let mut guard = timeline.layers.write(LayerManagerLockHolder::Testing).await;
+            let mut guard = timeline.layers.write().await;
             let layers = guard.likely_resident_layers().cloned().collect::<Vec<_>>();
             // remove the layers from layermap
             guard.open_mut().unwrap().finish_gc_timeline(&layers);

pageserver/src/tenant/timeline.rs

@@ -35,11 +35,7 @@ use fail::fail_point;
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
 use handle::ShardTimelineId;
-use layer_manager::{
-    LayerManagerLockHolder, LayerManagerReadGuard, LayerManagerWriteGuard, LockedLayerManager,
-    Shutdown,
-};
-
+use layer_manager::Shutdown;
 use offload::OffloadError;
 use once_cell::sync::Lazy;
 use pageserver_api::config::tenant_conf_defaults::DEFAULT_PITR_INTERVAL;
@@ -86,6 +82,7 @@ use wal_decoder::serialized_batch::{SerializedValueBatch, ValueMeta};
 use self::delete::DeleteTimelineFlow;
 pub(super) use self::eviction_task::EvictionTaskTenantState;
 use self::eviction_task::EvictionTaskTimelineState;
+use self::layer_manager::LayerManager;
 use self::logical_size::LogicalSize;
 use self::walreceiver::{WalReceiver, WalReceiverConf};
 use super::remote_timeline_client::RemoteTimelineClient;
@@ -184,13 +181,13 @@ impl std::fmt::Display for ImageLayerCreationMode {
 
 /// Temporary function for immutable storage state refactor, ensures we are dropping mutex guard instead of other things.
 /// Can be removed after all refactors are done.
-fn drop_layer_manager_rlock(rlock: LayerManagerReadGuard<'_>) {
+fn drop_rlock<T>(rlock: tokio::sync::RwLockReadGuard<T>) {
     drop(rlock)
 }
 
 /// Temporary function for immutable storage state refactor, ensures we are dropping mutex guard instead of other things.
 /// Can be removed after all refactors are done.
-fn drop_layer_manager_wlock(rlock: LayerManagerWriteGuard<'_>) {
+fn drop_wlock<T>(rlock: tokio::sync::RwLockWriteGuard<'_, T>) {
     drop(rlock)
 }
 
@@ -244,7 +241,7 @@ pub struct Timeline {
     ///
     /// In the future, we'll be able to split up the tuple of LayerMap and `LayerFileManager`,
     /// so that e.g. on-demand-download/eviction, and layer spreading, can operate just on `LayerFileManager`.
-    pub(crate) layers: LockedLayerManager,
+    pub(crate) layers: tokio::sync::RwLock<LayerManager>,
 
     last_freeze_at: AtomicLsn,
     // Atomic would be more appropriate here.
@@ -1538,10 +1535,7 @@ impl Timeline {
     /// This method makes no distinction between local and remote layers.
     /// Hence, the result **does not represent local filesystem usage**.
     pub(crate) async fn layer_size_sum(&self) -> u64 {
-        let guard = self
-            .layers
-            .read(LayerManagerLockHolder::GetLayerMapInfo)
-            .await;
+        let guard = self.layers.read().await;
         guard.layer_size_sum()
     }
 
@@ -1851,7 +1845,7 @@ impl Timeline {
         // time, and this was missed.
         // if write_guard.is_none() { return; }
 
-        let Ok(layers_guard) = self.layers.try_read(LayerManagerLockHolder::TryFreezeLayer) else {
+        let Ok(layers_guard) = self.layers.try_read() else {
             // Don't block if the layer lock is busy
             return;
         };
@@ -2164,7 +2158,7 @@ impl Timeline {
         if let ShutdownMode::FreezeAndFlush = mode {
             let do_flush = if let Some((open, frozen)) = self
                 .layers
-                .read(LayerManagerLockHolder::Shutdown)
+                .read()
                 .await
                 .layer_map()
                 .map(|lm| (lm.open_layer.is_some(), lm.frozen_layers.len()))
@@ -2268,10 +2262,7 @@ impl Timeline {
             // Allow any remaining in-memory layers to do cleanup -- until that, they hold the gate
             // open.
             let mut write_guard = self.write_lock.lock().await;
-            self.layers
-                .write(LayerManagerLockHolder::Shutdown)
-                .await
-                .shutdown(&mut write_guard);
+            self.layers.write().await.shutdown(&mut write_guard);
         }
 
         // Finally wait until any gate-holders are complete.
@@ -2374,10 +2365,7 @@ impl Timeline {
         &self,
         reset: LayerAccessStatsReset,
     ) -> Result<LayerMapInfo, layer_manager::Shutdown> {
-        let guard = self
-            .layers
-            .read(LayerManagerLockHolder::GetLayerMapInfo)
-            .await;
+        let guard = self.layers.read().await;
         let layer_map = guard.layer_map()?;
         let mut in_memory_layers = Vec::with_capacity(layer_map.frozen_layers.len() + 1);
         if let Some(open_layer) = &layer_map.open_layer {
@@ -3244,7 +3232,7 @@ impl Timeline {
 
     /// Initialize with an empty layer map. Used when creating a new timeline.
     pub(super) fn init_empty_layer_map(&self, start_lsn: Lsn) {
-        let mut layers = self.layers.try_write(LayerManagerLockHolder::Init).expect(
+        let mut layers = self.layers.try_write().expect(
             "in the context where we call this function, no other task has access to the object",
         );
         layers
@@ -3264,10 +3252,7 @@ impl Timeline {
         use init::Decision::*;
         use init::{Discovered, DismissedLayer};
 
-        let mut guard = self
-            .layers
-            .write(LayerManagerLockHolder::LoadLayerMap)
-            .await;
+        let mut guard = self.layers.write().await;
 
         let timer = self.metrics.load_layer_map_histo.start_timer();
 
@@ -3884,10 +3869,7 @@ impl Timeline {
         &self,
         layer_name: &LayerName,
     ) -> Result<Option<Layer>, layer_manager::Shutdown> {
-        let guard = self
-            .layers
-            .read(LayerManagerLockHolder::GetLayerMapInfo)
-            .await;
+        let guard = self.layers.read().await;
         let layer = guard
             .layer_map()?
             .iter_historic_layers()
@@ -3920,10 +3902,7 @@ impl Timeline {
             return None;
         }
 
-        let guard = self
-            .layers
-            .read(LayerManagerLockHolder::GenerateHeatmap)
-            .await;
+        let guard = self.layers.read().await;
 
         // Firstly, if there's any heatmap left over from when this location
         // was a secondary, take that into account. Keep layers that are:
@@ -4021,10 +4000,7 @@ impl Timeline {
     }
 
     pub(super) async fn generate_unarchival_heatmap(&self, end_lsn: Lsn) -> PreviousHeatmap {
-        let guard = self
-            .layers
-            .read(LayerManagerLockHolder::GenerateHeatmap)
-            .await;
+        let guard = self.layers.read().await;
 
         let now = SystemTime::now();
         let mut heatmap_layers = Vec::default();
@@ -4366,7 +4342,7 @@ impl Timeline {
         query: &VersionedKeySpaceQuery,
     ) -> Result<LayerFringe, GetVectoredError> {
         let mut fringe = LayerFringe::new();
-        let guard = self.layers.read(LayerManagerLockHolder::GetPage).await;
+        let guard = self.layers.read().await;
 
         match query {
             VersionedKeySpaceQuery::Uniform { keyspace, lsn } => {
@@ -4469,7 +4445,7 @@ impl Timeline {
             // required for correctness, but avoids visiting extra layers
             // which turns out to be a perf bottleneck in some cases.
             if !unmapped_keyspace.is_empty() {
-                let guard = timeline.layers.read(LayerManagerLockHolder::GetPage).await;
+                let guard = timeline.layers.read().await;
                 guard.update_search_fringe(&unmapped_keyspace, cont_lsn, &mut fringe)?;
 
                 // It's safe to drop the layer map lock after planning the next round of reads.
@@ -4579,10 +4555,7 @@ impl Timeline {
         _guard: &tokio::sync::MutexGuard<'_, Option<TimelineWriterState>>,
         ctx: &RequestContext,
     ) -> anyhow::Result<Arc<InMemoryLayer>> {
-        let mut guard = self
-            .layers
-            .write(LayerManagerLockHolder::GetLayerForWrite)
-            .await;
+        let mut guard = self.layers.write().await;
 
         let last_record_lsn = self.get_last_record_lsn();
         ensure!(
@@ -4624,10 +4597,7 @@ impl Timeline {
         write_lock: &mut tokio::sync::MutexGuard<'_, Option<TimelineWriterState>>,
     ) -> Result<u64, FlushLayerError> {
         let frozen = {
-            let mut guard = self
-                .layers
-                .write(LayerManagerLockHolder::TryFreezeLayer)
-                .await;
+            let mut guard = self.layers.write().await;
             guard
                 .open_mut()?
                 .try_freeze_in_memory_layer(at, &self.last_freeze_at, write_lock, &self.metrics)
@@ -4668,12 +4638,7 @@ impl Timeline {
         ctx: &RequestContext,
     ) {
         // Subscribe to L0 delta layer updates, for compaction backpressure.
-        let mut watch_l0 = match self
-            .layers
-            .read(LayerManagerLockHolder::FlushLoop)
-            .await
-            .layer_map()
-        {
+        let mut watch_l0 = match self.layers.read().await.layer_map() {
             Ok(lm) => lm.watch_level0_deltas(),
             Err(Shutdown) => return,
         };
@@ -4710,7 +4675,7 @@ impl Timeline {
 
                 // Fetch the next layer to flush, if any.
                 let (layer, l0_count, frozen_count, frozen_size) = {
-                    let layers = self.layers.read(LayerManagerLockHolder::FlushLoop).await;
+                    let layers = self.layers.read().await;
                     let Ok(lm) = layers.layer_map() else {
                         info!("dropping out of flush loop for timeline shutdown");
                         return;
@@ -5006,10 +4971,7 @@ impl Timeline {
         // in-memory layer from the map now. The flushed layer is stored in
         // the mapping in `create_delta_layer`.
         {
-            let mut guard = self
-                .layers
-                .write(LayerManagerLockHolder::FlushFrozenLayer)
-                .await;
+            let mut guard = self.layers.write().await;
 
             guard.open_mut()?.finish_flush_l0_layer(
                 delta_layer_to_add.as_ref(),
@@ -5224,7 +5186,7 @@ impl Timeline {
     async fn time_for_new_image_layer(&self, partition: &KeySpace, lsn: Lsn) -> bool {
         let threshold = self.get_image_creation_threshold();
 
-        let guard = self.layers.read(LayerManagerLockHolder::Compaction).await;
+        let guard = self.layers.read().await;
         let Ok(layers) = guard.layer_map() else {
             return false;
         };
@@ -5642,7 +5604,7 @@ impl Timeline {
             if let ImageLayerCreationMode::Force = mode {
                 // When forced to create image layers, we might try and create them where they already
                 // exist.  This mode is only used in tests/debug.
-                let layers = self.layers.read(LayerManagerLockHolder::Compaction).await;
+                let layers = self.layers.read().await;
                 if layers.contains_key(&PersistentLayerKey {
                     key_range: img_range.clone(),
                     lsn_range: PersistentLayerDesc::image_layer_lsn_range(lsn),
@@ -5767,7 +5729,7 @@ impl Timeline {
 
         let image_layers = batch_image_writer.finish(self, ctx).await?;
 
-        let mut guard = self.layers.write(LayerManagerLockHolder::Compaction).await;
+        let mut guard = self.layers.write().await;
 
         // FIXME: we could add the images to be uploaded *before* returning from here, but right
         // now they are being scheduled outside of write lock; current way is inconsistent with
@@ -5775,7 +5737,7 @@ impl Timeline {
         guard
             .open_mut()?
             .track_new_image_layers(&image_layers, &self.metrics);
-        drop_layer_manager_wlock(guard);
+        drop_wlock(guard);
         let duration = timer.stop_and_record();
 
         // Creating image layers may have caused some previously visible layers to be covered
@@ -6145,7 +6107,7 @@ impl Timeline {
         layers_to_remove: &[Layer],
     ) -> Result<(), CompactionError> {
         let mut guard = tokio::select! {
-            guard = self.layers.write(LayerManagerLockHolder::Compaction) => guard,
+            guard = self.layers.write() => guard,
             _ = self.cancel.cancelled() => {
                 return Err(CompactionError::ShuttingDown);
             }
@@ -6194,7 +6156,7 @@ impl Timeline {
         self.remote_client
             .schedule_compaction_update(&remove_layers, new_deltas)?;
 
-        drop_layer_manager_wlock(guard);
+        drop_wlock(guard);
 
         Ok(())
     }
@@ -6204,7 +6166,7 @@ impl Timeline {
         mut replace_layers: Vec<(Layer, ResidentLayer)>,
         mut drop_layers: Vec<Layer>,
     ) -> Result<(), CompactionError> {
-        let mut guard = self.layers.write(LayerManagerLockHolder::Compaction).await;
+        let mut guard = self.layers.write().await;
 
         // Trim our lists in case our caller (compaction) raced with someone else (GC) removing layers: we want
         // to avoid double-removing, and avoid rewriting something that was removed.
@@ -6555,10 +6517,7 @@ impl Timeline {
         // 5. newer on-disk image layers cover the layer's whole key range
         //
         // TODO holding a write lock is too agressive and avoidable
-        let mut guard = self
-            .layers
-            .write(LayerManagerLockHolder::GarbageCollection)
-            .await;
+        let mut guard = self.layers.write().await;
         let layers = guard.layer_map()?;
         'outer: for l in layers.iter_historic_layers() {
             result.layers_total += 1;
@@ -6860,10 +6819,7 @@ impl Timeline {
         use pageserver_api::models::DownloadRemoteLayersTaskState;
 
         let remaining = {
-            let guard = self
-                .layers
-                .read(LayerManagerLockHolder::GetLayerMapInfo)
-                .await;
+            let guard = self.layers.read().await;
             let Ok(lm) = guard.layer_map() else {
                 // technically here we could look into iterating accessible layers, but downloading
                 // all layers of a shutdown timeline makes no sense regardless.
@@ -6969,7 +6925,7 @@ impl Timeline {
 impl Timeline {
     /// Returns non-remote layers for eviction.
     pub(crate) async fn get_local_layers_for_disk_usage_eviction(&self) -> DiskUsageEvictionInfo {
-        let guard = self.layers.read(LayerManagerLockHolder::Eviction).await;
+        let guard = self.layers.read().await;
         let mut max_layer_size: Option<u64> = None;
 
         let resident_layers = guard
@@ -7070,7 +7026,7 @@ impl Timeline {
         let image_layer = Layer::finish_creating(self.conf, self, desc, &path)?;
         info!("force created image layer {}", image_layer.local_path());
         {
-            let mut guard = self.layers.write(LayerManagerLockHolder::Testing).await;
+            let mut guard = self.layers.write().await;
             guard
                 .open_mut()
                 .unwrap()
@@ -7133,7 +7089,7 @@ impl Timeline {
         let delta_layer = Layer::finish_creating(self.conf, self, desc, &path)?;
         info!("force created delta layer {}", delta_layer.local_path());
         {
-            let mut guard = self.layers.write(LayerManagerLockHolder::Testing).await;
+            let mut guard = self.layers.write().await;
             guard
                 .open_mut()
                 .unwrap()
@@ -7228,7 +7184,7 @@ impl Timeline {
 
         // Link the layer to the layer map
         {
-            let mut guard = self.layers.write(LayerManagerLockHolder::Testing).await;
+            let mut guard = self.layers.write().await;
             let layer_map = guard.open_mut().unwrap();
             layer_map.force_insert_in_memory_layer(Arc::new(layer));
         }
@@ -7245,7 +7201,7 @@ impl Timeline {
         io_concurrency: IoConcurrency,
     ) -> anyhow::Result<Vec<(Key, Bytes)>> {
         let mut all_data = Vec::new();
-        let guard = self.layers.read(LayerManagerLockHolder::Testing).await;
+        let guard = self.layers.read().await;
         for layer in guard.layer_map()?.iter_historic_layers() {
             if !layer.is_delta() && layer.image_layer_lsn() == lsn {
                 let layer = guard.get_from_desc(&layer);
@@ -7274,7 +7230,7 @@ impl Timeline {
         self: &Arc<Timeline>,
     ) -> anyhow::Result<Vec<super::storage_layer::PersistentLayerKey>> {
         let mut layers = Vec::new();
-        let guard = self.layers.read(LayerManagerLockHolder::Testing).await;
+        let guard = self.layers.read().await;
         for layer in guard.layer_map()?.iter_historic_layers() {
             layers.push(layer.key());
         }
@@ -7386,7 +7342,7 @@ impl TimelineWriter<'_> {
         let l0_count = self
             .tl
             .layers
-            .read(LayerManagerLockHolder::GetLayerMapInfo)
+            .read()
             .await
             .layer_map()?
             .level0_deltas()
@@ -7605,7 +7561,6 @@ mod tests {
     use crate::tenant::harness::{TenantHarness, test_img};
     use crate::tenant::layer_map::LayerMap;
     use crate::tenant::storage_layer::{Layer, LayerName, LayerVisibilityHint};
-    use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
     use crate::tenant::timeline::{DeltaLayerTestDesc, EvictionError};
     use crate::tenant::{PreviousHeatmap, Timeline};
 
@@ -7713,7 +7668,7 @@ mod tests {
         // Evict all the layers and stash the old heatmap in the timeline.
         // This simulates a migration to a cold secondary location.
 
-        let guard = timeline.layers.read(LayerManagerLockHolder::Testing).await;
+        let guard = timeline.layers.read().await;
         let mut all_layers = Vec::new();
         let forever = std::time::Duration::from_secs(120);
         for layer in guard.likely_resident_layers() {
@@ -7835,7 +7790,7 @@ mod tests {
             })));
 
         // Evict all the layers in the previous heatmap
-        let guard = timeline.layers.read(LayerManagerLockHolder::Testing).await;
+        let guard = timeline.layers.read().await;
         let forever = std::time::Duration::from_secs(120);
         for layer in guard.likely_resident_layers() {
             layer.evict_and_wait(forever).await.unwrap();
@@ -7898,10 +7853,7 @@ mod tests {
     }
 
     async fn find_some_layer(timeline: &Timeline) -> Layer {
-        let layers = timeline
-            .layers
-            .read(LayerManagerLockHolder::GetLayerMapInfo)
-            .await;
+        let layers = timeline.layers.read().await;
         let desc = layers
             .layer_map()
             .unwrap()

pageserver/src/tenant/timeline/analysis.rs

@@ -4,7 +4,6 @@ use std::ops::Range;
 use utils::lsn::Lsn;
 
 use super::Timeline;
-use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
 
 #[derive(serde::Serialize)]
 pub(crate) struct RangeAnalysis {
@@ -25,10 +24,7 @@ impl Timeline {
 
         let num_of_l0;
         let all_layer_files = {
-            let guard = self
-                .layers
-                .read(LayerManagerLockHolder::GetLayerMapInfo)
-                .await;
+            let guard = self.layers.read().await;
             num_of_l0 = guard.layer_map().unwrap().level0_deltas().len();
             guard.all_persistent_layers()
         };

pageserver/src/tenant/timeline/compaction.rs

@@ -9,7 +9,7 @@ use std::ops::{Deref, Range};
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
-use super::layer_manager::{LayerManagerLockHolder, LayerManagerReadGuard};
+use super::layer_manager::LayerManager;
 use super::{
     CompactFlags, CompactOptions, CompactionError, CreateImageLayersError, DurationRecorder,
     GetVectoredError, ImageLayerCreationMode, LastImageLayerCreationStatus, RecordedDuration,
@@ -62,7 +62,7 @@ use crate::tenant::storage_layer::{
 use crate::tenant::tasks::log_compaction_error;
 use crate::tenant::timeline::{
     DeltaLayerWriter, ImageLayerCreationOutcome, ImageLayerWriter, IoConcurrency, Layer,
-    ResidentLayer, drop_layer_manager_rlock,
+    ResidentLayer, drop_rlock,
 };
 use crate::tenant::{DeltaLayer, MaybeOffloaded};
 use crate::virtual_file::{MaybeFatalIo, VirtualFile};
@@ -314,10 +314,7 @@ impl GcCompactionQueue {
             .unwrap_or(Lsn::INVALID);
 
         let layers = {
-            let guard = timeline
-                .layers
-                .read(LayerManagerLockHolder::GetLayerMapInfo)
-                .await;
+            let guard = timeline.layers.read().await;
             let layer_map = guard.layer_map()?;
             layer_map.iter_historic_layers().collect_vec()
         };
@@ -411,10 +408,7 @@ impl GcCompactionQueue {
         timeline: &Arc<Timeline>,
         lsn: Lsn,
     ) -> Result<u64, CompactionError> {
-        let guard = timeline
-            .layers
-            .read(LayerManagerLockHolder::GetLayerMapInfo)
-            .await;
+        let guard = timeline.layers.read().await;
         let layer_map = guard.layer_map()?;
         let layers = layer_map.iter_historic_layers().collect_vec();
         let mut size = 0;
@@ -857,7 +851,7 @@ impl KeyHistoryRetention {
         }
         let layer_generation;
         {
-            let guard = tline.layers.read(LayerManagerLockHolder::Compaction).await;
+            let guard = tline.layers.read().await;
             if !guard.contains_key(key) {
                 return false;
             }
@@ -1288,10 +1282,7 @@ impl Timeline {
             // We do the repartition on the L0-L1 boundary. All data below the boundary
             // are compacted by L0 with low read amplification, thus making the `repartition`
             // function run fast.
-            let guard = self
-                .layers
-                .read(LayerManagerLockHolder::GetLayerMapInfo)
-                .await;
+            let guard = self.layers.read().await;
             guard
                 .all_persistent_layers()
                 .iter()
@@ -1470,7 +1461,7 @@ impl Timeline {
         let latest_gc_cutoff = self.get_applied_gc_cutoff_lsn();
         let pitr_cutoff = self.gc_info.read().unwrap().cutoffs.time;
 
-        let layers = self.layers.read(LayerManagerLockHolder::Compaction).await;
+        let layers = self.layers.read().await;
         let layers_iter = layers.layer_map()?.iter_historic_layers();
         let (layers_total, mut layers_checked) = (layers_iter.len(), 0);
         for layer_desc in layers_iter {
@@ -1731,10 +1722,7 @@ impl Timeline {
         // are implicitly left visible, because LayerVisibilityHint's default is Visible, and we never modify it here.
         // Note that L0 deltas _can_ be covered by image layers, but we consider them 'visible' because we anticipate that
         // they will be subject to L0->L1 compaction in the near future.
-        let layer_manager = self
-            .layers
-            .read(LayerManagerLockHolder::GetLayerMapInfo)
-            .await;
+        let layer_manager = self.layers.read().await;
         let layer_map = layer_manager.layer_map()?;
 
         let readable_points = {
@@ -1787,7 +1775,7 @@ impl Timeline {
             };
 
             let begin = tokio::time::Instant::now();
-            let phase1_layers_locked = self.layers.read(LayerManagerLockHolder::Compaction).await;
+            let phase1_layers_locked = self.layers.read().await;
             let now = tokio::time::Instant::now();
             stats.read_lock_acquisition_micros =
                 DurationRecorder::Recorded(RecordedDuration(now - begin), now);
@@ -1815,7 +1803,7 @@ impl Timeline {
     /// Level0 files first phase of compaction, explained in the [`Self::compact_legacy`] comment.
     async fn compact_level0_phase1<'a>(
         self: &'a Arc<Self>,
-        guard: LayerManagerReadGuard<'a>,
+        guard: tokio::sync::RwLockReadGuard<'a, LayerManager>,
         mut stats: CompactLevel0Phase1StatsBuilder,
         target_file_size: u64,
         force_compaction_ignore_threshold: bool,
@@ -2041,7 +2029,7 @@ impl Timeline {
             holes
         };
         stats.read_lock_held_compute_holes_micros = stats.read_lock_held_key_sort_micros.till_now();
-        drop_layer_manager_rlock(guard);
+        drop_rlock(guard);
 
         if self.cancel.is_cancelled() {
             return Err(CompactionError::ShuttingDown);
@@ -2481,7 +2469,7 @@ impl Timeline {
 
         // Find the top of the historical layers
         let end_lsn = {
-            let guard = self.layers.read(LayerManagerLockHolder::Compaction).await;
+            let guard = self.layers.read().await;
             let layers = guard.layer_map()?;
 
             let l0_deltas = layers.level0_deltas();
@@ -3020,7 +3008,7 @@ impl Timeline {
         }
         split_key_ranges.sort();
         let all_layers = {
-            let guard = self.layers.read(LayerManagerLockHolder::Compaction).await;
+            let guard = self.layers.read().await;
             let layer_map = guard.layer_map()?;
             layer_map.iter_historic_layers().collect_vec()
         };
@@ -3124,12 +3112,12 @@ impl Timeline {
                 .await?;
             let jobs_len = jobs.len();
             for (idx, job) in jobs.into_iter().enumerate() {
-                let sub_compaction_progress = format!("{}/{}", idx + 1, jobs_len);
+                info!(
+                    "running enhanced gc bottom-most compaction, sub-compaction {}/{}",
+                    idx + 1,
+                    jobs_len
+                );
                 self.compact_with_gc_inner(cancel, job, ctx, yield_for_l0)
-                    .instrument(info_span!(
-                        "sub_compaction",
-                        sub_compaction_progress = sub_compaction_progress
-                    ))
                     .await?;
             }
             if jobs_len == 0 {
@@ -3197,10 +3185,7 @@ impl Timeline {
         // 1. If a layer is in the selection, all layers below it are in the selection.
         // 2. Inferred from (1), for each key in the layer selection, the value can be reconstructed only with the layers in the layer selection.
         let job_desc = {
-            let guard = self
-                .layers
-                .read(LayerManagerLockHolder::GarbageCollection)
-                .await;
+            let guard = self.layers.read().await;
             let layers = guard.layer_map()?;
             let gc_info = self.gc_info.read().unwrap();
             let mut retain_lsns_below_horizon = Vec::new();
@@ -3971,10 +3956,7 @@ impl Timeline {
 
         // First, do a sanity check to ensure the newly-created layer map does not contain overlaps.
         let all_layers = {
-            let guard = self
-                .layers
-                .read(LayerManagerLockHolder::GarbageCollection)
-                .await;
+            let guard = self.layers.read().await;
             let layer_map = guard.layer_map()?;
             layer_map.iter_historic_layers().collect_vec()
         };
@@ -4038,10 +4020,7 @@ impl Timeline {
             let update_guard = self.gc_compaction_layer_update_lock.write().await;
             // Acquiring the update guard ensures current read operations end and new read operations are blocked.
             // TODO: can we use `latest_gc_cutoff` Rcu to achieve the same effect?
-            let mut guard = self
-                .layers
-                .write(LayerManagerLockHolder::GarbageCollection)
-                .await;
+            let mut guard = self.layers.write().await;
             guard
                 .open_mut()?
                 .finish_gc_compaction(&layer_selection, &compact_to, &self.metrics);
@@ -4109,11 +4088,7 @@ impl TimelineAdaptor {
 
     pub async fn flush_updates(&mut self) -> Result<(), CompactionError> {
         let layers_to_delete = {
-            let guard = self
-                .timeline
-                .layers
-                .read(LayerManagerLockHolder::Compaction)
-                .await;
+            let guard = self.timeline.layers.read().await;
             self.layers_to_delete
                 .iter()
                 .map(|x| guard.get_from_desc(x))
@@ -4158,11 +4133,7 @@ impl CompactionJobExecutor for TimelineAdaptor {
     ) -> anyhow::Result<Vec<OwnArc<PersistentLayerDesc>>> {
         self.flush_updates().await?;
 
-        let guard = self
-            .timeline
-            .layers
-            .read(LayerManagerLockHolder::Compaction)
-            .await;
+        let guard = self.timeline.layers.read().await;
         let layer_map = guard.layer_map()?;
 
         let result = layer_map
@@ -4201,11 +4172,7 @@ impl CompactionJobExecutor for TimelineAdaptor {
         // this is a lot more complex than a simple downcast...
         if layer.is_delta() {
             let l = {
-                let guard = self
-                    .timeline
-                    .layers
-                    .read(LayerManagerLockHolder::Compaction)
-                    .await;
+                let guard = self.timeline.layers.read().await;
                 guard.get_from_desc(layer)
             };
             let result = l.download_and_keep_resident(ctx).await?;

pageserver/src/tenant/timeline/detach_ancestor.rs

@@ -19,7 +19,7 @@ use utils::id::TimelineId;
 use utils::lsn::Lsn;
 use utils::sync::gate::GateError;
 
-use super::layer_manager::{LayerManager, LayerManagerLockHolder};
+use super::layer_manager::LayerManager;
 use super::{FlushLayerError, Timeline};
 use crate::context::{DownloadBehavior, RequestContext};
 use crate::task_mgr::TaskKind;
@@ -199,10 +199,7 @@ pub(crate) async fn generate_tombstone_image_layer(
     let image_lsn = ancestor_lsn;
 
     {
-        let layers = detached
-            .layers
-            .read(LayerManagerLockHolder::DetachAncestor)
-            .await;
+        let layers = detached.layers.read().await;
         for layer in layers.all_persistent_layers() {
             if !layer.is_delta
                 && layer.lsn_range.start == image_lsn
@@ -426,7 +423,7 @@ pub(super) async fn prepare(
         // we do not need to start from our layers, because they can only be layers that come
         // *after* ancestor_lsn
         let layers = tokio::select! {
-            guard = ancestor.layers.read(LayerManagerLockHolder::DetachAncestor) => guard,
+            guard = ancestor.layers.read() => guard,
             _ = detached.cancel.cancelled() => {
                 return Err(ShuttingDown);
             }
@@ -872,12 +869,7 @@ async fn remote_copy(
 
                 // Double check that the file is orphan (probably from an earlier attempt), then delete it
                 let key = file_name.clone().into();
-                if adoptee
-                    .layers
-                    .read(LayerManagerLockHolder::DetachAncestor)
-                    .await
-                    .contains_key(&key)
-                {
+                if adoptee.layers.read().await.contains_key(&key) {
                     // We are supposed to filter out such cases before coming to this function
                     return Err(Error::Prepare(anyhow::anyhow!(
                         "layer file {file_name} already present and inside layer map"

pageserver/src/tenant/timeline/eviction_task.rs

@@ -33,7 +33,6 @@ use crate::tenant::size::CalculateSyntheticSizeError;
 use crate::tenant::storage_layer::LayerVisibilityHint;
 use crate::tenant::tasks::{BackgroundLoopKind, BackgroundLoopSemaphorePermit, sleep_random};
 use crate::tenant::timeline::EvictionError;
-use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
 use crate::tenant::{LogicalSizeCalculationCause, TenantShard};
 
 #[derive(Default)]
@@ -209,7 +208,7 @@ impl Timeline {
 
         let mut js = tokio::task::JoinSet::new();
         {
-            let guard = self.layers.read(LayerManagerLockHolder::Eviction).await;
+            let guard = self.layers.read().await;
 
             guard
                 .likely_resident_layers()

pageserver/src/tenant/timeline/import_pgdata.rs

@@ -15,7 +15,6 @@ use super::{Timeline, TimelineDeleteProgress};
 use crate::context::RequestContext;
 use crate::controller_upcall_client::{StorageControllerUpcallApi, StorageControllerUpcallClient};
 use crate::tenant::metadata::TimelineMetadata;
-use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
 
 mod flow;
 mod importbucket_client;
@@ -164,10 +163,7 @@ async fn prepare_import(
     info!("wipe the slate clean");
     {
         // TODO: do we need to hold GC lock for this?
-        let mut guard = timeline
-            .layers
-            .write(LayerManagerLockHolder::ImportPgData)
-            .await;
+        let mut guard = timeline.layers.write().await;
         assert!(
             guard.layer_map()?.open_layer.is_none(),
             "while importing, there should be no in-memory layer" // this just seems like a good place to assert it

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

@@ -56,7 +56,6 @@ use crate::pgdatadir_mapping::{
 };
 use crate::task_mgr::TaskKind;
 use crate::tenant::storage_layer::{AsLayerDesc, ImageLayerWriter, Layer};
-use crate::tenant::timeline::layer_manager::LayerManagerLockHolder;
 
 pub async fn run(
     timeline: Arc<Timeline>,
@@ -985,10 +984,7 @@ impl ChunkProcessingJob {
             let (desc, path) = writer.finish(ctx).await?;
 
             {
-                let guard = timeline
-                    .layers
-                    .read(LayerManagerLockHolder::ImportPgData)
-                    .await;
+                let guard = timeline.layers.read().await;
                 let existing_layer = guard.try_get_from_key(&desc.key());
                 if let Some(layer) = existing_layer {
                     if layer.metadata().generation == timeline.generation {
@@ -1011,10 +1007,7 @@ impl ChunkProcessingJob {
         // certain that the existing layer is identical to the new one, so in that case
         // we replace the old layer with the one we just generated.
 
-        let mut guard = timeline
-            .layers
-            .write(LayerManagerLockHolder::ImportPgData)
-            .await;
+        let mut guard = timeline.layers.write().await;
 
         let existing_layer = guard
             .try_get_from_key(&resident_layer.layer_desc().key())
@@ -1043,7 +1036,7 @@ impl ChunkProcessingJob {
             }
         }
 
-        crate::tenant::timeline::drop_layer_manager_wlock(guard);
+        crate::tenant::timeline::drop_wlock(guard);
 
         timeline
             .remote_client

pageserver/src/tenant/timeline/layer_manager.rs

@@ -1,8 +1,5 @@
 use std::collections::HashMap;
-use std::mem::ManuallyDrop;
-use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
-use std::time::Duration;
 
 use anyhow::{Context, bail, ensure};
 use itertools::Itertools;
@@ -23,155 +20,6 @@ use crate::tenant::storage_layer::{
     PersistentLayerKey, ReadableLayerWeak, ResidentLayer,
 };
 
-/// Warn if the lock was held for longer than this threshold.
-/// It's very generous and we should bring this value down over time.
-const LAYER_MANAGER_LOCK_WARN_THRESHOLD: Duration = Duration::from_secs(5);
-const LAYER_MANAGER_LOCK_READ_WARN_THRESHOLD: Duration = Duration::from_secs(30);
-
-/// Describes the operation that is holding the layer manager lock
-#[derive(Debug, Clone, Copy, strum_macros::Display)]
-#[strum(serialize_all = "kebab_case")]
-pub(crate) enum LayerManagerLockHolder {
-    GetLayerMapInfo,
-    GenerateHeatmap,
-    GetPage,
-    Init,
-    LoadLayerMap,
-    GetLayerForWrite,
-    TryFreezeLayer,
-    FlushFrozenLayer,
-    FlushLoop,
-    Compaction,
-    GarbageCollection,
-    Shutdown,
-    ImportPgData,
-    DetachAncestor,
-    Eviction,
-    #[cfg(test)]
-    Testing,
-}
-
-/// Wrapper for the layer manager that tracks the amount of time during which
-/// it was held under read or write lock
-#[derive(Default)]
-pub(crate) struct LockedLayerManager {
-    locked: tokio::sync::RwLock<LayerManager>,
-}
-
-pub(crate) struct LayerManagerReadGuard<'a> {
-    guard: ManuallyDrop<tokio::sync::RwLockReadGuard<'a, LayerManager>>,
-    acquired_at: std::time::Instant,
-    holder: LayerManagerLockHolder,
-}
-
-pub(crate) struct LayerManagerWriteGuard<'a> {
-    guard: ManuallyDrop<tokio::sync::RwLockWriteGuard<'a, LayerManager>>,
-    acquired_at: std::time::Instant,
-    holder: LayerManagerLockHolder,
-}
-
-impl Drop for LayerManagerReadGuard<'_> {
-    fn drop(&mut self) {
-        // Drop the lock first, before potentially warning if it was held for too long.
-        // SAFETY: ManuallyDrop in Drop implementation
-        unsafe { ManuallyDrop::drop(&mut self.guard) };
-
-        let held_for = self.acquired_at.elapsed();
-        if held_for >= LAYER_MANAGER_LOCK_READ_WARN_THRESHOLD {
-            tracing::warn!(
-                holder=%self.holder,
-                "Layer manager read lock held for {}s",
-                held_for.as_secs_f64(),
-            );
-        }
-    }
-}
-
-impl Drop for LayerManagerWriteGuard<'_> {
-    fn drop(&mut self) {
-        // Drop the lock first, before potentially warning if it was held for too long.
-        // SAFETY: ManuallyDrop in Drop implementation
-        unsafe { ManuallyDrop::drop(&mut self.guard) };
-
-        let held_for = self.acquired_at.elapsed();
-        if held_for >= LAYER_MANAGER_LOCK_WARN_THRESHOLD {
-            tracing::warn!(
-                holder=%self.holder,
-                "Layer manager write lock held for {}s",
-                held_for.as_secs_f64(),
-            );
-        }
-    }
-}
-
-impl Deref for LayerManagerReadGuard<'_> {
-    type Target = LayerManager;
-
-    fn deref(&self) -> &Self::Target {
-        self.guard.deref()
-    }
-}
-
-impl Deref for LayerManagerWriteGuard<'_> {
-    type Target = LayerManager;
-
-    fn deref(&self) -> &Self::Target {
-        self.guard.deref()
-    }
-}
-
-impl DerefMut for LayerManagerWriteGuard<'_> {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        self.guard.deref_mut()
-    }
-}
-
-impl LockedLayerManager {
-    pub(crate) async fn read(&self, holder: LayerManagerLockHolder) -> LayerManagerReadGuard {
-        let guard = ManuallyDrop::new(self.locked.read().await);
-        LayerManagerReadGuard {
-            guard,
-            acquired_at: std::time::Instant::now(),
-            holder,
-        }
-    }
-
-    pub(crate) fn try_read(
-        &self,
-        holder: LayerManagerLockHolder,
-    ) -> Result<LayerManagerReadGuard, tokio::sync::TryLockError> {
-        let guard = ManuallyDrop::new(self.locked.try_read()?);
-
-        Ok(LayerManagerReadGuard {
-            guard,
-            acquired_at: std::time::Instant::now(),
-            holder,
-        })
-    }
-
-    pub(crate) async fn write(&self, holder: LayerManagerLockHolder) -> LayerManagerWriteGuard {
-        let guard = ManuallyDrop::new(self.locked.write().await);
-        LayerManagerWriteGuard {
-            guard,
-            acquired_at: std::time::Instant::now(),
-            holder,
-        }
-    }
-
-    pub(crate) fn try_write(
-        &self,
-        holder: LayerManagerLockHolder,
-    ) -> Result<LayerManagerWriteGuard, tokio::sync::TryLockError> {
-        let guard = ManuallyDrop::new(self.locked.try_write()?);
-
-        Ok(LayerManagerWriteGuard {
-            guard,
-            acquired_at: std::time::Instant::now(),
-            holder,
-        })
-    }
-}
-
 /// Provides semantic APIs to manipulate the layer map.
 pub(crate) enum LayerManager {
     /// Open as in not shutdown layer manager; we still have in-memory layers and we can manipulate

pgxn/neon/neon--1.6--1.5.sql

@@ -2,6 +2,6 @@ DROP FUNCTION IF EXISTS get_prewarm_info(out total_pages integer, out prewarmed_
 
 DROP FUNCTION IF EXISTS get_local_cache_state(max_chunks integer);
 
-DROP FUNCTION IF EXISTS prewarm_local_cache(state bytea, n_workers integer);
+DROP FUNCTION IF EXISTS prewarm_local_cache(state bytea, n_workers integer default 1);
 
 

poetry.lock

@@ -3051,19 +3051,19 @@ files = [
 
 [[package]]
 name = "requests"
-version = "2.32.4"
+version = "2.32.3"
 description = "Python HTTP for Humans."
 optional = false
 python-versions = ">=3.8"
 groups = ["main"]
 files = [
-    {file = "requests-2.32.4-py3-none-any.whl", hash = "sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c"},
-    {file = "requests-2.32.4.tar.gz", hash = "sha256:27d0316682c8a29834d3264820024b62a36942083d52caf2f14c0591336d3422"},
+    {file = "requests-2.32.3-py3-none-any.whl", hash = "sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6"},
+    {file = "requests-2.32.3.tar.gz", hash = "sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760"},
 ]
 
 [package.dependencies]
 certifi = ">=2017.4.17"
-charset_normalizer = ">=2,<4"
+charset-normalizer = ">=2,<4"
 idna = ">=2.5,<4"
 urllib3 = ">=1.21.1,<3"
 
@@ -3846,4 +3846,4 @@ cffi = ["cffi (>=1.11)"]
 [metadata]
 lock-version = "2.1"
 python-versions = "^3.11"
-content-hash = "bd93313f110110aa53b24a3ed47ba2d7f60e2c658a79cdff7320fed1bb1b57b5"
+content-hash = "7ab1e7b975af34b3271b7c6018fa22a261d3f73c7c0a0403b6b2bb86b5fbd36e"

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

proxy/src/auth/backend/console_redirect.rs

@@ -17,7 +17,7 @@ use crate::error::{ReportableError, UserFacingError};
 use crate::pqproto::BeMessage;
 use crate::proxy::NeonOptions;
 use crate::proxy::wake_compute::WakeComputeBackend;
-use crate::stream::PqStream;
+use crate::stream::PqFeStream;
 use crate::types::RoleName;
 use crate::{auth, compute, waiters};
 
@@ -96,7 +96,7 @@ impl ConsoleRedirectBackend {
         &self,
         ctx: &RequestContext,
         auth_config: &'static AuthenticationConfig,
-        client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+        client: &mut PqFeStream<impl AsyncRead + AsyncWrite + Unpin>,
     ) -> auth::Result<(ConsoleRedirectNodeInfo, AuthInfo, ComputeUserInfo)> {
         authenticate(ctx, auth_config, &self.console_uri, client)
             .await
@@ -122,7 +122,7 @@ async fn authenticate(
     ctx: &RequestContext,
     auth_config: &'static AuthenticationConfig,
     link_uri: &reqwest::Url,
-    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    client: &mut PqFeStream<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::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut stream::PqFeStream<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::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut stream::PqFeStream<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/jwt.rs

@@ -4,8 +4,6 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
 use arc_swap::ArcSwapOption;
-use base64::Engine as _;
-use base64::prelude::BASE64_URL_SAFE_NO_PAD;
 use clashmap::ClashMap;
 use jose_jwk::crypto::KeyInfo;
 use reqwest::{Client, redirect};
@@ -349,17 +347,17 @@ impl JwkCacheEntryLock {
             .split_once('.')
             .ok_or(JwtEncodingError::InvalidCompactForm)?;
 
-        let header = BASE64_URL_SAFE_NO_PAD.decode(header)?;
+        let header = base64::decode_config(header, base64::URL_SAFE_NO_PAD)?;
         let header = serde_json::from_slice::<JwtHeader<'_>>(&header)?;
 
-        let payloadb = BASE64_URL_SAFE_NO_PAD.decode(payload)?;
+        let payloadb = base64::decode_config(payload, base64::URL_SAFE_NO_PAD)?;
         let payload = serde_json::from_slice::<JwtPayload<'_>>(&payloadb)?;
 
         if let Some(iss) = &payload.issuer {
             ctx.set_jwt_issuer(iss.as_ref().to_owned());
         }
 
-        let sig = BASE64_URL_SAFE_NO_PAD.decode(signature)?;
+        let sig = base64::decode_config(signature, base64::URL_SAFE_NO_PAD)?;
 
         let kid = header.key_id.ok_or(JwtError::MissingKeyId)?;
 
@@ -798,6 +796,7 @@ mod tests {
     use std::net::SocketAddr;
     use std::time::SystemTime;
 
+    use base64::URL_SAFE_NO_PAD;
     use bytes::Bytes;
     use http::Response;
     use http_body_util::Full;
@@ -872,8 +871,9 @@ mod tests {
             key_id: Some(Cow::Owned(kid)),
         };
 
-        let header = BASE64_URL_SAFE_NO_PAD.encode(serde_json::to_string(&header).unwrap());
-        let body = BASE64_URL_SAFE_NO_PAD.encode(serde_json::to_string(&body).unwrap());
+        let header =
+            base64::encode_config(serde_json::to_string(&header).unwrap(), URL_SAFE_NO_PAD);
+        let body = base64::encode_config(serde_json::to_string(&body).unwrap(), URL_SAFE_NO_PAD);
 
         format!("{header}.{body}")
     }
@@ -883,7 +883,7 @@ mod tests {
 
         let payload = build_jwt_payload(kid, jose_jwa::Signing::Es256);
         let sig: Signature = SigningKey::from(key).sign(payload.as_bytes());
-        let sig = BASE64_URL_SAFE_NO_PAD.encode(sig.to_bytes());
+        let sig = base64::encode_config(sig.to_bytes(), URL_SAFE_NO_PAD);
 
         format!("{payload}.{sig}")
     }
@@ -893,7 +893,7 @@ mod tests {
 
         let payload = build_custom_jwt_payload(kid, body, jose_jwa::Signing::Es256);
         let sig: Signature = SigningKey::from(key).sign(payload.as_bytes());
-        let sig = BASE64_URL_SAFE_NO_PAD.encode(sig.to_bytes());
+        let sig = base64::encode_config(sig.to_bytes(), URL_SAFE_NO_PAD);
 
         format!("{payload}.{sig}")
     }
@@ -904,7 +904,7 @@ mod tests {
 
         let payload = build_jwt_payload(kid, jose_jwa::Signing::Rs256);
         let sig = SigningKey::<sha2::Sha256>::new(key).sign(payload.as_bytes());
-        let sig = BASE64_URL_SAFE_NO_PAD.encode(sig.to_bytes());
+        let sig = base64::encode_config(sig.to_bytes(), URL_SAFE_NO_PAD);
 
         format!("{payload}.{sig}")
     }

proxy/src/auth/backend/mod.rs

@@ -201,7 +201,7 @@ async fn auth_quirks(
     ctx: &RequestContext,
     api: &impl control_plane::ControlPlaneApi,
     user_info: ComputeUserInfoMaybeEndpoint,
-    client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut stream::PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     allow_cleartext: bool,
     config: &'static AuthenticationConfig,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
@@ -267,7 +267,7 @@ async fn authenticate_with_secret(
     ctx: &RequestContext,
     secret: AuthSecret,
     info: ComputeUserInfo,
-    client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+    client: &mut stream::PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
     unauthenticated_password: Option<Vec<u8>>,
     allow_cleartext: bool,
     config: &'static AuthenticationConfig,
@@ -318,7 +318,7 @@ impl<'a> Backend<'a, ComputeUserInfoMaybeEndpoint> {
     pub(crate) async fn authenticate(
         self,
         ctx: &RequestContext,
-        client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+        client: &mut stream::PqFeStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
         allow_cleartext: bool,
         config: &'static AuthenticationConfig,
         endpoint_rate_limiter: Arc<EndpointRateLimiter>,
@@ -446,7 +446,7 @@ mod tests {
     use crate::rate_limiter::EndpointRateLimiter;
     use crate::scram::ServerSecret;
     use crate::scram::threadpool::ThreadPool;
-    use crate::stream::{PqStream, Stream};
+    use crate::stream::{PqFeStream, Stream};
 
     struct Auth {
         ips: Vec<IpPattern>,
@@ -522,7 +522,7 @@ mod tests {
     #[tokio::test]
     async fn auth_quirks_scram() {
         let (mut client, server) = tokio::io::duplex(1024);
-        let mut stream = PqStream::new_skip_handshake(Stream::from_raw(server));
+        let mut stream = PqFeStream::new_skip_handshake(Stream::from_raw(server));
 
         let ctx = RequestContext::test();
         let api = Auth {
@@ -604,7 +604,7 @@ mod tests {
     #[tokio::test]
     async fn auth_quirks_cleartext() {
         let (mut client, server) = tokio::io::duplex(1024);
-        let mut stream = PqStream::new_skip_handshake(Stream::from_raw(server));
+        let mut stream = PqFeStream::new_skip_handshake(Stream::from_raw(server));
 
         let ctx = RequestContext::test();
         let api = Auth {
@@ -658,7 +658,7 @@ mod tests {
     #[tokio::test]
     async fn auth_quirks_password_hack() {
         let (mut client, server) = tokio::io::duplex(1024);
-        let mut stream = PqStream::new_skip_handshake(Stream::from_raw(server));
+        let mut stream = PqFeStream::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::{PqStream, Stream};
+use crate::stream::{PqFeStream, 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 PqStream<Stream<S>>,
+    stream: &'a mut PqFeStream<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 PqStream<Stream<S>>, method: M) -> Self {
+    pub(crate) fn new(stream: &'a mut PqFeStream<Stream<S>>, method: M) -> Self {
         let tls_server_end_point = stream.get_ref().tls_server_end_point();
 
         Self {

proxy/src/batch.rs

@@ -0,0 +1,145 @@
+//! 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, None)),
+        Arc::new(CancellationHandler::new(&config.connect_to_compute)),
         endpoint_rate_limiter,
     );
 

proxy/src/binary/pg_sni_router.rs

@@ -29,7 +29,7 @@ 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::{PqStream, Stream};
+use crate::stream::{PqFeStream, Stream};
 use crate::util::run_until_cancelled;
 
 project_git_version!(GIT_VERSION);
@@ -262,7 +262,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) = PqStream::parse_startup(Stream::from_raw(raw_stream)).await?;
+    let (mut stream, msg) = PqFeStream::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

@@ -11,19 +11,18 @@ use anyhow::Context;
 use anyhow::{bail, ensure};
 use arc_swap::ArcSwapOption;
 use futures::future::Either;
-use itertools::{Itertools, Position};
-use rand::{Rng, thread_rng};
 use remote_storage::RemoteStorageConfig;
 use tokio::net::TcpListener;
 use tokio::task::JoinSet;
 use tokio_util::sync::CancellationToken;
-use tracing::{Instrument, error, info, warn};
+use tracing::{Instrument, info, warn};
 use utils::sentry_init::init_sentry;
 use utils::{project_build_tag, project_git_version};
 
 use crate::auth::backend::jwt::JwkCache;
 use crate::auth::backend::{ConsoleRedirectBackend, MaybeOwned};
-use crate::cancellation::{CancellationHandler, handle_cancel_messages};
+use crate::batch::BatchQueue;
+use crate::cancellation::{CancellationHandler, CancellationProcessor};
 use crate::config::{
     self, AuthenticationConfig, CacheOptions, ComputeConfig, HttpConfig, ProjectInfoCacheOptions,
     ProxyConfig, ProxyProtocolV2, remote_storage_from_toml,
@@ -316,7 +315,7 @@ pub async fn run() -> anyhow::Result<()> {
     let jemalloc = match crate::jemalloc::MetricRecorder::new() {
         Ok(t) => Some(t),
         Err(e) => {
-            error!(error = ?e, "could not start jemalloc metrics loop");
+            tracing::error!(error = ?e, "could not start jemalloc metrics loop");
             None
         }
     };
@@ -392,13 +391,7 @@ pub async fn run() -> anyhow::Result<()> {
         .as_ref()
         .map(|redis_publisher| RedisKVClient::new(redis_publisher.clone(), redis_rps_limit));
 
-    // 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 cancellation_handler = Arc::new(CancellationHandler::new(&config.connect_to_compute));
 
     let endpoint_rate_limiter = Arc::new(EndpointRateLimiter::new_with_shards(
         RateBucketInfo::to_leaky_bucket(&args.endpoint_rps_limit)
@@ -522,44 +515,19 @@ pub async fn run() -> anyhow::Result<()> {
                 }
             }
 
-            // Try to connect to Redis 3 times with 1 + (0..0.1) second interval.
-            // This prevents immediate exit and pod restart,
-            // which can cause hammering of the redis in case of connection issues.
             if let Some(mut redis_kv_client) = redis_kv_client {
-                for attempt in (0..3).with_position() {
-                    match redis_kv_client.try_connect().await {
-                        Ok(()) => {
-                            info!("Connected to Redis KV client");
-                            maintenance_tasks.spawn(async move {
-                                handle_cancel_messages(
-                                    &mut redis_kv_client,
-                                    rx_cancel,
-                                    args.cancellation_batch_size,
-                                )
-                                .await?;
+                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,
+                    }));
 
-                                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`,
-                                // so let's wait forever instead.
-                                std::future::pending().await
-                            });
-                            break;
-                        }
-                        Err(e) => {
-                            error!("Failed to connect to Redis KV client: {e}");
-                            if matches!(attempt, Position::Last(_)) {
-                                bail!(
-                                    "Failed to connect to Redis KV client after {} attempts",
-                                    attempt.into_inner()
-                                );
-                            }
-                            let jitter = thread_rng().gen_range(0..100);
-                            tokio::time::sleep(Duration::from_millis(1000 + jitter)).await;
-                        }
-                    }
-                }
+                    // `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`,
+                    // so let's wait forever instead.
+                    std::future::pending().await
+                });
             }
 
             if let Some(regional_redis_client) = regional_redis_client {

proxy/src/cancellation.rs

@@ -1,19 +1,23 @@
+use std::convert::Infallible;
 use std::net::{IpAddr, SocketAddr};
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
+use std::time::Duration;
 
-use anyhow::{Context, anyhow};
+use anyhow::anyhow;
+use futures::FutureExt;
 use ipnet::{IpNet, Ipv4Net, Ipv6Net};
-use postgres_client::CancelToken;
+use postgres_client::RawCancelToken;
 use postgres_client::tls::MakeTlsConnect;
 use redis::{Cmd, FromRedisValue, Value};
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 use tokio::net::TcpStream;
-use tokio::sync::{mpsc, oneshot};
-use tracing::{debug, error, info, warn};
+use tokio::time::timeout;
+use tracing::{debug, error, info};
 
 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;
@@ -27,46 +31,36 @@ use crate::redis::kv_ops::RedisKVClient;
 
 type IpSubnetKey = IpNet;
 
-const CANCEL_KEY_TTL: i64 = 1_209_600; // 2 weeks cancellation key expire time
+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);
 
 // Message types for sending through mpsc channel
 pub enum CancelKeyOp {
     StoreCancelKey {
-        key: String,
-        field: String,
-        value: String,
-        resp_tx: Option<oneshot::Sender<anyhow::Result<()>>>,
-        _guard: CancelChannelSizeGuard<'static>,
-        expire: i64, // TTL for key
+        key: CancelKeyData,
+        value: Box<str>,
+        expire: std::time::Duration,
     },
     GetCancelData {
-        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>,
+        key: CancelKeyData,
     },
 }
 
 pub struct Pipeline {
     inner: redis::Pipeline,
-    replies: Vec<CancelReplyOp>,
+    replies: usize,
 }
 
 impl Pipeline {
     fn with_capacity(n: usize) -> Self {
         Self {
             inner: redis::Pipeline::with_capacity(n),
-            replies: Vec::with_capacity(n),
+            replies: 0,
         }
     }
 
-    async fn execute(&mut self, client: &mut RedisKVClient) {
-        let responses = self.replies.len();
+    async fn execute(self, client: &mut RedisKVClient) -> Vec<anyhow::Result<Value>> {
+        let responses = self.replies;
         let batch_size = self.inner.len();
 
         match client.query(&self.inner).await {
@@ -76,176 +70,73 @@ impl Pipeline {
                     batch_size,
                     responses, "successfully completed cancellation jobs",
                 );
-                for (value, reply) in std::iter::zip(values, self.replies.drain(..)) {
-                    reply.send_value(value);
-                }
+                values.into_iter().map(Ok).collect()
             }
             Ok(value) => {
                 error!(batch_size, ?value, "unexpected redis return value");
-                for reply in self.replies.drain(..) {
-                    reply.send_err(anyhow!("incorrect response type from redis"));
-                }
+                std::iter::repeat_with(|| Err(anyhow!("incorrect response type from redis")))
+                    .take(responses)
+                    .collect()
             }
             Err(err) => {
-                for reply in self.replies.drain(..) {
-                    reply.send_err(anyhow!("could not send cmd to redis: {err}"));
-                }
+                std::iter::repeat_with(|| Err(anyhow!("could not send cmd to redis: {err}")))
+                    .take(responses)
+                    .collect()
             }
         }
-
-        self.inner.clear();
-        self.replies.clear();
     }
 
-    fn add_command_with_reply(&mut self, cmd: Cmd, reply: CancelReplyOp) {
+    fn add_command_with_reply(&mut self, cmd: Cmd) {
         self.inner.add_command(cmd);
-        self.replies.push(reply);
+        self.replies += 1;
     }
 
     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,
-                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::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::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);
+            CancelKeyOp::GetCancelData { key } => {
+                let key = KeyPrefix::Cancel(*key).build_redis_key();
+                pipe.add_command_with_reply(Cmd::hget(key, "data"));
             }
         }
     }
 }
 
-// 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>,
-    },
+pub struct CancellationProcessor {
+    pub client: RedisKVClient,
+    pub batch_size: usize,
 }
 
-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();
-            }
-        }
+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
     }
 
-    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(());
-        }
+    async fn apply(&mut self, batch: Vec<Self::Req>) -> Vec<Self::Res> {
+        let mut pipeline = Pipeline::with_capacity(batch.len());
 
         let batch_size = batch.len();
         debug!(batch_size, "running cancellation jobs");
 
-        for msg in batch.drain(..) {
-            msg.register(&mut pipeline);
+        for (_, op) in &batch {
+            op.register(&mut pipeline);
         }
 
-        pipeline.execute(client).await;
+        pipeline.execute(&mut self.client).await
     }
 }
 
@@ -256,7 +147,7 @@ pub struct CancellationHandler {
     compute_config: &'static ComputeConfig,
     // rate limiter of cancellation requests
     limiter: Arc<std::sync::Mutex<LeakyBucketRateLimiter<IpSubnetKey>>>,
-    tx: Option<mpsc::Sender<CancelKeyOp>>, // send messages to the redis KV client task
+    tx: OnceLock<BatchQueue<CancellationProcessor>>, // send messages to the redis KV client task
 }
 
 #[derive(Debug, Error)]
@@ -296,13 +187,10 @@ impl ReportableError for CancelError {
 }
 
 impl CancellationHandler {
-    pub fn new(
-        compute_config: &'static ComputeConfig,
-        tx: Option<mpsc::Sender<CancelKeyOp>>,
-    ) -> Self {
+    pub fn new(compute_config: &'static ComputeConfig) -> Self {
         Self {
             compute_config,
-            tx,
+            tx: OnceLock::new(),
             limiter: Arc::new(std::sync::Mutex::new(
                 LeakyBucketRateLimiter::<IpSubnetKey>::new_with_shards(
                     LeakyBucketRateLimiter::<IpSubnetKey>::DEFAULT,
@@ -312,7 +200,14 @@ impl CancellationHandler {
         }
     }
 
-    pub(crate) fn get_key(self: &Arc<Self>) -> Session {
+    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 {
         // 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.
@@ -322,14 +217,10 @@ 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,
-            redis_key,
-            cancellation_handler: Arc::clone(self),
+            cancellation_handler: self,
         }
     }
 
@@ -337,62 +228,43 @@ impl CancellationHandler {
         &self,
         key: CancelKeyData,
     ) -> Result<Option<CancelClosure>, CancelError> {
-        let prefix_key: KeyPrefix = KeyPrefix::Cancel(key);
-        let redis_key = prefix_key.build_redis_key();
+        let guard = Metrics::get()
+            .proxy
+            .cancel_channel_size
+            .guard(RedisMsgKind::HGet);
+        let op = CancelKeyOp::GetCancelData { key };
 
-        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 {
+        let Some(tx) = self.tx.get() else {
             tracing::warn!("cancellation handler is not available");
             return Err(CancelError::InternalError);
         };
 
-        tx.try_send(op)
+        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
+            })?
             .map_err(|e| {
-                tracing::warn!("failed to send GetCancelData for {key}: {e}");
-            })
-            .map_err(|()| CancelError::InternalError)?;
+                tracing::warn!("failed to receive GetCancelData response: {e}");
+                CancelError::InternalError
+            })?;
 
-        let result = resp_rx.await.map_err(|e| {
+        let cancel_state_str = String::from_owned_redis_value(result).map_err(|e| {
             tracing::warn!("failed to receive GetCancelData response: {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);
-            }
-        };
+        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: 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)
+        Ok(Some(cancel_closure))
     }
+
     /// 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.
@@ -467,10 +339,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(Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CancelClosure {
     socket_addr: SocketAddr,
-    cancel_token: CancelToken,
+    cancel_token: RawCancelToken,
     hostname: String, // for pg_sni router
     user_info: ComputeUserInfo,
 }
@@ -478,7 +350,7 @@ pub struct CancelClosure {
 impl CancelClosure {
     pub(crate) fn new(
         socket_addr: SocketAddr,
-        cancel_token: CancelToken,
+        cancel_token: RawCancelToken,
         hostname: String,
         user_info: ComputeUserInfo,
     ) -> Self {
@@ -491,7 +363,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?;
@@ -512,7 +384,6 @@ impl CancelClosure {
 pub(crate) struct Session {
     /// The user-facing key identifying this session.
     key: CancelKeyData,
-    redis_key: String,
     cancellation_handler: Arc<CancellationHandler>,
 }
 
@@ -521,60 +392,66 @@ impl Session {
         &self.key
     }
 
-    // Send the store key op to the cancellation handler and set TTL for the key
-    pub(crate) fn write_cancel_key(
+    /// Ensure the cancel key is continously refreshed,
+    /// but stop when the channel is dropped.
+    pub(crate) async fn maintain_cancel_key(
         &self,
-        cancel_closure: CancelClosure,
-    ) -> Result<(), CancelError> {
-        let Some(tx) = &self.cancellation_handler.tx else {
-            tracing::warn!("cancellation handler is not available");
-            return Err(CancelError::InternalError);
-        };
+        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;
 
-        let closure_json = serde_json::to_string(&cancel_closure).map_err(|e| {
-            tracing::warn!("failed to serialize cancel closure: {e}");
-            CancelError::InternalError
-        })?;
-
-        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),
-            expire: CANCEL_KEY_TTL,
-        };
-
-        let _ = tx.try_send(op).map_err(|e| {
-            let key = self.key;
-            tracing::warn!("failed to send StoreCancelKey for {key}: {e}");
-        });
-        Ok(())
+        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"
+            );
+        }
     }
 
-    pub(crate) fn remove_cancel_key(&self) -> Result<(), CancelError> {
-        let Some(tx) = &self.cancellation_handler.tx else {
+    // 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 {
             tracing::warn!("cancellation handler is not available");
-            return Err(CancelError::InternalError);
+            // don't exit, as we only want to exit if cancelled externally.
+            std::future::pending().await
         };
 
-        let op = CancelKeyOp::RemoveCancelKey {
-            key: self.redis_key.clone(),
-            field: "data".to_string(),
-            resp_tx: None,
-            _guard: Metrics::get()
+        let closure_json = serde_json::to_string(&cancel_closure)
+            .expect("serialising to json string should not fail")
+            .into_boxed_str();
+
+        loop {
+            let guard = Metrics::get()
                 .proxy
                 .cancel_channel_size
-                .guard(RedisMsgKind::HDel),
-        };
+                .guard(RedisMsgKind::HSet);
+            let op = CancelKeyOp::StoreCancelKey {
+                key: self.key,
+                value: closure_json.clone(),
+                expire: CANCEL_KEY_TTL,
+            };
 
-        let _ = tx.try_send(op).map_err(|e| {
-            let key = self.key;
-            tracing::warn!("failed to send RemoveCancelKey for {key}: {e}");
-        });
-        Ok(())
+            tracing::debug!(
+                src=%self.key,
+                dest=?cancel_closure.cancel_token,
+                "registering cancellation key"
+            );
+
+            if tx.call((guard, op)).await.is_ok() {
+                tokio::time::sleep(CANCEL_KEY_REFRESH).await;
+            }
+        }
     }
 }

proxy/src/compute/authenticate.rs

@@ -0,0 +1,146 @@
+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,3 +1,4 @@
+mod authenticate;
 mod tls;
 
 use std::fmt::Debug;
@@ -9,15 +10,12 @@ 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, ComputeUserInfo};
+use crate::auth::backend::ComputeCredentialKeys;
 use crate::auth::parse_endpoint_param;
-use crate::cancellation::CancelClosure;
 use crate::compute::tls::TlsError;
 use crate::config::ComputeConfig;
 use crate::context::RequestContext;
@@ -28,6 +26,7 @@ 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";
@@ -37,7 +36,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] postgres_client::Error),
+    Postgres(#[from] PostgresError),
 
     #[error("{COULD_NOT_CONNECT}: {0}")]
     TlsError(#[from] TlsError),
@@ -54,20 +53,21 @@ impl UserFacingError for ConnectionError {
         match self {
             // This helps us drop irrelevant library-specific prefixes.
             // TODO: propagate severity level and other parameters.
-            ConnectionError::Postgres(err) => match err.as_db_error() {
-                Some(err) => {
-                    let msg = err.message();
+            ConnectionError::Postgres(PostgresError::Error(err)) => {
+                let (_code, msg) = err.parse();
+                let msg = String::from_utf8_lossy(msg);
 
-                    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()
-                    }
+                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()
                 }
-                None => err.to_string(),
-            },
+            }
+            ConnectionError::Postgres(err) => 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,10 +80,12 @@ impl UserFacingError for ConnectionError {
 impl ReportableError for ConnectionError {
     fn get_error_kind(&self) -> crate::error::ErrorKind {
         match self {
-            ConnectionError::Postgres(e) if e.as_db_error().is_some() => {
-                crate::error::ErrorKind::Postgres
-            }
-            ConnectionError::Postgres(_) => crate::error::ErrorKind::Compute,
+            ConnectionError::Postgres(PostgresError::Io(_)) => crate::error::ErrorKind::Compute,
+            ConnectionError::Postgres(
+                PostgresError::Error(_)
+                | PostgresError::InvalidAuthMessage
+                | PostgresError::Unexpected(_),
+            ) => crate::error::ErrorKind::Postgres,
             ConnectionError::TlsError(_) => crate::error::ErrorKind::Compute,
             ConnectionError::WakeComputeError(e) => e.get_error_kind(),
             ConnectionError::TooManyConnectionAttempts(e) => e.get_error_kind(),
@@ -162,18 +164,6 @@ 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,
@@ -213,7 +203,7 @@ impl ConnectInfo {
     async fn connect_raw(
         &self,
         config: &ComputeConfig,
-    ) -> Result<(SocketAddr, MaybeTlsStream<TcpStream, RustlsStream>), TlsError> {
+    ) -> Result<(SocketAddr, MaybeRustlsStream<TcpStream>), TlsError> {
         let timeout = config.timeout;
 
         // wrap TcpStream::connect with timeout
@@ -265,21 +255,19 @@ impl ConnectInfo {
     }
 }
 
-type RustlsStream = <ComputeConfig as MakeTlsConnect<tokio::net::TcpStream>>::Stream;
+pub type RustlsStream<S> = <ComputeConfig as MakeTlsConnect<S>>::Stream;
+pub type MaybeRustlsStream<S> = MaybeTlsStream<S, RustlsStream<S>>;
 
-pub(crate) struct PostgresConnection {
+pub struct PostgresConnection {
     /// Socket connected to a compute node.
-    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 stream: PqBeStream<MaybeRustlsStream<TcpStream>>,
 
-    _guage: NumDbConnectionsGuard<'static>,
+    pub socket_addr: SocketAddr,
+    pub hostname: String,
+    pub ssl_mode: SslMode,
+    pub aux: MetricsAuxInfo,
+
+    pub guage: NumDbConnectionsGuard<'static>,
 }
 
 impl ConnectInfo {
@@ -287,31 +275,18 @@ 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 connection = tmp_config.connect_raw(stream, NoTls).await?;
+        let stream =
+            authenticate::authenticate(stream, auth.auth.as_ref(), &auth.server_params).await?;
         drop(pause);
 
-        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("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!(
@@ -323,27 +298,13 @@ 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,
-            params: parameters,
-            delayed_notice,
-            cancel_closure,
-            aux,
-            _guage: Metrics::get().proxy.db_connections.guard(ctx.protocol()),
+            socket_addr,
+            hostname: self.host.to_string(),
+            ssl_mode: self.ssl_mode,
+            aux: aux.clone(),
+            guage: Metrics::get().proxy.db_connections.guard(ctx.protocol()),
         };
 
         Ok(connection)

proxy/src/console_redirect_proxy.rs

@@ -120,7 +120,7 @@ pub async fn task_main(
                 Ok(Some(p)) => {
                     ctx.set_success();
                     let _disconnect = ctx.log_connect();
-                    match p.proxy_pass(&config.connect_to_compute).await {
+                    match p.proxy_pass().await {
                         Ok(()) => {}
                         Err(ErrorSource::Client(e)) => {
                             error!(
@@ -177,7 +177,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 stream, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
+    let (mut client, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
         .await??
     {
         HandshakeData::Startup(stream, params) => (stream, params),
@@ -210,18 +210,17 @@ 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 stream)
+        .authenticate(ctx, &config.authentication_config, &mut client)
         .await
     {
         Ok(auth_result) => auth_result,
-        Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
+        Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
     };
     auth_info.set_startup_params(&params, true);
 
-    let node = connect_to_compute(
+    let mut node = connect_to_compute(
         ctx,
         &TcpMechanism {
-            user_info,
             auth: auth_info,
             locks: &config.connect_compute_locks,
         },
@@ -229,25 +228,41 @@ 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(stream.throw_error(e, Some(ctx)).await) })
+    .or_else(|e| async { Err(client.throw_error(e, Some(ctx)).await) })
     .await?;
 
-    let cancellation_handler_clone = Arc::clone(&cancellation_handler);
-    let session = cancellation_handler_clone.get_key();
+    let session = cancellation_handler.get_key();
 
-    session.write_cancel_key(node.cancel_closure.clone())?;
+    let cancel_closure =
+        prepare_client_connection(&mut node, session.key(), &mut client, user_info).await?;
 
-    prepare_client_connection(&node, *session.key(), &mut stream);
-    let stream = stream.flush_and_into_inner().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?;
 
     Ok(Some(ProxyPassthrough {
-        client: stream,
-        aux: node.aux.clone(),
+        client,
+        compute,
+
+        aux: node.aux,
         private_link_id: None,
-        compute: node,
-        session_id: ctx.session_id(),
-        cancel: session,
+
+        _cancel_on_shutdown: cancel_on_shutdown,
+
         _req: request_gauge,
         _conn: conn_gauge,
+        _db_conn: node.guage,
     }))
 }

proxy/src/control_plane/mod.rs

@@ -78,11 +78,8 @@ impl NodeInfo {
         ctx: &RequestContext,
         auth: &compute::AuthInfo,
         config: &ComputeConfig,
-        user_info: ComputeUserInfo,
     ) -> Result<compute::PostgresConnection, compute::ConnectionError> {
-        self.conn_info
-            .connect(ctx, self.aux.clone(), auth, config, user_info)
-            .await
+        self.conn_info.connect(ctx, &self.aux, auth, config).await
     }
 }
 

proxy/src/lib.rs

@@ -75,6 +75,7 @@
 pub mod binary;
 
 mod auth;
+mod batch;
 mod cache;
 mod cancellation;
 mod compute;

proxy/src/pglb/handshake.rs

@@ -12,7 +12,7 @@ use crate::pqproto::{
     BeMessage, CancelKeyData, FeStartupPacket, ProtocolVersion, StartupMessageParams,
 };
 use crate::proxy::TlsRequired;
-use crate::stream::{PqStream, Stream, StreamUpgradeError};
+use crate::stream::{PqFeStream, 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(PqStream<Stream<S>>, StartupMessageParams),
+    Startup(PqFeStream<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) = PqStream::parse_startup(Stream::from_raw(stream)).await?;
+    let (mut stream, mut msg) = PqFeStream::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) = PqStream::parse_startup(tls).await?;
+                        (stream, msg) = PqFeStream::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/passthrough.rs

@@ -1,15 +1,18 @@
-use futures::FutureExt;
+use std::convert::Infallible;
+
 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::cancellation;
-use crate::compute::PostgresConnection;
-use crate::config::ComputeConfig;
+use crate::compute::MaybeRustlsStream;
 use crate::control_plane::messages::MetricsAuxInfo;
-use crate::metrics::{Direction, Metrics, NumClientConnectionsGuard, NumConnectionRequestsGuard};
+use crate::metrics::{
+    Direction, Metrics, NumClientConnectionsGuard, NumConnectionRequestsGuard,
+    NumDbConnectionsGuard,
+};
 use crate::stream::Stream;
 use crate::usage_metrics::{Ids, MetricCounterRecorder, USAGE_METRICS};
 
@@ -64,40 +67,19 @@ pub(crate) async fn proxy_pass(
 
 pub(crate) struct ProxyPassthrough<S> {
     pub(crate) client: Stream<S>,
-    pub(crate) compute: PostgresConnection,
+    pub(crate) compute: MaybeRustlsStream<TcpStream>,
     pub(crate) aux: MetricsAuxInfo,
-    pub(crate) session_id: uuid::Uuid,
     pub(crate) private_link_id: Option<SmolStr>,
-    pub(crate) cancel: cancellation::Session,
+
+    pub(crate) _cancel_on_shutdown: tokio::sync::oneshot::Sender<Infallible>,
 
     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,
-        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
+    pub(crate) async fn proxy_pass(self) -> Result<(), ErrorSource> {
+        proxy_pass(self.client, self.compute, self.aux, self.private_link_id).await
     }
 }

proxy/src/pqproto.rs

@@ -11,11 +11,86 @@ use rand::distributions::{Distribution, Standard};
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 use zerocopy::{FromBytes, Immutable, IntoBytes, big_endian};
 
-pub type ErrorCode = [u8; 5];
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub struct ErrorCode(pub [u8; 5]);
 
-pub const FE_PASSWORD_MESSAGE: u8 = b'p';
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub struct FeTag(pub u8);
 
-pub const SQLSTATE_INTERNAL_ERROR: [u8; 5] = *b"XX000";
+#[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)),
+        }
+    }
+}
 
 /// The protocol version number.
 ///
@@ -280,6 +355,10 @@ 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 {
@@ -313,6 +392,19 @@ 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
@@ -353,7 +445,7 @@ impl WriteBuf {
 
             // Code: error_code
             buf.put_u8(b'C');
-            buf.put_slice(&error_code);
+            buf.put_slice(&error_code.0);
             buf.put_u8(0);
 
             // Message: msg
@@ -468,11 +560,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 {
@@ -494,17 +586,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, b'R', |buf| buf.put_i32(0));
+                buf.write_raw(1, BE_AUTH_MESSAGE.0, |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, b'R', |buf| buf.put_i32(3));
+                buf.write_raw(1, BE_AUTH_MESSAGE.0, |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, b'R', |buf| {
+                buf.write_raw(len + 2, BE_AUTH_MESSAGE.0, |buf| {
                     buf.put_i32(10); // Specifies that SASL auth method is used.
                     for method in methods {
                         buf.put_slice(method.as_bytes());
@@ -515,24 +607,19 @@ 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, b'R', |buf| {
+                buf.write_raw(extra.len() + 1, BE_AUTH_MESSAGE.0, |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, b'R', |buf| {
+                buf.write_raw(extra.len() + 1, BE_AUTH_MESSAGE.0, |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) => {
@@ -564,15 +651,16 @@ impl BeMessage<'_> {
                 });
             }
 
+            #[cfg(test)]
             // <https://www.postgresql.org/docs/current/protocol-message-formats.html#PROTOCOL-MESSAGE-FORMATS-NEGOTIATEPROTOCOLVERSION>
             BeMessage::ReadyForQuery => {
-                buf.write_raw(1, b'Z', |buf| buf.put_u8(b'I'));
+                buf.write_raw(1, BE_READY_MESSAGE.0, |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, b'v', |buf| {
+                buf.write_raw(8 + len, BE_NEGOTIATE_MESSAGE.0, |buf| {
                     buf.put_slice(version.as_bytes());
                     buf.put_u32(options.len() as u32);
                     for option in options {

proxy/src/proxy/connect_compute.rs

@@ -2,7 +2,6 @@ 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;
@@ -53,7 +52,6 @@ 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]
@@ -73,11 +71,7 @@ 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, self.user_info.clone())
-                .await,
-        )
+        permit.release_result(node_info.connect(ctx, &self.auth, config).await)
     }
 }
 

proxy/src/proxy/mod.rs

@@ -10,6 +10,7 @@ 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};
@@ -18,7 +19,8 @@ use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_util::sync::CancellationToken;
 use tracing::{Instrument, debug, error, info, warn};
 
-use crate::cancellation::{self, CancellationHandler};
+use crate::auth::backend::ComputeUserInfo;
+use crate::cancellation::{self, CancelClosure, CancellationHandler};
 use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig};
 use crate::context::RequestContext;
 use crate::error::{ReportableError, UserFacingError};
@@ -26,11 +28,11 @@ use crate::metrics::{Metrics, NumClientConnectionsGuard};
 pub use crate::pglb::copy_bidirectional::{ErrorSource, copy_bidirectional_client_compute};
 use crate::pglb::handshake::{HandshakeData, HandshakeError, handshake};
 use crate::pglb::passthrough::ProxyPassthrough;
-use crate::pqproto::{BeMessage, CancelKeyData, StartupMessageParams};
+use crate::pqproto::{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::{PqStream, Stream};
+use crate::stream::{PostgresError, PqFeStream, Stream};
 use crate::types::EndpointCacheKey;
 use crate::util::run_until_cancelled;
 use crate::{auth, compute};
@@ -155,7 +157,7 @@ pub async fn task_main(
                 Ok(Some(p)) => {
                     ctx.set_success();
                     let _disconnect = ctx.log_connect();
-                    match p.proxy_pass(&config.connect_to_compute).await {
+                    match p.proxy_pass().await {
                         Ok(()) => {}
                         Err(ErrorSource::Client(e)) => {
                             warn!(
@@ -253,7 +255,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     auth_backend: &'static auth::Backend<'static, ()>,
     ctx: &RequestContext,
     cancellation_handler: Arc<CancellationHandler>,
-    stream: S,
+    client: S,
     mode: ClientMode,
     endpoint_rate_limiter: Arc<EndpointRateLimiter>,
     conn_gauge: NumClientConnectionsGuard<'static>,
@@ -273,9 +275,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, stream, mode.handshake_tls(tls), record_handshake_error);
+    let do_handshake = handshake(ctx, client, mode.handshake_tls(tls), record_handshake_error);
 
-    let (mut stream, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
+    let (mut client, params) = match tokio::time::timeout(config.handshake_timeout, do_handshake)
         .await??
     {
         HandshakeData::Startup(stream, params) => (stream, params),
@@ -307,7 +309,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
 
     ctx.set_db_options(params.clone());
 
-    let hostname = mode.hostname(stream.get_ref());
+    let hostname = mode.hostname(client.get_ref());
 
     let common_names = tls.map(|tls| &tls.common_names);
 
@@ -319,14 +321,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(stream.throw_error(e, Some(ctx)).await)?,
+        Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
     };
 
     let user = user_info.get_user().to_owned();
     let user_info = match user_info
         .authenticate(
             ctx,
-            &mut stream,
+            &mut client,
             mode.allow_cleartext(),
             &config.authentication_config,
             endpoint_rate_limiter,
@@ -339,7 +341,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
             let app = params.get("application_name");
             let params_span = tracing::info_span!("", ?user, ?db, ?app);
 
-            return Err(stream
+            return Err(client
                 .throw_error(e, Some(ctx))
                 .instrument(params_span)
                 .await)?;
@@ -357,27 +359,40 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     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),
+        &auth::Backend::ControlPlane(cplane, creds.info.clone()),
         config.wake_compute_retry_config,
         &config.connect_to_compute,
     )
     .await;
 
-    let node = match res {
+    let mut node = match res {
         Ok(node) => node,
-        Err(e) => Err(stream.throw_error(e, Some(ctx)).await)?,
+        Err(e) => Err(client.throw_error(e, Some(ctx)).await)?,
     };
 
-    let cancellation_handler_clone = Arc::clone(&cancellation_handler);
-    let session = cancellation_handler_clone.get_key();
+    let session = cancellation_handler.get_key();
 
-    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 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?;
 
     let private_link_id = match ctx.extra() {
         Some(ConnectionInfoExtra::Aws { vpce_id }) => Some(vpce_id.clone()),
@@ -386,40 +401,75 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     };
 
     Ok(Some(ProxyPassthrough {
-        client: stream,
-        aux: node.aux.clone(),
+        client,
+        compute,
+
+        aux: node.aux,
         private_link_id,
-        compute: node,
-        session_id: ctx.session_id(),
-        cancel: session,
+
+        _cancel_on_shutdown: cancel_on_shutdown,
+
         _req: request_gauge,
         _conn: conn_gauge,
+        _db_conn: node.guage,
     }))
 }
 
 /// Finish client connection initialization: confirm auth success, send params, etc.
-pub(crate) fn 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());
-        });
+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?;
+        }
     }
 
-    // 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);
+    Ok(CancelClosure::new(
+        node.socket_addr,
+        RawCancelToken {
+            ssl_mode: node.ssl_mode,
+            process_id,
+            secret_key,
+        },
+        node.hostname.clone(),
+        user_info,
+    ))
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]

proxy/src/proxy/retry.rs

@@ -1,10 +1,12 @@
 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
@@ -96,10 +98,55 @@ 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(err) => err.could_retry(),
+            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::TlsError(err) => err.could_retry(),
             compute::ConnectionError::WakeComputeError(err) => err.could_retry(),
             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
@@ -109,7 +156,12 @@ impl CouldRetry for compute::ConnectionError {
 impl ShouldRetryWakeCompute for compute::ConnectionError {
     fn should_retry_wake_compute(&self) -> bool {
         match self {
-            compute::ConnectionError::Postgres(err) => err.should_retry_wake_compute(),
+            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,
             // the cache entry was not checked for validity
             compute::ConnectionError::TooManyConnectionAttempts(_) => false,
             _ => true,

proxy/src/proxy/tests/mod.rs

@@ -25,6 +25,7 @@ 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::tls::client_config::compute_client_config_with_certs;
 use crate::tls::server_config::CertResolver;
@@ -122,7 +123,7 @@ fn generate_tls_config<'a>(
 trait TestAuth: Sized {
     async fn authenticate<S: AsyncRead + AsyncWrite + Unpin + Send>(
         self,
-        stream: &mut PqStream<Stream<S>>,
+        stream: &mut PqFeStream<Stream<S>>,
     ) -> anyhow::Result<()> {
         stream.write_message(BeMessage::AuthenticationOk);
         Ok(())
@@ -151,7 +152,7 @@ impl Scram {
 impl TestAuth for Scram {
     async fn authenticate<S: AsyncRead + AsyncWrite + Unpin + Send>(
         self,
-        stream: &mut PqStream<Stream<S>>,
+        stream: &mut PqFeStream<Stream<S>>,
     ) -> anyhow::Result<()> {
         let outcome = auth::AuthFlow::new(stream, auth::Scram(&self.0, &RequestContext::test()))
             .authenticate()

proxy/src/redis/keys.rs

@@ -1,8 +1,4 @@
-use std::io::ErrorKind;
-
-use anyhow::Ok;
-
-use crate::pqproto::{CancelKeyData, id_to_cancel_key};
+use crate::pqproto::CancelKeyData;
 
 pub mod keyspace {
     pub const CANCEL_PREFIX: &str = "cancel";
@@ -23,39 +19,12 @@ 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]
@@ -65,16 +34,4 @@ 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,3 +1,6 @@
+use std::time::Duration;
+
+use futures::FutureExt;
 use redis::aio::ConnectionLike;
 use redis::{Cmd, FromRedisValue, Pipeline, RedisResult};
 
@@ -35,14 +38,11 @@ impl RedisKVClient {
     }
 
     pub async fn try_connect(&mut self) -> anyhow::Result<()> {
-        match self.client.connect().await {
-            Ok(()) => {}
-            Err(e) => {
-                tracing::error!("failed to connect to redis: {e}");
-                return Err(e);
-            }
-        }
-        Ok(())
+        self.client
+            .connect()
+            .boxed()
+            .await
+            .inspect_err(|e| tracing::error!("failed to connect to redis: {e}"))
     }
 
     pub(crate) async fn query<T: FromRedisValue>(
@@ -54,15 +54,25 @@ impl RedisKVClient {
             return Err(anyhow::anyhow!("Rate limit exceeded"));
         }
 
-        match q.query(&mut self.client).await {
+        let e = match q.query(&mut self.client).await {
             Ok(t) => return Ok(t),
-            Err(e) => {
-                tracing::error!("failed to run query: {e}");
+            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?;
             }
+            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/channel_binding.rs

@@ -1,8 +1,5 @@
 //! Definition and parser for channel binding flag (a part of the `GS2` header).
 
-use base64::Engine as _;
-use base64::prelude::BASE64_STANDARD;
-
 /// Channel binding flag (possibly with params).
 #[derive(Debug, PartialEq, Eq)]
 pub(crate) enum ChannelBinding<T> {
@@ -58,7 +55,7 @@ impl<T: std::fmt::Display> ChannelBinding<T> {
                 let mut cbind_input = vec![];
                 write!(&mut cbind_input, "p={mode},,",).unwrap();
                 cbind_input.extend_from_slice(get_cbind_data(mode)?);
-                BASE64_STANDARD.encode(&cbind_input).into()
+                base64::encode(&cbind_input).into()
             }
         })
     }
@@ -73,9 +70,9 @@ mod tests {
         use ChannelBinding::*;
 
         let cases = [
-            (NotSupportedClient, BASE64_STANDARD.encode("n,,")),
-            (NotSupportedServer, BASE64_STANDARD.encode("y,,")),
-            (Required("foo"), BASE64_STANDARD.encode("p=foo,,bar")),
+            (NotSupportedClient, base64::encode("n,,")),
+            (NotSupportedServer, base64::encode("y,,")),
+            (Required("foo"), base64::encode("p=foo,,bar")),
         ];
 
         for (cb, input) in cases {

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::PqStream;
+use crate::stream::PqFeStream;
 
 /// 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 PqStream<S>,
+    stream: &mut PqFeStream<S>,
     mechanism: F,
 ) -> super::Result<Outcome<M::Output>>
 where

proxy/src/scram/exchange.rs

@@ -2,8 +2,6 @@
 
 use std::convert::Infallible;
 
-use base64::Engine as _;
-use base64::prelude::BASE64_STANDARD;
 use hmac::{Hmac, Mac};
 use sha2::Sha256;
 
@@ -107,7 +105,7 @@ pub(crate) async fn exchange(
     secret: &ServerSecret,
     password: &[u8],
 ) -> sasl::Result<sasl::Outcome<super::ScramKey>> {
-    let salt = BASE64_STANDARD.decode(&secret.salt_base64)?;
+    let salt = base64::decode(&secret.salt_base64)?;
     let client_key = derive_client_key(pool, endpoint, password, &salt, secret.iterations).await;
 
     if secret.is_password_invalid(&client_key).into() {

proxy/src/scram/messages.rs

@@ -3,9 +3,6 @@
 use std::fmt;
 use std::ops::Range;
 
-use base64::Engine as _;
-use base64::prelude::BASE64_STANDARD;
-
 use super::base64_decode_array;
 use super::key::{SCRAM_KEY_LEN, ScramKey};
 use super::signature::SignatureBuilder;
@@ -91,7 +88,7 @@ impl<'a> ClientFirstMessage<'a> {
 
         let mut message = String::new();
         write!(&mut message, "r={}", self.nonce).unwrap();
-        BASE64_STANDARD.encode_string(nonce, &mut message);
+        base64::encode_config_buf(nonce, base64::STANDARD, &mut message);
         let combined_nonce = 2..message.len();
         write!(&mut message, ",s={salt_base64},i={iterations}").unwrap();
 
@@ -145,7 +142,11 @@ impl<'a> ClientFinalMessage<'a> {
         server_key: &ScramKey,
     ) -> String {
         let mut buf = String::from("v=");
-        BASE64_STANDARD.encode_string(signature_builder.build(server_key), &mut buf);
+        base64::encode_config_buf(
+            signature_builder.build(server_key),
+            base64::STANDARD,
+            &mut buf,
+        );
 
         buf
     }
@@ -250,7 +251,7 @@ mod tests {
             "iiYEfS3rOgn8S3rtpSdrOsHtPLWvIkdgmHxA0hf3JNOAG4dU"
         );
         assert_eq!(
-            BASE64_STANDARD.encode(msg.proof),
+            base64::encode(msg.proof),
             "SRpfsIVS4Gk11w1LqQ4QvCUBZYQmqXNSDEcHqbQ3CHI="
         );
     }

proxy/src/scram/mod.rs

@@ -15,8 +15,6 @@ mod secret;
 mod signature;
 pub mod threadpool;
 
-use base64::Engine as _;
-use base64::prelude::BASE64_STANDARD;
 pub(crate) use exchange::{Exchange, exchange};
 use hmac::{Hmac, Mac};
 pub(crate) use key::ScramKey;
@@ -34,7 +32,7 @@ pub(crate) const METHODS_WITHOUT_PLUS: &[&str] = &[SCRAM_SHA_256];
 fn base64_decode_array<const N: usize>(input: impl AsRef<[u8]>) -> Option<[u8; N]> {
     let mut bytes = [0u8; N];
 
-    let size = BASE64_STANDARD.decode_slice(input, &mut bytes).ok()?;
+    let size = base64::decode_config_slice(input, base64::STANDARD, &mut bytes).ok()?;
     if size != N {
         return None;
     }

proxy/src/scram/secret.rs

@@ -1,7 +1,5 @@
 //! Tools for SCRAM server secret management.
 
-use base64::Engine as _;
-use base64::prelude::BASE64_STANDARD;
 use subtle::{Choice, ConstantTimeEq};
 
 use super::base64_decode_array;
@@ -58,7 +56,7 @@ impl ServerSecret {
             // iteration count 1 for our generated passwords going forward.
             // PG16 users can set iteration count=1 already today.
             iterations: 1,
-            salt_base64: BASE64_STANDARD.encode(nonce),
+            salt_base64: base64::encode(nonce),
             stored_key: ScramKey::default(),
             server_key: ScramKey::default(),
             doomed: true,
@@ -90,7 +88,7 @@ mod tests {
         assert_eq!(parsed.iterations, iterations);
         assert_eq!(parsed.salt_base64, salt);
 
-        assert_eq!(BASE64_STANDARD.encode(parsed.stored_key), stored_key);
-        assert_eq!(BASE64_STANDARD.encode(parsed.server_key), server_key);
+        assert_eq!(base64::encode(parsed.stored_key), stored_key);
+        assert_eq!(base64::encode(parsed.server_key), server_key);
     }
 }

proxy/src/serverless/local_conn_pool.rs

@@ -16,8 +16,6 @@ use std::sync::atomic::AtomicUsize;
 use std::task::{Poll, ready};
 use std::time::Duration;
 
-use base64::Engine as _;
-use base64::prelude::BASE64_URL_SAFE_NO_PAD;
 use ed25519_dalek::{Signature, Signer, SigningKey};
 use futures::Future;
 use futures::future::poll_fn;
@@ -348,7 +346,7 @@ fn sign_jwt(sk: &SigningKey, payload: &[u8]) -> String {
     jwt.push_str("eyJhbGciOiJFZERTQSJ9.");
 
     // encode the jwt payload in-place
-    BASE64_URL_SAFE_NO_PAD.encode_string(payload, &mut jwt);
+    base64::encode_config_buf(payload, base64::URL_SAFE_NO_PAD, &mut jwt);
 
     // create the signature from the encoded header || payload
     let sig: Signature = sk.sign(jwt.as_bytes());
@@ -356,7 +354,7 @@ fn sign_jwt(sk: &SigningKey, payload: &[u8]) -> String {
     jwt.push('.');
 
     // encode the jwt signature in-place
-    BASE64_URL_SAFE_NO_PAD.encode_string(sig.to_bytes(), &mut jwt);
+    base64::encode_config_buf(sig.to_bytes(), base64::URL_SAFE_NO_PAD, &mut jwt);
 
     debug_assert_eq!(
         jwt.len(),

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(&config.connect_to_compute).await {
+            match p.proxy_pass().await {
                 Ok(()) => Ok(()),
                 Err(ErrorSource::Client(err)) => Err(err).context("client"),
                 Err(ErrorSource::Compute(err)) => Err(err).context("compute"),

proxy/src/stream.rs

@@ -1,351 +0,0 @@
-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

@@ -0,0 +1,168 @@
+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

@@ -0,0 +1,165 @@
+//! 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

@@ -0,0 +1,197 @@
+//! 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/tls/mod.rs

@@ -3,8 +3,6 @@ pub mod postgres_rustls;
 pub mod server_config;
 
 use anyhow::Context;
-use base64::Engine as _;
-use base64::prelude::BASE64_STANDARD;
 use rustls::pki_types::CertificateDer;
 use sha2::{Digest, Sha256};
 use tracing::{error, info};
@@ -60,7 +58,7 @@ impl TlsServerEndPoint {
         let oid = certificate.signature_algorithm.oid;
         if SHA256_OIDS.contains(&oid) {
             let tls_server_end_point: [u8; 32] = Sha256::new().chain_update(cert).finalize().into();
-            info!(%subject, tls_server_end_point = %BASE64_STANDARD.encode(tls_server_end_point), "determined channel binding");
+            info!(%subject, tls_server_end_point = %base64::encode(tls_server_end_point), "determined channel binding");
             Ok(Self::Sha256(tls_server_end_point))
         } else {
             error!(%subject, "unknown channel binding");

pyproject.toml

@@ -9,7 +9,7 @@ pytest = "^7.4.4"
 psycopg2-binary = "^2.9.10"
 typing-extensions = "^4.12.2"
 PyJWT = {version = "^2.1.0", extras = ["crypto"]}
-requests = "^2.32.4"
+requests = "^2.32.3"
 pytest-xdist = "^3.3.1"
 asyncpg = "^0.30.0"
 aiopg = "^1.4.0"

storage_controller/src/http.rs

@@ -1398,31 +1398,6 @@ async fn handle_timeline_import(req: Request<Body>) -> Result<Response<Body>, Ap
     )
 }
 
-async fn handle_tenant_timeline_locate(
-    service: Arc<Service>,
-    req: Request<Body>,
-) -> Result<Response<Body>, ApiError> {
-    let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
-    let timeline_id: TimelineId = parse_request_param(&req, "timeline_id")?;
-
-    check_permissions(&req, Scope::Admin)?;
-    maybe_rate_limit(&req, tenant_id).await;
-
-    match maybe_forward(req).await {
-        ForwardOutcome::Forwarded(res) => {
-            return res;
-        }
-        ForwardOutcome::NotForwarded(_req) => {}
-    };
-
-    json_response(
-        StatusCode::OK,
-        service
-            .tenant_timeline_locate(tenant_id, timeline_id)
-            .await?,
-    )
-}
-
 async fn handle_tenants_dump(req: Request<Body>) -> Result<Response<Body>, ApiError> {
     check_permissions(&req, Scope::Admin)?;
 
@@ -2164,16 +2139,6 @@ pub fn make_router(
                 )
             },
         )
-        .get(
-            "/debug/v1/tenant/:tenant_id/timeline/:timeline_id/locate",
-            |r| {
-                tenant_service_handler(
-                    r,
-                    handle_tenant_timeline_locate,
-                    RequestName("v1_tenant_timeline_locate"),
-                )
-            },
-        )
         .get("/debug/v1/scheduler", |r| {
             named_request_span(r, handle_scheduler_dump, RequestName("debug_v1_scheduler"))
         })

storage_controller/src/service.rs

@@ -2267,7 +2267,6 @@ impl Service {
                     // fail, and start from scratch, so it doesn't make sense for us to try and preserve
                     // the stale/multi states at this point.
                     mode: LocationConfigMode::AttachedSingle,
-                    stripe_size: shard.shard.stripe_size,
                 });
 
                 shard.generation = std::cmp::max(shard.generation, Some(new_gen));
@@ -2301,7 +2300,6 @@ impl Service {
                     id: *tenant_shard_id,
                     r#gen: None,
                     mode: LocationConfigMode::Secondary,
-                    stripe_size: shard.shard.stripe_size,
                 });
 
                 // We must not update observed, because we have no guarantee that our
@@ -6651,8 +6649,6 @@ impl Service {
 
     /// This is for debug/support only: assuming tenant data is already present in S3, we "create" a
     /// tenant with a very high generation number so that it will see the existing data.
-    /// It does not create timelines on safekeepers, because they might already exist on some
-    /// safekeeper set. So, the timelines are not storcon-managed after the import.
     pub(crate) async fn tenant_import(
         &self,
         tenant_id: TenantId,

storage_controller/src/service/safekeeper_service.rs

@@ -17,7 +17,7 @@ use pageserver_api::controller_api::{
     SafekeeperDescribeResponse, SkSchedulingPolicy, TimelineImportRequest,
 };
 use pageserver_api::models::{SafekeeperInfo, SafekeepersInfo, TimelineInfo};
-use safekeeper_api::membership::{MemberSet, SafekeeperGeneration, SafekeeperId};
+use safekeeper_api::membership::{MemberSet, SafekeeperId};
 use tokio::task::JoinSet;
 use tokio_util::sync::CancellationToken;
 use utils::id::{NodeId, TenantId, TimelineId};
@@ -26,13 +26,6 @@ use utils::lsn::Lsn;
 
 use super::Service;
 
-#[derive(serde::Serialize, serde::Deserialize, Clone)]
-pub struct TimelineLocateResponse {
-    pub generation: SafekeeperGeneration,
-    pub sk_set: Vec<NodeId>,
-    pub new_sk_set: Option<Vec<NodeId>>,
-}
-
 impl Service {
     /// Timeline creation on safekeepers
     ///
@@ -403,38 +396,6 @@ impl Service {
         Ok(())
     }
 
-    /// Locate safekeepers for a timeline.
-    /// Return the generation, sk_set and new_sk_set if present.
-    /// If the timeline is not storcon-managed, return NotFound.
-    pub(crate) async fn tenant_timeline_locate(
-        &self,
-        tenant_id: TenantId,
-        timeline_id: TimelineId,
-    ) -> Result<TimelineLocateResponse, ApiError> {
-        let timeline = self
-            .persistence
-            .get_timeline(tenant_id, timeline_id)
-            .await?;
-
-        let Some(timeline) = timeline else {
-            return Err(ApiError::NotFound(
-                anyhow::anyhow!("Timeline {}/{} not found", tenant_id, timeline_id).into(),
-            ));
-        };
-
-        Ok(TimelineLocateResponse {
-            generation: SafekeeperGeneration::new(timeline.generation as u32),
-            sk_set: timeline
-                .sk_set
-                .iter()
-                .map(|id| NodeId(*id as u64))
-                .collect(),
-            new_sk_set: timeline
-                .new_sk_set
-                .map(|sk_set| sk_set.iter().map(|id| NodeId(*id as u64)).collect()),
-        })
-    }
-
     /// Perform timeline deletion on safekeepers. Will return success: we persist the deletion into the reconciler.
     pub(super) async fn tenant_timeline_delete_safekeepers(
         self: &Arc<Self>,

test_runner/fixtures/neon_fixtures.py

@@ -2223,17 +2223,6 @@ class NeonStorageController(MetricsGetter, LogUtils):
         shards: list[dict[str, Any]] = body["shards"]
         return shards
 
-    def timeline_locate(self, tenant_id: TenantId, timeline_id: TimelineId):
-        """
-        :return: dict {"generation": int, "sk_set": [int], "new_sk_set": [int]}
-        """
-        response = self.request(
-            "GET",
-            f"{self.api}/debug/v1/tenant/{tenant_id}/timeline/{timeline_id}/locate",
-            headers=self.headers(TokenScope.ADMIN),
-        )
-        return response.json()
-
     def tenant_describe(self, tenant_id: TenantId):
         """
         :return: list of {"shard_id": "", "node_id": int, "listen_pg_addr": str, "listen_pg_port": int, "listen_http_addr: str, "listen_http_port: int, preferred_az_id: str}

test_runner/regress/test_compatibility.py

@@ -18,7 +18,6 @@ from fixtures.neon_fixtures import (
     NeonEnv,
     NeonEnvBuilder,
     PgBin,
-    Safekeeper,
     flush_ep_to_pageserver,
 )
 from fixtures.pageserver.http import PageserverApiException
@@ -27,7 +26,6 @@ from fixtures.pageserver.utils import (
 )
 from fixtures.pg_version import PgVersion
 from fixtures.remote_storage import RemoteStorageKind, S3Storage, s3_storage
-from fixtures.safekeeper.http import MembershipConfiguration
 from fixtures.workload import Workload
 
 if TYPE_CHECKING:
@@ -544,24 +542,6 @@ def test_historic_storage_formats(
     # All our artifacts should contain at least one timeline
     assert len(timelines) > 0
 
-    # Import tenant does not create the timeline on safekeepers,
-    # because it is a debug handler and the timeline may have already been
-    # created on some set of safekeepers.
-    # Create the timeline on safekeepers manually.
-    # TODO(diko): when we have the script/storcon handler to migrate
-    # the timeline to storcon, we can replace this code with it.
-    mconf = MembershipConfiguration(
-        generation=1,
-        members=Safekeeper.sks_to_safekeeper_ids([env.safekeepers[0]]),
-        new_members=None,
-    )
-    members_sks = Safekeeper.mconf_sks(env, mconf)
-
-    for timeline in timelines:
-        Safekeeper.create_timeline(
-            dataset.tenant_id, timeline["timeline_id"], env.pageserver, mconf, members_sks
-        )
-
     # TODO: ensure that the snapshots we're importing contain a sensible variety of content, at the very
     # least they should include a mixture of deltas and image layers.  Preferably they should also
     # contain some "exotic" stuff like aux files from logical replication.

test_runner/regress/test_tenant_delete.py

@@ -430,7 +430,6 @@ def test_tenant_delete_stale_shards(neon_env_builder: NeonEnvBuilder, pg_bin: Pg
     workload.init()
     workload.write_rows(256)
     workload.validate()
-    workload.stop()
 
     assert_prefix_not_empty(
         neon_env_builder.pageserver_remote_storage,

workspace_hack/Cargo.toml

@@ -20,7 +20,8 @@ anstream = { version = "0.6" }
 anyhow = { version = "1", features = ["backtrace"] }
 axum = { version = "0.8", features = ["ws"] }
 axum-core = { version = "0.5", default-features = false, features = ["tracing"] }
-base64 = { version = "0.21" }
+base64-594e8ee84c453af0 = { package = "base64", version = "0.13", features = ["alloc"] }
+base64-647d43efb71741da = { package = "base64", version = "0.21" }
 base64ct = { version = "1", default-features = false, features = ["std"] }
 bytes = { version = "1", features = ["serde"] }
 camino = { version = "1", default-features = false, features = ["serde1"] }