half-baked copy_file impl

.github/ISSUE_TEMPLATE/epic-template.md

@@ -17,9 +17,8 @@ assignees: ''
 ## Implementation ideas
 
 
-```[tasklist]
-### Tasks
-```
+## Tasks
+- [ ]
 
 
 ## Other related tasks and Epics

Cargo.lock

@@ -4064,7 +4064,6 @@ dependencies = [
  "aws-config",
  "aws-credential-types",
  "aws-sdk-s3",
- "aws-smithy-async",
  "aws-smithy-http",
  "aws-types",
  "azure_core",
@@ -6483,7 +6482,6 @@ dependencies = [
  "clap",
  "clap_builder",
  "crossbeam-utils",
- "dashmap",
  "either",
  "fail",
  "futures",

Cargo.toml

@@ -48,7 +48,6 @@ async-trait = "0.1"
 aws-config = { version = "0.56", default-features = false, features=["rustls"] }
 aws-sdk-s3 = "0.29"
 aws-smithy-http = "0.56"
-aws-smithy-async = { version = "0.56", default-features = false, features=["rt-tokio"] }
 aws-credential-types = "0.56"
 aws-types = "0.56"
 axum = { version = "0.6.20", features = ["ws"] }
@@ -67,7 +66,7 @@ comfy-table = "6.1"
 const_format = "0.2"
 crc32c = "0.6"
 crossbeam-utils = "0.8.5"
-dashmap = { version = "5.5.0", features = ["raw-api"] }
+dashmap = "5.5.0"
 either = "1.8"
 enum-map = "2.4.2"
 enumset = "1.0.12"

compute_tools/src/extension_server.rs

@@ -78,7 +78,7 @@ use regex::Regex;
 use remote_storage::*;
 use serde_json;
 use std::io::Read;
-use std::num::NonZeroUsize;
+use std::num::{NonZeroU32, NonZeroUsize};
 use std::path::Path;
 use std::str;
 use tar::Archive;
@@ -281,6 +281,8 @@ pub fn init_remote_storage(remote_ext_config: &str) -> anyhow::Result<GenericRem
         max_keys_per_list_response: None,
     };
     let config = RemoteStorageConfig {
+        max_concurrent_syncs: NonZeroUsize::new(100).expect("100 != 0"),
+        max_sync_errors: NonZeroU32::new(100).expect("100 != 0"),
         storage: RemoteStorageKind::AwsS3(config),
     };
     GenericRemoteStorage::from_config(&config)

libs/postgres_backend/src/lib.rs

@@ -17,7 +17,7 @@ use std::{fmt, io};
 use std::{future::Future, str::FromStr};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_rustls::TlsAcceptor;
-use tracing::{debug, error, info, trace, warn};
+use tracing::{debug, error, info, trace};
 
 use pq_proto::framed::{ConnectionError, Framed, FramedReader, FramedWriter};
 use pq_proto::{
@@ -35,11 +35,6 @@ pub enum QueryError {
     /// We were instructed to shutdown while processing the query
     #[error("Shutting down")]
     Shutdown,
-    /// Authentication failure
-    #[error("Unauthorized: {0}")]
-    Unauthorized(std::borrow::Cow<'static, str>),
-    #[error("Simulated Connection Error")]
-    SimulatedConnectionError,
     /// Some other error
     #[error(transparent)]
     Other(#[from] anyhow::Error),
@@ -54,9 +49,8 @@ impl From<io::Error> for QueryError {
 impl QueryError {
     pub fn pg_error_code(&self) -> &'static [u8; 5] {
         match self {
-            Self::Disconnected(_) | Self::SimulatedConnectionError => b"08006", // connection failure
+            Self::Disconnected(_) => b"08006", // connection failure
             Self::Shutdown => SQLSTATE_ADMIN_SHUTDOWN,
-            Self::Unauthorized(_) => SQLSTATE_INTERNAL_ERROR,
             Self::Other(_) => SQLSTATE_INTERNAL_ERROR, // internal error
         }
     }
@@ -616,7 +610,7 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
 
                     if let Err(e) = handler.check_auth_jwt(self, jwt_response) {
                         self.write_message_noflush(&BeMessage::ErrorResponse(
-                            &short_error(&e),
+                            &e.to_string(),
                             Some(e.pg_error_code()),
                         ))?;
                         return Err(e);
@@ -738,9 +732,6 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> PostgresBackend<IO> {
                 if let Err(e) = handler.process_query(self, query_string).await {
                     match e {
                         QueryError::Shutdown => return Ok(ProcessMsgResult::Break),
-                        QueryError::SimulatedConnectionError => {
-                            return Err(QueryError::SimulatedConnectionError)
-                        }
                         e => {
                             log_query_error(query_string, &e);
                             let short_error = short_error(&e);
@@ -975,8 +966,6 @@ pub fn short_error(e: &QueryError) -> String {
     match e {
         QueryError::Disconnected(connection_error) => connection_error.to_string(),
         QueryError::Shutdown => "shutdown".to_string(),
-        QueryError::Unauthorized(_e) => "JWT authentication error".to_string(),
-        QueryError::SimulatedConnectionError => "simulated connection error".to_string(),
         QueryError::Other(e) => format!("{e:#}"),
     }
 }
@@ -993,15 +982,9 @@ fn log_query_error(query: &str, e: &QueryError) {
         QueryError::Disconnected(other_connection_error) => {
             error!("query handler for '{query}' failed with connection error: {other_connection_error:?}")
         }
-        QueryError::SimulatedConnectionError => {
-            error!("query handler for query '{query}' failed due to a simulated connection error")
-        }
         QueryError::Shutdown => {
             info!("query handler for '{query}' cancelled during tenant shutdown")
         }
-        QueryError::Unauthorized(e) => {
-            warn!("query handler for '{query}' failed with authentication error: {e}");
-        }
         QueryError::Other(e) => {
             error!("query handler for '{query}' failed: {e:?}");
         }

libs/remote_storage/Cargo.toml

@@ -8,7 +8,6 @@ license.workspace = true
 anyhow.workspace = true
 async-trait.workspace = true
 once_cell.workspace = true
-aws-smithy-async.workspace = true
 aws-smithy-http.workspace = true
 aws-types.workspace = true
 aws-config.workspace = true

libs/remote_storage/src/azure_blob.rs

@@ -311,4 +311,8 @@ impl RemoteStorage for AzureBlobStorage {
         }
         Ok(())
     }
+
+    async fn copy_object(&self, src: &RemotePath, dst: &RemotePath) -> anyhow::Result<()> {
+        unimplemented!()
+    }
 }

libs/remote_storage/src/lib.rs

@@ -14,7 +14,13 @@ mod local_fs;
 mod s3_bucket;
 mod simulate_failures;
 
-use std::{collections::HashMap, fmt::Debug, num::NonZeroUsize, pin::Pin, sync::Arc};
+use std::{
+    collections::HashMap,
+    fmt::Debug,
+    num::{NonZeroU32, NonZeroUsize},
+    pin::Pin,
+    sync::Arc,
+};
 
 use anyhow::{bail, Context};
 use camino::{Utf8Path, Utf8PathBuf};
@@ -30,6 +36,12 @@ pub use self::{
 };
 use s3_bucket::RequestKind;
 
+/// How many different timelines can be processed simultaneously when synchronizing layers with the remote storage.
+/// During regular work, pageserver produces one layer file per timeline checkpoint, with bursts of concurrency
+/// during start (where local and remote timelines are compared and initial sync tasks are scheduled) and timeline attach.
+/// Both cases may trigger timeline download, that might download a lot of layers. This concurrency is limited by the clients internally, if needed.
+pub const DEFAULT_REMOTE_STORAGE_MAX_CONCURRENT_SYNCS: usize = 50;
+pub const DEFAULT_REMOTE_STORAGE_MAX_SYNC_ERRORS: u32 = 10;
 /// Currently, sync happens with AWS S3, that has two limits on requests per second:
 /// ~200 RPS for IAM services
 /// <https://docs.aws.amazon.com/AmazonRDS/latest/AuroraUserGuide/UsingWithRDS.IAMDBAuth.html>
@@ -102,7 +114,7 @@ impl RemotePath {
         self.0.file_name()
     }
 
-    pub fn join(&self, segment: &Utf8Path) -> Self {
+    pub fn join<P: AsRef<Utf8Path>>(&self, segment: P) -> Self {
         Self(self.0.join(segment))
     }
 
@@ -203,6 +215,8 @@ pub trait RemoteStorage: Send + Sync + 'static {
     async fn delete(&self, path: &RemotePath) -> anyhow::Result<()>;
 
     async fn delete_objects<'a>(&self, paths: &'a [RemotePath]) -> anyhow::Result<()>;
+
+    async fn copy_object(&self, src: &RemotePath, dst: &RemotePath) -> anyhow::Result<()>;
 }
 
 pub struct Download {
@@ -365,6 +379,15 @@ impl GenericRemoteStorage {
             Self::Unreliable(s) => s.delete_objects(paths).await,
         }
     }
+
+    pub async fn copy_object(&self, src: &RemotePath, dst: &RemotePath) -> anyhow::Result<()> {
+        match self {
+            Self::LocalFs(s) => s.copy_object(src, dst).await,
+            Self::AwsS3(s) => s.copy_object(src, dst).await,
+            Self::AzureBlob(s) => s.copy_object(src, dst).await,
+            Self::Unreliable(s) => s.copy_object(src, dst).await,
+        }
+    }
 }
 
 impl GenericRemoteStorage {
@@ -431,6 +454,10 @@ pub struct StorageMetadata(HashMap<String, String>);
 /// External backup storage configuration, enough for creating a client for that storage.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct RemoteStorageConfig {
+    /// Max allowed number of concurrent sync operations between the API user and the remote storage.
+    pub max_concurrent_syncs: NonZeroUsize,
+    /// Max allowed errors before the sync task is considered failed and evicted.
+    pub max_sync_errors: NonZeroU32,
     /// The storage connection configuration.
     pub storage: RemoteStorageKind,
 }
@@ -526,6 +553,18 @@ impl RemoteStorageConfig {
 
         let use_azure = container_name.is_some() && container_region.is_some();
 
+        let max_concurrent_syncs = NonZeroUsize::new(
+            parse_optional_integer("max_concurrent_syncs", toml)?
+                .unwrap_or(DEFAULT_REMOTE_STORAGE_MAX_CONCURRENT_SYNCS),
+        )
+        .context("Failed to parse 'max_concurrent_syncs' as a positive integer")?;
+
+        let max_sync_errors = NonZeroU32::new(
+            parse_optional_integer("max_sync_errors", toml)?
+                .unwrap_or(DEFAULT_REMOTE_STORAGE_MAX_SYNC_ERRORS),
+        )
+        .context("Failed to parse 'max_sync_errors' as a positive integer")?;
+
         let default_concurrency_limit = if use_azure {
             DEFAULT_REMOTE_STORAGE_AZURE_CONCURRENCY_LIMIT
         } else {
@@ -607,7 +646,11 @@ impl RemoteStorageConfig {
             }
         };
 
-        Ok(Some(RemoteStorageConfig { storage }))
+        Ok(Some(RemoteStorageConfig {
+            max_concurrent_syncs,
+            max_sync_errors,
+            storage,
+        }))
     }
 }
 

libs/remote_storage/src/local_fs.rs

@@ -393,6 +393,27 @@ impl RemoteStorage for LocalFs {
         }
         Ok(())
     }
+
+    async fn copy_object(&self, src: &RemotePath, dst: &RemotePath) -> anyhow::Result<()> {
+        let src_path = src.with_base(&self.storage_root);
+        let dst_path = dst.with_base(&self.storage_root);
+
+        // If the destination file already exists, we need to delete it first.
+        if dst_path.exists() {
+            fs::remove_file(&dst_path).await?;
+        }
+
+        // Copy the file.
+        fs::copy(&src_path, &dst_path).await?;
+
+        // Copy the metadata.
+        let metadata_path = storage_metadata_path(&src_path);
+        if metadata_path.exists() {
+            fs::copy(&metadata_path, storage_metadata_path(&dst_path)).await?;
+        }
+
+        Ok(())
+    }
 }
 
 fn storage_metadata_path(original_path: &Utf8Path) -> Utf8PathBuf {

libs/remote_storage/src/s3_bucket.rs

@@ -4,27 +4,23 @@
 //! allowing multiple api users to independently work with the same S3 bucket, if
 //! their bucket prefixes are both specified and different.
 
-use std::{borrow::Cow, sync::Arc};
+use std::borrow::Cow;
 
 use anyhow::Context;
 use aws_config::{
     environment::credentials::EnvironmentVariableCredentialsProvider,
-    imds::credentials::ImdsCredentialsProvider,
-    meta::credentials::CredentialsProviderChain,
-    provider_config::ProviderConfig,
-    retry::{RetryConfigBuilder, RetryMode},
-    web_identity_token::WebIdentityTokenCredentialsProvider,
+    imds::credentials::ImdsCredentialsProvider, meta::credentials::CredentialsProviderChain,
+    provider_config::ProviderConfig, web_identity_token::WebIdentityTokenCredentialsProvider,
 };
 use aws_credential_types::cache::CredentialsCache;
 use aws_sdk_s3::{
-    config::{AsyncSleep, Config, Region, SharedAsyncSleep},
+    config::{Config, Region},
     error::SdkError,
     operation::get_object::GetObjectError,
     primitives::ByteStream,
     types::{Delete, ObjectIdentifier},
     Client,
 };
-use aws_smithy_async::rt::sleep::TokioSleep;
 use aws_smithy_http::body::SdkBody;
 use hyper::Body;
 use scopeguard::ScopeGuard;
@@ -87,23 +83,10 @@ impl S3Bucket {
             .or_else("imds", ImdsCredentialsProvider::builder().build())
         };
 
-        // AWS SDK requires us to specify how the RetryConfig should sleep when it wants to back off
-        let sleep_impl: Arc<dyn AsyncSleep> = Arc::new(TokioSleep::new());
-
-        // We do our own retries (see [`backoff::retry`]).  However, for the AWS SDK to enable rate limiting in response to throttling
-        // responses (e.g. 429 on too many ListObjectsv2 requests), we must provide a retry config.  We set it to use at most one
-        // attempt, and enable 'Adaptive' mode, which causes rate limiting to be enabled.
-        let mut retry_config = RetryConfigBuilder::new();
-        retry_config
-            .set_max_attempts(Some(1))
-            .set_mode(Some(RetryMode::Adaptive));
-
         let mut config_builder = Config::builder()
             .region(region)
             .credentials_cache(CredentialsCache::lazy())
-            .credentials_provider(credentials_provider)
-            .sleep_impl(SharedAsyncSleep::from(sleep_impl))
-            .retry_config(retry_config.build());
+            .credentials_provider(credentials_provider);
 
         if let Some(custom_endpoint) = aws_config.endpoint.clone() {
             config_builder = config_builder
@@ -238,6 +221,8 @@ impl S3Bucket {
             )),
         }
     }
+
+
 }
 
 pin_project_lite::pin_project! {
@@ -532,6 +517,11 @@ impl RemoteStorage for S3Bucket {
         let paths = std::array::from_ref(path);
         self.delete_objects(paths).await
     }
+
+    async fn copy_object(&self, src: &RemotePath, dst: &RemotePath) -> anyhow::Result<()> {
+        unimplemented!()
+    }
+
 }
 
 /// On drop (cancellation) count towards [`metrics::BucketMetrics::cancelled_waits`].
@@ -614,4 +604,6 @@ mod tests {
             }
         }
     }
+
+
 }

libs/remote_storage/src/simulate_failures.rs

@@ -160,4 +160,8 @@ impl RemoteStorage for UnreliableWrapper {
         }
         Ok(())
     }
+
+    async fn copy_object(&self, src: &RemotePath, dst: &RemotePath) -> anyhow::Result<()> {
+        unimplemented!()
+    }
 }

libs/remote_storage/tests/test_real_azure.rs

@@ -1,6 +1,6 @@
 use std::collections::HashSet;
 use std::env;
-use std::num::NonZeroUsize;
+use std::num::{NonZeroU32, NonZeroUsize};
 use std::ops::ControlFlow;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -469,6 +469,8 @@ fn create_azure_client(
     let random = rand::thread_rng().gen::<u32>();
 
     let remote_storage_config = RemoteStorageConfig {
+        max_concurrent_syncs: NonZeroUsize::new(100).unwrap(),
+        max_sync_errors: NonZeroU32::new(5).unwrap(),
         storage: RemoteStorageKind::AzureContainer(AzureConfig {
             container_name: remote_storage_azure_container,
             container_region: remote_storage_azure_region,

libs/remote_storage/tests/test_real_s3.rs

@@ -1,6 +1,6 @@
 use std::collections::HashSet;
 use std::env;
-use std::num::NonZeroUsize;
+use std::num::{NonZeroU32, NonZeroUsize};
 use std::ops::ControlFlow;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -396,6 +396,8 @@ fn create_s3_client(
     let random = rand::thread_rng().gen::<u32>();
 
     let remote_storage_config = RemoteStorageConfig {
+        max_concurrent_syncs: NonZeroUsize::new(100).unwrap(),
+        max_sync_errors: NonZeroU32::new(5).unwrap(),
         storage: RemoteStorageKind::AwsS3(S3Config {
             bucket_name: remote_storage_s3_bucket,
             bucket_region: remote_storage_s3_region,

libs/utils/src/auth.rs

@@ -2,7 +2,7 @@
 
 use arc_swap::ArcSwap;
 use serde;
-use std::{borrow::Cow, fmt::Display, fs, sync::Arc};
+use std::{fs, sync::Arc};
 
 use anyhow::Result;
 use camino::Utf8Path;
@@ -11,7 +11,7 @@ use jsonwebtoken::{
 };
 use serde::{Deserialize, Serialize};
 
-use crate::{http::error::ApiError, id::TenantId};
+use crate::id::TenantId;
 
 /// Algorithm to use. We require EdDSA.
 const STORAGE_TOKEN_ALGORITHM: Algorithm = Algorithm::EdDSA;
@@ -54,7 +54,7 @@ impl SwappableJwtAuth {
     pub fn swap(&self, jwt_auth: JwtAuth) {
         self.0.swap(Arc::new(jwt_auth));
     }
-    pub fn decode(&self, token: &str) -> std::result::Result<TokenData<Claims>, AuthError> {
+    pub fn decode(&self, token: &str) -> Result<TokenData<Claims>> {
         self.0.load().decode(token)
     }
 }
@@ -65,24 +65,6 @@ impl std::fmt::Debug for SwappableJwtAuth {
     }
 }
 
-#[derive(Clone, PartialEq, Eq, Hash, Debug)]
-pub struct AuthError(pub Cow<'static, str>);
-
-impl Display for AuthError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.0)
-    }
-}
-
-impl From<AuthError> for ApiError {
-    fn from(_value: AuthError) -> Self {
-        // Don't pass on the value of the AuthError as a precautionary measure.
-        // Being intentionally vague in public error communication hurts debugability
-        // but it is more secure.
-        ApiError::Forbidden("JWT authentication error".to_string())
-    }
-}
-
 pub struct JwtAuth {
     decoding_keys: Vec<DecodingKey>,
     validation: Validation,
@@ -132,7 +114,7 @@ impl JwtAuth {
     /// The function tries the stored decoding keys in succession,
     /// and returns the first yielding a successful result.
     /// If there is no working decoding key, it returns the last error.
-    pub fn decode(&self, token: &str) -> std::result::Result<TokenData<Claims>, AuthError> {
+    pub fn decode(&self, token: &str) -> Result<TokenData<Claims>> {
         let mut res = None;
         for decoding_key in &self.decoding_keys {
             res = Some(decode(token, decoding_key, &self.validation));
@@ -141,9 +123,9 @@ impl JwtAuth {
             }
         }
         if let Some(res) = res {
-            res.map_err(|e| AuthError(Cow::Owned(e.to_string())))
+            res.map_err(anyhow::Error::new)
         } else {
-            Err(AuthError(Cow::Borrowed("no JWT decoding keys configured")))
+            anyhow::bail!("no JWT decoding keys configured")
         }
     }
 }
@@ -184,9 +166,9 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
 "#;
 
     #[test]
-    fn test_decode() {
+    fn test_decode() -> Result<(), anyhow::Error> {
         let expected_claims = Claims {
-            tenant_id: Some(TenantId::from_str("3d1f7595b468230304e0b73cecbcb081").unwrap()),
+            tenant_id: Some(TenantId::from_str("3d1f7595b468230304e0b73cecbcb081")?),
             scope: Scope::Tenant,
         };
 
@@ -205,24 +187,28 @@ MC4CAQAwBQYDK2VwBCIEID/Drmc1AA6U/znNRWpF3zEGegOATQxfkdWxitcOMsIH
         let encoded_eddsa = "eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJzY29wZSI6InRlbmFudCIsInRlbmFudF9pZCI6IjNkMWY3NTk1YjQ2ODIzMDMwNGUwYjczY2VjYmNiMDgxIiwiaXNzIjoibmVvbi5jb250cm9scGxhbmUiLCJleHAiOjE3MDkyMDA4NzksImlhdCI6MTY3ODQ0MjQ3OX0.U3eA8j-uU-JnhzeO3EDHRuXLwkAUFCPxtGHEgw6p7Ccc3YRbFs2tmCdbD9PZEXP-XsxSeBQi1FY0YPcT3NXADw";
 
         // Check it can be validated with the public key
-        let auth = JwtAuth::new(vec![DecodingKey::from_ed_pem(TEST_PUB_KEY_ED25519).unwrap()]);
-        let claims_from_token = auth.decode(encoded_eddsa).unwrap().claims;
+        let auth = JwtAuth::new(vec![DecodingKey::from_ed_pem(TEST_PUB_KEY_ED25519)?]);
+        let claims_from_token = auth.decode(encoded_eddsa)?.claims;
         assert_eq!(claims_from_token, expected_claims);
+
+        Ok(())
     }
 
     #[test]
-    fn test_encode() {
+    fn test_encode() -> Result<(), anyhow::Error> {
         let claims = Claims {
-            tenant_id: Some(TenantId::from_str("3d1f7595b468230304e0b73cecbcb081").unwrap()),
+            tenant_id: Some(TenantId::from_str("3d1f7595b468230304e0b73cecbcb081")?),
             scope: Scope::Tenant,
         };
 
-        let encoded = encode_from_key_file(&claims, TEST_PRIV_KEY_ED25519).unwrap();
+        let encoded = encode_from_key_file(&claims, TEST_PRIV_KEY_ED25519)?;
 
         // decode it back
-        let auth = JwtAuth::new(vec![DecodingKey::from_ed_pem(TEST_PUB_KEY_ED25519).unwrap()]);
-        let decoded = auth.decode(&encoded).unwrap();
+        let auth = JwtAuth::new(vec![DecodingKey::from_ed_pem(TEST_PUB_KEY_ED25519)?]);
+        let decoded = auth.decode(&encoded)?;
 
         assert_eq!(decoded.claims, claims);
+
+        Ok(())
     }
 }

libs/utils/src/http/endpoint.rs

@@ -1,4 +1,4 @@
-use crate::auth::{AuthError, Claims, SwappableJwtAuth};
+use crate::auth::{Claims, SwappableJwtAuth};
 use crate::http::error::{api_error_handler, route_error_handler, ApiError};
 use anyhow::Context;
 use hyper::header::{HeaderName, AUTHORIZATION};
@@ -400,11 +400,9 @@ pub fn auth_middleware<B: hyper::body::HttpBody + Send + Sync + 'static>(
                     })?;
                     let token = parse_token(header_value)?;
 
-                    let data = auth.decode(token).map_err(|err| {
-                        warn!("Authentication error: {err}");
-                        // Rely on From<AuthError> for ApiError impl
-                        err
-                    })?;
+                    let data = auth
+                        .decode(token)
+                        .map_err(|_| ApiError::Unauthorized("malformed jwt token".to_string()))?;
                     req.set_context(data.claims);
                 }
                 None => {
@@ -452,11 +450,12 @@ where
 
 pub fn check_permission_with(
     req: &Request<Body>,
-    check_permission: impl Fn(&Claims) -> Result<(), AuthError>,
+    check_permission: impl Fn(&Claims) -> Result<(), anyhow::Error>,
 ) -> Result<(), ApiError> {
     match req.context::<Claims>() {
-        Some(claims) => Ok(check_permission(&claims)
-            .map_err(|_err| ApiError::Forbidden("JWT authentication error".to_string()))?),
+        Some(claims) => {
+            Ok(check_permission(&claims).map_err(|err| ApiError::Forbidden(err.to_string()))?)
+        }
         None => Ok(()), // claims is None because auth is disabled
     }
 }

libs/utils/src/http/error.rs

@@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize};
 use std::borrow::Cow;
 use std::error::Error as StdError;
 use thiserror::Error;
-use tracing::{error, info, warn};
+use tracing::{error, info};
 
 #[derive(Debug, Error)]
 pub enum ApiError {
@@ -118,9 +118,6 @@ pub fn api_error_handler(api_error: ApiError) -> Response<Body> {
     // Print a stack trace for Internal Server errors
 
     match api_error {
-        ApiError::Forbidden(_) | ApiError::Unauthorized(_) => {
-            warn!("Error processing HTTP request: {api_error:#}")
-        }
         ApiError::ResourceUnavailable(_) => info!("Error processing HTTP request: {api_error:#}"),
         ApiError::NotFound(_) => info!("Error processing HTTP request: {api_error:#}"),
         ApiError::InternalServerError(_) => error!("Error processing HTTP request: {api_error:?}"),

libs/utils/src/seqwait.rs

@@ -125,9 +125,6 @@ where
             // Wake everyone with an error.
             let mut internal = self.internal.lock().unwrap();
 
-            // Block any future waiters from starting
-            internal.shutdown = true;
-
             // This will steal the entire waiters map.
             // When we drop it all waiters will be woken.
             mem::take(&mut internal.waiters)

libs/walproposer/src/api_bindings.rs

@@ -188,7 +188,6 @@ extern "C" fn recovery_download(
     }
 }
 
-#[allow(clippy::unnecessary_cast)]
 extern "C" fn wal_read(
     sk: *mut Safekeeper,
     buf: *mut ::std::os::raw::c_char,
@@ -422,7 +421,6 @@ impl std::fmt::Display for Level {
 }
 
 /// Take ownership of `Vec<u8>` from StringInfoData.
-#[allow(clippy::unnecessary_cast)]
 pub(crate) fn take_vec_u8(pg: &mut StringInfoData) -> Option<Vec<u8>> {
     if pg.data.is_null() {
         return None;

libs/walproposer/src/walproposer.rs

@@ -186,7 +186,7 @@ impl Wrapper {
             .unwrap()
             .into_bytes_with_nul();
         assert!(safekeepers_list_vec.len() == safekeepers_list_vec.capacity());
-        let safekeepers_list = safekeepers_list_vec.as_mut_ptr() as *mut std::ffi::c_char;
+        let safekeepers_list = safekeepers_list_vec.as_mut_ptr() as *mut i8;
 
         let callback_data = Box::into_raw(Box::new(api)) as *mut ::std::os::raw::c_void;
 

pageserver/src/auth.rs

@@ -1,21 +1,22 @@
-use utils::auth::{AuthError, Claims, Scope};
+use anyhow::{bail, Result};
+use utils::auth::{Claims, Scope};
 use utils::id::TenantId;
 
-pub fn check_permission(claims: &Claims, tenant_id: Option<TenantId>) -> Result<(), AuthError> {
+pub fn check_permission(claims: &Claims, tenant_id: Option<TenantId>) -> Result<()> {
     match (&claims.scope, tenant_id) {
-        (Scope::Tenant, None) => Err(AuthError(
-            "Attempt to access management api with tenant scope. Permission denied".into(),
-        )),
+        (Scope::Tenant, None) => {
+            bail!("Attempt to access management api with tenant scope. Permission denied")
+        }
         (Scope::Tenant, Some(tenant_id)) => {
             if claims.tenant_id.unwrap() != tenant_id {
-                return Err(AuthError("Tenant id mismatch. Permission denied".into()));
+                bail!("Tenant id mismatch. Permission denied")
             }
             Ok(())
         }
         (Scope::PageServerApi, None) => Ok(()), // access to management api for PageServerApi scope
         (Scope::PageServerApi, Some(_)) => Ok(()), // access to tenant api using PageServerApi scope
-        (Scope::SafekeeperData, _) => Err(AuthError(
-            "SafekeeperData scope makes no sense for Pageserver".into(),
-        )),
+        (Scope::SafekeeperData, _) => {
+            bail!("SafekeeperData scope makes no sense for Pageserver")
+        }
     }
 }

pageserver/src/config.rs

@@ -1314,6 +1314,12 @@ broker_endpoint = '{broker_endpoint}'
             assert_eq!(
                 parsed_remote_storage_config,
                 RemoteStorageConfig {
+                    max_concurrent_syncs: NonZeroUsize::new(
+                        remote_storage::DEFAULT_REMOTE_STORAGE_MAX_CONCURRENT_SYNCS
+                    )
+                        .unwrap(),
+                    max_sync_errors: NonZeroU32::new(remote_storage::DEFAULT_REMOTE_STORAGE_MAX_SYNC_ERRORS)
+                        .unwrap(),
                     storage: RemoteStorageKind::LocalFs(local_storage_path.clone()),
                 },
                 "Remote storage config should correctly parse the local FS config and fill other storage defaults"
@@ -1374,6 +1380,8 @@ broker_endpoint = '{broker_endpoint}'
             assert_eq!(
                 parsed_remote_storage_config,
                 RemoteStorageConfig {
+                    max_concurrent_syncs,
+                    max_sync_errors,
                     storage: RemoteStorageKind::AwsS3(S3Config {
                         bucket_name: bucket_name.clone(),
                         bucket_region: bucket_region.clone(),

pageserver/src/deletion_queue.rs

@@ -893,6 +893,14 @@ mod test {
         std::fs::create_dir_all(remote_fs_dir)?;
         let remote_fs_dir = harness.conf.workdir.join("remote_fs").canonicalize_utf8()?;
         let storage_config = RemoteStorageConfig {
+            max_concurrent_syncs: std::num::NonZeroUsize::new(
+                remote_storage::DEFAULT_REMOTE_STORAGE_MAX_CONCURRENT_SYNCS,
+            )
+            .unwrap(),
+            max_sync_errors: std::num::NonZeroU32::new(
+                remote_storage::DEFAULT_REMOTE_STORAGE_MAX_SYNC_ERRORS,
+            )
+            .unwrap(),
             storage: RemoteStorageKind::LocalFs(remote_fs_dir.clone()),
         };
         let storage = GenericRemoteStorage::from_config(&storage_config).unwrap();

pageserver/src/http/openapi_spec.yml

@@ -352,8 +352,7 @@ paths:
           in: query
           required: true
           schema:
-            type: string
-            format: hex
+            type: integer
           description: A LSN to get the timestamp
       responses:
         "200":

pageserver/src/http/routes.rs

@@ -303,7 +303,11 @@ async fn build_timeline_info(
         // we're executing this function, we will outlive the timeline on-disk state.
         info.current_logical_size_non_incremental = Some(
             timeline
-                .get_current_logical_size_non_incremental(info.last_record_lsn, ctx)
+                .get_current_logical_size_non_incremental(
+                    info.last_record_lsn,
+                    CancellationToken::new(),
+                    ctx,
+                )
                 .await?,
         );
     }
@@ -1667,8 +1671,6 @@ where
         );
 
         match handle.await {
-            // TODO: never actually return Err from here, always Ok(...) so that we can log
-            // spanned errors. Call api_error_handler instead and return appropriate Body.
             Ok(result) => result,
             Err(e) => {
                 // The handler task panicked. We have a global panic handler that logs the

pageserver/src/metrics.rs

@@ -1225,6 +1225,15 @@ pub(crate) static WAL_REDO_TIME: Lazy<Histogram> = Lazy::new(|| {
     .expect("failed to define a metric")
 });
 
+pub(crate) static WAL_REDO_WAIT_TIME: Lazy<Histogram> = Lazy::new(|| {
+    register_histogram!(
+        "pageserver_wal_redo_wait_seconds",
+        "Time spent waiting for access to the Postgres WAL redo process",
+        redo_histogram_time_buckets!(),
+    )
+    .expect("failed to define a metric")
+});
+
 pub(crate) static WAL_REDO_RECORDS_HISTOGRAM: Lazy<Histogram> = Lazy::new(|| {
     register_histogram!(
         "pageserver_wal_redo_records_histogram",
@@ -1919,6 +1928,7 @@ pub fn preinitialize_metrics() {
         &READ_NUM_FS_LAYERS,
         &WAIT_LSN_TIME,
         &WAL_REDO_TIME,
+        &WAL_REDO_WAIT_TIME,
         &WAL_REDO_RECORDS_HISTOGRAM,
         &WAL_REDO_BYTES_HISTOGRAM,
     ]

pageserver/src/page_service.rs

@@ -218,27 +218,9 @@ async fn page_service_conn_main(
     // no write timeout is used, because the kernel is assumed to error writes after some time.
     let mut socket = tokio_io_timeout::TimeoutReader::new(socket);
 
-    let default_timeout_ms = 10 * 60 * 1000; // 10 minutes by default
-    let socket_timeout_ms = (|| {
-        fail::fail_point!("simulated-bad-compute-connection", |avg_timeout_ms| {
-            // Exponential distribution for simulating
-            // poor network conditions, expect about avg_timeout_ms to be around 15
-            // in tests
-            if let Some(avg_timeout_ms) = avg_timeout_ms {
-                let avg = avg_timeout_ms.parse::<i64>().unwrap() as f32;
-                let u = rand::random::<f32>();
-                ((1.0 - u).ln() / (-avg)) as u64
-            } else {
-                default_timeout_ms
-            }
-        });
-        default_timeout_ms
-    })();
-
-    // A timeout here does not mean the client died, it can happen if it's just idle for
-    // a while: we will tear down this PageServerHandler and instantiate a new one if/when
-    // they reconnect.
-    socket.set_timeout(Some(std::time::Duration::from_millis(socket_timeout_ms)));
+    // timeout should be lower, but trying out multiple days for
+    // <https://github.com/neondatabase/neon/issues/4205>
+    socket.set_timeout(Some(std::time::Duration::from_secs(60 * 60 * 24 * 3)));
     let socket = std::pin::pin!(socket);
 
     // XXX: pgbackend.run() should take the connection_ctx,
@@ -512,11 +494,7 @@ impl PageServerHandler {
             };
 
             if let Err(e) = &response {
-                // Requests may fail as soon as we are Stopping, even if the Timeline's cancellation token wasn't fired yet,
-                // because wait_lsn etc will drop out
-                // is_stopping(): [`Timeline::flush_and_shutdown`] has entered
-                // is_canceled(): [`Timeline::shutdown`]` has entered
-                if timeline.cancel.is_cancelled() || timeline.is_stopping() {
+                if timeline.cancel.is_cancelled() {
                     // If we fail to fulfil a request during shutdown, which may be _because_ of
                     // shutdown, then do not send the error to the client.  Instead just drop the
                     // connection.
@@ -920,7 +898,7 @@ impl PageServerHandler {
 
     // when accessing management api supply None as an argument
     // when using to authorize tenant pass corresponding tenant id
-    fn check_permission(&self, tenant_id: Option<TenantId>) -> Result<(), QueryError> {
+    fn check_permission(&self, tenant_id: Option<TenantId>) -> anyhow::Result<()> {
         if self.auth.is_none() {
             // auth is set to Trust, nothing to check so just return ok
             return Ok(());
@@ -932,7 +910,7 @@ impl PageServerHandler {
             .claims
             .as_ref()
             .expect("claims presence already checked");
-        check_permission(claims, tenant_id).map_err(|e| QueryError::Unauthorized(e.0))
+        check_permission(claims, tenant_id)
     }
 
     /// Shorthand for getting a reference to a Timeline of an Active tenant.
@@ -971,17 +949,16 @@ where
             .auth
             .as_ref()
             .unwrap()
-            .decode(str::from_utf8(jwt_response).context("jwt response is not UTF-8")?)
-            .map_err(|e| QueryError::Unauthorized(e.0))?;
+            .decode(str::from_utf8(jwt_response).context("jwt response is not UTF-8")?)?;
 
         if matches!(data.claims.scope, Scope::Tenant) && data.claims.tenant_id.is_none() {
-            return Err(QueryError::Unauthorized(
-                "jwt token scope is Tenant, but tenant id is missing".into(),
-            ));
+            return Err(QueryError::Other(anyhow::anyhow!(
+                "jwt token scope is Tenant, but tenant id is missing"
+            )));
         }
 
-        debug!(
-            "jwt scope check succeeded for scope: {:#?} by tenant id: {:?}",
+        info!(
+            "jwt auth succeeded for scope: {:#?} by tenant id: {:?}",
             data.claims.scope, data.claims.tenant_id,
         );
 
@@ -1003,13 +980,9 @@ where
         pgb: &mut PostgresBackend<IO>,
         query_string: &str,
     ) -> Result<(), QueryError> {
-        fail::fail_point!("simulated-bad-compute-connection", |_| {
-            info!("Hit failpoint for bad connection");
-            Err(QueryError::SimulatedConnectionError)
-        });
-
         let ctx = self.connection_ctx.attached_child();
         debug!("process query {query_string:?}");
+
         if query_string.starts_with("pagestream ") {
             let (_, params_raw) = query_string.split_at("pagestream ".len());
             let params = params_raw.split(' ').collect::<Vec<_>>();

pageserver/src/pgdatadir_mapping.rs

@@ -21,6 +21,7 @@ use serde::{Deserialize, Serialize};
 use std::collections::{hash_map, HashMap, HashSet};
 use std::ops::ControlFlow;
 use std::ops::Range;
+use tokio_util::sync::CancellationToken;
 use tracing::{debug, trace, warn};
 use utils::{bin_ser::BeSer, lsn::Lsn};
 
@@ -577,6 +578,7 @@ impl Timeline {
     pub async fn get_current_logical_size_non_incremental(
         &self,
         lsn: Lsn,
+        cancel: CancellationToken,
         ctx: &RequestContext,
     ) -> Result<u64, CalculateLogicalSizeError> {
         crate::tenant::debug_assert_current_span_has_tenant_and_timeline_id();
@@ -588,7 +590,7 @@ impl Timeline {
         let mut total_size: u64 = 0;
         for (spcnode, dbnode) in dbdir.dbdirs.keys() {
             for rel in self.list_rels(*spcnode, *dbnode, lsn, ctx).await? {
-                if self.cancel.is_cancelled() {
+                if cancel.is_cancelled() {
                     return Err(CalculateLogicalSizeError::Cancelled);
                 }
                 let relsize_key = rel_size_to_key(rel);

pageserver/src/tenant.rs

@@ -1841,13 +1841,7 @@ impl Tenant {
             timelines.values().for_each(|timeline| {
                 let timeline = Arc::clone(timeline);
                 let span = Span::current();
-                js.spawn(async move {
-                    if freeze_and_flush {
-                        timeline.flush_and_shutdown().instrument(span).await
-                    } else {
-                        timeline.shutdown().instrument(span).await
-                    }
-                });
+                js.spawn(async move { timeline.shutdown(freeze_and_flush).instrument(span).await });
             })
         };
         tracing::info!("Waiting for timelines...");
@@ -3534,6 +3528,10 @@ pub(crate) mod harness {
             let remote_fs_dir = conf.workdir.join("localfs");
             std::fs::create_dir_all(&remote_fs_dir).unwrap();
             let config = RemoteStorageConfig {
+                // TODO: why not remote_storage::DEFAULT_REMOTE_STORAGE_MAX_CONCURRENT_SYNCS,
+                max_concurrent_syncs: std::num::NonZeroUsize::new(2_000_000).unwrap(),
+                // TODO: why not remote_storage::DEFAULT_REMOTE_STORAGE_MAX_SYNC_ERRORS,
+                max_sync_errors: std::num::NonZeroU32::new(3_000_000).unwrap(),
                 storage: RemoteStorageKind::LocalFs(remote_fs_dir.clone()),
             };
             let remote_storage = GenericRemoteStorage::from_config(&config).unwrap();
@@ -4733,7 +4731,7 @@ mod tests {
             // Keeps uninit mark in place
             let raw_tline = tline.raw_timeline().unwrap();
             raw_tline
-                .shutdown()
+                .shutdown(false)
                 .instrument(info_span!("test_shutdown", tenant_id=%raw_tline.tenant_id))
                 .await;
             std::mem::forget(tline);

pageserver/src/tenant/blob_io.rs

@@ -327,7 +327,7 @@ mod tests {
                 let mut sz: u16 = rng.gen();
                 // Make 50% of the arrays small
                 if rng.gen() {
-                    sz &= 63;
+                    sz |= 63;
                 }
                 random_array(sz.into())
             })

pageserver/src/tenant/mgr.rs

@@ -566,10 +566,8 @@ pub(crate) async fn shutdown_all_tenants() {
 async fn shutdown_all_tenants0(tenants: &std::sync::RwLock<TenantsMap>) {
     use utils::completion;
 
-    let mut join_set = JoinSet::new();
-
-    // Atomically, 1. create the shutdown tasks and 2. prevent creation of new tenants.
-    let (total_in_progress, total_attached) = {
+    // Atomically, 1. extract the list of tenants to shut down and 2. prevent creation of new tenants.
+    let (in_progress_ops, tenants_to_shut_down) = {
         let mut m = tenants.write().unwrap();
         match &mut *m {
             TenantsMap::Initializing => {
@@ -579,67 +577,78 @@ async fn shutdown_all_tenants0(tenants: &std::sync::RwLock<TenantsMap>) {
             }
             TenantsMap::Open(tenants) => {
                 let mut shutdown_state = HashMap::new();
-                let mut total_in_progress = 0;
-                let mut total_attached = 0;
+                let mut in_progress_ops = Vec::new();
+                let mut tenants_to_shut_down = Vec::new();
 
-                for (tenant_id, v) in tenants.drain() {
+                for (k, v) in tenants.drain() {
                     match v {
                         TenantSlot::Attached(t) => {
-                            shutdown_state.insert(tenant_id, TenantSlot::Attached(t.clone()));
-                            join_set.spawn(
-                                async move {
-                                    let freeze_and_flush = true;
-
-                                    let res = {
-                                        let (_guard, shutdown_progress) = completion::channel();
-                                        t.shutdown(shutdown_progress, freeze_and_flush).await
-                                    };
-
-                                    if let Err(other_progress) = res {
-                                        // join the another shutdown in progress
-                                        other_progress.wait().await;
-                                    }
-
-                                    // we cannot afford per tenant logging here, because if s3 is degraded, we are
-                                    // going to log too many lines
-                                    debug!("tenant successfully stopped");
-                                }
-                                .instrument(info_span!("shutdown", %tenant_id)),
-                            );
-
-                            total_attached += 1;
+                            tenants_to_shut_down.push(t.clone());
+                            shutdown_state.insert(k, TenantSlot::Attached(t));
                         }
                         TenantSlot::Secondary => {
-                            shutdown_state.insert(tenant_id, TenantSlot::Secondary);
+                            shutdown_state.insert(k, TenantSlot::Secondary);
                         }
                         TenantSlot::InProgress(notify) => {
                             // InProgress tenants are not visible in TenantsMap::ShuttingDown: we will
                             // wait for their notifications to fire in this function.
-                            join_set.spawn(async move {
-                                notify.wait().await;
-                            });
-
-                            total_in_progress += 1;
+                            in_progress_ops.push(notify);
                         }
                     }
                 }
                 *m = TenantsMap::ShuttingDown(shutdown_state);
-                (total_in_progress, total_attached)
+                (in_progress_ops, tenants_to_shut_down)
             }
             TenantsMap::ShuttingDown(_) => {
+                // TODO: it is possible that detach and shutdown happen at the same time. as a
+                // result, during shutdown we do not wait for detach.
                 error!("already shutting down, this function isn't supposed to be called more than once");
                 return;
             }
         }
     };
 
-    let started_at = std::time::Instant::now();
-
     info!(
         "Waiting for {} InProgress tenants and {} Attached tenants to shut down",
-        total_in_progress, total_attached
+        in_progress_ops.len(),
+        tenants_to_shut_down.len()
     );
 
+    for barrier in in_progress_ops {
+        barrier.wait().await;
+    }
+
+    info!(
+        "InProgress tenants shut down, waiting for {} Attached tenants to shut down",
+        tenants_to_shut_down.len()
+    );
+    let started_at = std::time::Instant::now();
+    let mut join_set = JoinSet::new();
+    for tenant in tenants_to_shut_down {
+        let tenant_id = tenant.get_tenant_id();
+        join_set.spawn(
+            async move {
+                let freeze_and_flush = true;
+
+                let res = {
+                    let (_guard, shutdown_progress) = completion::channel();
+                    tenant.shutdown(shutdown_progress, freeze_and_flush).await
+                };
+
+                if let Err(other_progress) = res {
+                    // join the another shutdown in progress
+                    other_progress.wait().await;
+                }
+
+                // we cannot afford per tenant logging here, because if s3 is degraded, we are
+                // going to log too many lines
+
+                debug!("tenant successfully stopped");
+            }
+            .instrument(info_span!("shutdown", %tenant_id)),
+        );
+    }
+
     let total = join_set.len();
     let mut panicked = 0;
     let mut buffering = true;
@@ -652,7 +661,7 @@ async fn shutdown_all_tenants0(tenants: &std::sync::RwLock<TenantsMap>) {
                 match joined {
                     Ok(()) => {}
                     Err(join_error) if join_error.is_cancelled() => {
-                        unreachable!("we are not cancelling any of the tasks");
+                        unreachable!("we are not cancelling any of the futures");
                     }
                     Err(join_error) if join_error.is_panic() => {
                         // cannot really do anything, as this panic is likely a bug

pageserver/src/tenant/size.rs

@@ -6,6 +6,7 @@ use std::sync::Arc;
 use anyhow::{bail, Context};
 use tokio::sync::oneshot::error::RecvError;
 use tokio::sync::Semaphore;
+use tokio_util::sync::CancellationToken;
 
 use crate::context::RequestContext;
 use crate::pgdatadir_mapping::CalculateLogicalSizeError;
@@ -349,6 +350,10 @@ async fn fill_logical_sizes(
     // our advantage with `?` error handling.
     let mut joinset = tokio::task::JoinSet::new();
 
+    let cancel = tokio_util::sync::CancellationToken::new();
+    // be sure to cancel all spawned tasks if we are dropped
+    let _dg = cancel.clone().drop_guard();
+
     // For each point that would benefit from having a logical size available,
     // spawn a Task to fetch it, unless we have it cached already.
     for seg in segments.iter() {
@@ -366,8 +371,15 @@ async fn fill_logical_sizes(
                 let parallel_size_calcs = Arc::clone(limit);
                 let ctx = ctx.attached_child();
                 joinset.spawn(
-                    calculate_logical_size(parallel_size_calcs, timeline, lsn, cause, ctx)
-                        .in_current_span(),
+                    calculate_logical_size(
+                        parallel_size_calcs,
+                        timeline,
+                        lsn,
+                        cause,
+                        ctx,
+                        cancel.child_token(),
+                    )
+                    .in_current_span(),
                 );
             }
             e.insert(cached_size);
@@ -475,13 +487,14 @@ async fn calculate_logical_size(
     lsn: utils::lsn::Lsn,
     cause: LogicalSizeCalculationCause,
     ctx: RequestContext,
+    cancel: CancellationToken,
 ) -> Result<TimelineAtLsnSizeResult, RecvError> {
     let _permit = tokio::sync::Semaphore::acquire_owned(limit)
         .await
         .expect("global semaphore should not had been closed");
 
     let size_res = timeline
-        .spawn_ondemand_logical_size_calculation(lsn, cause, ctx)
+        .spawn_ondemand_logical_size_calculation(lsn, cause, ctx, cancel)
         .instrument(info_span!("spawn_ondemand_logical_size_calculation"))
         .await?;
     Ok(TimelineAtLsnSizeResult(timeline, lsn, size_res))

pageserver/src/tenant/storage_layer/layer.rs

@@ -251,7 +251,6 @@ impl Layer {
 
         layer
             .get_value_reconstruct_data(key, lsn_range, reconstruct_data, &self.0, ctx)
-            .instrument(tracing::info_span!("get_value_reconstruct_data", layer=%self))
             .await
     }
 
@@ -1212,10 +1211,8 @@ impl DownloadedLayer {
             // this will be a permanent failure
             .context("load layer");
 
-            if let Err(e) = res.as_ref() {
+            if res.is_err() {
                 LAYER_IMPL_METRICS.inc_permanent_loading_failures();
-                // TODO(#5815): we are not logging all errors, so temporarily log them here as well
-                tracing::error!("layer loading failed permanently: {e:#}");
             }
             res
         };
@@ -1294,7 +1291,6 @@ impl ResidentLayer {
     }
 
     /// Loads all keys stored in the layer. Returns key, lsn and value size.
-    #[tracing::instrument(skip_all, fields(layer=%self))]
     pub(crate) async fn load_keys<'a>(
         &'a self,
         ctx: &RequestContext,

pageserver/src/tenant/timeline.rs

@@ -36,6 +36,7 @@ use std::time::{Duration, Instant, SystemTime};
 use crate::context::{
     AccessStatsBehavior, DownloadBehavior, RequestContext, RequestContextBuilder,
 };
+use crate::deletion_queue::DeletionQueueClient;
 use crate::tenant::storage_layer::delta_layer::DeltaEntry;
 use crate::tenant::storage_layer::{
     AsLayerDesc, DeltaLayerWriter, EvictionError, ImageLayerWriter, InMemoryLayer, Layer,
@@ -49,7 +50,6 @@ use crate::tenant::{
     metadata::{save_metadata, TimelineMetadata},
     par_fsync,
 };
-use crate::{deletion_queue::DeletionQueueClient, tenant::remote_timeline_client::StopError};
 
 use crate::config::PageServerConf;
 use crate::keyspace::{KeyPartitioning, KeySpace, KeySpaceRandomAccum};
@@ -247,7 +247,7 @@ pub struct Timeline {
     /// the flush finishes. You can use that to wait for the flush to finish.
     layer_flush_start_tx: tokio::sync::watch::Sender<u64>,
     /// to be notified when layer flushing has finished, subscribe to the layer_flush_done channel
-    layer_flush_done_tx: tokio::sync::watch::Sender<(u64, Result<(), FlushLayerError>)>,
+    layer_flush_done_tx: tokio::sync::watch::Sender<(u64, anyhow::Result<()>)>,
 
     /// Layer removal lock.
     /// A lock to ensure that no layer of the timeline is removed concurrently by other tasks.
@@ -374,19 +374,6 @@ pub enum PageReconstructError {
     WalRedo(anyhow::Error),
 }
 
-#[derive(thiserror::Error, Debug)]
-enum FlushLayerError {
-    /// Timeline cancellation token was cancelled
-    #[error("timeline shutting down")]
-    Cancelled,
-
-    #[error(transparent)]
-    PageReconstructError(#[from] PageReconstructError),
-
-    #[error(transparent)]
-    Other(#[from] anyhow::Error),
-}
-
 impl std::fmt::Debug for PageReconstructError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
         match self {
@@ -904,16 +891,15 @@ impl Timeline {
         self.launch_eviction_task(background_jobs_can_start);
     }
 
-    /// Graceful shutdown, may do a lot of I/O as we flush any open layers to disk and then
-    /// also to remote storage.  This method can easily take multiple seconds for a busy timeline.
-    ///
-    /// While we are flushing, we continue to accept read I/O.
     #[instrument(skip_all, fields(timeline_id=%self.timeline_id))]
-    pub(crate) async fn flush_and_shutdown(&self) {
+    pub async fn shutdown(self: &Arc<Self>, freeze_and_flush: bool) {
         debug_assert_current_span_has_tenant_and_timeline_id();
 
-        // Stop ingesting data, so that we are not still writing to an InMemoryLayer while
-        // trying to flush
+        // Signal any subscribers to our cancellation token to drop out
+        tracing::debug!("Cancelling CancellationToken");
+        self.cancel.cancel();
+
+        // prevent writes to the InMemoryLayer
         tracing::debug!("Waiting for WalReceiverManager...");
         task_mgr::shutdown_tasks(
             Some(TaskKind::WalReceiverManager),
@@ -922,70 +908,40 @@ impl Timeline {
         )
         .await;
 
-        // Since we have shut down WAL ingest, we should not let anyone start waiting for the LSN to advance
-        self.last_record_lsn.shutdown();
-
         // now all writers to InMemory layer are gone, do the final flush if requested
-        match self.freeze_and_flush().await {
-            Ok(_) => {
-                // drain the upload queue
-                if let Some(client) = self.remote_client.as_ref() {
-                    // if we did not wait for completion here, it might be our shutdown process
-                    // didn't wait for remote uploads to complete at all, as new tasks can forever
-                    // be spawned.
-                    //
-                    // what is problematic is the shutting down of RemoteTimelineClient, because
-                    // obviously it does not make sense to stop while we wait for it, but what
-                    // about corner cases like s3 suddenly hanging up?
-                    if let Err(e) = client.wait_completion().await {
-                        // Non-fatal.  Shutdown is infallible.  Failures to flush just mean that
-                        // we have some extra WAL replay to do next time the timeline starts.
-                        warn!("failed to flush to remote storage: {e:#}");
-                    }
+        if freeze_and_flush {
+            match self.freeze_and_flush().await {
+                Ok(()) => {}
+                Err(e) => {
+                    warn!("failed to freeze and flush: {e:#}");
+                    return; // TODO: should probably drain remote timeline client anyways?
                 }
             }
-            Err(e) => {
-                // Non-fatal.  Shutdown is infallible.  Failures to flush just mean that
-                // we have some extra WAL replay to do next time the timeline starts.
-                warn!("failed to freeze and flush: {e:#}");
+
+            // drain the upload queue
+            let res = if let Some(client) = self.remote_client.as_ref() {
+                // if we did not wait for completion here, it might be our shutdown process
+                // didn't wait for remote uploads to complete at all, as new tasks can forever
+                // be spawned.
+                //
+                // what is problematic is the shutting down of RemoteTimelineClient, because
+                // obviously it does not make sense to stop while we wait for it, but what
+                // about corner cases like s3 suddenly hanging up?
+                client.wait_completion().await
+            } else {
+                Ok(())
+            };
+
+            if let Err(e) = res {
+                warn!("failed to await for frozen and flushed uploads: {e:#}");
             }
         }
 
-        self.shutdown().await;
-    }
-
-    /// Shut down immediately, without waiting for any open layers to flush to disk.  This is a subset of
-    /// the graceful [`Timeline::flush_and_shutdown`] function.
-    pub(crate) async fn shutdown(&self) {
-        // Signal any subscribers to our cancellation token to drop out
-        tracing::debug!("Cancelling CancellationToken");
-        self.cancel.cancel();
-
         // Page request handlers might be waiting for LSN to advance: they do not respect Timeline::cancel
         // while doing so.
         self.last_record_lsn.shutdown();
 
-        // Shut down the layer flush task before the remote client, as one depends on the other
-        task_mgr::shutdown_tasks(
-            Some(TaskKind::LayerFlushTask),
-            Some(self.tenant_id),
-            Some(self.timeline_id),
-        )
-        .await;
-
-        // Shut down remote timeline client: this gracefully moves its metadata into its Stopping state in
-        // case our caller wants to use that for a deletion
-        if let Some(remote_client) = self.remote_client.as_ref() {
-            match remote_client.stop() {
-                Ok(()) => {}
-                Err(StopError::QueueUninitialized) => {
-                    // Shutting down during initialization is legal
-                }
-            }
-        }
-
         tracing::debug!("Waiting for tasks...");
-
         task_mgr::shutdown_tasks(None, Some(self.tenant_id), Some(self.timeline_id)).await;
 
         // Finally wait until any gate-holders are complete
@@ -1029,12 +985,7 @@ impl Timeline {
             reason,
             backtrace: backtrace_str,
         };
-        self.set_state(broken_state);
-
-        // Although the Broken state is not equivalent to shutdown() (shutdown will be called
-        // later when this tenant is detach or the process shuts down), firing the cancellation token
-        // here avoids the need for other tasks to watch for the Broken state explicitly.
-        self.cancel.cancel();
+        self.set_state(broken_state)
     }
 
     pub fn current_state(&self) -> TimelineState {
@@ -1790,8 +1741,12 @@ impl Timeline {
                 // delay will be terminated by a timeout regardless.
                 let _completion = { self_clone.initial_logical_size_attempt.lock().expect("unexpected initial_logical_size_attempt poisoned").take() };
 
+                // no extra cancellation here, because nothing really waits for this to complete compared
+                // to spawn_ondemand_logical_size_calculation.
+                let cancel = CancellationToken::new();
+
                 let calculated_size = match self_clone
-                    .logical_size_calculation_task(lsn, LogicalSizeCalculationCause::Initial, &background_ctx)
+                    .logical_size_calculation_task(lsn, LogicalSizeCalculationCause::Initial, &background_ctx, cancel)
                     .await
                 {
                     Ok(s) => s,
@@ -1860,6 +1815,7 @@ impl Timeline {
         lsn: Lsn,
         cause: LogicalSizeCalculationCause,
         ctx: RequestContext,
+        cancel: CancellationToken,
     ) -> oneshot::Receiver<Result<u64, CalculateLogicalSizeError>> {
         let (sender, receiver) = oneshot::channel();
         let self_clone = Arc::clone(self);
@@ -1880,7 +1836,7 @@ impl Timeline {
             false,
             async move {
                 let res = self_clone
-                    .logical_size_calculation_task(lsn, cause, &ctx)
+                    .logical_size_calculation_task(lsn, cause, &ctx, cancel)
                     .await;
                 let _ = sender.send(res).ok();
                 Ok(()) // Receiver is responsible for handling errors
@@ -1896,28 +1852,58 @@ impl Timeline {
         lsn: Lsn,
         cause: LogicalSizeCalculationCause,
         ctx: &RequestContext,
+        cancel: CancellationToken,
     ) -> Result<u64, CalculateLogicalSizeError> {
         span::debug_assert_current_span_has_tenant_and_timeline_id();
 
-        let _guard = self.gate.enter();
-
+        let mut timeline_state_updates = self.subscribe_for_state_updates();
         let self_calculation = Arc::clone(self);
 
         let mut calculation = pin!(async {
+            let cancel = cancel.child_token();
             let ctx = ctx.attached_child();
             self_calculation
-                .calculate_logical_size(lsn, cause, &ctx)
+                .calculate_logical_size(lsn, cause, cancel, &ctx)
                 .await
         });
+        let timeline_state_cancellation = async {
+            loop {
+                match timeline_state_updates.changed().await {
+                    Ok(()) => {
+                        let new_state = timeline_state_updates.borrow().clone();
+                        match new_state {
+                            // we're running this job for active timelines only
+                            TimelineState::Active => continue,
+                            TimelineState::Broken { .. }
+                            | TimelineState::Stopping
+                            | TimelineState::Loading => {
+                                break format!("aborted because timeline became inactive (new state: {new_state:?})")
+                            }
+                        }
+                    }
+                    Err(_sender_dropped_error) => {
+                        // can't happen, the sender is not dropped as long as the Timeline exists
+                        break "aborted because state watch was dropped".to_string();
+                    }
+                }
+            }
+        };
+
+        let taskmgr_shutdown_cancellation = async {
+            task_mgr::shutdown_watcher().await;
+            "aborted because task_mgr shutdown requested".to_string()
+        };
 
         tokio::select! {
             res = &mut calculation => { res }
-            _ = self.cancel.cancelled() => {
-                debug!("cancelling logical size calculation for timeline shutdown");
+            reason = timeline_state_cancellation => {
+                debug!(reason = reason, "cancelling calculation");
+                cancel.cancel();
                 calculation.await
             }
-            _ = task_mgr::shutdown_watcher() => {
-                debug!("cancelling logical size calculation for task shutdown");
+            reason = taskmgr_shutdown_cancellation => {
+                debug!(reason = reason, "cancelling calculation");
+                cancel.cancel();
                 calculation.await
             }
         }
@@ -1931,6 +1917,7 @@ impl Timeline {
         &self,
         up_to_lsn: Lsn,
         cause: LogicalSizeCalculationCause,
+        cancel: CancellationToken,
         ctx: &RequestContext,
     ) -> Result<u64, CalculateLogicalSizeError> {
         info!(
@@ -1973,7 +1960,7 @@ impl Timeline {
         };
         let timer = storage_time_metrics.start_timer();
         let logical_size = self
-            .get_current_logical_size_non_incremental(up_to_lsn, ctx)
+            .get_current_logical_size_non_incremental(up_to_lsn, cancel, ctx)
             .await?;
         debug!("calculated logical size: {logical_size}");
         timer.stop_and_record();
@@ -2386,10 +2373,6 @@ impl Timeline {
         info!("started flush loop");
         loop {
             tokio::select! {
-                _ = self.cancel.cancelled() => {
-                    info!("shutting down layer flush task");
-                    break;
-                },
                 _ = task_mgr::shutdown_watcher() => {
                     info!("shutting down layer flush task");
                     break;
@@ -2401,14 +2384,6 @@ impl Timeline {
             let timer = self.metrics.flush_time_histo.start_timer();
             let flush_counter = *layer_flush_start_rx.borrow();
             let result = loop {
-                if self.cancel.is_cancelled() {
-                    info!("dropping out of flush loop for timeline shutdown");
-                    // Note: we do not bother transmitting into [`layer_flush_done_tx`], because
-                    // anyone waiting on that will respect self.cancel as well: they will stop
-                    // waiting at the same time we as drop out of this loop.
-                    return;
-                }
-
                 let layer_to_flush = {
                     let guard = self.layers.read().await;
                     guard.layer_map().frozen_layers.front().cloned()
@@ -2417,18 +2392,9 @@ impl Timeline {
                 let Some(layer_to_flush) = layer_to_flush else {
                     break Ok(());
                 };
-                match self.flush_frozen_layer(layer_to_flush, ctx).await {
-                    Ok(()) => {}
-                    Err(FlushLayerError::Cancelled) => {
-                        info!("dropping out of flush loop for timeline shutdown");
-                        return;
-                    }
-                    err @ Err(
-                        FlushLayerError::Other(_) | FlushLayerError::PageReconstructError(_),
-                    ) => {
-                        error!("could not flush frozen layer: {err:?}");
-                        break err;
-                    }
+                if let Err(err) = self.flush_frozen_layer(layer_to_flush, ctx).await {
+                    error!("could not flush frozen layer: {err:?}");
+                    break Err(err);
                 }
             };
             // Notify any listeners that we're done
@@ -2477,17 +2443,7 @@ impl Timeline {
                 }
             }
             trace!("waiting for flush to complete");
-            tokio::select! {
-                rx_e = rx.changed() => {
-                    rx_e?;
-                },
-                // Cancellation safety: we are not leaving an I/O in-flight for the flush, we're just ignoring
-                // the notification from [`flush_loop`] that it completed.
-                _ = self.cancel.cancelled() => {
-                    tracing::info!("Cancelled layer flush due on timeline shutdown");
-                    return Ok(())
-                }
-            };
+            rx.changed().await?;
             trace!("done")
         }
     }
@@ -2502,7 +2458,7 @@ impl Timeline {
         self: &Arc<Self>,
         frozen_layer: Arc<InMemoryLayer>,
         ctx: &RequestContext,
-    ) -> Result<(), FlushLayerError> {
+    ) -> anyhow::Result<()> {
         // As a special case, when we have just imported an image into the repository,
         // instead of writing out a L0 delta layer, we directly write out image layer
         // files instead. This is possible as long as *all* the data imported into the
@@ -2527,11 +2483,6 @@ impl Timeline {
                 let (partitioning, _lsn) = self
                     .repartition(self.initdb_lsn, self.get_compaction_target_size(), ctx)
                     .await?;
-
-                if self.cancel.is_cancelled() {
-                    return Err(FlushLayerError::Cancelled);
-                }
-
                 // For image layers, we add them immediately into the layer map.
                 (
                     self.create_image_layers(&partitioning, self.initdb_lsn, true, ctx)
@@ -2563,10 +2514,6 @@ impl Timeline {
                 )
             };
 
-        if self.cancel.is_cancelled() {
-            return Err(FlushLayerError::Cancelled);
-        }
-
         let disk_consistent_lsn = Lsn(lsn_range.end.0 - 1);
         let old_disk_consistent_lsn = self.disk_consistent_lsn.load();
 
@@ -2576,10 +2523,6 @@ impl Timeline {
         let metadata = {
             let mut guard = self.layers.write().await;
 
-            if self.cancel.is_cancelled() {
-                return Err(FlushLayerError::Cancelled);
-            }
-
             guard.finish_flush_l0_layer(delta_layer_to_add.as_ref(), &frozen_layer, &self.metrics);
 
             if disk_consistent_lsn != old_disk_consistent_lsn {

pageserver/src/tenant/timeline/eviction_task.rs

@@ -326,7 +326,8 @@ impl Timeline {
         match state.last_layer_access_imitation {
             Some(ts) if ts.elapsed() < inter_imitate_period => { /* no need to run */ }
             _ => {
-                self.imitate_timeline_cached_layer_accesses(ctx).await;
+                self.imitate_timeline_cached_layer_accesses(cancel, ctx)
+                    .await;
                 state.last_layer_access_imitation = Some(tokio::time::Instant::now())
             }
         }
@@ -366,12 +367,21 @@ impl Timeline {
 
     /// Recompute the values which would cause on-demand downloads during restart.
     #[instrument(skip_all)]
-    async fn imitate_timeline_cached_layer_accesses(&self, ctx: &RequestContext) {
+    async fn imitate_timeline_cached_layer_accesses(
+        &self,
+        cancel: &CancellationToken,
+        ctx: &RequestContext,
+    ) {
         let lsn = self.get_last_record_lsn();
 
         // imitiate on-restart initial logical size
         let size = self
-            .calculate_logical_size(lsn, LogicalSizeCalculationCause::EvictionTaskImitation, ctx)
+            .calculate_logical_size(
+                lsn,
+                LogicalSizeCalculationCause::EvictionTaskImitation,
+                cancel.clone(),
+                ctx,
+            )
             .instrument(info_span!("calculate_logical_size"))
             .await;
 

pageserver/src/walredo.rs

@@ -44,6 +44,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use crate::config::PageServerConf;
 use crate::metrics::{
     WAL_REDO_BYTES_HISTOGRAM, WAL_REDO_RECORDS_HISTOGRAM, WAL_REDO_RECORD_COUNTER, WAL_REDO_TIME,
+    WAL_REDO_WAIT_TIME,
 };
 use crate::pgdatadir_mapping::{key_to_rel_block, key_to_slru_block};
 use crate::repository::Key;
@@ -206,8 +207,11 @@ impl PostgresRedoManager {
     ) -> anyhow::Result<Bytes> {
         let (rel, blknum) = key_to_rel_block(key).context("invalid record")?;
         const MAX_RETRY_ATTEMPTS: u32 = 1;
+        let start_time = Instant::now();
         let mut n_attempts = 0u32;
         loop {
+            let lock_time = Instant::now();
+
             // launch the WAL redo process on first use
             let proc: Arc<WalRedoProcess> = {
                 let proc_guard = self.redo_process.read().unwrap();
@@ -232,7 +236,7 @@ impl PostgresRedoManager {
                 }
             };
 
-            let started_at = std::time::Instant::now();
+            WAL_REDO_WAIT_TIME.observe(lock_time.duration_since(start_time).as_secs_f64());
 
             // Relational WAL records are applied using wal-redo-postgres
             let buf_tag = BufferTag { rel, blknum };
@@ -240,7 +244,8 @@ impl PostgresRedoManager {
                 .apply_wal_records(buf_tag, &base_img, records, wal_redo_timeout)
                 .context("apply_wal_records");
 
-            let duration = started_at.elapsed();
+            let end_time = Instant::now();
+            let duration = end_time.duration_since(lock_time);
 
             let len = records.len();
             let nbytes = records.iter().fold(0, |acumulator, record| {

proxy/src/bin/proxy.rs

@@ -80,9 +80,6 @@ struct ProxyCliArgs {
     /// cache for `wake_compute` api method (use `size=0` to disable)
     #[clap(long, default_value = config::CacheOptions::DEFAULT_OPTIONS_NODE_INFO)]
     wake_compute_cache: String,
-    /// lock for `wake_compute` api method. example: "shards=32,permits=4,epoch=10m,timeout=1s". (use `permits=0` to disable).
-    #[clap(long, default_value = config::WakeComputeLockOptions::DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK)]
-    wake_compute_lock: String,
     /// Allow self-signed certificates for compute nodes (for testing)
     #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
     allow_self_signed_compute: bool,
@@ -223,23 +220,10 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
                 node_info: console::caches::NodeInfoCache::new("node_info_cache", size, ttl),
             }));
 
-            let config::WakeComputeLockOptions {
-                shards,
-                permits,
-                epoch,
-                timeout,
-            } = args.wake_compute_lock.parse()?;
-            info!(permits, shards, ?epoch, "Using NodeLocks (wake_compute)");
-            let locks = Box::leak(Box::new(
-                console::locks::ApiLocks::new("wake_compute_lock", permits, shards, timeout)
-                    .unwrap(),
-            ));
-            tokio::spawn(locks.garbage_collect_worker(epoch));
-
             let url = args.auth_endpoint.parse()?;
             let endpoint = http::Endpoint::new(url, http::new_client());
 
-            let api = console::provider::neon::Api::new(endpoint, caches, locks);
+            let api = console::provider::neon::Api::new(endpoint, caches);
             auth::BackendType::Console(Cow::Owned(api), ())
         }
         AuthBackend::Postgres => {

proxy/src/config.rs

@@ -264,79 +264,6 @@ impl FromStr for CacheOptions {
     }
 }
 
-/// Helper for cmdline cache options parsing.
-pub struct WakeComputeLockOptions {
-    /// The number of shards the lock map should have
-    pub shards: usize,
-    /// The number of allowed concurrent requests for each endpoitn
-    pub permits: usize,
-    /// Garbage collection epoch
-    pub epoch: Duration,
-    /// Lock timeout
-    pub timeout: Duration,
-}
-
-impl WakeComputeLockOptions {
-    /// Default options for [`crate::console::provider::ApiLocks`].
-    pub const DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK: &'static str = "permits=0";
-
-    // pub const DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK: &'static str = "shards=32,permits=4,epoch=10m,timeout=1s";
-
-    /// Parse lock options passed via cmdline.
-    /// Example: [`Self::DEFAULT_OPTIONS_WAKE_COMPUTE_LOCK`].
-    fn parse(options: &str) -> anyhow::Result<Self> {
-        let mut shards = None;
-        let mut permits = None;
-        let mut epoch = None;
-        let mut timeout = None;
-
-        for option in options.split(',') {
-            let (key, value) = option
-                .split_once('=')
-                .with_context(|| format!("bad key-value pair: {option}"))?;
-
-            match key {
-                "shards" => shards = Some(value.parse()?),
-                "permits" => permits = Some(value.parse()?),
-                "epoch" => epoch = Some(humantime::parse_duration(value)?),
-                "timeout" => timeout = Some(humantime::parse_duration(value)?),
-                unknown => bail!("unknown key: {unknown}"),
-            }
-        }
-
-        // these dont matter if lock is disabled
-        if let Some(0) = permits {
-            timeout = Some(Duration::default());
-            epoch = Some(Duration::default());
-            shards = Some(2);
-        }
-
-        let out = Self {
-            shards: shards.context("missing `shards`")?,
-            permits: permits.context("missing `permits`")?,
-            epoch: epoch.context("missing `epoch`")?,
-            timeout: timeout.context("missing `timeout`")?,
-        };
-
-        ensure!(out.shards > 1, "shard count must be > 1");
-        ensure!(
-            out.shards.is_power_of_two(),
-            "shard count must be a power of two"
-        );
-
-        Ok(out)
-    }
-}
-
-impl FromStr for WakeComputeLockOptions {
-    type Err = anyhow::Error;
-
-    fn from_str(options: &str) -> Result<Self, Self::Err> {
-        let error = || format!("failed to parse cache lock options '{options}'");
-        Self::parse(options).with_context(error)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -361,42 +288,4 @@ mod tests {
 
         Ok(())
     }
-
-    #[test]
-    fn test_parse_lock_options() -> anyhow::Result<()> {
-        let WakeComputeLockOptions {
-            epoch,
-            permits,
-            shards,
-            timeout,
-        } = "shards=32,permits=4,epoch=10m,timeout=1s".parse()?;
-        assert_eq!(epoch, Duration::from_secs(10 * 60));
-        assert_eq!(timeout, Duration::from_secs(1));
-        assert_eq!(shards, 32);
-        assert_eq!(permits, 4);
-
-        let WakeComputeLockOptions {
-            epoch,
-            permits,
-            shards,
-            timeout,
-        } = "epoch=60s,shards=16,timeout=100ms,permits=8".parse()?;
-        assert_eq!(epoch, Duration::from_secs(60));
-        assert_eq!(timeout, Duration::from_millis(100));
-        assert_eq!(shards, 16);
-        assert_eq!(permits, 8);
-
-        let WakeComputeLockOptions {
-            epoch,
-            permits,
-            shards,
-            timeout,
-        } = "permits=0".parse()?;
-        assert_eq!(epoch, Duration::ZERO);
-        assert_eq!(timeout, Duration::ZERO);
-        assert_eq!(shards, 2);
-        assert_eq!(permits, 0);
-
-        Ok(())
-    }
 }

proxy/src/console.rs

@@ -13,10 +13,5 @@ pub mod caches {
     pub use super::provider::{ApiCaches, NodeInfoCache};
 }
 
-/// Various cache-related types.
-pub mod locks {
-    pub use super::provider::ApiLocks;
-}
-
 /// Console's management API.
 pub mod mgmt;

proxy/src/console/provider.rs

@@ -8,13 +8,7 @@ use crate::{
     compute, scram,
 };
 use async_trait::async_trait;
-use dashmap::DashMap;
-use std::{sync::Arc, time::Duration};
-use tokio::{
-    sync::{OwnedSemaphorePermit, Semaphore},
-    time::Instant,
-};
-use tracing::info;
+use std::sync::Arc;
 
 pub mod errors {
     use crate::{
@@ -155,9 +149,6 @@ pub mod errors {
 
         #[error(transparent)]
         ApiError(ApiError),
-
-        #[error("Timeout waiting to acquire wake compute lock")]
-        TimeoutError,
     }
 
     // This allows more useful interactions than `#[from]`.
@@ -167,17 +158,6 @@ pub mod errors {
         }
     }
 
-    impl From<tokio::sync::AcquireError> for WakeComputeError {
-        fn from(_: tokio::sync::AcquireError) -> Self {
-            WakeComputeError::TimeoutError
-        }
-    }
-    impl From<tokio::time::error::Elapsed> for WakeComputeError {
-        fn from(_: tokio::time::error::Elapsed) -> Self {
-            WakeComputeError::TimeoutError
-        }
-    }
-
     impl UserFacingError for WakeComputeError {
         fn to_string_client(&self) -> String {
             use WakeComputeError::*;
@@ -187,8 +167,6 @@ pub mod errors {
                 BadComputeAddress(_) => REQUEST_FAILED.to_owned(),
                 // However, API might return a meaningful error.
                 ApiError(e) => e.to_string_client(),
-
-                TimeoutError => "timeout while acquiring the compute resource lock".to_owned(),
             }
         }
     }
@@ -255,145 +233,3 @@ pub struct ApiCaches {
     /// Cache for the `wake_compute` API method.
     pub node_info: NodeInfoCache,
 }
-
-/// Various caches for [`console`](super).
-pub struct ApiLocks {
-    name: &'static str,
-    node_locks: DashMap<Arc<str>, Arc<Semaphore>>,
-    permits: usize,
-    timeout: Duration,
-    registered: prometheus::IntCounter,
-    unregistered: prometheus::IntCounter,
-    reclamation_lag: prometheus::Histogram,
-    lock_acquire_lag: prometheus::Histogram,
-}
-
-impl ApiLocks {
-    pub fn new(
-        name: &'static str,
-        permits: usize,
-        shards: usize,
-        timeout: Duration,
-    ) -> prometheus::Result<Self> {
-        let registered = prometheus::IntCounter::with_opts(
-            prometheus::Opts::new(
-                "semaphores_registered",
-                "Number of semaphores registered in this api lock",
-            )
-            .namespace(name),
-        )?;
-        prometheus::register(Box::new(registered.clone()))?;
-        let unregistered = prometheus::IntCounter::with_opts(
-            prometheus::Opts::new(
-                "semaphores_unregistered",
-                "Number of semaphores unregistered in this api lock",
-            )
-            .namespace(name),
-        )?;
-        prometheus::register(Box::new(unregistered.clone()))?;
-        let reclamation_lag = prometheus::Histogram::with_opts(
-            prometheus::HistogramOpts::new(
-                "reclamation_lag_seconds",
-                "Time it takes to reclaim unused semaphores in the api lock",
-            )
-            .namespace(name)
-            // 1us -> 65ms
-            // benchmarks on my mac indicate it's usually in the range of 256us and 512us
-            .buckets(prometheus::exponential_buckets(1e-6, 2.0, 16)?),
-        )?;
-        prometheus::register(Box::new(reclamation_lag.clone()))?;
-        let lock_acquire_lag = prometheus::Histogram::with_opts(
-            prometheus::HistogramOpts::new(
-                "semaphore_acquire_seconds",
-                "Time it takes to reclaim unused semaphores in the api lock",
-            )
-            .namespace(name)
-            // 0.1ms -> 6s
-            .buckets(prometheus::exponential_buckets(1e-4, 2.0, 16)?),
-        )?;
-        prometheus::register(Box::new(lock_acquire_lag.clone()))?;
-
-        Ok(Self {
-            name,
-            node_locks: DashMap::with_shard_amount(shards),
-            permits,
-            timeout,
-            lock_acquire_lag,
-            registered,
-            unregistered,
-            reclamation_lag,
-        })
-    }
-
-    pub async fn get_wake_compute_permit(
-        &self,
-        key: &Arc<str>,
-    ) -> Result<WakeComputePermit, errors::WakeComputeError> {
-        if self.permits == 0 {
-            return Ok(WakeComputePermit { permit: None });
-        }
-        let now = Instant::now();
-        let semaphore = {
-            // get fast path
-            if let Some(semaphore) = self.node_locks.get(key) {
-                semaphore.clone()
-            } else {
-                self.node_locks
-                    .entry(key.clone())
-                    .or_insert_with(|| {
-                        self.registered.inc();
-                        Arc::new(Semaphore::new(self.permits))
-                    })
-                    .clone()
-            }
-        };
-        let permit = tokio::time::timeout_at(now + self.timeout, semaphore.acquire_owned()).await;
-
-        self.lock_acquire_lag
-            .observe((Instant::now() - now).as_secs_f64());
-
-        Ok(WakeComputePermit {
-            permit: Some(permit??),
-        })
-    }
-
-    pub async fn garbage_collect_worker(&self, epoch: std::time::Duration) {
-        if self.permits == 0 {
-            return;
-        }
-
-        let mut interval = tokio::time::interval(epoch / (self.node_locks.shards().len()) as u32);
-        loop {
-            for (i, shard) in self.node_locks.shards().iter().enumerate() {
-                interval.tick().await;
-                // temporary lock a single shard and then clear any semaphores that aren't currently checked out
-                // race conditions: if strong_count == 1, there's no way that it can increase while the shard is locked
-                // therefore releasing it is safe from race conditions
-                info!(
-                    name = self.name,
-                    shard = i,
-                    "performing epoch reclamation on api lock"
-                );
-                let mut lock = shard.write();
-                let timer = self.reclamation_lag.start_timer();
-                let count = lock
-                    .extract_if(|_, semaphore| Arc::strong_count(semaphore.get_mut()) == 1)
-                    .count();
-                drop(lock);
-                self.unregistered.inc_by(count as u64);
-                timer.observe_duration()
-            }
-        }
-    }
-}
-
-pub struct WakeComputePermit {
-    // None if the lock is disabled
-    permit: Option<OwnedSemaphorePermit>,
-}
-
-impl WakeComputePermit {
-    pub fn should_check_cache(&self) -> bool {
-        self.permit.is_some()
-    }
-}

proxy/src/console/provider/neon.rs

@@ -3,12 +3,12 @@
 use super::{
     super::messages::{ConsoleError, GetRoleSecret, WakeCompute},
     errors::{ApiError, GetAuthInfoError, WakeComputeError},
-    ApiCaches, ApiLocks, AuthInfo, CachedNodeInfo, ConsoleReqExtra, NodeInfo,
+    ApiCaches, AuthInfo, CachedNodeInfo, ConsoleReqExtra, NodeInfo,
 };
 use crate::{auth::ClientCredentials, compute, http, scram};
 use async_trait::async_trait;
 use futures::TryFutureExt;
-use std::{net::SocketAddr, sync::Arc};
+use std::net::SocketAddr;
 use tokio::time::Instant;
 use tokio_postgres::config::SslMode;
 use tracing::{error, info, info_span, warn, Instrument};
@@ -17,17 +17,12 @@ use tracing::{error, info, info_span, warn, Instrument};
 pub struct Api {
     endpoint: http::Endpoint,
     caches: &'static ApiCaches,
-    locks: &'static ApiLocks,
     jwt: String,
 }
 
 impl Api {
     /// Construct an API object containing the auth parameters.
-    pub fn new(
-        endpoint: http::Endpoint,
-        caches: &'static ApiCaches,
-        locks: &'static ApiLocks,
-    ) -> Self {
+    pub fn new(endpoint: http::Endpoint, caches: &'static ApiCaches) -> Self {
         let jwt: String = match std::env::var("NEON_PROXY_TO_CONTROLPLANE_TOKEN") {
             Ok(v) => v,
             Err(_) => "".to_string(),
@@ -35,7 +30,6 @@ impl Api {
         Self {
             endpoint,
             caches,
-            locks,
             jwt,
         }
     }
@@ -169,22 +163,9 @@ impl super::Api for Api {
             return Ok(cached);
         }
 
-        let key: Arc<str> = key.into();
-
-        let permit = self.locks.get_wake_compute_permit(&key).await?;
-
-        // after getting back a permit - it's possible the cache was filled
-        // double check
-        if permit.should_check_cache() {
-            if let Some(cached) = self.caches.node_info.get(&key) {
-                info!(key = &*key, "found cached compute node info");
-                return Ok(cached);
-            }
-        }
-
         let node = self.do_wake_compute(extra, creds).await?;
-        let (_, cached) = self.caches.node_info.insert(key.clone(), node);
-        info!(key = &*key, "created a cache entry for compute node info");
+        let (_, cached) = self.caches.node_info.insert(key.into(), node);
+        info!(key = key, "created a cache entry for compute node info");
 
         Ok(cached)
     }

proxy/src/proxy.rs

@@ -570,7 +570,6 @@ fn report_error(e: &WakeComputeError, retry: bool) {
             "api_console_other_server_error"
         }
         WakeComputeError::ApiError(ApiError::Console { .. }) => "api_console_other_error",
-        WakeComputeError::TimeoutError => "timeout_error",
     };
     NUM_WAKEUP_FAILURES.with_label_values(&[retry, kind]).inc();
 }

safekeeper/src/auth.rs

@@ -1,20 +1,19 @@
-use utils::auth::{AuthError, Claims, Scope};
+use anyhow::{bail, Result};
+use utils::auth::{Claims, Scope};
 use utils::id::TenantId;
 
-pub fn check_permission(claims: &Claims, tenant_id: Option<TenantId>) -> Result<(), AuthError> {
+pub fn check_permission(claims: &Claims, tenant_id: Option<TenantId>) -> Result<()> {
     match (&claims.scope, tenant_id) {
-        (Scope::Tenant, None) => Err(AuthError(
-            "Attempt to access management api with tenant scope. Permission denied".into(),
-        )),
+        (Scope::Tenant, None) => {
+            bail!("Attempt to access management api with tenant scope. Permission denied")
+        }
         (Scope::Tenant, Some(tenant_id)) => {
             if claims.tenant_id.unwrap() != tenant_id {
-                return Err(AuthError("Tenant id mismatch. Permission denied".into()));
+                bail!("Tenant id mismatch. Permission denied")
             }
             Ok(())
         }
-        (Scope::PageServerApi, _) => Err(AuthError(
-            "PageServerApi scope makes no sense for Safekeeper".into(),
-        )),
+        (Scope::PageServerApi, _) => bail!("PageServerApi scope makes no sense for Safekeeper"),
         (Scope::SafekeeperData, _) => Ok(()),
     }
 }

safekeeper/src/handler.rs

@@ -6,7 +6,7 @@ use std::str::FromStr;
 use std::str::{self};
 use std::sync::Arc;
 use tokio::io::{AsyncRead, AsyncWrite};
-use tracing::{debug, info, info_span, Instrument};
+use tracing::{info, info_span, Instrument};
 
 use crate::auth::check_permission;
 use crate::json_ctrl::{handle_json_ctrl, AppendLogicalMessage};
@@ -165,27 +165,26 @@ impl<IO: AsyncRead + AsyncWrite + Unpin + Send> postgres_backend::Handler<IO>
             .auth
             .as_ref()
             .expect("auth_type is configured but .auth of handler is missing");
-        let data = auth
-            .decode(str::from_utf8(jwt_response).context("jwt response is not UTF-8")?)
-            .map_err(|e| QueryError::Unauthorized(e.0))?;
+        let data =
+            auth.decode(str::from_utf8(jwt_response).context("jwt response is not UTF-8")?)?;
 
         // The handler might be configured to allow only tenant scope tokens.
         if matches!(allowed_auth_scope, Scope::Tenant)
             && !matches!(data.claims.scope, Scope::Tenant)
         {
-            return Err(QueryError::Unauthorized(
-                "passed JWT token is for full access, but only tenant scope is allowed".into(),
-            ));
+            return Err(QueryError::Other(anyhow::anyhow!(
+                "passed JWT token is for full access, but only tenant scope is allowed"
+            )));
         }
 
         if matches!(data.claims.scope, Scope::Tenant) && data.claims.tenant_id.is_none() {
-            return Err(QueryError::Unauthorized(
-                "jwt token scope is Tenant, but tenant id is missing".into(),
-            ));
+            return Err(QueryError::Other(anyhow::anyhow!(
+                "jwt token scope is Tenant, but tenant id is missing"
+            )));
         }
 
-        debug!(
-            "jwt scope check succeeded for scope: {:#?} by tenant id: {:?}",
+        info!(
+            "jwt auth succeeded for scope: {:#?} by tenant id: {:?}",
             data.claims.scope, data.claims.tenant_id,
         );
 
@@ -264,7 +263,7 @@ impl SafekeeperPostgresHandler {
 
     // when accessing management api supply None as an argument
     // when using to authorize tenant pass corresponding tenant id
-    fn check_permission(&self, tenant_id: Option<TenantId>) -> Result<(), QueryError> {
+    fn check_permission(&self, tenant_id: Option<TenantId>) -> anyhow::Result<()> {
         if self.auth.is_none() {
             // auth is set to Trust, nothing to check so just return ok
             return Ok(());
@@ -276,7 +275,7 @@ impl SafekeeperPostgresHandler {
             .claims
             .as_ref()
             .expect("claims presence already checked");
-        check_permission(claims, tenant_id).map_err(|e| QueryError::Unauthorized(e.0))
+        check_permission(claims, tenant_id)
     }
 
     async fn handle_timeline_status<IO: AsyncRead + AsyncWrite + Unpin>(

test_runner/regress/test_auth.py

@@ -25,7 +25,7 @@ def assert_client_authorized(env: NeonEnv, http_client: PageserverHttpClient):
 def assert_client_not_authorized(env: NeonEnv, http_client: PageserverHttpClient):
     with pytest.raises(
         PageserverApiException,
-        match="Forbidden: JWT authentication error",
+        match="Unauthorized: malformed jwt token",
     ):
         assert_client_authorized(env, http_client)
 
@@ -56,7 +56,9 @@ def test_pageserver_auth(neon_env_builder: NeonEnvBuilder):
     assert_client_authorized(env, pageserver_http_client)
 
     # fail to create branch using token with different tenant_id
-    with pytest.raises(PageserverApiException, match="Forbidden: JWT authentication error"):
+    with pytest.raises(
+        PageserverApiException, match="Forbidden: Tenant id mismatch. Permission denied"
+    ):
         assert_client_authorized(env, invalid_tenant_http_client)
 
     # create tenant using management token
@@ -65,7 +67,7 @@ def test_pageserver_auth(neon_env_builder: NeonEnvBuilder):
     # fail to create tenant using tenant token
     with pytest.raises(
         PageserverApiException,
-        match="Forbidden: JWT authentication error",
+        match="Forbidden: Attempt to access management api with tenant scope. Permission denied",
     ):
         tenant_http_client.tenant_create(TenantId.generate())
 
@@ -92,7 +94,6 @@ def test_compute_auth_to_pageserver(neon_env_builder: NeonEnvBuilder):
 def test_pageserver_multiple_keys(neon_env_builder: NeonEnvBuilder):
     neon_env_builder.auth_enabled = True
     env = neon_env_builder.init_start()
-    env.pageserver.allowed_errors.append(".*Authentication error: InvalidSignature.*")
     env.pageserver.allowed_errors.append(".*Unauthorized: malformed jwt token.*")
 
     pageserver_token_old = env.auth_keys.generate_pageserver_token()
@@ -145,7 +146,6 @@ def test_pageserver_multiple_keys(neon_env_builder: NeonEnvBuilder):
 def test_pageserver_key_reload(neon_env_builder: NeonEnvBuilder):
     neon_env_builder.auth_enabled = True
     env = neon_env_builder.init_start()
-    env.pageserver.allowed_errors.append(".*Authentication error: InvalidSignature.*")
     env.pageserver.allowed_errors.append(".*Unauthorized: malformed jwt token.*")
 
     pageserver_token_old = env.auth_keys.generate_pageserver_token()
@@ -162,7 +162,7 @@ def test_pageserver_key_reload(neon_env_builder: NeonEnvBuilder):
     # Next attempt fails as we use the old auth token
     with pytest.raises(
         PageserverApiException,
-        match="Forbidden: JWT authentication error",
+        match="Unauthorized: malformed jwt token",
     ):
         pageserver_http_client_old.reload_auth_validation_keys()
 

test_runner/regress/test_bad_connection.py

@@ -1,60 +0,0 @@
-import random
-import time
-
-from fixtures.log_helper import log
-from fixtures.neon_fixtures import NeonEnvBuilder
-
-
-def test_compute_pageserver_connection_stress(neon_env_builder: NeonEnvBuilder):
-    env = neon_env_builder.init_start()
-    env.pageserver.allowed_errors.append(".*simulated connection error.*")
-
-    pageserver_http = env.pageserver.http_client()
-    env.neon_cli.create_branch("test_compute_pageserver_connection_stress")
-    endpoint = env.endpoints.create_start("test_compute_pageserver_connection_stress")
-
-    # Enable failpoint after starting everything else up so that loading initial
-    # basebackup doesn't fail
-    pageserver_http.configure_failpoints(("simulated-bad-compute-connection", "50%return(15)"))
-
-    pg_conn = endpoint.connect()
-    cur = pg_conn.cursor()
-
-    # Create table, and insert some rows. Make it big enough that it doesn't fit in
-    # shared_buffers, otherwise the SELECT after restart will just return answer
-    # from shared_buffers without hitting the page server, which defeats the point
-    # of this test.
-    cur.execute("CREATE TABLE foo (t text)")
-    cur.execute(
-        """
-        INSERT INTO foo
-            SELECT 'long string to consume some space' || g
-            FROM generate_series(1, 100000) g
-        """
-    )
-
-    # Verify that the table is larger than shared_buffers
-    cur.execute(
-        """
-        select setting::int * pg_size_bytes(unit) as shared_buffers, pg_relation_size('foo') as tbl_size
-        from pg_settings where name = 'shared_buffers'
-        """
-    )
-    row = cur.fetchone()
-    assert row is not None
-    log.info(f"shared_buffers is {row[0]}, table size {row[1]}")
-    assert int(row[0]) < int(row[1])
-
-    cur.execute("SELECT count(*) FROM foo")
-    assert cur.fetchone() == (100000,)
-
-    end_time = time.time() + 30
-    times_executed = 0
-    while time.time() < end_time:
-        if random.random() < 0.5:
-            cur.execute("INSERT INTO foo VALUES ('stas'), ('heikki')")
-        else:
-            cur.execute("SELECT t FROM foo ORDER BY RANDOM() LIMIT 10")
-            cur.fetchall()
-        times_executed += 1
-    log.info(f"Workload executed {times_executed} times")

test_runner/regress/test_broken_timeline.py

@@ -26,7 +26,6 @@ def test_local_corruption(neon_env_builder: NeonEnvBuilder):
             ".*will not become active. Current state: Broken.*",
             ".*failed to load metadata.*",
             ".*load failed.*load local timeline.*",
-            ".*layer loading failed permanently: load layer: .*",
         ]
     )
 

workspace_hack/Cargo.toml

@@ -25,7 +25,6 @@ chrono = { version = "0.4", default-features = false, features = ["clock", "serd
 clap = { version = "4", features = ["derive", "string"] }
 clap_builder = { version = "4", default-features = false, features = ["color", "help", "std", "string", "suggestions", "usage"] }
 crossbeam-utils = { version = "0.8" }
-dashmap = { version = "5", default-features = false, features = ["raw-api"] }
 either = { version = "1" }
 fail = { version = "0.5", default-features = false, features = ["failpoints"] }
 futures = { version = "0.3" }